Refactor SuffixMatchNode using a SuffixMatchTree

diff --git a/pdns/dnsname.hh b/pdns/dnsname.hh
index 89f5ee91164ccaf84807b19ab5069d774a579f07..b3482350182a7ac7a502458549ceb217b8adbf1f 100644 (file)
--- a/pdns/dnsname.hh
+++ b/pdns/dnsname.hh
@@ -227,76 +227,6 @@ inline DNSName operator+(const DNSName& lhs, const DNSName& rhs)
   return ret;
 }
 
-/* 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) : d_name(name_), endNode(endNode_)
-  {}
-  std::string d_name;
-  std::string d_human;
-  mutable std::set<SuffixMatchNode> children;
-  mutable bool endNode;
-  bool operator<(const SuffixMatchNode& rhs) const
-  {
-    return strcasecmp(d_name.c_str(), rhs.d_name.c_str()) < 0;
-  }
-
-  void add(const DNSName& dnsname)
-  {
-    if(!d_human.empty())
-      d_human.append(", ");
-    d_human += dnsname.toString();
-    add(dnsname.getRawLabels());
-  }
-
-  void add(std::vector<std::string> labels) const
-  {
-    if(labels.empty()) { // this allows insertion of the root
-      endNode=true;
-    }
-    else if(labels.size()==1) {
-      auto res=children.insert(SuffixMatchNode(*labels.begin(), true));
-      if(!res.second) {
-        if(!res.first->endNode) {
-          res.first->endNode = true;
-        }
-      }
-    }
-    else {
-      auto res=children.insert(SuffixMatchNode(*labels.rbegin(), false));
-      labels.pop_back();
-      res.first->add(labels);
-    }
-  }
-
-  bool check(const DNSName& dnsname)  const
-  {
-    if(children.empty()) // speed up empty set
-      return endNode;
-    return check(dnsname.getRawLabels());
-  }
-
-  bool check(std::vector<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);
-  }
-  
-  std::string toString() const
-  {
-    return d_human;
-  }
-
-};
-
 template<typename T>
 struct SuffixMatchTree
 {
@@ -390,6 +320,41 @@ struct SuffixMatchTree
 
 };
 
+/* 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()
+  {}
+  std::string d_human;
+  SuffixMatchTree<bool> d_tree;
+
+  void add(const DNSName& dnsname)
+  {
+    if(!d_human.empty())
+      d_human.append(", ");
+    d_human += dnsname.toString();
+
+    d_tree.add(dnsname, true);
+  }
+
+  void add(std::vector<std::string> labels)
+  {
+    d_tree.add(labels, true);
+  }
+
+  bool check(const DNSName& dnsname) const
+  {
+    return d_tree.lookup(dnsname) != nullptr;
+  }
+
+  std::string toString() const
+  {
+    return d_human;
+  }
+
+};
+
 std::ostream & operator<<(std::ostream &os, const DNSName& d);
 namespace std {
     template <>