dnsdist: Refactoring of EDNS code

diff --git a/pdns/dnsdist-ecs.cc b/pdns/dnsdist-ecs.cc
index 251eae661acc0a19e5721e7b080d31ef6150870f..dc29efd4842b581dbad647e7113e49a7b8db3b16 100644 (file)
--- a/pdns/dnsdist-ecs.cc
+++ b/pdns/dnsdist-ecs.cc
@@ -251,20 +251,25 @@ static int getEDNSOption(char* optRR, const size_t len, const uint16_t wantedOpt
   return ENOENT;
 }
 
+void generateEDNSOption(uint16_t optionCode, const std::string& payload, std::string& res)
+{
+  const uint16_t ecsOptionCode = htons(optionCode);
+  const uint16_t payloadLen = htons(payload.length());
+  res.append((const char *) &ecsOptionCode, sizeof ecsOptionCode);
+  res.append((const char *) &payloadLen, sizeof payloadLen);
+  res.append(payload);
+}
+
 static void generateECSOption(const ComboAddress& source, string& res)
 {
-  const uint16_t ecsOptionCode = htons(EDNS0_OPTION_CODE_ECS);
   Netmask sourceNetmask(source, source.sin4.sin_family == AF_INET ? g_ECSSourcePrefixV4 : g_ECSSourcePrefixV6);
   EDNSSubnetOpts ecsOpts;
   ecsOpts.source = sourceNetmask;
   string payload = makeEDNSSubnetOptsString(ecsOpts);
-  const uint16_t payloadLen = htons(payload.length());
-  res.append((const char *) &ecsOptionCode, sizeof ecsOptionCode);
-  res.append((const char *) &payloadLen, sizeof payloadLen);
-  res.append(payload);
+  generateEDNSOption(EDNS0_OPTION_CODE_ECS, payload, res);
 }
 
-static void generateECSOptRR(const ComboAddress& source, string & res)
+void generateOptRR(const std::string& optRData, string& res)
 {
   const uint8_t name = 0;
   dnsrecordheader dh;
@@ -276,8 +281,6 @@ static void generateECSOptRR(const ComboAddress& source, string & res)
   dh.d_type = htons(QType::OPT);
   dh.d_class = htons(q_EdnsUDPPayloadSize);
   memcpy(&dh.d_ttl, &edns0, sizeof edns0);
-  string optRData;
-  generateECSOption(source, optRData);
   dh.d_clen = htons((uint16_t) optRData.length());
   res.assign((const char *) &name, sizeof name);
   res.append((const char *) &dh, sizeof dh);
@@ -389,7 +392,9 @@ void handleEDNSClientSubnet(char * const packet, const size_t packetSize, const
     /* we need to add a EDNS0 RR with one EDNS0 ECS option, fixing the AR count */
     string EDNSRR;
     struct dnsheader* dh = (struct dnsheader*) packet;
-    generateECSOptRR(remote, EDNSRR);
+    string optRData;
+    generateECSOption(remote, optRData);
+    generateOptRR(optRData, EDNSRR);
     uint16_t arcount = ntohs(dh->arcount);
     arcount++;
     dh->arcount = htons(arcount);

diff --git a/pdns/dnsdist-ecs.hh b/pdns/dnsdist-ecs.hh
index 74138fdda6485c455a72e693e77c41dfa5c40b62..eb782968a46c139a9430623f2e316f57d1c26d9f 100644 (file)
--- a/pdns/dnsdist-ecs.hh
+++ b/pdns/dnsdist-ecs.hh
@@ -3,6 +3,8 @@
 int rewriteResponseWithoutEDNS(const char * packet, size_t len, vector<uint8_t>& newContent);
 int locateEDNSOptRR(const char * packet, size_t len, const char ** optStart, size_t * optLen, bool * last);
 void handleEDNSClientSubnet(char * packet, size_t packetSize, unsigned int consumed, uint16_t * len, string& largerPacket, bool * ednsAdded, const ComboAddress& remote);
+void generateOptRR(const std::string& optRData, string& res);
+void generateEDNSOption(uint16_t optionCode, const std::string& payload, std::string& res);
 
 
 

diff --git a/pdns/dnsrulactions.hh b/pdns/dnsrulactions.hh
index b047d7beb563347224fd92d1229b3754415331d0..c32655d255c6638d7f8bbb2bff50a09c2506cc84 100644 (file)
--- a/pdns/dnsrulactions.hh
+++ b/pdns/dnsrulactions.hh
@@ -1,4 +1,5 @@
 #include "dnsdist.hh"
+#include "dnsdist-ecs.hh"
 #include "dnsname.hh"
 #include "dolog.hh"
 
@@ -521,32 +522,24 @@ public:
     if(dq->dh->arcount)
       return Action::None;
 
-    char* dest = ((char*)dq->dh) + dq->len;    
-    dnsrecordheader dh;
-    EDNS0Record edns0;
-    edns0.extRCode = 0;
-    edns0.version = 0;
-    edns0.Z = 0;
-    
-    dh.d_type = htons(QType::OPT);
-    dh.d_class = htons(1500);
-    memcpy(&dh.d_ttl, &edns0, sizeof edns0);
+    string mac = getMACAddress(*dq->remote);
+    if(mac.empty())
+      return Action::None;
+
     string optRData;
-    const uint16_t optionCode=htons(d_code), optionLen=htons(6);
-    optRData.append((const char*)&optionCode, 2);
-    optRData.append((const char*)&optionLen, 2);
-    string mac=getMACAddress(*dq->remote);
-    if(!mac.empty()) {
-      dq->dh->arcount=ntohs(1);
-      optRData.append(mac);
-      dh.d_clen = htons((uint16_t) optRData.length());
-      uint8_t name=0;
-      string res((const char *) &name, sizeof name);
-      res.append((const char *) &dh, sizeof dh);
-      res.append(optRData.c_str(), optRData.length());
-      memcpy(dest, res.c_str(), res.length());
-      dq->len+=res.length();
-    }
+    generateEDNSOption(d_code, mac, optRData);
+
+    string res;
+    generateOptRR(optRData, res);
+
+    if ((dq->size - dq->len) < res.length())
+      return Action::None;
+
+    dq->dh->arcount = htons(1);
+    char* dest = ((char*)dq->dh) + dq->len;
+    memcpy(dest, res.c_str(), res.length());
+    dq->len += res.length();
+
     return Action::None;
   }  
   string toString() const override