From: Remi Gacogne <remi.gacogne@powerdns.com>
Date: Tue, 15 Aug 2017 12:38:42 +0000 (+0200)
Subject: Better support for deleting entries in NetmaskTree and NetmaskGroup
X-Git-Tag: auth-4.1.0-rc3~23^2
X-Git-Url: https://granicus.if.org/sourcecode?a=commitdiff_plain;h=11f4719bc5ed0657aa0aed7ea806f073b9324e32;p=pdns

Better support for deleting entries in NetmaskTree and NetmaskGroup

- NetmaskTree erase() method now has optional cleanup parameter
  - cleanup_tree() method will remove "hanging" branches of the tree
- NetmaskGroup now has deleteMask() methods that use NetmaskTree:erase()

Grabbed from the weakforced tree:
- https://github.com/PowerDNS/weakforced/commit/d5c916a2d00ec3bec3a09055de5709976e184c95
- https://github.com/PowerDNS/weakforced/commit/1170d8da692c0033e0a80c22447d0c00fb0cb7b2
---

diff --git a/pdns/iputils.hh b/pdns/iputils.hh
index bf6ba4c9f..e5ac78471 100644
--- a/pdns/iputils.hh
+++ b/pdns/iputils.hh
@@ -550,10 +550,13 @@ private:
   };
 
 public:
-  NetmaskTree() noexcept {
+  NetmaskTree() noexcept : NetmaskTree(false) {
   }
 
-  NetmaskTree(const NetmaskTree& rhs) {
+  NetmaskTree(bool cleanup) noexcept : d_cleanup_tree(cleanup) {
+  }
+
+  NetmaskTree(const NetmaskTree& rhs): d_cleanup_tree(rhs.d_cleanup_tree) {
     // it is easier to copy the nodes than tree.
     // also acts as handy compactor
     for(auto const& node: rhs._nodes)
@@ -711,6 +714,24 @@ public:
     return ret;
   }
 
+  void cleanup_tree(TreeNode* node)
+  {
+    // only cleanup this node if it has no children and node4 and node6 are both empty
+    if (!(node->left || node->right || node->node6 || node->node4)) {
+      // get parent node ptr
+      TreeNode* parent = node->parent;
+      // delete this node
+      if (parent) {
+	if (parent->left.get() == node)
+	  parent->left.reset();
+	else
+	  parent->right.reset();
+	// now recurse up to the parent
+	cleanup_tree(parent);
+      }
+    }
+  }
+
   //<! Removes key from TreeMap. This does not clean up the tree.
   void erase(const key_type& key) {
     TreeNode *node = root.get();
@@ -738,7 +759,11 @@ public:
           else
             it++;
         }
+
         node->node4.reset();
+
+        if (d_cleanup_tree)
+          cleanup_tree(node);
       }
     } else {
       uint64_t* addr = (uint64_t*)key.getNetwork().sin6.sin6_addr.s6_addr;
@@ -765,6 +790,9 @@ public:
         }
 
         node->node6.reset();
+
+        if (d_cleanup_tree)
+          cleanup_tree(node);
       }
     }
   }
@@ -807,6 +835,7 @@ public:
 private:
   unique_ptr<TreeNode> root; //<! Root of our tree
   std::vector<node_type*> _nodes; //<! Container for actual values
+  bool d_cleanup_tree; //<! Whether or not to cleanup the tree on erase
 };
 
 /** This class represents a group of supplemental Netmask classes. An IP address matchs
@@ -815,6 +844,14 @@ private:
 class NetmaskGroup
 {
 public:
+  //! By default, initialise the tree to cleanup
+  NetmaskGroup() noexcept : NetmaskGroup(true) {
+  }
+
+  //! This allows control over whether to cleanup or not
+  NetmaskGroup(bool cleanup) noexcept : tree(cleanup) {
+  }
+
   //! If this IP address is matched by any of the classes within
 
   bool match(const ComboAddress *ip) const
@@ -829,6 +866,23 @@ public:
     return match(&ip);
   }
 
+  bool lookup(const ComboAddress* ip, Netmask* nmp) const
+  {
+    const auto &ret = tree.lookup(*ip);
+    if (ret) {
+      if (nmp != nullptr)
+        *nmp = ret->first;
+
+      return ret->second;
+    }
+    return false;
+  }
+
+  bool lookup(const ComboAddress& ip, Netmask* nmp) const
+  {
+    return lookup(&ip, nmp);
+  }
+
   //! Add this string to the list of possible matches
   void addMask(const string &ip, bool positive=true)
   {
@@ -845,6 +899,18 @@ public:
     tree.insert(nm).second=positive;
   }
 
+  //! Delete this Netmask from the list of possible matches
+  void deleteMask(const Netmask& nm)
+  {
+    tree.erase(nm);
+  }
+
+  void deleteMask(const std::string& ip)
+  {
+    if (!ip.empty())
+      deleteMask(Netmask(ip));
+  }
+
   void clear()
   {
     tree.clear();
diff --git a/pdns/test-nmtree.cc b/pdns/test-nmtree.cc
index b929acad4..a2fd11bc8 100644
--- a/pdns/test-nmtree.cc
+++ b/pdns/test-nmtree.cc
@@ -97,7 +97,7 @@ BOOST_AUTO_TEST_CASE(test_scale) {
 
 BOOST_AUTO_TEST_CASE(test_removal) {
   std::string prefix = "192.";
-  NetmaskTree<int> nmt;
+  NetmaskTree<int> nmt(true);
 
   size_t count = 0;
   for(unsigned int i = 0; i < 256; ++i) {