New methods to add and remove individual policy tags

diff --git a/pdns/rec-protobuf.cc b/pdns/rec-protobuf.cc
index f12d793510cad3b47100825553b34773f86b9032..c7eec6943dc2138328c1090e4a6e0ce5136ecb3d 100644 (file)
--- a/pdns/rec-protobuf.cc
+++ b/pdns/rec-protobuf.cc
@@ -179,6 +179,39 @@ void RecProtoBufMessage::setPolicyTags(const std::vector<std::string>& policyTag
 #endif /* HAVE_PROTOBUF */
 }
 
+void RecProtoBufMessage::addPolicyTag(const std::string& policyTag)
+{
+#ifdef HAVE_PROTOBUF
+  PBDNSMessage_DNSResponse* response = d_message.mutable_response();
+  if (response) {
+    response->add_tags(policyTag);
+  }
+#endif
+}
+
+void RecProtoBufMessage::removePolicyTag(const std::string& policyTag)
+{
+#ifdef HAVE_PROTOBUF
+  PBDNSMessage_DNSResponse* response = d_message.mutable_response();
+  if (response) {
+    const int count = response->tags_size();
+    int keep = 0;
+    for (int idx = 0; idx < count; ++idx) {
+      auto tagp = response->mutable_tags(idx);
+      if (tagp->compare(policyTag) == 0) {        
+      }
+      else {
+        if (keep < idx) {
+          response->mutable_tags()->SwapElements(idx, keep);
+        }
+        ++keep;
+      }
+    }
+    response->mutable_tags()->DeleteSubrange(keep, count - keep);
+  }  
+#endif
+}
+
 std::string RecProtoBufMessage::getAppliedPolicy() const
 {
   std::string result;

diff --git a/pdns/rec-protobuf.hh b/pdns/rec-protobuf.hh
index 34386283d60ebdbed631de1194a3eb0796bc4532..e4f76b843cb7d3e0750fe44fcceb347341557c60 100644 (file)
--- a/pdns/rec-protobuf.hh
+++ b/pdns/rec-protobuf.hh
@@ -53,6 +53,8 @@ public:
   void setAppliedPolicy(const std::string& policy);
   void setAppliedPolicyType(const DNSFilterEngine::PolicyType& policyType);
   void setPolicyTags(const std::vector<std::string>& policyTags);
+  void addPolicyTag(const std::string& policyTag);
+  void removePolicyTag(const std::string& policyTag);
   std::string getAppliedPolicy() const;
   std::vector<std::string> getPolicyTags() const;
 };