/* ISC license. */

#include <string.h>
#include <stdint.h>
#include <limits.h>
#include <stdlib.h>
#include <errno.h>

#include <skalibs/tai.h>
#include <skalibs/stralloc.h>
#include <skalibs/genalloc.h>
#include <skalibs/iopause.h>
#include <skalibs/ip46.h>
#include <skalibs/random.h>
#include <skalibs/prog.h>

#include <s6-dns/s6dns.h>
#include <s6-dns/skadns.h>

#include "qmailr.h"
#include "qmail-remote.h"

typedef struct cnameinfo_s cnameinfo, *cnameinfo_ref ;
struct cnameinfo_s
{
  stralloc sa ;
  size_t atpos ;
  uint16_t id ;
  uint16_t count ;
} ;

typedef struct mxipinfo_s mxipinfo, mxipinfo_ref ;
struct mxipinfo_s
{
  stralloc ip4 ;
  stralloc ip6 ;
  size_t pos ;
  uint16_t id4 ;
  uint16_t id6 ;
} ;
#define MXIPINFO_ZERO { .ip4 = STRALLOC_ZERO, .ip6 = STRALLOC_ZERO, .id4 = UINT16_MAX, .id6 = UINT16_MAX }

static int mx_cmp (void const *a, void const *b)
{
  s6dns_message_rr_mx_t const *aa = a ;
  s6dns_message_rr_mx_t const *bb = b ;
  return aa->preference < bb-> preference ? -1 : aa->preference > bb->preference ;
}

static unsigned int use_host_as_mx (skadns_t *a, char const *host, genalloc *mxip, stralloc *storage, tain const *deadline)
{
  size_t hostlen = strlen(host) ;
  unsigned int newreqs = 0 ;
  mxipinfo info = MXIPINFO_ZERO ;
  s6dns_domain_t q ;
  if (!s6dns_domain_fromstring_noqualify_encode(&q, host, hostlen))
  {
    skadns_end(a) ;
    qmailr_tempusys("DNS-encode host domain") ;
  }
  info.pos = storage->len ;
  if (!stralloc_catb(storage, host, hostlen+1))
  {
    skadns_end(a) ;
    dienomem() ;
  }
  if (hostlen > 1 && storage->s[storage->len - 2] == '.') storage->s[--storage->len - 1] = 0 ;
  if (!skadns_send_g(a, &info.id4, &q, S6DNS_T_A, deadline, deadline))
  {
    skadns_end(a) ;
    qmailr_tempusys("send ", "A", " DNS query") ;
  }
  newreqs++ ;
#ifdef SKALIBS_IPV6_ENABLED
  if (!skadns_send_g(a, &info.id6, &q, S6DNS_T_AAAA, deadline, deadline))
  {
    skadns_end(a) ;
    qmailr_tempusys("send ", "AAAA", " DNS query") ;
  }
  newreqs++ ;
#endif
  if (!genalloc_catb(mxipinfo, mxip, &info, 1))
  {
    skadns_end(a) ;
    dienomem() ;
  }
  return newreqs ;
}

 /*
   The point of this monster here is to do all the DNS resolutions in parallel,
   to avoid compounding network latency. One of the many things that could never
   be done by patching the original qmail-remote.
   1 sender + n-1 recipients are given in eaddr.
   - loop around CNAME until we get the canonical name, for the n eaddrs
   - either lookup the MX for the host then find all the A and AAAAs of all the MXes,
     or get the A and AAAAs of the host directly (if smtproutes)
   - do not keep the As and AAAAs listed in ipme
   - sort the set of IPs by MX preference
   When done, addrmangle (i.e. quote if needed) all the boxnames in eaddr.
   Shove everything in storage and return the indices:
   in eaddrpos for sender+recipients, in mxipind for the IPs to connect to.

   Also, fuck DNS for requiring so many small allocations and data copies.

   Also, fuck DNS.
 */


#define ddienomem() do { skadns_end(&a) ; dienomem() ; } while (0)
#define qmailr_dperm(...) do { skadns_end(&a) ; qmailr_perm(__VA_ARGS__) ; } while (0)
#define qmailr_dtemp(...) do { skadns_end(&a) ; qmailr_temp(__VA_ARGS__) ; } while (0)
#define qmailr_dtempsys(...) do { skadns_end(&a) ; qmailr_tempsys(__VA_ARGS__) ; } while (0)
#define qmailr_dtempusys(...) do { skadns_end(&a) ; qmailr_tempusys(__VA_ARGS__) ; } while (0)

