/*
Copyright (c) 2012, 2013, Felix J. Ogris
All rights reserved.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met: 

1. Redistributions of source code must retain the above copyright notice, this
   list of conditions and the following disclaimer. 
2. Redistributions in binary form must reproduce the above copyright notice,
   this list of conditions and the following disclaimer in the documentation
   and/or other materials provided with the distribution. 

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR
ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*/

/* mdnsproxy $Rev: 237 $ - listens for DNS requests on 224.0.0.251:5353 and forwards
 * them to the nameservers listed in /etc/resolv.conf
 *
 * compile: cc -W -Wall -O3 -s -pipe -o mdnsproxy mdnsproxy.c
 * start:   ./mdnsproxy
 * stop:    killall mdnsproxy
 * help:    ./mdnsproxy -h
 */

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

#define TIMEOUT 10

#define LOGX(level, msg, params ...) do { \
  if (log_to_syslog) \
    syslog(level, "[%s:%03i] " msg "\n", __FILE__, __LINE__, ## params); \
  else \
    fprintf(stderr, "[%s:%03i] " msg "\n", __FILE__, __LINE__, ## params); \
} while (0)

#define ERRDIEX(msg, params ...) do { \
  LOGX(LOG_ERR, msg, ## params); \
  exit(1); \
} while (0)
#define LOGWARNX(msg, params ...) LOGX(LOG_WARNING, msg, ## params)
#define LOGINFOX(msg, params ...) LOGX(LOG_INFO, msg, ## params)
#define LOGDEBUGX(msg, params ...) do { if (debug >= 2) \
  LOGX(LOG_DEBUG, msg, ## params); \
} while (0)

#define ERRDIE(msg, params ...) do { \
  LOGX(LOG_ERR, msg ": %s", ## params, strerror(errno)); \
  exit(1); \
} while (0)
#define LOGWARN(msg, params ...) LOGX(LOG_WARNING, msg ": %s", \
                                      ## params, strerror(errno))

struct {
  struct sockaddr_in clntaddr;
  struct sockaddr_in servaddr;
  int servsock;
  time_t lastaction;
  int nameserver;
} conns[512];

int debug = 0;
int log_to_syslog = 0;
int running = 1;
int freeslot;

void close_slot(int slot)
{
  LOGDEBUGX("#%i: close", slot);
  close(conns[slot].servsock);
  conns[slot].servsock = -1;
  freeslot = slot;
}

/* load nameserver from line #nsidx of /etc/resolv.conf to servaddr at
   connection #slot */
int loadns(int slot, int nsidx)
{
  FILE *f;
  char buf[1024];
  char *c;
  int idx = 0;

  if ((f = fopen("/etc/resolv.conf", "r")) == NULL) {
    LOGWARN("fopen(/etc/resolv.conf, r)");
    return -1;
  }

  conns[slot].nameserver = nsidx;

  while (fgets(buf, sizeof(buf), f) != NULL) {
    if ((strncasecmp("nameserver", buf, strlen("nameserver")) == 0) &&
        (idx++ >= nsidx)) {
      /* remove trailing whitespace */
      c = buf + strlen(buf) - 1;
      while ((c >= buf) && ((*c==' ')||(*c=='\t')||(*c=='\r')||(*c=='\n'))) {
        *c-- = '\0';
      }
      /* skip leading whitespace */
      for (c = buf + strlen("nameserver"); (*c==' ')||(*c=='\t'); c++)
        ;;
      
      memset(&conns[slot].servaddr, 0, sizeof(conns[slot].servaddr));
      conns[slot].servaddr.sin_addr.s_addr = inet_addr(c);
      if (conns[slot].servaddr.sin_addr.s_addr != INADDR_NONE) {
        conns[slot].servaddr.sin_port = htons(53);
        conns[slot].servaddr.sin_family = PF_INET;

        fclose(f);
        LOGDEBUGX("#%i: ns %s", slot,
                  inet_ntoa(conns[slot].servaddr.sin_addr));

        return slot;
      }
    }
  }

  fclose(f);
  if (nsidx == 0) LOGWARNX("#%i: no ns", slot);

  return -1;
}

void kill_handler(int sig __attribute__((unused)))
{
  running = 0;
  LOGINFOX("received SIGKILL, going down...");
}

int sockaddr_in_equal(socklen_t sinlen, struct sockaddr_in *s1, struct sockaddr_in *s2)
{
  return ((sinlen == sizeof(*s2)) &&
          (s1->sin_family == s2->sin_family) &&
          (s1->sin_addr.s_addr == s2->sin_addr.s_addr) &&
          (s1->sin_port == s2->sin_port));
}

int main(int argc, char **argv)
{
  char ch;
  char *pidfile = "/var/run/mdnsproxy.pid";
  gid_t gid = 65535;
  uid_t uid = 65535;
  struct passwd *pw;
  pid_t pid;
  struct stat st;
  FILE *f;
  int mcastsock;
  int i;
  struct sockaddr_in sin;
  socklen_t sinlen;
  struct ip_mreq ipm;
  struct sigaction sa;
  fd_set fds;
  int maxsock;
  struct timeval timeout;
  char buf[1024];
  ssize_t recvlen;
  ssize_t sentlen;

  while ((ch = getopt(argc, argv, "dhi:u:")) != -1) {
    switch (ch) {
      case 'd': debug++; break;
      case 'i': pidfile = optarg; break;
      case 'u':
        if ((pw = getpwnam(optarg)) == NULL)
          ERRDIEX("no such user: %s", optarg);
        gid = pw->pw_gid;
        uid = pw->pw_uid;
        break;
      default:
        fputs(
"mdnsproxy $Rev: 237 $ - forwards multicast dns requests to unicast nameservers\n\n"
"usage: mdnsproxy [-d [-d]] [-i <pidfile>] [-u <user>]\n"
"       mdnsproxy -h\n\n"
"  -d             run in foreground, don't write pidfile\n"
"  -d -d          run in foreground, don't write pidfile, debug log to stderr,\n"
"                 don't drop privileges\n"
"  -i <pidfile>   write pid to <pidfile> instead of /var/run/mdnsproxy.pid\n"
"  -u <user>      run as user <user> instead of uid 65535 and gid 65535\n"
"  -h             show this help ;-)\n\n",
              (ch == 'h' ? stdout : stderr));
        return (ch != 'h');
    }
  }

  close(0);
  close(1);

  if (chdir("/") != 0) ERRDIE("chdir(/)");
  umask(0022);

  if (debug <= 0) {
    /* check pidfile */
    if (lstat(pidfile, &st) == 0) {
      if (!S_ISREG(st.st_mode))
        ERRDIEX("%s exists, but is not a regular file", pidfile);
      if ((f = fopen(pidfile, "r")) == NULL) ERRDIE("fopen(%s, r)", pidfile);
      if (fgets(buf, sizeof(buf), f) == NULL) ERRDIE("fgets(%s)", pidfile);
      pid = atoi(buf);
      fclose(f);
      if (kill(pid, 0) == 0) ERRDIEX("still running (pid=%li)", (long)pid);
    } else if (errno != ENOENT) ERRDIE("%s", pidfile);

    /* daemonize */
    pid = fork();
    if (pid < 0) ERRDIE("fork()");
    if (pid > 0) return 0;
    if (setsid() == -1) ERRDIE("setsid()");

    /* save pid */
    if ((f = fopen(pidfile, "w")) == NULL) ERRDIE("fopen(%s, w)", pidfile);
    fprintf(f, "%li\n", (long)getpid());
    if (fclose(f) != 0) ERRDIE("fclose(%s)", pidfile);;
  }

  if (debug <= 1) {
    /* drop privileges */
    if (setgid(gid) != 0) ERRDIE("setgid(%lu)", (long unsigned)gid);
    if (setuid(uid) != 0) ERRDIE("setuid(%lu)", (long unsigned)uid);
  }

  /* join multicast group */
  memset(&sin, 0, sizeof(sin));
  sin.sin_addr.s_addr = inet_addr("224.0.0.251");
  sin.sin_port = htons(5353);
  sin.sin_family = PF_INET;

  memset(&ipm, 0, sizeof(ipm));
  ipm.imr_multiaddr.s_addr = sin.sin_addr.s_addr;
  ipm.imr_interface.s_addr = htonl(INADDR_ANY);

  if ((mcastsock = socket(AF_INET, SOCK_DGRAM, 0)) < 0)
    ERRDIE("socket()");

  i = 0;
  if (setsockopt(mcastsock, IPPROTO_IP, IP_MULTICAST_LOOP, &i, sizeof(i)))
    ERRDIE("setsockopt(IP_MULTICAST_LOOP)");

  i = 1;
  if (setsockopt(mcastsock, IPPROTO_IP, IP_MULTICAST_TTL, &i, sizeof(i)))
    ERRDIE("setsockopt(IP_MULTICAST_TTL)");

  if (bind(mcastsock, (struct sockaddr*) &sin, sizeof(sin)))
    ERRDIE("bind()");

  if (setsockopt(mcastsock, IPPROTO_IP, IP_ADD_MEMBERSHIP, &ipm, sizeof(ipm)))
    ERRDIE("setsockopt(IP_ADD_MEMBERSHIP");

  /* set up signal handlers */
  memset(&sa, 0, sizeof(sa));
  sa.sa_handler = kill_handler;
  if (sigaction(SIGTERM, &sa, NULL) < 0)
    ERRDIE("sigaction(SIGTERM)");

  if (debug <= 1) {
    /* open syslog */
    close(2);
    openlog("mdnsproxy", LOG_PID, LOG_DAEMON);
    log_to_syslog = 1;
  }

  for (i = sizeof(conns)/sizeof(conns[0]) - 1; i >= 0; i--)
    conns[i].servsock = -1;

  LOGINFOX("starting...");

  while (running) {
    FD_ZERO(&fds);
    FD_SET(mcastsock, &fds);
    maxsock = mcastsock;
    freeslot = -1;

    /* fill fds */
    for (i = sizeof(conns)/sizeof(conns[0]) - 1; i >= 0; i--) {
      if (conns[i].servsock == -1) { /* no-op */ }
      else if (conns[i].lastaction + TIMEOUT < time(NULL)) close_slot(i);
      else {
        FD_SET(conns[i].servsock, &fds);
        if (conns[i].servsock > maxsock) maxsock = conns[i].servsock;
      }
    }

    timeout.tv_sec = TIMEOUT;
    timeout.tv_usec = 0;
    i = select(maxsock + 1, &fds, NULL, NULL, &timeout);

    if ((i < 0) && (errno != EINTR)) ERRDIE("select()");
    if (i <= 0) continue;

    /* check for data from nameservers */
    for (i = sizeof(conns)/sizeof(conns[0]) - 1; i >= 0; i--) {
      sinlen = sizeof(sin);

      if (conns[i].servsock == -1) {
        /* reuse this connection slot */
        freeslot = i;
      }
      else if (!FD_ISSET(conns[i].servsock, &fds)) {
        /* no-op, no data */
      }
      else if ((recvlen = recvfrom(conns[i].servsock, buf, sizeof(buf), 0,
                                   (struct sockaddr*) &sin, &sinlen)) < 0) {
        /* recvfrom() failed */
        LOGWARN("#%i: recvfrom()", i);
        close_slot(i);
      }
      else if (!sockaddr_in_equal(sinlen, &sin, &conns[i].servaddr)) {
        /* some bogus nameserver sent an udp packet */
        LOGWARNX("#%i: martian %s", i, inet_ntoa(sin.sin_addr));
      }
      else if ((sentlen = sendto(mcastsock, buf, recvlen, 0,
                                 (struct sockaddr*) &conns[i].clntaddr,
                                 sizeof(conns[i].clntaddr))) < 0) {
        /* sendto() failed */
        LOGWARN("#%i: sendto() on mcast socket", i);
        close_slot(i);
      }
      else if (sentlen != recvlen) {
        /* sendto() failed */
        LOGWARNX("#%i: sendto() on mcast socket didn't send all data "
                 "(wanted=%li, sent=%li)", i, recvlen, sentlen);
        close_slot(i);
      }
      else {
        /* got data from nameserver and sent it to the client */
        LOGDEBUGX("#%i: s->c %li bytes", i, sentlen);
        conns[i].lastaction = time(NULL);
      }
    }

    /* check for data from clients */
    if (FD_ISSET(mcastsock, &fds)) {
      sinlen = sizeof(sin);
      recvlen = recvfrom(mcastsock, buf, sizeof(buf), 0,
                         (struct sockaddr*) &sin, &sinlen);

      if (recvlen < 0) {
        /* recvfrom() failed */
        if (errno != EINTR) ERRDIE("recvfrom() on mcast socket");
        continue;
      }

      for (i = sizeof(conns)/sizeof(conns[0]) - 1; i >= 0; i--)
        /* do we know this client already? */
        if ((conns[i].servsock != -1) &&
            sockaddr_in_equal(sinlen, &sin, &conns[i].clntaddr))
          break;

      if (i >= 0) {
        /* found returning client */
        LOGDEBUGX("#%i: old %s", i, inet_ntoa(sin.sin_addr));

        /* if first call to loadns() succeeds, then i is >= 0 */
        if (loadns(i, conns[i].nameserver + 1) < 0) {
          /* otherwise i becomes the result of loadns(),
             which is i itself or < 0 (=error) */
          i = loadns(i, 0);
        }
      }
      else if (freeslot < 0) {
        /* i is < 0 */
        LOGWARNX("no free slot");
      }
      else {
        /* new client, i < 0, freeslot >= 0 */
        LOGDEBUGX("#%i: new %s", freeslot, inet_ntoa(sin.sin_addr));

        memcpy(&conns[freeslot].clntaddr, &sin, sinlen);

        /* if socket() fails, i remains < 0
           otherwise i becomes the result of loadns(),
           which is freeslot or < 0 (=error) */
        if ((conns[freeslot].servsock = socket(AF_INET, SOCK_DGRAM, 0)) < 0)
          LOGWARN("#%i: socket()", freeslot);
        else i = loadns(freeslot, 0);
      }
      
      if (i < 0) {
        /* no-op, error */
      }
      else if ((sentlen = sendto(conns[i].servsock, buf, recvlen, 0,
                                 (struct sockaddr*) &conns[i].servaddr,
                                 sizeof(conns[i].servaddr))) < 0) {
        /* sento() failed */
        LOGWARN("#%i: sendto()", i);
        close_slot(i);
      }
      else if (sentlen != recvlen) {
        /* sento() failed */
        LOGWARNX("#%i: sendto() didn't send all data (wanted=%li, "
                 "sent=%li)", i, recvlen, sentlen);
        close_slot(i);
      }
      else {
        /* got data from client and sent it to the nameserver */
        LOGDEBUGX("#%i: c->s %li bytes", i, sentlen);
        conns[i].lastaction = time(NULL);
      }
    }
  }

  /* cleanup */
  for (i = sizeof(conns)/sizeof(conns[0]) - 1; i >= 0; i--)
    if (conns[i].servsock != -1) close_slot(i);

  /* leave multicast group */
  if (setsockopt(mcastsock, IPPROTO_IP, IP_DROP_MEMBERSHIP,
                 &ipm, sizeof(ipm)) < 0)
    LOGWARNX("setsockopt(IP_DROP_MEMBERSHIP): %s (Are you missing a route "
             "for multicast, e.g. `ip route add 224.0.0.0/4 dev eth0`?)",
             strerror(errno));

  close(mcastsock);

  LOGINFOX("exiting...");
  if (log_to_syslog) closelog();

  return 0;
}