rec: Allow setting the protobuf `requestorId` from Lua hooks

diff --git a/docs/markdown/recursor/scripting.md b/docs/markdown/recursor/scripting.md
index dfd7a5747961686a6c3ea04855474c8c1afadf66..087e9db2a35d88592bf9e365a4ab4f3cb3511e93 100644 (file)
--- a/docs/markdown/recursor/scripting.md
+++ b/docs/markdown/recursor/scripting.md
@@ -100,6 +100,7 @@ The DNSQuestion object contains at least the following fields:
      * policyTTL: The TTL in seconds for the `pdns.policyactions.Custom` response
 * wantsRPZ - A boolean that indicates the use of the Policy Engine, can be set to `false` in `prerpz` to disable RPZ for this query
 * data - a Lua object reference that is persistent throughout the lifetime of the `dq` object for a single query. It can be used to store custom data. Most scripts initialise this to an empty table early on so they can store multiple items.
+* requestorId - a string that will be used to set the `requestorId` field in protobuf messages (introduced in 4.1).
 
 It also supports the following methods:
 
@@ -162,7 +163,8 @@ In addition to this integer, this function can return a table of policy tags.
 The resulting tag number can be accessed via `dq.tag` in the `preresolve` hook,
 and the policy tags via `dq:getPolicyTags()` in every hook.
 Starting with 4.1.0, it can also return a table whose keys and values are strings
-to fill the upcoming `DNSQuestion`'s `data` table.
+to fill the upcoming `DNSQuestion`'s `data` table, as well as a `requestorId`
+value to fill the upcoming `DNSQuestion`'s `requestorId` field.
 
 The tagged packetcache can e.g. be used to answer queries from cache that have
 e.g. been filtered for certain IPs (this logic should be implemented in the

diff --git a/pdns/lua-recursor4.cc b/pdns/lua-recursor4.cc
index bb78ae7935282e526eb6f127e54468badcec9172..469878e980a3b3fb1c6a443d71e4bec933e22d4e 100644 (file)
--- a/pdns/lua-recursor4.cc
+++ b/pdns/lua-recursor4.cc
@@ -343,6 +343,7 @@ RecursorLua4::RecursorLua4(const std::string& fname)
 
   d_lw->registerMember("rcode", &DNSQuestion::rcode);
   d_lw->registerMember("tag", &DNSQuestion::tag);
+  d_lw->registerMember("requestorId", &DNSQuestion::requestorId);
   d_lw->registerMember("followupFunction", &DNSQuestion::followupFunction);
   d_lw->registerMember("followupPrefix", &DNSQuestion::followupPrefix);
   d_lw->registerMember("followupName", &DNSQuestion::followupName);
@@ -587,7 +588,7 @@ bool RecursorLua4::ipfilter(const ComboAddress& remote, const ComboAddress& loca
   return false; // don't block
 }
 
-unsigned int RecursorLua4::gettag(const ComboAddress& remote, const Netmask& ednssubnet, const ComboAddress& local, const DNSName& qname, uint16_t qtype, std::vector<std::string>* policyTags, LuaContext::LuaObject& data, const std::map<uint16_t, EDNSOptionView>& ednsOptions, bool tcp)
+unsigned int RecursorLua4::gettag(const ComboAddress& remote, const Netmask& ednssubnet, const ComboAddress& local, const DNSName& qname, uint16_t qtype, std::vector<std::string>* policyTags, LuaContext::LuaObject& data, const std::map<uint16_t, EDNSOptionView>& ednsOptions, bool tcp, std::string& requestorId)
 {
   if(d_gettag) {
     auto ret = d_gettag(remote, ednssubnet, local, qname, qtype, ednsOptions, tcp);
@@ -604,6 +605,10 @@ unsigned int RecursorLua4::gettag(const ComboAddress& remote, const Netmask& edn
     if (dataret) {
       data = *dataret;
     }
+    const auto reqIdret = std::get<3>(ret);
+    if (reqIdret) {
+      requestorId = *reqIdret;
+    }
     return std::get<0>(ret);
   }
   return 0;

diff --git a/pdns/lua-recursor4.hh b/pdns/lua-recursor4.hh
index 227c1d16e4b18d43ba80587cda2b1a4d1478e4d9..adfeaedfcc31f02fad0c3ddf2e30586edb792662 100644 (file)
--- a/pdns/lua-recursor4.hh
+++ b/pdns/lua-recursor4.hh
@@ -70,6 +70,7 @@ public:
     DNSFilterEngine::Policy* appliedPolicy{nullptr};
     std::vector<std::string>* policyTags{nullptr};
     std::unordered_map<std::string,bool>* discardedPolicies{nullptr};
+    std::string requestorId;
     bool& variable;
     bool& wantsRPZ;
     unsigned int tag{0};
@@ -101,7 +102,7 @@ public:
     DNSName followupName;
   };
 
-  unsigned int gettag(const ComboAddress& remote, const Netmask& ednssubnet, const ComboAddress& local, const DNSName& qname, uint16_t qtype, std::vector<std::string>* policyTags, LuaContext::LuaObject& data, const std::map<uint16_t, EDNSOptionView>&, bool tcp);
+  unsigned int gettag(const ComboAddress& remote, const Netmask& ednssubnet, const ComboAddress& local, const DNSName& qname, uint16_t qtype, std::vector<std::string>* policyTags, LuaContext::LuaObject& data, const std::map<uint16_t, EDNSOptionView>&, bool tcp, std::string& requestorId);
 
   bool prerpz(DNSQuestion& dq, int& ret);
   bool preresolve(DNSQuestion& dq, int& ret);
@@ -121,7 +122,7 @@ public:
             d_postresolve);
   }
 
