Description: dnscache: merge similar outgoing udp packets
 This patch takes a slightly different approach to merging
 than the previous qmerge patch: rather than merging between
 the "query" and "dns_transmit" layers, it actually
 incorporates the merging into the dns_transmit layer.
 .
 This makes it a much more invasive and messy patch, but
 means that it should correctly handle all situations
 (in particular, ones where sending to two different IPs will
 gives different results, generally because one server is
 down or lame).
 .
 The general strategy is this:
 .
   - when a dns_transmit is about to send() a UDP query, it
     checks for an equivalent inprogress query.
 .
     If it finds one, it marks the inprogress dns_transmit as
     its "master", and itself as a "slave" of the master.
 .
     If it doesn't, it sends the packet and registers itself
     as inprogress.
 .
   - when a dns_transmit recv()s a UDP packet, it creates a
     copy of the packet for each of its slaves
 .
 Two outgoing packets are considered equivalent if:
 .
   1. they are going to the same IP
   2. they have the same qtype
   3. they have the same qname
 .
 Because this change affects the dns library itself, this
 change can potentially affect not just dnscache, but all of
 the diagnostic tools. To address this, merging must be
 enabled explicitly by the caller; dnscache will enable
 merging if the MERGEQUERIES environment variable is set.
 .
 I tried to keep the patch as small and simple as possible so
 that its correctness could be verified by readers.  There
 are a few places where performance might benefit from making
 it more complex:
 .
   - dns_transmit now knows the value of MAXUDP, since it is
     an upper bound on the number of slaves and inprogress
     queries. As a result:
 .
       - each non-merging program which uses the dns library
         wastes MAXUDP*sizeof(pointer) bytes of static memory
         for the inprogress list
 .
       - each dns_transmit uses an extra MAXUDP*sizeof(pointer)
         bytes for its slaves list. For dnscache, this
         translates to 160K total on a 32-bit platform with
         default MAXUDP.
 .
     Both could be avoided by using dynamic allocation.
 .
   - We have to do an O(MAXUDP) linear search to find similar
     inprogress queries (non-merge-enabled users of the
     library avoid paying this price, though).
 .
     This could be avoided by using a data structure with a
     fast key lookup for inprogress.
 .
 This patch is only lightly tested. Use on production servers at your own
 risk (and please report to the list if you have success using it).
Author: Jeff King <peff () peff ! net>
Date: Wed, 1 Apr 2009 14:10:54 +0000
Debian-Bug: https://bugs.debian.org/516394
Last-Update: 2020-07-26

diff --git a/clients.h b/clients.h
new file mode 100644
index 0000000..983a4ad
--- /dev/null
+++ b/clients.h
@@ -0,0 +1,7 @@
+#ifndef CLIENTS_H
+#define CLIENTS_H
+
+#define MAXUDP 200
+#define MAXTCP 20
+
+#endif /* CLIENTS_H */
diff --git a/dns.h b/dns.h
index 3849f4c..d1e2ffc 100644
--- a/dns.h
+++ b/dns.h
@@ -4,6 +4,7 @@
 #include "stralloc.h"
 #include "iopause.h"
 #include "taia.h"
+#include "clients.h"
 
 #define DNS_C_IN "\0\1"
 #define DNS_C_ANY "\0\377"
@@ -38,8 +39,14 @@ struct dns_transmit {
   const char *servers;
   char localip[4];
   char qtype[2];
+  struct dns_transmit *master;
+  struct dns_transmit *slaves[MAXUDP];
+  int nslaves;
 } ;
 
+extern void dns_enable_merge(void (*logger)(const char *, const char *,
+      const char *));
+
 extern void dns_random_init(const char *);
 extern unsigned int dns_random(unsigned int);
 
diff --git a/dns_transmit.c b/dns_transmit.c
index 4d6e39f..3984776 100644
--- a/dns_transmit.c
+++ b/dns_transmit.c
@@ -7,6 +7,61 @@
 #include "byte.h"
 #include "uint16.h"
 #include "dns.h"
+#include "strerr.h"
+
+static int merge_enable;
+static void (*merge_logger)(const char *, const char *, const char *);
+void dns_enable_merge(void (*f)(const char *, const char *, const char *))
+{
+  merge_enable = 1;
+  merge_logger = f;
+}
+
+static int merge_equal(struct dns_transmit *a, struct dns_transmit *b)
+{
+  const char *ip1 = a->servers + 4 * a->curserver;
+  const char *ip2 = b->servers + 4 * b->curserver;
+  return
+    byte_equal(ip1, 4, ip2) &&
+    byte_equal(a->qtype, 2, b->qtype) &&
+    dns_domain_equal(a->query + 14, b->query + 14);
+}
+
+struct dns_transmit *inprogress[MAXUDP];
+
+static int try_merge(struct dns_transmit *d)
+{
+  int i;
+  for (i = 0; i < MAXUDP; i++) {
+    if (!inprogress[i]) continue;
+    if (!merge_equal(d, inprogress[i])) continue;
+    d->master = inprogress[i];
+    inprogress[i]->slaves[inprogress[i]->nslaves++] = d;
+    return 1;
+  }
+  return 0;
+}
+
+static void register_inprogress(struct dns_transmit *d)
+{
+  int i;
+  for (i = 0; i < MAXUDP; i++) {
+    if (!inprogress[i]) {
+      inprogress[i] = d;
+      return;
+    }
+  }
+  strerr_die1x(100, "BUG: out of inprogress slots");
+}
+
+static void unregister_inprogress(struct dns_transmit *d)
+{
+  int i;
+  for (i = 0; i < MAXUDP; i++) {
+    if (inprogress[i] == d)
+      inprogress[i] = 0;
+  }
+}
 
 static int serverwantstcp(const char *buf,unsigned int len)
 {
@@ -59,8 +114,28 @@ static void packetfree(struct dns_transmit *d)
   d->packet = 0;
 }
 
