Skip to content

Commit a26a2d8

Browse files
cmcfarlenclaude
andauthored
Support per-remap geo DB handles in header_rewrite (#13042)
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. (cherry picked from commit 8713ad4) Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
1 parent a65725d commit a26a2d8

7 files changed

Lines changed: 126 additions & 77 deletions

File tree

plugins/header_rewrite/conditions.cc

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -831,14 +831,14 @@ ConditionNow::eval(const Resources &res)
831831
}
832832

833833
std::string
834-
ConditionGeo::get_geo_string(const sockaddr * /* addr ATS_UNUSED */) const
834+
ConditionGeo::get_geo_string(const sockaddr * /* addr ATS_UNUSED */, void * /* geo_handle ATS_UNUSED */) const
835835
{
836836
TSError("[%s] No Geo library available!", PLUGIN_NAME);
837837
return "";
838838
}
839839

840840
int64_t
841-
ConditionGeo::get_geo_int(const sockaddr * /* addr ATS_UNUSED */) const
841+
ConditionGeo::get_geo_int(const sockaddr * /* addr ATS_UNUSED */, void * /* geo_handle ATS_UNUSED */) const
842842
{
843843
TSError("[%s] No Geo library available!", PLUGIN_NAME);
844844
return 0;
@@ -891,9 +891,9 @@ void
891891
ConditionGeo::append_value(std::string &s, const Resources &res)
892892
{
893893
if (is_int_type()) {
894-
s += std::to_string(get_geo_int(getClientAddr(res.state.txnp, _txn_private_slot)));
894+
s += std::to_string(get_geo_int(getClientAddr(res.state.txnp, _txn_private_slot), res.geo_handle));
895895
} else {
896-
s += get_geo_string(getClientAddr(res.state.txnp, _txn_private_slot));
896+
s += get_geo_string(getClientAddr(res.state.txnp, _txn_private_slot), res.geo_handle);
897897
}
898898
Dbg(pi_dbg_ctl, "Appending GEO() to evaluation value -> %s", s.c_str());
899899
}
@@ -905,7 +905,7 @@ ConditionGeo::eval(const Resources &res)
905905

906906
Dbg(pi_dbg_ctl, "Evaluating GEO()");
907907
if (is_int_type()) {
908-
int64_t geo = get_geo_int(getClientAddr(res.state.txnp, _txn_private_slot));
908+
int64_t geo = get_geo_int(getClientAddr(res.state.txnp, _txn_private_slot), res.geo_handle);
909909

910910
ret = static_cast<const Matchers<int64_t> *>(_matcher.get())->test(geo, res);
911911
} else {

plugins/header_rewrite/conditions.h

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -477,8 +477,8 @@ class ConditionGeo : public Condition
477477
}
478478

479479
private:
480-
virtual int64_t get_geo_int(const sockaddr *addr) const;
481-
virtual std::string get_geo_string(const sockaddr *addr) const;
480+
virtual int64_t get_geo_int(const sockaddr *addr, void *geo_handle) const;
481+
virtual std::string get_geo_string(const sockaddr *addr, void *geo_handle) const;
482482

483483
protected:
484484
bool

plugins/header_rewrite/conditions_geo.h

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -27,10 +27,10 @@ class MMConditionGeo : public ConditionGeo
2727
MMConditionGeo() {}
2828
virtual ~MMConditionGeo() {}
2929

30-
static void initLibrary(const std::string &path);
30+
static void *initLibrary(const std::string &path);
3131

32-
virtual int64_t get_geo_int(const sockaddr *addr) const override;
33-
virtual std::string get_geo_string(const sockaddr *addr) const override;
32+
int64_t get_geo_int(const sockaddr *addr, void *geo_handle) const override;
33+
std::string get_geo_string(const sockaddr *addr, void *geo_handle) const override;
3434
};
3535

3636
class GeoIPConditionGeo : public ConditionGeo
@@ -39,8 +39,8 @@ class GeoIPConditionGeo : public ConditionGeo
3939
GeoIPConditionGeo() {}
4040
virtual ~GeoIPConditionGeo() {}
4141

42-
static void initLibrary(const std::string &path);
42+
static void *initLibrary(const std::string &path);
4343

