From: Remi Gacogne <remi.gacogne@powerdns.com>
Date: Fri, 12 Jul 2019 08:18:24 +0000 (+0200)
Subject: dnsdist: Initial implementation of a CDB-based Key Value Store
X-Git-Tag: dnsdist-1.4.0-rc2~9^2~21
X-Git-Url: https://granicus.if.org/sourcecode?a=commitdiff_plain;h=90fe8ae602dcfc08ec902cc092d36013de027000;p=pdns

dnsdist: Initial implementation of a CDB-based Key Value Store
---

diff --git a/configure.ac b/configure.ac
index c2f7c6d7e..912361bc6 100644
--- a/configure.ac
+++ b/configure.ac
@@ -203,6 +203,7 @@ for a in $modules $dynmodules; do
       have_remotebackend=yes
       ;;
     tinydns)
+      needcdb=yes
       PDNS_CHECK_CDB
       ;;
     geoip)
@@ -252,6 +253,15 @@ AS_IF([test "x$needlmdb" = "xyes"], [
   ])
 ])
 
+AM_CONDITIONAL([CDB], [test "x$needcdb" != "x"])
+AM_CONDITIONAL([HAVE_CDB], [test "x$CDB_LIBS" != "x"])
+
+AS_IF([test "x$needcdb" = "xyes"], [
+  AS_IF([test "x$HAVE_CDB" != "x1"], [
+    AC_MSG_ERROR([cdb not found via pkg-config, please install cdb])
+  ])
+])
+
 for a in $modules; do
   AC_MSG_CHECKING([whether we can build module "${a}"])
   if [[ -d "$srcdir/modules/${a}backend" ]]; then
diff --git a/m4/pdns_check_cdb.m4 b/m4/pdns_check_cdb.m4
index fdc00e578..3d1e3498e 100644
--- a/m4/pdns_check_cdb.m4
+++ b/m4/pdns_check_cdb.m4
@@ -1,14 +1,21 @@
 AC_DEFUN([PDNS_CHECK_CDB],[
-  PKG_CHECK_MODULES([CDB], [libcdb],
-    [],
+  PKG_CHECK_MODULES([CDB], [libcdb], [
+    AC_DEFINE([HAVE_CDB], [1], [Define to 1 if you have CDB])
+    [HAVE_CDB=1]
+    ],
     [AC_CHECK_HEADERS([cdb.h],
       [AC_CHECK_LIB([cdb], [cdb_find],
-        [CDB_LIBS="-lcdb"],
-        [AC_MSG_ERROR([Could not find libcdb])]
+        [
+          CDB_LIBS="-lcdb"
+          AC_DEFINE([HAVE_CDB], [1], [Define to 1 if you have CDB])
+          [HAVE_CDB=1]
+        ],
+        [:]
       )],
-      [AC_MSG_ERROR([Could not find cdb.h])]
+      [:]
     )]
   )
   AC_SUBST(CDB_LIBS)
   AC_SUBST(CDB_CFLAGS)
+  AM_CONDITIONAL([HAVE_CDB], [test "x$CDB_LIBS" != "x"])
 ])
diff --git a/m4/systemd.m4 b/m4/systemd.m4
index be2bf7bc4..77919fcd2 100644
--- a/m4/systemd.m4
+++ b/m4/systemd.m4
@@ -192,6 +192,7 @@ AC_DEFUN([AX_CHECK_SYSTEMD_FEATURES], [
         AM_CONDITIONAL([HAVE_SYSTEMD_RESTRICT_ADDRESS_FAMILIES], [ test x"$systemd_restrict_address_families" = "xy" ])
         AM_CONDITIONAL([HAVE_SYSTEMD_RESTRICT_NAMESPACES], [ test x"$systemd_restrict_namespaces" = "xy" ])
         AM_CONDITIONAL([HAVE_SYSTEMD_RESTRICT_REALTIME], [ test x"$systemd_restrict_realtime" = "xy" ])
+        AM_CONDITIONAL([HAVE_SYSTEMD_RESTRICT_SUIDSGID], [ test x"$systemd_restrict_suidsgid" = "xy" ])
         AM_CONDITIONAL([HAVE_SYSTEMD_SYSTEM_CALL_ARCHITECTURES], [ test x"$systemd_system_call_architectures" = "xy" ])
         AM_CONDITIONAL([HAVE_SYSTEMD_SYSTEM_CALL_FILTER], [ test x"$systemd_system_call_filter" = "xy" ])
 ])
diff --git a/pdns/cdb.cc b/pdns/cdb.cc
index 2fa1cd7b7..e2490eff2 100644
--- a/pdns/cdb.cc
+++ b/pdns/cdb.cc
@@ -35,7 +35,7 @@ CDB::CDB(const string &cdbfile)
   d_fd = open(cdbfile.c_str(), O_RDONLY);
   if (d_fd < 0)
   {
-    throw std::runtime_error("Failed to open cdb database file '"+cdbfile+"'. Error: " + stringerror());
+    throw std::runtime_error("Failed to open cdb database file '"+cdbfile+"': " + stringerror());
   }
 
   memset(&d_cdbf,0,sizeof(struct cdb_find));
@@ -44,7 +44,7 @@ CDB::CDB(const string &cdbfile)
   {
     close(d_fd);
     d_fd = -1;
-    throw std::runtime_error("Failed to initialize cdb structure. ErrorNt: '" + std::to_string(cdbinit) + "'");
+    throw std::runtime_error("Failed to initialize cdb structure for database '+cdbfile+': '" + std::to_string(cdbinit) + "'");
   }
 }
 
