rec: Export the server ID in protobuf messages

[thirdparty/pdns.git] / pdns / pdns_recursor.cc

diff --git a/pdns/pdns_recursor.cc b/pdns/pdns_recursor.cc
index 4c768fe84c120ef320376fd117f050f7a851abb3..2413e30c22a10179f8632ba5d7656ca94dca5a95 100644 (file)
--- a/pdns/pdns_recursor.cc
+++ b/pdns/pdns_recursor.cc
@@ -146,6 +146,7 @@ static AtomicCounter counter;
 static std::shared_ptr<SyncRes::domainmap_t> g_initialDomainMap; // new threads needs this to be setup
 static std::shared_ptr<NetmaskGroup> g_initialAllowFrom; // new thread needs to be setup with this
 static size_t g_tcpMaxQueriesPerConn;
+static size_t s_maxUDPQueriesPerRound;
 static uint64_t g_latencyStatSize;
 static uint32_t g_disthashseed;
 static unsigned int g_maxTCPPerClient;
@@ -245,9 +246,9 @@ ArgvMap &arg()
   return theArg;
 }
 
-unsigned int getRecursorThreadId()
+int getRecursorThreadId()
 {
-  return static_cast<unsigned int>(t_id);
+  return t_id;
 }
 
 int getMTaskerTID()
