Handle NOD flag in rr data

diff --git a/pdns/pdns_recursor.cc b/pdns/pdns_recursor.cc
index 21464655999b11f46d8483609e62f262d9636093..3afc6255357c3da76490402af29a405c6cab1f08 100644 (file)
--- a/pdns/pdns_recursor.cc
+++ b/pdns/pdns_recursor.cc
@@ -895,10 +895,10 @@ static void protobufLogQuery(uint8_t maskV4, uint8_t maskV6, const boost::uuids:
   m.setDeviceName(deviceName);
 
   if (!policyTags.empty()) {
-    m.setPolicyTags(policyTags);
+    m.addPolicyTags(policyTags);
   }
 
-  std::string msg(m.movebuf());
+  std::string msg(m.finishAndMoveBuf());
   for (auto& server : *t_protobufServers) {
     server->queueData(msg);
   }
@@ -910,7 +910,7 @@ static void protobufLogResponse(pdns::ProtoZero::Message& message)
     return;
   }
 
-  std::string msg(message.movebuf());
+  std::string msg(message.finishAndMoveBuf());
   for (auto& server : *t_protobufServers) {
     server->queueData(msg);
   }
@@ -1867,15 +1867,17 @@ static void startDoResolve(void *p)
         pbMessage->setAppliedPolicyTrigger(appliedPolicy.d_trigger);
         pbMessage->setAppliedPolicyHit(appliedPolicy.d_hit);
       }
-      // XXX if (nod) 
-      pbMessage->setPolicyTags(dc->d_policyTags);
+      pbMessage->addPolicyTags(dc->d_policyTags);
       pbMessage->setInBytes(packet.size());
 
+      // Take s snap of the current protobuf buffer state to store in the PC
       pbDataForCache = boost::make_optional(RecursorPacketCache::PBData{
-        pbMessage->getmsgbuf(), 
-        pbMessage->getrspbuf(),
+        pbMessage->getMessageBuf(),
+        pbMessage->getResponseBuf(),
         !appliedPolicy.getName().empty() || !dc->d_policyTags.empty()});
-
+#ifdef NOD_ENABLED
+      pbMessage->clearUDR(pbDataForCache->d_response);
+#endif
       // Below are the fields that are not stored in the packet cache and will be appended here and on a cache hit
       if (g_useKernelTimestamp && dc->d_kernelTimestamp.tv_sec) {
         pbMessage->setQueryTime(dc->d_kernelTimestamp.tv_sec, dc->d_kernelTimestamp.tv_usec);
@@ -1913,18 +1915,8 @@ static void startDoResolve(void *p)
       if (dc->d_logResponse) {
         protobufLogResponse(*pbMessage);
       }
-#ifdef NOD_ENABLED
-      if (g_nodEnabled) {
-        pbMessage->setNewlyObservedDomain(false);
-        pbMessage->clearUDR();
-        if (nod)
-          pbMessage->removePolicyTag(g_nod_pbtag);
-        if (hasUDR)
-          pbMessage->removePolicyTag(g_udr_pbtag);
-      }
-#endif /* NOD_ENABLED */
     }
-#endif
+#endif /* HAVE_PROTOBUF */
     if(!dc->d_tcp) {
       struct msghdr msgh;
       struct iovec iov;
@@ -2749,6 +2741,11 @@ static string* doProcessUDPQuestion(const std::string& question, const ComboAddr
         pbMessage->setDeviceName(deviceName);
         pbMessage->setFromPort(source.getPort());
         pbMessage->setToPort(destination.getPort());
+#ifdef NOD_ENABLED
+        if (g_nodEnabled) {
+          pbMessage->setNewlyObservedDomain(false);
+        }
+#endif
         protobufLogResponse(*pbMessage);
       }
 #endif /* HAVE_PROTOBUF */

diff --git a/pdns/protozero.cc b/pdns/protozero.cc
index 43842e2274ebc99c41752e037cba88b13fd24d48..88dafab3437dbc975239918eb705451853cbc47c 100644 (file)
--- a/pdns/protozero.cc
+++ b/pdns/protozero.cc
@@ -163,5 +163,19 @@ void pdns::ProtoZero::Message::addRR(const DNSRecord& record, const std::set<uin
   }
 #ifdef NOD_ENABLED
   pbf_rr.add_bool(6, udr);
+  pbf_rr.commit();
+
+  // Save the offset of the byte containing the just added bool. We can do this since
+  // we know a bit about how protobuf's encoding works.
+  offsets.push_back(d_rspbuf.length() - 1);
 #endif
 }
