From: Remi Gacogne <remi.gacogne@powerdns.com>
Date: Fri, 12 Jul 2019 15:38:10 +0000 (+0200)
Subject: dnsdist: Add KVS suffix matching
X-Git-Tag: dnsdist-1.4.0-rc2~9^2~20
X-Git-Url: https://granicus.if.org/sourcecode?a=commitdiff_plain;h=77eebb379b6803e9db437abccc25b5d20f760c44;p=pdns

dnsdist: Add KVS suffix matching
---

diff --git a/pdns/dnsdist-lua-actions.cc b/pdns/dnsdist-lua-actions.cc
index 773491034..3f7c3a638 100644
--- a/pdns/dnsdist-lua-actions.cc
+++ b/pdns/dnsdist-lua-actions.cc
@@ -1116,9 +1116,13 @@ public:
 
   DNSAction::Action operator()(DNSQuestion* dq, std::string* ruleresult) const override
   {
-    std::string key = d_key->getKey(*dq);
+    std::vector<std::string> keys = d_key->getKeys(*dq);
     std::string result;
-    d_kvs->getValue(key, result);
+    for (const auto& key : keys) {
+      if (d_kvs->getValue(key, result) == true) {
+        break;
+      }
+    }
 
     if (!dq->qTag) {
       dq->qTag = std::make_shared<QTag>();
diff --git a/pdns/dnsdist-lua-bindings.cc b/pdns/dnsdist-lua-bindings.cc
index bc83d808d..dc6483646 100644
--- a/pdns/dnsdist-lua-bindings.cc
+++ b/pdns/dnsdist-lua-bindings.cc
@@ -722,6 +722,9 @@ void setupLuaBindings(bool client)
   g_lua.writeFunction("KeyValueLookupKeyQName", []() {
     return std::shared_ptr<KeyValueLookupKey>(new KeyValueLookupKeyQName());
   });
+  g_lua.writeFunction("KeyValueLookupKeySuffix", []() {
+    return std::shared_ptr<KeyValueLookupKey>(new KeyValueLookupKeySuffix());
+  });
   g_lua.writeFunction("KeyValueLookupKeyTag", [](const std::string& tag) {
     return std::shared_ptr<KeyValueLookupKey>(new KeyValueLookupKeyTag(tag));
   });
diff --git a/pdns/dnsdistdist/dnsdist-kvs.hh b/pdns/dnsdistdist/dnsdist-kvs.hh
index 2767bb7c5..fc2fa2a4c 100644
--- a/pdns/dnsdistdist/dnsdist-kvs.hh
+++ b/pdns/dnsdistdist/dnsdist-kvs.hh
@@ -29,23 +29,25 @@ public:
   virtual ~KeyValueLookupKey()
   {
   }
-  virtual std::string getKey(const DNSQuestion&) = 0;
+  virtual std::vector<std::string> getKeys(const DNSQuestion&) = 0;
   virtual std::string toString() const = 0;
 };
 
 class KeyValueLookupKeySourceIP: public KeyValueLookupKey
 {
 public:
-  std::string getKey(const DNSQuestion& dq) override
+  std::vector<std::string> getKeys(const DNSQuestion& dq) override
   {
-    std::string key;
+    std::vector<std::string> result;
+
     if (dq.remote->sin4.sin_family == AF_INET) {
-      key = std::string(reinterpret_cast<const char*>(&dq.remote->sin4.sin_addr.s_addr), sizeof(dq.remote->sin4.sin_addr.s_addr));
+      result.emplace_back(reinterpret_cast<const char*>(&dq.remote->sin4.sin_addr.s_addr), sizeof(dq.remote->sin4.sin_addr.s_addr));
     }
     else if (dq.remote->sin4.sin_family == AF_INET6) {
-      key = std::string(reinterpret_cast<const char*>(&dq.remote->sin6.sin6_addr.s6_addr), sizeof(dq.remote->sin6.sin6_addr.s6_addr));
+      result.emplace_back(reinterpret_cast<const char*>(&dq.remote->sin6.sin6_addr.s6_addr), sizeof(dq.remote->sin6.sin6_addr.s6_addr));
     }
-    return key;
+
+    return result;
   }
 
   std::string toString() const override
@@ -57,9 +59,31 @@ public:
 class KeyValueLookupKeyQName: public KeyValueLookupKey
 {
 public:
-  std::string getKey(const DNSQuestion& dq) override
+  std::vector<std::string> getKeys(const DNSQuestion& dq) override
+  {
+    return {dq.qname->toDNSStringLC()};
+  }
+
+  std::string toString() const override
+  {
+    return "qname";
+  }
+};
+
+class KeyValueLookupKeySuffix: public KeyValueLookupKey
+{
+public:
+  std::vector<std::string> getKeys(const DNSQuestion& dq) override
   {
-    return dq.qname->toDNSStringLC();
+    auto lowerQName = dq.qname->makeLowerCase();
+    std::vector<std::string> result(lowerQName.countLabels());
+
+    do {
+      result.emplace_back(lowerQName.toDNSString());
+    }
+    while (lowerQName.chopOff());
+
+    return result;
   }
 
   std::string toString() const override
@@ -75,16 +99,15 @@ public:
   {
   }
 
-  std::string getKey(const DNSQuestion& dq) override
+  std::vector<std::string> getKeys(const DNSQuestion& dq) override
   {
-    std::string key;
     if (dq.qTag) {
       const auto& it = dq.qTag->find(d_tag);
       if (it != dq.qTag->end()) {
-        key = it->second;
+        return { it->second };
       }
     }
-    return key;
+    return {};
   }
 
   std::string toString() const override
diff --git a/pdns/dnsdistdist/test-dnsdistkvs_cc.cc b/pdns/dnsdistdist/test-dnsdistkvs_cc.cc
index 8cab16a74..ccd40983e 100644
--- a/pdns/dnsdistdist/test-dnsdistkvs_cc.cc
+++ b/pdns/dnsdistdist/test-dnsdistkvs_cc.cc
@@ -11,9 +11,6 @@ BOOST_AUTO_TEST_SUITE(dnsdistkvs_cc)
 #ifdef HAVE_LMDB
 BOOST_AUTO_TEST_CASE(test_LMDB) {
 
-  auto lmdb = make_unique<LMDBKVStore>("/data/Dumps/lmdb", "db-name");
-  auto key = make_unique<KeyValueLookupKeySourceIP>();
-
   DNSName qname("powerdns.com.");
   uint16_t qtype = QType::A;
   uint16_t qclass = QClass::IN;
@@ -32,13 +29,19 @@ BOOST_AUTO_TEST_CASE(test_LMDB) {
 
   DNSQuestion dq(&qname, qtype, qclass, qname.wirelength(), &lc, &rem, &dh, bufferSize, queryLen, isTcp, &queryRealTime);
 
+  auto lmdb = make_unique<LMDBKVStore>("/data/Dumps/lmdb", "db-name");
+  auto lookupKey = make_unique<KeyValueLookupKeySourceIP>();
+
   std::string value;
   DTime dt;
   dt.set();
   for (size_t idx = 0; idx < 10000000; idx++) {
-    value.clear();
-    BOOST_CHECK_EQUAL(lmdb->getValue(key->getKey(dq), value), true);
-    BOOST_CHECK_EQUAL(value, "this is the value of the tag");
+    auto keys = lookupKey->getKeys(dq);
+    for (const auto& key : keys) {
+      value.clear();
+      BOOST_CHECK_EQUAL(lmdb->getValue(key, value), true);
+      BOOST_CHECK_EQUAL(value, "this is the value of the tag");
+    }
   }
   cerr<<dt.udiff()/1000/1000<<endl;
 }
@@ -73,14 +76,17 @@ BOOST_AUTO_TEST_CASE(test_CDB) {
   writer.close();
 
   auto cdb = make_unique<CDBKVStore>(db);
-  auto key = make_unique<KeyValueLookupKeySourceIP>();
+  auto lookupKey = make_unique<KeyValueLookupKeySourceIP>();
 
   std::string value;
   DTime dt;
   dt.set();
   for (size_t idx = 0; idx < 10000000; idx++) {
-    BOOST_CHECK_EQUAL(cdb->getValue(key->getKey(dq), value), true);
-    BOOST_CHECK_EQUAL(value, "this is the value of the tag");
+    auto keys = lookupKey->getKeys(dq);
+    for (const auto& key : keys) {
+      BOOST_CHECK_EQUAL(cdb->getValue(key, value), true);
+      BOOST_CHECK_EQUAL(value, "this is the value of the tag");
+    }
   }
   cerr<<dt.udiff()/1000/1000<<endl;
 }