From: Pieter Lexis <pieter.lexis@powerdns.com>
Date: Wed, 13 Feb 2019 15:10:56 +0000 (+0100)
Subject: SuffixMatchTree: add remove() functions
X-Git-Tag: rec-4.2.0-beta1~7^2~16
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=a364fbfb6b70c4de3a13fcb66c791f35b9094f33;p=thirdparty%2Fpdns.git

SuffixMatchTree: add remove() functions
---

diff --git a/pdns/dnsname.hh b/pdns/dnsname.hh
index f2d9eee21c..f55d6cbf68 100644
--- 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 d003b8ddcc..ac81e6679c 100644
--- 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);
 }