@@ -102,7 +102,10 @@ bool CDB::readNext(pair<string, string> &value) {
 
     std::string key;
     key.resize(len);
-    cdb_read(&d_cdb, &key[0], len, pos);
+    int ret = cdb_read(&d_cdb, &key[0], len, pos);
+    if (ret < 0) {
+      throw std::runtime_error("Error while reading key for key '" + key + "' from CDB database: " + std::to_string(ret));
+    }
 
     if (d_searchType == SearchSuffix) {
       char *p = strstr(const_cast<char*>(key.c_str()), d_key.c_str());
@@ -115,7 +118,10 @@ bool CDB::readNext(pair<string, string> &value) {
     len = cdb_datalen(&d_cdb);
     std::string val;
     val.resize(len);
-    cdb_read(&d_cdb, &val[0], len, pos);
+    ret = cdb_read(&d_cdb, &val[0], len, pos);
+    if (ret < 0) {
+      throw std::runtime_error("Error while reading value for key '" + key + "' from CDB database: " + std::to_string(ret));
+    }
 
     value = make_pair(std::move(key), std::move(val));
     return true;
@@ -134,7 +140,11 @@ vector<string> CDB::findall(string &key)
   vector<string> ret;
   struct cdb_find cdbf;
 
-  cdb_findinit(&cdbf, &d_cdb, key.c_str(), key.size());
+  int res = cdb_findinit(&cdbf, &d_cdb, key.c_str(), key.size());
+  if (res < 0) {
+    throw std::runtime_error("Error looking up key '" + key + "' from CDB database: " + std::to_string(res));
+  }
+
   int x=0;
   while(cdb_findnext(&cdbf) > 0) {
     x++;
@@ -142,9 +152,67 @@ vector<string> CDB::findall(string &key)
     unsigned int vlen = cdb_datalen(&d_cdb);
     std::string val;
     val.resize(vlen);
-    cdb_read(&d_cdb, &val[0], vlen, vpos);
+    res = cdb_read(&d_cdb, &val[0], vlen, vpos);
+    if (res < 0) {
+      throw std::runtime_error("Error while reading value for key '" + key + "' from CDB database: " + std::to_string(res));
+    }
     ret.push_back(std::move(val));
   }
 
   return ret;
 }
+
+bool CDB::findOne(const string& key, string& value)
+{
+  int ret = cdb_find(&d_cdb, key.c_str(), key.size());
+  if (ret < 0) {
+    throw std::runtime_error("Error while looking up key '" + key + "' from CDB database: " + std::to_string(ret));
+  }
+  if (ret == 0) {
+    /* no such key */
+    return false;
+  }
+
+  unsigned int vpos = cdb_datapos(&d_cdb);
+  unsigned int vlen = cdb_datalen(&d_cdb);
+  value.resize(vlen);
+  ret = cdb_read(&d_cdb, &value[0], vlen, vpos);
+  if (ret < 0) {
+    throw std::runtime_error("Error while reading value for key '" + key + "' from CDB database: " + std::to_string(ret));
+  }
+
+  return true;
+}
+
+CDBWriter::CDBWriter(int fd): d_fd(fd)
+{
+  cdb_make_start(&d_cdbm, d_fd);
+}
+
+CDBWriter::~CDBWriter()
+{
+  close();
+}
+
+void CDBWriter::close()
+{
+  if (d_fd >= 0) {
+    cdb_make_finish(&d_cdbm);
+    ::close(d_fd);
+    d_fd = -1;
+  }
+}
+
+bool CDBWriter::addEntry(const std::string& key, const std::string& value)
+{
+  if (d_fd < 0) {
+    throw std::runtime_error("Can't add an entry to a closed CDB database");
+  }
+
+  int ret = cdb_make_add(&d_cdbm, key.c_str(), key.size(), value.c_str(), value.size());
+  if (ret != 0) {
+    throw std::runtime_error("Error adding key '" + key + "' to CDB database: " + std::to_string(ret));
+  }
+
+  return true;
+}
diff --git a/pdns/cdb.hh b/pdns/cdb.hh
index f1f85a4c1..8179005b1 100644
--- a/pdns/cdb.hh
+++ b/pdns/cdb.hh
@@ -24,7 +24,7 @@
 
 #include <cdb.h>
 
-#include "pdns/misc.hh"
+#include "misc.hh"
 
 // This class is responsible for the reading of a CDB file.
 // The constructor opens the CDB file, the destructor closes it, so make sure you call that.
@@ -34,11 +34,14 @@ public:
   CDB(const string &cdbfile);
   ~CDB();
 
+  /* Return negative value on error or non-negative value on success.
+     Values can be retrieved via readNext() */
   int searchKey(const string &key);
   bool searchSuffix(const string &key);
   void searchAll();
   bool readNext(pair<string, string> &value);
   vector<string> findall(string &key);
+  bool findOne(const string& key, string& value);
 
 private:
   bool moveToNext();
@@ -51,4 +54,20 @@ private:
   enum SearchType { SearchSuffix, SearchKey, SearchAll } d_searchType{SearchKey};
 };
 
+class CDBWriter
+{
+public:
+  /* we own the fd after this call, don't ever touch it */
+  CDBWriter(int fd);
+  ~CDBWriter();
+
+  bool addEntry(const std::string& key, const std::string& value);
+  /* finalize the database and close the fd, the only thing you can do now is to call the destructor */
+  void close();
+
+private:
+  struct cdb_make d_cdbm;
+  int d_fd{-1};
+};
+
 #endif // CDB_HH
diff --git a/pdns/dnsdist-lua-actions.cc b/pdns/dnsdist-lua-actions.cc
index b6526088c..773491034 100644
--- a/pdns/dnsdist-lua-actions.cc
+++ b/pdns/dnsdist-lua-actions.cc
@@ -1117,7 +1117,8 @@ public:
   DNSAction::Action operator()(DNSQuestion* dq, std::string* ruleresult) const override
   {
     std::string key = d_key->getKey(*dq);
-    std::string result = d_kvs->getValue(key);
+    std::string result;
+    d_kvs->getValue(key, result);
 
     if (!dq->qTag) {
       dq->qTag = std::make_shared<QTag>();
diff --git a/pdns/dnsdist-lua-bindings.cc b/pdns/dnsdist-lua-bindings.cc
index 4dad8fbe4..bc83d808d 100644
--- a/pdns/dnsdist-lua-bindings.cc
+++ b/pdns/dnsdist-lua-bindings.cc
@@ -735,4 +735,13 @@ void setupLuaBindings(bool client)
   });
 #endif /* HAVE_LMDB */
 
+#ifdef HAVE_CDB
+  g_lua.writeFunction("newCDBKVStore", [client](const std::string& fname) {
+    if (client) {
+      return std::shared_ptr<KeyValueStore>(nullptr);
+    }
+    return std::shared_ptr<KeyValueStore>(new CDBKVStore(fname));
+  });
+#endif /* HAVE_CDB */
+
 }