unsigned int dns_stuff (char const *helohost, char *heloip4, char *heloip6, char const *host, char const *const *eaddr, unsigned int n, size_t *eaddrpos, genalloc *mxipind, stralloc *storage, unsigned int timeoutdns, char const *ipme4, unsigned int n4, char const *ipme6, unsigned int n6, uint32_t flags)
{
  skadns_t a = SKADNS_ZERO ;
  genalloc mxipi = GENALLOC_ZERO ;  /* mxipinfo */
  unsigned int pending = 0 ;
  unsigned int mxn = 0 ;
  stralloc helosa = STRALLOC_ZERO ;
  uint16_t mxid = UINT16_MAX ;
  uint16_t heloid4 = UINT16_MAX ;
#ifdef SKALIBS_IPV6_ENABLED
  uint16_t heloid6 = UINT16_MAX ;
#endif
  tain deadline ;
  cnameinfo cnames[n] ;

  qdeadline(&deadline, timeoutdns) ;
  if (!skadns_startf_g(&a, &deadline))
    qmailr_tempusys("start asynchronous DNS helper") ;

  {
    s6dns_domain_t q ;
    if (!s6dns_domain_fromstring_noqualify_encode(&q, helohost, strlen(helohost)))
      qmailr_dtempusys("DNS-encode helo string") ;
    if (!skadns_send_g(&a, &heloid4, &q, S6DNS_T_A, &deadline, &deadline))
      qmailr_dtempusys("send ", "A", " DNS query") ;
    pending++ ;
#ifdef SKALIBS_IPV6_ENABLED
    if (!skadns_send_g(&a, &heloid6, &q, S6DNS_T_AAAA, &deadline, &deadline))
      qmailr_dtempusys("send ", "AAAA", " DNS query") ;
    pending++ ;
#endif
  }

  for (unsigned int i = 0 ; i < n ; i++)
  {
    char const *at = strrchr(eaddr[i], '@') ;
    cnames[i].sa = stralloc_zero ;
    cnames[i].id = UINT16_MAX ;
    if (at)
    {
      s6dns_domain_t q ;
      size_t len = strlen(at+1) ;
      cnames[i].atpos = at - eaddr[i] ;
      if (!stralloc_catb(&cnames[i].sa, at+1, len)) ddienomem() ;
      cnames[i].count = 1 ;
      if (at[1] != '[')
      {
        if (!s6dns_domain_fromstring_noqualify_encode(&q, at+1, len))
          qmailr_dtempusys("DNS-encode recipient domain") ;
        if (!skadns_send_g(&a, &cnames[i].id, &q, S6DNS_T_CNAME, &deadline, &deadline))
          qmailr_dtempusys("send ", "CNAME", " DNS query") ;
        pending++ ;
      }
    }
    else
    {
      cnames[i].count = 0 ;
      cnames[i].atpos = strlen(eaddr[i]) ;
    }
  }

  if (flags & 1)
  {
    s6dns_domain_t q ;
    if (!s6dns_domain_fromstring_noqualify_encode(&q, host, strlen(host)))
      qmailr_dtempusys("DNS-encode host domain") ;
    if (!skadns_send_g(&a, &mxid, &q, S6DNS_T_MX, &deadline, &deadline))
      qmailr_dtempusys("send ", "MX", " DNS query") ;
    pending++ ;
  }
  else
  {
    mxn = 1 ;
    pending += use_host_as_mx(&a, host, &mxipi, storage, &deadline) ;
  }

  while (pending)
  {
    uint16_t *ids ;
    iopause_fd x = { .fd = skadns_fd(&a), .events = IOPAUSE_READ } ;
    int r = iopause_g(&x, 1, &deadline) ;
    if (r == -1) qmailr_dtempusys("iopause") ;
    if (!r) qmailr_dtempsys("Timed out waiting for DNS") ;
    r = skadns_update(&a) ;
    if (r == -1) qmailr_dtempusys("read DNS answers") ;
    ids = genalloc_s(uint16_t, &a.list) ;
    for (size_t j = 0 ; j < genalloc_len(uint16_t, &a.list) ; j++)
    {
      char const *packet = skadns_packet(&a, ids[j]) ;
      uint16_t packetlen = skadns_packetlen(&a, ids[j]) ;
      if (!packet) qmailr_dtempsys("DNS packet reading error") ;

      if (ids[j] == heloid4)  /* ipv4 for the helohost */
      {
        s6dns_message_header_t h ;
        r = s6dns_message_parse(&h, packet, packetlen, &s6dns_message_parse_answer_a, &helosa) ;
        if (r == -1) qmailr_dtempsys("DNS packet parsing error") ;
        if (!r)
        {
          if (errno == EBUSY || errno == EIO) qmailr_dtemp("Temporary DNS error while resolving ", "A", "for helohost") ;
          else qmailr_dperm("DNS ", "A", " resolution error") ;
        }
        skadns_release(&a, heloid4) ;
        pending-- ;
        heloid4 = UINT16_MAX ;
        if (helosa.len >= 4) memcpy(heloip4, helosa.s, 4) ;
        helosa.len = 0 ;
      }

#ifdef SKALIBS_IPV6_ENABLED
      if (ids[j] == heloid6)  /* ipv6 for the helohost */
      {
        s6dns_message_header_t h ;
        r = s6dns_message_parse(&h, packet, packetlen, &s6dns_message_parse_answer_aaaa, &helosa) ;
        if (r == -1) qmailr_dtempsys("DNS packet parsing error") ;
        if (!r)
        {
          if (errno == EBUSY || errno == EIO) qmailr_dtemp("Temporary DNS error while resolving ", "AAAA", "for helohost") ;
          else qmailr_dperm("DNS ", "AAAA", " resolution error") ;
        }
        skadns_release(&a, heloid6) ;
        pending-- ;
        heloid6 = UINT16_MAX ;
        if (helosa.len >= 16) memcpy(heloip6, helosa.s, 16) ;
        helosa.len = 0 ;
      }
#endif

      if (ids[j] == mxid)  /* return from MX query */
      {
        s6dns_message_header_t h ;
        genalloc mxes = GENALLOC_ZERO ;  /* s6dns_message_rr_mx_t */

        r = s6dns_message_parse(&h, packet, packetlen, &s6dns_message_parse_answer_mx, &mxes) ;
        if (r == -1) qmailr_dtempsys("DNS packet parsing error") ;
        if (!r)
        {
          if (errno == EBUSY || errno == EIO) qmailr_dtemp("Temporary DNS error while resolving ", "MX") ;
          else qmailr_dperm("DNS ", "CNAME", " resolution error") ;
        }
        skadns_release(&a, ids[j]) ;
        pending-- ;
        mxid = UINT16_MAX ;
        if (r >= 2)  /* we have MXes, ask for their IPs */
        {
          s6dns_message_rr_mx_t *mxs = genalloc_s(s6dns_message_rr_mx_t, &mxes) ;
          mxn = genalloc_len(s6dns_message_rr_mx_t, &mxes) ;
          if (!genalloc_readyplus(mxipinfo, &mxipi, mxn)) ddienomem() ;
          qsort(mxs, mxn, sizeof(s6dns_message_rr_mx_t), &mx_cmp) ;
          for (unsigned int i = 0 ; i < mxn ; i++)
          {
            mxipinfo *p = genalloc_s(mxipinfo, &mxipi) + i ;
            unsigned int len ;
            if (!stralloc_readyplus(storage, 256)) dienomem() ;
            p->pos = storage->len ;
            len = s6dns_domain_tostring(storage->s + p->pos, 256, &mxs[i].exchange) ;
            if (!len) qmailr_perm("invalid MX name") ;
            storage->len += len ;
            if (storage->s[storage->len - 1] == '.') storage->len-- ;
            storage->s[storage->len++] = 0 ;
            p->ip4 = p->ip6 = stralloc_zero ;
            s6dns_domain_encode(&mxs[i].exchange) ;
            if (!skadns_send_g(&a, &p->id4, &mxs[i].exchange, S6DNS_T_A, &deadline, &deadline))
              qmailr_dtempusys("send ", "A", " DNS query") ;
            pending++ ;
#ifdef SKALIBS_IPV6_ENABLED
            if (!skadns_send_g(&a, &p->id6, &mxs[i].exchange, S6DNS_T_AAAA, &deadline, &deadline))
              qmailr_dtempusys("send ", "AAAA", " DNS query") ;
            pending++ ;
#endif
          }
          genalloc_free(s6dns_message_rr_mx_t, &mxes) ;
        }
        else
        {
          mxn = 1 ;
          pending += use_host_as_mx(&a, host, &mxipi, storage, &deadline) ;
        }
        continue ;
      }

      for (unsigned int i = 0 ; i < n ; i++) if (ids[j] == cnames[i].id)  /* return from CNAME query */
      {
        s6dns_message_header_t h ;
        s6dns_dpag_t dlist = { .ds = GENALLOC_ZERO, .rtype = S6DNS_T_CNAME } ;
        r = s6dns_message_parse(&h, packet, packetlen, &s6dns_message_parse_answer_domain, &dlist) ;
        if (r == -1) qmailr_dtempsys("DNS packet parsing error") ;
        if (!r)
        {
          if (errno == EBUSY || errno == EIO) qmailr_temp("Temporary DNS error while resolving ", "CNAME") ;
          else qmailr_dperm("DNS ", "CNAME", " resolution error") ;
        }
        skadns_release(&a, ids[j]) ;
        pending-- ;
        if (r >= 2)  /* it's a CNAME, loop on it */
        {
          s6dns_domain_t *domain = genalloc_s(s6dns_domain_t, &dlist.ds) ;
          if (cnames[i].count++ >= 100) qmailr_perm("DNS CNAME loop") ;
          if (!skadns_send_g(&a, &cnames[i].id, domain, S6DNS_T_CNAME, &deadline, &deadline))
            qmailr_dtempusys("send ", "CNAME", " DNS query") ;
          pending++ ;
          if (!stralloc_ready(&cnames[i].sa, 256)) ddienomem() ;
          s6dns_domain_decode(domain) ;
          cnames[i].sa.len = s6dns_domain_tostring(cnames[i].sa.s, 256, domain) ;
          genalloc_free(s6dns_domain_t, &dlist.ds) ;
        }
        else cnames[i].id = UINT16_MAX ;  /* we have the canonical host in cnames[i].sa */
        continue ;
      }

      for (unsigned int i = 0 ; i < mxn ; i++)
      {
        mxipinfo *p = genalloc_s(mxipinfo, &mxipi) + i ;
        if (ids[j] == p->id4)
        {
          s6dns_message_header_t h ;
          r = s6dns_message_parse(&h, packet, packetlen, &s6dns_message_parse_answer_a, &p->ip4) ;
          if (r == -1) qmailr_dtempsys("DNS packet parsing error") ;
          if (!r)
          {
            if (errno == EBUSY || errno == EIO) qmailr_dtemp("Temporary DNS error while resolving ", "A") ;
            else qmailr_dperm("DNS ", "A", " resolution error") ;
          }
          skadns_release(&a, ids[j]) ;
          pending-- ;
          p->id4 = UINT16_MAX ;
          for (unsigned int k = 0 ; k < p->ip4.len ; k += 4)
          {
            if (bsearch(p->ip4.s + k, ipme4, n4, 4, &qmailr_memcmp4))
            {
              memmove(p->ip4.s + k, p->ip4.s + p->ip4.len - 4, 4) ;
              p->ip4.len -= 4 ;
              k -= 4 ;
            }
          }
        }
#ifdef SKALIBS_IPV6_ENABLED
        else if (ids[j] == p->id6)
        {
          s6dns_message_header_t h ;
          r = s6dns_message_parse(&h, packet, packetlen, &s6dns_message_parse_answer_aaaa, &p->ip6) ;
          if (r == -1) qmailr_dtempsys("DNS packet parsing error") ;
          if (!r)
          {
            if (errno == EBUSY || errno == EIO) qmailr_dtemp("Temporary DNS error while resolving ", "AAAA") ;
            else qmailr_dperm("DNS ", "AAAA", " resolution error") ;
          }
          skadns_release(&a, ids[j]) ;
          pending-- ;
          p->id6 = UINT16_MAX ;
          for (unsigned int k = 0 ; k < p->ip6.len ; k += 16)
          {
            if (bsearch(p->ip6.s + k, ipme6, n6, 16, &qmailr_memcmp16))
            {
              memmove(p->ip6.s + k, p->ip6.s + p->ip6.len - 16, 16) ;
              p->ip6.len -= 16 ;
              k -= 16 ;
            }
          }
        }
#endif
      }
    }
  }
  skadns_end(&a) ;  /* we done, buddy */

  stralloc_free(&helosa) ;

  for (unsigned int i = 0 ; i < n ; i++)
  {
    eaddrpos[i] = storage->len ;
    if (!qmailr_box_encode(eaddr[i], cnames[i].atpos, storage)) dienomem() ;
    if (cnames[i].count)
    {
      if (!stralloc_catb(storage, "@", 1)) dienomem() ;
      if (!stralloc_catb(storage, cnames[i].sa.s, cnames[i].sa.len)) dienomem() ;
      stralloc_free(&cnames[i].sa) ;
    }
    if (!stralloc_0(storage)) dienomem() ;
  }

  if (!genalloc_readyplus(mxip, mxipind, mxn)) dienomem() ;
  for (unsigned int i = 0 ; i < mxn ; i++)
  {
    mxip data ;
    mxipinfo *p = genalloc_s(mxipinfo, &mxipi) + i ;
    data.namepos = p->pos ;
    data.n4 = p->ip4.len >> 2 ;
    data.pos4 = storage->len ;
    if (!stralloc_catb(storage, p->ip4.s, p->ip4.len)) dienomem() ;
    stralloc_free(&p->ip4) ;
#ifdef SKALIBS_IPV6_ENABLED
    data.n6 = p->ip6.len >> 4 ;
    data.pos6 = storage->len ;
    if (!stralloc_catb(storage, p->ip6.s, p->ip6.len)) dienomem() ;
    stralloc_free(&p->ip6) ;
    genalloc_catb(mxip, mxipind, &data, 1) ;
#endif
  }
  genalloc_free(mxipinfo, &mxipi) ;
  return mxn ;
}