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 e83cdcc060d6e91959aed1aab645b9354230627b..1c97ae572c755ff9d5ec0ce03637c044df4ca4fc 100644 (file)
--- 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 c36a326aeed486b97f9d407f7659647d43182b0c..51a33d202a479eb932b60150457b4678bf950c05 100644 (file)
--- 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 758104a941a923626073570defce078fc4f310a9..92c027eb869a4583458d0ab6364d96c48fc08220 100644 (file)
--- 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;