#include <sys/types.h>
#include <unistd.h>
#include <sys/param.h>
#include <netdb.h>
#include "uint_t.h"
#include "str.h"
#include "byte.h"
#include "fmt.h"
#include "scan.h"
#include "ip.h"
#include "fd.h"
#include "exit.h"
#include "env.h"
#include "prot.h"
#include "open.h"
#include "wait.h"
#include "genalloc.h"
#include "stralloc.h"
#include "alloc.h"
#include "buffer.h"
#include "logmsg.h"
#include "getoptb.h"
#include "pathexec.h"
#include "socket_if.h"
#include "ndelay.h"
#include "remoteinfo.h"
#include "rules.h"
#include "sig.h"
#include "dnsresolv.h"

#define WHO "tcpserver"

int verbosity = 1;
int flagkillopts = 1;
int flagdelay = 1;
char *banner = "";
int flagremoteinfo = 0;
int flagremotehost = 1;
int flagparanoid = 0;
unsigned long timeout = 26;
uint32 netif = 0;
int flagdualstack = 0;

static stralloc tcpremoteinfo;

uint16 localport;
char localportstr[FMT_ULONG];
char localip[16];
char localipstr[IP6_FMT];
static stralloc localhostsa;
char *localhost = 0;
const char *thishost = "0.0.0.0";

uint16 remoteport;
char remoteportstr[FMT_ULONG];
char remoteip[16];
char remoteipstr[IP6_FMT];
static stralloc remotehostsa;
char *remotehost = 0;

static char strnum[FMT_ULONG];
static char strnum2[FMT_ULONG];
static char strnum3[FMT_ULONG];

static stralloc tmp;
static stralloc fqdn;
static stralloc addresses;

unsigned long limit = 40;
unsigned long numchildren = 0;
unsigned long ipchildren = 0;
unsigned long maxconip = 0;

char bspace[16];
buffer bo;

void drop_nomem(void)
{
  logmsg(WHO,111,FATAL,"out of memory");
}

/* ---------------------------- per ip limit */

struct child {
  char ipaddr[16];
  uint32 num;
};

GEN_ALLOC_typedef(child_alloc,struct child,c,len,a)
GEN_ALLOC_readyplus(child_alloc,struct child,c,len,a,i,n,x,24,child_readyplus)
GEN_ALLOC_append(child_alloc,struct child,c,len,a,i,n,x,24,child_readyplus,child_append)

child_alloc children = {0};

void ipchild_append(char ip[16],unsigned long n)
{
  struct child *ipchild = 0;
  int i;

  for (i = 0; i <= n; ++i) {
    ipchild = &children.c[i];
    if (byte_equal(ipchild->ipaddr,16,ip)) {
      ++ipchild->num;
      break;
    } else {
      byte_copy(ipchild->ipaddr,16,ip);
      ++ipchild->num;
      break; 
    }
  }
}

void ipchild_clear(char ip[16])
{ 
  struct child *ipchild = 0;
  int i; 

  for (i = 0; i <= children.len; ++i) {
    ipchild = &children.c[i];
    if (byte_equal(ipchild->ipaddr,16,ip)) {
      if (ipchild->num) --ipchild->num;
      break;
    }
  }
}

int ipchild_limit(char ip[16],unsigned long n)
{
  int i;

  for (i = 0; i <= n; ++i)
    if (byte_equal(children.c[i].ipaddr,16,ip)) 
      return children.c[i].num;

  return 0;
}
       
/* ---------------------------- child */

int flagdeny = 0;
int flagallownorules = 0;
char *fnrules = 0;

void cats(char *s)
{
  if (!stralloc_cats(&tmp,s)) drop_nomem();
}
void append(char *ch)
{
  if (!stralloc_append(&tmp,ch)) drop_nomem();
}
void safecats(char *s)
{
  char ch;
  int i;

  for (i = 0; i < 100; ++i) {
    ch = s[i];
    if (!ch) return;
    if (ch < 33) ch = '?';
    if (ch > 126) ch = '?';
    if (ch == '%') ch = '?'; /* logger stupidity */
    append(&ch);
  }
  cats("...");
}
void env(const char *s,const char *t)
{
  if (!pathexec_env(s,t)) drop_nomem();
}
void drop_rules(void)
{
  logmsg(WHO,110,DROP,B("unable to read: ",fnrules));
}

void found(char *data,unsigned int datalen)
{
  unsigned int next0;
  unsigned int split;

  while ((next0 = byte_chr(data,datalen,0)) < datalen) {
    switch(data[0]) {
      case 'D':
        flagdeny = 1;
        break;
      case '+':
        split = str_chr(data + 1,'=');
        if (data[1 + split] == '=') {
          data[1 + split] = 0;
          env(data + 1,data + 1 + split + 1);
          if (!str_diff(data + 1,"MAXCONIP")) {
            scan_ulong(data + 1 + split + 1,&maxconip);
            if (limit && maxconip > limit) maxconip = limit;
            if (ipchildren >= maxconip) { flagdeny = 2; }
          }
        }
        break;
    }
    ++next0;
    data += next0; datalen -= next0;
  }
}

