From 70ba95d89bf220fe64387188465f1f8446181e76 Mon Sep 17 00:00:00 2001
From: Remi Gacogne <remi.gacogne@powerdns.com>
Date: Wed, 7 Aug 2019 12:04:38 +0200
Subject: [PATCH] dnsdist: Take ComboAddress and DNSName in the KVS lookup
 binding

---
 pdns/dnsdist-lua-bindings.cc            | 32 +++++++++++--------------
 pdns/dnsdistdist/docs/reference/kvs.rst |  8 +++----
 2 files changed, 18 insertions(+), 22 deletions(-)

diff --git a/pdns/dnsdist-lua-bindings.cc b/pdns/dnsdist-lua-bindings.cc
index 1c97ae572..0b005d74e 100644
--- a/pdns/dnsdist-lua-bindings.cc
+++ b/pdns/dnsdist-lua-bindings.cc
@@ -747,14 +747,14 @@ 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) {
+  g_lua.registerFunction<std::string(std::shared_ptr<KeyValueStore>::*)(const boost::variant<ComboAddress, DNSName, std::string>)>("lookup", [](std::shared_ptr<KeyValueStore>& kvs, const boost::variant<ComboAddress, DNSName, std::string> keyVar) {
     std::string result;
     if (!kvs) {
       return result;
     }
 
-    try {
-      ComboAddress ca(keyStr);
+    if (keyVar.type() == typeid(ComboAddress)) {
+      const auto ca = *boost::get<ComboAddress>(&keyVar);
       KeyValueLookupKeySourceIP lookup;
       for (const auto& key : lookup.getKeys(ca)) {
         if (kvs->getValue(key, result)) {
@@ -762,33 +762,29 @@ void setupLuaBindings(bool client)
         }
       }
     }
-    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;
-          }
+    else if (keyVar.type() == typeid(DNSName)) {
+      DNSName dn = *boost::get<DNSName>(&keyVar);
+      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);
-      }
+    }
+    else if (keyVar.type() == typeid(std::string)) {
+      std::string keyStr = *boost::get<std::string>(&keyVar);
+      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) {
+  g_lua.registerFunction<std::string(std::shared_ptr<KeyValueStore>::*)(const DNSName&)>("lookupSuffix", [](std::shared_ptr<KeyValueStore>& kvs, const DNSName& dn) {
     std::string result;
     if (!kvs) {
       return result;
     }
 
-    DNSName dn(keyStr);
     KeyValueLookupKeySuffix lookup;
     for (const auto& key : lookup.getKeys(dn)) {
       if (kvs->getValue(key, result)) {
diff --git a/pdns/dnsdistdist/docs/reference/kvs.rst b/pdns/dnsdistdist/docs/reference/kvs.rst
index bb79f63e4..00585f53b 100644
--- a/pdns/dnsdistdist/docs/reference/kvs.rst
+++ b/pdns/dnsdistdist/docs/reference/kvs.rst
@@ -46,16 +46,16 @@ If the value found in the LMDB database for the key '\8powerdns\3com\0' was 'thi
   .. method:: KeyValueStore:lookup(key)
 
     Does a lookup into the corresponding key value store, and return the result as a string.
-    The key is first parsed as a network address, and if that fails into a DNS name. If that also fails the raw string is used for the lookup.
+    The key can be a :class:`ComboAddress` obtained via the :func:`newCA`, a :class:`DNSName` obtained via the :func:`newDNSName` function, or a raw string.
 
-    :param string key: The key to look up
+    :param ComboAddress, DNSName or string key: The key to look up
 
   .. method:: KeyValueStore:lookupSuffix(key)
 
     Does a suffix-based lookup into the corresponding key value store, and return the result as a string.
-    The key is parsed as a DNS name, and several lookups will be done, removing one label from the name at a time until a match has been found or there is no label left.
+    The key should be a :class:`DNSName` object obtained via the :func:`newDNSName` function, and several lookups will be done, removing one label from the name at a time until a match has been found or there is no label left.
 
-    :param string key: The name to look up
+    :param DNSName key: The name to look up
 
   .. method:: KeyValueStore:reload()
 
-- 
2.40.0