/**
Small and simple TCP relay
**/

#include <netinet/in.h>
#include <sys/time.h>
#include <unistd.h>
#include <string.h>
#include <inttypes.h>
#include <fcntl.h>
#include <stdlib.h>
#include <sys/socket.h>
#include <stdio.h>
#include <arpa/inet.h>
#include <netdb.h>
#include <signal.h>
#include <errno.h>

static time_t timeout=3600;
static bool debug=false;

static in_port_t str2port(const char *s) {
	if(strchr("0123456789",s[0]))
		return htons(strtol(s,0,0));
	struct servent *se = getservbyname(s,"tcp");
	if(!se) return 0;
	return se->s_port;
}

static void do_endpoint(const char *endpoint_str, int *read_fd, int *write_fd) {
	if(strcmp(endpoint_str,"-")==0) {
		if(debug) printf("using stdin/stdout\n");
		// Use stdin/stdout
		*read_fd = 0;
		*write_fd = 1;
		return;
	}
	
	if(strchr(endpoint_str,':')==0) {
		if(debug) printf("listening\n");
		// No colon - listen
		int listen_fd = socket(AF_INET,SOCK_STREAM,0);
		if(listen_fd<0) {
			perror("socket()");
			exit(EXIT_FAILURE);
		}
		struct sockaddr_in sa;
		memset(&sa,0,sizeof(sa));
		sa.sin_family = AF_INET;
		sa.sin_addr.s_addr = INADDR_ANY;
		sa.sin_port = str2port(endpoint_str);
	
		int reuse=1;
		setsockopt(listen_fd, SOL_SOCKET, SO_REUSEADDR, (char*)&reuse, sizeof(reuse));
		
		if(bind(listen_fd,(sockaddr*)(void*)&sa,sizeof(sa))!=0) {
			fprintf(stderr,"bind() failed, errno=%d (%s)\n",errno,strerror(errno));
			exit(EXIT_FAILURE);
		}

		listen(listen_fd,10);
		
		int cfd = accept(listen_fd,0,0);
		
		if(cfd<0) {
			fprintf(stderr,"accept() failed, errno=%d (%s)\n",errno,strerror(errno));
			exit(EXIT_FAILURE);
		}
			
		close(listen_fd);
		
		*read_fd = cfd;
		*write_fd = cfd;
		return;
	}
	
	if(debug) printf("connect\n");
	//colon - connect
	char host[128];
	char port[128];
	const char *p=strchr(endpoint_str,':');
	memcpy(host,endpoint_str,p-endpoint_str);
	host[p-endpoint_str]='\0';
	strcpy(port,p+1);
		
	int fd = socket(AF_INET,SOCK_STREAM,0);
	if(fd<0) {
		perror("socket()");
		exit(EXIT_FAILURE);
	}
	
	struct sockaddr_in sa;
	memset(&sa,0,sizeof(sa));
	sa.sin_family = AF_INET;
	sa.sin_addr.s_addr = inet_addr(host);
	sa.sin_port = str2port(port);
	if(debug) printf("connecting to %s\n",endpoint_str);
	if(connect(fd,(sockaddr*)(void*)&sa,sizeof(sa))<0) {
		fprintf(stderr,"connect(%s) failed, errno=%d (%s)\n",endpoint_str,errno,strerror(errno));
		exit(EXIT_FAILURE);
	}
	
	*read_fd = fd;
	*write_fd = fd;
	return;
}


