SuffixMatchTree: add remove() functions

diff --git a/pdns/dnsname.hh b/pdns/dnsname.hh
index f2d9eee21c196b3ff12a42d86fa6d0176a924bf9..f55d6cbf6881b0bc3a273b3916a6c9e3238f505e 100644 (file)
--- a/pdns/dnsname.hh
+++ b/pdns/dnsname.hh
@@ -286,6 +286,39 @@ struct SuffixMatchTree
     }
   }
 
+  void remove(const DNSName &name) const
+  {
+    remove(name.getRawLabels());
+  }
+
+  /* Removes the node at `labels`, also make sure that no empty
+   * children will be left behind in memory
+   */
+  void remove(std::vector<std::string> labels) const
+  {
+    SuffixMatchTree smt(*labels.rbegin());
+    auto child = children.find(smt);
+    if (child == children.end()) {
+      // No subnode found, we're done
+      return;
+    }
+
+    // We have found a child
+    labels.pop_back();
+    if (labels.empty()) {
+      // The child is no longer an endnode
+      child->endNode = false;
+      // If the child has no further children, just remove it from the set.
+      if (child->children.empty()) {
+        children.erase(child);
+      }
+      return;
+    }
+
+    // We are not at the end, let the child figure out what to do
+    child->remove(labels);
+  }
+
   T* lookup(const DNSName& name)  const
   {
     if(children.empty()) { // speed up empty set
@@ -340,6 +373,16 @@ struct SuffixMatchNode
     d_tree.add(labels, true);
   }
 
+  void remove(const DNSName& name)
+  {
+    d_tree.remove(name);
+  }
+
+  void remove(std::vector<std::string> labels)
+  {
+    d_tree.remove(labels);
+  }
+
   bool check(const DNSName& dnsname) const
   {
     return d_tree.lookup(dnsname) != nullptr;

diff --git a/pdns/test-dnsname_cc.cc b/pdns/test-dnsname_cc.cc
index d003b8ddcc0c901dbb3df0b6bcc219cf8dcb3247..ac81e6679cf2cad2efd44ce91299cd1223a37158 100644 (file)
--- a/pdns/test-dnsname_cc.cc
+++ b/pdns/test-dnsname_cc.cc
@@ -507,6 +507,11 @@ BOOST_AUTO_TEST_CASE(test_suffixmatch) {
   smn.add(net);
   BOOST_CHECK(smn.check(examplenet));
   BOOST_CHECK(smn.check(net));
+
+  // Remove .net and check that example.net still exists
+  smn.remove(net);
+  BOOST_CHECK_EQUAL(smn.check(net), false);
+  BOOST_CHECK(smn.check(examplenet));
 }
 
 BOOST_AUTO_TEST_CASE(test_suffixmatch_tree) {
@@ -558,6 +563,11 @@ BOOST_AUTO_TEST_CASE(test_suffixmatch_tree) {
   BOOST_CHECK_EQUAL(*smt.lookup(examplenet), examplenet);
   BOOST_REQUIRE(smt.lookup(net));
   BOOST_CHECK_EQUAL(*smt.lookup(net), net);
+
+  // Remove .net, and check that example.net remains
+  smt.remove(net);
+  BOOST_CHECK(smt.lookup(net) == nullptr);
+  BOOST_CHECK_EQUAL(*smt.lookup(examplenet), examplenet);
 }