-  typedef std::function<std::tuple<unsigned int,boost::optional<std::unordered_map<int,string> >,boost::optional<LuaContext::LuaObject> >(ComboAddress, Netmask, ComboAddress, DNSName, uint16_t, const std::map<uint16_t, EDNSOptionView>&, bool)> gettag_t;
+  typedef std::function<std::tuple<unsigned int,boost::optional<std::unordered_map<int,string> >,boost::optional<LuaContext::LuaObject>,boost::optional<std::string> >(ComboAddress, Netmask, ComboAddress, DNSName, uint16_t, const std::map<uint16_t, EDNSOptionView>&, bool)> gettag_t;
   gettag_t d_gettag; // public so you can query if we have this hooked
 
 private:

diff --git a/pdns/pdns_recursor.cc b/pdns/pdns_recursor.cc
index e6d53b41fffe310571e8e60f3dc49ffa82d32573..53fe9f0c8f9da6a0670bca9df4e2ba8edfc0ba59 100644 (file)
--- a/pdns/pdns_recursor.cc
+++ b/pdns/pdns_recursor.cc
@@ -205,6 +205,7 @@ struct DNSComboWriter {
   ComboAddress d_remote, d_local;
 #ifdef HAVE_PROTOBUF
   boost::uuids::uuid d_uuid;
+  string d_requestorId;
 #endif
   EDNSSubnetOpts d_ednssubnet;
   bool d_ecsFound{false};
@@ -642,12 +643,13 @@ catch(...)
 }
 
 #ifdef HAVE_PROTOBUF
