#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <string.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <sys/time.h>
#ifdef sun
#include "/usr/include/netdb.h"	/* for gethostbyname */
#else
#include <netdb.h>		/* for gethostbyname */
#endif
#include <netinet/in.h>
#include <signal.h>
#include <sys/wait.h>
#include <errno.h>
#include <sys/resource.h>
#include <arpa/inet.h>

#define VERSION "2.1.1"

#ifndef SOMAXCONN
#define SOMAXCONN 5
#endif

#ifndef INADDR_NONE
#define INADDR_NONE 0xffffffff	/* -1 return */
#endif /* INADDR_NONE */

char* progname;
char* client_name = NULL;
char* relay_host = NULL;
int relay_port;
int debug_flag = 0;

void usage(int status)
{
    fprintf(stderr, "Usage: %s [-d] [-h] [-v] <connection_port> <relay_host:port>"
	    " [<relay_host:port> ...]\n", progname);
    exit(status);
}


#if 0
FILE* logfp = NULL;
#define LOGFILENAME "tcprelay.log"
void open_log_file(const char* fname)
{
    if(fname == NULL || !strcmp(fname, "-"))
	logfp = stdin;
    else
    {
	if((logfp = fopen(fname, "a+")) == NULL)
	{
	    perror(fname);
	    exit(1);
	}
    }
}
#endif

/* Return internet IP address in network byte order */
char** host_ip_addrs(const char* address)
{
    unsigned int addr;
    struct hostent *hp;
    static char *ipbuf[2];

    if((addr = inet_addr(address)) != INADDR_NONE) {
	ipbuf[0] = (char *)&addr;
	ipbuf[1] = NULL;
	return ipbuf;
    }
    if((hp = gethostbyname(address)) == NULL)
	return NULL;
    return hp->h_addr_list;
}

int connect_to_server(const char* hostname, int tcp_port, int cnt)
{
    char **addrs;
    int i, j, n;
    
    addrs = host_ip_addrs(hostname);
    if(addrs == NULL)
	return -1;

    for(n = 0; addrs[n] != NULL; n++)
	;

    for(j = 0; j < n; j++)
    {
	int fd;
	struct sockaddr_in in;

	i = (cnt + j) % n;
	memset(&in, 0, sizeof(in));
	in.sin_family = AF_INET;
	in.sin_port   = htons(tcp_port);
	memcpy(&in.sin_addr.s_addr, addrs[i], 4);

	if((fd = socket(AF_INET, SOCK_STREAM, 0)) < 0)
	    return -1;

	fprintf(stderr, "Try to connect to %d.%d.%d.%d\n",
		addrs[i][0]&0xff, addrs[i][1]&0xff,
		addrs[i][2]&0xff, addrs[i][3]&0xff);

	if(connect(fd, &in, sizeof(in)) < 0)
	{
	    close(fd);
	    continue;
	}

	return fd;
    }
    return -1;
}

int open_server_socket(int port)
{
    int sfd;
    struct sockaddr_in server;

    if((sfd = socket(AF_INET, SOCK_STREAM, 0)) < 0)
	return -1;

    memset(&server, 0, sizeof(server));
    server.sin_port        = htons(port);
    server.sin_family      = AF_INET;
    server.sin_addr.s_addr = htonl(INADDR_ANY);

#ifdef SO_REUSEADDR
{int on=1; setsockopt(sfd, SOL_SOCKET, SO_REUSEADDR, (caddr_t)&on, sizeof(on));}
#endif /* SO_REUSEADDR */


    if(bind(sfd, (struct sockaddr *)&server, sizeof(server)) < 0)
    {
	close(sfd);
	return -1;
    }

    /* Set it up to wait for connections. */
    if(listen(sfd, SOMAXCONN) < 0)
    {
	close(sfd);
	return -1;
    }

    return sfd;
}

int server_sock_fd = -1;

void terminate_handler(int sig)
{
    printf("\n%s is terminated. (pid=%d, sig=%d)\n", progname, getpid(), sig);
    if(server_sock_fd != -1)
	shutdown(server_sock_fd, 2);
    exit(0);
}


void sig_chld(int sig)
{
    struct rusage dummy;
    int pid, status;

    while((pid = wait3(&status, WNOHANG, &dummy)) > 0)
	printf("%d exit(0x%x)/0x%x\n",
	       pid, (status & 0xff00) >> 8, status & 0xff);
    signal(SIGCHLD, sig_chld);
}

char* sockaddr_hostname(struct sockaddr_in* addr)
{
    struct hostent *hp;
    unsigned char s[4];

    memcpy(s, &addr->sin_addr, 4);
    hp = gethostbyaddr(s, 4, AF_INET);
    if(hp == NULL)
    {
	static char buff[32];
	sprintf(buff, "%d.%d.%d.%d", s[0], s[1], s[2], s[3]);
	return buff;
    }
    return hp->h_name;
}

