From e3a3fa098539e28beec5e86c324221f220977c1f Mon Sep 17 00:00:00 2001 From: Chris McFarlen Date: Fri, 27 Mar 2026 11:35:30 -0500 Subject: [PATCH] Support per-remap geo DB handles in header_rewrite Replace the single global geo DB handle with a cache of handles keyed by path, allowing different remap rules to use different MMDB files via --geo-db-path. The handle is threaded through RulesConfig -> Resources -> ConditionGeo::get_geo_*() methods. The cache deduplicates when multiple rules share the same path. Co-Authored-By: Claude Opus 4.6 (cherry picked from commit 8713ad495b4113853216885e31b927b1ccd1999c) --- plugins/header_rewrite/conditions.cc | 10 +-- plugins/header_rewrite/conditions.h | 4 +- plugins/header_rewrite/conditions_geo.h | 12 ++-- .../header_rewrite/conditions_geo_geoip.cc | 71 ++++++++++++------- .../header_rewrite/conditions_geo_maxmind.cc | 65 ++++++++++------- plugins/header_rewrite/header_rewrite.cc | 40 +++++++---- plugins/header_rewrite/resources.h | 1 + 7 files changed, 126 insertions(+), 77 deletions(-) diff --git a/plugins/header_rewrite/conditions.cc b/plugins/header_rewrite/conditions.cc index faf4a19a52a..35497aaeeb9 100644 --- a/plugins/header_rewrite/conditions.cc +++ b/plugins/header_rewrite/conditions.cc @@ -831,14 +831,14 @@ ConditionNow::eval(const Resources &res) } std::string -ConditionGeo::get_geo_string(const sockaddr * /* addr ATS_UNUSED */) const +ConditionGeo::get_geo_string(const sockaddr * /* addr ATS_UNUSED */, void * /* geo_handle ATS_UNUSED */) const { TSError("[%s] No Geo library available!", PLUGIN_NAME); return ""; } int64_t -ConditionGeo::get_geo_int(const sockaddr * /* addr ATS_UNUSED */) const +ConditionGeo::get_geo_int(const sockaddr * /* addr ATS_UNUSED */, void * /* geo_handle ATS_UNUSED */) const { TSError("[%s] No Geo library available!", PLUGIN_NAME); return 0; @@ -891,9 +891,9 @@ void ConditionGeo::append_value(std::string &s, const Resources &res) { if (is_int_type()) { - s += std::to_string(get_geo_int(getClientAddr(res.state.txnp, _txn_private_slot))); + s += std::to_string(get_geo_int(getClientAddr(res.state.txnp, _txn_private_slot), res.geo_handle)); } else { - s += get_geo_string(getClientAddr(res.state.txnp, _txn_private_slot)); + s += get_geo_string(getClientAddr(res.state.txnp, _txn_private_slot), res.geo_handle); } Dbg(pi_dbg_ctl, "Appending GEO() to evaluation value -> %s", s.c_str()); } @@ -905,7 +905,7 @@ ConditionGeo::eval(const Resources &res) Dbg(pi_dbg_ctl, "Evaluating GEO()"); if (is_int_type()) { - int64_t geo = get_geo_int(getClientAddr(res.state.txnp, _txn_private_slot)); + int64_t geo = get_geo_int(getClientAddr(res.state.txnp, _txn_private_slot), res.geo_handle); ret = static_cast *>(_matcher.get())->test(geo, res); } else { diff --git a/plugins/header_rewrite/conditions.h b/plugins/header_rewrite/conditions.h index 4393ecbab09..c393e602663 100644 --- a/plugins/header_rewrite/conditions.h +++ b/plugins/header_rewrite/conditions.h @@ -477,8 +477,8 @@ class ConditionGeo : public Condition } private: - virtual int64_t get_geo_int(const sockaddr *addr) const; - virtual std::string get_geo_string(const sockaddr *addr) const; + virtual int64_t get_geo_int(const sockaddr *addr, void *geo_handle) const; + virtual std::string get_geo_string(const sockaddr *addr, void *geo_handle) const; protected: bool diff --git a/plugins/header_rewrite/conditions_geo.h b/plugins/header_rewrite/conditions_geo.h index 6894ca3a590..224cff554b6 100644 --- a/plugins/header_rewrite/conditions_geo.h +++ b/plugins/header_rewrite/conditions_geo.h @@ -27,10 +27,10 @@ class MMConditionGeo : public ConditionGeo MMConditionGeo() {} virtual ~MMConditionGeo() {} - static void initLibrary(const std::string &path); + static void *initLibrary(const std::string &path); - virtual int64_t get_geo_int(const sockaddr *addr) const override; - virtual std::string get_geo_string(const sockaddr *addr) const override; + int64_t get_geo_int(const sockaddr *addr, void *geo_handle) const override; + std::string get_geo_string(const sockaddr *addr, void *geo_handle) const override; }; class GeoIPConditionGeo : public ConditionGeo @@ -39,8 +39,8 @@ class GeoIPConditionGeo : public ConditionGeo GeoIPConditionGeo() {} virtual ~GeoIPConditionGeo() {} - static void initLibrary(const std::string &path); + static void *initLibrary(const std::string &path); - virtual int64_t get_geo_int(const sockaddr *addr) const override; - virtual std::string get_geo_string(const sockaddr *addr) const override; + int64_t get_geo_int(const sockaddr *addr, void *geo_handle) const override; + std::string get_geo_string(const sockaddr *addr, void *geo_handle) const override; }; diff --git a/plugins/header_rewrite/conditions_geo_geoip.cc b/plugins/header_rewrite/conditions_geo_geoip.cc index 1bde085cb41..5e2ffe0756b 100644 --- a/plugins/header_rewrite/conditions_geo_geoip.cc +++ b/plugins/header_rewrite/conditions_geo_geoip.cc @@ -24,6 +24,7 @@ #include #include #include +#include #include "ts/ts.h" @@ -31,49 +32,65 @@ #include -GeoIP *gGeoIP[NUM_DB_TYPES]; +struct GeoIPHandleSet { + GeoIP *dbs[NUM_DB_TYPES] = {}; +}; -void +static std::mutex gGeoIPCacheMutex; +static GeoIPHandleSet *gGeoIPHandleSet = nullptr; + +void * GeoIPConditionGeo::initLibrary(const std::string &) { + std::lock_guard lock(gGeoIPCacheMutex); + + if (gGeoIPHandleSet != nullptr) { + return gGeoIPHandleSet; + } + + gGeoIPHandleSet = new GeoIPHandleSet; + GeoIPDBTypes dbs[] = {GEOIP_COUNTRY_EDITION, GEOIP_COUNTRY_EDITION_V6, GEOIP_ASNUM_EDITION, GEOIP_ASNUM_EDITION_V6}; for (auto &db : dbs) { - if (!gGeoIP[db] && GeoIP_db_avail(db)) { - // GEOIP_STANDARD seems to break threaded apps... - gGeoIP[db] = GeoIP_open_type(db, GEOIP_MMAP_CACHE); + if (!gGeoIPHandleSet->dbs[db] && GeoIP_db_avail(db)) { + gGeoIPHandleSet->dbs[db] = GeoIP_open_type(db, GEOIP_MMAP_CACHE); - char *db_info = GeoIP_database_info(gGeoIP[db]); + char *db_info = GeoIP_database_info(gGeoIPHandleSet->dbs[db]); Dbg(pi_dbg_ctl, "initialized GeoIP-DB[%d] %s", db, db_info); free(db_info); } } + + return gGeoIPHandleSet; } std::string -GeoIPConditionGeo::get_geo_string(const sockaddr *addr) const +GeoIPConditionGeo::get_geo_string(const sockaddr *addr, void *geo_handle) const { std::string ret = "(unknown)"; int v = 4; - if (addr) { + auto *handle = static_cast(geo_handle); + + if (addr && handle) { switch (_geo_qual) { // Country database case GEO_QUAL_COUNTRY: switch (addr->sa_family) { case AF_INET: - if (gGeoIP[GEOIP_COUNTRY_EDITION]) { + if (handle->dbs[GEOIP_COUNTRY_EDITION]) { uint32_t ip = ntohl(reinterpret_cast(addr)->sin_addr.s_addr); - ret = GeoIP_country_code_by_ipnum(gGeoIP[GEOIP_COUNTRY_EDITION], ip); + ret = GeoIP_country_code_by_ipnum(handle->dbs[GEOIP_COUNTRY_EDITION], ip); } break; case AF_INET6: { - if (gGeoIP[GEOIP_COUNTRY_EDITION_V6]) { + if (handle->dbs[GEOIP_COUNTRY_EDITION_V6]) { geoipv6_t ip = reinterpret_cast(addr)->sin6_addr; v = 6; - ret = GeoIP_country_code_by_ipnum_v6(gGeoIP[GEOIP_COUNTRY_EDITION_V6], ip); + ret = GeoIP_country_code_by_ipnum_v6(handle->dbs[GEOIP_COUNTRY_EDITION_V6], ip); } } break; default: @@ -86,18 +103,18 @@ GeoIPConditionGeo::get_geo_string(const sockaddr *addr) const case GEO_QUAL_ASN_NAME: switch (addr->sa_family) { case AF_INET: - if (gGeoIP[GEOIP_ASNUM_EDITION]) { + if (handle->dbs[GEOIP_ASNUM_EDITION]) { uint32_t ip = ntohl(reinterpret_cast(addr)->sin_addr.s_addr); - ret = GeoIP_name_by_ipnum(gGeoIP[GEOIP_ASNUM_EDITION], ip); + ret = GeoIP_name_by_ipnum(handle->dbs[GEOIP_ASNUM_EDITION], ip); } break; case AF_INET6: { - if (gGeoIP[GEOIP_ASNUM_EDITION_V6]) { + if (handle->dbs[GEOIP_ASNUM_EDITION_V6]) { geoipv6_t ip = reinterpret_cast(addr)->sin6_addr; v = 6; - ret = GeoIP_name_by_ipnum_v6(gGeoIP[GEOIP_ASNUM_EDITION_V6], ip); + ret = GeoIP_name_by_ipnum_v6(handle->dbs[GEOIP_ASNUM_EDITION_V6], ip); } } break; default: @@ -114,12 +131,14 @@ GeoIPConditionGeo::get_geo_string(const sockaddr *addr) const } int64_t -GeoIPConditionGeo::get_geo_int(const sockaddr *addr) const +GeoIPConditionGeo::get_geo_int(const sockaddr *addr, void *geo_handle) const { int64_t ret = -1; int v = 4; - if (!addr) { + auto *handle = static_cast(geo_handle); + + if (!addr || !handle) { return 0; } @@ -128,18 +147,18 @@ GeoIPConditionGeo::get_geo_int(const sockaddr *addr) const case GEO_QUAL_COUNTRY_ISO: switch (addr->sa_family) { case AF_INET: - if (gGeoIP[GEOIP_COUNTRY_EDITION]) { + if (handle->dbs[GEOIP_COUNTRY_EDITION]) { uint32_t ip = ntohl(reinterpret_cast(addr)->sin_addr.s_addr); - ret = GeoIP_id_by_ipnum(gGeoIP[GEOIP_COUNTRY_EDITION], ip); + ret = GeoIP_id_by_ipnum(handle->dbs[GEOIP_COUNTRY_EDITION], ip); } break; case AF_INET6: { - if (gGeoIP[GEOIP_COUNTRY_EDITION_V6]) { + if (handle->dbs[GEOIP_COUNTRY_EDITION_V6]) { geoipv6_t ip = reinterpret_cast(addr)->sin6_addr; v = 6; - ret = GeoIP_id_by_ipnum_v6(gGeoIP[GEOIP_COUNTRY_EDITION_V6], ip); + ret = GeoIP_id_by_ipnum_v6(handle->dbs[GEOIP_COUNTRY_EDITION_V6], ip); } } break; default: @@ -153,18 +172,18 @@ GeoIPConditionGeo::get_geo_int(const sockaddr *addr) const switch (addr->sa_family) { case AF_INET: - if (gGeoIP[GEOIP_ASNUM_EDITION]) { + if (handle->dbs[GEOIP_ASNUM_EDITION]) { uint32_t ip = ntohl(reinterpret_cast(addr)->sin_addr.s_addr); - asn_name = GeoIP_name_by_ipnum(gGeoIP[GEOIP_ASNUM_EDITION], ip); + asn_name = GeoIP_name_by_ipnum(handle->dbs[GEOIP_ASNUM_EDITION], ip); } break; case AF_INET6: - if (gGeoIP[GEOIP_ASNUM_EDITION_V6]) { + if (handle->dbs[GEOIP_ASNUM_EDITION_V6]) { geoipv6_t ip = reinterpret_cast(addr)->sin6_addr; v = 6; - asn_name = GeoIP_name_by_ipnum_v6(gGeoIP[GEOIP_ASNUM_EDITION_V6], ip); + asn_name = GeoIP_name_by_ipnum_v6(handle->dbs[GEOIP_ASNUM_EDITION_V6], ip); } break; } diff --git a/plugins/header_rewrite/conditions_geo_maxmind.cc b/plugins/header_rewrite/conditions_geo_maxmind.cc index f9168a5a23a..81e34ed1d9f 100644 --- a/plugins/header_rewrite/conditions_geo_maxmind.cc +++ b/plugins/header_rewrite/conditions_geo_maxmind.cc @@ -23,6 +23,9 @@ #include #include +#include +#include +#include #include "ts/ts.h" @@ -30,10 +33,15 @@ #include -MMDB_s *gMaxMindDB = nullptr; - enum class MmdbSchema { NESTED, FLAT }; -static MmdbSchema gMmdbSchema = MmdbSchema::NESTED; + +struct MmdbHandle { + MMDB_s db; + MmdbSchema schema = MmdbSchema::NESTED; +}; + +static std::map gMmdbCache; +static std::mutex gMmdbCacheMutex; // Detect whether the MMDB uses nested (GeoLite2) or flat (vendor) field layout // by probing for the nested country path on a lookup result. @@ -57,56 +65,63 @@ detect_schema(MMDB_entry_s *entry) static const char *probe_ips[] = {"8.8.8.8", "1.1.1.1", "128.0.0.1"}; -void +void * MMConditionGeo::initLibrary(const std::string &path) { if (path.empty()) { Dbg(pi_dbg_ctl, "Empty MaxMind db path specified. Not initializing!"); - return; + return nullptr; } - if (gMaxMindDB != nullptr) { - Dbg(pi_dbg_ctl, "Maxmind library already initialized"); - return; + std::lock_guard lock(gMmdbCacheMutex); + + auto it = gMmdbCache.find(path); + if (it != gMmdbCache.end()) { + Dbg(pi_dbg_ctl, "Maxmind library already initialized for %s", path.c_str()); + return it->second; } - gMaxMindDB = new MMDB_s; + auto *handle = new MmdbHandle; + int status = MMDB_open(path.c_str(), MMDB_MODE_MMAP, &handle->db); - int status = MMDB_open(path.c_str(), MMDB_MODE_MMAP, gMaxMindDB); if (MMDB_SUCCESS != status) { Dbg(pi_dbg_ctl, "Cannot open %s - %s", path.c_str(), MMDB_strerror(status)); - delete gMaxMindDB; - gMaxMindDB = nullptr; - return; + delete handle; + return nullptr; } // Probe the database schema at load time so we know which field paths to // use for country lookups. Try a few well-known IPs until one hits. for (auto *ip : probe_ips) { int gai_error, mmdb_error; - MMDB_lookup_result_s result = MMDB_lookup_string(gMaxMindDB, ip, &gai_error, &mmdb_error); + MMDB_lookup_result_s result = MMDB_lookup_string(&handle->db, ip, &gai_error, &mmdb_error); if (gai_error == 0 && MMDB_SUCCESS == mmdb_error && result.found_entry) { - gMmdbSchema = detect_schema(&result.entry); - Dbg(pi_dbg_ctl, "Loaded %s (schema: %s)", path.c_str(), gMmdbSchema == MmdbSchema::FLAT ? "flat" : "nested"); - return; + handle->schema = detect_schema(&result.entry); + Dbg(pi_dbg_ctl, "Loaded %s (schema: %s)", path.c_str(), handle->schema == MmdbSchema::FLAT ? "flat" : "nested"); + gMmdbCache[path] = handle; + return handle; } } Dbg(pi_dbg_ctl, "Loaded %s (schema: defaulting to nested, no probe IPs matched)", path.c_str()); + gMmdbCache[path] = handle; + return handle; } std::string -MMConditionGeo::get_geo_string(const sockaddr *addr) const +MMConditionGeo::get_geo_string(const sockaddr *addr, void *geo_handle) const { std::string ret = "(unknown)"; int mmdb_error; - if (gMaxMindDB == nullptr) { + auto *handle = static_cast(geo_handle); + + if (handle == nullptr) { Dbg(pi_dbg_ctl, "MaxMind not initialized; using default value"); return ret; } - MMDB_lookup_result_s result = MMDB_lookup_sockaddr(gMaxMindDB, addr, &mmdb_error); + MMDB_lookup_result_s result = MMDB_lookup_sockaddr(&handle->db, addr, &mmdb_error); if (MMDB_SUCCESS != mmdb_error) { Dbg(pi_dbg_ctl, "Error during sockaddr lookup: %s", MMDB_strerror(mmdb_error)); @@ -123,7 +138,7 @@ MMConditionGeo::get_geo_string(const sockaddr *addr) const switch (_geo_qual) { case GEO_QUAL_COUNTRY: - if (gMmdbSchema == MmdbSchema::FLAT) { + if (handle->schema == MmdbSchema::FLAT) { status = MMDB_get_value(&result.entry, &entry_data, "country_code", NULL); } else { status = MMDB_get_value(&result.entry, &entry_data, "country", "iso_code", NULL); @@ -150,17 +165,19 @@ MMConditionGeo::get_geo_string(const sockaddr *addr) const } int64_t -MMConditionGeo::get_geo_int(const sockaddr *addr) const +MMConditionGeo::get_geo_int(const sockaddr *addr, void *geo_handle) const { int64_t ret = -1; int mmdb_error; - if (gMaxMindDB == nullptr) { + auto *handle = static_cast(geo_handle); + + if (handle == nullptr) { Dbg(pi_dbg_ctl, "MaxMind not initialized; using default value"); return ret; } - MMDB_lookup_result_s result = MMDB_lookup_sockaddr(gMaxMindDB, addr, &mmdb_error); + MMDB_lookup_result_s result = MMDB_lookup_sockaddr(&handle->db, addr, &mmdb_error); if (MMDB_SUCCESS != mmdb_error) { Dbg(pi_dbg_ctl, "Error during sockaddr lookup: %s", MMDB_strerror(mmdb_error)); diff --git a/plugins/header_rewrite/header_rewrite.cc b/plugins/header_rewrite/header_rewrite.cc index a2601be4c59..16f8b0c6170 100644 --- a/plugins/header_rewrite/header_rewrite.cc +++ b/plugins/header_rewrite/header_rewrite.cc @@ -41,7 +41,6 @@ // Debugs namespace header_rewrite_ns { -std::once_flag initGeoLibs; std::once_flag initPlugin; PluginFactory plugin_factory; } // namespace header_rewrite_ns @@ -52,19 +51,21 @@ initPluginFactory() header_rewrite_ns::plugin_factory.setRuntimeDir(RecConfigReadRuntimeDir()).addSearchDir(RecConfigReadPluginDir()); } -static void +static void * initGeoLibraries(const std::string &dbPath) { if (dbPath.empty()) { - return; + return nullptr; } Dbg(pi_dbg_ctl, "Loading geo db %s", dbPath.c_str()); #if TS_USE_HRW_GEOIP - GeoIPConditionGeo::initLibrary(dbPath); + return GeoIPConditionGeo::initLibrary(dbPath); #elif TS_USE_HRW_MAXMINDDB - MMConditionGeo::initLibrary(dbPath); + return MMConditionGeo::initLibrary(dbPath); +#else + return nullptr; #endif } @@ -76,7 +77,8 @@ static int cont_rewrite_headers(TSCont, TSEvent, void *); class RulesConfig { public: - RulesConfig(int timezone, int inboundIpSource) : _timezone(timezone), _inboundIpSource(inboundIpSource) + RulesConfig(int timezone, int inboundIpSource, void *geo_handle = nullptr) + : _timezone(timezone), _inboundIpSource(inboundIpSource), _geo_handle(geo_handle) { Dbg(dbg_ctl, "RulesConfig CTOR"); _cont = TSContCreate(cont_rewrite_headers, nullptr); @@ -119,6 +121,12 @@ class RulesConfig return _inboundIpSource; } + [[nodiscard]] void * + geo_handle() const + { + return _geo_handle; + } + bool parse_config(const std::string &fname, TSHttpHookID default_hook, char *from_url = nullptr, char *to_url = nullptr); private: @@ -130,6 +138,8 @@ class RulesConfig int _timezone = 0; int _inboundIpSource = 0; + + void *_geo_handle = nullptr; }; void @@ -509,6 +519,8 @@ cont_rewrite_headers(TSCont contp, TSEvent event, void *edata) RuleSet *rule = conf->rule(hook); Resources res(txnp, contp); + res.geo_handle = conf->geo_handle(); + // Get the resources necessary to process this event res.gather(conf->resid(hook), hook); @@ -602,12 +614,12 @@ TSPluginInit(int argc, const char *argv[]) Dbg(pi_dbg_ctl, "Global geo db %s", geoDBpath.c_str()); - std::call_once(initGeoLibs, [&geoDBpath]() { initGeoLibraries(geoDBpath); }); + void *geo_handle = initGeoLibraries(geoDBpath); std::call_once(initPlugin, initPluginFactory); // Parse the global config file(s). All rules are just appended // to the "global" Rules configuration. - auto *conf = new RulesConfig(timezone, inboundIpSource); + auto *conf = new RulesConfig(timezone, inboundIpSource, geo_handle); bool got_config = false; for (int i = optind; i < argc; ++i) { @@ -711,21 +723,19 @@ TSRemapNewInstance(int argc, char *argv[], void **ih, char * /* errbuf ATS_UNUSE } } + void *geo_handle = nullptr; + if (!geoDBpath.empty()) { if (!geoDBpath.starts_with('/')) { geoDBpath = std::string(TSConfigDirGet()) + '/' + geoDBpath; } Dbg(pi_dbg_ctl, "Remap geo db %s", geoDBpath.c_str()); - - // This MUST be called only if the geoDBpath is set. If called without a geoDBPath (i.e. outside of this if) then - // NO hrw remap rule can load a mmdb file. - // The call_once applies to every remap instance as its a plugin global - std::call_once(initGeoLibs, [&geoDBpath]() { initGeoLibraries(geoDBpath); }); + geo_handle = initGeoLibraries(geoDBpath); } std::call_once(initPlugin, initPluginFactory); - auto *conf = new RulesConfig(timezone, inboundIpSource); + auto *conf = new RulesConfig(timezone, inboundIpSource, geo_handle); for (int i = optind; i < argc; ++i) { Dbg(pi_dbg_ctl, "Loading remap configuration file %s", argv[i]); @@ -781,6 +791,8 @@ TSRemapDoRemap(void *ih, TSHttpTxn rh, TSRemapRequestInfo *rri) RuleSet *rule = conf->rule(TS_REMAP_PSEUDO_HOOK); Resources res(rh, rri); + res.geo_handle = conf->geo_handle(); + if (rule) { res.gather(conf->resid(TS_REMAP_PSEUDO_HOOK), TS_REMAP_PSEUDO_HOOK); diff --git a/plugins/header_rewrite/resources.h b/plugins/header_rewrite/resources.h index 2fb45b08c44..80d1268106c 100644 --- a/plugins/header_rewrite/resources.h +++ b/plugins/header_rewrite/resources.h @@ -127,6 +127,7 @@ class Resources TransactionState state; // Without cripts, txnp / ssnp goes here #endif TSHttpStatus resp_status = TS_HTTP_STATUS_NONE; + void *geo_handle = nullptr; struct LifetimeExtension { std::string subject_storage;