From 9969d936f355e938555fba0fd60b6f6eb6470821 Mon Sep 17 00:00:00 2001
From: Aki Tuomi <cmouse@cmouse.fi>
Date: Sun, 1 May 2022 21:36:25 +0300
Subject: [PATCH] geoipbackend: Add zones_dir key for geoipbackend file

Allows loading zones from a directory.

Based on code by mason-chase <sid@moontius.com>
---
 configure.ac                         |  5 +++++
 m4/ax_cxx_fs.m4                      | 31 ++++++++++++++++++++++++++++
 modules/geoipbackend/geoipbackend.cc | 26 +++++++++++++++++++++++
 modules/geoipbackend/geoipbackend.hh |  1 +
 4 files changed, 63 insertions(+)
 create mode 100644 m4/ax_cxx_fs.m4

diff --git a/configure.ac b/configure.ac
index 66480b71aa..8f6fde4604 100644
--- a/configure.ac
+++ b/configure.ac
@@ -40,6 +40,7 @@ LT_INIT([disable-static dlopen])
 PDNS_CHECK_OS
 PTHREAD_SET_NAME
 AC_FUNC_STRERROR_R
+AX_CXX_CXXFS
 
 PDNS_WITH_LUA([mandatory])
 PDNS_CHECK_LUA_HPP
@@ -308,6 +309,10 @@ LDFLAGS="$RELRO_LDFLAGS $LDFLAGS"
 CFLAGS="$PIE_CFLAGS $CFLAGS"
 CXXFLAGS="$PIE_CFLAGS $CXXFLAGS"
 PROGRAM_LDFLAGS="$PIE_LDFLAGS $PROGRAM_LDFLAGS"
+AS_IF([test "$ax_cxx_cv_filesystem_lib" != "none"],
+ [PROGRAM_LDFLAGS="$PROGRAM_LDFLAGS -l$ax_cxx_cv_filesystem_lib"],
+ []
+)
 AC_SUBST([PROGRAM_LDFLAGS])
 
 PDNS_ENABLE_COVERAGE
diff --git a/m4/ax_cxx_fs.m4 b/m4/ax_cxx_fs.m4
new file mode 100644
index 0000000000..fca0d05764
--- /dev/null
+++ b/m4/ax_cxx_fs.m4
@@ -0,0 +1,31 @@
+AC_DEFUN([AX_CXX_CXXFS], [
+   AC_LANG_PUSH([C++])
+   old_LIBS="$LIBS"
+   dnl * Test first if it can be used without anything, then -lstdc++fs and -lc++fs
+   AC_CACHE_CHECK([for library with std::filesystem], [ax_cxx_cv_filesystem_lib], [
+      ax_cxx_cv_filesystem_lib=none
+      AC_LINK_IFELSE([AC_LANG_PROGRAM(
+        [[#include <iostream>
+          #include <filesystem>]],
+        [[std::filesystem::path path(".");
+          std::filesystem::status(path);]])],
+        [], [
+           LIBS="$LIBS -lstdc++fs"
+           AC_LINK_IFELSE([AC_LANG_PROGRAM(
+             [[#include <iostream>
+               #include <filesystem>]],
+             [[std::filesystem::path path(".");
+               std::filesystem::status(path);]])],
+             [ax_cxx_cv_filesystem_lib=stdc++fs], [
+               LIBS="$old_LIBS -lc++fs"
+               AC_LINK_IFELSE([AC_LANG_PROGRAM(
+                 [[#include <iostream>
+                   #include <filesystem>]],
+                 [[std::filesystem::path path(".");
+                   std::filesystem::status(path);]])],
+                 [ax_cxx_cv_filesystem_lib=c++fs], [AC_MSG_ERROR([Cannot find std::filesystem library])])
+      ])])
+      LIBS="$old_LIBS"
+   ])
+   AC_LANG_POP()
+])
diff --git a/modules/geoipbackend/geoipbackend.cc b/modules/geoipbackend/geoipbackend.cc
index de939804fe..ada8093b41 100644
--- a/modules/geoipbackend/geoipbackend.cc
+++ b/modules/geoipbackend/geoipbackend.cc
@@ -31,6 +31,7 @@
 #include <boost/algorithm/string/replace.hpp>
 #include <boost/format.hpp>
 #include <fstream>
+#include <filesystem>
 #include <yaml-cpp/yaml.h>
 
 ReadWriteLock GeoIPBackend::s_state_lock;
@@ -333,6 +334,28 @@ bool GeoIPBackend::loadDomain(const YAML::Node& domain, unsigned int id, GeoIPDo
   return true;
 }
 
+void GeoIPBackend::loadDomainsFromDirectory(const std::string& dir, vector<GeoIPDomain>& domains)
+{
+  vector<std::filesystem::path> paths;
+  for (const std::filesystem::path& p : std::filesystem::directory_iterator(std::filesystem::path(dir)))
+    if (std::filesystem::is_regular_file(p) && p.has_extension() && (p.extension() == ".yaml" || p.extension() == ".yml"))
+      paths.push_back(p);
+  std::sort(paths.begin(), paths.end());
+  for (const auto& p : paths) {
+    try {
+      GeoIPDomain dom;
+      const auto& zoneRoot = YAML::LoadFile(p.string());
+      // expect zone key
+      const auto& zone = zoneRoot["zone"];
+      if (loadDomain(zone, domains.size(), dom))
+        domains.push_back(dom);
+    }
+    catch (std::exception& ex) {
+      g_log << Logger::Warning << "Cannot load zone from " << p << ": " << ex.what() << endl;
+    }
+  }
+}
+
 void GeoIPBackend::initialize()
 {
   YAML::Node config;
@@ -378,6 +401,9 @@ void GeoIPBackend::initialize()
       tmp_domains.push_back(std::move(dom));
   }
 
+  if (YAML::Node domain_dir = config["zones_dir"])
+    loadDomainsFromDirectory(domain_dir.as<string>(), tmp_domains);
+
   s_domains.clear();
   std::swap(s_domains, tmp_domains);
 
diff --git a/modules/geoipbackend/geoipbackend.hh b/modules/geoipbackend/geoipbackend.hh
index 949481dfd2..c7e988a2f8 100644
--- a/modules/geoipbackend/geoipbackend.hh
+++ b/modules/geoipbackend/geoipbackend.hh
@@ -83,6 +83,7 @@ private:
   bool hasDNSSECkey(const DNSName& name);
   bool lookup_static(const GeoIPDomain& dom, const DNSName& search, const QType& qtype, const DNSName& qdomain, const Netmask& addr, GeoIPNetmask& gl);
   bool loadDomain(const YAML::Node& domain, unsigned int id, GeoIPDomain& dom);
+  void loadDomainsFromDirectory(const std::string& dir, vector<GeoIPDomain>& domains);
   vector<DNSResourceRecord> d_result;
   vector<GeoIPInterface> d_files;
   std::vector<std::string> d_global_mapping_lookup_formats;
-- 
2.39.2