From 41ff82c903b55df3bd4b6430730b75ef113d454e Mon Sep 17 00:00:00 2001
From: Pieter Lexis <pieter.lexis@powerdns.com>
Date: Tue, 25 Jun 2019 15:36:31 +0200
Subject: [PATCH] SuffixMatchNode:add: Accept more types

The SuffixMatchNode now also accepts a table of DNSNames, a table of
string, and singular strings.
---
 pdns/dnsdist-lua-bindings.cc                 | 27 +++++++++++++++++++-
 pdns/dnsdistdist/docs/reference/config.rst   |  4 +++
 pdns/dnsname.hh                              |  5 ++++
 regression-tests.dnsdist/test_CheckConfig.py |  4 +++
 4 files changed, 39 insertions(+), 1 deletion(-)

diff --git a/pdns/dnsdist-lua-bindings.cc b/pdns/dnsdist-lua-bindings.cc
index d29081e10..322ae318e 100644
--- a/pdns/dnsdist-lua-bindings.cc
+++ b/pdns/dnsdist-lua-bindings.cc
@@ -205,7 +205,32 @@ void setupLuaBindings(bool client)
   g_lua.registerFunction("empty",(bool (DNSNameSet::*)()) &DNSNameSet::empty);
 
   /* SuffixMatchNode */
-  g_lua.registerFunction("add",(void (SuffixMatchNode::*)(const DNSName&)) &SuffixMatchNode::add);
+  g_lua.registerFunction<void (SuffixMatchNode::*)(const boost::variant<DNSName, string, vector<pair<int, DNSName>>, vector<pair<int, string>>> &name)>("add", [](SuffixMatchNode &smn, const boost::variant<DNSName, string, vector<pair<int, DNSName>>, vector<pair<int, string>>> &name) {
+      if (name.type() == typeid(DNSName)) {
+          auto n = boost::get<DNSName>(name);
+          smn.add(n);
+          return;
+      }
+      if (name.type() == typeid(string)) {
+          auto n = boost::get<string>(name);
+          smn.add(n);
+          return;
+      }
+      if (name.type() == typeid(vector<pair<int, DNSName>>)) {
+          auto names = boost::get<vector<pair<int, DNSName>>>(name);
+          for (auto const n : names) {
+            smn.add(n.second);
+          }
+          return;
+      }
+      if (name.type() == typeid(vector<pair<int, string>>)) {
+          auto names = boost::get<vector<pair<int, string>>>(name);
+          for (auto const n : names) {
+            smn.add(n.second);
+          }
+          return;
+      }
+  });
   g_lua.registerFunction("check",(bool (SuffixMatchNode::*)(const DNSName&) const) &SuffixMatchNode::check);
 
   /* NetmaskGroup */
diff --git a/pdns/dnsdistdist/docs/reference/config.rst b/pdns/dnsdistdist/docs/reference/config.rst
index 34f876e7b..958843125 100644
--- a/pdns/dnsdistdist/docs/reference/config.rst
+++ b/pdns/dnsdistdist/docs/reference/config.rst
@@ -1072,10 +1072,14 @@ If you are looking for exact name matching, your might want to consider using a
   Represent a set of DNS suffixes for quick matching.
 
   .. method:: SuffixMatchNode:add(name)
+    .. versionchanged:: 1.4.0
+      This method now accepts strings, lists of DNSNames and lists of strings.
 
     Add a suffix to the current set.
 
     :param DNSName name: The suffix to add to the set.
+    :param string name: The suffix to add to the set.
+    :param table name: The suffixes to add to the set. Elements of the table should be of the same type, either DNSName or string.
 
   .. method:: SuffixMatchNode:check(name) -> bool
 
diff --git a/pdns/dnsname.hh b/pdns/dnsname.hh
index d4e429c9c..5b571085e 100644
--- a/pdns/dnsname.hh
+++ b/pdns/dnsname.hh
@@ -394,6 +394,11 @@ struct SuffixMatchNode
       d_nodes.insert(dnsname);
     }
 
+    void add(const std::string& name)
+    {
+      add(DNSName(name));
+    }
+
     void add(std::vector<std::string> labels)
     {
       d_tree.add(labels, true);
diff --git a/regression-tests.dnsdist/test_CheckConfig.py b/regression-tests.dnsdist/test_CheckConfig.py
index ce9345c05..835d7ae8d 100644
--- a/regression-tests.dnsdist/test_CheckConfig.py
+++ b/regression-tests.dnsdist/test_CheckConfig.py
@@ -37,6 +37,10 @@ class TestCheckConfig(unittest.TestCase):
             addAction(RegexRule("evil[0-9]{4,}\\\\.regex\\\\.tests\\\\.powerdns\\\\.com$"), RCodeAction(DNSRCode.REFUSED))
             mySMN = newSuffixMatchNode()
             mySMN:add(newDNSName("nameAndQtype.tests.powerdns.com."))
+            mySMN:add("string.smn.tests.powerdns.com.")
+            mySMN:add("string-no-dot.smn.tests.powerdns.com")
+            mySMN:add({"string-one.smn.tests.powerdns.com", "string-two.smn.tests.powerdns.com"})
+            mySMN:add({newDNSName("dnsname-one.smn.tests.powerdns.com"), newDNSName("dnsname-two.smn.tests.powerdns.com")})
             addAction(AndRule{SuffixMatchNodeRule(mySMN), QTypeRule("TXT")}, RCodeAction(DNSRCode.NOTIMP))
             addAction(makeRule("drop.test.powerdns.com."), DropAction())
         """
-- 
2.40.0