/*
 * rum - version 0.9
 * LICENSE: beerware
 * author: stano@websupport.sk (https://websupport.sk/~stanojr/)
 * one process tcp redirector with sockfile support, multiple listen ports and with nice statistics ! :)
 * linux 2.6 only because we use epoll syscall
 *
 * usage:
 * ./rum -s tcp:host:port [-s tcp:host:port [-s sock:path]] -d tcp:host:port [-b] [-m tcp:host:port]
 *         -s - where to listen host:port or sockfile (host muste be some ip address from interface or 0.0.0.0 for all inerfaces)
 *         -d - destination host:port
 *         -b - goto background
 *         -m - statistics port
 *
 *
 * #./rum -s tcp:0.0.0.0:3306 -s sock:/tmp/mysql.sock -d tcp:1.2.3.4:3306 -m 127.0.0.1:666 -d
 * # telnet localhost 666          
 * Trying 127.0.0.1...
 * Connected to localhost.
 * Escape character is '^]'.
 * [              source] [                  bytes] [         destination] [         connections]
 * [    tcp:0.0.0.0:3306] [-->                   0] [    tcp:1.2.3.4:3306] [                   0]
 *                        [<--                   0]                                              
 * 
 * [sock:/tmp/mysql.sock] [-->                3424] [    tcp:1.2.3.4:3306] [                   2]
 *                        [<--              814830]                                              
 *
 * ---
 * greetz 42448 :)
 */

#define BUFLEN 16384

#define SOCK_SERVER     0
#define SOCK_STATS      1
#define SOCK_IN         2
#define SOCK_OUT        3
#define SOCK_FLUSHEND   4
#define SOCK_CONNECT  5

#define RUM_EPOLL_EVENTS 16

#include <stdio.h>
#include <stdlib.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <unistd.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <sys/un.h>
#include <sys/epoll.h>
#include <getopt.h>
#include <netdb.h>
#include <errno.h>
#include <signal.h>
#include <sys/ioctl.h>
#include <netinet/tcp.h>
#include <linux/sockios.h>

typedef struct {
  int fd;
  int type; // 0 - normal, 1 - statistics
  char *s; // string (tcp:blah:blah)

  unsigned int nr_allconn,nr_conn;
  unsigned int input_bytes,output_bytes; // input_bytes are from client->redir->target output are target->redir->client
} listeners;

struct core {
  int nl;
  listeners l[16];
} core;

struct dst {
  int t;
  struct sockaddr_in sin;
  struct sockaddr_un sun;
  int addrlen;
  char *s;
} dst;

struct ed {
  int fd;
  int op;
  int type;
  int z;
  void *ed2;
  char *bufptr;
  int len;
  char buf[BUFLEN];
};

void usage() {
  printf("./rum -s tcp:host:port [-s tcp:host:port [-s sock:path]] -d tcp:host:port [-b] [-m tcp:host:port]\n\t-s - where to listen host:port or sockfile (host muste be some ip address from interface or 0.0.0.0 for all inerfaces)\n\t-d - destination host:port\n\t-b - goto background\n\t-m - statistics port\n");
  _exit(-1);
}

