dnsdist: add MMDB class impl

Signed-off-by: Ensar Sarajčić <dev@ensarsarajcic.com>

diff --git a/pdns/dnsdistdist/mmdb.hh b/pdns/dnsdistdist/mmdb.hh
index c4a542c3e2f13f3d4807982f244e77a8d5f89534..2bdb98590c090c30c96ffb7a50b336b4a77c9217 100644 (file)
--- a/pdns/dnsdistdist/mmdb.hh
+++ b/pdns/dnsdistdist/mmdb.hh
@@ -20,3 +20,156 @@
  * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
  */
 #pragma once
+
+#include "dolog.hh"
+#include <maxminddb.h>
+#include <string>
+
+class MMDB
+{
+public:
+  MMDB(const std::string& fname, const std::string& modeStr)
+  {
+    int ec;
+    int flags = 0;
+    if (modeStr == "")
+      /* for the benefit of ifdef */
+      ;
+#ifdef HAVE_MMAP
+    else if (modeStr == "mmap")
+      flags |= MMDB_MODE_MMAP;
+#endif
+    else
+      throw std::runtime_error(std::string("Unsupported mode ") + modeStr + ("for mmdb"));
+    memset(&d_s, 0, sizeof(d_s));
+    if ((ec = MMDB_open(fname.c_str(), flags, &d_s)) < 0)
+      throw std::runtime_error(std::string("Cannot open ") + fname + std::string(": ") + std::string(MMDB_strerror(ec)));
+    VERBOSESLOG(infolog("Opened MMDB database %s (type: %s version: %d.%d)", fname, d_db.metadata.database_type, d_db.metadata.binary_format_major_version, d_db.metadata.binary_format_minor_version),
+                dnsdist::logging::getTopLogger("mmdb")->info(Logr::Info, "Opened MMDB database", "path", Logging::Loggable(fname), "type", Logging::Loggable(d_db.metadata.database_type), "version", Logging::Loggable(std::to_string(d_db.metadata.binary_format_major_version) + "." + std::to_string(d_db.metadata.binary_format_minor_version))));
+  }
+
+  bool queryCountry(string& ret, const string& ip)
+  {
+    MMDB_entry_data_s data;
+    MMDB_lookup_result_s res;
+    if (!mmdbLookup(ip, res))
+      return false;
+    if (MMDB_get_value(&res.entry, &data, "country", "iso_code", NULL) != MMDB_SUCCESS || !data.has_data)
+      return false;
+    ret = string(data.utf8_string, data.data_size);
+    return true;
+  }
+
+  bool queryContinent(string& ret, const string& ip)
+  {
+    MMDB_entry_data_s data;
+    MMDB_lookup_result_s res;
+    if (!mmdbLookup(ip, res))
+      return false;
+    if (MMDB_get_value(&res.entry, &data, "continent", "code", NULL) != MMDB_SUCCESS || !data.has_data)
+      return false;
+    ret = string(data.utf8_string, data.data_size);
+    return true;
+  }
+
+  bool queryAS(string& ret, const string& ip)
+  {
+    MMDB_entry_data_s data;
+    MMDB_lookup_result_s res;
+    if (!mmdbLookup(ip, res))
+      return false;
+    if (MMDB_get_value(&res.entry, &data, "autonomous_system_organization", NULL) != MMDB_SUCCESS || !data.has_data)
+      return false;
+    ret = string(data.utf8_string, data.data_size);
+    return true;
+  }
+
+  bool queryASN(string& ret, const string& ip)
+  {
+    MMDB_entry_data_s data;
+    MMDB_lookup_result_s res;
+    if (!mmdbLookup(ip, res))
+      return false;
+    if (MMDB_get_value(&res.entry, &data, "autonomous_system_number", NULL) != MMDB_SUCCESS || !data.has_data)
+      return false;
+    ret = std::to_string(data.uint32);
+    return true;
+  }
+
+  bool queryRegion(string& ret, const string& ip)
+  {
+    MMDB_entry_data_s data;
+    MMDB_lookup_result_s res;
+    if (!mmdbLookup(ip, res))
+      return false;
+    if (MMDB_get_value(&res.entry, &data, "subdivisions", "0", "iso_code", NULL) != MMDB_SUCCESS || !data.has_data)
+      return false;
+    ret = string(data.utf8_string, data.data_size);
+    return true;
+  }
+
+  bool queryCity(string& ret, const string& ip, const string& language)
+  {
+    MMDB_entry_data_s data;
+    MMDB_lookup_result_s res;
+    if (!mmdbLookup(ip, res))
+      return false;
+    if ((MMDB_get_value(&res.entry, &data, "cities", "0", NULL) != MMDB_SUCCESS || !data.has_data) && (MMDB_get_value(&res.entry, &data, "city", "names", language.c_str(), NULL) != MMDB_SUCCESS || !data.has_data))
+      return false;
+    ret = string(data.utf8_string, data.data_size);
+    return true;
+  }
+
+  bool queryLocation(const string& ip,
+                     double& latitude, double& longitude,
+                     int& prec)
+  {
+    MMDB_entry_data_s data;
+    MMDB_lookup_result_s res;
+    if (!mmdbLookup(ip, res))
+      return false;
+    if (MMDB_get_value(&res.entry, &data, "location", "latitude", NULL) != MMDB_SUCCESS || !data.has_data)
+      return false;
+    latitude = data.double_value;
+    if (MMDB_get_value(&res.entry, &data, "location", "longitude", NULL) != MMDB_SUCCESS || !data.has_data)
+      return false;
+    longitude = data.double_value;
+    if (MMDB_get_value(&res.entry, &data, "location", "accuracy_radius", NULL) != MMDB_SUCCESS || !data.has_data)
+      return false;
+    prec = data.uint16;
+    return true;
+  }
+
+  bool exists(const string& ip)
+  {
+    MMDB_lookup_result_s res;
+    return mmdbLookup(ip, res);
+  }
+
+  ~MMDB() { MMDB_close(&d_s); };
+
+private:
+  MMDB_s d_s;
+
+  bool mmdbLookup(const string& ip, MMDB_lookup_result_s& res)
+  {
+    int gai_ec = 0, mmdb_ec = 0;
+    res = MMDB_lookup_string(&d_s, ip.c_str(), &gai_ec, &mmdb_ec);
+
+    if (gai_ec != 0) {
+      VERBOSESLOG(infolog("mmdbLookup(%s) failed: %s", ip.toString(), MMDB_strerror(mmdb_ec)), dnsdist::logging::getTopLogger("mmdb")->error(Logr::Info, gai_strerror(mmdb_ec), "mmdbLookup failed", "ip", Logging::Loggable(ip)));
+    }
+    else if (mmdb_ec != MMDB_SUCCESS) {
+      VERBOSESLOG(infolog("mmdbLookup(%s) failed: %s", ip.toString(), MMDB_strerror(mmdb_ec)), dnsdist::logging::getTopLogger("mmdb")->error(Logr::Info, MMDB_strerror(mmdb_ec), "mmdbLookup failed", "ip", Logging::Loggable(ip)));
+    }
+    else if (res.found_entry) {
+      // gl.netmask = res.netmask;
+      // /* If it's a IPv6 database, IPv4 netmasks are reduced from 128, so we need to deduct
+      //    96 to get from [96,128] => [0,32] range */
+      // if (!v6 && gl.netmask > 32)
+      //   gl.netmask -= 96;
+      return true;
+    }
+    return false;
+  }
+};