dnsdist: Support exporting tags in the Protocol Buffer 'tags' field

diff --git a/pdns/dnsdist-lua-actions.cc b/pdns/dnsdist-lua-actions.cc
index c4d8e9c87ab6fe5148660ab91b7047654009f8a9..e43d88d2a781f39a5dc9eac27faaee16085c7bc4 100644 (file)
--- a/pdns/dnsdist-lua-actions.cc
+++ b/pdns/dnsdist-lua-actions.cc
@@ -156,7 +156,7 @@ private:
   bool d_addECS{false};
 };
 
-TeeAction::TeeAction(const ComboAddress& rca, const boost::optional<ComboAddress>& lca, bool addECS) 
+TeeAction::TeeAction(const ComboAddress& rca, const boost::optional<ComboAddress>& lca, bool addECS)
   : d_remote(rca), d_addECS(addECS)
 {
   d_fd=SSocket(d_remote.sin4.sin_family, SOCK_DGRAM, 0);
@@ -1489,11 +1489,31 @@ static void addMetaDataToProtobuf(DNSDistProtoBufMessage& message, const DNSQues
   }
 }
 
+static void addTagsToProtobuf(DNSDistProtoBufMessage& message, const DNSQuestion& dq, const std::unordered_set<std::string>& allowed)
+{
+  if (!dq.ids.qTag) {
+    return;
+  }
+
+  for (const auto& [key, value] : *dq.ids.qTag) {
+    if (!allowed.empty() && allowed.count(key) == 0) {
+      continue;
+    }
+
+    if (value.empty()) {
+      message.addTag(key);
+    }
+    else {
+      message.addTag(key + ":" + value);
+    }
+  }
+}
+
 class RemoteLogAction : public DNSAction, public boost::noncopyable
 {
 public:
   // this action does not stop the processing
-  RemoteLogAction(std::shared_ptr<RemoteLoggerInterface>& logger, boost::optional<std::function<void(DNSQuestion*, DNSDistProtoBufMessage*)> > alterFunc, const std::string& serverID, const std::string& ipEncryptKey, std::vector<std::pair<std::string, ProtoBufMetaKey>>& metas): d_metas(std::move(metas)), d_logger(logger), d_alterFunc(alterFunc), d_serverID(serverID), d_ipEncryptKey(ipEncryptKey)
+  RemoteLogAction(std::shared_ptr<RemoteLoggerInterface>& logger, boost::optional<std::function<void(DNSQuestion*, DNSDistProtoBufMessage*)> > alterFunc, const std::string& serverID, const std::string& ipEncryptKey, std::vector<std::pair<std::string, ProtoBufMetaKey>>&& metas, std::optional<std::unordered_set<std::string>>&& tagsToExport): d_tagsToExport(std::move(tagsToExport)), d_metas(std::move(metas)), d_logger(logger), d_alterFunc(alterFunc), d_serverID(serverID), d_ipEncryptKey(ipEncryptKey)
   {
   }
 
@@ -1515,6 +1535,10 @@ public:
     }
 #endif /* HAVE_IPCIPHER */
 
+    if (d_tagsToExport) {
+      addTagsToProtobuf(message, *dq, *d_tagsToExport);
+    }
+
     addMetaDataToProtobuf(message, *dq, d_metas);
 
     if (d_alterFunc) {
@@ -1534,6 +1558,7 @@ public:
     return "remote log to " + (d_logger ? d_logger->toString() : "");
   }
 private:
+  std::optional<std::unordered_set<std::string>> d_tagsToExport;
   std::vector<std::pair<std::string, ProtoBufMetaKey>> d_metas;
   std::shared_ptr<RemoteLoggerInterface> d_logger;
   boost::optional<std::function<void(DNSQuestion*, DNSDistProtoBufMessage*)> > d_alterFunc;
