From: Remi Gacogne <rgacogne-github@coredump.fr>
Date: Wed, 18 Nov 2015 16:38:46 +0000 (+0100)
Subject: Add setCD()/getCD(). Restore initial flags in the response (RD/CD for now).
X-Git-Tag: dnsdist-1.0.0-alpha1~210^2~17^2~1
X-Git-Url: https://granicus.if.org/sourcecode?a=commitdiff_plain;h=aeb36780467ddc24e4b374105c31ad516f808238;p=pdns

Add setCD()/getCD(). Restore initial flags in the response (RD/CD for now).
---

diff --git a/pdns/README-dnsdist.md b/pdns/README-dnsdist.md
index cd1db7848..5a4a75c41 100644
--- a/pdns/README-dnsdist.md
+++ b/pdns/README-dnsdist.md
@@ -640,6 +640,8 @@ Here are all functions:
      * member `setRD(bool)`: set recursion desired flag
      * member `setTC(bool)`: set truncation flag (TC)
      * member `setQR(bool)`: set Query Response flag (setQR(true) indicates an *answer* packet)
+     * member `getCD()`: get checking disabled flag
+     * member `setCD(bool)`: set checking disabled flag
    * NetmaskGroup related
      * nothing yet
    * QPSLimiter related:
diff --git a/pdns/dns.hh b/pdns/dns.hh
index b013beb24..9295b5363 100644
--- a/pdns/dns.hh
+++ b/pdns/dns.hh
@@ -217,6 +217,18 @@ struct dnsheader {
         unsigned        arcount :16;    /* number of resource entries */
 };
 
+inline uint16_t * getFlagsFromDNSHeader(struct dnsheader * dh)
+{
+  return (uint16_t*) (((char *) dh) + sizeof(uint16_t));
+}
+
+#if BYTE_ORDER == BIG_ENDIAN
+#define FLAGS_RD_OFFSET (8)
+#define FLAGS_CD_OFFSET (12)
+#elif BYTE_ORDER == LITTLE_ENDIAN || BYTE_ORDER == PDP_ENDIAN
+#define FLAGS_RD_OFFSET (0)
+#define FLAGS_CD_OFFSET (12)
+#endif
 
 #define L theL()
 extern time_t s_starttime;
diff --git a/pdns/dnsdist-lua.cc b/pdns/dnsdist-lua.cc
index b1ddce066..32c12f848 100644
--- a/pdns/dnsdist-lua.cc
+++ b/pdns/dnsdist-lua.cc
@@ -513,6 +513,14 @@ vector<std::function<void(void)>> setupLua(bool client, const std::string& confi
       return (bool)dh.rd;
     });
 
