#include <gnutls/gnutls.h>
#include <gnutls/dtls.h>
#include <unistd.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <netinet/in.h>
#include <fcntl.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <errno.h>
#include <poll.h>
#include <time.h>
#include <wait.h>

// {{{ types

typedef struct {
	int count;
} filter_packet_state_t;

typedef struct {
	gnutls_datum_t packets[3];
	int* order;
	int count;
} filter_permute_state_t;

typedef void (*filter_fn)(gnutls_transport_ptr_t, const unsigned char*, size_t);

typedef int (*match_fn)(const unsigned char*, size_t);

enum role { SERVER, CLIENT };

// }}}

// {{{ static data

static int permutations2[2][2]
	= { { 0, 1 }, { 1, 0 } };

static const char* permutation_names2[]
	= { "01", "10", 0 };

static int permutations3[6][3]
	= { { 0, 1, 2 }, { 0, 2, 1 },
		{ 1, 0, 2 }, { 1, 2, 0 },
		{ 2, 0, 1 }, { 2, 1, 0 } };

static const char* permutation_names3[]
	= { "012", "021", "102", "120", "201", "210", 0 };

static const char* filter_names[8]
	= { "SHello",
		"SKeyExchange",
		"SHelloDone",
		"CKeyExchange",
		"CChangeCipherSpec",
		"CFinished",
		"SChangeCipherSpec",
		"SFinished" };

// }}}

// {{{ other global state

enum role role;

#define role_name (role == SERVER ? "server" : "client")

int debug;
int nonblock;

int run_id;

// }}}

// {{{ logging and error handling

static void logfn(int level, const char* s)
{
	if (debug) {
		fprintf(stdout, "%i %s|<%i> %s", run_id, role_name, level, s);
	}
}

static void auditfn(gnutls_session_t session, const char* s)
{
	if (debug) {
		fprintf(stdout, "%i %s| %s", run_id, role_name, s);
	}
}

static void drop(const char* packet)
{
	if (debug) {
		fprintf(stdout, "%i %s| dropping %s\n", run_id, role_name, packet);
	}
}

static int _process_error(int loc, int code, int die)
{
	if (code < 0 && (die || code != GNUTLS_E_AGAIN)) {
		fprintf(stdout, "%i <%s tls> line %i: %s", run_id, role_name, loc, gnutls_strerror(code));
		if (gnutls_error_is_fatal(code) || die) {
			fprintf(stdout, " (fatal)\n");
			exit(1);
		} else {
			fprintf(stdout, "\n");
		}
	}
	return code;
}

#define die_on_error(code) _process_error(__LINE__, code, 1)
#define process_error(code) _process_error(__LINE__, code, 0)

static void _process_error_or_timeout(int loc, int err, time_t tdiff)
{
	if (err < 0) {
		if (err != GNUTLS_E_TIMEDOUT || tdiff >= 60) {
			_process_error(loc, err, 0);
		} else {
			fprintf(stdout, "%i %s| line %i: {spurious timeout} (fatal)", run_id, role_name, loc);
			exit(1);
		}
	}
}

#define process_error_or_timeout(code, tdiff) _process_error_or_timeout(__LINE__, code, tdiff)

// }}}

// {{{ init, shared, and teardown code and data for packet stream filters

filter_packet_state_t state_packet_ServerHello = { 0 };
filter_packet_state_t state_packet_ServerKeyExchange = { 0 };
filter_packet_state_t state_packet_ServerHelloDone = { 0 };
filter_packet_state_t state_packet_ClientKeyExchange = { 0 };
filter_packet_state_t state_packet_ClientChangeCipherSpec = { 0 };
filter_packet_state_t state_packet_ClientFinished = { 0 };
filter_packet_state_t state_packet_ServerChangeCipherSpec = { 0 };
filter_packet_state_t state_packet_ServerFinished = { 0 };

filter_permute_state_t state_permute_ServerHello = { { { 0, 0 }, { 0, 0 }, { 0, 0 } }, 0, 0 };
filter_permute_state_t state_permute_ServerFinished = { { { 0, 0 }, { 0, 0 }, { 0, 0 } }, 0, 0 };
filter_permute_state_t state_permute_ClientFinished = { { { 0, 0 }, { 0, 0 }, { 0, 0 } }, 0, 0 };