+
+#ifdef NOD_ENABLED
+void pdns::ProtoZero::Message::clearUDR(std::string& str)
+{
+  for (auto i : offsets) {
+    str.at(i) = 0;
+  }
+}
+#endif

diff --git a/pdns/protozero.hh b/pdns/protozero.hh
index 7cd15055d987b1e1b4f4848cfeb3d9d1df30769b..08332b3037a82ea36f5d5faee7a6cc9d1ed9e4a7 100644 (file)
--- a/pdns/protozero.hh
+++ b/pdns/protozero.hh
@@ -37,35 +37,31 @@ namespace pdns {
       // Start a new messagebug, constaining separate data for the response part
       Message(std::string::size_type sz1, std::string::size_type sz2) : d_message{d_msgbuf}, d_response{d_rspbuf}
       {
-        if (d_msgbuf.capacity() < sz1) {
-          d_msgbuf.reserve(sz1);
-        }
-        if (d_rspbuf.capacity() < sz2) {
-          d_rspbuf.reserve(sz2);
-        }
+        // This is extra space in addition to what's already there
+        // Different from what string.reserve() does
+        d_message.reserve(sz1);
+        d_response.reserve(sz2);
       }
 
       // Construct a Message with (partially) constructed content
       Message(const std::string& buf1, const std::string& buf2, std::string::size_type sz1, std::string::size_type sz2) :
         d_msgbuf{buf1}, d_rspbuf{buf2}, d_message{d_msgbuf}, d_response{d_rspbuf}
       {
-        // We expect to grow the buffer
-        if (d_msgbuf.capacity() - d_msgbuf.length() < sz1) {
-          d_msgbuf.reserve(d_msgbuf.length() + sz1);
-        }
-        if (d_rspbuf.capacity() - d_rspbuf.length() < sz2) {
-          d_msgbuf.reserve(d_rspbuf.length() + sz2);
-        }
+        // We expect to grow the buffers
+        // This is extra space in addition to what's already there
+        // Different from what string.reserve() does
+        d_message.reserve(sz1);
+        d_response.reserve(sz2);
       }
-      const std::string& getmsgbuf() const
+      const std::string& getMessageBuf() const
       {
         return d_msgbuf;
       }
-      const std::string& getrspbuf() const
+      const std::string& getResponseBuf() const
       {
         return d_rspbuf;
       }
-      std::string&& movebuf()
+      std::string&& finishAndMoveBuf()
       {
         if (!d_rspbuf.empty()) {
           d_message.add_message(13, d_rspbuf);
@@ -184,14 +180,11 @@ namespace pdns {
 #else
       void addRR(const DNSRecord& record, const std::set<uint16_t>& exportTypes);
 #endif
-      void clearUDR()
-      {
-      }
       void setAppliedPolicy(const std::string& policy)
       {
         d_response.add_string(3, policy);
       }
-      void setPolicyTags(const std::unordered_set<std::string>& tags)
+      void addPolicyTags(const std::unordered_set<std::string>& tags)
       {
         for (const auto& tag : tags) {
           d_response.add_string(4, tag);
@@ -199,9 +192,7 @@ namespace pdns {
       }
       void addPolicyTag(const string& tag)
       {
-      }
-      void removePolicyTag(const string& tag)
-      {
+        d_response.add_string(4, tag);
       }
       void setQueryTime(uint32_t sec, uint32_t usec)
       {
@@ -245,11 +236,19 @@ namespace pdns {
         d_response.add_string(9, hit);
       }
 
+#ifdef NOD_ENABLED      
+      void clearUDR(std::string&);
+#endif
+
     private:
       std::string d_msgbuf;
       std::string d_rspbuf;
       protozero::pbf_writer d_message;
       protozero::pbf_writer d_response;
+
+#ifdef NOD_ENABLED
+      vector<std::string::size_type> offsets;
+#endif
     };
   };
 };