int makeserver(char *wwtf) {
  int sock,sockopt;
  char *wtf;

  wtf=strdup(wwtf);

  if (strstr(wtf,"tcp:")==wtf) {
    char *host,*port,*tmp;
    struct hostent *h;
    struct sockaddr_in s;
    int p;

    wtf+=4;
    tmp=strstr(wtf,":");
    if (tmp!=NULL) {
      host=wtf;
      *tmp='\0';
      tmp++;
      port=tmp;
    } else {
      usage();
    }

    if ((h=gethostbyname(host))==NULL) {
      herror("gethostbyname");
      fflush(stdout);
      fflush(stderr);
      _exit(-1);
    }
    p=atoi(port);

    memset(&s,0,sizeof(struct sockaddr_in));
    memcpy(&s.sin_addr,h->h_addr_list[0],sizeof(in_addr_t));
    s.sin_port=htons(p);
    s.sin_family=AF_INET;

    if ((sock=socket(PF_INET,SOCK_STREAM,0))==-1) {
      perror("socket");
      _exit(-1);
    }

    sockopt = 1;
    setsockopt(sock, SOL_SOCKET, SO_REUSEADDR, &sockopt, sizeof(sockopt));


    if (bind(sock,(struct sockaddr *)&s,sizeof(struct sockaddr_in))==-1) {
      perror("bind");
      _exit(-1);
    }

    if (listen(sock,255)==-1) {
      perror("listen");
      _exit(-1);
    }

    printf("listening on tcp:%s:%s\n",host,port);

    return sock;

  } else if (strstr(wtf,"sock:")==wtf) {
    char *sockf;
    struct sockaddr_un s;

    wtf+=5;
    sockf=wtf;

    memset(&s,0,sizeof(struct sockaddr_un));
    s.sun_family=AF_UNIX;
    memcpy(s.sun_path,sockf,strlen(sockf));

    if (!access(sockf,F_OK)) {
      if (unlink(sockf)) {
        perror("unlink");
        _exit(-1);
      }
    }

    if ((sock=socket(PF_UNIX,SOCK_STREAM,0))==-1) {
      perror("socket");
      _exit(-1);
    }

    sockopt = 1;
    setsockopt(sock, SOL_SOCKET, SO_REUSEADDR, &sockopt, sizeof(sockopt));

    if (bind(sock,(struct sockaddr *)&s,sizeof(struct sockaddr_un))==-1) {
      perror("bind");
      _exit(-1);
    }

    if (listen(sock,255)==-1) {
      perror("listen");
      _exit(-1);
    }

    chmod(sockf,0777);

    printf("listening on sock:%s\n",sockf);

    return sock;

  } else {
    usage();
  }

  _exit(-1);
}

void prepareclient(char *wwtf) {
  char *wtf;

  wtf=strdup(wwtf);

  dst.s=strdup(wtf);

  if (strstr(wtf,"tcp:")==wtf) {
    char *host,*port,*tmp;
    struct hostent *h;
    int p;

    wtf+=4;
    tmp=strstr(wtf,":");
    if (tmp!=NULL) {
      host=wtf;
      *tmp='\0';
      tmp++;
      port=tmp;
    } else {
      usage();
    }

    if ((h=gethostbyname(host))==NULL) {
      herror("gethostbyname");
      _exit(-1);
    }
    p=atoi(port);

    memset(&dst.sin,0,sizeof(struct sockaddr_in));
    memcpy(&dst.sin.sin_addr,h->h_addr_list[0],sizeof(in_addr_t));
    dst.sin.sin_port=htons(p);
    dst.sin.sin_family=AF_INET;
    dst.addrlen=sizeof(struct sockaddr_in);
    dst.t='t';
  } else if (strstr(wtf,"sock:")==wtf) {
    char *sockf;

    wtf+=5;
    sockf=wtf;

    memset(&dst.sun,0,sizeof(struct sockaddr_un));
    dst.sun.sun_family=AF_UNIX;
    memcpy(dst.sun.sun_path,sockf,strlen(sockf));
    dst.addrlen=sizeof(struct sockaddr_un);
    dst.t='s';
  }
}

int makeclientsock() {
  int sock;
  //struct linger l;

  if (dst.t=='t') {
    if ((sock=socket(PF_INET,SOCK_STREAM,0))==-1) {
      perror("socket");
      return -1;
    }
  } else if (dst.t=='s') {
    if ((sock=socket(PF_UNIX,SOCK_STREAM,0))==-1) {
      perror("socket");
      return -1;
    }
  }
  fcntl(sock,F_SETFL,O_NONBLOCK);

  return sock;
}

int makeclientconnect(int sock) {
  int val;
  int valsize=sizeof(val);

  if (dst.t=='t') {
    connect(sock,(struct sockaddr *)&dst.sin,dst.addrlen);
  } else if (dst.t=='s') {
    connect(sock,(struct sockaddr *)&dst.sun,dst.addrlen);
  }

  getsockopt(sock,SOL_SOCKET,SO_ERROR,&val,&valsize);

  return val;
}