filter_fn filter_chain[32];
int filter_current_idx;

static void filter_permute_state_free_buffer(filter_permute_state_t* state)
{
	int i;

	for (i = 0; i < 3; i++) {
		free(state->packets[i].data);
		state->packets[i].data = NULL;
	}
}

static void filter_clear_state()
{
	filter_current_idx = 0;

	filter_permute_state_free_buffer(&state_permute_ServerHello);
	filter_permute_state_free_buffer(&state_permute_ServerFinished);
	filter_permute_state_free_buffer(&state_permute_ClientFinished);

	memset(&state_packet_ServerHello, 0, sizeof(state_packet_ServerHello));
	memset(&state_packet_ServerKeyExchange, 0, sizeof(state_packet_ServerKeyExchange));
	memset(&state_packet_ServerHelloDone, 0, sizeof(state_packet_ServerHelloDone));
	memset(&state_packet_ClientKeyExchange, 0, sizeof(state_packet_ClientKeyExchange));
	memset(&state_packet_ClientChangeCipherSpec, 0, sizeof(state_packet_ClientChangeCipherSpec));
	memset(&state_packet_ClientFinished, 0, sizeof(state_packet_ClientFinished));
	memset(&state_packet_ServerChangeCipherSpec, 0, sizeof(state_packet_ServerChangeCipherSpec));
	memset(&state_packet_ServerFinished, 0, sizeof(state_packet_ServerFinished));
	memset(&state_permute_ServerHello, 0, sizeof(state_permute_ServerHello));
	memset(&state_permute_ServerFinished, 0, sizeof(state_permute_ServerFinished));
	memset(&state_permute_ClientFinished, 0, sizeof(state_permute_ClientFinished));
}

static void filter_run_next(gnutls_transport_ptr_t fd,
		const unsigned char* buffer, size_t len)
{
	filter_fn fn = filter_chain[filter_current_idx];
	filter_current_idx++;
	if (fn) {
		fn(fd, buffer, len);
	} else {
		send((intptr_t) fd, buffer, len, 0);
	}
	filter_current_idx--;
}

// }}}

// {{{ packet match functions

static int match_ServerHello(const unsigned char* buffer, size_t len)
{
	return role == SERVER && len >= 13 + 1 && buffer[0] == 22 && buffer[13] == 2;
}

static int match_ServerKeyExchange(const unsigned char* buffer, size_t len)
{
	return role == SERVER && len >= 13 + 1 && buffer[0] == 22 && buffer[13] == 12;
}

static int match_ServerHelloDone(const unsigned char* buffer, size_t len)
{
	return role == SERVER && len >= 13 + 1 && buffer[0] == 22 && buffer[13] == 14;
}

static int match_ClientKeyExchange(const unsigned char* buffer, size_t len)
{
	return role == CLIENT && len >= 13 + 1 && buffer[0] == 22 && buffer[13] == 16;
}

static int match_ClientChangeCipherSpec(const unsigned char* buffer, size_t len)
{
	return role == CLIENT && len >= 13 && buffer[0] == 20;
}

static int match_ClientFinished(const unsigned char* buffer, size_t len)
{
	return role == CLIENT && len >= 13 && buffer[0] == 22 && buffer[4] == 1;
}

static int match_ServerChangeCipherSpec(const unsigned char* buffer, size_t len)
{
	return role == SERVER && len >= 13 && buffer[0] == 20;
}

static int match_ServerFinished(const unsigned char* buffer, size_t len)
{
	return role == SERVER && len >= 13 && buffer[0] == 22 && buffer[4] == 1;
}

// }}}

// {{{ packet drop filters

