/*
 * This file is licensed under the GPL v2, or a later version
 * at the discretion of Eric Wong <normalperson@yhbt.net>.
 */

/*
 * The real number of connections is MAX_CHILDREN * NR_CONN, we fork
 * children to get around file descriptor limits and to take advantage
 * of multiple processors/cores
 *
 * NR_CONN needs to be <= (FD_SETSIZE - 4)
 */
static int MAX_CHILDREN = 1;
static int NR_CONN = 1000;

/*
 * the delay between subsequent send/recv() calls to each socket,
 * assuming the socket is write/readable
 */
static int DELAY_SECONDS = 1;

/* the maximum number of bytes we write at once */
static int WR_BYTES = 8;

/* the maximum number of bytes we read at once */
static int RD_BYTES = 8;

/*
 * set this to 1 to enable printing of requests and responses to stdout
 */
static int DEBUG = 0;

/* -------- end of configuration section --------- */


#include <stdio.h>
#include <sys/time.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <unistd.h>
#include <stdlib.h>
#include <netdb.h>
#include <time.h>
#include <errno.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include <sys/wait.h>
#include <fcntl.h>
#include <stdarg.h>
#include <string.h>
#include <netinet/in.h>
#include <arpa/inet.h>

static char *http_request;
static size_t http_request_size;
static struct addrinfo *ap;
static int max_fd;
#define CONNECT_OFFSET 0x7fffffff

#ifdef TCP_NODELAY
static int can_tcp_nodelay = 1;
#else
static int can_tcp_nodelay;
#define TCP_NODELAY 0 /* just to get it to compile */
#endif

struct interface {
	int fd;

	/* offset:
	 *  CONNECT_OFFSET indicates connect()
	 *  < 0 indicates writing
	 *  >= 0 indicates reading
	 */
	off_t offset;

	time_t last;
};

static struct interface *interfaces;
static pid_t *children_pids;

static void die(const char *fmt, ...)
{
	va_list args;
	va_start(args, fmt);
	fprintf(stderr, fmt, args);
	va_end(args);
	fputc('\n', stderr);
	exit(EXIT_FAILURE);
}

static void mmap_request(const char *path)
{
	struct stat sb;
	int fd = open(path, O_RDONLY);

	if (fd < 0 || fstat(fd, &sb) < 0)
		die("Couldn't open or stat: %s", strerror(errno));

	http_request_size = sb.st_size;
	http_request = mmap(NULL, http_request_size,
	                    PROT_READ, MAP_PRIVATE, fd, 0);
}

static void setup_addrinfo(void)
{
	struct addrinfo *ans;
	struct addrinfo hints;
	char *http_host = getenv("HTTP_HOST");
	char *http_port = getenv("HTTP_PORT");

	if (!http_host)
		die("HTTP_HOST must be defined in the environment");
	if (!http_port)
		http_port = "80";

	hints.ai_flags = 0;
	hints.ai_family = PF_UNSPEC;
	hints.ai_socktype = SOCK_STREAM;
	hints.ai_protocol = IPPROTO_TCP;
	hints.ai_addrlen = 0;
	hints.ai_addr = NULL;
	hints.ai_canonname = NULL;
	hints.ai_next = NULL;

	if (getaddrinfo(http_host, http_port, &hints, &ans) != 0)
		die("Couldn't get address");

	for (ap = ans; ap != NULL; ap = ap->ai_next) {
		int fd;
		fd = socket(ap->ai_family, ap->ai_socktype, ap->ai_protocol);
		if (fd < 0)
			die("Couldn't connect: %s", strerror(errno));
		if (connect(fd, ap->ai_addr, ap->ai_addrlen) >= 0
		    || errno == EINPROGRESS) {
			close(fd);
			break;
		}
		close(fd);
	}
}

static void init_interface(const int i)
{
	int fd;
	socklen_t optval;
	fd = socket(ap->ai_family, ap->ai_socktype, ap->ai_protocol);
	if (fd > max_fd)
		max_fd = fd;
	if (fd < 0) {
		fprintf(stderr,
			"Couldn't get socket: %s", strerror(errno));
		return;
	}
	if (fcntl(fd, F_SETFL, fcntl(fd, F_GETFL) | O_NONBLOCK) < 0)
		die("Couldn't set nonblocking: %s", strerror(errno));
	if (connect(fd, ap->ai_addr, ap->ai_addrlen) < 0) {
		if (errno != EINPROGRESS)
			die("connect error: %s", strerror(errno));
	}

	/* we want buffers to be as small as possible,
	 * and disable Nagle algorithm buffering, too. */
	optval = 1;
	if (setsockopt(fd, SOL_SOCKET, SO_SNDBUF,
		       &optval, sizeof(optval)) < 0)
		die("[%d] Couldn't set SNDBUF: %s", i, strerror(errno));
	optval = 1;
	if (setsockopt(fd, SOL_SOCKET, SO_RCVBUF,
		       &optval, sizeof(optval)) < 0)
		die("[%d] Couldn't set RCVBUF: %s", i, strerror(errno));
	if (can_tcp_nodelay) {
		optval = 1;
		if (setsockopt(fd, IPPROTO_TCP, TCP_NODELAY,
			       &optval, sizeof(optval)) < 0) {
			can_tcp_nodelay = 0;
			fprintf(stderr, "Couldn't set TCP_NODELAY: %s\n",
			        strerror(errno));
		}
	}
	interfaces[i].fd = fd;
	interfaces[i].offset = CONNECT_OFFSET;
	/* fprintf(stderr, "[%d] ready fd: %d\n", i, fd); */
}

