Merge pull request #7015 from rgacogne/dnsdist-protobuf-serverid

[thirdparty/pdns.git] / pdns / dnsdist-lua-actions.cc

diff --git a/pdns/dnsdist-lua-actions.cc b/pdns/dnsdist-lua-actions.cc
index 7d7a07f2a52d9dcc2be2f76d7c8c3af6edeb8256..f59ec3353fb240f6642a5b9e3df190fafee016cf 100644 (file)
--- a/pdns/dnsdist-lua-actions.cc
+++ b/pdns/dnsdist-lua-actions.cc
@@ -751,7 +751,7 @@ private:
 class RemoteLogAction : public DNSAction, public boost::noncopyable
 {
 public:
-  RemoteLogAction(std::shared_ptr<RemoteLoggerInterface>& logger, boost::optional<std::function<void(const DNSQuestion&, DNSDistProtoBufMessage*)> > alterFunc): d_logger(logger), d_alterFunc(alterFunc)
+  RemoteLogAction(std::shared_ptr<RemoteLoggerInterface>& logger, boost::optional<std::function<void(const DNSQuestion&, DNSDistProtoBufMessage*)> > alterFunc, const std::string& serverID): d_logger(logger), d_alterFunc(alterFunc), d_serverID(serverID)
   {
   }
   DNSAction::Action operator()(DNSQuestion* dq, string* ruleresult) const override
@@ -762,12 +762,15 @@ public:
     }
 
     DNSDistProtoBufMessage message(*dq);
-    {
-      if (d_alterFunc) {
-        std::lock_guard<std::mutex> lock(g_luamutex);
-        (*d_alterFunc)(*dq, &message);
-      }
+    if (!d_serverID.empty()) {
+      message.setServerIdentity(d_serverID);
     }
+
+    if (d_alterFunc) {
+      std::lock_guard<std::mutex> lock(g_luamutex);
+      (*d_alterFunc)(*dq, &message);
+    }
+
     std::string data;
     message.serialize(data);
     d_logger->queueData(data);
@@ -781,6 +784,7 @@ public:
 private:
   std::shared_ptr<RemoteLoggerInterface> d_logger;
   boost::optional<std::function<void(const DNSQuestion&, DNSDistProtoBufMessage*)> > d_alterFunc;
+  std::string d_serverID;
 };
 
 class SNMPTrapAction : public DNSAction
@@ -867,7 +871,7 @@ private:
 class RemoteLogResponseAction : public DNSResponseAction, public boost::noncopyable
 {
 public:
-  RemoteLogResponseAction(std::shared_ptr<RemoteLoggerInterface>& logger, boost::optional<std::function<void(const DNSResponse&, DNSDistProtoBufMessage*)> > alterFunc, bool includeCNAME): d_logger(logger), d_alterFunc(alterFunc), d_includeCNAME(includeCNAME)
+  RemoteLogResponseAction(std::shared_ptr<RemoteLoggerInterface>& logger, boost::optional<std::function<void(const DNSResponse&, DNSDistProtoBufMessage*)> > alterFunc, const std::string& serverID, bool includeCNAME): d_logger(logger), d_alterFunc(alterFunc), d_serverID(serverID), d_includeCNAME(includeCNAME)
   {
   }
   DNSResponseAction::Action operator()(DNSResponse* dr, string* ruleresult) const override
@@ -878,12 +882,15 @@ public:
     }
 
     DNSDistProtoBufMessage message(*dr, d_includeCNAME);
-    {
-      if (d_alterFunc) {
-        std::lock_guard<std::mutex> lock(g_luamutex);
-        (*d_alterFunc)(*dr, &message);
-      }
+    if (!d_serverID.empty()) {
+      message.setServerIdentity(d_serverID);
+    }
+
+    if (d_alterFunc) {
+      std::lock_guard<std::mutex> lock(g_luamutex);
+      (*d_alterFunc)(*dr, &message);
     }
+
     std::string data;
     message.serialize(data);
     d_logger->queueData(data);
@@ -897,6 +904,7 @@ public:
 private:
   std::shared_ptr<RemoteLoggerInterface> d_logger;
   boost::optional<std::function<void(const DNSResponse&, DNSDistProtoBufMessage*)> > d_alterFunc;
+  std::string d_serverID;
   bool d_includeCNAME;
 };
 