+static void mergefree(struct dns_transmit *d)
+{
+  int i;
+  if (merge_enable)
+    unregister_inprogress(d);
+  /* unregister us from our master */
+  if (d->master) {
+    for (i = 0; i < d->master->nslaves; i++)
+      if (d->master->slaves[i] == d)
+        d->master->slaves[i] = 0;
+  }
+  /* and unregister all of our slaves from us */
+  for (i = 0; i < d->nslaves; i++) {
+    if (d->slaves[i])
+      d->slaves[i]->master = NULL;
+  }
+  d->nslaves = 0;
+}
+
 static void queryfree(struct dns_transmit *d)
 {
+  mergefree(d);
   if (!d->query) return;
   alloc_free(d->query);
   d->query = 0;
@@ -99,11 +174,18 @@ static int thisudp(struct dns_transmit *d)
   const char *ip;
 
   socketfree(d);
+  mergefree(d);
 
   while (d->udploop < 4) {
     for (;d->curserver < 16;++d->curserver) {
       ip = d->servers + 4 * d->curserver;
       if (byte_diff(ip,4,"\0\0\0\0")) {
+        if (merge_enable && try_merge(d)) {
+          if (merge_logger)
+            merge_logger(ip, d->qtype, d->query + 14);
+          return 0;
+        }
+
 	d->query[2] = dns_random(256);
 	d->query[3] = dns_random(256);
   
@@ -118,6 +200,8 @@ static int thisudp(struct dns_transmit *d)
             taia_uint(&d->deadline,timeouts[d->udploop]);
             taia_add(&d->deadline,&d->deadline,&now);
             d->tcpstate = 0;
+            if (merge_enable)
+              register_inprogress(d);
             return 0;
           }
   
@@ -226,8 +310,12 @@ void dns_transmit_io(struct dns_transmit *d,iopause_fd *x,struct taia *deadline)
   x->fd = d->s1 - 1;
 
   switch(d->tcpstate) {
-    case 0: case 3: case 4: case 5:
-      x->events = IOPAUSE_READ;
+    case 0:
+      if (d->master) return;
+      if (d->packet) { taia_now(deadline); return; }
+      /* otherwise, fall through */
+    case 3: case 4: case 5:
+        x->events = IOPAUSE_READ;
       break;
     case 1: case 2:
       x->events = IOPAUSE_WRITE;
@@ -244,10 +332,14 @@ int dns_transmit_get(struct dns_transmit *d,const iopause_fd *x,const struct tai
   unsigned char ch;
   int r;
   int fd;
+  int i;
 
   errno = error_io;
   fd = d->s1 - 1;
 
+  if (d->tcpstate == 0 && d->master) return 0;
+  if (d->tcpstate == 0 && d->packet) return 1;
+
   if (!x->revents) {
     if (taia_less(when,&d->deadline)) return 0;
     errno = error_timeout;
@@ -279,6 +371,15 @@ have sent query to curserver on UDP socket s
     d->packet = alloc(d->packetlen);
     if (!d->packet) { dns_transmit_free(d); return -1; }
     byte_copy(d->packet,d->packetlen,udpbuf);
+
+    for (i = 0; i < d->nslaves; i++) {
+      if (!d->slaves[i]) continue;
+      d->slaves[i]->packetlen = d->packetlen;
+      d->slaves[i]->packet = alloc(d->packetlen);
+      if (!d->slaves[i]->packet) { dns_transmit_free(d->slaves[i]); continue; }
+      byte_copy(d->slaves[i]->packet,d->packetlen,udpbuf);
+    }
+
     queryfree(d);
     return 1;
   }
diff --git a/dnscache.c b/dnscache.c
index 8c899a3..c8db179 100644
--- a/dnscache.c
+++ b/dnscache.c
@@ -54,7 +54,6 @@ uint64 numqueries = 0;
 
 static int udp53;
 
-#define MAXUDP 200
 static struct udpclient {
   struct query q;
   struct taia start;
@@ -131,7 +130,6 @@ void u_new(void)
 
 static int tcp53;
 
-#define MAXTCP 20
 struct tcpclient {
   struct query q;
   struct taia start;
@@ -435,6 +433,8 @@ int main()
     response_hidettl();
   if (env_get("FORWARDONLY"))
     query_forwardonly();
+  if (env_get("MERGEQUERIES"))
+    dns_enable_merge(log_merge);
 
   if (!roots_init())
     strerr_die2sys(111,FATAL,"unable to read servers: ");
diff --git a/log.c b/log.c
index c43e8b0..3e1c674 100644
--- a/log.c
+++ b/log.c
@@ -150,6 +150,12 @@ void log_tx(const char *q,const char qtype[2],const char *control,const char ser
   line();
 }
 
+void log_merge(const char *addr, const char qtype[2], const char *q)
+{
+  string("merge "); ip(addr); space(); logtype(qtype); space(); name(q);
+  line();
+}
+
 void log_cachedanswer(const char *q,const char type[2])
 {
   string("cached "); logtype(type); space();
diff --git a/log.h b/log.h
index fe62fa3..dd6408a 100644
--- a/log.h
+++ b/log.h
@@ -18,6 +18,7 @@ extern void log_cachednxdomain(const char *);
 extern void log_cachedns(const char *,const char *);
 
 extern void log_tx(const char *,const char *,const char *,const char *,unsigned int);
+extern void log_merge(const char *, const char *, const char *);
 
 extern void log_nxdomain(const char *,const char *,unsigned int);
 extern void log_nodata(const char *,const char *,const char *,unsigned int);