#include <stdlib.h>
#include <stdio.h>
#include <unistd.h>
#include <errno.h>
#include <signal.h>
#include <string.h>
#include <wait.h>
#include "net.h"

/*
 * Creates a daemon listening on a port.
 * When it receives a connection it connects to the specified socks server and
 * sets up the communication to the target using the socks4a protocol,
 * then it bounces data beetween the client and the socks.
 *
 * Useful for eg Tor where the client does not know socks
 *
 * Copyright (c) Yoann Guillot 2006
 *
 * Redistributes under the terms of the GPL v2
 *
 */

static volatile int run = 1;
static void sigh(int sign)
{
	if (sign == SIGCHLD)
		while (waitpid(0, 0, WNOHANG) > 0);
	else
		run = 0;
	printf("sig %d, run=%d\n", sign, run);
}

static int sock_neg(int fd, char *host, unsigned short port)
{
	unsigned char buf[256];

	buf[0] = 4;	// socks version
	buf[1] = 1;	// 1 = connect, 2 = bind

	buf[2] = port >> 8; buf[3] = port;  // port number, network byte order

	buf[4] = buf[5] = buf[6] = 0;
	buf[7] = 1;     // addr = 0.0.0.x lets the socks resolve the hostname
	buf[8] = 0;     // strz = login
	strncpy(buf+9, host, 246);
	buf[255] = 0;   // strz = hostname

#ifdef DEBUG
	printf("starting socks negociation\n");
#endif

	write(fd, buf, 9 + strlen(buf+9) + 1);

	fd_set rfd;
	FD_ZERO(&rfd);
	FD_SET(fd, &rfd);
	select(fd+1, &rfd, 0, 0, 0);
	read(fd, buf, 8);

#ifdef DEBUG
	static char *s_ans[] = {"access granted", "acces failed", "failed-noident", "failed-badindent"};
	printf("socks answer: zero %d, %s %x\n", buf[0], s_ans[buf[1] - 0x5a], buf[1]);
#endif

	return buf[1] != 0x5a;
}

static void bounce(int fd1, int fd2)
{
	fd_set rfds;
	static char buf[4096];
	ssize_t ret;
	struct timeval tv;
	int mfd = fd1 > fd2 ? fd1+1 : fd2+1;

	while (run && fd1 >= 0 && fd2 >= 0) {
		FD_ZERO(&rfds);
		FD_SET(fd1, &rfds);
		FD_SET(fd2, &rfds);

		tv.tv_sec  = 1;
		tv.tv_usec = 0;

		switch (select(mfd, &rfds, 0, 0, &tv)) {
		case -1:
			if (errno == EINTR)
				continue;
			return;
		case 0: continue;
		}

		if (FD_ISSET(fd1, &rfds)) {
			ret = read(fd1, buf, 4096);
			if (ret <= 0)
				return;
			if (ret != write(fd2, buf, ret))
				return;
		}
		if (FD_ISSET(fd2, &rfds)) {
			ret = read(fd2, buf, 4096);
			if (ret <= 0)
				return;
			if (ret != write(fd1, buf, ret))
				return;
		}
	}
}

static void main_loop(int lfd, char *sockh, char *sockp, char *targeth, unsigned short targetp)
{
	fd_set lfds;
	struct timeval tv;
	while (run) {
		FD_ZERO(&lfds);
		FD_SET(lfd, &lfds);

		tv.tv_sec  = 1;
		tv.tv_usec = 0;

		switch (select(lfd+1, &lfds, 0, 0, &tv)) {
		case -1:
			if (errno == EINTR)
				continue;
			socket_close(lfd);
			exit(1);
		case 0:
			continue;
		}

		int afd = socket_accept(lfd, 0, 0, 0, 0);
		if (afd == -1) {
			sleep(1);
			continue;
		}

		switch (fork()) {
		case -1:
			socket_close(afd);
			sleep(1);
			break;
		case 0:
			close(lfd);
			int cfd = socket_connect(sockh, sockp);
			if (cfd >= 0 && !sock_neg(cfd, targeth, targetp))
				bounce(afd, cfd);
			socket_close(afd);
			socket_close(cfd);
			exit(0);
		}
		close(afd);
	}
}

int main(int argc, char **argv)
{
	int listener;

	if (argc < 7) {
		fprintf(stderr, "forwards incoming connections to target using socks4a\n");
		fprintf(stderr, "usage: %s <bind host> <bind port> <socks host> <socks port> <target host> <target port number>\n", *argv);
		exit(1);
	}

	signal(SIGINT, sigh);
	signal(SIGTERM, sigh);
	signal(SIGCHLD, sigh);

	listener = socket_listen(argv[1], argv[2]);
	if (listener == -1) {
		perror("listen");
		exit(-1);
	}
#ifndef DEBUG
	pid_t pid;
	switch(pid = fork()) {
	case -1:
		perror("fork");
		exit(1);
	case 0:
		close(0);
		close(1);
		close(2);
#endif
		main_loop(listener, argv[3], argv[4], argv[5], atoi(argv[6]));
		socket_close(listener);
		exit(0);
#ifndef DEBUG
	default:
		printf("backgrounded to pid %d\n", pid);
		return 0;
	}
#endif
}