From: bert hubert <bert.hubert@netherlabs.nl>
Date: Fri, 10 Jun 2016 16:59:36 +0000 (+0200)
Subject: dnsdist query node infra
X-Git-Tag: auth-4.0.0-rc1~19^2~4
X-Git-Url: https://granicus.if.org/sourcecode?a=commitdiff_plain;h=71c946759511bf01716a639ffd85faded43ae728;p=pdns

dnsdist query node infra
---

diff --git a/pdns/dnsdist-console.cc b/pdns/dnsdist-console.cc
index eb1476dcd..dabd5e700 100644
--- a/pdns/dnsdist-console.cc
+++ b/pdns/dnsdist-console.cc
@@ -153,7 +153,8 @@ void doConsole()
             boost::variant<
               string, 
               shared_ptr<DownstreamState>,
-              std::unordered_map<string, double>
+              std::unordered_map<string, double>,
+              std::vector<pair<unsigned int, std::unordered_map<string, string> > >
               >
             >
           >(withReturn ? ("return "+line) : line);
@@ -173,6 +174,20 @@ void doConsole()
             Json out = o;
             cout<<out.dump()<<endl;
           }
+          else if(const auto um = boost::get<std::vector<pair<unsigned int, std::unordered_map<string, string> > >>(&*ret)) {
+            cerr<<"So.. "<<um->size()<<endl;
+            using namespace json11;
+
+            Json::array a;
+            for(auto& e : *um) {
+              Json::object o;
+              o["qname"]=e.second["qname"];
+              o["rcode"]=e.second["rcode"];
+              a.push_back(o);
+            }
+            Json out=a;
+            cout<<out.dump()<<endl;
+          }
         }
         else 
           cout << g_outputBuffer;
diff --git a/pdns/dnsdist-lua2.cc b/pdns/dnsdist-lua2.cc
index 17be04bac..657f6b68b 100644
--- a/pdns/dnsdist-lua2.cc
+++ b/pdns/dnsdist-lua2.cc
@@ -10,6 +10,7 @@
 #include <map>
 #include <fstream>
 #include <boost/logic/tribool.hpp>
+#include "statnode.hh"
 
 boost::tribool g_noLuaSideEffect;
 
@@ -55,6 +56,61 @@ map<ComboAddress,int> filterScore(const map<ComboAddress, unsigned int,ComboAddr
 }
 
 
+typedef std::function<void(const StatNode&, const StatNode::Stat&, const StatNode::Stat&)> statvisitor_t;
+
+void statNodeRespRing(statvisitor_t visitor)
+{
+  std::lock_guard<std::mutex> lock(g_rings.respMutex);
+  
+  StatNode root;
+  for(const auto& c : g_rings.respRing) {
+    root.submit(c.name, c.dh.rcode, c.requestor);
+  }
+  StatNode::Stat node;
+  try {
+    root.visit([&visitor](const StatNode* node, const StatNode::Stat& self, const StatNode::Stat& children) {
+        visitor(*node, self, children);}
+      , node);  
+  }
+  catch(const LuaContext::ExecutionErrorException& e) {
+    std::cerr << e.what(); 
+    try {
+      std::rethrow_if_nested(e);
+      std::cerr << std::endl;
+    } catch(const std::exception& e) {
+      // e is the exception that was thrown from inside the lambda
+      std::cerr << ": " << e.what() << std::endl;      
+    }
+    catch(const PDNSException& e) {
+      // e is the exception that was thrown from inside the lambda
+      std::cerr << ": " << e.reason << std::endl;      
+    }
+  }
+  catch(std::exception& e) {
+    cerr<<"error: "<<e.what()<<endl;
+  }
+
+}
+
+vector<pair<unsigned int, std::unordered_map<string,string> > > getRespRing(boost::optional<int> rcode) 
+{
+  typedef std::unordered_map<string,string>  entry_t;
+  vector<pair<unsigned int, entry_t > > ret;
+  std::lock_guard<std::mutex> lock(g_rings.respMutex);
+  
+  entry_t e;
+  unsigned int count=1;
+  for(const auto& c : g_rings.respRing) {
+    if(rcode && (rcode.get() != c.dh.rcode))
+      continue;
+    e["qname"]=c.name.toString();
+    e["rcode"]=std::to_string(c.dh.rcode);
+    ret.push_back(std::make_pair(count,e));
+    count++;
+  }
+  return ret;
+}
+
 typedef   map<ComboAddress, unsigned int,ComboAddress::addressOnlyLessThan > counts_t;
 map<ComboAddress,int> exceedRespGen(int rate, int seconds, std::function<void(counts_t&, const Rings::Response&)> T) 
 {
@@ -120,6 +176,7 @@ map<ComboAddress,int> exceedRespByterate(int rate, int seconds)
 }
 
 
