From f4a5db3a8de9b52133c31f4caaafe85a674fc7b1 Mon Sep 17 00:00:00 2001
From: Aki Tuomi <cmouse@cmouse.fi>
Date: Sat, 30 Dec 2017 17:54:28 +0200
Subject: [PATCH] geoipbackend: Add mmdb interface

---
 modules/geoipbackend/Makefile.am            |   7 +-
 modules/geoipbackend/OBJECTFILES            |   2 +-
 modules/geoipbackend/OBJECTLIBS             |   2 +-
 modules/geoipbackend/geoipbackend.hh        |   1 -
 modules/geoipbackend/geoipinterface-mmdb.cc | 240 ++++++++++++++++++++
 modules/geoipbackend/geoipinterface.cc      |   4 +-
 modules/geoipbackend/geoipinterface.hh      |   1 +
 7 files changed, 250 insertions(+), 7 deletions(-)
 create mode 100644 modules/geoipbackend/geoipinterface-mmdb.cc

diff --git a/modules/geoipbackend/Makefile.am b/modules/geoipbackend/Makefile.am
index 7aad6f6f6..be692ef08 100644
--- a/modules/geoipbackend/Makefile.am
+++ b/modules/geoipbackend/Makefile.am
@@ -1,4 +1,4 @@
-AM_CPPFLAGS += $(YAML_CFLAGS) $(GEOIP_CFLAGS)
+AM_CPPFLAGS += $(YAML_CFLAGS) $(GEOIP_CFLAGS) $(MMDB_CFLAGS)
 
 EXTRA_DIST = OBJECTFILES OBJECTLIBS
 
@@ -6,6 +6,7 @@ pkglib_LTLIBRARIES = libgeoipbackend.la
 
 libgeoipbackend_la_SOURCES = geoipbackend.cc geoipbackend.hh \
 	geoipinterface.cc geoipinterface.hh \
-	geoipinterface-dat.cc
+	geoipinterface-dat.cc \
+	geoipinterface-mmdb.cc
 libgeoipbackend_la_LDFLAGS = -module -avoid-version
-libgeoipbackend_la_LIBADD = $(YAML_LIBS) $(GEOIP_LIBS)
+libgeoipbackend_la_LIBADD = $(YAML_LIBS) $(GEOIP_LIBS) $(MMDB_LIBS)
diff --git a/modules/geoipbackend/OBJECTFILES b/modules/geoipbackend/OBJECTFILES
index b2245db98..fc3a1ad31 100644
--- a/modules/geoipbackend/OBJECTFILES
+++ b/modules/geoipbackend/OBJECTFILES
@@ -1 +1 @@
-geoipbackend.lo geoipinterface-dat.lo  geoipinterface.lo
+geoipbackend.lo geoipinterface-dat.lo geoipinterface.lo geoipinterface-mmdb.lo
diff --git a/modules/geoipbackend/OBJECTLIBS b/modules/geoipbackend/OBJECTLIBS
index d87d58b84..dfc26246b 100644
--- a/modules/geoipbackend/OBJECTLIBS
+++ b/modules/geoipbackend/OBJECTLIBS
@@ -1 +1 @@
-$(YAML_LIBS) $(GEOIP_LIBS)
+$(YAML_LIBS) $(GEOIP_LIBS) $(MMDB_LIBS)
diff --git a/modules/geoipbackend/geoipbackend.hh b/modules/geoipbackend/geoipbackend.hh
index b224e4174..56774da7b 100644
--- a/modules/geoipbackend/geoipbackend.hh
+++ b/modules/geoipbackend/geoipbackend.hh
@@ -31,7 +31,6 @@
 #include <sys/types.h>
 #include <dirent.h>
 
-#include "GeoIP.h"
 #include "pdns/dnspacket.hh"
 #include "pdns/dns.hh"
 #include "pdns/dnsbackend.hh"