void doit(int t)
{
  int j;
 
  if (socket_local(t,localip,&localport,&netif) == -1)
    logmsg(WHO,111,FATAL,"unable to get local address");
  if (flagkillopts) 
    socket_ipoptionskill(t);
  if (!flagdelay)
    socket_tcpnodelay(t);

  if (ip6_isv4mapped(remoteip)) {
    remoteipstr[ip4_fmt(remoteipstr,remoteip + 12)] = 0;
    localipstr[ip4_fmt(localipstr,localip + 12)] = 0;
  } else {
    remoteipstr[ip6_fmt(remoteipstr,remoteip)] = 0;
    localipstr[ip6_fmt(localipstr,localip)] = 0;
  }

  if (verbosity >= 2) {
    strnum[fmt_ulong(strnum,getpid())] = 0;
    log_who(WHO,B("pid ",strnum," from ",remoteipstr));
  }

  if (*banner) {
    buffer_init(&bo,write,t,bspace,sizeof(bspace));
    if (buffer_putsflush(&bo,banner) == -1)
      logmsg(WHO,111,FATAL,"unable to print banner");
  }

  if (!localhost)
    if (dns_name(&localhostsa,localip) >= 0) 
      if (localhostsa.len) {
        if (!stralloc_0(&localhostsa)) drop_nomem();
        localhost = localhostsa.s;
      }

  remoteportstr[fmt_ulong(remoteportstr,remoteport)] = 0;

  if (flagremotehost)
    if (dns_name(&remotehostsa,remoteip) >= 0)
      if (remotehostsa.len) {
        if (flagparanoid) {
          if (dns_ip6(&tmp,&remotehostsa) >= 0)
            for (j = 0; j + 16 <= tmp.len; j += 16)
              if (byte_equal(remoteip,16,tmp.s + j)) {
                flagparanoid = 0;
                break;
              }
          if (dns_ip4(&tmp,&remotehostsa) >= 0)
            for (j = 0; j + 4 <= tmp.len; j += 4)
              if (byte_equal(remoteip + 12,4,tmp.s + j)) {
                flagparanoid = 0;
                break;
              }
        }
        if (!flagparanoid) {
          if (!stralloc_0(&remotehostsa)) drop_nomem();
            remotehost = remotehostsa.s;
        }
      }

  if (flagremoteinfo) {
    if (remoteinfo(&tcpremoteinfo,remoteip,remoteport,localip,localport,timeout,netif) == -1)
      flagremoteinfo = 0;
    if (!stralloc_0(&tcpremoteinfo)) drop_nomem();
  }

  if (fnrules) {
    int fdrules;
    fdrules = open_read(fnrules);
    if (fdrules == -1) {
      if (errno != ENOENT) drop_rules();
      if (!flagallownorules) drop_rules();
    }
    else {
      if (rules(found,fdrules,remoteipstr,remotehost,flagremoteinfo ? tcpremoteinfo.s : 0) == -1) drop_rules();
      close(fdrules);
    }
  }

  if (verbosity >= 2) {
    strnum[fmt_ulong(strnum,getpid())] = 0;
    strnum2[fmt_ulong(strnum2,maxconip)] = 0;
    if (!stralloc_copys(&tmp,"tcpserver: ")) drop_nomem();
    safecats(flagdeny ? "deny" : "ok");
    cats(" "); safecats(strnum);
    cats(" "); if (localhost) safecats(localhost);
    cats(":"); safecats(localipstr);
    cats(":"); safecats(localportstr);
    cats(" "); if (remotehost) safecats(remotehost);
    cats(":"); safecats(remoteipstr);
    cats(":"); if (flagremoteinfo) safecats(tcpremoteinfo.s);
    cats(":"); safecats(remoteportstr);
    if (flagdeny == 2) { cats(" ip connection limit:"); cats(strnum2); cats(" exceeded"); }
    cats("\n");
    buffer_putflush(buffer_2,tmp.s,tmp.len);
  }

  if (flagdeny) _exit(100);


  /* Set up environment late */

  env("PROTO",ip6_isv4mapped(remoteip)? "TCP":"TCP6");
  env("TCPLOCALIP",localipstr);
  env("TCPLOCALPORT",localportstr);
  env("TCPLOCALHOST",localhost);
  env("TCPREMOTEIP",remoteipstr);
  env("TCPREMOTEPORT",remoteportstr);
  env("TCPREMOTEHOST",remotehost);
  if (!ip6_isv4mapped(remoteip)) {
    env("TCP6LOCALIP",localipstr);
    env("TCP6LOCALHOST",localhost);
    env("TCP6LOCALPORT",localportstr);
    if (netif) env("TCP6INTERFACE",socket_getifname(netif));
    env("TCP6REMOTEIP",remoteipstr);
    env("TCP6REMOTEPORT",remoteportstr);
    env("TCP6REMOTEHOST",remotehost);
  }
  env("TCPREMOTEINFO",flagremoteinfo ? tcpremoteinfo.s : 0);

}


