#include <sys/types.h>
#include <unistd.h>
#include <stdlib.h>
#include <stdarg.h>
#include <stdio.h>
#include <string.h>
#include <errno.h>
#include <pwd.h>
#include <grp.h>
#include <sys/socket.h>
#include <netdb.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <signal.h>
#include <sys/stat.h>
#include <syslog.h>
#include <time.h>

#ifndef INADDR_NONE
# define INADDR_NONE ((in_addr_t)-1)
#endif

#include "rbldnsd.h"
#include "dns.h"
#include "mempool.h"

/* TODO:
 * maybe multi-children environment
 * review all names to be consistent
 * review log/error routines to be buffer-overflow friendly
 * implement TCP mode (in separate process, to reduce number of syscalls)
 * implement AXFR query - stupid idea but AXFR is widely used
 * use config file with:
 *  - NS records
 *  - SOA records
 *  - ability to construct a zone from several files
 * maybe: AXFR client to convert DNS zones into rbldnsd fomat
 */

static unsigned ttl = 2048;	/* default record TTL */
static int recheck = 60;	/* interval between checks for reload */

/* a list of zonetypes. */
static struct zonetype {
  const char *name;	/* name of the type */
  z_queryfn *queryfn;	/* routine to perform query */
  z_loadfn *loadfn;	/* routine to load zone data */
  z_freefn *freefn;	/* routine to free zone data */
  const char *descr;	/* short description of a zone type */
} zonetypes[] = {
  /* first type is the default */
#define zonetype(t) { #t, t##_query, t##_load, t##_free, t##_descr }
  zonetype(ip4vset),
  zonetype(ip4set),
  zonetype(dnvset),
  zonetype(dnset),
#undef zonetype
  { NULL, NULL, NULL, NULL, NULL }
};

static char *progname; /* limited to 32 chars */
static int logto;
#define LOGTO_STDOUT 0x01
#define LOGTO_SYSLOG 0x02

static void error(int errnum, const char *fmt, ...) {
 /* Note: all usages of this routine should ensure
  * there is no buffer overflow here */
  char buf[128];
  int l, pl;
  va_list ap;
  pl = sprintf(buf, "%.50s: ", progname);
  va_start(ap, fmt);
  l = pl + vsprintf(buf + pl, fmt, ap);
  if (errnum)
    l += sprintf(buf + l, ": %.50s", strerror(errnum));
  if (logto & LOGTO_SYSLOG) {
    fmt = buf + pl;
    syslog(LOG_ERR, strchr(fmt, '%') ? "%s" : fmt, fmt);
  }
  buf[l++] = '\n';
  write(2, buf, l);
  _exit(1);
}

void log(int level, const char *fmt, ...) {
  va_list ap;
  va_start(ap, fmt);
  if (logto & LOGTO_SYSLOG)
    vsyslog(level, fmt, ap);
  if (logto & LOGTO_STDOUT) {
    char buf[1024];
    int l = snprintf(buf, sizeof(buf), "%s: ", progname);
    l += vsnprintf(buf + l, sizeof(buf) - l, fmt, ap);
    buf[l++] = '\n';
    write(1, buf, l);
  }
  va_end(ap);
}

struct zone {
  const char *name;	/* name of a zone */
  unsigned char *base;	/* base dn */
  unsigned blen;	/* length of base dn */
  z_queryfn *queryfn;	/* for faster access */
  const struct zonetype *type;
  time_t loaded;	/* when it was last loaded */
  struct zonedata *data; /* type-specific data */
};

static struct zone *curzone;

void oom() {
  log(LOG_ERR,
      "out of memory loading zone %s (zone will be empty)",
      curzone ? curzone->name : "(none)");
}

void zlog(int level, int lineno, const char *fmt, ...) {
  char buf[128];
  va_list ap;
  va_start(ap, fmt);
  vsnprintf(buf, sizeof(buf), fmt, ap);
  va_end(ap);
  if (lineno)
    log(level, "zone %s(%d): %s", curzone->name, lineno, buf);
  else
    log(level, "zone %s: %s", curzone->name, buf);
}

void zloaded(const char *fmt, ...) {
  char buf[128];
  va_list ap;
  struct tm *tm = gmtime(&curzone->loaded);
  va_start(ap, fmt);
  vsnprintf(buf, sizeof(buf), fmt, ap);
  va_end(ap);
  log(LOG_INFO, "zone %s %04d-%02d-%02d %02d:%02d:%02d: %s",
      curzone->name,
      tm->tm_year + 1900, tm->tm_mon + 1, tm->tm_mday,
      tm->tm_hour, tm->tm_min, tm->tm_sec,
      buf);
}