static void open_interfaces(void)
{
	int i;

	interfaces = malloc(sizeof(struct interface) * NR_CONN);
	for (i = 0; i < NR_CONN; ++i)
		init_interface(i);
}

static void do_io(void)
{
	struct timeval tv = { 1, 0 };
	time_t now;
	fd_set wfds;
	fd_set rfds;
	FD_ZERO(&wfds);
	FD_ZERO(&rfds);
	char buf[RD_BYTES];
	ssize_t nr;
	int i;

	now = time(NULL);
	for (i = 0; i < NR_CONN; ++i) {
		if (interfaces[i].fd < 0)
			continue;
		if (DELAY_SECONDS &&
		    (now - interfaces[i].last) <= DELAY_SECONDS)
			continue;
		if (interfaces[i].offset < 0 ||
		    interfaces[i].offset == CONNECT_OFFSET) {
			FD_SET(interfaces[i].fd, &wfds);
		} else {
			FD_SET(interfaces[i].fd, &rfds);
		}
	}

	i = select(max_fd + 1, &rfds, &wfds, NULL, &tv);
	if (i < 0)
		die("select returned %i: %s %d", i, strerror(errno), NR_CONN);

	for (i = 0; i < NR_CONN; ++i) {
		if (interfaces[i].fd < 0)
			continue;
		if (FD_ISSET(interfaces[i].fd, &wfds)) {
			ssize_t wr_bytes;
			if (interfaces[i].offset == CONNECT_OFFSET) {
				int err;
				size_t err_len = sizeof(err);
				getsockopt(interfaces[i].fd, SOL_SOCKET,
				           SO_ERROR, &err, &err_len);
				if (err) {
					close(interfaces[i].fd);
					init_interface(i);
					continue;
				}
				interfaces[i].offset = -http_request_size;
			}
			/* if offset == -2 and WR_BYTES == 8,
			 * only write 2 bytes here: */
			if (interfaces[i].offset > -WR_BYTES)
				wr_bytes = -interfaces[i].offset;
			else
				wr_bytes = WR_BYTES;

			nr = send(interfaces[i].fd,
			          http_request +
			            http_request_size + interfaces[i].offset,
			          wr_bytes, MSG_DONTWAIT);
			if (DEBUG && nr > 0) {
				printf("W%d:", i);
				fwrite(http_request +
				         http_request_size +
				         interfaces[i].offset,
				       1, wr_bytes, stdout);
				puts("");
			}
		} else if (FD_ISSET(interfaces[i].fd, &rfds)) {
			nr = recv(interfaces[i].fd, buf,
			          RD_BYTES, MSG_DONTWAIT);
			if (DEBUG && nr > 0) {
				printf("R%d:", i);
				fwrite(buf, 1, nr, stdout);
				puts("");
			}
		} else
			continue;

		if (nr > 0) {
			now = interfaces[i].last = time(NULL);
			interfaces[i].offset += nr;
		} else {
			switch (errno) {
			case EAGAIN:
			case EINTR:
			case EINPROGRESS:
				break;
			default:
				fprintf(stderr, "[%d] closing for error: %s\n",
					i, strerror(errno));
				close(interfaces[i].fd);
				init_interface(i);
			}
		}
	}
}

static void kill_children(void)
{
	int i;
	for (i = 0; i < MAX_CHILDREN; ++i)
		kill(children_pids[i], SIGTERM);
}

int main(int argc, const char *argv[])
{
	int i;

	if (argc < 2)
		return -1;

	mmap_request(argv[1]);
	setup_addrinfo();
	children_pids = malloc(sizeof(pid_t) * MAX_CHILDREN);

	for (i = 0; i < MAX_CHILDREN; ++i) {
		fflush(NULL);
		pid_t pid = fork();
		if (!pid) {
			open_interfaces();
			while (1)
				do_io();
			return 0;
		}
		if (pid < 0)
			die("Couldn't fork: %s", strerror(errno));
		children_pids[i] = pid;
	}
	atexit(kill_children);

	{
		int too_lazy_to_check;
		for (i = 0; i < MAX_CHILDREN; ++i)
			wait(&too_lazy_to_check);
	}

	return 0;
}