#include <sys/types.h>
#include <sys/time.h>
#include <sys/resource.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <netdb.h>
#include <unistd.h>
#include <stdlib.h>
#include <string.h>
#include <stdio.h>
#include <signal.h>
#include <errno.h>

#ifdef __i386__
#define rdtscl(low) \
     __asm__ __volatile__ ("rdtsc" : "=A" (low))
#else
#define rdtscl(val) \
do {						\
	struct timespec now;			\
	clock_gettime(CLOCK_REALTIME, &now);	\
	val  = (u_int64_t)now.tv_sec;		\
	val *= 1000000000LL;			\
	val += (u_int64_t)now.tv_nsec;		\
} while(0)
#endif

void
usage(void)
{
	extern char *__progname;
	fprintf(stderr, "%s [-n num] [-a [-l lqueue] [IP] port | -b | "
	    "-c [-C bindaddr [-m mask]] [IP] port | -f | -s]\n", __progname);
	exit(1);
}

void
display(char *msg, int i, u_int64_t a, u_int64_t b)
{
	printf("%s\t%d\t%lld\n", msg, i, b - a);
}

void
socket_ben(int n, int do_bind)
{
	int i, *sockets;
        struct sockaddr_in sin;
	u_int64_t a, b;

	sockets = calloc(n, sizeof(int));
	if (sockets == NULL) {
		fprintf(stderr, "malloc: %s\n", strerror(errno));
		exit(1);
	}
        sin.sin_family = AF_INET;
        sin.sin_addr.s_addr = INADDR_ANY;
        sin.sin_port = 0;

	for (i = 0; i < n; i++) {
		rdtscl(a);
		if ((sockets[i] = socket(AF_INET, SOCK_STREAM, 0)) < 0) {
			fprintf(stderr, "socket %d: %s\n", i, strerror(errno));
			n = i;
			break;
		}
		rdtscl(b);
		display("socket", i, a, b);

		if (do_bind == 0)
			continue;

		rdtscl(a);
		if (bind(sockets[i], (struct sockaddr *)&sin, sizeof(sin))) {
			fprintf(stderr, "bind %d: %s\n", i, strerror(errno));
			exit(2);
		}
		rdtscl(b);
		display("bind", i, a, b);
	}
}

void
fork_ben(int n)
{
	int i, quit = 0, pfd[2];
	char buf[1];
	pid_t *proc;
	u_int64_t a, b;

	proc = calloc(n, sizeof(pid_t));
	if (proc == NULL) {
		fprintf(stderr, "malloc: %s\n", strerror(errno));
		exit(1);
	}
	if (pipe(pfd) < 0) {
		fprintf(stderr, "pipe: %s\n", strerror(errno));
		exit(1);
	}
	for (i = 0; i < n; i++) {
		rdtscl(a);
		proc[i] = fork();
		switch(proc[i]) {
		case -1:
			n = i;
			quit = 1;
			break;
		case 0:
			/* child */
			write(pfd[1], "+", 1);
			close(pfd[1]);
			sleep(5*60);
			_exit(0);
			break;
		default:
			break;
		}
		if (quit || read(pfd[0], buf, 1) != 1)
			break;	
		rdtscl(b);
		display("fork", i, a, b);
	}

	for (i = 0; i < n; i++)
		if (proc[i] > 1)
			kill(proc[i], SIGTERM);
}

void
accept_ben(int n, struct in_addr addr, u_int16_t port, int lnum)
{
	int i, bindfd, *sockets, sinlen;
        struct sockaddr_in sin;
	u_int64_t a, b;

	sockets = calloc(n, sizeof(int));
	if (sockets == NULL) {
		fprintf(stderr, "malloc: %s\n", strerror(errno));
		exit(1);
	}
        sin.sin_family = AF_INET;
        sin.sin_addr = addr;
        sin.sin_port = port;

	bindfd = socket(AF_INET, SOCK_STREAM, 0);
	if (bindfd == -1) {
		fprintf(stderr, "socket %d: %s\n", i, strerror(errno));
		exit(2);
	}
	if (bind(bindfd, (struct sockaddr *)&sin,
		    sizeof(struct sockaddr_in)) == -1) {
		fprintf(stderr, "bind %d: %s\n", i, strerror(errno));
		exit(2);
	}
	
	if (listen(bindfd, lnum) == -1) {
		fprintf(stderr, "listen %d: %s\n", i, strerror(errno));
		exit(2);
	}

	fprintf(stderr, "accepting connections ...");
	for (i = 0; i < n; i++) {
		sinlen = sizeof(sin);
		rdtscl(a);
		sockets[i] = accept(bindfd, (struct sockaddr *)&sin, &sinlen);
		if (sockets[i] == -1) {
			fprintf(stderr, "accept %d: %s\n", i, strerror(errno));
			exit(2);
		}
		rdtscl(b);
		display("accept", i, a, b);
	}
	fprintf(stderr, "done\n");
}