static sigset_t sigset_block; /* signals to block during zone reload */

static int loadzones(struct zone *zone) {
  int r = 1;
  struct stat st;
  FILE *f;
  sigset_t ssrestore;
  sigprocmask(SIG_BLOCK, &sigset_block, &ssrestore);
  for(; zone->name; ++zone) {
    if ((f = fopen(zone->name, "r")) == NULL ||
        fstat(fileno(f), &st) < 0) {
      log(LOG_ERR, "unable to open zonefile %s: %s",
          zone->name, strerror(errno));
      if (f) fclose(f);
      r = 0;
      continue;
    }
    if (st.st_mtime == zone->loaded) {
      fclose(f);
      continue;
    }
    curzone = zone;
    zone->type->freefn(zone->data);
    zone->loaded = st.st_mtime;
    if (!(zone->data = zone->type->loadfn(f))) {
      r = 0;
      zone->loaded = 0;
    }
    else if (ferror(f)) {
      zlog(LOG_ERR, 0, "read error: %s", strerror(errno));
      zone->type->freefn(zone->data); /*XX maybe not free it? */
      zone->data = NULL;
      zone->loaded = 0;
    }
    fclose(f);
  }
  curzone = NULL;
  sigprocmask(SIG_SETMASK, &ssrestore, NULL);
  return r;
}

static int satoi(const char *s) {
  int n = 0;
  if (*s < '0' || *s > '9') return -1;
  do n = n * 10 + (*s++ - '0');
  while (*s >= '0' && *s <= '9');
  return *s ? -1 : n;
}

static void usage() {
   int i;
   printf(
"%s: rbl dns daemon version " VERSION "\n"
"Usage is: %s [options] zone...\n"
"where options are:\n"
" -u user[:group] - run as this user:group (rbldns)\n"
" -r rootdir - chroot to this directory\n"
" -w workdir - working directory with zone files\n"
" -b [address][:port] - bind to (listen on) this address (*:53)\n"
" -t ttl - TTL value set in answers (2048)\n"
" -c check - check for file updates every `check' secs (60)\n"
" -p pidfile - write backgrounded pid to specified file\n"
" -n - do not become a daemon\n"
"each zone specified using `type:name' syntax, where type is\n"
"a type of input zone file format and name is both a name of\n"
"a zone and a name of a file to load it from.\n"
"Available zone types:\n"
, progname, progname);
   for(i = 0; zonetypes[i].name; ++i)
     printf(" %s %s\n", zonetypes[i].name, zonetypes[i].descr);
   exit(0);
}