-static void protobufLogQuery(const std::shared_ptr<RemoteLogger>& logger, uint8_t maskV4, uint8_t maskV6, const boost::uuids::uuid& uniqueId, const ComboAddress& remote, const ComboAddress& local, const Netmask& ednssubnet, bool tcp, uint16_t id, size_t len, const DNSName& qname, uint16_t qtype, uint16_t qclass, const std::vector<std::string>& policyTags)
+static void protobufLogQuery(const std::shared_ptr<RemoteLogger>& logger, uint8_t maskV4, uint8_t maskV6, const boost::uuids::uuid& uniqueId, const ComboAddress& remote, const ComboAddress& local, const Netmask& ednssubnet, bool tcp, uint16_t id, size_t len, const DNSName& qname, uint16_t qtype, uint16_t qclass, const std::vector<std::string>& policyTags, const std::string& requestorId)
 {
   Netmask requestorNM(remote, remote.sin4.sin_family == AF_INET ? maskV4 : maskV6);
   const ComboAddress& requestor = requestorNM.getMaskedNetwork();
   RecProtoBufMessage message(DNSProtoBufMessage::Query, uniqueId, &requestor, &local, qname, qtype, qclass, id, tcp, len);
   message.setEDNSSubnet(ednssubnet, ednssubnet.isIpv4() ? maskV4 : maskV6);
+  message.setRequestorId(requestorId);
 
   if (!policyTags.empty()) {
     message.setPolicyTags(policyTags);
@@ -801,6 +803,9 @@ static void startDoResolve(void *p)
     dq.currentRecords = &ret;
     dq.dh = &dc->d_mdp.d_header;
     dq.data = dc->d_data;
+#ifdef HAVE_PROTOBUF
+    dq.requestorId = dc->d_requestorId;
+#endif
 
     if(dc->d_mdp.d_qtype==QType::ANY && !dc->d_tcp && g_anyToTcp) {
       pw.getHeader()->tc = 1;
@@ -1151,6 +1156,7 @@ static void startDoResolve(void *p)
       }
       pbMessage.setPolicyTags(dc->d_policyTags);
       pbMessage.setQueryTime(dc->d_now.tv_sec, dc->d_now.tv_usec);
+      pbMessage.setRequestorId(dq.requestorId);
       protobufLogResponse(luaconfsLocal->protobufServer, pbMessage);
     }
 #endif
@@ -1416,6 +1422,7 @@ static void handleRunningTCPQuestion(int fd, FDMultiplexer::funcparam_t& var)
       uint16_t qtype=0;
       uint16_t qclass=0;
       bool needECS = false;
+      string requestorId;
 #ifdef HAVE_PROTOBUF
       auto luaconfsLocal = g_luaconfs.getLocal();
       if (luaconfsLocal->protobufServer) {
@@ -1432,7 +1439,7 @@ static void handleRunningTCPQuestion(int fd, FDMultiplexer::funcparam_t& var)
 
           if(t_pdl->get() && (*t_pdl)->d_gettag) {
             try {
-              dc->d_tag = (*t_pdl)->gettag(conn->d_remote, dc->d_ednssubnet.source, dest, qname, qtype, &dc->d_policyTags, dc->d_data, ednsOptions, true);
+              dc->d_tag = (*t_pdl)->gettag(conn->d_remote, dc->d_ednssubnet.source, dest, qname, qtype, &dc->d_policyTags, dc->d_data, ednsOptions, true, requestorId);
             }
             catch(std::exception& e)  {
               if(g_logCommonErrors)
@@ -1448,6 +1455,7 @@ static void handleRunningTCPQuestion(int fd, FDMultiplexer::funcparam_t& var)
       }
 #ifdef HAVE_PROTOBUF
       if(luaconfsLocal->protobufServer || luaconfsLocal->outgoingProtobufServer) {
+        dc->d_requestorId = requestorId;
         dc->d_uuid = (*t_uuidGenerator)();
       }
 
@@ -1456,7 +1464,7 @@ static void handleRunningTCPQuestion(int fd, FDMultiplexer::funcparam_t& var)
           const struct dnsheader* dh = (const struct dnsheader*) conn->data;
 
           if (!luaconfsLocal->protobufTaggedOnly) {
-            protobufLogQuery(luaconfsLocal->protobufServer, luaconfsLocal->protobufMaskV4, luaconfsLocal->protobufMaskV6, dc->d_uuid, conn->d_remote, dest, dc->d_ednssubnet.source, true, dh->id, conn->qlen, qname, qtype, qclass, dc->d_policyTags);
+            protobufLogQuery(luaconfsLocal->protobufServer, luaconfsLocal->protobufMaskV4, luaconfsLocal->protobufMaskV6, dc->d_uuid, conn->d_remote, dest, dc->d_ednssubnet.source, true, dh->id, conn->qlen, qname, qtype, qclass, dc->d_policyTags, dc->d_requestorId);
           }
         }
         catch(std::exception& e) {
@@ -1565,6 +1573,7 @@ static string* doProcessUDPQuestion(const std::string& question, const ComboAddr
   bool needECS = false;
   std::vector<std::string> policyTags;
   LuaContext::LuaObject data;
+  string requestorId;
 #ifdef HAVE_PROTOBUF
   boost::uuids::uuid uniqueId;
   auto luaconfsLocal = g_luaconfs.getLocal();
@@ -1605,7 +1614,7 @@ static string* doProcessUDPQuestion(const std::string& question, const ComboAddr
 
         if(t_pdl->get() && (*t_pdl)->d_gettag) {
           try {
-            ctag=(*t_pdl)->gettag(fromaddr, ednssubnet.source, destaddr, qname, qtype, &policyTags, data, ednsOptions, false);
+            ctag=(*t_pdl)->gettag(fromaddr, ednssubnet.source, destaddr, qname, qtype, &policyTags, data, ednsOptions, false, requestorId);
           }
           catch(std::exception& e)  {
             if(g_logCommonErrors)
@@ -1625,7 +1634,7 @@ static string* doProcessUDPQuestion(const std::string& question, const ComboAddr
 #ifdef HAVE_PROTOBUF
     if(luaconfsLocal->protobufServer) {
       if (!luaconfsLocal->protobufTaggedOnly || !policyTags.empty()) {
-        protobufLogQuery(luaconfsLocal->protobufServer, luaconfsLocal->protobufMaskV4, luaconfsLocal->protobufMaskV6, uniqueId, fromaddr, destaddr, ednssubnet.source, false, dh->id, question.size(), qname, qtype, qclass, policyTags);
+        protobufLogQuery(luaconfsLocal->protobufServer, luaconfsLocal->protobufMaskV4, luaconfsLocal->protobufMaskV6, uniqueId, fromaddr, destaddr, ednssubnet.source, false, dh->id, question.size(), qname, qtype, qclass, policyTags, requestorId);
       }
     }
 #endif /* HAVE_PROTOBUF */
@@ -1645,6 +1654,7 @@ static string* doProcessUDPQuestion(const std::string& question, const ComboAddr
         pbMessage.update(uniqueId, &requestor, &destaddr, false, dh->id);
         pbMessage.setEDNSSubnet(ednssubnet.source, ednssubnet.source.isIpv4() ? luaconfsLocal->protobufMaskV4 : luaconfsLocal->protobufMaskV6);
         pbMessage.setQueryTime(g_now.tv_sec, g_now.tv_usec);
+        pbMessage.setRequestorId(requestorId);
         protobufLogResponse(luaconfsLocal->protobufServer, pbMessage);
       }
 #endif /* HAVE_PROTOBUF */
@@ -1714,6 +1724,7 @@ static string* doProcessUDPQuestion(const std::string& question, const ComboAddr
   if (luaconfsLocal->protobufServer || luaconfsLocal->outgoingProtobufServer) {
     dc->d_uuid = uniqueId;
   }
+  dc->d_requestorId = requestorId;
 #endif
 
   MT->makeThread(startDoResolve, (void*) dc); // deletes dc

diff --git a/pdns/protobuf.cc b/pdns/protobuf.cc
index 3d4e359f32eebe49cd17edc0039ff11ac598fc76..c881e0084b1e2edbe20e9dd8e08c22be53de465c 100644 (file)
--- a/pdns/protobuf.cc
+++ b/pdns/protobuf.cc
@@ -196,6 +196,13 @@ void DNSProtoBufMessage::setRequestor(const ComboAddress& requestor)
 #endif /* HAVE_PROTOBUF */
 }
 
+void DNSProtoBufMessage::setRequestorId(const std::string& requestorId)
+{
+#ifdef HAVE_PROTOBUF
+  d_message.set_requestorid(requestorId);
+#endif /* HAVE_PROTOBUF */
+}
+
 void DNSProtoBufMessage::setResponder(const std::string& responder)
 {
 #ifdef HAVE_PROTOBUF

diff --git a/pdns/protobuf.hh b/pdns/protobuf.hh
index 9834a99ced877e26b7882c6aeab027b35c28832b..776bc10a7fffd9b5813ebcf0ba719574748f7f99 100644 (file)
--- a/pdns/protobuf.hh
+++ b/pdns/protobuf.hh
@@ -68,6 +68,7 @@ public:
   void setRequestor(const ComboAddress& requestor);
   void setResponder(const std::string& responder);
   void setResponder(const ComboAddress& responder);
+  void setRequestorId(const std::string& requestorId);
   std::string toDebugString() const;
 
 #ifdef HAVE_PROTOBUF