/*
 * Copyright (c) 2021 Maurizio Lombardi
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 *
 * - Redistributions of source code must retain the above copyright
 *   notice, this list of conditions and the following disclaimer.
 * - Redistributions in binary form must reproduce the above copyright
 *   notice, this list of conditions and the following disclaimer in the
 *   documentation and/or other materials provided with the distribution.
 * - The name of the author may not be used to endorse or promote products
 *   derived from this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR
 * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
 * OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
 * IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT,
 * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT
 * NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
 * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
 * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF
 * THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

#include <stdio.h>
#include <sys/socket.h>
#include <sys/select.h>
#include <netinet/in.h>
#include <netinet/tcp.h>
#include <unistd.h>
#include <string.h>
#include <stdlib.h>
#include <sys/types.h>
#include <sys/uio.h>
#include <sys/stat.h>
#include <sys/mman.h>
#include <arpa/inet.h>
#include <pthread.h>
#include <fcntl.h>
#include <dirent.h>
#include <libgen.h>
#include <ctype.h>
#include <signal.h>
#include <netdb.h>

#define BUFSIZE  4096

struct th_data {
	int in_conn_fd;
	int out_conn_fd;
};

static void* handle_connection(void *data);

static void usage(void)
{
	printf("Usage: proxy -l local_port -h remote_address -p remote_port [-d] (daemonize)\n");
}

int main(int argc, char **argv)
{
	int s, s6, r, opt;
	int nargs = 0;
	int local_port = -1;
	int daemonize = 0;
	int one = 1;
	char *remote_addr_str = NULL;
	struct sockaddr_in srv_addr;
	struct sockaddr_in6 srv_addr6;
	struct sockaddr_storage client_addr;
	pthread_attr_t th_attr;
	char *remote_port_str;

	while ((opt = getopt(argc, argv, "l:h:p:d")) != -1) {
		switch (opt) {
		case 'l':
			local_port = atoi(optarg);
			nargs++;
			break;
		case 'h':
			remote_addr_str = optarg;
			nargs++;
			break;
		case 'p':
			remote_port_str = optarg;
			nargs++;
			break;
		case 'd':
			daemonize = 1;
			break;
		default:
			usage();
			return -1;
		}

	}

	if (nargs != 3) {
		usage();
		return -1;
	}

	sigaction(SIGPIPE, &(struct sigaction) {SIG_IGN}, NULL);

	if (daemonize) {
		int pid;
		int sid;

		pid = fork();
		if (pid < 0) {
			printf("Fork failed!\n");
			return pid;
		}

		if (pid > 0)
			exit(0);

		umask(0);
		sid = setsid();
		if (sid < 0)
			exit(sid);
		chdir("/");
		close(STDIN_FILENO);
		close(STDOUT_FILENO);
		close(STDERR_FILENO);
	}

	s = socket(AF_INET, SOCK_STREAM, 0);
	if (s < 0)
		return s;

	s6 = socket(AF_INET6, SOCK_STREAM, 0);
	if (s6 < 0)
		return s6;

	setsockopt(s6, IPPROTO_IPV6, IPV6_V6ONLY, &one, sizeof(one));

	/* Set up IPv4 server address structure */
	memset(&srv_addr, 0, sizeof(srv_addr));
	srv_addr.sin_family = AF_INET;
	srv_addr.sin_addr.s_addr = htonl(INADDR_ANY);
	srv_addr.sin_port = htons(local_port);

	/* Set up IPv6 server address structure */
	memset(&srv_addr6, 0, sizeof(srv_addr6));
	srv_addr6.sin6_family = AF_INET6;
	srv_addr6.sin6_addr = in6addr_any;
	srv_addr6.sin6_port = htons(local_port);

	/* Bind the server's IPv4 socket */
	r = bind(s, (struct sockaddr *) &srv_addr, sizeof(srv_addr));
	if (r < 0)
		return r;

	/* Bind the server's IPv6 socket */
	r = bind(s6, (struct sockaddr *)&srv_addr6, sizeof(srv_addr6));
	if (r < 0)
		return r;

	r = listen(s, 128);
	if (r < 0)
		return r;

	r = listen(s6, 128);
	if (r < 0)
		return r;

	pthread_attr_init(&th_attr);
	pthread_attr_setdetachstate(&th_attr, PTHREAD_CREATE_DETACHED);

	while (1) {
		pthread_t th;
		struct th_data *data = NULL;
		struct addrinfo hints, *res = NULL, *p;
		socklen_t sin_size = sizeof(struct sockaddr_storage);
		fd_set readfds;
		int in_fd = -1, out_fd = -1;
		int r;

		FD_ZERO(&readfds);
		FD_SET(s, &readfds);
		FD_SET(s6, &readfds);

		r = select(FD_SETSIZE, &readfds, NULL, NULL, NULL);
		if (r < 0)
			return r;

		if (FD_ISSET(s, &readfds)) {
			in_fd = accept(s, (struct sockaddr *) &client_addr,
					&sin_size);
		} else if (FD_ISSET(s6, &readfds)) {
			in_fd = accept(s6, (struct sockaddr *) &client_addr,
					&sin_size);
		} else {
			continue;
		}

		if (in_fd < 0)
			continue;

		data = malloc(sizeof(*data));
		if (!data)
			goto err;

		memset(&hints, 0, sizeof(hints));
		hints.ai_family = AF_UNSPEC;
		hints.ai_socktype = SOCK_STREAM;

		r = getaddrinfo(remote_addr_str, remote_port_str, &hints, &res);
		if (r < 0) {
			printf("getaddrinfo() failed\n");
			goto err;
		}

		for (p = res; p != NULL; p = p->ai_next) {
			out_fd = socket(p->ai_family, p->ai_socktype, p->ai_protocol);
			if (out_fd < 0)
				continue;

			if (connect(out_fd, p->ai_addr, p->ai_addrlen) == -1) {
				close(out_fd);
				out_fd = -1;
				continue;
			}
			break;
		}

		if (out_fd < 0)
			goto err;

		setsockopt(out_fd, IPPROTO_TCP, TCP_NODELAY, &one, sizeof(one));
		setsockopt(in_fd, IPPROTO_TCP, TCP_NODELAY, &one, sizeof(one));

		data->out_conn_fd = out_fd;
		data->in_conn_fd = in_fd;

		r = pthread_create(&th, &th_attr, handle_connection, data);
		if (r)
			goto err;

		freeaddrinfo(res);
		continue;

err:
		if (in_fd >= 0)
			close(in_fd);
		if (out_fd >= 0)
			close(out_fd);
		free(data);
		if (res)
			freeaddrinfo(res);
	}

	return 0;
}