static int init(int argc, char **argv, struct zone **zonep) {
  int c;
  char *p;
  char *user = "rbldns", *bindaddr = "";
  char *root = NULL, *workdir = NULL, *pidfile = NULL;
  FILE *fpid = NULL;
  uid_t uid = 0;
  gid_t gid = 0;
  struct sockaddr_in sin;
  int fd;
  struct zone *z;
  int nodaemon = 0;

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

  if (argc <= 1) usage();

  while((c = getopt(argc, argv, "u:r:b:w:t:c:p:nh")) != EOF)
    switch(c) {
    case 'u': user = optarg; break;
    case 'r': root = optarg; break;
    case 'b': bindaddr = optarg; break;
    case 'w': workdir = optarg; break;
    case 'p': pidfile = optarg; break;
    case 't':
      if ((c = satoi(optarg)) < 0)
        error(0, "invalid ttl (-t) value `%.50s'", optarg);
      ttl = c;
      break;
    case 'c':
      if ((c = satoi(optarg)) < 0)
        error(0, "invalid check interval (-c) value `%.50s'", optarg);
      recheck = c;
      break;
    case 'n': nodaemon = 1; break;
    case 'h': usage(); break;
    default: error(0, "type `%.50s -h' for help", progname);
    }

  if (!(argc -= optind))
    error(0, "no zone(s) to service specified (-h for help)");
  argv += optind;

  if (nodaemon)
    logto = LOGTO_STDOUT;
  else {
    tzset();
    openlog(progname, LOG_PID|LOG_NDELAY, LOG_DAEMON);
    logto = LOGTO_STDOUT | LOGTO_SYSLOG;
  }

  fd = socket(PF_INET, SOCK_DGRAM, IPPROTO_UDP);
  if (fd < 0)
    error(errno, "unable to create listening socket");
  memset(&sin, 0, sizeof(sin));
  sin.sin_family = AF_INET;
  if (bindaddr) {
    if ((p = strchr(bindaddr, ':')) != NULL)
      *p++ = '\0';
    if (*bindaddr &&
        (sin.sin_addr.s_addr = inet_addr(bindaddr)) == INADDR_NONE) {
      struct hostent *he = gethostbyname(bindaddr);
      if (!he)
        error(0, "invalid bind address specified: `%.50s'", bindaddr);
      if (he->h_addrtype != AF_INET || he->h_length != sizeof(sin.sin_addr))
        error(0, "unexpected type of listening address `%.50s'", bindaddr);
      memcpy(&sin.sin_addr, he->h_addr_list[0], sizeof(sin.sin_addr));
      endhostent();
    }
  }
  else
    p = NULL;
  if (p && *p) {
    if ((c = satoi(p)) > 0)
      ;
    else {
      struct servent *se = getservbyname(p, "udp");
      if (!se)
        error(0, "%.50s/udp: unknown service", p);
      c = se->s_port;
      endservent();
    }
    p[-1] = ':';
  }
  else
    c = DNS_PORT;
  sin.sin_port = htons(c);
  if (bind(fd, (struct sockaddr *)&sin, sizeof(sin)) < 0)
    error(errno, "unable to bind to %.50s:%d", inet_ntoa(sin.sin_addr), c);

  c = 65536;
  do
    if (setsockopt(fd, SOL_SOCKET, SO_RCVBUF, &c, sizeof c) == 0)
      break;
  while ((c -= (c >> 5)) >= 1024);

  if ((p = strchr(user, ':')) != NULL)
    *p++ = '\0';
  if ((c = satoi(user)) >= 0)
    uid = c, gid = c;
  else {
    struct passwd *pw = getpwnam(user);
    if (!pw)
      error(0, "invalid user `%.50s' specified", user);
    uid = pw->pw_uid;
    gid = pw->pw_gid;
    endpwent();
  }
  if (!uid)
    error(0, "daemon should not run as root, specify -u option");
  if (p) {
    if ((c = satoi(p)) >= 0)
      gid = c;
    else {
      struct group *gr = getgrnam(p);
      if (!gr)
        error(0, "invalid group `%.50s' specified", p);
      gid = gr->gr_gid;
      endgrent();
    }
    p[-1] = ':';
  }

  if (pidfile && (fpid = fopen(pidfile, "w")) == NULL)
    error(errno, "unable to write pidfile");

  if (root && (chdir(root) < 0 || chroot(root) < 0))
    error(errno, "unable to chroot");
  if (workdir && chdir(workdir) < 0)
    error(errno, "unable to chdir to %.50s", workdir);

  if (setgroups(1, &gid) < 0 || setgid(gid) < 0 || setuid(uid) < 0)
    error(errno, "unable to setuid(%d:%d)", uid, gid);

  z = (struct zone*)malloc(sizeof(struct zone) * (argc + 1));
  if (!z)
    error(0, "out of memory");
  memset(z, 0, sizeof(struct zone) * (argc + 1));
  for(c = 0; c < argc; ++c) {
    unsigned i;
    unsigned char dn[DNS_MAXDN];
    char *name = strchr(argv[c], ':');
    if (name) {
      p = argv[c];
      *name++ = '\0';
      for(i = 0;; ++i)
        if (!zonetypes[i].name)
          error(0, "unknown zone type `%.50s'", p);
        else if (strcmp(p, zonetypes[i].name) == 0)
          break;
    }
    else {
      name = argv[c];
      i = 0;
      p = NULL;
    }
    z[c].name = name;
    z[c].type = &zonetypes[i];
    z[c].queryfn = zonetypes[i].queryfn;
    if (p == argv[c]) name[-1] = ':'; /* restore argv */
    i = dns_ptodn(name, dn, sizeof(dn));
    if (!i || strlen(name) > 200) /* random constraint */
      error(errno, "invalid zone name `%.90s'", name);
    z[c].blen = i;
    if (!(z[c].base = (unsigned char*)malloc(i)))
      error(0, "out of memory");
    memcpy(z[c].base, dn, i);
  }
  loadzones(z);
  *zonep = z;

  if (!nodaemon) {
    if (fork() > 0) exit(0);
    close(0); close(1); close(2);
    setsid();
    logto = LOGTO_SYSLOG;
  }
  if (fpid) {
    fprintf(fpid, "%lu\n", (unsigned long)getpid());
    fclose(fpid);
  }
  log(LOG_INFO, "started");

  return fd;
}