/* ---------------------------- parent */

void usage(void)
{
  logmsg(WHO,100,USAGE,"tcpserver \
[ -46UxXpPhHrRoOdDqQv ] \
[ -c limit ] \
[ -x rules.cdb ] \
[ -B banner ] \
[ -g gid ] \
[ -u uid ] \
[ -b backlog ] \
[ -l localname ] \
[ -t timeout ] \
[ -I interface ] \
host port program");
}

int flag1 = 0;
unsigned long backlog = 20;
unsigned long uid = 0;
unsigned long gid = 0;

void printstatus(void)
{
  if (verbosity < 2) return;
  strnum[fmt_ulong(strnum,numchildren)] = 0;
  strnum2[fmt_ulong(strnum2,limit)] = 0;
  strnum3[fmt_ulong(strnum3,maxconip)] = 0;
  log_who(WHO,B("status: ",strnum,"/",strnum2,"/",strnum3));
}

void sigterm(void)
{
  _exit(0);
}

void sigchld(void)
{
  int wstat;
  int pid;
 
  while ((pid = wait_nohang(&wstat)) > 0) {
    if (verbosity >= 2) {
      strnum[fmt_ulong(strnum,pid)] = 0;
      strnum2[fmt_ulong(strnum2,wstat)] = 0;
      log_who(WHO,B("end ",strnum," status ",strnum2));
    }
    if (maxconip) ipchild_clear(remoteip);
    if (numchildren) --numchildren; 
    printstatus();
  }
}

