From: bert hubert <bert.hubert@netherlabs.nl>
Date: Thu, 19 Feb 2015 11:49:23 +0000 (+0100)
Subject: implement SuffixMatchNode, unit tests, expose it to Lua and use it in dnsdistconf.lua
X-Git-Tag: dnsdist-1.0.0-alpha1~248^2~88^2~134
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=ceee6652f020f8df3376d45522a833b660e2c99a;p=thirdparty%2Fpdns.git

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 ebd09c9cb8..57d3256f2e 100644
--- 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 c1d6346250..37dc8255ab 100644
--- 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 d6b1c4450f..c478ce9849 100644
--- 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 440bf61f14..91fa0b49ea 100644
--- 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()