diff --git a/modules/geoipbackend/geoipinterface-mmdb.cc b/modules/geoipbackend/geoipinterface-mmdb.cc
new file mode 100644
index 000000000..488b97f13
--- /dev/null
+++ b/modules/geoipbackend/geoipinterface-mmdb.cc
@@ -0,0 +1,240 @@
+/*
+ * This file is part of PowerDNS or dnsdist.
+ * Copyright -- PowerDNS.COM B.V. and its contributors
+ *
+ * This program is free software; you can redistribute it and/or modify
+ * it under the terms of version 2 of the GNU General Public License as
+ * published by the Free Software Foundation.
+ *
+ * In addition, for the avoidance of any doubt, permission is granted to
+ * link this program with OpenSSL and to (re)distribute the binaries
+ * produced as the result of such linking.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program; if not, write to the Free Software
+ * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
+ */
+#ifdef HAVE_CONFIG_H
+#include "config.h"
+#endif
+#include "geoipbackend.hh"
+#include "geoipinterface.hh"
+
+#ifdef HAVE_MMDB
+
+#include "maxminddb.h"
+
+class GeoIPInterfaceMMDB : public GeoIPInterface {
+public:
+  GeoIPInterfaceMMDB(const string &fname, const string &modeStr, const string& language) {
+    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 PDNSException(string("Unsupported mode ") + modeStr + ("for geoipbackend-mmdb"));
+    memset(&d_s, 0, sizeof(d_s));
+    if ((ec = MMDB_open(fname.c_str(), flags, &d_s)) < 0)
+      throw PDNSException(string("Cannot open ") + fname + string(": ") + string(MMDB_strerror(ec)));
+    d_lang = language;
+    L<<Logger::Debug<<"Opened MMDB database "<<fname<<"(type: "<<d_s.metadata.database_type<<
+                      " version: "<<d_s.metadata.binary_format_major_version << "." <<
+                      d_s.metadata.binary_format_minor_version << ")" << endl;
+  }
+
+  bool queryCountry(string &ret, GeoIPNetmask& gl, const string &ip) override {
+    MMDB_entry_data_s data;
+    MMDB_lookup_result_s res;
+    if (!mmdbLookup(ip, false, gl, 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 queryCountryV6(string &ret, GeoIPNetmask& gl, const string &ip) override {
+    MMDB_entry_data_s data;
+    MMDB_lookup_result_s res;
+    if (!mmdbLookup(ip, true, gl, 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 queryCountry2(string &ret, GeoIPNetmask& gl, const string &ip) override {
+    return queryCountry(ret, gl, ip);
+  }
+
+  bool queryCountry2V6(string &ret, GeoIPNetmask& gl, const string &ip) override {
+    return queryCountryV6(ret, gl, ip);
+  }
+
+  bool queryContinent(string &ret, GeoIPNetmask& gl, const string &ip) override {
+    MMDB_entry_data_s data;
+    MMDB_lookup_result_s res;
+    if (!mmdbLookup(ip, false, gl, 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 queryContinentV6(string &ret, GeoIPNetmask& gl, const string &ip) override {
+    MMDB_entry_data_s data;
+    MMDB_lookup_result_s res;
+    if (!mmdbLookup(ip, true, gl, 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 queryName(string &ret, GeoIPNetmask& gl, const string &ip) override {
+    MMDB_entry_data_s data;
+    MMDB_lookup_result_s res;
+    if (!mmdbLookup(ip, false, gl, 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 queryNameV6(string &ret, GeoIPNetmask& gl, const string &ip) override {
+    MMDB_entry_data_s data;
+    MMDB_lookup_result_s res;
+    if (!mmdbLookup(ip, true, gl, 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 queryASnum(string &ret, GeoIPNetmask& gl, const string &ip) override {
+    MMDB_entry_data_s data;
+    MMDB_lookup_result_s res;
+    if (!mmdbLookup(ip, false, gl, 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 queryASnumV6(string &ret, GeoIPNetmask& gl, const string &ip) override {
+    MMDB_entry_data_s data;
+    MMDB_lookup_result_s res;
+    if (!mmdbLookup(ip, true, gl, 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, GeoIPNetmask& gl, const string &ip) override {
+    MMDB_entry_data_s data;
+    MMDB_lookup_result_s res;
+    if (!mmdbLookup(ip, false, gl, 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 queryRegionV6(string &ret, GeoIPNetmask& gl, const string &ip) override {
+    MMDB_entry_data_s data;
+    MMDB_lookup_result_s res;
+    if (!mmdbLookup(ip, true, gl, 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, GeoIPNetmask& gl, const string &ip) override {
+    MMDB_entry_data_s data;
+    MMDB_lookup_result_s res;
+    if (!mmdbLookup(ip, false, gl, 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", d_lang.c_str(), NULL) != MMDB_SUCCESS || !data.has_data))
+      return false;
+    ret = string(data.utf8_string, data.data_size);
+    return true;
+  }
+
+  bool queryCityV6(string &ret, GeoIPNetmask& gl, const string &ip) override {
+    MMDB_entry_data_s data;
+    MMDB_lookup_result_s res;
+    if (!mmdbLookup(ip, true, gl, 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", d_lang.c_str(), NULL) != MMDB_SUCCESS || !data.has_data))
+      return false;
+    ret = string(data.utf8_string, data.data_size);
+    return true;
+  }
+
+  ~GeoIPInterfaceMMDB() { MMDB_close(&d_s); };
+private:
+  MMDB_s d_s;
+  string d_lang;
+
+  bool mmdbLookup(const string &ip, bool v6, GeoIPNetmask& gl, 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)
+      L<<Logger::Warning<<"MMDB_lookup_string("<<ip<<") failed: "<<gai_strerror(gai_ec)<<endl;
+    else if (mmdb_ec != MMDB_SUCCESS)
+      L<<Logger::Warning<<"MMDB_lookup_string("<<ip<<") failed: "<<MMDB_strerror(mmdb_ec)<<endl;
+    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;
+  }
+};
+
+unique_ptr<GeoIPInterface> GeoIPInterface::makeMMDBInterface(const string &fname, const map<string, string>& opts) {
+  string mode = "";
+  string language = "en";
+  const auto &opt_mode = opts.find("mode");
+  if (opt_mode != opts.end())
+    mode = opt_mode->second;
+  const auto &opt_lang = opts.find("language");
+  if (opt_lang != opts.end())
+    language = opt_lang->second;
+  return unique_ptr<GeoIPInterface>(new GeoIPInterfaceMMDB(fname, mode, language));
+}
+
+#else
+
+unique_ptr<GeoIPInterface> GeoIPInterface::makeMMDBInterface(const string &fname, const map<string, string>& opts) {
+  throw PDNSException("libmaxminddb support not compiled in");
+}
+
+#endif
diff --git a/modules/geoipbackend/geoipinterface.cc b/modules/geoipbackend/geoipinterface.cc
index 31c6eac3e..0f0f8a240 100644
--- a/modules/geoipbackend/geoipinterface.cc
+++ b/modules/geoipbackend/geoipinterface.cc
@@ -39,7 +39,7 @@ unique_ptr<GeoIPInterface> GeoIPInterface::makeInterface(const string& dbStr) {
     filename = parts2[0];
     size_t pos = filename.find_last_of(".");
     if (pos != string::npos)
-      driver = driver.substr(pos+1);
+      driver = filename.substr(pos+1);
     else
       driver = "unknown";
   } else {
@@ -59,6 +59,8 @@ unique_ptr<GeoIPInterface> GeoIPInterface::makeInterface(const string& dbStr) {
 
   if (driver == "dat") {
      return makeDATInterface(filename, opts);
+  } else if (driver == "mmdb") {
+     return makeMMDBInterface(filename, opts);
   } else {
      throw PDNSException(string("Unsupported file type '") + driver + string("' (use type: prefix to force type)"));
   }
diff --git a/modules/geoipbackend/geoipinterface.hh b/modules/geoipbackend/geoipinterface.hh
index 20ebb05e8..97dd5ee48 100644
--- a/modules/geoipbackend/geoipinterface.hh
+++ b/modules/geoipbackend/geoipinterface.hh
@@ -53,6 +53,7 @@ public:
 
   static unique_ptr<GeoIPInterface> makeInterface(const string& dbStr);
 private:
+  static unique_ptr<GeoIPInterface> makeMMDBInterface(const string &fname, const map<string, string>& opts);
   static unique_ptr<GeoIPInterface> makeDATInterface(const string& fname, const map<string, string>& opts);
 };
 
-- 
2.40.0