int main(int argc,char **argv)
{
  const char *hostname;
  int opt;
  struct servent *se;
  char *x;
  unsigned long u;
  int j;
  int s;
  int t;
  int ipflag = 0;

  while ((opt = getopt(argc,argv,"146dDvqQhHrRUXx:t:u:g:l:b:B:c:I:pPoO")) != opteof) {
    switch(opt) {
      case '1': flag1 = 1; break;
      case '4': ipflag = 1; break;
      case '6': ipflag = 2; break;
      case 'd': flagdelay = 1; break;
      case 'D': flagdelay = 0; break;
      case 'v': verbosity = 2; break;
      case 'q': verbosity = 0; break;
      case 'Q': verbosity = 1; break;
      case 'h': flagremotehost = 1; break;
      case 'H': flagremotehost = 0; break;
      case 'r': flagremoteinfo = 1; break;
      case 'R': flagremoteinfo = 0; break;
      case 'U': x = env_get("UID"); if (x) scan_ulong(x,&uid);
                x = env_get("GID"); if (x) scan_ulong(x,&gid); break;
      case 'x': fnrules = optarg; break;
      case 'X': flagallownorules = 1; break;
      case 't': scan_ulong(optarg,&timeout); break;
      case 'u': scan_ulong(optarg,&uid); break;
      case 'g': scan_ulong(optarg,&gid); break;
      case 'l': localhost = optarg; break;
      case 'b': scan_ulong(optarg,&backlog); break;
      case 'B': banner = optarg; break;
      case 'c': scan_ulong(optarg,&limit); break;
      case 'I': netif = socket_getifidx(optarg); break;
      case 'p': flagparanoid = 1; break;
      case 'P': flagparanoid = 0; break;
      case 'o': flagkillopts = 0; break;
      case 'O': flagkillopts = 1; break;
      default: usage();
    }
  }
  argc -= optind;
  argv += optind;

  if (!verbosity)
    buffer_2->fd = -1;
 
  hostname = *argv++;
  if (!hostname || str_equal((char *)hostname,"")) usage();
  if (str_equal((char *)hostname,"0")) hostname = thishost;
  else if (str_equal((char *)hostname,":0")) {
    ipflag = 2;
    flagdualstack = 1;
    hostname = "::";
  }

  x = *argv++;
  if (!x) usage();
  if (!x[scan_ulong(x,&u)])
    localport = u;
  else {
    se = getservbyname(x,"tcp");
    if (!se)
      logmsg(WHO,111,FATAL,B("unable to figure out port number for: ",x));
    uint16_unpack_big((char*)&se->s_port,&localport);
  }

  if (!*argv) usage();

  if ((x = env_get("MAXCONIP"))) { scan_ulong(x,&u); maxconip = u; }
  if (limit && maxconip > limit) maxconip = limit;
  if (!child_readyplus(&children,limit)) drop_nomem();

  sig_block(sig_child);
  sig_catch(sig_child,sigchld);
  sig_catch(sig_term,sigterm);
  sig_ignore(sig_pipe);

  /* IP address only */

  if (ip4_scan(hostname,localip)) {
    if (!stralloc_copyb(&addresses,(char *)V4mappedprefix,12)) drop_nomem();
    if (!stralloc_catb(&addresses,localip,4)) drop_nomem();
    byte_copy(localip,16,addresses.s);
  } else if (ip6_scan(hostname,localip)) {
      if (!stralloc_copyb(&addresses,localip,16)) drop_nomem();
      byte_copy(localip,16,addresses.s);
  }

  /* Asynchronous DNS IPv4/IPv6 Name qualification */

  if (!addresses.len) {
    if (!stralloc_copys(&tmp,hostname)) drop_nomem();
    if (dns_ip_qualify(&addresses,&fqdn,&tmp) < 0)
      logmsg(WHO,111,FATAL,B("temporarily unable to figure out IP address for: ",(char *)hostname));

    byte_copy(localip,16,addresses.s);

    for (j = 0; j < addresses.len; j += 16) {  // Select best matching IP address
      if (ipflag == 1 && !ip6_isv4mapped(addresses.s + j)) continue;
      if (ipflag == 2 && ip6_isv4mapped(addresses.s + j)) continue;
      byte_copy(localip,16,addresses.s + j);
    }

  } 
  if (addresses.len < 16)
    logmsg(WHO,111,FATAL,B("no IP address for: ",(char *)hostname));

  if (ip6_isv4mapped(localip))
    s = socket_tcp4();
  else
    s = socket_tcp6();
  if (s == -1)
    logmsg(WHO,111,FATAL,"unable to create socket");
  if (flagdualstack)
    socket_dualstack(s);
  if (socket_bind_reuse(s,localip,localport,netif) == -1)
    logmsg(WHO,111,FATAL,"unable to bind");
  if (socket_local(s,localip,&localport,&netif) == -1)
    logmsg(WHO,111,FATAL,"unable to get local address");
  if (socket_listen(s,backlog) == -1)
    logmsg(WHO,111,FATAL,"unable to listen");
  ndelay_off(s);

  /* Swap user and permissions */

  if (gid) if (prot_gid(gid) == -1)
    logmsg(WHO,111,FATAL,"unable to set gid");
  if (uid) if (prot_uid(uid) == -1)
    logmsg(WHO,111,FATAL,"unable to set uid");

  if (ip6_isv4mapped(localip))
    localipstr[ip4_fmt(localipstr,localip + 12)] = 0;
  else
    localipstr[ip6_fmt(localipstr,localip)] = 0;
  localportstr[fmt_ulong(localportstr,localport)] = 0;

  /* Initial setup */

  if (flag1) {
    buffer_init(&bo,write,1,bspace,sizeof(bspace));
    buffer_puts(&bo,localipstr);
    buffer_puts(&bo," : ");
    buffer_puts(&bo,localportstr);
    buffer_puts(&bo,"\n");
    buffer_flush(&bo);
  }

  close(0);
  close(1);
  printstatus();

  for (;;) {

    while (numchildren >= limit) sig_pause();

    sig_unblock(sig_child);
    t = socket_accept(s,remoteip,&remoteport,&netif);
    sig_block(sig_child);
    if (t == -1) continue;

    ipchildren = ipchild_limit(remoteip,numchildren);
    if (maxconip) {
      if (ipchildren >= maxconip) {
        if (ip6_isv4mapped(remoteip))
          remoteipstr[ip4_fmt(remoteipstr,remoteip + 12)] = 0;
        else 
          remoteipstr[ip6_fmt(remoteipstr,remoteip)] = 0;
        strnum[fmt_ulong(strnum,maxconip)] = 0;
        logmsg(WHO,100,WARN,B("ip connection limit of ",strnum," exceeded for: ",remoteipstr));
        close(t);
        continue;
      }
    }
    ipchild_append(remoteip,numchildren); // needs to happen in parent
    ++numchildren; 
    printstatus();

    switch(fork()) {
      case 0:
        close(s);
        doit(t);
        if ((fd_move(0,t) == -1) || (fd_copy(1,0) == -1))
          logmsg(WHO,111,FATAL,"unable to set up descriptors");
        sig_uncatch(sig_child);
        sig_unblock(sig_child);
        sig_uncatch(sig_term);
        sig_uncatch(sig_pipe);
        pathexec(argv);
        logmsg(WHO,111,FATAL,B("unable to run: ",*argv));
      case -1:
        if (maxconip) ipchild_clear(remoteip); --numchildren;
        logmsg(WHO,111,FATAL,"unable to fork");
    }
    close(t);
  }
}
