vibe-core/tests/0-tcpproxy.d

/+ dub.sdl:
	name "tests"
	description "TCP proxy test"
	copyright "Copyright © 2015-2020, Sönke Ludwig"
	dependency "vibe-core" path=".."
+/
module tests;

import vibe.core.core;
import vibe.core.log;
import vibe.core.net;
import vibe.core.stream;
import std.exception;
import std.string;
import std.range.primitives : front;
import core.time;


void testProtocol(TCPConnection server, bool terminate)
{
	foreach (i; 0 .. 1) {
		foreach (j; 0 .. 1) {
			auto str = format("Hello, World #%s", i*100+j);
			server.write(str);
			server.write("\n");
			auto reply = server.readLine();
			assert(reply == format("Hash: %08X", typeid(string).getHash(&str)), "Unexpected reply");
		}
		sleep(10.msecs);
	}

	assert(!server.dataAvailableForRead, "Still data available.");

	if (terminate) {
		// forcefully close connection
		server.close();
	} else {
		server.write("quit\n");
		enforce(server.readLine() == "Bye bye!");
		// should have closed within 500 ms
		enforce(!server.waitForData(500.msecs));
		assert(server.empty, "Server still connected.");
	}
}

void runTest()
{
	import std.algorithm : find;
	import std.socket : AddressFamily;

	// server for a simple line based protocol
	auto l1 = listenTCP(0, (client) {
		while (!client.empty) {
			auto ln = client.readLine();
			if (ln == "quit") {
				client.write("Bye bye!\n");
				client.close();
				break;
			}

			client.write(format("Hash: %08X\n", typeid(string).getHash(&ln)));
		}
	}, "127.0.0.1");
	scope (exit) l1.stopListening;

	// proxy server
	auto l2 = listenTCP(0, (client) {
		auto server = connectTCP(l1.bindAddress);

		// pipe server to client as long as the server connection is alive
		auto t = runTask!(TCPConnection, TCPConnection)((client, server) {
			scope (exit) client.close();
			pipe(server, client);
			logInfo("Proxy 2 out");
		}, client, server);

		// pipe client to server as long as the client connection is alive
		scope (exit) {
			server.close();
			t.join();
		}
		pipe(client, server);
		logInfo("Proxy out");
	}, "127.0.0.1");
	scope (exit) l2.stopListening;

	// test server
	logInfo("Test protocol implementation on server");
	testProtocol(connectTCP(l2.bindAddress), false);
	logInfo("Test protocol implementation on server with forced disconnect");
	testProtocol(connectTCP(l2.bindAddress), true);

	// test proxy
	logInfo("Test protocol implementation on proxy");
	testProtocol(connectTCP(l2.bindAddress), false);
	logInfo("Test protocol implementation on proxy with forced disconnect");
	testProtocol(connectTCP(l2.bindAddress), true);
}

int main()
{
	int ret = 0;
	runTask({
		try runTest();
		catch (Throwable th) {
			logError("Test failed: %s", th.msg);
			logDiagnostic("Full error: %s", th);
			ret = 1;
		} finally exitEventLoop(true);
	});
	runEventLoop();
	return ret;
}

string readLine(TCPConnection c)
{
	import std.string : indexOf;

	string ret;
	while (!c.empty) {
		auto buf = cast(char[])c.peek();
		auto idx = buf.indexOf('\n');
		if (idx < 0) {
			ret ~= buf;
			c.skip(buf.length);
		} else {
			ret ~= buf[0 .. idx];
			c.skip(idx+1);
			break;
		}
	}
	return ret;
}