/*
 * As its sole author, I explicitly place this program in the public domain.
 * It may be used by anyone in any way for any purpose, though I would
 * appreciate credit where it's due.
 *		der Mouse, mouse@rodents.montreal.qc.ca, 1998-08-05
 */

#include <poll.h>
#include <stdio.h>
#include <netdb.h>
#include <errno.h>
#include <paths.h>
#include <stdlib.h>
#include <signal.h>
#include <unistd.h>
#include <strings.h>
#include <sys/time.h>
#include <sys/wait.h>
#include <sys/ioctl.h>
#include <arpa/inet.h>
#include <netinet/in.h>
#include <sys/socket.h>

extern const char *__progname;

typedef struct str STR;

struct str {
  unsigned int dead : 1;
  unsigned int nowait : 1;
  unsigned int noshut : 1;
  unsigned int timeout;
  struct timeval lastinput;
  int ifd;
  int ofd;
  char buf[10240];
  char *bwp;
  int nbuf;
  const char *itag;
  const char *otag;
  } ;

static const char *arg1 = 0;
static const char *arg2 = 0;

static int nowait_stdin = 0;
static int nowait_net = 0;
static int hold_stdin = 0;
static int hold_net = 0;
static int time_stdin = 0;
static int time_net = 0;
static int reuse_a = 0;
static int reuse_p = 0;
static int keepalive = 0;
static const char *srcspec = 0;
static const char *cmd = 0;
static int server = 0; /* 1 for -server1, 2 for -server */
static int protos = 0;
#define PROTO_IPV4 1
#define PROTO_IPV6 2

static int net;
static STR r;
static STR w;

#ifdef NO_INET_ATON
static int inet_aton(const char *cp, struct in_addr *inp)
{
 unsigned long int a;

 a = inet_addr(cp);
 if (a == (unsigned long int)-1UL) return(0);
 inp->s_addr = a;
 return(1);
}
#endif

static void handleargs(int ac, char **av)
{
 int skip;
 int errs;

 skip = 0;
 errs = 0;
 for (ac--,av++;ac;ac--,av++)
  { if (skip > 0)
     { skip --;
       continue;
     }
    if (**av != '-')
     { if (! arg1)
	{ arg1 = *av;
	}
       else if (! arg2)
	{ arg2 = *av;
	}
       else
	{ fprintf(stderr,"%s: extra argument `%s'\n",__progname,*av);
	  errs ++;
	}
       continue;
     }
    if (0)
     {
needarg:;
       fprintf(stderr,"%s: %s needs a following argument\n",__progname,*av);
       errs ++;
       continue;
     }
#define WANTARG() do { if (++skip >= ac) goto needarg; } while (0)
    if (!strcmp(*av,"-stdin"))
     { nowait_stdin = 1;
       continue;
     }
    if (!strcmp(*av,"-net"))
     { nowait_net = 1;
       continue;
     }
    if (!strcmp(*av,"-holdstdin"))
     { hold_stdin = 1;
       continue;
     }
    if (!strcmp(*av,"-holdnet"))
     { hold_net = 1;
       continue;
     }
    if (!strcmp(*av,"-timestdin"))
     { WANTARG();
       time_stdin = atoi(av[skip]);
       continue;
     }
    if (!strcmp(*av,"-timenet"))
     { WANTARG();
       time_net = atoi(av[skip]);
       continue;
     }
    if (!strcmp(*av,"-ra"))
     { reuse_a = 1;
       continue;
     }
    if (!strcmp(*av,"-rp"))
     { reuse_p = 1;
       continue;
     }
    if (!strcmp(*av,"-ka"))
     { keepalive = 1;
       continue;
     }
    if (!strcmp(*av,"-4"))
     { protos |= PROTO_IPV4;
       continue;
     }
    if (!strcmp(*av,"-6"))
     { protos |= PROTO_IPV6;
       continue;
     }
    if (!strcmp(*av,"-any"))
     { protos = 0;
       continue;
     }
    if (!strcmp(*av,"-src"))
     { WANTARG();
       srcspec = av[skip];
       continue;
     }
    if (!strcmp(*av,"-cmd"))
     { WANTARG();
       cmd = av[skip];
       continue;
     }
    if (!strcmp(*av,"-server"))
     { WANTARG();
       cmd = av[skip];
       server = 2;
       continue;
     }
    if (!strcmp(*av,"-server1"))
     { WANTARG();
       cmd = av[skip];
       server = 1;
       continue;
     }
#undef WANTARG
    fprintf(stderr,"%s: unrecognized option `%s'\n",__progname,*av);
    errs ++;
  }
 if (errs)
  { exit(1);
  }
}

