/*
socksproxy - Peter Benie <pjb1008@eng.cam.ac.uk> 
v4 - 2013-09-01
*/

#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <sys/uio.h>
#include <netinet/in.h>
#include <string.h>
#include <dlfcn.h>
#include <pwd.h>
#include <errno.h>
#include <fcntl.h>
#include <netdb.h>

static const char *progname="socksproxy";

typedef int connect_prototype(int fd, const struct sockaddr *addr, 
			      socklen_t len);
static connect_prototype *real_connect=NULL;
static int configurations=0;
static char **usernames;
static int *socksports;
static in_addr_t *networks;
static in_addr_t *netmasks;
static in_addr_t *socksservers;
static int live=0;

static int count(const char *arg)
{
	if (!arg) return 0;

	int i=1;
	while ((arg=index(arg, ':'))) i++, arg++;
	return i;
}

static void init(void)
{
	const char *portname, *networkspec, *hostname;
	int portno, i;
	in_addr_t thissocksserver;
	char *username;
 
	/* Illegal cast, but okay on gcc-linux-i386 and x86_64*/
	real_connect=(connect_prototype *)dlsym(RTLD_NEXT, "connect");
	if (!real_connect)
	{
		fprintf(stderr, "%s: dlsym: %s\n", progname, dlerror());
		exit(1);
	}
	username=getenv("SOCKS_PROXY_USERNAME");
	portname=getenv("SOCKS_PROXY_PORT");
	hostname=getenv("SOCKS_PROXY_SERVER");
	networkspec=getenv("SOCKS_PROXY_NETWORK");

	configurations=count(networkspec);

	if (!configurations) return;

	usernames=malloc(sizeof(char *) * configurations);
	socksports=malloc(sizeof(int) * configurations);
	networks=malloc(sizeof(in_addr_t) * configurations);
	netmasks=malloc(sizeof(in_addr_t) * configurations);
	socksservers=malloc(sizeof(in_addr_t) * configurations);
	if (!socksports || !networks || !netmasks || !socksservers)
	{
		fprintf(stderr, "%s: malloc: %s\n", progname, strerror(errno));
		exit(1);
	}

	if (!username)
	{
		struct passwd *pwd=getpwuid(getuid());
		if (!pwd)
		{
			fprintf(stderr, "%s: Don't know who you are - "
				"try setting SOCKS_PROXY_USERNAME\n", progname);
			exit(1);
		}
		username=pwd->pw_name;
		if (!username)
		{
			fprintf(stderr, "%s: strdup: %s\n", progname, strerror(errno));
			exit(1);
		}
	}

	for (i=0; i<configurations; i++)
	{
		char *tail=index(username, ':');
		if (!tail) tail=username+strlen(username);

		usernames[i]=malloc(tail-username+1);
		if (!usernames[i])
		{
			fprintf(stderr, "%s: malloc %lu: %s\n", 
				progname, 
				(unsigned long)(tail-username),
				strerror(errno));
			exit(1);
		}
		memcpy(usernames[i], username, tail-username);
		usernames[i][tail-username]=0;
		if (*tail) username=tail+1;
	}

	portno=1080;
	for (i=0; i<configurations; i++)
	{
		if (portname && *portname)
		{
			if (*portname!=':')
				portno=strtoul(portname, (char **)&portname, 10);
			if (*portname==':') portname++;
		}
		socksports[i]=htons((portno>0 && portno<65536)? portno:1080);
	}

	thissocksserver=htonl(INADDR_LOOPBACK);
	for (i=0; i<configurations; i++)
	{
		if (hostname && *hostname)
		{
			struct hostent ent, *result;
			int my_h_errno=0;
			char buf[5000];
			char *tail;
			tail=index(hostname, ':');
			if (tail) *tail=0;
			if (gethostbyname_r(hostname, &ent, buf, sizeof(buf), &result, &my_h_errno))
			{
				if (errno==ERANGE)
				{
					fprintf(stderr, "%s: gethostbyname_r: entry for %s too large - recompile with buf larger than %lu\n", progname, hostname, (unsigned long)sizeof(buf));
					exit(1);
				}
				fprintf(stderr, "%s: gethostbyname_r %s failed: error=%s\n", progname, hostname, strerror(errno));
				exit(1);
			}
			else if (!result)
			{
				fprintf(stderr, "%s: gethostbyname_r %s failed: h_error=%s\n", progname, hostname, hstrerror(my_h_errno));
				exit(1);
			}
			else if (result->h_addrtype != AF_INET || !result->h_addr_list[0])
			{
				fprintf(stderr, "%s: gethostbyname_r %s did not return an IPv4 address\n", progname, hostname);
				exit(1);
			}
			else
			{
				memcpy(&thissocksserver, result->h_addr_list[0], 4);
			}
			if (tail) *tail=':';
		}
		memcpy(&socksservers[i], &thissocksserver, 4);
	}

	for (i=0; i<configurations; i++)
	{
		char *endptr;
		int a,b=0,c=0,d=0,e;

		if (*networkspec=='/' || *networkspec=='.') return;
		a=strtoul(networkspec, &endptr, 10);
		if (*endptr=='/') goto slash;
		if (*endptr!='.') return; else networkspec=endptr+1;
		if (*networkspec=='/' || *networkspec=='.') return;
		b=strtoul(networkspec, &endptr, 10);
		if (*endptr=='/') goto slash;
		if (*endptr!='.') return; else networkspec=endptr+1;
		if (*networkspec=='/' || *networkspec=='.') return;
		c=strtoul(networkspec, &endptr, 10);
		if (*endptr=='/') goto slash;
		if (*endptr!='.') return; else networkspec=endptr+1;
		if (*networkspec=='/') return;
		d=strtoul(networkspec, &endptr, 10);
		if (*endptr!='/') return; 
	slash:
		networkspec=endptr+1;
		if (*networkspec=='0') return;
		e=strtoul(networkspec, &endptr, 10);
		if (*endptr!=0 && *endptr!=':') return;
		if (a<256 && b<256 && c<256 && d<256 && e>=0 && e<=32)
		{
			networks[i]=htonl(a<<24 | b<<16 | c<<8 | d);
			netmasks[i]=htonl((e==32)? 0: (0xffffffffL<<(32-e)));
		}
		networkspec=*endptr? endptr+1:NULL;
	}
	live=1;
}