static void exchange(int read_fd_1, int write_fd_1, int read_fd_2, int write_fd_2) {
	if(debug) printf("read1=%d, write1=%d, read2=%d, write2=%d\n",read_fd_1,write_fd_1,read_fd_2,write_fd_2);
	int maxfd=-1;
	
	fd_set rfds;
	FD_ZERO(&rfds);
	FD_SET(read_fd_1,&rfds);
	if(read_fd_1>maxfd) maxfd=read_fd_1;
	FD_SET(read_fd_2,&rfds);
	if(read_fd_2>maxfd) maxfd=read_fd_2;
	
	fd_set wfds;
	FD_ZERO(&wfds);
	FD_CLR(write_fd_1,&wfds);
	if(write_fd_1>maxfd) maxfd=write_fd_1;
	FD_CLR(write_fd_2,&wfds);
	if(write_fd_2>maxfd) maxfd=write_fd_2;
	
	char buf_1_to_2[4096];
	size_t bytes_1_to_2=0;
	char buf_2_to_1[4096];
	size_t bytes_2_to_1=0;
	
	for(;;) {
		struct timeval tv={3600,0};
		fd_set tmp_rfds=rfds;
		fd_set tmp_wfds=wfds;
		if(debug) {
			printf("testing readability of: ");
			for(int i=0; i<=maxfd; i++) if(FD_ISSET(i,&tmp_rfds)) printf("%d ",i);
			printf("\n");
			printf("testing writability of: ");
			for(int i=0; i<=maxfd; i++) if(FD_ISSET(i,&tmp_wfds)) printf("%d ",i);
			printf("\n");
		}
		if(debug) { printf("selecting..."); fflush(stdout); }
		int rc = select(maxfd+1, &tmp_rfds, &tmp_wfds, 0, &tv);
		if(debug) printf("select returned %d\n",rc);
		if(rc<=0) return;
		if(FD_ISSET(read_fd_1,&tmp_rfds)) {
			if(debug) printf("read_fd_1 is ready\n");
			long b = read(read_fd_1, buf_1_to_2+bytes_1_to_2, sizeof(buf_1_to_2)-bytes_1_to_2);
			if(b<=0) return;
			if(bytes_1_to_2==0)
				FD_SET(write_fd_2,&wfds);
			bytes_1_to_2 += b;
			if(bytes_1_to_2==sizeof(buf_1_to_2))
				FD_CLR(read_fd_1,&rfds);
		}
		if(FD_ISSET(read_fd_2,&tmp_rfds)) {
			if(debug) printf("read_fd_2 is ready\n");
			long b = read(read_fd_2, buf_2_to_1+bytes_2_to_1, sizeof(buf_2_to_1)-bytes_2_to_1);
			if(b<=0) return;
			if(bytes_2_to_1==0)
				FD_SET(write_fd_1,&wfds);
			bytes_2_to_1 += b;
			if(bytes_2_to_1==sizeof(buf_2_to_1))
				FD_CLR(read_fd_2,&rfds);
		}
		if(FD_ISSET(write_fd_1,&tmp_wfds)) {
			if(debug) printf("write_fd_1 is ready\n");
			long b = write(write_fd_1, buf_2_to_1, bytes_2_to_1);
			if(b<=0) return;
			if(bytes_2_to_1==sizeof(buf_2_to_1))
				FD_SET(read_fd_2,&rfds);
			memmove(buf_2_to_1,buf_2_to_1+b,bytes_2_to_1-b);
			bytes_2_to_1 -= b;
			if(bytes_2_to_1==0)
				FD_CLR(write_fd_1,&wfds);
		}
		if(FD_ISSET(write_fd_2,&tmp_wfds)) {
			if(debug) printf("write_fd_2 is ready\n");
			long b = write(write_fd_2, buf_1_to_2, bytes_1_to_2);
			if(b<=0) return;
			if(bytes_1_to_2==sizeof(buf_1_to_2))
				FD_SET(read_fd_1,&rfds);
			memmove(buf_1_to_2,buf_1_to_2+b,bytes_1_to_2-b);
			bytes_1_to_2 -= b;
			if(bytes_1_to_2==0)
				FD_CLR(write_fd_2,&wfds);
		}
	}
}


int main(int argc, char **argv) {
	int c;
	while((c=getopt(argc,argv,"t:d"))!=EOF) {
		switch(c) {
			case 't':
				timeout = (time_t)strtol(optarg,0,0);
				break;
			case 'd':
				debug=true;
				break;
		}
	}

	if(argc-optind!=2) {
		fprintf(stderr,"%s: 2 non-option arguments expected\n",argv[0]);
		fprintf(stderr,"Usage: %s [-td] <endpoint-1> <endpoint-2>\n"
			       " -t<timeout>   specifies timeout. default is 3600 seconds\n"
			       " -d            enables debugging\n"
		               "Endpoint is:\n"
			       "  <port>           for listen operation\n"
			       "  <ip-addr>:<port> for connect operation\n"
			       "   -               for use of stdin/stdout (inetd)\n"
		       ,argv[0]
		       );
		return EXIT_FAILURE;
	}
	
	if(debug) printf("setting endpoint 1 (%s)\n",argv[optind]);
	int read_fd_1;
	int write_fd_1;
	do_endpoint(argv[optind],&read_fd_1,&write_fd_1);
	
	if(debug) printf("setting endpoint 2 (%s)\n",argv[optind+1]);
	int read_fd_2;
	int write_fd_2;
	do_endpoint(argv[optind+1],&read_fd_2,&write_fd_2);
	
	exchange(read_fd_1,write_fd_1,read_fd_2,write_fd_2);
	
	return EXIT_SUCCESS;
}