/* DNS packet:
 * 0:1   identifier (client supplied)
 * 2     flags1
 *  0x80 QR response flag
 *  0x78 opcode: 0 query
 *  0x04 AA auth answer
 *  0x02 TC truncation flag
 *  0x01 RD recursion desired (may be set in query)
 * 3     flags2
 *  0x80 RA recursion available
 *  0x70 Z reserved
 *  0x0f RCODE:
 *     0 ok, 1 format error, 2 servfail, 3 nxdomain, 4 notimpl, 5 refused
 * 4:5   qdcount (numqueries)
 * 6:7   ancount (numanswers)
 * 8:9   nscount (numauthority)
 * 10:11 arcount (numadditional)
 */

struct dnspacket {
  unsigned char *c; /* current pointer */
  unsigned char *sans; /* start of answers */
  unsigned nans; /* number of answers */
  unsigned char *p; /* packet, where it begin */
  unsigned char *e; /* pointer to end-of-packet */
};

static unsigned process(struct dnspacket *p, const struct zone *zone) {
  unsigned char *q, *x;
  unsigned qlen, qcls, qtyp, qflg;
  unsigned char query[DNS_MAXDN+1];
  q = p->p + 12;
  x = p->c;
  if (q >= x) return 0; /* short packet */
  if (p->p[2] & 0x80) return 0; /* response?! */
  if (p->p[4] || p->p[5] != 1) return 0; /* qdcount != 1 */
  /* check query's DN */
  x -= 4;
  while(*q)
    if (q + *q >= x) return 0;
    else q += *q + 1;
  qlen = ++q - (p->p + 12);
  dns_dntol(p->p + 12, query);
  qtyp = ((unsigned)(q[0]) << 8) | q[1];
  qcls = ((unsigned)(q[2]) << 8) | q[3];
  q += 4;

  /* from now on, we see (almost?) valid dns query, should reply */
  p->c = p->sans = q;
  p->nans = 0;

#define refuse(p,code) (p->p[2] = 0x84, p->p[3] = (code), (p->sans - p->p))
#define nxdomain(p) (p->p[2] = 0x84, p->p[3] = 3, (p->sans - p->p))

  /* construct reply packet */
  /* identifier already in place */
  /* flags will be set up later */
  /* p[4:5] (qdcount) already set up in query */
  p->p[6] = p->p[7] = 0; /* ancount */
  p->p[8] = p->p[9] = 0; /* nscount */
  p->p[10] = p->p[11] = 0; /* arcount */

  if (qcls != DNS_C_IN && qcls != DNS_C_ANY)
    return refuse(p,1); /* format error */
  if (p->p[2] & 126)
    return refuse(p,4); /* not implemented */
  if (qtyp == DNS_T_ANY) qflg = QUERY_A|QUERY_TXT;
  else if (qtyp == DNS_T_A) qflg = QUERY_A;
  else if (qtyp == DNS_T_TXT) qflg = QUERY_TXT;
  else
    return refuse(p,5); /* refused */

  p->p[2] = 0x80; /* 0x81?! */
  if (qcls == DNS_C_IN) p->p[2] |= 0x04; /* AA */
  p->p[3] = 0;
  for(; zone->name; ++zone) {
    if (!zone->data) continue;
    if (zone->blen > qlen) continue;
    x = query + qlen - zone->blen;
    if (memcmp(zone->base, x, zone->blen)) continue;
    q = query;
    while(q < x) q += *q + 1;
    if (q != x) continue;
    *q = 0;
    if (zone->queryfn(zone->data, p, query, qflg)) {
      p->p[6] = p->nans >> 8; p->p[7] = p->nans;
      return p->c - p->p;
    }
    else
      return nxdomain(p);
  }

  return refuse(p, 5); /* refused */

}

static int aexists(const struct dnspacket *p,
	unsigned typ, void *val, unsigned vlen) {
  const unsigned char *c, *e;
  for(c = p->sans, e = p->c; c < e; c = c + 12 + c[11]) {
    if (c[2] == (typ>>8) && c[3] == (typ&255) &&
        c[11] == vlen && memcmp(c + 12, val, vlen) == 0)
      return 1;
  }
  return 0;
}

int
addrec_a(struct dnspacket *p, ip4addr_t aip) {
  if (p->c + 16 >= p->e)
    return 0;
  aip = htonl(aip);
  if (aexists(p, DNS_C_IN, &aip, 4)) return 0;
  p->nans++;
  *p->c++ = 192; *p->c++ = 12;  /* jump after header */
  *p->c++ = DNS_T_A>>8; *p->c++ = DNS_T_A; /* type */
  *p->c++ = DNS_C_IN>>8; *p->c++ = DNS_C_IN; /* class */
  *p->c++ = ttl>>24; *p->c++ = ttl>>16; *p->c++ = ttl>>8; *p->c++ = ttl;
  *p->c++ = 0; *p->c++ = 4;
  memcpy(p->c, &aip, 4);
  p->c += 4;
  return 1;
}