static ssize_t writev_full(int fd, struct iovec *vector, int count)
{
	while (count)
	{
		ssize_t bytes=writev(fd, vector, count);
		if (bytes==-1 && errno==EINTR) continue;
		if (bytes==-1 && errno==EWOULDBLOCK) { sleep(1); continue; }
		if (bytes==-1) return -1; 
		while (count && vector->iov_len<=(size_t)bytes)
		{
			bytes-=vector->iov_len;
			vector++, count--;
		}
		if (count)
		{
			vector->iov_base+=bytes;
			vector->iov_len-=bytes;
		}
	}
	return 0;
}

static ssize_t read_full(int fd, void *buf, size_t count)
{
	while (count)
	{
		ssize_t bytes=read(fd, buf, count);
		if (bytes==-1 && errno==EINTR) continue;
		if (bytes==-1 && errno==EWOULDBLOCK) { sleep(1); continue; }
		if (bytes==-1) return -1;
		if (bytes==0) { errno=EIO; return -1; }
		buf+=bytes;
		count-=bytes;
	}
	return 0;
}

static int socks_connect(int socksfd, struct sockaddr_in *dst, socklen_t len,
	int i)
{
	struct sockaddr_in socks;
	char buf[8]={4,1};
	struct iovec vec[2];
	int flags;

	if (!real_connect) init();

	if (len<sizeof(struct sockaddr_in)) {errno=EINVAL; return -1;}

	socks.sin_family=AF_INET;
	socks.sin_port=socksports[i];
	socks.sin_addr.s_addr=socksservers[i];
	
	/* Temporarily turn off non-blocking mode */
	flags=fcntl(socksfd, F_GETFL);
	if (flags==-1) return -1;
	if (flags & O_NONBLOCK)
		if (fcntl(socksfd, F_SETFL, flags & ~O_NONBLOCK)) return -1;

connect_again:
	if ((*real_connect)(socksfd, (struct sockaddr *)&socks, sizeof(socks)))
	{
		if (errno==EINTR) {sleep(1); goto connect_again;}
/*              Non-blocking mode not implemented yet */
/* 		if (errno==EINPROGRESS) return -1; */
		if (errno==ECONNREFUSED)
			fprintf(stderr, "%s: Connection to SOCKS4 server "
				"failed: %s\n", progname, strerror(errno));
		return -1;
	}

	memcpy(buf+2, &dst->sin_port, 2);
	memcpy(buf+4, &dst->sin_addr, 4);
	vec[0].iov_base=buf;
	vec[0].iov_len=8;
	vec[1].iov_base=usernames[i];
	vec[1].iov_len=strlen(usernames[i])+1;
/* SOCKS4 request: version command port addr username NUL */
	if (writev_full(socksfd, vec, 2)) goto fail;
/* SOCKS4 reply: version status port addr */
	if (read_full(socksfd, buf, 8)) goto fail;
	if (buf[0]!=0) goto fail;
	switch (buf[1])
	{
	case 90:
		break;
	case 91:
		fprintf(stderr, "%s: SOCKS4 request rejected or failed\n", 
			progname);
		goto fail;
	case 92:
		fprintf(stderr, "%s: SOCKS4 request rejected becasue SOCKS "
			"server cannot connect to identd on the client\n", 
			progname);
		goto fail;
	case 93:
		fprintf(stderr, "%s: SOCKS4 request request rejected - "
			"requested userid does not match ident response\n",
			progname);
		goto fail;
	default:
		fprintf(stderr, "%s: SOCKS4 request request rejected - "
			"undefined response code %i\n", 
			progname, (unsigned char)buf[1]);
		goto fail;
	}
	/* Turn non-blocking mode back on if necessary */
	if (flags & O_NONBLOCK)
		if (fcntl(socksfd, F_SETFL, flags)) goto fail;
	errno=0;
	return 0;

fail:
	shutdown(socksfd, SHUT_RDWR);
	errno=ECONNREFUSED;
	return -1;
}

int connect(int fd, const struct sockaddr *addr, socklen_t len)
{
	int i;

	if (!real_connect) init();
	if (!live) goto out;

	/* Check that address is a candidate for redirection */
	if (addr->sa_family!=AF_INET ||
	    ((struct sockaddr_in *)addr)->sin_port==htons(22) ||
	    ((struct sockaddr_in *)addr)->sin_port==htons(53))
		goto out;

	for (i=0; i<configurations; i++)
	{
		if ((((struct sockaddr_in *)addr)->sin_addr.s_addr 
		     & netmasks[i])==networks[i])
		{
			int type;
			socklen_t typelen=sizeof(type);
			/* Check that socket type is also a candidate */
			if (getsockopt(fd, SOL_SOCKET, SO_TYPE, 
				       &type, &typelen))
				return -1;
			if (type==SOCK_STREAM)
				return socks_connect
					(fd, (struct sockaddr_in *)addr, len,
					 i);
		}
	}
	
out:
	return (*real_connect)(fd, addr, len);
}