#include <errno.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

#include <sys/ioctl.h>
#include <sys/poll.h>
#include <sys/socket.h>
#include <sys/time.h>

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


#define DEFAULT_PORT 22222
#define BUFFER_SIZE 4096
#define MAX_FDS 200

void
run_server(int listen_port, int use_tfo, int send_size)
{
	int len, rc, on = 1;
	int listen_fd = -1, new_fd = -1;
	int end_server = 0, compress_array = 0;
	int close_conn;
	char buffer[BUFFER_SIZE];
	struct sockaddr_in addr;
	struct pollfd fds[MAX_FDS];
	unsigned int bytes_received[MAX_FDS];
	int nfds = 1, current_size = 0, i, j;

	printf("Listening on port %u on all interfaces, TFO %s, send_size=%d\n", listen_port, use_tfo ? "enabled" : "disabled", send_size); 
	
	listen_fd = socket(AF_INET, SOCK_STREAM, 0);
	if (listen_fd < 0) {
		perror("socket() failed");
		exit(1);
	}

	rc = setsockopt(listen_fd, SOL_SOCKET, SO_REUSEADDR,
			&on, sizeof(on));
	if (rc < 0) {
		perror("setsockopt() failed");
		close(listen_fd);
		exit(1);
	}

	if (use_tfo) {
		rc = setsockopt(listen_fd, IPPROTO_TCP, TCP_FASTOPEN,
				&on, sizeof(on));
		if (rc < 0) {
			perror("setsockopt(TCP_FASTOPEN) failed");
			close(listen_fd);
			exit(1);
		}
	}

	rc = ioctl(listen_fd, FIONBIO, (char *)&on);
	if (rc < 0) {
		perror("ioctl() failed");
		close(listen_fd);
		exit(1);
	}

	memset(&addr, 0, sizeof(addr));
	addr.sin_family = AF_INET;
	addr.sin_addr.s_addr = INADDR_ANY;
	addr.sin_port = htons(listen_port);
	rc = bind(listen_fd, (struct sockaddr *)&addr, sizeof(addr));
	if (rc < 0) {
		perror("bind() failed");
		close(listen_fd);
		exit(1);
	}

	rc = listen(listen_fd, 32);
	if (rc < 0) {
		perror("listen() failed");
		close(listen_fd);
		exit(1);
	}

	memset(fds, 0 , sizeof(fds));
	fds[0].fd = listen_fd;
	fds[0].events = POLLIN;

	do {
		printf("Waiting on poll()...\n");
		rc = poll(fds, nfds, -1);
		if (rc < 0) {
			perror("  poll() failed");
			break;
		}

		current_size = nfds;
		for (i = 0; i < current_size; i++) {
			if (fds[i].revents == 0)
				continue;

			if (fds[i].revents != POLLIN) {
				if (fds[i].fd == listen_fd) {
					printf("  Error! revents = %d\n", fds[i].revents);
					end_server = 1;
					break;
				} else {
					printf("  Connection closed\n");
					close(fds[i].fd);
					fds[i].fd = -1;
					compress_array = 1;
					continue;
				}
			}

			if (fds[i].fd == listen_fd) {
				printf("  Listening socket is readable\n");

				do {
					new_fd = accept(listen_fd, NULL, NULL);
					if (new_fd < 0) {
						if (errno != EWOULDBLOCK) {
							perror("  accept() failed");
							end_server = 1;
						}
						break;
					}

					printf("  New incoming connection - %d\n", new_fd);
					fds[nfds].fd = new_fd;
					fds[nfds].events = POLLIN;
					bytes_received[nfds] = 0;
					nfds++;
				} while (new_fd != -1);
			} else {
				printf("  Descriptor %d is readable\n", fds[i].fd);
				close_conn = 0;

				do {
					rc = recv(fds[i].fd, buffer, sizeof(buffer), 0);
					if (rc < 0) {
						if (errno != EWOULDBLOCK) {
							perror("  recv() failed");
							close_conn = 1;
						}
						break;
					}

					if (rc == 0) {
						printf("  Connection closed\n");
						close_conn = 1;
						break;
					}

					len = rc;
					bytes_received[i] += rc;
					printf("  %d bytes received on fd %d\n", len, fds[i].fd);

					if (send_size != 0) {
						int bytes_to_send;
						bytes_to_send = (send_size < 0) ? len : send_size;
						if (bytes_to_send > len) {
							for (j = len; j < send_size; j++)
								buffer[j] = 'A' + (j % 26);
						}
						rc = send(fds[i].fd, buffer, bytes_to_send, 0);
						if (rc < 0) {
							perror("  send() failed");
							close_conn = 1;
							break;
						}
						printf("  %d bytes sent on fd %d\n", bytes_to_send, fds[i].fd);
					}

				} while(1);

				if (close_conn) {
					close(fds[i].fd);
					fds[i].fd = -1;
					compress_array = 1;
				}
			}
		}

		if (compress_array) {
			compress_array = 0;
			for (i = 0; i < nfds; i++) {
				if (fds[i].fd == -1) {
					if (i < nfds - 1)
						memcpy(&fds[i], &fds[i + 1], (nfds - i - 1) * sizeof(fds[0]));
					i--;
					nfds--;
				}
			}
		}
	} while (!end_server);

	for (i = 0; i < nfds; i++) {
		if (fds[i].fd >= 0)
			close(fds[i].fd);
	}
}

void
usage(const char *progname)
{
	printf("usage: %s [options]\n", progname);
	printf("\n");
	printf("Options:\n");
	printf("  -n           Use standard (non-TFO) listen (default is TFO)\n");
	printf("  -p port      Port number to listen on (default %u)\n", DEFAULT_PORT);
	printf("  -s size      Specify amount of data to send in response to each receive (default min(receive_size, %u))\n", BUFFER_SIZE);
	printf("\n");
}

int
main (int argc, char **argv)
{
	int c;
	int use_tfo;
	unsigned int port;
	int send_size;
	
	use_tfo = 1;
	port = DEFAULT_PORT;
	send_size = -1;
	while ((c = getopt (argc, argv, "np:s:")) != -1) {
		switch (c)
			{
			case 'n':
				use_tfo = 0;
				break;
			case 'p':
				port = strtoul(optarg, NULL, 10);
				break;
			case 's':
				send_size = strtoul(optarg, NULL, 10);
				if (send_size > BUFFER_SIZE) {
					printf("Limiting send size to %u\n", BUFFER_SIZE);
					send_size = BUFFER_SIZE;
				}
				break;
			default:
			case '?':
				usage(argv[0]);
				return (1);
			}
	}

	if (argc != optind) {
		usage(argv[0]);
		return (1);
	}

	run_server(port, use_tfo, send_size);
	
	return (0);
}