int
addrec_txt(struct dnspacket *p, const char *txt, const char *subst) {
  unsigned sl;
  char sb[256];
  char *lp, *s, *const e = sb + 254;
  if (!txt) return 0;
  if (p->c + 13 >= p->e) return 0;
  lp = sb + 1;
  if (!subst) subst = "$";
  while(lp < e) {
    if ((s = strchr(txt, '$')) == NULL)
      s = (char*)txt + strlen(txt);
    sl = s - txt;
    if (lp + sl > e)
      sl = e - lp;
    memcpy(lp, txt, sl);
    lp += sl;
    if (!*s) break;
    sl = strlen(subst);
    if (lp + sl > e)
      sl = e - lp;
    memcpy(lp, subst, sl);
    lp += sl;
    txt = s + 1;
  }
  sl = lp - sb;
  sb[0] = sl - 1;
  if (p->c + 12 + sl >= p->e) return 0;
  if (aexists(p, DNS_T_TXT, sb, sl)) return 0;
  p->nans++;
  *p->c++ = 192; *p->c++ = 12; /* jump to after header */
  *p->c++ = DNS_T_TXT>>8; *p->c++ = DNS_T_TXT;
  *p->c++ = DNS_C_IN>>8; *p->c++ = DNS_C_IN;
  *p->c++ = ttl>>24; *p->c++ = ttl>>16; *p->c++ = ttl>>8; *p->c++ = ttl;
  *p->c++ = 0; *p->c++ = sl; /* length of record */
  memcpy(p->c, sb, sl);
  p->c += sl;
  return 1;
}

static volatile int reload;
static void sighandler(int sig) {
  reload = 1;
  switch(sig) {
  case SIGALRM:
    alarm(recheck);
    reload = 1;
    break;
  case SIGHUP:
    log(LOG_INFO, "reloading");
    reload = 1;
    break;
  case SIGTERM:
  case SIGINT:
    log(LOG_INFO, "terminating");
    exit(0);
  }
}

static void setup_signals() {
  struct sigaction sa;
  memset(&sa, 0, sizeof(sa));
  sa.sa_handler = sighandler;
  sigemptyset(&sigset_block);
  sigaction(SIGHUP, &sa, NULL);
  sigaddset(&sigset_block, SIGHUP);
  sigaction(SIGALRM, &sa, NULL);
  sigaddset(&sigset_block, SIGALRM);
  sigaction(SIGTERM, &sa, NULL);
  sigaction(SIGINT, &sa, NULL);
}

int main(int argc, char **argv) {
  int fd;
  struct zone *zone;
  struct dnspacket p;
  unsigned char buf[DNS_MAXPACKET+1];
  fd = init(argc, argv, &zone);
  setup_signals();
  alarm(recheck);
  p.p = buf;
  p.e = buf + sizeof(buf);
  for(;;) {
    struct sockaddr_in sin;
    socklen_t sinl = sizeof(sin);
    int r;
    if (reload) {
      loadzones(zone);
      reload = 0;
    }
    r = recvfrom(fd, buf, sizeof(buf), 0, (struct sockaddr *)&sin, &sinl);
    if (r < 0)
      continue;
    p.c = buf + r;
    if (!(r = process(&p, zone)))
      continue;
    while(sendto(fd, buf, r, 0, (struct sockaddr *)&sin, sizeof(sin)) < 0)
      if (errno != EINTR) break;
  }
  return 0;
}

void *emalloc(unsigned size) {
  void *ptr = malloc(size);
  if (!ptr)
    oom();
  return ptr;
}

void *erealloc(void *ptr, unsigned size) {
  void *nptr = realloc(ptr, size);
  if (!nptr)
    oom();
  return nptr;
}

char *estrdup(const char *str) {
  char *s = strdup(str);
  if (!s) oom();
  return s;
}

void *mp_ealloc(struct mempool *mp, unsigned size) {
  void *p = mp_alloc(mp, size);
  if (!p) oom();
  return p;
}

char *mp_estrdup(struct mempool *mp, const char *str) {
  str = mp_strdup(mp, str);
  if (!str) oom();
  return (char*)str;
}

const char *sp_estrdup(struct strpool *sp, const char *str) {
  str = sp_strdup(sp, str);
  if (!str) oom();
  return str;
}