int wait_and_fork_client(int service_port, int *cnt)
{
    int sfd, fd, pid;
    extern int errno;

    signal(SIGPIPE, SIG_IGN);
    signal(SIGCHLD, sig_chld);
    signal(SIGINT, terminate_handler);
    signal(SIGTERM, terminate_handler);

    if((sfd = open_server_socket(service_port)) < 0)
    {
	perror(progname);
	exit(1);
    }
    server_sock_fd = sfd;

    if(debug_flag)
	printf("%s is ready. (pid=%d)\n", progname, getpid());

    *cnt = 0;
    while(1)
    {
	struct sockaddr_in in;
	int addrlen;

	addrlen = sizeof(in);
	memset(&in, 0, addrlen);

	if((fd = accept(sfd, (struct sockaddr *)&in, &addrlen)) < 0)
	{
	    if(errno == EINTR)
		continue;
	    perror(progname);
	    close(sfd);
	    exit(1);
	}

	(*cnt)++;
	if((pid = fork()) < 0)
	{
	    perror("fork");
	    close(sfd);
	    close(fd);
	    exit(1);
	}
	if(pid == 0)		/* child */
	{
	    client_name = sockaddr_hostname(&in);

	    printf("%s: connection established. ID=%d\n",
		   client_name, getpid());
	    close(sfd);
	    server_sock_fd = -1;
	    return fd;
	}

	close(fd);
    }
}

void write_quote_message(int fd, const unsigned char* buff, int n)
{
    int i;
    char s[16];
    char* p;

    for(i = 0; i < n; i++)
    {
	switch(buff[i])
	{
	  case '\n': p = "\\n\\\n"; break;
	  case '\t': p = "\\t"; break;
	  case '\v': p = "\\v"; break;
	  case '\b': p = "\\b"; break;
	  case '\r': p = "\\r"; break;
	  case '\f': p = "\\f"; break;
	  case 0x07: p = "\\a"; break;
	  case '\\': p = "\\\\"; break;
	  default:
	    p = s;
	    if(buff[i] < ' ' || buff[i] >= 127)
		sprintf(p, "\\%03o", buff[i]);
	    else
	    {
		s[0] = buff[i];
		s[1] = '\0';
	    }
	    break;
	}
	write(fd, p, strlen(p));
    }
}

void relay_proc(int fd1, int fd2)
{
    fd_set rmask;
    int lastfd;
    int cnt, n;
    char buff[BUFSIZ];

    dup2(fd1, 2);

    FD_ZERO(&rmask);
    if(fd1 < fd2)
	lastfd = fd2;
    else
	lastfd = fd1;

    while(1)
    {
	FD_SET(fd1, &rmask);
	FD_SET(fd2, &rmask);

	cnt = select(lastfd + 1, &rmask, NULL, NULL, NULL);
	if(cnt < 0)
	{
	    perror("select");
	    break;
	}

	if(FD_ISSET(fd1, &rmask))
	{
	    if((n = read(fd1, buff, BUFSIZ)) < 0)
	    {
		perror("read");
		break;
	    }
	    if(n == 0)
		break;
	    if(debug_flag)
	    {
		write(1, "C>", 2);
		write_quote_message(1, (const unsigned char *)buff, n);
		write(1, "\n", 1);
	    }
	    if(write(fd2, buff, n) < 0)
	    {
		perror("write");
		break;
	    }
	}
	if(FD_ISSET(fd2, &rmask))
	{
	    if((n = read(fd2, buff, BUFSIZ)) < 0)
	    {
		perror("read");
		break;
	    }
	    if(n == 0)
		break;
	    if(debug_flag)
	    {
		write(1, "S>", 2);
		write_quote_message(1, (const unsigned char *)buff, n);
		write(1, "\n", 1);
	    }
	    if(write(fd1, buff, n) < 0)
	    {
		perror("write");
		break;
	    }
	}
    }
    fprintf(stderr, "Connection closed.\n");
}

int str2port(char *name) {
    int port;
    if('0' <= *name && *name <= '9') {
	port = atoi(name);
    } else {
	struct servent *sp;
	if((sp = getservbyname(name, "tcp")) == NULL) {
	    fprintf(stderr, "Unknown service: %s/tcp\n", name);
	}
	port = sp->s_port;
    }
    return port;
}

char *split_hostport(char *host, int *port) {
    char *p;

    if((p = strchr(host, ':')) == NULL) {
	fprintf(stderr, "Host syntax error: %s\n", host);
	return NULL;
    }
    *p++ = '\0';
    if((*port = str2port(p)) == 0)
	return NULL;
    return host;
}

int main(int argc, char** argv)
{
    int fd, relay_fd = -1;
    char *host, *host_and_port;
    int port;
    int i;
    int open_port;
    int cnt;

    if((progname = strrchr(argv[0], '/')) != NULL)
	progname++;
    else
	progname = argv[0];

    for(i = 1; i < argc; i++)
    {
	if(!strcmp(argv[i], "-d"))
	    debug_flag = 1;
	else if(!strcmp(argv[i], "-h"))
	    usage(0);
	else if(!strcmp(argv[i], "-v"))
	{
	    fprintf(stderr, "%s %s\n", progname, VERSION);
	    return 0;
	}    
	else
	{
	    if(argv[i][0] != '-')
		break;
	    usage(1);
	}
    }

    if(i + 1 >= argc)
	usage(1);
    if((open_port = str2port(argv[i++])) == 0)
	usage(1);
    fd = wait_and_fork_client(open_port, &cnt);

    for(; i < argc; i++)
    {
	host_and_port = argv[i];
	if((host = split_hostport(host_and_port, &port)) == NULL)
	    continue;
	relay_fd = connect_to_server(host, port, cnt);
	if(relay_fd == -1)
	{
	    char* msg = "Can't connect.\n";
	    write(fd, msg, strlen(msg));
	    if(errno == ECONNRESET)
		continue;
	    shutdown(fd, 2);
	    return 1;
	}
	relay_proc(fd, relay_fd);
	shutdown(relay_fd, 2);
	shutdown(fd, 2);
	return 0;
    }
    shutdown(relay_fd, 2);
    shutdown(fd, 2);
    return 0;
}