static void* handle_connection(void *data)
{
	struct th_data *th_data = data;
	int in_fd = th_data->in_conn_fd;
	int out_fd = th_data->out_conn_fd;
	pthread_attr_t th_attr;
	struct timeval timeout;
	char *buf = NULL;
	int r;
	fd_set rd;

	buf = malloc(BUFSIZE);
	if (!buf)
		goto exit;

	while (1) {
		FD_ZERO(&rd);
		FD_SET(in_fd, &rd);
		FD_SET(out_fd, &rd);

		timeout.tv_sec = 3600 * 4;
		timeout.tv_usec = 0;

		r = select(FD_SETSIZE, &rd, NULL, NULL, &timeout);
		if (r <= 0)
			goto exit;

		if (FD_ISSET(in_fd, &rd)) {
			r = read(in_fd, buf, BUFSIZE);
			if (r <= 0)
				goto exit;

			r = write(out_fd, buf, r);
			if (r < 0)
				goto exit;
		}

		if (FD_ISSET(out_fd, &rd)) {
			r = read(out_fd, buf, BUFSIZE);
			if (r <= 0)
				goto exit;

			r = write(in_fd, buf, r);
			if (r < 0)
				goto exit;
		}
	}

exit:
	close(in_fd);
	close(out_fd);
	free(data);
	free(buf);
	pthread_exit((void *) 0);
}