From 0b184c82201dcd3cd54e79254c0d1eb19478ac34 Mon Sep 17 00:00:00 2001
From: Nicko Dehaine <nicko@threatstop.com>
Date: Fri, 5 Feb 2021 16:43:57 +0000
Subject: [PATCH] Additional adjustments to migration to protozero library

---
 pdns/protozero.hh                  | 32 ++++--------------------------
 pdns/recursordist/rec-protozero.hh | 29 +++++++++++++++++++++++++++
 2 files changed, 33 insertions(+), 28 deletions(-)

diff --git a/pdns/protozero.hh b/pdns/protozero.hh
index 04dd69eab9..a5e3e0e081 100644
--- a/pdns/protozero.hh
+++ b/pdns/protozero.hh
@@ -35,7 +35,7 @@ namespace pdns {
       enum class MessageType : int32_t { DNSQueryType = 1, DNSResponseType = 2, DNSOutgoingQueryType = 3, DNSIncomingResponseType = 4 };
       enum class Field : protozero::pbf_tag_type { type = 1, messageId = 2, serverIdentity = 3, socketFamily = 4, socketProtocol = 5, from = 6, to = 7, inBytes = 8, timeSec = 9, timeUsec = 10, id = 11, question = 12, response = 13, originalRequestorSubnet = 14, requestorId = 15, initialRequestId = 16, deviceId = 17, newlyObservedDomain = 18, deviceName = 19, fromPort = 20, toPort = 21 };
       enum class QuestionField : protozero::pbf_tag_type { qName = 1, qType = 2, qClass = 3};
-      enum class ResponseField : protozero::pbf_tag_type { rcode = 1, rrs = 2, appliedPolicy = 3, tags = 4, queryTimeSec = 5, queryTimeUsec = 6, appliedPolicyType = 7, appliedPolicyTrigger = 8, appliedPolicyHit = 9 };
+      enum class ResponseField : protozero::pbf_tag_type { rcode = 1, rrs = 2, appliedPolicy = 3, tags = 4, queryTimeSec = 5, queryTimeUsec = 6, appliedPolicyType = 7, appliedPolicyTrigger = 8, appliedPolicyHit = 9, appliedPolicyKind = 10 };
       enum class RRField : protozero::pbf_tag_type { name = 1, type = 2, class_ = 3, ttl = 4, rdata = 5, udr = 6 };
 
       Message(std::string& buffer): d_buffer(buffer), d_message{d_buffer}
@@ -235,33 +235,9 @@ namespace pdns {
         writer.add_uint64(static_cast<protozero::pbf_tag_type>(type), value);
       }
 
-      void setAppliedPolicyKind(const DNSFilterEngine::PolicyKind& kind)
-      {
-        uint32_t k;
-
-        switch(kind) {
-	case DNSFilterEngine::PolicyKind::NoAction:
-          k = 1;
-          break;
-	case DNSFilterEngine::PolicyKind::Drop:
-          k = 2;
-          break;
-	case DNSFilterEngine::PolicyKind::NXDOMAIN:
-          k = 3;
-          break;
-	case DNSFilterEngine::PolicyKind::NODATA:
-          k = 4;
-          break;
-	case DNSFilterEngine::PolicyKind::Truncate:
-          k = 5;
-          break;
-	case DNSFilterEngine::PolicyKind::Custom:
-          k = 6;
-          break;
-        default:
-          throw std::runtime_error("Unsupported protobuf policy kind");
-        }
-        d_response.add_uint32(10, k);
+      static void add_bytes(protozero::pbf_writer& writer, Field type, const char* data, size_t len)
+      {
+        writer.add_bytes(static_cast<protozero::pbf_tag_type>(type), data, len);
       }
 
       static void add_string(protozero::pbf_writer& writer, Field type, const std::string& str)
diff --git a/pdns/recursordist/rec-protozero.hh b/pdns/recursordist/rec-protozero.hh
index d63e2397ff..efa77150f8 100644
--- a/pdns/recursordist/rec-protozero.hh
+++ b/pdns/recursordist/rec-protozero.hh
@@ -141,6 +141,35 @@ namespace ProtoZero
       d_response.add_string(static_cast<protozero::pbf_tag_type>(ResponseField::appliedPolicyHit), hit);
     }
 
+    void setAppliedPolicyKind(const DNSFilterEngine::PolicyKind& kind)
+    {
+        uint32_t k;
+
+        switch(kind) {
+	case DNSFilterEngine::PolicyKind::NoAction:
+          k = 1;
+          break;
+	case DNSFilterEngine::PolicyKind::Drop:
+          k = 2;
+          break;
+	case DNSFilterEngine::PolicyKind::NXDOMAIN:
+          k = 3;
+          break;
+	case DNSFilterEngine::PolicyKind::NODATA:
+          k = 4;
+          break;
+	case DNSFilterEngine::PolicyKind::Truncate:
+          k = 5;
+          break;
+	case DNSFilterEngine::PolicyKind::Custom:
+          k = 6;
+          break;
+        default:
+          throw std::runtime_error("Unsupported protobuf policy kind");
+        }
+        d_response.add_uint32(static_cast<protozero::pbf_tag_type>(ResponseField::appliedPolicyKind), k);
+    }
+
 #ifdef NOD_ENABLED
     void clearUDR(std::string&);
 #endif
-- 
2.47.2