@@ -1178,29 +1186,45 @@ void setupLuaActions()
       return std::shared_ptr<DNSResponseAction>(new LuaResponseAction(func));
     });
 
-  g_lua.writeFunction("RemoteLogAction", [](std::shared_ptr<RemoteLoggerInterface> logger, boost::optional<std::function<void(const DNSQuestion&, DNSDistProtoBufMessage*)> > alterFunc) {
+  g_lua.writeFunction("RemoteLogAction", [](std::shared_ptr<RemoteLoggerInterface> logger, boost::optional<std::function<void(const DNSQuestion&, DNSDistProtoBufMessage*)> > alterFunc, boost::optional<std::unordered_map<std::string, std::string>> vars) {
       // avoids potentially-evaluated-expression warning with clang.
       RemoteLoggerInterface& rl = *logger.get();
       if (typeid(rl) != typeid(RemoteLogger)) {
         // We could let the user do what he wants, but wrapping PowerDNS Protobuf inside a FrameStream tagged as dnstap is logically wrong.
         throw std::runtime_error(std::string("RemoteLogAction only takes RemoteLogger. For other types, please look at DnstapLogAction."));
       }
+
+      std::string serverID;
+      if (vars) {
+        if (vars->count("serverID")) {
+          serverID = boost::get<std::string>((*vars)["serverID"]);
+        }
+      }
+
 #ifdef HAVE_PROTOBUF
-      return std::shared_ptr<DNSAction>(new RemoteLogAction(logger, alterFunc));
+      return std::shared_ptr<DNSAction>(new RemoteLogAction(logger, alterFunc, serverID));
 #else
       throw std::runtime_error("Protobuf support is required to use RemoteLogAction");
 #endif
     });
 
-  g_lua.writeFunction("RemoteLogResponseAction", [](std::shared_ptr<RemoteLoggerInterface> logger, boost::optional<std::function<void(const DNSResponse&, DNSDistProtoBufMessage*)> > alterFunc, boost::optional<bool> includeCNAME) {
+  g_lua.writeFunction("RemoteLogResponseAction", [](std::shared_ptr<RemoteLoggerInterface> logger, boost::optional<std::function<void(const DNSResponse&, DNSDistProtoBufMessage*)> > alterFunc, boost::optional<bool> includeCNAME, boost::optional<std::unordered_map<std::string, std::string>> vars) {
       // avoids potentially-evaluated-expression warning with clang.
       RemoteLoggerInterface& rl = *logger.get();
       if (typeid(rl) != typeid(RemoteLogger)) {
         // We could let the user do what he wants, but wrapping PowerDNS Protobuf inside a FrameStream tagged as dnstap is logically wrong.
         throw std::runtime_error("RemoteLogResponseAction only takes RemoteLogger. For other types, please look at DnstapLogResponseAction.");
       }
+
+      std::string serverID;
+      if (vars) {
+        if (vars->count("serverID")) {
+          serverID = boost::get<std::string>((*vars)["serverID"]);
+        }
+      }
+
 #ifdef HAVE_PROTOBUF
-      return std::shared_ptr<DNSResponseAction>(new RemoteLogResponseAction(logger, alterFunc, includeCNAME ? *includeCNAME : false));
+      return std::shared_ptr<DNSResponseAction>(new RemoteLogResponseAction(logger, alterFunc, serverID, includeCNAME ? *includeCNAME : false));
 #else
       throw std::runtime_error("Protobuf support is required to use RemoteLogResponseAction");
 #endif