+
 void moreLua(bool client)
 {
   typedef NetmaskTree<DynBlock> nmts_t;
@@ -143,17 +200,29 @@ void moreLua(bool client)
       struct timespec now;
       gettime(&now);
       boost::format fmt("%-24s %8d %8d %s\n");
-      g_outputBuffer = (fmt % "Netmask" % "Seconds" % "Blocks" % "Reason").str();
+      g_outputBuffer = (fmt % "What" % "Seconds" % "Blocks" % "Reason").str();
       for(const auto& e: slow) {
 	if(now < e->second.until)
 	  g_outputBuffer+= (fmt % e->first.toString() % (e->second.until.tv_sec - now.tv_sec) % e->second.blocks % e->second.reason).str();
       }
+      auto slow2 = g_dynblockSMT.getCopy();
+      slow2.visit([&now, &fmt](const SuffixMatchTree<DynBlock>& node) {
+          if(now <node.d_value.until) {
+            string dom("empty");
+            if(!node.d_value.domain.empty())
+              dom = node.d_value.domain.toString();
+            g_outputBuffer+= (fmt % dom % (node.d_value.until.tv_sec - now.tv_sec) % node.d_value.blocks % node.d_value.reason).str();
+          }
+        });
+
     });
 
   g_lua.writeFunction("clearDynBlocks", []() {
       setLuaSideEffect();
       nmts_t nmg;
       g_dynblockNMG.setState(nmg);
+      SuffixMatchTree<DynBlock> smt;
+      g_dynblockSMT.setState(smt);
     });
 
   g_lua.writeFunction("addDynBlocks", 
@@ -186,6 +255,40 @@ void moreLua(bool client)
 			   g_dynblockNMG.setState(slow);
 			 });
 
+  g_lua.writeFunction("addDynBlockSMT", 
+                      [](const vector<pair<unsigned int, string> >&names, const std::string& msg, boost::optional<int> seconds) { 
+                           setLuaSideEffect();
+			   auto slow = g_dynblockSMT.getCopy();
+			   struct timespec until, now;
+			   gettime(&now);
+			   until=now;
+                           int actualSeconds = seconds ? *seconds : 10;
+			   until.tv_sec += actualSeconds; 
+
+ 			   for(const auto& capair : names) {
+			     unsigned int count = 0;
+                             DNSName domain(capair.second);
+                             auto got = slow.lookup(domain);
+                             bool expired=false;
+			     if(got) {
+			       if(until < got->until) // had a longer policy
+				 continue;
+			       if(now < got->until) // only inherit count on fresh query we are extending
+				 count=got->blocks;
+                               else
+                                 expired=true;
+			     }
+
+			     DynBlock db{msg,until,domain};
+			     db.blocks=count;
+                             if(!got || expired)
+                               warnlog("Inserting dynamic block for %s for %d seconds: %s", domain, actualSeconds, msg);
+			     slow.add(domain, db);
+			   }
+			   g_dynblockSMT.setState(slow);
+			 });
+
+
 
   g_lua.registerFunction<bool(nmts_t::*)(const ComboAddress&)>("match", 
 								     [](nmts_t& s, const ComboAddress& ca) { return s.match(ca); });
