/*
socksproxy - Peter Benie <pjb1008@eng.cam.ac.uk> - 2004-02-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>

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 char *username=NULL;
static int socksport=1080;
static in_addr_t network=0;
static in_addr_t netmask=0;
static int live=0;

static void init(void)
{
	const char *portname, *networkspec;

	/* Illegal cast, but okay on gcc-linux-i386 */
	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");
	if (!username) username=getenv("USER");
	if (!username)
	{
		struct passwd *pwd=getpwuid(getuid());
		if (pwd) username=pwd->pw_name;
	}
	if (!username)
	{
		fprintf(stderr, "%s: Don't know who you are - "
			"try setting SOCKS_PROXY_USERNAME\n", progname);
		exit(1);
	}
	username=strdup(username);
	if (!username) abort();

	portname=getenv("SOCKS_PROXY_PORT");
	if (portname)
	{
		int port=atoi(portname);
		if (port>0 && port<65536) socksport=port;
	}

	networkspec=getenv("SOCKS_PROXY_NETWORK");
	if (networkspec)
	{
		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) return;
		if (a<256 && b<256 && c<256 && d<256 && e>=0 && e<=32)
		{
			network=a<<24 | b<<16 | c<<8 | d;
			netmask=(e==32)? 0: (0xffffffffL<<(32-e));
		}
		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 (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)
{
	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=htons(socksport);
	socks.sin_addr.s_addr=htonl(INADDR_LOOPBACK);
	
	/* 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=username;
	vec[1].iov_len=strlen(username)+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)
{
	if (!real_connect) init();

	/* Check that address is a candidate for redirection */
	if (live && addr->sa_family==AF_INET &&
	    ((struct sockaddr_in *)addr)->sin_port!=htons(22) &&
	    ((struct sockaddr_in *)addr)->sin_port!=htons(53) &&
	    (((struct sockaddr_in *)addr)->sin_addr.s_addr & 
	     htonl(netmask))==htonl(network))
	{
		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);
	}
	return (*real_connect)(fd, addr, len);
}