44-
virtual int64_t get_geo_int(const sockaddr *addr) const override;
45-
virtual std::string get_geo_string(const sockaddr *addr) const override;
44+
int64_t get_geo_int(const sockaddr *addr, void *geo_handle) const override;
45+
std::string get_geo_string(const sockaddr *addr, void *geo_handle) const override;
4646
};

plugins/header_rewrite/conditions_geo_geoip.cc

Lines changed: 45 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -24,56 +24,73 @@
2424
#include <unistd.h>
2525
#include <arpa/inet.h>
2626
#include <cctype>
27+
#include <mutex>
2728

2829
#include "ts/ts.h"
2930

3031
#include "conditions_geo.h"
3132

3233
#include <GeoIP.h>
3334

34-
GeoIP *gGeoIP[NUM_DB_TYPES];
35+
struct GeoIPHandleSet {
36+
GeoIP *dbs[NUM_DB_TYPES] = {};
37+
};
3538

36-
void
39+
static std::mutex gGeoIPCacheMutex;
40+
static GeoIPHandleSet *gGeoIPHandleSet = nullptr;
41+
42+
void *
3743
GeoIPConditionGeo::initLibrary(const std::string &)
3844
{
45+
std::lock_guard<std::mutex> lock(gGeoIPCacheMutex);
46+
47+
if (gGeoIPHandleSet != nullptr) {
48+
return gGeoIPHandleSet;
49+
}
50+
51+
gGeoIPHandleSet = new GeoIPHandleSet;
52+
3953
GeoIPDBTypes dbs[] = {GEOIP_COUNTRY_EDITION, GEOIP_COUNTRY_EDITION_V6, GEOIP_ASNUM_EDITION, GEOIP_ASNUM_EDITION_V6};
4054

4155
for (auto &db : dbs) {
42-
if (!gGeoIP[db] && GeoIP_db_avail(db)) {
43-
// GEOIP_STANDARD seems to break threaded apps...
44-
gGeoIP[db] = GeoIP_open_type(db, GEOIP_MMAP_CACHE);
56+
if (!gGeoIPHandleSet->dbs[db] && GeoIP_db_avail(db)) {
57+
gGeoIPHandleSet->dbs[db] = GeoIP_open_type(db, GEOIP_MMAP_CACHE);
4558

46-
char *db_info = GeoIP_database_info(gGeoIP[db]);
59+
char *db_info = GeoIP_database_info(gGeoIPHandleSet->dbs[db]);
4760
Dbg(pi_dbg_ctl, "initialized GeoIP-DB[%d] %s", db, db_info);
4861
free(db_info);
4962
}
5063
}
64+
65+
return gGeoIPHandleSet;
5166
}
5267

5368
std::string
54-
GeoIPConditionGeo::get_geo_string(const sockaddr *addr) const
69+
GeoIPConditionGeo::get_geo_string(const sockaddr *addr, void *geo_handle) const
5570
{
5671
std::string ret = "(unknown)";
5772
int v = 4;
5873

59-
if (addr) {
74+
auto *handle = static_cast<GeoIPHandleSet *>(geo_handle);
75+
76+
if (addr && handle) {
6077
switch (_geo_qual) {
6178
// Country database
6279
case GEO_QUAL_COUNTRY:
6380
switch (addr->sa_family) {
6481
case AF_INET:
65-
if (gGeoIP[GEOIP_COUNTRY_EDITION]) {
82+
if (handle->dbs[GEOIP_COUNTRY_EDITION]) {
6683
uint32_t ip = ntohl(reinterpret_cast<const struct sockaddr_in *>(addr)->sin_addr.s_addr);
6784

68-
ret = GeoIP_country_code_by_ipnum(gGeoIP[GEOIP_COUNTRY_EDITION], ip);
85+
ret = GeoIP_country_code_by_ipnum(handle->dbs[GEOIP_COUNTRY_EDITION], ip);
6986
}
7087
break;
7188
case AF_INET6: {
72-
if (gGeoIP[GEOIP_COUNTRY_EDITION_V6]) {
89+
if (handle->dbs[GEOIP_COUNTRY_EDITION_V6]) {
7390
geoipv6_t ip = reinterpret_cast<const struct sockaddr_in6 *>(addr)->sin6_addr;
7491

7592
v = 6;
76-
ret = GeoIP_country_code_by_ipnum_v6(gGeoIP[GEOIP_COUNTRY_EDITION_V6], ip);
93+
ret = GeoIP_country_code_by_ipnum_v6(handle->dbs[GEOIP_COUNTRY_EDITION_V6], ip);
7794
}
7895
} break;
7996
default:
@@ -86,18 +103,18 @@ GeoIPConditionGeo::get_geo_string(const sockaddr *addr) const
86103
case GEO_QUAL_ASN_NAME:
87104
switch (addr->sa_family) {
88105
case AF_INET:
89-
if (gGeoIP[GEOIP_ASNUM_EDITION]) {
106+
if (handle->dbs[GEOIP_ASNUM_EDITION]) {
90107
uint32_t ip = ntohl(reinterpret_cast<const struct sockaddr_in *>(addr)->sin_addr.s_addr);
91108

92-
ret = GeoIP_name_by_ipnum(gGeoIP[GEOIP_ASNUM_EDITION], ip);
109+
ret = GeoIP_name_by_ipnum(handle->dbs[GEOIP_ASNUM_EDITION], ip);
93110
}
94111
break;
95112
case AF_INET6: {
96-
if (gGeoIP[GEOIP_ASNUM_EDITION_V6]) {
113+
if (handle->dbs[GEOIP_ASNUM_EDITION_V6]) {
97114
geoipv6_t ip = reinterpret_cast<const struct sockaddr_in6 *>(addr)->sin6_addr;
98115

99116
v = 6;
100-
ret = GeoIP_name_by_ipnum_v6(gGeoIP[GEOIP_ASNUM_EDITION_V6], ip);
117+
ret = GeoIP_name_by_ipnum_v6(handle->dbs[GEOIP_ASNUM_EDITION_V6], ip);
101118
}
102119
} break;
103120
default:
@@ -114,12 +131,14 @@ GeoIPConditionGeo::get_geo_string(const sockaddr *addr) const
114131
}
115132

116133
int64_t
117-
GeoIPConditionGeo::get_geo_int(const sockaddr *addr) const
134+
GeoIPConditionGeo::get_geo_int(const sockaddr *addr, void *geo_handle) const
118135
{
119136
int64_t ret = -1;
120137
int v = 4;
121138

122-
if (!addr) {
139+
auto *handle = static_cast<GeoIPHandleSet *>(geo_handle);
140+
141+
if (!addr || !handle) {
123142
return 0;
124143
}
125144

@@ -128,18 +147,18 @@ GeoIPConditionGeo::get_geo_int(const sockaddr *addr) const
128147
case GEO_QUAL_COUNTRY_ISO:
129148
switch (addr->sa_family) {
130149
case AF_INET:
131-
if (gGeoIP[GEOIP_COUNTRY_EDITION]) {
150+
if (handle->dbs[GEOIP_COUNTRY_EDITION]) {
132151
uint32_t ip = ntohl(reinterpret_cast<const struct sockaddr_in *>(addr)->sin_addr.s_addr);
133152

134-
ret = GeoIP_id_by_ipnum(gGeoIP[GEOIP_COUNTRY_EDITION], ip);
153+
ret = GeoIP_id_by_ipnum(handle->dbs[GEOIP_COUNTRY_EDITION], ip);
135154
}
136155
break;
137156
case AF_INET6: {
138-
if (gGeoIP[GEOIP_COUNTRY_EDITION_V6]) {
157+
if (handle->dbs[GEOIP_COUNTRY_EDITION_V6]) {
139158
geoipv6_t ip = reinterpret_cast<const struct sockaddr_in6 *>(addr)->sin6_addr;
140159

141160
v = 6;
142-
ret = GeoIP_id_by_ipnum_v6(gGeoIP[GEOIP_COUNTRY_EDITION_V6], ip);
161+
ret = GeoIP_id_by_ipnum_v6(handle->dbs[GEOIP_COUNTRY_EDITION_V6], ip);
143162
}
144163
} break;
145164
default:
@@ -153,18 +172,18 @@ GeoIPConditionGeo::get_geo_int(const sockaddr *addr) const
153172

154173
switch (addr->sa_family) {
155174
case AF_INET:
156-
if (gGeoIP[GEOIP_ASNUM_EDITION]) {
175+
if (handle->dbs[GEOIP_ASNUM_EDITION]) {
157176
uint32_t ip = ntohl(reinterpret_cast<const struct sockaddr_in *>(addr)->sin_addr.s_addr);
158177

159-
asn_name = GeoIP_name_by_ipnum(gGeoIP[GEOIP_ASNUM_EDITION], ip);
178+
asn_name = GeoIP_name_by_ipnum(handle->dbs[GEOIP_ASNUM_EDITION], ip);
160179
}
161180
break;
162181
case AF_INET6:
163-
if (gGeoIP[GEOIP_ASNUM_EDITION_V6]) {
182+
if (handle->dbs[GEOIP_ASNUM_EDITION_V6]) {
164183
geoipv6_t ip = reinterpret_cast<const struct sockaddr_in6 *>(addr)->sin6_addr;
165184

166185
v = 6;
167-
asn_name = GeoIP_name_by_ipnum_v6(gGeoIP[GEOIP_ASNUM_EDITION_V6], ip);
186+
asn_name = GeoIP_name_by_ipnum_v6(handle->dbs[GEOIP_ASNUM_EDITION_V6], ip);
168187
}
169188
break;
170189
}

plugins/header_rewrite/conditions_geo_maxmind.cc

Lines changed: 41 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -23,17 +23,25 @@
2323

2424
#include <unistd.h>
2525
#include <arpa/inet.h>
26+
#include <map>
27+
#include <mutex>
28+
#include <string>
2629

2730
#include "ts/ts.h"
2831

2932
#include "conditions_geo.h"
3033

3134
#include <maxminddb.h>
3235

33-
MMDB_s *gMaxMindDB = nullptr;
34-
3536
enum class MmdbSchema { NESTED, FLAT };
36-
static MmdbSchema gMmdbSchema = MmdbSchema::NESTED;
37+
38+
struct MmdbHandle {
39+
MMDB_s db;
40+
MmdbSchema schema = MmdbSchema::NESTED;
41+
};
42+
43+
static std::map<std::string, MmdbHandle *> gMmdbCache;
44+
static std::mutex gMmdbCacheMutex;
3745

3846
// Detect whether the MMDB uses nested (GeoLite2) or flat (vendor) field layout
3947
// by probing for the nested country path on a lookup result.
@@ -57,56 +65,63 @@ detect_schema(MMDB_entry_s *entry)
5765

5866
static const char *probe_ips[] = {"8.8.8.8", "1.1.1.1", "128.0.0.1"};
5967

60-
void
68+
void *
6169
MMConditionGeo::initLibrary(const std::string &path)
6270
{
6371
if (path.empty()) {
6472
Dbg(pi_dbg_ctl, "Empty MaxMind db path specified. Not initializing!");
65-
return;
73+
return nullptr;
6674
}
6775

68-
if (gMaxMindDB != nullptr) {
69-
Dbg(pi_dbg_ctl, "Maxmind library already initialized");
70-
return;
76+
std::lock_guard<std::mutex> lock(gMmdbCacheMutex);
77+
78+
auto it = gMmdbCache.find(path);
79+
if (it != gMmdbCache.end()) {
80+
Dbg(pi_dbg_ctl, "Maxmind library already initialized for %s", path.c_str());
81+
return it->second;
7182
}
7283

73-
gMaxMindDB = new MMDB_s;
84+
auto *handle = new MmdbHandle;
85+
int status = MMDB_open(path.c_str(), MMDB_MODE_MMAP, &handle->db);
7486

75-
int status = MMDB_open(path.c_str(), MMDB_MODE_MMAP, gMaxMindDB);
7687
if (MMDB_SUCCESS != status) {
7788
Dbg(pi_dbg_ctl, "Cannot open %s - %s", path.c_str(), MMDB_strerror(status));
78-
delete gMaxMindDB;
79-
gMaxMindDB = nullptr;
80-
return;
89+
delete handle;
90+
return nullptr;
8191
}
8292

8393
// Probe the database schema at load time so we know which field paths to
8494
// use for country lookups. Try a few well-known IPs until one hits.
8595
for (auto *ip : probe_ips) {
8696
int gai_error, mmdb_error;
87-
MMDB_lookup_result_s result = MMDB_lookup_string(gMaxMindDB, ip, &gai_error, &mmdb_error);
97+
MMDB_lookup_result_s result = MMDB_lookup_string(&handle->db, ip, &gai_error, &mmdb_error);
8898
if (gai_error == 0 && MMDB_SUCCESS == mmdb_error && result.found_entry) {
89-
gMmdbSchema = detect_schema(&result.entry);
90-
Dbg(pi_dbg_ctl, "Loaded %s (schema: %s)", path.c_str(), gMmdbSchema == MmdbSchema::FLAT ? "flat" : "nested");
91-
return;
99+
handle->schema = detect_schema(&result.entry);
100+
Dbg(pi_dbg_ctl, "Loaded %s (schema: %s)", path.c_str(), handle->schema == MmdbSchema::FLAT ? "flat" : "nested");
101+
gMmdbCache[path] = handle;
102+
return handle;
92103
}
93104
}
94105

95106
Dbg(pi_dbg_ctl, "Loaded %s (schema: defaulting to nested, no probe IPs matched)", path.c_str());
107+
gMmdbCache[path] = handle;
108+
return handle;
96109
}
97110

98111
std::string
99-
MMConditionGeo::get_geo_string(const sockaddr *addr) const
112+
MMConditionGeo::get_geo_string(const sockaddr *addr, void *geo_handle) const
100113
{
101114
std::string ret = "(unknown)";
102115
int mmdb_error;
103116

104-
if (gMaxMindDB == nullptr) {
117+
auto *handle = static_cast<MmdbHandle *>(geo_handle);
118+
119+
if (handle == nullptr) {
105120
Dbg(pi_dbg_ctl, "MaxMind not initialized; using default value");
106121
return ret;
107122
}
108123

109-
MMDB_lookup_result_s result = MMDB_lookup_sockaddr(gMaxMindDB, addr, &mmdb_error);
124+
MMDB_lookup_result_s result = MMDB_lookup_sockaddr(&handle->db, addr, &mmdb_error);
110125

111126
if (MMDB_SUCCESS != mmdb_error) {
112127
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
123138

124139
switch (_geo_qual) {
125140
case GEO_QUAL_COUNTRY:
126-
if (gMmdbSchema == MmdbSchema::FLAT) {
141+
if (handle->schema == MmdbSchema::FLAT) {
127142
status = MMDB_get_value(&result.entry, &entry_data, "country_code", NULL);
128143
} else {
129144
status = MMDB_get_value(&result.entry, &entry_data, "country", "iso_code", NULL);
@@ -150,17 +165,19 @@ MMConditionGeo::get_geo_string(const sockaddr *addr) const
150165
}
151166

152167
int64_t
153-
MMConditionGeo::get_geo_int(const sockaddr *addr) const
168+
MMConditionGeo::get_geo_int(const sockaddr *addr, void *geo_handle) const
154169
{
155170
int64_t ret = -1;
156171
int mmdb_error;
157172

158-
if (gMaxMindDB == nullptr) {
173+
auto *handle = static_cast<MmdbHandle *>(geo_handle);
174+
175+
if (handle == nullptr) {
159176
Dbg(pi_dbg_ctl, "MaxMind not initialized; using default value");
160177
return ret;
161178
}
162179

163-
MMDB_lookup_result_s result = MMDB_lookup_sockaddr(gMaxMindDB, addr, &mmdb_error);
180+
MMDB_lookup_result_s result = MMDB_lookup_sockaddr(&handle->db, addr, &mmdb_error);
164181

165182
if (MMDB_SUCCESS != mmdb_error) {
166183
Dbg(pi_dbg_ctl, "Error during sockaddr lookup: %s", MMDB_strerror(mmdb_error));

0 commit comments

Comments
 (0)