void
connect_ben(int n, struct in_addr addr, u_int16_t port,
    struct in_addr base, u_int32_t mask)
{
	int i, *sockets;
        struct sockaddr_in sbind, sconn;
	u_int64_t a, b;

	sockets = calloc(n, sizeof(int));
	if (sockets == NULL) {
		fprintf(stderr, "malloc: %s\n", strerror(errno));
		exit(1);
	}
	bzero(&sconn, sizeof(sconn));
	sconn.sin_len = sizeof(sconn);
        sconn.sin_family = AF_INET;
        sconn.sin_port = port;
        sconn.sin_addr = addr;

	for (i = 0; i < n; i++) {
		rdtscl(a);
		if ((sockets[i] = socket(AF_INET, SOCK_STREAM, 0)) < 0) {
			fprintf(stderr, "socket %d: %s\n", i, strerror(errno));
			n = i;
			break;
		}
		rdtscl(b);
		display("socket", i, a, b);

		if (mask != 0) {
			sbind.sin_len = sizeof(sbind);
			sbind.sin_family = AF_INET;
			sbind.sin_port = 0;
			sbind.sin_addr.s_addr = htonl(
			    (ntohl(base.s_addr) & mask) |
			    (arc4random() & ~mask));
			rdtscl(a);
			if (bind(sockets[i],
				    (struct sockaddr *)&sbind, sizeof(sbind))) {
				fprintf(stderr, "bind %d: %s\n",
				    i, strerror(errno));
				exit(2);
			}
			rdtscl(b);
			display("bind", i, a, b);
		}

		rdtscl(a);
		if (connect(sockets[i],
			    (struct sockaddr *)&sconn, sizeof(sconn))) {
			fprintf(stderr, "connect %d: %s [%i]\n",
			    i, strerror(errno), errno);
			exit(2);
		}
		rdtscl(b);
		display("connect", i, a, b);

	}
}

int
my_atoi(char *s, int min, int max)
{
	char *ep;
	long lval;

	lval = strtol(s, &ep, 0);
	if (s[0] == '\0' || *ep != '\0') {
		fprintf(stderr, "argument not a number.");
		usage();
	}
	if (lval < (long)min || lval > (long)max) {
		fprintf(stderr, "number out of range (%d - %d).", min, max);
		usage();
	}
	return lval;
}

int
main(int argc, char **argv)
{
	struct rlimit rl;
        struct in_addr addr, maddr, base = {0};
	int l, n, p, ch, a_flag, b_flag, c_flag, s_flag, f_flag;
	u_int32_t mask = 0;

	a_flag = b_flag = c_flag = f_flag = s_flag = 0;
	n = 100; l = 10;

	while ((ch = getopt(argc, argv, "abcC:fl:m:n:s")) != -1) {
		switch (ch) {
		case 'a':
			a_flag = 1;
			break;
		case 'b':
			b_flag = 1;
			break;
		case 'c':
			c_flag = 1;
			break;
		case 'C':
			if (inet_aton(optarg, &base) != 1) {
				fprintf(stderr, "bad network address\n");
				usage();
			}
			break;
		case 'f':
			f_flag = 1;
			break;
		case 'l':
			if (a_flag == 0) usage();
			l = my_atoi(optarg, 1, 128);
			break;
		case 'n':
			n = my_atoi(optarg, 1, INT_MAX);
			break;
		case 'm':
			if (inet_aton(optarg, &maddr) != 1) {
				fprintf(stderr, "bad network mask\n");
				usage();
			}
			mask = ntohl(maddr.s_addr);
			break;
		case 's':
			s_flag = 1;
			break;
		case '?':
		default:
			usage();
		}
	}
	argc -= optind;
	argv += optind;

	if (c_flag || a_flag) {
		if (argc == 0 || argc > 2) usage();
		if (argc != 2)
			addr.s_addr = 0;
		else
			if (inet_aton(*argv++, &addr) != 1) {
				fprintf(stderr, "bad network address");
				usage();
			}
		p = htons(my_atoi(*argv, 1, 65535));
	}

	rl.rlim_cur = RLIM_INFINITY;
	rl.rlim_max = RLIM_INFINITY;
	setrlimit(RLIMIT_NOFILE, &rl);
	setrlimit(RLIMIT_NPROC, &rl);

	if (s_flag)
		socket_ben(n, 0);
	else if (b_flag)
		socket_ben(n, 1);
	else if (f_flag)
		fork_ben(n);
	else if (c_flag)
		connect_ben(n, addr, p, base, mask);
	else if (a_flag)
		accept_ben(n, addr, p, l);
	else
		fprintf(stderr, "no benchmark\n");
	exit(0);
}