/*
 * net_utils.c:
 *
 */

#include "common.h"
#include "linked_list.h"
#include "net_utils.h"

#include <stdint.h>
#include <stdlib.h>
#include <string.h>

#include <errno.h>
#include <fcntl.h>
#include <netdb.h>
#include <pthread.h>
#include <unistd.h>
#include <arpa/inet.h>
#include <sys/select.h>
#include <sys/socket.h>

int test_connection(int sockfd)
{
	ssize_t ret;
	char c;
	int old_errno = errno;

	errno = 0;
	ret = recv(sockfd, &c, 1, MSG_DONTWAIT | MSG_PEEK);
	if (ret <= 0 && errno != EAGAIN)
		return -1;

	errno = old_errno;
	return 0;
}

void *resolve_host(const char *hostname, struct sockaddr_in *addr)
{
	struct hostent *hp = gethostbyname(hostname);

	if (hp == NULL) {
		addr->sin_addr.s_addr = inet_addr(hostname);
		hp = gethostbyaddr(&(addr->sin_addr.s_addr),
			sizeof(addr->sin_addr.s_addr), AF_INET);
	}

	return hp;
}

int tcp_server_socket(uint16_t port, int backlog)
{
	int sock;
	struct sockaddr_in addr;

	/* Create the socket. */
	sock = socket(AF_INET, SOCK_STREAM, 0);

	/* Create the address of the server. */
	addr.sin_family = AF_INET;
	addr.sin_port = htons(port);
	addr.sin_addr.s_addr = htonl(INADDR_ANY); /* Use the wildcard address.*/

	const int reuse = 1;
	setsockopt(sock, SOL_SOCKET, SO_REUSEADDR, &reuse, sizeof reuse);

	/* Bind the socket to the address. */
	if (bind(sock, (struct sockaddr *)&addr, sizeof addr)) {
		print_error("cannot bind the socket.\n");
		return -1;
	}

	/* Listen for connections.  */
	listen(sock, backlog);

	return sock;
}

int tcp_client_socket(const char *hostname, uint16_t port)
{
	int sock, len, flags;
	struct hostent *hp;
	struct sockaddr_in addr;
	struct timeval timeout = {.tv_sec = 5, .tv_usec = 0};
	fd_set working_fds;

	/* Look up our host's network address.*/
	if ((hp = resolve_host(hostname, &addr)) == NULL) {
		print_error("Can't find host %s\n", hostname);
		return -2;
	}

	/* Create a socket in the INET domain.*/
	sock = socket(AF_INET, SOCK_STREAM, 0);

	/* Set socket in non-blocking mode */
	flags = fcntl(sock, F_GETFL, 0);
	fcntl(sock, F_SETFL, flags | O_NONBLOCK);

	/* Create the address of the server. */
	addr.sin_family = AF_INET;
	addr.sin_port = htons(port);
	memcpy(&addr.sin_addr, hp->h_addr_list[0], hp->h_length);
	len = sizeof(struct sockaddr_in);

	/* Connect to the server. */
	connect(sock, (struct sockaddr *)&addr, len);

	/* Wait for timeout */
	FD_ZERO(&working_fds);
	FD_SET(sock, &working_fds);
	if (select(sock + 1, NULL, &working_fds, NULL, &timeout) == 1) {
		int so_error;
		socklen_t so_len = sizeof so_error;

		getsockopt(sock, SOL_SOCKET, SO_ERROR, &so_error, &so_len);
		if (so_error == 0) {
			flags = fcntl(sock, F_GETFL);
			fcntl(sock, F_SETFL, flags & ~O_NONBLOCK);
			return sock;
		}
	}

	/* Timed out */
	return -1;
}

void *tcp_relay(void *arg)
{
	struct conn_context *cc = (struct conn_context *)arg;

	int maxfd, retval;
	ssize_t nbr, nbw;

	int sockets[2] = {
		cc->sockets[0], 
		cc->sockets[1]
	};

	char buf[RELAY_BUFSZ];

	fd_set master_fds, working_fds;
	FD_ZERO(&master_fds);
	FD_SET(sockets[0], &master_fds);
	FD_SET(sockets[1], &master_fds);

	maxfd = 1 + max(sockets[0], sockets[1]);

	while (cc->status == S_RUNNING) {
		working_fds = master_fds;
		retval = select(maxfd, &working_fds, NULL, NULL, NULL);

		if (retval > 0) {
			if (FD_ISSET(sockets[0], &working_fds)) {
				if ((nbr = recv(sockets[0], buf, RELAY_BUFSZ, 0)) <= 0)
					break;

				if ((nbw = send(sockets[1], buf, nbr, MSG_NOSIGNAL)) <= 0)
					break;
			}

			if (FD_ISSET(sockets[1], &working_fds)) {
				if ((nbr = recv(sockets[1], buf,  RELAY_BUFSZ, 0)) <= 0)
					break;

				if ((nbw = send(sockets[0], buf, nbr, MSG_NOSIGNAL)) <= 0)
					break;
			}
		}
	}

	cc->status = S_TERMINATED;
	return NULL;
}

struct conn_context *find_unpaired_node(struct node *head)
{
	struct node *ni;
	struct conn_context *cc;

	for (ni = head; ni; ni = ni->next) {
		cc = ni->data;
		if (cc->status == S_UNPAIRED)
			return cc;
	}

	return NULL;
}

void mark_all_connections(struct node *head, int new_status)
{
	struct node *ni;
	struct conn_context *cc;

	for (ni = head; ni; ni = ni->next) {
		cc = ni->data;
		cc->status = new_status;
	}
}

void run_connection_scheduler(struct node **head_ref)
{
	struct node *ni, *nn;
	struct conn_context *cc;

	if (!head_ref)
		return;

	for (ni = *head_ref; ni != NULL; ni = nn) {
		nn = ni->next; /* Saved first in case of node removal. */
		cc = ni->data;

		switch (cc->status) {
		case S_CREATED:
			break;
		case S_UNPAIRED:
			if (get_timestamp_ms() - cc->timestamp > 5000) {
				print_info("Removing unpaired connection <client: %d>\n",
					cc->sockets[0]);
				/* Close client socket */
				close(cc->sockets[0]);
				/* Remove current node */
				node_remove(head_ref, ni);
			}
			break;
		case S_RUNNING:
			break;
		case S_WAITING:
			print_info("Starting a new tunnel <client: %d, relay: %d>\n",
				cc->sockets[0], cc->sockets[1]);

			/* Run thread */
			cc->status = S_RUNNING;
			pthread_create(&cc->tid, NULL, &tcp_relay, cc);

			/* When a detached thread terminates, its resources are
			 * automatically released. */
			pthread_detach(cc->tid);
			break;
		case S_TERMINATED:
			print_info("Closing connections <client: %d, relay: %d>\n",
				cc->sockets[0], cc->sockets[1]);

			/* Close descriptors */
			close(cc->sockets[0]);
			close(cc->sockets[1]);

			/* Critical section: removing while iterating. */
			node_remove(head_ref, ni);
			break;
		}
	}
}
