#include <u.h>
#include <libc.h>
#include <thread.h>

enum
{
	STACK = 32768,
	BUFSIZE = 4096
};

typedef struct Conn Conn;

struct Conn {
	int lcfd;
	char ldir[40];
};

int verbose;
int nproc;
char *laddr;
char *raddr;

char adir[40];

void mainproc(void*);
void connproc(void*);
void listenproc(void*);
void proxyproc(void*);
void rwproc(void*);
int* mkfd2(int, int);

void
usage(void)
{
	fprint(2, "usage: tcpproxy [-v] [-n nproc] laddr lport raddr rport\n");
	threadexitsall("usage");
}

void
threadmain(int argc, char **argv)
{
	nproc = 1;
	verbose = 0;

	ARGBEGIN{
	default:
		usage();
	case 'n':
		nproc = atoi(EARGF(usage()));
		break;
	case 'v':
		verbose++;
		break;
	}ARGEND

	if(argc != 4){
		usage();
	}

	laddr = smprint("tcp!%s!%d", argv[0], atoi(argv[1]));
	raddr = smprint("tcp!%s!%d", argv[2], atoi(argv[3]));

	proccreate(mainproc, nil, STACK);
}

void
mainproc(void *v)
{
	int acfd, i;

	if((acfd = announce(laddr, adir)) < 0){
		fprint(2, "announce: %s: %r\n", laddr);
		threadexitsall(nil);
	}

	for(i=0; i<nproc; i++)
		proccreate(listenproc, nil, STACK);

	threadexits(0);
}

void
listenproc(void *v)
{
	Conn *c;

	for(;;){
		c = mallocz(sizeof(Conn), 0);
		c->lcfd = listen(adir, c->ldir);
		if(c->lcfd < 0){
			fprint(2, "listen %s: %r\n", c->ldir);
			threadexitsall(nil);
		}
		proccreate(connproc, (void*)(uintptr)c, STACK);
	}
}

void
connproc(void *v)
{
	Conn *c;
	NetConnInfo *nci;
	int dfd;

	c = v;
	dfd = accept(c->lcfd, c->ldir);
	if(dfd < 0){
		fprint(2, "accept %s: %r\n", c->ldir);
		threadexits(nil);
	}

	close(c->lcfd);


	if(verbose){
		nci = getnetconninfo(nil, dfd);
		if(nci == nil)
			fprint(2, "can't get connection info\n");
		else{
			fprint(2, "connection from %s to %s\n", nci->raddr, nci->laddr);
			freenetconninfo(nci);
		}
	}

	proccreate(proxyproc, (void*)(uintptr)dfd, STACK);
}

void
proxyproc(void *v)
{
	int fd, remotefd;

	fd = (int)(uintptr)v;
	if((remotefd = dial(raddr, nil, nil, nil)) < 0){
		fprint(2, "dial: %s: %r\n", raddr);
		close(fd);
		threadexits(nil);
	}

	if(verbose)
		fprint(2, "proxy to %s\n", raddr);

	proccreate(rwproc, mkfd2(fd, remotefd), STACK);
	proccreate(rwproc, mkfd2(remotefd, fd), STACK);
}

void
rwproc(void *v)
{
	int *a, rfd, wfd, n;
	char buf[BUFSIZE];

	a = v;
	rfd = a[0];
	wfd = a[1];
	free(a);

	while((n = read(rfd, buf, sizeof buf)) > 0)
		write(wfd, buf, n);

#ifdef PLAN9PORT
	shutdown(wfd, 1);
#else
	close(wfd);
#endif
	close(rfd);
}

int*
mkfd2(int fd1, int fd2)
{
	int *a;
	
	a = malloc(2*sizeof a[0]);
	if(a == 0){
		fprint(2, "out of memory\n");
		abort();
	}
	a[0] = fd1;
	a[1] = fd2;
	return a;
}