#define FILTER_DROP_COUNT 3
#define DECLARE_FILTER(packet) \
	static void filter_packet_##packet(gnutls_transport_ptr_t fd, \
			const unsigned char* buffer, size_t len) \
	{ \
		if (match_##packet(buffer, len) && (state_packet_##packet).count++ < FILTER_DROP_COUNT) { \
			drop(#packet); \
		} else { \
			filter_run_next(fd, buffer, len); \
		} \
	}

DECLARE_FILTER(ServerHello)
DECLARE_FILTER(ServerKeyExchange)
DECLARE_FILTER(ServerHelloDone)
DECLARE_FILTER(ClientKeyExchange)
DECLARE_FILTER(ClientChangeCipherSpec)
DECLARE_FILTER(ClientFinished)
DECLARE_FILTER(ServerChangeCipherSpec)
DECLARE_FILTER(ServerFinished)

// }}}

// {{{ flight permutation filters

static void filter_permute_state_run(filter_permute_state_t* state, int packetCount,
		gnutls_transport_ptr_t fd, const unsigned char* buffer, size_t len)
{
	unsigned char* data = malloc(len);
	int packet = state->order[state->count];

	memcpy(data, buffer, len);
	state->packets[packet].data = data;
	state->packets[packet].size = len;
	state->count++;

	if (state->count == packetCount) {
		for (packet = 0; packet < packetCount; packet++) {
			filter_run_next(fd, state->packets[packet].data,
					state->packets[packet].size);
		}
		filter_permute_state_free_buffer(state);
		state->count = 0;
	}
}

#define DECLARE_PERMUTE(flight) \
	static void filter_permute_##flight(gnutls_transport_ptr_t fd, \
			const unsigned char* buffer, size_t len) \
	{ \
		int count = sizeof(permute_match_##flight) / sizeof(permute_match_##flight[0]); \
		int i; \
		for (i = 0; i < count; i++) { \
			if (permute_match_##flight[i](buffer, len)) { \
				filter_permute_state_run(&state_permute_##flight, count, fd, buffer, len); \
				return; \
			} \
		} \
		filter_run_next(fd, buffer, len); \
	}

static match_fn permute_match_ServerHello[] = { match_ServerHello, match_ServerKeyExchange, match_ServerHelloDone };
static match_fn permute_match_ServerFinished[] = { match_ServerChangeCipherSpec, match_ServerFinished };
static match_fn permute_match_ClientFinished[] = { match_ClientKeyExchange, match_ClientChangeCipherSpec, match_ClientFinished };

DECLARE_PERMUTE(ServerHello)
DECLARE_PERMUTE(ServerFinished)
DECLARE_PERMUTE(ClientFinished)
	
// }}}

// {{{ emergency deadlock resolution time bomb

timer_t killtimer_tid = 0;

static void killtimer_set()
{
	struct sigevent sig;
	struct itimerspec tout = { { 0, 0 }, { 120, 0 } };

	if (killtimer_tid != 0) {
		timer_delete(killtimer_tid);
	}

	memset(&sig, 0, sizeof(sig));
	sig.sigev_notify = SIGEV_SIGNAL;
	sig.sigev_signo = 15;
	timer_create(CLOCK_MONOTONIC, &sig, &killtimer_tid);

	timer_settime(killtimer_tid, 0, &tout, 0);
}

// }}}

// {{{ actual gnutls operations

gnutls_session_t session;

static ssize_t writefn(gnutls_transport_ptr_t fd, const void* buffer, size_t len)
{
	filter_run_next(fd, (const unsigned char*) buffer, len);
	return len;
}

static void await(int fd)
{
	if (nonblock) {
		struct pollfd p = { fd, POLLIN, 0 };
		poll(&p, 1, 100);
	}
}

static void session_init(int sock, int server)
{
	gnutls_init(&session, GNUTLS_DATAGRAM | (server ? GNUTLS_SERVER : GNUTLS_CLIENT)
			| GNUTLS_NONBLOCK * nonblock);
	gnutls_priority_set_direct(session, "NORMAL:+ANON-ECDH", 0);
	gnutls_transport_set_ptr(session, (gnutls_transport_ptr_t) (intptr_t) sock);

	if (server) {
		gnutls_anon_server_credentials_t cred;
		gnutls_anon_allocate_server_credentials(&cred);
		gnutls_credentials_set(session, GNUTLS_CRD_ANON, cred);
	} else {
		gnutls_anon_client_credentials_t cred;
		gnutls_anon_allocate_client_credentials(&cred);
		gnutls_credentials_set(session, GNUTLS_CRD_ANON, cred);
	}

	gnutls_transport_set_push_function(session, writefn);

	gnutls_dtls_set_mtu(session, 1400);
}

static void client(int sock)
{
	session_init(sock, 0);
	int err = 0;
	time_t started = time(0);
	const char* line = "foobar!";
	char buffer[8192];
	int len;

	killtimer_set();
	do {
		await(sock);
		err = process_error(gnutls_handshake(session));
	} while (err != 0);
	process_error_or_timeout(err, time(0) - started);
	
	killtimer_set();
	die_on_error(gnutls_record_send(session, line, strlen(line)));
	
	do {
		await(sock);
		len = process_error(gnutls_record_recv(session, buffer, sizeof(buffer)));
	} while (len < 0);

	if (len > 0 && strcmp(line, buffer) == 0) {
		exit(0);
	} else {
		exit(1);
	}
}

static void server(int sock)
{ 
	session_init(sock, 1);
	int err;
	time_t started = time(0);
	char buffer[8192];
	int len;

	write(sock, &sock, 1);

	killtimer_set();
	do {
		await(sock);
		err = process_error(gnutls_handshake(session));
	} while (err != 0);
	process_error_or_timeout(err, time(0) - started);

	killtimer_set();
	do {
		await(sock);
		len = process_error(gnutls_record_recv(session, buffer, sizeof(buffer)));
	} while (len < 0);

	die_on_error(gnutls_record_send(session, buffer, len));
	exit(0);
}

// }}}

static void udp_sockpair(int* socks)
{
	struct sockaddr_in6 sa = { AF_INET6, htons(30000), 0, in6addr_loopback, 0 };
	struct sockaddr_in6 sb = { AF_INET6, htons(20000), 0, in6addr_loopback, 0 };

	socks[0] = socket(AF_INET6, SOCK_DGRAM, 0);
	socks[1] = socket(AF_INET6, SOCK_DGRAM, 0);

	bind(socks[0], (struct sockaddr*) &sa, sizeof(sa));
	bind(socks[1], (struct sockaddr*) &sb, sizeof(sb));

	connect(socks[1], (struct sockaddr*) &sa, sizeof(sa));
	connect(socks[0], (struct sockaddr*) &sb, sizeof(sb));
}

static int run_test()
{
	int fds[2];
	int pid1, pid2;
	int status2;

	socketpair(AF_LOCAL, SOCK_DGRAM, 0, fds);

	if (nonblock) {
		fcntl(fds[0], F_SETFL, (long) O_NONBLOCK);
		fcntl(fds[1], F_SETFL, (long) O_NONBLOCK);
	}

	if (!(pid1 = fork())) {
		role = SERVER;
		server(fds[1]);
	}
	read(fds[0], &status2, sizeof(status2));
	if (!(pid2 = fork())) {
		role = CLIENT;
		client(fds[0]);
	}
	waitpid(pid2, &status2, 0);
	kill(pid1, 15);
	waitpid(pid1, 0, 0);

	close(fds[0]);
	close(fds[1]);

	if (WIFEXITED(status2)) {
		return !!WEXITSTATUS(status2);
	} else {
		return 2;
	}
}

static filter_fn filters[8]
	= { filter_packet_ServerHello,
		filter_packet_ServerKeyExchange,
		filter_packet_ServerHelloDone,
		filter_packet_ClientKeyExchange,
		filter_packet_ClientChangeCipherSpec,
		filter_packet_ClientFinished,
		filter_packet_ServerChangeCipherSpec,
		filter_packet_ServerFinished };

static int run_one_test(int dropMode, int serverFinishedPermute, int serverHelloPermute, int clientFinishedPermute)
{
	int fnIdx = 0;
	int res, filterIdx;
	run_id = ((dropMode * 2 + serverFinishedPermute) * 6 + serverHelloPermute) * 6 + clientFinishedPermute;

	filter_clear_state();

	filter_chain[fnIdx++] = filter_permute_ServerHello;
	state_permute_ServerHello.order = permutations3[serverHelloPermute];

	filter_chain[fnIdx++] = filter_permute_ServerFinished;
	state_permute_ServerFinished.order = permutations2[serverFinishedPermute];
  
	filter_chain[fnIdx++] = filter_permute_ClientFinished;
	state_permute_ClientFinished.order = permutations3[clientFinishedPermute];

	if (dropMode) {
		for (filterIdx = 0; filterIdx < 8; filterIdx++) {
			if (dropMode & (1 << filterIdx)) {
				filter_chain[fnIdx++] = filters[filterIdx];
			}
		}
	}
	filter_chain[fnIdx++] = NULL;

	res = run_test();

	switch (res) {
		case 0:
			fprintf(stdout, "%i ++ ", run_id);
			break;
		case 1:
			fprintf(stdout, "%i -- ", run_id);
			break;
		case 2:
			fprintf(stdout, "%i !! ", run_id);
			break;
	}

	fprintf(stdout, "SHello(%s), ", permutation_names3[serverHelloPermute]);
	fprintf(stdout, "SFinished(%s), ", permutation_names2[serverFinishedPermute]);
	fprintf(stdout, "CFinished(%s) :- ", permutation_names3[clientFinishedPermute]);
	if (dropMode) {
		for (filterIdx = 0; filterIdx < 8; filterIdx++) {
			if (dropMode & (1 << filterIdx)) {
				if (dropMode & ((1 << filterIdx) - 1)) {
					fprintf(stdout, ", ");
				}
				fprintf(stdout, "%s", filter_names[filterIdx]);
			}
		}
	}
	fprintf(stdout, "\n");

	return res;
}

static void run_tests(int childcount)
{
	int children = 0;
	int dropMode, serverFinishedPermute, serverHelloPermute, clientFinishedPermute;

	for (dropMode = 0; dropMode != 1 << 8; dropMode++)
	for (serverFinishedPermute = 0; serverFinishedPermute < 2; serverFinishedPermute++)
	for (serverHelloPermute = 0; serverHelloPermute < 6; serverHelloPermute++)
	for (clientFinishedPermute = 0; clientFinishedPermute < 6; clientFinishedPermute++) {
		usleep(10000);
		if (!fork()) {
			exit(run_one_test(dropMode, serverFinishedPermute, serverHelloPermute, clientFinishedPermute));
		} else {
			children++;
			while (children >= childcount) {
				waitpid(0, 0, 0);
				children--;
			}
		}
	}

	while (children > 0) {
		waitpid(0, 0, 0);
		children--;
	}
}



static void parse_permutation(const char* arg, const char* permutations[], int* val)
{
	*val = 0;
	while (permutations[*val]) {
		if (strcmp(permutations[*val], arg) == 0) {
			return;
		} else {
			*val += 1;
		}
	}
	fprintf(stderr, "Unknown permutation %s\n", arg);
	exit(1);
}

int main(int argc, const char* argv[])
{
	setlinebuf(stdout);
	gnutls_global_init();
	gnutls_global_set_log_function(logfn);
	gnutls_global_set_audit_log_function(auditfn);
	gnutls_global_set_log_level(+3);

	nonblock = 0;
	debug = 1;

	if (argc == 1) {
		run_tests(100);
	} else {
		int dropMode = 0;
		int serverFinishedPermute = 0;
		int serverHelloPermute = 0;
		int clientFinishedPermute = 0;
		int arg;

#define NEXT_ARG(name) \
		do { \
			if (++arg >= argc) { \
				fprintf(stderr, "No arg to -" #name "\n"); \
				exit(1); \
			} \
		} while (0);

		for (arg = 1; arg < argc; arg++) {
			if (strcmp("-shello", argv[arg]) == 0) {
				NEXT_ARG(shello);
				parse_permutation(argv[arg], permutation_names3, &serverHelloPermute);
			} else if (strcmp("-sfinished", argv[arg]) == 0) {
				NEXT_ARG(sfinished);
				parse_permutation(argv[arg], permutation_names2, &serverFinishedPermute);
			} else if (strcmp("-cfinished", argv[arg]) == 0) {
				NEXT_ARG(cfinished);
				parse_permutation(argv[arg], permutation_names3, &clientFinishedPermute);
			} else {
				int drop;
				for (drop = 0; drop < 8; drop++) {
					if (strcmp(filter_names[drop], argv[arg]) == 0) {
						dropMode |= (1 << drop);
						break;
					}
				}
				if (drop == 8) {
					fprintf(stderr, "Unknown packet %s\n", argv[arg]);
					exit(1);
				}
			}
		}

		run_one_test(dropMode, serverFinishedPermute, serverHelloPermute, clientFinishedPermute);
	}
}

// vim: foldmethod=marker