From 1a083287851fb12c48504f6fc3c4fcac9a1256c7 Mon Sep 17 00:00:00 2001
From: Remi Gacogne <remi.gacogne@powerdns.com>
Date: Thu, 6 Oct 2016 14:35:09 +0200
Subject: [PATCH] dnsdist: Fix insertion issues in SuffixMatchTree

---
 pdns/dnsname.hh         | 17 +++++++-------
 pdns/test-dnsname_cc.cc | 51 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 60 insertions(+), 8 deletions(-)

diff --git a/pdns/dnsname.hh b/pdns/dnsname.hh
index f951dfaf15..492def3dd7 100644
--- a/pdns/dnsname.hh
+++ b/pdns/dnsname.hh
@@ -301,13 +301,11 @@ struct SuffixMatchTree
   SuffixMatchTree(const SuffixMatchTree& rhs)
   {
     name = rhs.name;
-    d_human = rhs.d_human;
     children = rhs.children;
     endNode = rhs.endNode;
     d_value = rhs.d_value;
   }
   std::string name;
-  std::string d_human;
   mutable std::set<SuffixMatchTree> children;
   mutable bool endNode;
   mutable T d_value;
@@ -338,16 +336,19 @@ struct SuffixMatchTree
     }
     else if(labels.size()==1) {
       SuffixMatchTree newChild(*labels.begin(), true);
-      newChild.d_value=value;
-      children.insert(newChild);
+      auto res=children.insert(newChild);
+      if(!res.second) {
+        // we might already have had the node as an
+        // intermediary one, but it's now an end node
+        if(!res.first->endNode) {
+          res.first->endNode = true;
+        }
+      }
+      res.first->d_value = value;
     }
     else {
       SuffixMatchTree newnode(*labels.rbegin(), false);
       auto res=children.insert(newnode);
-      if(!res.second) {
-        children.erase(newnode);
-        res=children.insert(newnode);
-      }
       labels.pop_back();
       res.first->add(labels, value);
     }
diff --git a/pdns/test-dnsname_cc.cc b/pdns/test-dnsname_cc.cc
index 4995ee45e8..ea76cae70c 100644
--- a/pdns/test-dnsname_cc.cc
+++ b/pdns/test-dnsname_cc.cc
@@ -493,6 +493,57 @@ BOOST_AUTO_TEST_CASE(test_suffixmatch) {
   BOOST_CHECK(smn.check(DNSName("a.root-servers.net.")));
 }
 
+BOOST_AUTO_TEST_CASE(test_suffixmatch_tree) {
+  SuffixMatchTree<DNSName> smt;
+  DNSName ezdns("ezdns.it.");
+  smt.add(ezdns, ezdns);
+
+  smt.add(DNSName("org.").getRawLabels(), DNSName("org."));
+
+  DNSName wwwpowerdnscom("www.powerdns.com.");
+  DNSName wwwezdnsit("www.ezdns.it.");
+  BOOST_REQUIRE(smt.lookup(wwwezdnsit));
+  BOOST_CHECK_EQUAL(*smt.lookup(wwwezdnsit), ezdns);
+  BOOST_CHECK(smt.lookup(wwwpowerdnscom) == nullptr);
+
+  BOOST_REQUIRE(smt.lookup(DNSName("www.powerdns.org.")));
+  BOOST_CHECK_EQUAL(*smt.lookup(DNSName("www.powerdns.org.")), DNSName("org."));
+  BOOST_REQUIRE(smt.lookup(DNSName("www.powerdns.oRG.")));
+  BOOST_CHECK_EQUAL(*smt.lookup(DNSName("www.powerdns.oRG.")), DNSName("org."));
+
+  smt.add(DNSName("news.bbc.co.uk."), DNSName("news.bbc.co.uk."));
+  BOOST_REQUIRE(smt.lookup(DNSName("news.bbc.co.uk.")));
+  BOOST_CHECK_EQUAL(*smt.lookup(DNSName("news.bbc.co.uk.")), DNSName("news.bbc.co.uk."));
+  BOOST_REQUIRE(smt.lookup(DNSName("www.news.bbc.co.uk.")));
+  BOOST_CHECK_EQUAL(*smt.lookup(DNSName("www.news.bbc.co.uk.")), DNSName("news.bbc.co.uk."));
+  BOOST_REQUIRE(smt.lookup(DNSName("www.www.www.www.www.news.bbc.co.uk.")));
+  BOOST_CHECK_EQUAL(*smt.lookup(DNSName("www.www.www.www.www.news.bbc.co.uk.")), DNSName("news.bbc.co.uk."));
+  BOOST_CHECK(smt.lookup(DNSName("images.bbc.co.uk.")) == nullptr);
+  BOOST_CHECK(smt.lookup(DNSName("www.news.gov.uk.")) == nullptr);
+
+  smt.add(g_rootdnsname, g_rootdnsname); // block the root
+  BOOST_REQUIRE(smt.lookup(DNSName("a.root-servers.net.")));
+  BOOST_CHECK_EQUAL(*smt.lookup(DNSName("a.root-servers.net.")), g_rootdnsname);
+
+  DNSName apowerdnscom("a.powerdns.com.");
+  DNSName bpowerdnscom("b.powerdns.com.");
+  smt.add(apowerdnscom, apowerdnscom);
+  smt.add(bpowerdnscom, bpowerdnscom);
+  BOOST_REQUIRE(smt.lookup(apowerdnscom));
+  BOOST_CHECK_EQUAL(*smt.lookup(apowerdnscom), apowerdnscom);
+  BOOST_REQUIRE(smt.lookup(bpowerdnscom));
+  BOOST_CHECK_EQUAL(*smt.lookup(bpowerdnscom), bpowerdnscom);
+
+  DNSName examplenet("example.net.");
+  DNSName net("net.");
+  smt.add(examplenet, examplenet);
+  smt.add(net, net);
+  BOOST_REQUIRE(smt.lookup(examplenet));
+  BOOST_CHECK_EQUAL(*smt.lookup(examplenet), examplenet);
+  BOOST_REQUIRE(smt.lookup(net));
+  BOOST_CHECK_EQUAL(*smt.lookup(net), net);
+}
+
 
 BOOST_AUTO_TEST_CASE(test_concat) {
   DNSName first("www."), second("powerdns.com.");
-- 
2.47.2