diff --git a/pdns/dnsdist.cc b/pdns/dnsdist.cc
index 6089b854e..367a89646 100644
--- a/pdns/dnsdist.cc
+++ b/pdns/dnsdist.cc
@@ -2510,6 +2510,9 @@ try
       cout<<"dnsdist "<<VERSION<<" ("<<LUA_RELEASE<<")"<<endl;
 #endif
       cout<<"Enabled features: ";
+#ifdef HAVE_CDB
+      cout<<"cdb ";
+#endif
 #ifdef HAVE_DNS_OVER_TLS
       cout<<"dns-over-tls(";
 #ifdef HAVE_GNUTLS
diff --git a/pdns/dnsdistdist/Makefile.am b/pdns/dnsdistdist/Makefile.am
index 98c662965..7ea3678d0 100644
--- a/pdns/dnsdistdist/Makefile.am
+++ b/pdns/dnsdistdist/Makefile.am
@@ -254,6 +254,13 @@ testrunner_LDADD = \
 	$(SANITIZER_FLAGS) \
 	$(LIBCAP_LIBS)
 
+if HAVE_CDB
+dnsdist_LDADD += $(CDB_LDFLAGS) $(CDB_LIBS)
+testrunner_LDADD += $(CDB_LDFLAGS) $(CDB_LIBS)
+dnsdist_SOURCES += cdb.cc cdb.hh
+testrunner_SOURCES += cdb.cc cdb.hh
+endif
+
 if HAVE_RE2
 dnsdist_LDADD += $(RE2_LIBS)
 endif
