implement SuffixMatchNode, unit tests, expose it to Lua and use it in dnsdistconf.lua

diff --git a/pdns/dnsdist.cc b/pdns/dnsdist.cc
index ebd09c9cb8453484e85ed02eb94c2bc10a58cb6b..57d3256f2e1cee9729c1c619fb75abed60dc4902 100644 (file)
--- a/pdns/dnsdist.cc
+++ b/pdns/dnsdist.cc
@@ -563,6 +563,10 @@ void setupLua()
   g_lua.registerFunction("isPartOf", &DNSName::isPartOf);
   g_lua.registerFunction("tostring", &DNSName::toString);
   g_lua.writeFunction("newDNSName", [](const std::string& name) { return DNSName(name); });
+  g_lua.writeFunction("newSuffixNode", []() { return SuffixMatchNode(); });
+
+  g_lua.registerFunction("add",(void (SuffixMatchNode::*)(const DNSName&)) &SuffixMatchNode::add);
+  g_lua.registerFunction("check",(bool (SuffixMatchNode::*)(const DNSName&) const) &SuffixMatchNode::check);
 
   g_lua.executeCode(ifs);
 }

diff --git a/pdns/dnsdistconf.lua b/pdns/dnsdistconf.lua
index c1d63462508a33df30e51ad2ce3d9bc95df67405..37dc8255ab8d90c44cee632bdf12662f3063c67b 100644 (file)
--- a/pdns/dnsdistconf.lua
+++ b/pdns/dnsdistconf.lua
@@ -21,21 +21,23 @@ end
 
 counter=0
 
-block=newDNSName("ezdns.it.")
+block=newSuffixNode()
+block:add(newDNSName("ezdns.it."))
+block:add(newDNSName("xxx."))
 
 -- called to pick a downstream server
 function pickServer(remote, qname, qtype) 
-       print("qname: ",qname:tostring())
-       local servers
-       if(qname:isPartOf(block))
-       then 
-               servers=abuse 
-       else
+        local servers
+                if(block:check(qname))
+                then 
+                   print("Sending to abuse pool: ",qname:tostring())   
+           servers=abuse 
+        else
                servers=good
-       end
+        end
 
-       counter=counter+1;
-       return servers[1 + (counter % #servers)]
+        counter=counter+1;
+        return servers[1 + (counter % #servers)]
 end
 
 

diff --git a/pdns/dnsname.hh b/pdns/dnsname.hh
index d6b1c4450fd2e841f9244fe196dcffdacd2fd8dc..c478ce9849b421eb18b2fb589e7b834c8414a498 100644 (file)
--- a/pdns/dnsname.hh
+++ b/pdns/dnsname.hh
@@ -1,6 +1,8 @@
 #pragma once
 #include <string>
 #include <deque>
+#include <set>
+#include <strings.h>
 
 /* Quest in life: 
      accept escaped ascii presentations of DNS names and store them "natively"
@@ -37,3 +39,60 @@ private:
   static std::string escapeLabel(const std::string& orig);
   static std::string unescapeLabel(const std::string& orig);
 };
+
+
+/* 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 */
+struct SuffixMatchNode
+{
+  SuffixMatchNode(const std::string& name_="", bool endNode_=false) : name(name_), endNode(endNode_)
+  {}
+  std::string name;
+  mutable bool endNode;
+  mutable std::set<SuffixMatchNode> children;
+  bool operator<(const SuffixMatchNode& rhs) const
+  {
+    return strcasecmp(name.c_str(), rhs.name.c_str()) < 0;
+  }
+
+  void add(const DNSName& name) 
+  {
+    add(name.getRawLabels());
+  }
+
+  void add(std::deque<std::string> labels) const
+  {
+    if(labels.empty()) { // this allows insertion of the root
+      endNode=true;
+    }
+    else if(labels.size()==1) {
+      children.insert({*labels.begin(), true});
+    }
+    else {
+      auto res=children.insert({*labels.rbegin(), false});
+      labels.pop_back();
+      res.first->add(labels);
+    }
+  }
+
+  bool check(const DNSName& name)  const
+  {
+    return check(name.getRawLabels());
+  }
+
+
+  bool check(std::deque<std::string> labels) const
+  {
+    if(labels.empty()) // optimization
+      return endNode; 
+
+    SuffixMatchNode smn({*labels.rbegin()});
+    auto child = children.find(smn);
+    if(child == children.end())
+      return endNode;
+    labels.pop_back();
+    return child->check(labels);
+  }
+            
+
+};

diff --git a/pdns/test-dnsname_cc.cc b/pdns/test-dnsname_cc.cc
index 440bf61f14e66fbc7811efbb88ecb944fb590618..91fa0b49ea2df4bdc70903349e54f958a5513138 100644 (file)
--- a/pdns/test-dnsname_cc.cc
+++ b/pdns/test-dnsname_cc.cc
@@ -117,4 +117,32 @@ BOOST_AUTO_TEST_CASE(test_packetParse) {
   BOOST_CHECK_EQUAL(qtype, QType::AAAA);
 }
 
+BOOST_AUTO_TEST_CASE(test_suffixmatch) {
+  SuffixMatchNode smn;
+  DNSName ezdns("ezdns.it.");
+  smn.add(ezdns.getRawLabels());
+
+  smn.add(DNSName("org.").getRawLabels());
+
+  DNSName wwwpowerdnscom("www.powerdns.com.");
+  DNSName wwwezdnsit("www.ezdns.it.");
+  BOOST_CHECK(smn.check(wwwezdnsit));
+  BOOST_CHECK(!smn.check(wwwpowerdnscom));
+
+  BOOST_CHECK(smn.check(DNSName("www.powerdns.org.")));
+  BOOST_CHECK(smn.check(DNSName("www.powerdns.oRG.")));
+
+  smn.add(DNSName("news.bbc.co.uk."));
+  BOOST_CHECK(smn.check(DNSName("news.bbc.co.uk.")));
+  BOOST_CHECK(smn.check(DNSName("www.news.bbc.co.uk.")));
+  BOOST_CHECK(smn.check(DNSName("www.www.www.www.www.news.bbc.co.uk.")));
+  BOOST_CHECK(!smn.check(DNSName("images.bbc.co.uk.")));
+
+  BOOST_CHECK(!smn.check(DNSName("www.news.gov.uk.")));
+
+  smn.add(DNSName()); // block the root
+  BOOST_CHECK(smn.check(DNSName("a.root-servers.net.")));
+
+
+}
 BOOST_AUTO_TEST_SUITE_END()