@@ -1630,7 +1655,7 @@ class RemoteLogResponseAction : public DNSResponseAction, public boost::noncopya
 {
 public:
   // this action does not stop the processing
-  RemoteLogResponseAction(std::shared_ptr<RemoteLoggerInterface>& logger, boost::optional<std::function<void(DNSResponse*, DNSDistProtoBufMessage*)> > alterFunc, const std::string& serverID, const std::string& ipEncryptKey, bool includeCNAME, std::vector<std::pair<std::string, ProtoBufMetaKey>>& metas): d_metas(std::move(metas)), d_logger(logger), d_alterFunc(alterFunc), d_serverID(serverID), d_ipEncryptKey(ipEncryptKey), d_includeCNAME(includeCNAME)
+  RemoteLogResponseAction(std::shared_ptr<RemoteLoggerInterface>& logger, boost::optional<std::function<void(DNSResponse*, DNSDistProtoBufMessage*)> > alterFunc, const std::string& serverID, const std::string& ipEncryptKey, bool includeCNAME, std::vector<std::pair<std::string, ProtoBufMetaKey>>&& metas, std::optional<std::unordered_set<std::string>>&& tagsToExport): d_tagsToExport(std::move(tagsToExport)), d_metas(std::move(metas)), d_logger(logger), d_alterFunc(alterFunc), d_serverID(serverID), d_ipEncryptKey(ipEncryptKey), d_includeCNAME(includeCNAME)
   {
   }
   DNSResponseAction::Action operator()(DNSResponse* dr, std::string* ruleresult) const override
@@ -1651,6 +1676,10 @@ public:
     }
 #endif /* HAVE_IPCIPHER */
 
+    if (d_tagsToExport) {
+      addTagsToProtobuf(message, *dr, *d_tagsToExport);
+    }
+
     addMetaDataToProtobuf(message, *dr, d_metas);
 
     if (d_alterFunc) {
@@ -1670,6 +1699,7 @@ public:
     return "remote log response to " + (d_logger ? d_logger->toString() : "");
   }
 private:
+  std::optional<std::unordered_set<std::string>> d_tagsToExport;
   std::vector<std::pair<std::string, ProtoBufMetaKey>> d_metas;
   std::shared_ptr<RemoteLoggerInterface> d_logger;
   boost::optional<std::function<void(DNSResponse*, DNSDistProtoBufMessage*)> > d_alterFunc;
@@ -2474,8 +2504,10 @@ void setupLuaActions(LuaContext& luaCtx)
 
       std::string serverID;
       std::string ipEncryptKey;
+      std::string tags;
       getOptionalValue<std::string>(vars, "serverID", serverID);
       getOptionalValue<std::string>(vars, "ipEncryptKey", ipEncryptKey);
+      getOptionalValue<std::string>(vars, "exportTags", tags);
 
       std::vector<std::pair<std::string, ProtoBufMetaKey>> metaOptions;
       if (metas) {
@@ -2484,9 +2516,21 @@ void setupLuaActions(LuaContext& luaCtx)
         }
       }
 
+      std::optional<std::unordered_set<std::string>> tagsToExport{std::nullopt};
+      if (!tags.empty()) {
+        tagsToExport = std::unordered_set<std::string>();
+        if (tags != "*") {
+          std::vector<std::string> tokens;
+          stringtok(tokens, tags, ",");
+          for (auto& token : tokens) {
+            tagsToExport->insert(std::move(token));
+          }
+        }
+      }
+
       checkAllParametersConsumed("RemoteLogAction", vars);
 