int main(int ac, char *av[]) {
  int csock,len,ret,i,z,y,ch,daemonize=0,log=0;
  struct sockaddr csa;
  int epfd;
  struct epoll_event ev,ev2,events[RUM_EPOLL_EVENTS];
  struct ed *sed,*sed2,*sed3,*ed,*ed2;
  char *logfile;

  signal(SIGPIPE,SIG_IGN);

  if (ac==1) {
    usage();
  }

  core.nl=0;

  while ((ch = getopt(ac, av, "bd:s:m:l:")) != -1) {
    switch (ch) {
      case 'b':
        daemonize=1;
      break;
      case 's':
        core.l[core.nl].s=strdup(optarg);
        core.l[core.nl].fd=makeserver(optarg);
        core.l[core.nl].type=SOCK_SERVER;
        core.l[core.nl].nr_conn=0;
        core.l[core.nl].nr_allconn=0;
        core.l[core.nl].input_bytes=0;
        core.l[core.nl].output_bytes=0;
        core.nl++;
      break;
      case 'm':
        core.l[core.nl].fd=makeserver(optarg);
        core.l[core.nl].type=SOCK_STATS;
        core.nl++;
      break;
      case 'd':
        prepareclient(optarg);
      break;
      case 'l':
        log=1;
        logfile=strdup(optarg);
      break;
    }
  }

  if (daemonize) {
    if (log) {
      daemon(0,1);
      close(0);
      close(1);
      close(2);
      ret=open(logfile,O_WRONLY|O_CREAT|O_APPEND,S_IRUSR|S_IWUSR);
      if (ret!=-1) {
        dup2(ret,1);
        dup2(ret,2);
      }
      
    } else {
      daemon(0,0);
    }
  }

  epfd=epoll_create(RUM_EPOLL_EVENTS);

  memset(&ev,0,sizeof(ev));
  for (i=0;i<core.nl;i++) {
    ed=malloc(sizeof(struct ed));
    ev.events=EPOLLIN;
    ev.data.ptr=ed;
    ed->op=EPOLLIN;
    ed->fd=core.l[i].fd;
    ed->ed2=NULL;
    if (core.l[i].type==SOCK_SERVER) {
      ed->type=SOCK_SERVER; // listener
    } else if (core.l[i].type==SOCK_STATS) {
      ed->type=SOCK_STATS; // statistics
    }

    epoll_ctl(epfd,EPOLL_CTL_ADD,core.l[i].fd,&ev);
  }

  while (1) {
    ret=epoll_wait(epfd,(struct epoll_event *)&events,RUM_EPOLL_EVENTS,-1);
    if (ret<0 && errno!=EINTR) {
      _exit(0);
    } else if (ret<0 && errno==EINTR) {
      continue;
    } else if (ret>0) {
      /* check for new connections */
      for (i=0;i<ret;i++) {
        if (events[i].data.ptr==NULL) {
          continue;
        }

        sed=events[i].data.ptr;
        sed2=sed->ed2;

        if ((events[i].events&EPOLLIN)) {
          if (sed->type==SOCK_SERVER) {
            // accept new connection
            for (z=0;z<core.nl;z++) {
              if (core.l[z].fd==sed->fd) {
                len=sizeof(csa);
                csock=accept(core.l[z].fd,&csa,&len);
                fcntl(csock,F_SETFL,O_NONBLOCK);
                memset(&ev,0,sizeof(ev));
                ed=malloc(sizeof(struct ed));
                ev.data.ptr=ed;
                ev.events=EPOLLIN;
                ed->op=ev.events;
                ed->fd=csock;
                ed->type=SOCK_IN; // in conn
                ed->z=z; // what listener
                ed->len=0;

                csock=makeclientsock();
                if (csock==-1) {
                  close(ed->fd);
                  break;
                }
                //fcntl(csock,F_SETFL,O_NONBLOCK);
                memset(&ev2,0,sizeof(ev2));
                ed2=malloc(sizeof(struct ed));
                ed->ed2=ed2;
                //epoll_ctl(epfd,EPOLL_CTL_ADD,ed->fd,&ev);

                ed2->fd=csock;
                ev2.events=EPOLLOUT;
                ed2->op=ev2.events;
                ed2->type=SOCK_CONNECT; // out conn
                ed2->z=z; // what listener
                ed2->ed2=ed;
                ev2.data.ptr=ed2;
                ed2->len=0;

                epoll_ctl(epfd,EPOLL_CTL_ADD,ed2->fd,&ev2);

                core.l[z].nr_conn++;
                core.l[z].nr_allconn++;

                break;
              }
            }

          } else if (sed->type==SOCK_STATS) {
            // accept, print statistics and close
            for (z=0;z<core.nl;z++) {
              if (core.l[z].fd==sed->fd) {
                FILE *fp;
                len=sizeof(csa);
                csock=accept(core.l[z].fd,&csa,&len);
                fp=fdopen(csock, "w");
                if (fp!=NULL) {
                  fprintf(fp,"[%20s] [   %10s] [%20s] [%15s] [%18s]\n","source","bytes","destination","all connections","actual connections");
                  for(y=0;y<core.nl;y++) {
                    if (core.l[y].type==SOCK_STATS)
                      break;
                    fprintf(fp,"[%20s] [-->%10u] [%20s] [%15u] [%18u]\n", core.l[y].s, core.l[y].input_bytes, dst.s, core.l[y].nr_allconn,core.l[y].nr_conn);
                    fprintf(fp," %20s  [<--%10u]  %20s   %15s   %18s\n\n", "", core.l[y].output_bytes,"","","");
                  }
                  fflush(fp);
                }
                close(csock);
              }
            }
          } else if (sed->type==SOCK_IN || sed->type==SOCK_OUT) {
            // copy data
            if (sed->len<=0) {
              sed->len=read(sed->fd,sed->buf,BUFLEN);
              sed->bufptr=sed->buf;
            }
            if (sed->len==0) {
              epoll_ctl(epfd,EPOLL_CTL_DEL,sed->fd,NULL);
              epoll_ctl(epfd,EPOLL_CTL_DEL,sed2->fd,NULL);
              for (z=i;z<ret;z++) {
                if (sed2==events[z].data.ptr) {
                  events[z].data.ptr=NULL;
                  break;
                }
              }
              core.l[sed->z].nr_conn--;
              close(sed->fd);
              close(sed2->fd);
              free(sed2);
              free(sed);
            } else if (sed->len==-1 && errno==EAGAIN) {
            } else if (sed->len==-1 && errno!=EAGAIN) {
              /* napr conn resed by peer*/
              epoll_ctl(epfd,EPOLL_CTL_DEL,sed->fd,NULL);
              epoll_ctl(epfd,EPOLL_CTL_DEL,sed2->fd,NULL);
              for (z=i;z<ret;z++) {
                if (sed2==events[z].data.ptr) {
                  events[z].data.ptr=NULL;
                  break;
                }
              }
              core.l[sed->z].nr_conn--;
              close(sed->fd);
              close(sed2->fd);
              free(sed2);
              free(sed);
              break;
            } else {
              if (sed->type==SOCK_IN) {
                core.l[sed->z].input_bytes+=len;
              } else if (sed->type==SOCK_OUT) {
                core.l[sed->z].output_bytes+=len;
              }
              len=write(sed2->fd,sed->bufptr,sed->len);
              if (len==-1 && errno!=EAGAIN){
                epoll_ctl(epfd,EPOLL_CTL_DEL,sed->fd,NULL);
                epoll_ctl(epfd,EPOLL_CTL_DEL,sed2->fd,NULL);
                for (z=i;z<ret;z++) {
                  if (sed2==events[z].data.ptr) {
                    events[z].data.ptr=NULL;
                    break;
                  }
                }
                core.l[sed->z].nr_conn--;
                close(sed->fd);
                close(sed2->fd);
                free(sed2);
                free(sed);
              } else if (len==-1 && errno==EAGAIN) {
                sed2->op|=EPOLLOUT;
                ev.events=sed2->op;
                ev.data.ptr=sed2;
                epoll_ctl(epfd,EPOLL_CTL_MOD,sed2->fd,&ev);
                sed->op^=EPOLLIN;
                ev.events=sed->op;
                ev.data.ptr=sed;
                epoll_ctl(epfd,EPOLL_CTL_MOD,sed->fd,&ev);
              } else if (len<sed->len) {
                sed->len-=len;
                sed->bufptr=sed->bufptr+len;
                sed2->op|=EPOLLOUT;
                ev.events=sed2->op;
                ev.data.ptr=sed2;
                epoll_ctl(epfd,EPOLL_CTL_MOD,sed2->fd,&ev);
                sed->op^=EPOLLIN;
                ev.events=sed->op;
                ev.data.ptr=sed;
                epoll_ctl(epfd,EPOLL_CTL_MOD,sed->fd,&ev);
              } else {
                sed->len=0;
              }
            }
          } else {
            printf("hmmmmm??? sed->type: %d\n",sed->type);
          }
        } else if ((events[i].events&EPOLLOUT) || (events[i].events&EPOLLHUP)) {
          if (sed->type==SOCK_IN || sed->type==SOCK_OUT) {
            if (sed2->len>0) {
              len=write(sed->fd,sed2->bufptr,sed2->len);
              if (len==-1 && errno!=EAGAIN){
                epoll_ctl(epfd,EPOLL_CTL_DEL,sed->fd,NULL);
                epoll_ctl(epfd,EPOLL_CTL_DEL,sed2->fd,NULL);
                for (z=i;z<ret;z++) {
                  if (sed2==events[z].data.ptr) {
                    events[z].data.ptr=NULL;
                    break;
                  }
                }
                core.l[sed->z].nr_conn--;
                close(sed->fd);
                close(sed2->fd);
                free(sed2);
                free(sed);
              } else if (len==-1 && errno==EAGAIN) {
              } else if (len<sed2->len) {
                sed2->len-=len;
                sed2->bufptr=sed2->bufptr+len;
              } else {
                sed2->len=0;

                sed->op^=EPOLLOUT;
                ev.events=sed->op;
                ev.data.ptr=sed;
                epoll_ctl(epfd,EPOLL_CTL_MOD,sed->fd,&ev);

                sed2->op|=EPOLLIN;
                ev.events=sed2->op;
                ev.data.ptr=sed2;
                epoll_ctl(epfd,EPOLL_CTL_MOD,sed2->fd,&ev);
              }

            } else {
              printf("uoaaaa, len is %d\n",sed2->len);
            }
          } else if (sed->type==SOCK_CONNECT) {
            errno=0;
            z=makeclientconnect(sed->fd);
            if (z==EINPROGRESS || z==EALREADY) {
              continue;
            } else if (z==0) {
              fcntl(sed->fd,F_SETFL,O_NONBLOCK);
              ev.events=EPOLLIN;
              ev.data.ptr=sed2;
              sed2->op=ev.events;
              sed2->type=SOCK_IN;
              if (epoll_ctl(epfd,EPOLL_CTL_ADD,sed2->fd,&ev)==-1) {
                epoll_ctl(epfd,EPOLL_CTL_DEL,sed->fd,NULL);
                close(sed->fd);
                close(sed2->fd);
                free(sed2);
                free(sed);
                continue;
              }
              ev.events=EPOLLIN;
              ev.data.ptr=sed;
              sed->op=ev.events;
              sed->type=SOCK_OUT;
              epoll_ctl(epfd,EPOLL_CTL_MOD,sed->fd,&ev);
            } else {
              printf("connect error: %s\n",strerror(z));
              fflush(stdout);
              epoll_ctl(epfd,EPOLL_CTL_DEL,sed->fd,NULL);
              close(sed->fd);
              close(sed2->fd);
              free(sed2);
              free(sed);
            }
          } else {
            printf("errrrrrr fd:%d\n",sed->fd);
          }
        }
      }
    }
  }
}