@@ -665,6 +666,7 @@ static void protobufLogQuery(const std::shared_ptr<RemoteLogger>& logger, uint8_
   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.setServerIdentity(SyncRes::s_serverID);
   message.setEDNSSubnet(ednssubnet, ednssubnet.isIpv4() ? maskV4 : maskV6);
   message.setRequestorId(requestorId);
   message.setDeviceId(deviceId);
@@ -776,6 +778,7 @@ static void startDoResolve(void *p)
     if (luaconfsLocal->protobufServer) {
       Netmask requestorNM(dc->d_remote, dc->d_remote.sin4.sin_family == AF_INET ? luaconfsLocal->protobufMaskV4 : luaconfsLocal->protobufMaskV6);
       const ComboAddress& requestor = requestorNM.getMaskedNetwork();
+      pbMessage.setServerIdentity(SyncRes::s_serverID);
       pbMessage.update(dc->d_uuid, &requestor, &dc->d_local, dc->d_tcp, dc->d_mdp.d_header.id);
       pbMessage.setEDNSSubnet(dc->d_ednssubnet.source, dc->d_ednssubnet.source.isIpv4() ? luaconfsLocal->protobufMaskV4 : luaconfsLocal->protobufMaskV6);
       pbMessage.setQuestion(dc->d_mdp.d_qname, dc->d_mdp.d_qtype, dc->d_mdp.d_qclass);
@@ -1343,6 +1346,7 @@ static void startDoResolve(void *p)
   }
   catch(...) {
     L<<Logger::Error<<"Any other exception in a resolver context "<< makeLoginfo(dc) <<endl;
+    delete dc;
   }
 
   g_stats.maxMThreadStackUsage = max(MT->getMaxStackUsage(), g_stats.maxMThreadStackUsage);
@@ -1715,6 +1719,7 @@ static string* doProcessUDPQuestion(const std::string& question, const ComboAddr
     bool cacheHit = false;
     RecProtoBufMessage pbMessage(DNSProtoBufMessage::DNSProtoBufMessageType::Response);
 #ifdef HAVE_PROTOBUF
+    pbMessage.setServerIdentity(SyncRes::s_serverID);
     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, requestorId, deviceId);
@@ -1836,7 +1841,7 @@ static void handleNewUDPQuestion(int fd, FDMultiplexer::funcparam_t& var)
   fromaddr.sin6.sin6_family=AF_INET6; // this makes sure fromaddr is big enough
   fillMSGHdr(&msgh, &iov, cbuf, sizeof(cbuf), data, sizeof(data), &fromaddr);
 
-  for(;;)
+  for(size_t counter = 0; counter < s_maxUDPQueriesPerRound; counter++)
   if((len=recvmsg(fd, &msgh, 0)) >= 0) {
 
     firstQuery = false;
@@ -2959,9 +2964,10 @@ static int serviceMain(int argc, char*argv[])
 
   g_maxCacheEntries = ::arg().asNum("max-cache-entries");
   g_maxPacketCacheEntries = ::arg().asNum("max-packetcache-entries");
-  
+
+  luaConfigDelayedThreads delayedLuaThreads;
   try {
-    loadRecursorLuaConfig(::arg()["lua-config-file"], ::arg().mustDo("daemon"));
+    loadRecursorLuaConfig(::arg()["lua-config-file"], delayedLuaThreads);
   }
   catch (PDNSException &e) {
     L<<Logger::Error<<"Cannot load Lua configuration: "<<e.reason<<endl;
@@ -3125,7 +3131,6 @@ static int serviceMain(int argc, char*argv[])
     L<<Logger::Warning<<"Calling daemonize, going to background"<<endl;
     L.toConsole(Logger::Critical);
     daemonize();
-    loadRecursorLuaConfig(::arg()["lua-config-file"], false);
   }
   signal(SIGUSR1,usr1Handler);
   signal(SIGUSR2,usr2Handler);
@@ -3178,11 +3183,14 @@ static int serviceMain(int argc, char*argv[])
 
   Utility::dropUserPrivs(newuid);
 
+  startLuaConfigDelayedThreads(delayedLuaThreads);
+
   makeThreadPipes();
 
   g_tcpTimeout=::arg().asNum("client-tcp-timeout");
   g_maxTCPPerClient=::arg().asNum("max-tcp-per-client");
   g_tcpMaxQueriesPerConn=::arg().asNum("max-tcp-queries-per-connection");
+  s_maxUDPQueriesPerRound=::arg().asNum("max-udp-queries-per-round");
 
   if (::arg().mustDo("snmp-agent")) {
     g_snmpAgent = std::make_shared<RecursorSNMPAgent>("recursor", ::arg()["snmp-master-socket"]);
@@ -3236,15 +3244,17 @@ try
 #endif
   L<<Logger::Warning<<"Done priming cache with root hints"<<endl;
 
-  try {
-    if(!::arg()["lua-dns-script"].empty()) {
-      t_pdl = std::make_shared<RecursorLua4>(::arg()["lua-dns-script"]);
-      L<<Logger::Warning<<"Loaded 'lua' script from '"<<::arg()["lua-dns-script"]<<"'"<<endl;
+  if(worker && (!g_weDistributeQueries || t_id != s_distributorThreadID)) {
+    try {
+      if(!::arg()["lua-dns-script"].empty()) {
+        t_pdl = std::make_shared<RecursorLua4>(::arg()["lua-dns-script"]);
+        L<<Logger::Warning<<"Loaded 'lua' script from '"<<::arg()["lua-dns-script"]<<"'"<<endl;
+      }
+    }
+    catch(std::exception &e) {
+      L<<Logger::Error<<"Failed to load 'lua' script from '"<<::arg()["lua-dns-script"]<<"': "<<e.what()<<endl;
+      _exit(99);
     }
-  }
-  catch(std::exception &e) {
-    L<<Logger::Error<<"Failed to load 'lua' script from '"<<::arg()["lua-dns-script"]<<"': "<<e.what()<<endl;
-    _exit(99);
   }
 
   unsigned int ringsize=::arg().asNum("stats-ringbuffer-entries") / g_numWorkerThreads;
@@ -3498,6 +3508,7 @@ int main(int argc, char **argv)
     ::arg().set("max-qperq", "Maximum outgoing queries per query")="50";
     ::arg().set("max-total-msec", "Maximum total wall-clock time per query in milliseconds, 0 for unlimited")="7000";
     ::arg().set("max-recursion-depth", "Maximum number of internal recursion calls per query, 0 for unlimited")="40";
+    ::arg().set("max-udp-queries-per-round", "Maximum number of UDP queries processed per recvmsg() round, before returning back to normal processing")="10000";
 
     ::arg().set("include-dir","Include *.conf files from this directory")="";
     ::arg().set("security-poll-suffix","Domain name from which to query security update notifications")="secpoll.powerdns.com.";