static int proto_not_listed(int pf)
{
 if (protos == 0) return(0);
 switch (pf)
  { case PF_INET:
       if (protos & PROTO_IPV4) return(0);
       break;
    case PF_INET6:
       if (protos & PROTO_IPV6) return(0);
       break;
  }
 return(1);
}

static void setup_connect(const char *host, const char *port)
{
 struct addrinfo hints;
 struct addrinfo *ai0;
 struct addrinfo *ai;
 int e;
 int se;
 int s;
 char *ssslash;
 const char *ssh;
 const char *ssp;
 struct addrinfo *ssai0;
 struct addrinfo *ssai;
 char hnf[NI_MAXHOST];
 char pnf[NI_MAXSERV];
 char hnt[NI_MAXHOST];
 char pnt[NI_MAXSERV];

 if (srcspec == 0)
  { ssh = 0;
    ssp = 0;
  }
 else
  { char *f;
    int l;
    f = 0;
    ssslash = index(srcspec,'/');
    if (ssslash)
     { l = ssslash - srcspec;
       if (l == 0)
	{ ssh = 0;
	}
       else
	{ f = malloc(l+1);
	  bcopy(srcspec,f,l);
	  f[l] = '\0';
	  ssh = f;
	}
       ssp = ssslash + 1;
       if (! *ssp) ssp = 0;
     }
    else
     { ssh = srcspec;
       ssp = 0;
     }
  }
 if (ssh || ssp)
  { hints.ai_flags = AI_PASSIVE;
    hints.ai_family = PF_UNSPEC;
    hints.ai_socktype = SOCK_STREAM;
    hints.ai_protocol = 0;
    hints.ai_addrlen = 0;
    hints.ai_canonname = 0;
    hints.ai_addr = 0;
    hints.ai_next = 0;
    e = getaddrinfo(host,port,&hints,&ssai0);
    if (e)
     { fprintf(stderr,"%s: %s/%s: %s\n",__progname,ssh?:"",ssp?:"",gai_strerror(e));
       exit(1);
     }
    if (! ssai0)
     { fprintf(stderr,"%s: %s/%s: successful lookup but no addresses?\n",__progname,ssh?:"",ssp?:"");
       exit(1);
     }
  }
 else
  { ssai0 = 0;
  }
 hints.ai_flags = 0;
 hints.ai_family = PF_UNSPEC;
 hints.ai_socktype = SOCK_STREAM;
 hints.ai_protocol = 0;
 hints.ai_addrlen = 0;
 hints.ai_canonname = 0;
 hints.ai_addr = 0;
 hints.ai_next = 0;
 e = getaddrinfo(host,port,&hints,&ai0);
 if (e)
  { fprintf(stderr,"%s: %s/%s: %s\n",__progname,host,port,gai_strerror(e));
    exit(1);
  }
 if (! ai0)
  { fprintf(stderr,"%s: %s: successful lookup but no addresses?\n",__progname,strerror(se));
    exit(1);
  }
 se = 0;
 for (ai=ai0;ai;ai=ai->ai_next)
  { if (ssai0)
     { if (proto_not_listed(ai->ai_family)) continue;
       for (ssai=ssai0;ssai;ssai=ssai->ai_next)
	{ if (ssai->ai_family != ai->ai_family) continue;
	  s = socket(ai->ai_family,ai->ai_socktype,ai->ai_protocol);
	  if (s < 0)
	   { if (se == 0) se = errno;
	     continue;
	   }
	  se = -1;
	  if (bind(s,ssai->ai_addr,ssai->ai_addrlen) < 0)
	   { e = errno;
	     if (getnameinfo(ssai->ai_addr,ssai->ai_addrlen,&hnf[0],NI_MAXHOST,&pnf[0],NI_MAXSERV,NI_NUMERICHOST|NI_NUMERICSERV))
	      { fprintf(stderr,"%s: %s/%s: bind failed [%s], can't get numeric hostname info [%s]\n",__progname,ssh?:"",ssp?:"",strerror(e),strerror(errno));
	      }
	     else
	      { fprintf(stderr,"%s: bind to %s/%s: %s\n",__progname,&hnf[0],&pnf[0],strerror(e));
	      }
	     close(s);
	     continue;
	   }
	  if (connect(s,ai->ai_addr,ai->ai_addrlen) < 0)
	   { char *fs;
	     char *ts;
	     e = errno;
	     if (getnameinfo(ssai->ai_addr,ssai->ai_addrlen,&hnf[0],NI_MAXHOST,&pnf[0],NI_MAXSERV,NI_NUMERICHOST|NI_NUMERICSERV))
	      { asprintf(&fs,"%s/%s[%s]",ssh?:"",ssp?:"",strerror(errno));
	      }
	     else
	      { asprintf(&fs,"%s/%s",&hnf[0],&pnf[0]);
	      }
	     if (getnameinfo(ai->ai_addr,ai->ai_addrlen,&hnt[0],NI_MAXHOST,&pnt[0],NI_MAXSERV,NI_NUMERICHOST|NI_NUMERICSERV))
	      { asprintf(&ts,"%s%s%s[%s]",host,port?"/":"",port?:"",strerror(errno));
	      }
	     else
	      { asprintf(&ts,"%s/%s",&hnf[0],&pnf[0]);
	      }
	     fprintf(stderr,"%s: connect %s -> %s: %s\n",__progname,fs,ts,strerror(e));
	     free(fs);
	     free(ts);
	     close(s);
	     continue;
	   }
	  freeaddrinfo(ssai0);
	  freeaddrinfo(ai0);
	  net = s;
	  return;
	}
     }
    else
     { if (proto_not_listed(ai->ai_family)) continue;
       s = socket(ai->ai_family,ai->ai_socktype,ai->ai_protocol);
       if (s < 0)
	{ if (se == 0) se = errno;
	  continue;
	}
       se = -1;
       if (connect(s,ai->ai_addr,ai->ai_addrlen) < 0)
	{ char hnbuf[NI_MAXHOST];
	  char pnbuf[NI_MAXSERV];
	  e = errno;
	  if (getnameinfo(ai->ai_addr,ai->ai_addrlen,&hnbuf[0],NI_MAXHOST,&pnbuf[0],NI_MAXSERV,NI_NUMERICHOST|NI_NUMERICSERV))
	   { fprintf(stderr,"%s: %s%s%s: connect failed [%s], can't get numeric hostname info [%s]\n",__progname,host,port?"/":"",port?:"",strerror(e),strerror(errno));
	   }
	  else
	   { if (ai0->ai_next)
	      { fprintf(stderr,"%s: connect %s [%s/%s]: %s\n",__progname,host,&hnbuf[0],&pnbuf[0],strerror(e));
	      }
	     else
	      { fprintf(stderr,"%s: connect %s/%s: %s\n",__progname,&hnbuf[0],&pnbuf[0],strerror(e));
	      }
	   }
	  close(s);
	  continue;
	}
       freeaddrinfo(ai0);
       net = s;
       return;
     }
  }
 if (se > 0)
  { fprintf(stderr,"%s: socket: %s\n",__progname,strerror(se));
  }
 else if (se == 0)
  { fprintf(stderr,"%s: source and destination incompatible\n",__progname);
  }
 exit(1);
}

