dnsdist: Reduce the number of string alloc/copy during ECS processing

diff --git a/pdns/dnsdist-ecs.cc b/pdns/dnsdist-ecs.cc
index 0df7c210cc2927317ba5137bebaf8a559cfbfaed..05bfd53aa9a2d6b7b5243aec5828fd3ff1c32ced 100644 (file)
--- a/pdns/dnsdist-ecs.cc
+++ b/pdns/dnsdist-ecs.cc
@@ -36,11 +36,10 @@ uint16_t g_ECSSourcePrefixV6 = 56;
 
 bool g_ECSOverride{false};
 
-int rewriteResponseWithoutEDNS(const char * packet, const size_t len, vector<uint8_t>& newContent)
+int rewriteResponseWithoutEDNS(const std::string& initialPacket, vector<uint8_t>& newContent)
 {
-  assert(packet != NULL);
-  assert(len >= sizeof(dnsheader));
-  const struct dnsheader* dh = (const struct dnsheader*) packet;
+  assert(initialPacket.size() >= sizeof(dnsheader));
+  const struct dnsheader* dh = reinterpret_cast<const struct dnsheader*>(initialPacket.data());
 
   if (ntohs(dh->arcount) == 0)
     return ENOENT;
@@ -48,9 +47,8 @@ int rewriteResponseWithoutEDNS(const char * packet, const size_t len, vector<uin
   if (ntohs(dh->qdcount) == 0)
     return ENOENT;
 
-  std::string packetStr(packet, len);
-  PacketReader pr(packetStr);
-  
+  PacketReader pr(initialPacket);
+
   size_t idx = 0;
   DNSName rrname;
   uint16_t qdcount = ntohs(dh->qdcount);
@@ -124,19 +122,17 @@ int rewriteResponseWithoutEDNS(const char * packet, const size_t len, vector<uin
   return 0;
 }
 
-int locateEDNSOptRR(char * packet, const size_t len, char ** optStart, size_t * optLen, bool * last)
+int locateEDNSOptRR(const std::string& packet, uint16_t * optStart, size_t * optLen, bool * last)
 {
-  assert(packet != NULL);
   assert(optStart != NULL);
   assert(optLen != NULL);
   assert(last != NULL);
-  const struct dnsheader* dh = (const struct dnsheader*) packet;
+  const struct dnsheader* dh = reinterpret_cast<const struct dnsheader*>(packet.data());
 
   if (ntohs(dh->arcount) == 0)
     return ENOENT;
 
-  std::string packetStr(packet, len);
-  PacketReader pr(packetStr);
+  PacketReader pr(packet);
   size_t idx = 0;
   DNSName rrname;
   uint16_t qdcount = ntohs(dh->qdcount);
@@ -170,10 +166,10 @@ int locateEDNSOptRR(char * packet, const size_t len, char ** optStart, size_t *
     pr.getDnsrecordheader(ah);
 
     if (ah.d_type == QType::OPT) {
-      *optStart = packet + start;
+      *optStart = start;
       *optLen = (pr.getPosition() - start) + ah.d_clen;
 
-      if ((packet + len) < (*optStart + *optLen)) {
+      if (packet.size() < (*optStart + *optLen)) {
         throw std::range_error("Opt record overflow");
       }
 
@@ -431,11 +427,10 @@ int removeEDNSOptionFromOPT(char* optStart, size_t* optLen, const uint16_t optio
   return 0;
 }
 
-int rewriteResponseWithoutEDNSOption(const char * packet, const size_t len, const uint16_t optionCodeToSkip, vector<uint8_t>& newContent)
+int rewriteResponseWithoutEDNSOption(const std::string& initialPacket, const uint16_t optionCodeToSkip, vector<uint8_t>& newContent)
 {
-  assert(packet != NULL);
-  assert(len >= sizeof(dnsheader));
-  const struct dnsheader* dh = (const struct dnsheader*) packet;
+  assert(initialPacket.size() >= sizeof(dnsheader));
+  const struct dnsheader* dh = reinterpret_cast<const struct dnsheader*>(initialPacket.data());
 
   if (ntohs(dh->arcount) == 0)
     return ENOENT;
@@ -443,8 +438,7 @@ int rewriteResponseWithoutEDNSOption(const char * packet, const size_t len, cons
   if (ntohs(dh->qdcount) == 0)
     return ENOENT;
 
-  std::string packetStr(packet, len);
-  PacketReader pr(packetStr);
+  PacketReader pr(initialPacket);
 
   size_t idx = 0;
   DNSName rrname;

diff --git a/pdns/dnsdist-ecs.hh b/pdns/dnsdist-ecs.hh
index a4e145a344b00a0280e2754d369078de37f624d4..d280ee54ffccc9b55fa35ae8bd2eb4d2dd144efb 100644 (file)
--- a/pdns/dnsdist-ecs.hh
+++ b/pdns/dnsdist-ecs.hh
@@ -21,10 +21,10 @@
  */
 #pragma once
 
-int rewriteResponseWithoutEDNS(const char * packet, size_t len, vector<uint8_t>& newContent);
-int locateEDNSOptRR(char * packet, size_t len, char ** optStart, size_t * optLen, bool * last);
+int rewriteResponseWithoutEDNS(const std::string& initialPacket, vector<uint8_t>& newContent);
+int locateEDNSOptRR(const std::string& packet, uint16_t * optStart, size_t * optLen, bool * last);
 bool handleEDNSClientSubnet(char * packet, size_t packetSize, unsigned int consumed, uint16_t * len, bool* ednsAdded, bool* ecsAdded, const ComboAddress& remote, bool overrideExisting, uint16_t ecsPrefixLength);
 void generateOptRR(const std::string& optRData, string& res);
 int removeEDNSOptionFromOPT(char* optStart, size_t* optLen, const uint16_t optionCodeToRemove);
-int rewriteResponseWithoutEDNSOption(const char * packet, const size_t len, const uint16_t optionCodeToSkip, vector<uint8_t>& newContent);
+int rewriteResponseWithoutEDNSOption(const std::string& initialPacket, const uint16_t optionCodeToSkip, vector<uint8_t>& newContent);
 int getEDNSOptionsStart(char* packet, const size_t offset, const size_t len, char ** optRDLen, size_t * remaining);

diff --git a/pdns/dnsdist.cc b/pdns/dnsdist.cc
index 45862626ecbfd0c06477939565997bf673161f5a..e4c0cd2c9d3c420a332467768c5207b47c98ff66 100644 (file)
--- a/pdns/dnsdist.cc
+++ b/pdns/dnsdist.cc
@@ -273,11 +273,12 @@ bool fixUpResponse(char** response, uint16_t* responseLen, size_t* responseSize,
   }
 
   if (ednsAdded || ecsAdded) {
-    char * optStart = NULL;
+    uint16_t optStart = NULL;
     size_t optLen = 0;
     bool last = false;
 
-    int res = locateEDNSOptRR(*response, *responseLen, &optStart, &optLen, &last);
+    const std::string responseStr(*response, *responseLen);
+    int res = locateEDNSOptRR(responseStr, &optStart, &optLen, &last);
 
     if (res == 0) {
       if (ednsAdded) {
@@ -292,7 +293,7 @@ bool fixUpResponse(char** response, uint16_t* responseLen, size_t* responseSize,
         }
         else {
           /* Removing an intermediary RR could lead to compression error */
-          if (rewriteResponseWithoutEDNS(*response, *responseLen, rewrittenResponse) == 0) {
+          if (rewriteResponseWithoutEDNS(responseStr, rewrittenResponse) == 0) {
             *responseLen = rewrittenResponse.size();
             if (addRoom && (UINT16_MAX - *responseLen) > addRoom) {
               rewrittenResponse.reserve(*responseLen + addRoom);
@@ -312,12 +313,12 @@ bool fixUpResponse(char** response, uint16_t* responseLen, size_t* responseSize,
           /* nothing after the OPT RR, we can simply remove the
              ECS option */
           size_t existingOptLen = optLen;
-          removeEDNSOptionFromOPT(optStart, &optLen, EDNSOptionCode::ECS);
+          removeEDNSOptionFromOPT(*response + optStart, &optLen, EDNSOptionCode::ECS);
           *responseLen -= (existingOptLen - optLen);
         }
         else {
           /* Removing an intermediary RR could lead to compression error */
-          if (rewriteResponseWithoutEDNSOption(*response, *responseLen, EDNSOptionCode::ECS, rewrittenResponse) == 0) {
+          if (rewriteResponseWithoutEDNSOption(responseStr, EDNSOptionCode::ECS, rewrittenResponse) == 0) {
             *responseLen = rewrittenResponse.size();
             if (addRoom && (UINT16_MAX - *responseLen) > addRoom) {
               rewrittenResponse.reserve(*responseLen + addRoom);

diff --git a/pdns/dnsdistdist/dnsdist-rules.hh b/pdns/dnsdistdist/dnsdist-rules.hh
index ed1e28ca1c69b8ad7ad3ff9bf335bc935b2cb746..4746ce2f2b9e9e63361f1431601a6de2bf3672c8 100644 (file)
--- a/pdns/dnsdistdist/dnsdist-rules.hh
+++ b/pdns/dnsdistdist/dnsdist-rules.hh
@@ -846,10 +846,12 @@ public:
       return false;
     }
 
-    char * optStart = NULL;
+    uint16_t optStart;
     size_t optLen = 0;
     bool last = false;
-    int res = locateEDNSOptRR(const_cast<char*>(reinterpret_cast<const char*>(dq->dh)), dq->len, &optStart, &optLen, &last);
+    const char * packet = reinterpret_cast<const char*>(dq->dh);
+    std::string packetStr(packet, dq->len);
+    int res = locateEDNSOptRR(packetStr, &optStart, &optLen, &last);
     if (res != 0) {
       // no EDNS OPT RR
       return d_extrcode == 0;
@@ -860,14 +862,14 @@ public:
       return false;
     }
 
-    if (*optStart != 0) {
+    if (optStart < dq->len && packet[optStart] != 0) {
       // OPT RR Name != '.'
       return false;
     }
     EDNS0Record edns0;
     static_assert(sizeof(EDNS0Record) == sizeof(uint32_t), "sizeof(EDNS0Record) must match sizeof(uint32_t) AKA RR TTL size");
     // copy out 4-byte "ttl" (really the EDNS0 record), after root label (1) + type (2) + class (2).
-    memcpy(&edns0, optStart + 5, sizeof edns0);
+    memcpy(&edns0, packet + optStart + 5, sizeof edns0);
 
     return d_extrcode == edns0.extRCode;
   }

diff --git a/pdns/test-dnsdist_cc.cc b/pdns/test-dnsdist_cc.cc
index b60f063a3b8d4f7b0bb2bb7c595a0b5298c6f1c5..75d0a5a62e1a16472bbbb7cd814314ef8950a23d 100644 (file)
--- a/pdns/test-dnsdist_cc.cc
+++ b/pdns/test-dnsdist_cc.cc
@@ -295,7 +295,7 @@ BOOST_AUTO_TEST_CASE(removeEDNSWhenFirst) {
   pw.commit();
 
   vector<uint8_t> newResponse;
-  int res = rewriteResponseWithoutEDNS((const char *) response.data(), response.size(), newResponse);
+  int res = rewriteResponseWithoutEDNS(std::string((const char *) response.data(), response.size()), newResponse);
   BOOST_CHECK_EQUAL(res, 0);
 
   unsigned int consumed = 0;
@@ -327,7 +327,7 @@ BOOST_AUTO_TEST_CASE(removeEDNSWhenIntermediary) {
   pw.commit();
 
   vector<uint8_t> newResponse;
-  int res = rewriteResponseWithoutEDNS((const char *) response.data(), response.size(), newResponse);
+  int res = rewriteResponseWithoutEDNS(std::string((const char *) response.data(), response.size()), newResponse);
   BOOST_CHECK_EQUAL(res, 0);
 
   unsigned int consumed = 0;
@@ -357,7 +357,7 @@ BOOST_AUTO_TEST_CASE(removeEDNSWhenLast) {
   pw.commit();
 
   vector<uint8_t> newResponse;
-  int res = rewriteResponseWithoutEDNS((const char *) response.data(), response.size(), newResponse);
+  int res = rewriteResponseWithoutEDNS(std::string((const char *) response.data(), response.size()), newResponse);
 
   BOOST_CHECK_EQUAL(res, 0);
 
@@ -394,18 +394,18 @@ BOOST_AUTO_TEST_CASE(removeECSWhenOnlyOption) {
   pw.addOpt(512, 0, 0, opts);
   pw.commit();
 
-  char * optStart = NULL;
+  uint16_t optStart;
   size_t optLen = 0;
   bool last = false;
 
-  int res = locateEDNSOptRR((char *) response.data(), response.size(), &optStart, &optLen, &last);
+  int res = locateEDNSOptRR(std::string((char *) response.data(), response.size()), &optStart, &optLen, &last);
   BOOST_CHECK_EQUAL(res, 0);
   BOOST_CHECK_EQUAL(last, true);
 
   size_t responseLen = response.size();
   size_t existingOptLen = optLen;
   BOOST_CHECK(existingOptLen < responseLen);
-  res = removeEDNSOptionFromOPT(optStart, &optLen, EDNSOptionCode::ECS);
+  res = removeEDNSOptionFromOPT(reinterpret_cast<char *>(response.data()) + optStart, &optLen, EDNSOptionCode::ECS);
   BOOST_CHECK_EQUAL(res, 0);
   BOOST_CHECK_EQUAL(optLen, existingOptLen - (origECSOptionStr.size() + 4));
   responseLen -= (existingOptLen - optLen);
@@ -446,18 +446,18 @@ BOOST_AUTO_TEST_CASE(removeECSWhenFirstOption) {
   pw.addOpt(512, 0, 0, opts);
   pw.commit();
 
-  char * optStart = NULL;
+  uint16_t optStart;
   size_t optLen = 0;
   bool last = false;
 
-  int res = locateEDNSOptRR((char *) response.data(), response.size(), &optStart, &optLen, &last);
+  int res = locateEDNSOptRR(std::string((char *) response.data(), response.size()), &optStart, &optLen, &last);
   BOOST_CHECK_EQUAL(res, 0);
   BOOST_CHECK_EQUAL(last, true);
 
   size_t responseLen = response.size();
   size_t existingOptLen = optLen;
   BOOST_CHECK(existingOptLen < responseLen);
-  res = removeEDNSOptionFromOPT(optStart, &optLen, EDNSOptionCode::ECS);
+  res = removeEDNSOptionFromOPT(reinterpret_cast<char *>(response.data()) + optStart, &optLen, EDNSOptionCode::ECS);
   BOOST_CHECK_EQUAL(res, 0);
   BOOST_CHECK_EQUAL(optLen, existingOptLen - (origECSOptionStr.size() + 4));
   responseLen -= (existingOptLen - optLen);
@@ -502,18 +502,18 @@ BOOST_AUTO_TEST_CASE(removeECSWhenIntermediaryOption) {
   pw.addOpt(512, 0, 0, opts);
   pw.commit();
 
-  char * optStart = NULL;
+  uint16_t optStart;
   size_t optLen = 0;
   bool last = false;
 
-  int res = locateEDNSOptRR((char *) response.data(), response.size(), &optStart, &optLen, &last);
+  int res = locateEDNSOptRR(std::string((char *) response.data(), response.size()), &optStart, &optLen, &last);
   BOOST_CHECK_EQUAL(res, 0);
   BOOST_CHECK_EQUAL(last, true);
 
   size_t responseLen = response.size();
   size_t existingOptLen = optLen;
   BOOST_CHECK(existingOptLen < responseLen);
-  res = removeEDNSOptionFromOPT(optStart, &optLen, EDNSOptionCode::ECS);
+  res = removeEDNSOptionFromOPT(reinterpret_cast<char *>(response.data()) + optStart, &optLen, EDNSOptionCode::ECS);
   BOOST_CHECK_EQUAL(res, 0);
   BOOST_CHECK_EQUAL(optLen, existingOptLen - (origECSOptionStr.size() + 4));
   responseLen -= (existingOptLen - optLen);
@@ -554,18 +554,18 @@ BOOST_AUTO_TEST_CASE(removeECSWhenLastOption) {
   pw.addOpt(512, 0, 0, opts);
   pw.commit();
 
-  char * optStart = NULL;
+  uint16_t optStart;
   size_t optLen = 0;
   bool last = false;
 
-  int res = locateEDNSOptRR((char *) response.data(), response.size(), &optStart, &optLen, &last);
+  int res = locateEDNSOptRR(std::string((char *) response.data(), response.size()), &optStart, &optLen, &last);
   BOOST_CHECK_EQUAL(res, 0);
   BOOST_CHECK_EQUAL(last, true);
 
   size_t responseLen = response.size();
   size_t existingOptLen = optLen;
   BOOST_CHECK(existingOptLen < responseLen);
-  res = removeEDNSOptionFromOPT(optStart, &optLen, EDNSOptionCode::ECS);
+  res = removeEDNSOptionFromOPT(reinterpret_cast<char *>(response.data()) + optStart, &optLen, EDNSOptionCode::ECS);
   BOOST_CHECK_EQUAL(res, 0);
   BOOST_CHECK_EQUAL(optLen, existingOptLen - (origECSOptionStr.size() + 4));
   responseLen -= (existingOptLen - optLen);
@@ -602,7 +602,7 @@ BOOST_AUTO_TEST_CASE(rewritingWithoutECSWhenOnlyOption) {
   pw.commit();
 
   vector<uint8_t> newResponse;
-  int res = rewriteResponseWithoutEDNSOption((const char *) response.data(), response.size(), EDNSOptionCode::ECS, newResponse);
+  int res = rewriteResponseWithoutEDNSOption(std::string((const char *) response.data(), response.size()), EDNSOptionCode::ECS, newResponse);
   BOOST_CHECK_EQUAL(res, 0);
 
   BOOST_CHECK_EQUAL(newResponse.size(), response.size() - (origECSOptionStr.size() + 4));
@@ -644,7 +644,7 @@ BOOST_AUTO_TEST_CASE(rewritingWithoutECSWhenFirstOption) {
   pw.commit();
 
   vector<uint8_t> newResponse;
-  int res = rewriteResponseWithoutEDNSOption((const char *) response.data(), response.size(), EDNSOptionCode::ECS, newResponse);
+  int res = rewriteResponseWithoutEDNSOption(std::string((const char *) response.data(), response.size()), EDNSOptionCode::ECS, newResponse);
   BOOST_CHECK_EQUAL(res, 0);
 
   BOOST_CHECK_EQUAL(newResponse.size(), response.size() - (origECSOptionStr.size() + 4));
@@ -688,7 +688,7 @@ BOOST_AUTO_TEST_CASE(rewritingWithoutECSWhenIntermediaryOption) {
   pw.commit();
 
   vector<uint8_t> newResponse;
-  int res = rewriteResponseWithoutEDNSOption((const char *) response.data(), response.size(), EDNSOptionCode::ECS, newResponse);
+  int res = rewriteResponseWithoutEDNSOption(std::string((const char *) response.data(), response.size()), EDNSOptionCode::ECS, newResponse);
   BOOST_CHECK_EQUAL(res, 0);
 
   BOOST_CHECK_EQUAL(newResponse.size(), response.size() - (origECSOptionStr.size() + 4));
@@ -730,7 +730,7 @@ BOOST_AUTO_TEST_CASE(rewritingWithoutECSWhenLastOption) {
   pw.commit();
 
   vector<uint8_t> newResponse;
-  int res = rewriteResponseWithoutEDNSOption((const char *) response.data(), response.size(), EDNSOptionCode::ECS, newResponse);
+  int res = rewriteResponseWithoutEDNSOption(std::string((const char *) response.data(), response.size()), EDNSOptionCode::ECS, newResponse);
   BOOST_CHECK_EQUAL(res, 0);
 
   BOOST_CHECK_EQUAL(newResponse.size(), response.size() - (origECSOptionStr.size() + 4));