@@ -221,6 +324,20 @@ void moreLua(bool client)
 	});
     });
 
+  g_lua.writeFunction("getRespRing", getRespRing);
+
+  g_lua.registerFunction<StatNode, unsigned int()>("numChildren", 
+                                                      [](StatNode& sn) -> unsigned int {
+
+                                                        return sn.children.size();
+                                                      } );
+
+  g_lua.registerMember("fullname", &StatNode::fullname);
+  g_lua.registerMember("servfails", &StatNode::Stat::servfails);
+  g_lua.registerMember("nxdomains", &StatNode::Stat::nxdomains);
+  g_lua.registerMember("queries", &StatNode::Stat::queries);
+
+  g_lua.writeFunction("statNodeRespRing", statNodeRespRing);
 
   g_lua.writeFunction("getTopBandwidth", [](unsigned int top) {
       setLuaNoSideEffect();
diff --git a/pdns/dnsdist-tcp.cc b/pdns/dnsdist-tcp.cc
index 4adb768b7..34e8a1430 100644
--- a/pdns/dnsdist-tcp.cc
+++ b/pdns/dnsdist-tcp.cc
@@ -165,6 +165,7 @@ void* tcpClientThread(int pipefd)
   auto localRulactions = g_rulactions.getLocal();
   auto localRespRulactions = g_resprulactions.getLocal();
   auto localDynBlockNMG = g_dynblockNMG.getLocal();
+  auto localDynBlockSMT = g_dynblockSMT.getLocal();
   auto localPools = g_pools.getLocal();
 #ifdef HAVE_PROTOBUF
   boost::uuids::random_generator uuidGenerator;
@@ -267,7 +268,7 @@ void* tcpClientThread(int pipefd)
 	struct timespec now;
 	gettime(&now);
 
-	if (!processQuery(localDynBlockNMG, localRulactions, blockFilter, dq, poolname, &delayMsec, now)) {
+	if (!processQuery(localDynBlockNMG, localDynBlockSMT, localRulactions, blockFilter, dq, poolname, &delayMsec, now)) {
 	  goto drop;
 	}
 
diff --git a/pdns/dnsdist.cc b/pdns/dnsdist.cc
index a496e2b6e..120f98d7e 100644
--- a/pdns/dnsdist.cc
+++ b/pdns/dnsdist.cc
@@ -119,6 +119,7 @@ Rings g_rings;
 
 GlobalStateHolder<servers_t> g_dstates;
 GlobalStateHolder<NetmaskTree<DynBlock>> g_dynblockNMG;
+GlobalStateHolder<SuffixMatchTree<DynBlock>> g_dynblockSMT;
 int g_tcpRecvTimeout{2};
 int g_tcpSendTimeout{2};
 
@@ -729,14 +730,16 @@ void spoofResponseFromString(DNSQuestion& dq, const string& spoofContent)
   }
 }
 
-bool processQuery(LocalStateHolder<NetmaskTree<DynBlock> >& localDynBlock, LocalStateHolder<vector<pair<std::shared_ptr<DNSRule>, std::shared_ptr<DNSAction> > > >& localRulactions, blockfilter_t blockFilter, DNSQuestion& dq, string& poolname, int* delayMsec, const struct timespec& now)
+bool processQuery(LocalStateHolder<NetmaskTree<DynBlock> >& localDynNMGBlock, 
+                  LocalStateHolder<SuffixMatchTree<DynBlock> >& localDynSMTBlock,
+                  LocalStateHolder<vector<pair<std::shared_ptr<DNSRule>, std::shared_ptr<DNSAction> > > >& localRulactions, blockfilter_t blockFilter, DNSQuestion& dq, string& poolname, int* delayMsec, const struct timespec& now)
 {
   {
     WriteLock wl(&g_rings.queryLock);
     g_rings.queryRing.push_back({now,*dq.remote,*dq.qname,dq.len,dq.qtype,*dq.dh});
   }
 
-  if(auto got=localDynBlock->lookup(*dq.remote)) {
+  if(auto got=localDynNMGBlock->lookup(*dq.remote)) {
     if(now < got->second.until) {
       vinfolog("Query from %s dropped because of dynamic block", dq.remote->toStringWithPort());
       g_stats.dynBlocked++;
@@ -745,6 +748,16 @@ bool processQuery(LocalStateHolder<NetmaskTree<DynBlock> >& localDynBlock, Local
     }
   }
 
+  if(auto got=localDynSMTBlock->lookup(*dq.qname)) {
+    if(now < got->until) {
+      vinfolog("Query from %s for %s dropped because of dynamic block", dq.remote->toStringWithPort(), dq.qname->toString());
+      g_stats.dynBlocked++;
+      got->blocks++;
+      return false;
+    }
+  }
+
+
   if(blockFilter) {
     std::lock_guard<std::mutex> lock(g_luamutex);
 
@@ -851,7 +864,8 @@ try
   auto localPolicy = g_policy.getLocal();
   auto localRulactions = g_rulactions.getLocal();
   auto localServers = g_dstates.getLocal();
-  auto localDynBlock = g_dynblockNMG.getLocal();
+  auto localDynNMGBlock = g_dynblockNMG.getLocal();
+  auto localDynSMTBlock = g_dynblockSMT.getLocal();
   auto localPools = g_pools.getLocal();
   struct msghdr msgh;
   struct iovec iov;
@@ -945,7 +959,7 @@ try
       struct timespec now;
       gettime(&now);
 
-      if (!processQuery(localDynBlock, localRulactions, blockFilter, dq, poolname, &delayMsec, now))
+      if (!processQuery(localDynNMGBlock, localDynSMTBlock, localRulactions, blockFilter, dq, poolname, &delayMsec, now))
       {
         continue;
       }
diff --git a/pdns/dnsdist.hh b/pdns/dnsdist.hh
index 5b14ca955..5709238ef 100644
--- a/pdns/dnsdist.hh
+++ b/pdns/dnsdist.hh
@@ -29,12 +29,14 @@ struct DynBlock
   {
     reason=rhs.reason;
     until=rhs.until;
+    domain=rhs.domain;
     blocks.store(rhs.blocks);
     return *this;
   }
-  
+
   string reason;
   struct timespec until;
+  DNSName domain;
   mutable std::atomic<unsigned int> blocks;
 };
 
@@ -468,6 +470,102 @@ enum ednsHeaderFlags {
   EDNS_HEADER_FLAG_DO = 32768
 };
 
+/* Quest in life: serve as a rapid block list. If you add a DNSName to a root SuffixMatchNode, 
+   anything part of that domain will return 'true' in check */
+template<typename T>
+struct SuffixMatchTree
+{
+  SuffixMatchTree(const std::string& name_="", bool endNode_=false) : name(name_), endNode(endNode_)
+  {}
+
+  SuffixMatchTree(const SuffixMatchTree& rhs)
+  {
+    name = rhs.name;
+    d_human = rhs.d_human;
+    children = rhs.children;
+    endNode = rhs.endNode;
+    d_value = rhs.d_value;
+  }
+  std::string name;
+  std::string d_human;
+  mutable std::set<SuffixMatchTree> children;
+  mutable bool endNode;
+  mutable T d_value;
+  bool operator<(const SuffixMatchTree& rhs) const
+  {
+    return strcasecmp(name.c_str(), rhs.name.c_str()) < 0;
+  }
+  typedef SuffixMatchTree value_type;
+
+  template<typename V>
+  void visit(const V& v) const {
+    for(const auto& c : children) 
+      c.visit(v);
+    if(endNode)
+      v(*this);
+  }
+
+  void add(const DNSName& name, const T& t) 
+  {
+    add(name.getRawLabels(), t);
+  }
+
+  void add(std::vector<std::string> labels, const T& value) const
+  {
+    if(labels.empty()) { // this allows insertion of the root
+      endNode=true;
+      d_value=value;
+    }
+    else if(labels.size()==1) {
+      SuffixMatchTree newChild(*labels.begin(), true);
+      newChild.d_value=value;
+      children.insert(newChild);
+    }
+    else {
+      SuffixMatchTree newnode(*labels.rbegin(), false);
+      auto res=children.insert(newnode);
+      if(!res.second) {
+        children.erase(newnode);
+        res=children.insert(newnode);
+      }
+      labels.pop_back();
+      res.first->add(labels, value);
+    }
+  }
+
+  T* lookup(const DNSName& name)  const
+  {
+    if(children.empty()) { // speed up empty set
+      if(endNode)
+        return &d_value;
+      return 0;
+    }
+    return lookup(name.getRawLabels());
+  }
+
+  T* lookup(std::vector<std::string> labels) const
+  {
+    if(labels.empty()) { // optimization
+      if(endNode)
+        return &d_value;
+      return 0;
+    }
+
+    SuffixMatchTree smn(*labels.rbegin());
+    auto child = children.find(smn);
+    if(child == children.end()) {
+      if(endNode)
+        return &d_value;
+      return 0;
+    }
+    labels.pop_back();
+    return child->lookup(labels);
+  }
+  
+};
+
+extern GlobalStateHolder<SuffixMatchTree<DynBlock>> g_dynblockSMT;
+
 extern GlobalStateHolder<vector<CarbonConfig> > g_carbon;
 extern GlobalStateHolder<ServerPolicy> g_policy;
 extern GlobalStateHolder<servers_t> g_dstates;
@@ -531,7 +629,8 @@ bool getLuaNoSideEffect(); // set if there were only explicit declarations of _n
 void resetLuaSideEffect(); // reset to indeterminate state
 
 bool responseContentMatches(const char* response, const uint16_t responseLen, const DNSName& qname, const uint16_t qtype, const uint16_t qclass, const ComboAddress& remote);
-bool processQuery(LocalStateHolder<NetmaskTree<DynBlock> >& localDynBlock, LocalStateHolder<vector<pair<std::shared_ptr<DNSRule>, std::shared_ptr<DNSAction> > > >& localRulactions, blockfilter_t blockFilter, DNSQuestion& dq, string& poolname, int* delayMsec, const struct timespec& now);
+bool processQuery(LocalStateHolder<NetmaskTree<DynBlock> >& localDynBlockNMG,
+                  LocalStateHolder<SuffixMatchTree<DynBlock> >& localDynBlockSMT, LocalStateHolder<vector<pair<std::shared_ptr<DNSRule>, std::shared_ptr<DNSAction> > > >& localRulactions, blockfilter_t blockFilter, DNSQuestion& dq, string& poolname, int* delayMsec, const struct timespec& now);
 bool processResponse(LocalStateHolder<vector<pair<std::shared_ptr<DNSRule>, std::shared_ptr<DNSResponseAction> > > >& localRespRulactions, DNSQuestion& dq);
 bool fixUpResponse(char** response, uint16_t* responseLen, size_t* responseSize, const DNSName& qname, uint16_t origFlags, bool ednsAdded, bool ecsAdded, std::vector<uint8_t>& rewrittenResponse, uint16_t addRoom);
 void restoreFlags(struct dnsheader* dh, uint16_t origFlags);
diff --git a/pdns/dnsdistdist/Makefile.am b/pdns/dnsdistdist/Makefile.am
index 1b97108fb..01565fc91 100644
--- a/pdns/dnsdistdist/Makefile.am
+++ b/pdns/dnsdistdist/Makefile.am
@@ -93,6 +93,7 @@ dnsdist_SOURCES = \
 	sholder.hh \
 	sodcrypto.cc sodcrypto.hh \
 	sstuff.hh \
+	statnode.cc statnode.hh \
 	ext/luawrapper/include/LuaContext.hpp \
 	ext/json11/json11.cpp \
 	ext/json11/json11.hpp \