static void setup_accept(const char *port)
{
 typedef struct acc ACC;
 struct acc {
   ACC *link;
   char *txt;
   int fd;
   int px;
   } ;
 struct addrinfo hints;
 struct addrinfo *ai0;
 struct addrinfo *ai;
 struct sockaddr_storage from;
 socklen_t fromlen;
 const char *hostpart;
 ACC *alist;
 ACC *a;
 int nacc;
 struct pollfd *pfds;
 int s;
 int se;
 int i;
 int haderr;
 pid_t pid;
 pid_t pid2;
 char hn[NI_MAXHOST];
 char pn[NI_MAXSERV];

 hints.ai_flags = AI_PASSIVE;
 hints.ai_family = PF_UNSPEC;
 hints.ai_socktype = SOCK_STREAM;
 hints.ai_protocol = 0;
 hints.ai_addrlen = 0;
 hints.ai_canonname = 0;
 hints.ai_addr = 0;
 hints.ai_next = 0;
 if (srcspec)
  { char *slash;
    slash = index(srcspec,'/');
    if (slash)
     { char *t;
       i = slash - srcspec;
       t = malloc(i+1);
       bcopy(srcspec,t,i);
       t[i] = '\0';
       if (getaddrinfo(t,port,&hints,&ai0))
	{ fprintf(stderr,"%s: %s/%s: %s\n",__progname,t,port,gai_strerror(errno));
	  exit(1);
	}
       hostpart = t;
     }
    else
     { if (getaddrinfo(srcspec,port,&hints,&ai0))
	{ fprintf(stderr,"%s: %s/%s: %s\n",__progname,srcspec,port,gai_strerror(errno));
	  exit(1);
	}
       hostpart = srcspec;
     }
  }
 else
  { if (getaddrinfo(0,port,&hints,&ai0))
     { fprintf(stderr,"%s: %s: %s\n",__progname,port,gai_strerror(errno));
       exit(1);
     }
    hostpart = 0;
  }
 if (! ai0)
  { fprintf(stderr,"%s: %s: successful lookup but no addresses?\n",__progname,port);
    exit(1);
  }
 haderr = 0;
 alist = 0;
 nacc = 0;
 se = 0;
 for (ai=ai0;ai;ai=ai->ai_next)
  { if (proto_not_listed(ai->ai_family)) continue;
    s = socket(ai->ai_family,ai->ai_socktype,ai->ai_protocol);
    if (s < 0)
     { if (se == 0) se = errno;
       continue;
     }
    se = -1;
    if (getnameinfo(ai->ai_addr,ai->ai_addrlen,&hn[0],NI_MAXHOST,&pn[0],NI_MAXSERV,NI_NUMERICHOST|NI_NUMERICSERV))
     { fprintf(stderr,"%s: can't get printable form: %s\n",__progname,strerror(errno));
       exit(1);
     }
    if (reuse_a)
     { i = 1;
       if (setsockopt(s,SOL_SOCKET,SO_REUSEADDR,&i,sizeof(i)) < 0)
	{ fprintf(stderr,"%s: setsockopt SO_REUSEADDR: %s\n",__progname,strerror(errno));
	  exit(1);
	}
     }
    if (reuse_p)
     { i = 1;
       if (setsockopt(s,SOL_SOCKET,SO_REUSEPORT,&i,sizeof(i)) < 0)
	{ fprintf(stderr,"%s: setsockopt SO_REUSEPORT: %s\n",__progname,strerror(errno));
	  exit(1);
	}
     }
    if (bind(s,ai->ai_addr,ai->ai_addrlen) < 0)
     { fprintf(stderr,"%s: bind: %s\n",__progname,strerror(errno));
       close(s);
       haderr = 1;
       continue;
     }
    listen(s,10);
    a = malloc(sizeof(ACC));
    a->link = alist;
    alist = a;
    a->fd = s;
    if (hostpart)
     { asprintf(&a->txt,"%s/%s",&hn[0],&pn[0]);
     }
    else
     { asprintf(&a->txt,"/%s",&pn[0]);
     }
    nacc ++;
  }
 if (nacc < 1)
  { if (! haderr)
     { fprintf(stderr,"%s: no addresses to listen on!\n",__progname);
     }
    exit(1);
  }
 else
  { if (haderr)
     { fprintf(stderr,"%s: note: listening on other %s\n",__progname,alist->link?"addresses":"address");
     }
  }
 pfds = malloc(nacc*sizeof(struct pollfd));
 while (1)
  { for (i=0,a=alist;a;i++,a=a->link)
     { a->px = i;
       if (i >= nacc) abort();
       pfds[i].fd = a->fd;
       pfds[i].events = POLLIN | POLLRDNORM;
     }
    if (i != nacc) abort();
    i = poll(pfds,nacc,INFTIM);
    if (i < 0)
     { if (errno == EINTR) continue;
       fprintf(stderr,"%s: poll: %s\n",__progname,strerror(errno));
       exit(1);
     }
    for (a=alist;a;a=a->link)
     { if (pfds[a->px].revents & (POLLIN|POLLRDNORM))
	{ fromlen = sizeof(from);
	  i = accept(a->fd,(struct sockaddr *)&from,&fromlen);
	  if (i < 0)
	   { fprintf(stderr,"%s: accept %s: %s\n",__progname,a->txt,strerror(errno));
	     exit(1);
	   }
	  switch (server)
	   { case 0:
		pid = 0;
		break;
	     case 1:
		pid = fork();
		break;
	     case 2:
		pid = fork();
		if (pid == 0)
		 { pid2 = fork();
		   if (pid2 == 0) break;
		   exit(0);
		 }
		break;
	   }
	  if (pid == 0)
	   { for (a=alist;a;a=a->link) close(a->fd);
	     free(pfds);
	     net = i;
	     return;
	   }
	  close(i);
	  while (1)
	   { pid2 = wait(0);
	     if (pid2 == pid) break;
	     if (pid2 < 0)
	      { if (errno == EINTR) continue;
		if (errno == ECHILD) break;
		fprintf(stderr,"%s: wait: %s\n",__progname,strerror(errno));
		exit(1);
	      }
	   }
	}
     }
  }
}