-      return std::shared_ptr<DNSAction>(new RemoteLogAction(logger, alterFunc, serverID, ipEncryptKey, metaOptions));
+      return std::shared_ptr<DNSAction>(new RemoteLogAction(logger, alterFunc, serverID, ipEncryptKey, std::move(metaOptions), std::move(tagsToExport)));
     });
 
   luaCtx.writeFunction("RemoteLogResponseAction", [](std::shared_ptr<RemoteLoggerInterface> logger, boost::optional<std::function<void(DNSResponse*, DNSDistProtoBufMessage*)> > alterFunc, boost::optional<bool> includeCNAME, boost::optional<LuaAssociativeTable<std::string>> vars, boost::optional<LuaAssociativeTable<std::string>> metas) {
@@ -2501,8 +2545,10 @@ void setupLuaActions(LuaContext& luaCtx)
 
       std::string serverID;
       std::string ipEncryptKey;
+      std::string tags;
       getOptionalValue<std::string>(vars, "serverID", serverID);
       getOptionalValue<std::string>(vars, "ipEncryptKey", ipEncryptKey);
+      getOptionalValue<std::string>(vars, "exportTags", tags);
 
       std::vector<std::pair<std::string, ProtoBufMetaKey>> metaOptions;
       if (metas) {
@@ -2511,9 +2557,21 @@ void setupLuaActions(LuaContext& luaCtx)
         }
       }
 
+      std::optional<std::unordered_set<std::string>> tagsToExport{std::nullopt};
+      if (!tags.empty()) {
+        tagsToExport = std::unordered_set<std::string>();
+        if (tags != "*") {
+          std::vector<std::string> tokens;
+          stringtok(tokens, tags, ",");
+          for (auto& token : tokens) {
+            tagsToExport->insert(std::move(token));
+          }
+        }
+      }
+
       checkAllParametersConsumed("RemoteLogResponseAction", vars);
 
-      return std::shared_ptr<DNSResponseAction>(new RemoteLogResponseAction(logger, alterFunc, serverID, ipEncryptKey, includeCNAME ? *includeCNAME : false, metaOptions));
+      return std::shared_ptr<DNSResponseAction>(new RemoteLogResponseAction(logger, alterFunc, serverID, ipEncryptKey, includeCNAME ? *includeCNAME : false, std::move(metaOptions), std::move(tagsToExport)));
     });
 
   luaCtx.writeFunction("DnstapLogAction", [](const std::string& identity, std::shared_ptr<RemoteLoggerInterface> logger, boost::optional<std::function<void(DNSQuestion*, DnstapMessage*)> > alterFunc) {

diff --git a/pdns/dnsdistdist/docs/rules-actions.rst b/pdns/dnsdistdist/docs/rules-actions.rst
index 8db6d1cb36cdd12ebb7acd554bd6e8e977403358..7cb76469342606b9c9930707e290e5d47351b298 100644 (file)
--- a/pdns/dnsdistdist/docs/rules-actions.rst
+++ b/pdns/dnsdistdist/docs/rules-actions.rst
@@ -1323,6 +1323,7 @@ The following actions exist.
 
   .. versionchanged:: 1.8.0
     ``metas`` optional parameter added.
+    ``exportTags`` optional key added to the options table.
 
   Send the content of this query to a remote logger via Protocol Buffer.
   ``alterFunction`` is a callback, receiving a :class:`DNSQuestion` and a :class:`DNSDistProtoBufMessage`, that can be used to modify the Protocol Buffer content, for example for anonymization purposes.
@@ -1353,6 +1354,7 @@ The following actions exist.
 
   * ``serverID=""``: str - Set the Server Identity field.
   * ``ipEncryptKey=""``: str - A key, that can be generated via the :func:`makeIPCipherKey` function, to encrypt the IP address of the requestor for anonymization purposes. The encryption is done using ipcrypt for IPv4 and a 128-bit AES ECB operation for IPv6.
+  * ``exportTags=""``: str - The comma-separated list of keys of internal tags to export into the ``tags`` Protocol Buffer field, as "key:value" strings. Note that a tag with an empty value will be exported as "<key>", not "<key>:". An empty string means that no internal tag will be exported. The special value ``*`` means that all tags will be exported.
 
 .. function:: RemoteLogResponseAction(remoteLogger[, alterFunction[, includeCNAME [, options [, metas]]]])
 
@@ -1361,6 +1363,7 @@ The following actions exist.
 
   .. versionchanged:: 1.8.0
     ``metas`` optional parameter added.
+    ``exportTags`` optional key added to the options table.
 
   Send the content of this response to a remote logger via Protocol Buffer.
   ``alterFunction`` is the same callback that receiving a :class:`DNSQuestion` and a :class:`DNSDistProtoBufMessage`, that can be used to modify the Protocol Buffer content, for example for anonymization purposes.
@@ -1379,6 +1382,7 @@ The following actions exist.
 
   * ``serverID=""``: str - Set the Server Identity field.
   * ``ipEncryptKey=""``: str - A key, that can be generated via the :func:`makeIPCipherKey` function, to encrypt the IP address of the requestor for anonymization purposes. The encryption is done using ipcrypt for IPv4 and a 128-bit AES ECB operation for IPv6.
+  * ``exportTags=""``: str - The comma-separated list of keys of internal tags to export into the ``tags`` Protocol Buffer field, as "key:value" strings. Note that a tag with an empty value will be exported as "<key>", not "<key>:". An empty string means that no internal tag will be exported. The special value ``*`` means that all tags will be exported.
 
 .. function:: SetAdditionalProxyProtocolValueAction(type, value)
 

diff --git a/regression-tests.dnsdist/test_Protobuf.py b/regression-tests.dnsdist/test_Protobuf.py
index 125788915b3651cc0574c939797b89da40e809b4..3bbcb794668d9b9f7e91615ec1663ba105d4107e 100644 (file)
--- a/regression-tests.dnsdist/test_Protobuf.py
+++ b/regression-tests.dnsdist/test_Protobuf.py
@@ -419,9 +419,9 @@ class TestProtobufMetaTags(DNSDistProtobufTest):
 
     addAction(AllRule(), SetTagAction('my-tag-key', 'my-tag-value'))
     addAction(AllRule(), SetTagAction('my-empty-key', ''))
-    addAction(AllRule(), RemoteLogAction(rl, nil, {serverID='dnsdist-server-1'}, {b64='b64-content', ['my-tag-export-name']='tag:my-tag-key'}))
+    addAction(AllRule(), RemoteLogAction(rl, nil, {serverID='dnsdist-server-1', exportTags='*'}, {b64='b64-content', ['my-tag-export-name']='tag:my-tag-key'}))
     addResponseAction(AllRule(), SetTagResponseAction('my-tag-key2', 'my-tag-value2'))
-    addResponseAction(AllRule(), RemoteLogResponseAction(rl, nil, false, {serverID='dnsdist-server-1'}, {['my-tag-export-name']='tags'}))
+    addResponseAction(AllRule(), RemoteLogResponseAction(rl, nil, false, {serverID='dnsdist-server-1', exportTags='my-empty-key,my-tag-key2'}, {['my-tag-export-name']='tags'}))
     """
 
     def testProtobufMeta(self):
@@ -452,6 +452,11 @@ class TestProtobufMetaTags(DNSDistProtobufTest):
         msg = self.getFirstProtobufMessage()
 
         self.checkProtobufQuery(msg, dnsmessage_pb2.PBDNSMessage.UDP, query, dns.rdataclass.IN, dns.rdatatype.A, name)
+        # regular tags
+        self.assertEqual(len(msg.response.tags), 2)
+        self.assertIn('my-tag-key:my-tag-value', msg.response.tags)
+        self.assertIn('my-empty-key', msg.response.tags)
+        # meta tags
         self.assertEqual(len(msg.meta), 2)
         tags = {}
         for entry in msg.meta:
@@ -467,6 +472,11 @@ class TestProtobufMetaTags(DNSDistProtobufTest):
         # check the protobuf message corresponding to the UDP response
         msg = self.getFirstProtobufMessage()
         self.checkProtobufResponse(msg, dnsmessage_pb2.PBDNSMessage.UDP, response)
+        # regular tags
+        self.assertEqual(len(msg.response.tags), 2)
+        self.assertIn('my-tag-key2:my-tag-value2', msg.response.tags)
+        self.assertIn('my-empty-key', msg.response.tags)
+        # meta tags
         self.assertEqual(len(msg.meta), 1)
         self.assertEqual(msg.meta[0].key, 'my-tag-export-name')
         self.assertEqual(len(msg.meta[0].value.stringVal), 3)