+  g_lua.registerFunction<void(dnsheader::*)(bool)>("setCD", [](dnsheader& dh, bool v) {
+      dh.cd=v;
+    });
+
+  g_lua.registerFunction<bool(dnsheader::*)()>("getCD", [](dnsheader& dh) {
+      return (bool)dh.cd;
+    });
+
 
   g_lua.registerFunction<void(dnsheader::*)(bool)>("setTC", [](dnsheader& dh, bool v) {
       dh.tc=v;
diff --git a/pdns/dnsdist-tcp.cc b/pdns/dnsdist-tcp.cc
index 30287084b..a69c0ae72 100644
--- a/pdns/dnsdist-tcp.cc
+++ b/pdns/dnsdist-tcp.cc
@@ -134,6 +134,9 @@ void* tcpClientThread(int pipefd)
 
     uint16_t qlen, rlen;
     string pool; 
+    const uint16_t rdMask = 1 << FLAGS_RD_OFFSET;
+    const uint16_t cdMask = 1 << FLAGS_CD_OFFSET;
+    const uint16_t restoreFlagsMask = UINT16_MAX & ~(rdMask | cdMask);
 
     shared_ptr<DownstreamState> ds;
     if (!setNonBlocking(ci.fd))
@@ -150,6 +153,8 @@ void* tcpClientThread(int pipefd)
 	DNSName qname(query, qlen, 12, false, &qtype);
 	string ruleresult;
 	struct dnsheader* dh =(dnsheader*)query;
+	const uint16_t * flags = getFlagsFromDNSHeader(dh);
+	uint16_t origFlags = *flags;
 	
         if(blockFilter) {
 	  std::lock_guard<std::mutex> lock(g_luamutex);
@@ -260,7 +265,15 @@ void* tcpClientThread(int pipefd)
 
         char answerbuffer[rlen];
         readn2WithTimeout(dsock, answerbuffer, rlen, ds->tcpRecvTimeout);
-      
+        struct dnsheader* responseHeaders = (struct dnsheader*)answerbuffer;
+        uint16_t * responseFlags = getFlagsFromDNSHeader(responseHeaders);
+        /* clear the flags we are about to restore */
+        *responseFlags &= restoreFlagsMask;
+        /* only keep the flags we want to restore */
+        origFlags &= ~restoreFlagsMask;
+        /* set the saved flags as they were */
+        *responseFlags |= origFlags;
+
         if (putNonBlockingMsgLen(ci.fd, rlen, ds->tcpSendTimeout))
           writen2WithTimeout(ci.fd, answerbuffer, rlen, ds->tcpSendTimeout);
       }
diff --git a/pdns/dnsdist.cc b/pdns/dnsdist.cc
index 9f5d60fe0..5e2137743 100644
--- a/pdns/dnsdist.cc
+++ b/pdns/dnsdist.cc
@@ -143,6 +143,9 @@ DelayPipe<DelayedPacket> * g_delay = 0;
 void* responderThread(std::shared_ptr<DownstreamState> state)
 {
   char packet[4096];
+  const uint16_t rdMask = 1 << FLAGS_RD_OFFSET;
+  const uint16_t cdMask = 1 << FLAGS_CD_OFFSET;
+  const uint16_t restoreFlagsMask = UINT16_MAX & ~(rdMask | cdMask);
   
   struct dnsheader* dh = (struct dnsheader*)packet;
   int len;
@@ -164,6 +167,15 @@ void* responderThread(std::shared_ptr<DownstreamState> state)
       truncateTC(packet, (unsigned int*)&len);
     }
 
+    uint16_t * flags = getFlagsFromDNSHeader(dh);
+    uint16_t origFlags = ids->origFlags;
+    /* clear the flags we are about to restore */
+    *flags &= restoreFlagsMask;
+    /* only keep the flags we want to restore */
+    origFlags &= ~restoreFlagsMask;
+    /* set the saved flags as they were */
+    *flags |= origFlags;
+
     dh->id = ids->origID;
     g_stats.responses++;
 
@@ -426,6 +438,8 @@ try
       if(dh->qr)    // don't respond to responses
 	continue;
       
+      const uint16_t * flags = getFlagsFromDNSHeader(dh);
+      const uint16_t origFlags = *flags;
       DNSName qname(packet, len, 12, false, &qtype);
       g_rings.queryRing.push_back(qname);
             
@@ -525,6 +539,7 @@ try
       ids->qtype = qtype;
       ids->origDest.sin4.sin_family=0;
       ids->delayMsec = delayMsec;
+      ids->origFlags = origFlags;
       HarvestDestinationAddress(&msgh, &ids->origDest);
       
       dh->id = idOffset;
diff --git a/pdns/dnsdist.hh b/pdns/dnsdist.hh
index 871abbbef..d3450e1fb 100644
--- a/pdns/dnsdist.hh
+++ b/pdns/dnsdist.hh
@@ -168,6 +168,7 @@ struct IDState
   std::atomic<uint16_t> age;                                  // 4
   uint16_t qtype;                                             // 2
   uint16_t origID;                                            // 2
+  uint16_t origFlags;                                         // 2
   int delayMsec;
 };