static void init_str(STR *s, int rfd, int wfd, int nowait, int hold, int timeout, const char *it, const char *ot)
{
 s->dead = 0;
 s->nowait = !!nowait;
 s->noshut = !!hold;
 s->timeout = timeout;
 if (timeout) gettimeofday(&s->lastinput,0);
 s->nbuf = 0;
 s->ifd = rfd;
 s->ofd = wfd;
 s->itag = it;
 s->otag = ot;
}

static void setselect(STR *s, fd_set *rf, fd_set *wf)
{
 if (! s->dead)
  { if (s->nbuf > 0)
     { FD_SET(s->ofd,wf);
     }
    else
     { FD_SET(s->ifd,rf);
     }
  }
}

static void doshut(STR *s)
{
 if (shutdown(s->ifd,0) < 0) close(s->ifd);
 if (!s->noshut && (shutdown(s->ofd,1) < 0)) close(s->ofd);
 s->dead = 1;
}

static void tryio(STR *s, fd_set *rf, fd_set *wf)
{
 if (s->nbuf > 0)
  { if (FD_ISSET(s->ofd,wf))
     { int n;
       n = write(s->ofd,s->bwp,s->nbuf);
       if (n < 0)
	{ switch (errno)
	   { case EINTR:
#ifdef EWOULDBLOCK
	     case EWOULDBLOCK:
#endif
#if defined(EAGAIN) && (!defined(EWOULDBLOCK) || (EWOULDBLOCK != EAGAIN))
	     case EAGAIN:
#endif
		return;
		break;
	     case EPIPE:
		doshut(s);
		return;
		break;
	   }
	  fprintf(stderr,"%s: write error on %s: %s\n",__progname,s->otag,strerror(errno));
	  doshut(s);
	  return;
	}
       s->nbuf -= n;
       s->bwp += n;
     }
  }
 else
  { if (FD_ISSET(s->ifd,rf))
     { int n;
       n = read(s->ifd,&s->buf[0],sizeof(s->buf));
       if (n < 0)
	{ switch (errno)
	   { case EINTR:
#ifdef EWOULDBLOCK
	     case EWOULDBLOCK:
#endif
#if defined(EAGAIN) && (!defined(EWOULDBLOCK) || (EWOULDBLOCK != EAGAIN))
	     case EAGAIN:
#endif
		return;
	   }
	  fprintf(stderr,"%s: read error on %s: %s\n",__progname,s->itag,strerror(errno));
	  doshut(s);
	  return;
	}
       if (n == 0)
	{ doshut(s);
	  return;
	}
       s->bwp = &s->buf[0];
       s->nbuf = n;
       if (s->timeout) gettimeofday(&s->lastinput,0);
     }
  }
}