diff --git a/pdns/dnsdistdist/cdb.cc b/pdns/dnsdistdist/cdb.cc
new file mode 120000
index 000000000..915796f65
--- /dev/null
+++ b/pdns/dnsdistdist/cdb.cc
@@ -0,0 +1 @@
+../cdb.cc
\ No newline at end of file
diff --git a/pdns/dnsdistdist/cdb.hh b/pdns/dnsdistdist/cdb.hh
new file mode 120000
index 000000000..4ab150b4c
--- /dev/null
+++ b/pdns/dnsdistdist/cdb.hh
@@ -0,0 +1 @@
+../cdb.hh
\ No newline at end of file
diff --git a/pdns/dnsdistdist/configure.ac b/pdns/dnsdistdist/configure.ac
index 34b484fdd..8a27a917e 100644
--- a/pdns/dnsdistdist/configure.ac
+++ b/pdns/dnsdistdist/configure.ac
@@ -64,6 +64,7 @@ PDNS_CHECK_LUA_HPP
 AM_CONDITIONAL([HAVE_GNUTLS], [false])
 AM_CONDITIONAL([HAVE_LIBSSL], [false])
 AM_CONDITIONAL([HAVE_LMDB], [false])
+AM_CONDITIONAL([HAVE_CDB], [false])
 
 PDNS_CHECK_LIBCRYPTO
 
@@ -93,6 +94,7 @@ AS_IF([test "x$enable_dns_over_https" != "xno"], [
   ])
 ])
 
+PDNS_CHECK_CDB
 PDNS_CHECK_LMDB
 
 AX_CXX_COMPILE_STDCXX_11([ext], [mandatory])
@@ -216,6 +218,10 @@ AS_IF([test "x$enable_dns_over_tls" != "xno" -o "x$enable_dns_over_https" != "xn
     [AC_MSG_NOTICE([OpenSSL: no])]
   )]
 )
+AS_IF([test "x$CDB_LIBS" != "x"],
+  [AC_MSG_NOTICE([cdb: yes])],
+  [AC_MSG_NOTICE([cdb: no])]
+)
 AS_IF([test "x$LMDB_LIBS" != "x"],
   [AC_MSG_NOTICE([lmdb: yes])],
   [AC_MSG_NOTICE([lmdb: no])]
diff --git a/pdns/dnsdistdist/dnsdist-kvs.cc b/pdns/dnsdistdist/dnsdist-kvs.cc
index f29e3d644..80def8d4a 100644
--- a/pdns/dnsdistdist/dnsdist-kvs.cc
+++ b/pdns/dnsdistdist/dnsdist-kvs.cc
@@ -25,9 +25,7 @@
 
 #ifdef HAVE_LMDB
 
-#include "lmdb-safe.hh"
-
-std::string LMDBKVStore::getValue(const std::string& key)
+bool LMDBKVStore::getValue(const std::string& key, std::string& value)
 {
   string_view result;
   try {
@@ -35,16 +33,34 @@ std::string LMDBKVStore::getValue(const std::string& key)
     auto dbi = transaction.openDB(d_dbName, 0);
     int rc = transaction.get(dbi, MDBInVal(key), result);
     if (rc == 0) {
-      return result.to_string();
+      value = result.to_string();
+      return true;
     }
     else if (rc == MDB_NOTFOUND) {
-      return std::string();
+      return false;
     }
   }
   catch(const std::exception& e) {
     warnlog("Error while looking up key '%s' from LMDB file '%s', database '%s': %s", key, d_fname, d_dbName);
   }
-  return std::string();
+  return false;
 }
 
 #endif /* HAVE_LMDB */
+
+#ifdef HAVE_CDB
+
+bool CDBKVStore::getValue(const std::string& key, std::string& value)
+{
+  try {
+    if (d_cdb.findOne(key, value)) {
+      return true;
+    }
+  }
+  catch(const std::exception& e) {
+    warnlog("Error while looking up key '%s' from CDB file '%s': %s", key, d_fname);
+  }
+  return false;
+}
+
+#endif /* HAVE_CDB */
diff --git a/pdns/dnsdistdist/dnsdist-kvs.hh b/pdns/dnsdistdist/dnsdist-kvs.hh
index aa5654536..2767bb7c5 100644
--- a/pdns/dnsdistdist/dnsdist-kvs.hh
+++ b/pdns/dnsdistdist/dnsdist-kvs.hh
@@ -103,7 +103,7 @@ public:
   {
   }
 
-  virtual std::string getValue(const std::string& key) = 0;
+  virtual bool getValue(const std::string& key, std::string& value) = 0;
 };
 
 #ifdef HAVE_LMDB
