#include <sys/types.h>
#include <sys/socket.h>
#include <sys/time.h>

#include <netinet/in.h>
#include <netinet/tcp.h>

#include <err.h>
#include <errno.h>
#include <limits.h>
#include <netdb.h>
#include <stdarg.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <time.h>

#if defined(TCP_CORK) && !defined(TCP_NOPUSH)
#define TCP_NOPUSH TCP_CORK
#endif

static void
report(const char *fmt, ...) {
	struct timeval tv;
	va_list ap;
	char now[10], msg[80];

	va_start(ap, fmt);
	if(gettimeofday(&tv, NULL)) err(1, "gettimeofday");
	strftime(now, sizeof(now), "%T", gmtime(&tv.tv_sec));
	vsnprintf(msg, sizeof(msg), fmt, ap);
	printf("%s.%06ld %s\n", now, tv.tv_usec, msg);
	va_end(ap);
}

static void
reportsa(const struct sockaddr *sa, socklen_t salen) {
	int r;
	char host[NI_MAXHOST];
	char serv[NI_MAXSERV];

	r = getnameinfo(sa, salen, host, sizeof(host), serv, sizeof(serv), NI_NAMEREQD);
	if(r == 0) report("hostname %s port %s", host, serv);
	r = getnameinfo(sa, salen, host, sizeof(host), serv, sizeof(serv), NI_NUMERICHOST);
	if(r) errx(1, "getnameinfo: %s", gai_strerror(r));
	report("address %s port %s", host, serv);
}

static int
reader(int s, size_t buffsize) {
	int r;
	char *buff;
	if(buffsize == 0) return(0);
	buff = malloc(buffsize);
	if(buff == NULL) err(1, "malloc");
	do {
		r = read(s, buff, buffsize);
		if(r < 0) err(1, "read");
		report("read %d", r);
	} while(r && buff[r-1] == 0);
	free(buff);
	return(r);
}

static void
writer(int s, int nopush, size_t buffsize, int count) {
	int i, r;
	char *buff;
	if(buffsize == 0) return;
	buff = malloc(buffsize);
	if(buff == NULL) err(1, "malloc");
	memset(buff, 0, buffsize);
	if(nopush) {
		i = 1; r = setsockopt(s, IPPROTO_TCP, TCP_NOPUSH, &i, sizeof(i));
		if(r < 0) err(1, "setsockopt TCP_NOPUSH %d", i);
	}
	for(i = 0; i < count; i++) {
		if(i == count-1) buff[buffsize-1] = 1;
		r = write(s, buff, buffsize);
		if(r < 0) err(1, "write");
		report("write %d", r);
	}
	if(nopush) {
		i = 0; r = setsockopt(s, IPPROTO_TCP, TCP_NOPUSH, &i, sizeof(i));
		if(r < 0) err(1, "setsockopt TCP_NOPUSH %d", i);
	}
	free(buff);
}

static void
usage(void) {
	errx(1, "usage: a.out [-ns] [-d delay] [-h hostname] [-p port]"
	     " [-r readbuff] [-R requests] [-w writebuff] [-W writecount]");
}

int
main(int argc, char *argv[]) {
	int i, r, s;
	int nopush, server;
	const char *hostname, *port;
	size_t readbuff, writebuff, sockbuff;
	int writecount, requests;
	u_int delay;
	struct addrinfo ai_hints, *ai0, *ai;
	struct sockaddr_storage sa;
	socklen_t salen;

	/*
	 * defaults
	 */
	nopush = 0;
	server = 0;
	hostname = port = NULL;
	delay = 0;
	readbuff = 65536;
	requests = 1;
	writebuff = 0;
	writecount = 1;
	sockbuff = 0;

	while((r = getopt(argc, argv, "d:h:np:r:R:sS:w:W:")) != -1)
		switch(r) {
		case('d'):
			delay = atoi(optarg);
			break;
		case('h'):
			hostname = optarg;
			break;
		case('n'):
			nopush = 1;
			break;
		case('p'):
			port = optarg;
			break;
		case('r'):
			readbuff = atoi(optarg);
			break;
		case('R'):
			requests = atoi(optarg);
			break;
		case('s'):
			server++;
			break;
		case('S'):
			sockbuff = atoi(optarg);
			break;
		case('w'):
			writebuff = atoi(optarg);
			break;
		case('W'):
			writecount = atoi(optarg);
			break;
		default:
			usage();
		}
	argc -= optind;
	if(argc) usage();

	setbuf(stdout, NULL);

	report("init");

	memset(&ai_hints, 0, sizeof(ai_hints));
	if(server)
		ai_hints.ai_flags = AI_PASSIVE;
	ai_hints.ai_family = PF_UNSPEC;
	ai_hints.ai_socktype = SOCK_STREAM;
	r = getaddrinfo(hostname, port, &ai_hints, &ai0);
	if(r) errx(1, "getaddrinfo(%s,%s,...): %s", hostname, port, gai_strerror(r));

	r = -1;
	errno = EADDRNOTAVAIL;
	for(ai = ai0; ai; ai = ai->ai_next) {

		report("trying flags=%d family=%d type=%d proto=%d",
		       ai->ai_flags, ai->ai_family, ai->ai_socktype, ai->ai_protocol);
		reportsa(ai->ai_addr, ai->ai_addrlen);

		r = s = socket(ai->ai_family, ai->ai_socktype, ai->ai_protocol);
		if(r < 0) { warn("socket"); continue; }

		if(sockbuff) {
			r = setsockopt(s, SOL_SOCKET, SO_SNDBUF, &sockbuff, sizeof(sockbuff));
			if(r < 0) err(1, "setsockopt SO_SNDBUF %lu", (unsigned long)sockbuff);
			r = setsockopt(s, SOL_SOCKET, SO_RCVBUF, &sockbuff, sizeof(sockbuff));
			if(r < 0) err(1, "setsockopt SO_RCVBUF %lu", (unsigned long)sockbuff);
		}
		if(server) {
			i = 1; r = setsockopt(s, SOL_SOCKET, SO_REUSEADDR, &i, sizeof(i));
			if(r < 0) err(1, "setsockopt SO_REUSEADDR");
			r = bind(s, ai->ai_addr, ai->ai_addrlen);
			if(r < 0) { warn("bind"); continue; }
			r = listen(s, SOMAXCONN);
			if(r < 0) err(1, "listen");
			report("waiting for connections");
		} else {
			r = connect(s, ai->ai_addr, ai->ai_addrlen);
			if(r < 0) { close(s); warn("connect"); continue; }
			report("connected");
		}
		break;
	}
	if(r < 0) errx(1, "all addresses failed");
	freeaddrinfo(ai0);

	sleep(delay);

	if(server) do {
		salen = sizeof(sa);
		r = accept(s, (struct sockaddr *)&sa, &salen);
		if(r < 0) err(1, "accept");
		report("connection received");
		reportsa((struct sockaddr *)&sa, salen);
		switch(server > 1 ? fork() : 0) {
		case(-1):
			err(1, "fork");
		case(0):
			close(s);
			s = r;
			server = 1;
			break;
		default:
			close(r);
			continue;
		}
	} while(server > 1);

	if(server && readbuff) {
		r = reader(s, readbuff);
		while(r) {
			writer(s, nopush, writebuff, writecount);
			r = reader(s, readbuff);
		}
	} else if(server) {
		writer(s, nopush, writebuff, writecount);
		sleep(delay);
	} else {
		for(i = 0; i < requests; i++) {
			writer(s, nopush, writebuff, writecount);
			reader(s, readbuff);
			sleep(delay);
		}
	}

	report("exit");
	exit(0);
}