static void set_nonblocking(int fd)
{
 static int on = 1;

 ioctl(fd,FIONBIO,&on);
}

static __inline__ int tvcmp(struct timeval a, struct timeval b)
{
 if (a.tv_sec < b.tv_sec) return(-1);
 if (a.tv_sec > b.tv_sec) return(1);
 if (a.tv_usec < b.tv_usec) return(-1);
 if (a.tv_usec > b.tv_usec) return(1);
 return(0);
}

static void run_cmd(void) __attribute__((__noreturn__));
static void run_cmd(void)
{
 const char *shell;

 if (net != 0)
  { dup2(net,0);
    close(net);
  }
 dup2(0,1);
 shell = getenv("SHELL");
 if (shell == 0) shell = _PATH_BSHELL;
 execl(shell,shell,"-c",cmd,(char *)0);
 fprintf(stderr,"%s: can't exec %s for command: %s\n",__progname,shell,strerror(errno));
 exit(1);
}

int main(int, char **);
int main(int ac, char **av)
{
 signal(SIGPIPE,SIG_IGN);
 handleargs(ac,av);
 if (! arg1)
  { fprintf(stderr,"Usage: %s port [for accept mode]\n",__progname);
    fprintf(stderr,"   or: %s host port [for connect mode]\n",__progname);
    fprintf(stderr,"flags: -stdin, -net, -holdstdin, -holdnet,\n");
    fprintf(stderr,"       -src [addr][:[port]]\n");
    exit(1);
  }
 if (arg2)
  { if (server != 0)
     { fprintf(stderr,"%s: -server or -server1 work only when accepting\n",__progname);
       exit(1);
     }
    setup_connect(arg1,arg2);
  }
 else
  { setup_accept(arg1);
  }
 if (keepalive)
  { int v;
    v = 1;
    if (setsockopt(net,SOL_SOCKET,SO_KEEPALIVE,&v,sizeof(v)) < 0)
     { fprintf(stderr,"%s: setsockopt SO_KEEPALIVE: %s\n",__progname,strerror(errno));
     }
  }
 if (cmd) run_cmd();
 set_nonblocking(1);
 set_nonblocking(net);
 init_str(&r,0,net,nowait_stdin,hold_stdin,time_stdin,"stdin","network");
 init_str(&w,net,1,nowait_net,hold_net,time_net,"network","stdout");
 while (1)
  { fd_set rfds;
    fd_set wfds;
    struct timeval now;
    struct timeval timeout;
    struct timeval *tvp;
    static int set_timeout(STR *s)
     { struct timeval tv;
       if (s->dead || !s->timeout) return(0);
       if (now.tv_sec == 0) gettimeofday(&now,0);
       tv = s->lastinput;
       tv.tv_sec += s->timeout;
       if (!tvp || (tvcmp(tv,timeout) < 0))
	{ if (tvcmp(tv,now) <= 0)
	   { doshut(s);
	     return(1);
	   }
	  timeout = tv;
	  tvp = &timeout;
	}
       return(0);
     }
    if ((r.dead||r.nowait) && (w.dead||w.nowait)) break;
    FD_ZERO(&rfds);
    FD_ZERO(&wfds);
    setselect(&r,&rfds,&wfds);
    setselect(&w,&rfds,&wfds);
    now.tv_sec = 0;
    tvp = 0;
    if (set_timeout(&r) || set_timeout(&w)) continue;
    if (tvp)
     { if (now.tv_usec > timeout.tv_usec)
	{ timeout.tv_usec += 1000000 - now.tv_usec;
	  timeout.tv_sec -= now.tv_sec + 1;
	}
       else
	{ timeout.tv_usec -= now.tv_usec;
	  timeout.tv_sec -= now.tv_sec;
	}
     }
    if (select(net+1,&rfds,&wfds,0,tvp) < 0)
     { if (errno == EINTR) continue;
       fprintf(stderr,"%s: select: %s\n",__progname,strerror(errno));
       exit(1);
     }
    tryio(&r,&rfds,&wfds);
    tryio(&w,&rfds,&wfds);
  }
 exit(0);
}