From: Remi Gacogne <remi.gacogne@powerdns.com>
Date: Wed, 17 Jul 2019 16:08:49 +0000 (+0200)
Subject: dnsdist: Add KVS lookup, lookupSuffix and reload bindings to Lua
X-Git-Tag: dnsdist-1.4.0-rc2~9^2~16
X-Git-Url: https://granicus.if.org/sourcecode?a=commitdiff_plain;h=14de749b6b75f7b32ff0ed7493b150954d062407;p=pdns

dnsdist: Add KVS lookup, lookupSuffix and reload bindings to Lua
---

diff --git a/pdns/dnsdist-lua-bindings.cc b/pdns/dnsdist-lua-bindings.cc
index e83cdcc06..1c97ae572 100644
--- a/pdns/dnsdist-lua-bindings.cc
+++ b/pdns/dnsdist-lua-bindings.cc
@@ -747,4 +747,64 @@ void setupLuaBindings(bool client)
   });
 #endif /* HAVE_CDB */
 
+  g_lua.registerFunction<std::string(std::shared_ptr<KeyValueStore>::*)(const std::string&)>("lookup", [](std::shared_ptr<KeyValueStore>& kvs, const std::string& keyStr) {
+    std::string result;
+    if (!kvs) {
+      return result;
+    }
+
+    try {
+      ComboAddress ca(keyStr);
+      KeyValueLookupKeySourceIP lookup;
+      for (const auto& key : lookup.getKeys(ca)) {
+        if (kvs->getValue(key, result)) {
+          return result;
+        }
+      }
+    }
+    catch(const std::exception& e) {
+      /* not a valid address, treating it as a DNSName */
+      try {
+        DNSName dn(keyStr);
+        KeyValueLookupKeyQName lookup;
+        for (const auto& key : lookup.getKeys(dn)) {
+          if (kvs->getValue(key, result)) {
+            return result;
+          }
+        }
+      }
+      catch (const std::exception& e) {
+        /* not a valid name, trying to pass it as it is */
+        kvs->getValue(keyStr, result);
+      }
+    }
+
+    return result;
+  });
+
+  g_lua.registerFunction<std::string(std::shared_ptr<KeyValueStore>::*)(const std::string&)>("lookupSuffix", [](std::shared_ptr<KeyValueStore>& kvs, const std::string& keyStr) {
+    std::string result;
+    if (!kvs) {
+      return result;
+    }
+
+    DNSName dn(keyStr);
+    KeyValueLookupKeySuffix lookup;
+    for (const auto& key : lookup.getKeys(dn)) {
+      if (kvs->getValue(key, result)) {
+        return result;
+      }
+    }
+
+    return result;
+  });
+
+  g_lua.registerFunction<bool(std::shared_ptr<KeyValueStore>::*)()>("reload", [](std::shared_ptr<KeyValueStore>& kvs) {
+    if (!kvs) {
+      return false;
+    }
+
+    return kvs->reload();
+  });
+
 }
diff --git a/pdns/dnsdistdist/dnsdist-kvs.cc b/pdns/dnsdistdist/dnsdist-kvs.cc
index c36a326ae..51a33d202 100644
--- a/pdns/dnsdistdist/dnsdist-kvs.cc
+++ b/pdns/dnsdistdist/dnsdist-kvs.cc
@@ -25,15 +25,15 @@
 
 #include <sys/stat.h>
 
-std::vector<std::string> KeyValueLookupKeySourceIP::getKeys(const DNSQuestion& dq)
+std::vector<std::string> KeyValueLookupKeySourceIP::getKeys(const ComboAddress& addr)
 {
   std::vector<std::string> result;
 
-  if (dq.remote->sin4.sin_family == AF_INET) {
-    result.emplace_back(reinterpret_cast<const char*>(&dq.remote->sin4.sin_addr.s_addr), sizeof(dq.remote->sin4.sin_addr.s_addr));
+  if (addr.sin4.sin_family == AF_INET) {
+    result.emplace_back(reinterpret_cast<const char*>(&addr.sin4.sin_addr.s_addr), sizeof(addr.sin4.sin_addr.s_addr));
   }
-  else if (dq.remote->sin4.sin_family == AF_INET6) {
-    result.emplace_back(reinterpret_cast<const char*>(&dq.remote->sin6.sin6_addr.s6_addr), sizeof(dq.remote->sin6.sin6_addr.s6_addr));
+  else if (addr.sin4.sin_family == AF_INET6) {
+    result.emplace_back(reinterpret_cast<const char*>(&addr.sin6.sin6_addr.s6_addr), sizeof(addr.sin6.sin6_addr.s6_addr));
   }
 
   return result;
@@ -99,6 +99,29 @@ CDBKVStore::CDBKVStore(const std::string& fname, time_t refreshDelay): d_fname(f
   refreshDBIfNeeded(now);
 }
 
+bool CDBKVStore::reload(const struct stat& st)
+{
+  auto newCDB = make_unique<CDB>(d_fname);
+  {
+    WriteLock wl(&d_lock);
+    d_cdb = std::move(newCDB);
+  }
+  d_mtime = st.st_mtime;
+  return true;
+}
+
+bool CDBKVStore::reload()
+{
+  struct stat st;
+  if (stat(d_fname.c_str(), &st) == 0) {
+    return reload(st);
+  }
+  else {
+    warnlog("Error while retrieving the last modification time of CDB database '%s': %s", d_fname, stringerror());
+    return false;
+  }
+}
+
 void CDBKVStore::refreshDBIfNeeded(time_t now)
 {
   if (d_refreshing.test_and_set()) {
@@ -110,12 +133,7 @@ void CDBKVStore::refreshDBIfNeeded(time_t now)
     struct stat st;
     if (stat(d_fname.c_str(), &st) == 0) {
       if (st.st_mtime > d_mtime) {
-        auto newCDB = make_unique<CDB>(d_fname);
-        {
-          WriteLock wl(&d_lock);
-          d_cdb = std::move(newCDB);
-        }
-        d_mtime = st.st_mtime;
+        reload(st);
       }
     }
     else {
diff --git a/pdns/dnsdistdist/dnsdist-kvs.hh b/pdns/dnsdistdist/dnsdist-kvs.hh
index 758104a94..92c027eb8 100644
--- a/pdns/dnsdistdist/dnsdist-kvs.hh
+++ b/pdns/dnsdistdist/dnsdist-kvs.hh
@@ -36,7 +36,12 @@ public:
 class KeyValueLookupKeySourceIP: public KeyValueLookupKey
 {
 public:
-  std::vector<std::string> getKeys(const DNSQuestion& dq) override;
+  std::vector<std::string> getKeys(const ComboAddress& addr);
+
+  std::vector<std::string> getKeys(const DNSQuestion& dq) override
+  {
+    return getKeys(*dq.remote);
+  }
 
   std::string toString() const override
   {
@@ -115,6 +120,10 @@ public:
   }
 
   virtual bool getValue(const std::string& key, std::string& value) = 0;
+  virtual bool reload()
+  {
+    return false;
+  }
 };
 
 #ifdef HAVE_LMDB
@@ -148,9 +157,11 @@ public:
   CDBKVStore(const std::string& fname, time_t refreshDelay);
 
   bool getValue(const std::string& key, std::string& value) override;
+  bool reload() override;
 
 private:
   void refreshDBIfNeeded(time_t now);
+  bool reload(const struct stat& st);
 
   std::unique_ptr<CDB> d_cdb{nullptr};
   std::string d_fname;