@@ -117,7 +117,7 @@ public:
   {
   }
 
-  std::string getValue(const std::string& key) override;
+  bool getValue(const std::string& key, std::string& value) override;
 
 private:
   MDBEnv d_env;
@@ -126,3 +126,23 @@ private:
 };
 
 #endif /* HAVE_LMDB */
+
+#ifdef HAVE_CDB
+
+#include "cdb.hh"
+
+class CDBKVStore: public KeyValueStore
+{
+public:
+  CDBKVStore(const std::string& fname): d_cdb(fname), d_fname(fname)
+  {
+  }
+
+  bool getValue(const std::string& key, std::string& value) override;
+
+private:
+  CDB d_cdb;
+  std::string d_fname;
+};
+
+#endif /* HAVE_LMDB */
diff --git a/pdns/dnsdistdist/m4/pdns_check_cdb.m4 b/pdns/dnsdistdist/m4/pdns_check_cdb.m4
new file mode 120000
index 000000000..02cc19493
--- /dev/null
+++ b/pdns/dnsdistdist/m4/pdns_check_cdb.m4
@@ -0,0 +1 @@
+../../../m4/pdns_check_cdb.m4
\ No newline at end of file
diff --git a/pdns/dnsdistdist/test-dnsdistkvs_cc.cc b/pdns/dnsdistdist/test-dnsdistkvs_cc.cc
index 00e530013..8cab16a74 100644
--- a/pdns/dnsdistdist/test-dnsdistkvs_cc.cc
+++ b/pdns/dnsdistdist/test-dnsdistkvs_cc.cc
@@ -32,14 +32,58 @@ BOOST_AUTO_TEST_CASE(test_LMDB) {
 
   DNSQuestion dq(&qname, qtype, qclass, qname.wirelength(), &lc, &rem, &dh, bufferSize, queryLen, isTcp, &queryRealTime);
 
+  std::string value;
   DTime dt;
   dt.set();
   for (size_t idx = 0; idx < 10000000; idx++) {
-    std::string value = lmdb->getValue(key->getKey(dq));
+    value.clear();
+    BOOST_CHECK_EQUAL(lmdb->getValue(key->getKey(dq), value), true);
     BOOST_CHECK_EQUAL(value, "this is the value of the tag");
   }
   cerr<<dt.udiff()/1000/1000<<endl;
 }
 #endif /* HAVE_LMDB */
 
+#ifdef HAVE_CDB
+BOOST_AUTO_TEST_CASE(test_CDB) {
+
+  DNSName qname("powerdns.com.");
+  uint16_t qtype = QType::A;
+  uint16_t qclass = QClass::IN;
+  ComboAddress lc("127.0.0.1:53");
+  ComboAddress rem("127.0.0.1:42");
+  struct dnsheader dh;
+  memset(&dh, 0, sizeof(dh));
+  size_t bufferSize = 0;
+  size_t queryLen = 0;
+  bool isTcp = false;
+  struct timespec queryRealTime;
+  gettime(&queryRealTime, true);
+  struct timespec expiredTime;
+  /* the internal QPS limiter does not use the real time */
+  gettime(&expiredTime);
+
+  DNSQuestion dq(&qname, qtype, qclass, qname.wirelength(), &lc, &rem, &dh, bufferSize, queryLen, isTcp, &queryRealTime);
+
+  char db[] = "/tmp/test_cdb.XXXXXX";
+  int fd = mkstemp(db);
+  BOOST_REQUIRE(fd >= 0);
+  CDBWriter writer(fd);
+  BOOST_REQUIRE(writer.addEntry(std::string(reinterpret_cast<const char*>(&dq.remote->sin4.sin_addr.s_addr), sizeof(dq.remote->sin4.sin_addr.s_addr)), "this is the value of the tag"));
+  writer.close();
+
+  auto cdb = make_unique<CDBKVStore>(db);
+  auto key = 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");
+  }
+  cerr<<dt.udiff()/1000/1000<<endl;
+}
+#endif /* HAVE_CDB */
+
 BOOST_AUTO_TEST_SUITE_END()