@@ -23,6 +23,7 @@ static constexpr absl::string_view CITY_DB_TYPE = "city_db";
2323static constexpr absl::string_view ISP_DB_TYPE = " isp_db" ;
2424static constexpr absl::string_view ANON_DB_TYPE = " anon_db" ;
2525static constexpr absl::string_view ASN_DB_TYPE = " asn_db" ;
26+ static constexpr absl::string_view COUNTRY_DB_TYPE = " country_db" ;
2627} // namespace
2728
2829GeoipProviderConfig::GeoipProviderConfig (
@@ -36,6 +37,9 @@ GeoipProviderConfig::GeoipProviderConfig(
3637 : absl::nullopt),
3738 asn_db_path_(!config.asn_db_path().empty() ? absl::make_optional(config.asn_db_path())
3839 : absl::nullopt),
40+ country_db_path_(!config.country_db_path().empty()
41+ ? absl::make_optional(config.country_db_path())
42+ : absl::nullopt),
3943 stats_scope_(scope.createScope(absl::StrCat(stat_prefix, " maxmind." ))),
4044 stat_name_set_(stats_scope_->symbolTable ().makeSet(" Maxmind" )) {
4145 auto geo_headers_to_add = config.common_provider_config ().geo_headers_to_add ();
@@ -75,9 +79,9 @@ GeoipProviderConfig::GeoipProviderConfig(
7579 apple_private_relay_header_ = !geo_headers_to_add.apple_private_relay ().empty ()
7680 ? absl::make_optional (geo_headers_to_add.apple_private_relay ())
7781 : absl::nullopt ;
78- if (!city_db_path_ && !anon_db_path_ && !asn_db_path_ && !isp_db_path_) {
82+ if (!city_db_path_ && !anon_db_path_ && !asn_db_path_ && !isp_db_path_ && !country_db_path_ ) {
7983 throw EnvoyException (" At least one geolocation database path needs to be configured: "
80- " city_db_path, isp_db_path, asn_db_path or anon_db_path " );
84+ " city_db_path, isp_db_path, asn_db_path, anon_db_path or country_db_path " );
8185 }
8286 if (city_db_path_) {
8387 registerGeoDbStats (CITY_DB_TYPE);
@@ -91,6 +95,9 @@ GeoipProviderConfig::GeoipProviderConfig(
9195 if (asn_db_path_) {
9296 registerGeoDbStats (ASN_DB_TYPE);
9397 }
98+ if (country_db_path_) {
99+ registerGeoDbStats (COUNTRY_DB_TYPE);
100+ }
94101};
95102
96103void GeoipProviderConfig::registerGeoDbStats (const absl::string_view& db_type) {
@@ -126,6 +133,9 @@ GeoipProvider::GeoipProvider(Event::Dispatcher& dispatcher, Api::Api& api,
126133 config_->anonDbPath () ? initMaxmindDb (config_->anonDbPath ().value (), ANON_DB_TYPE) : nullptr ;
127134 asn_db_ =
128135 config_->asnDbPath () ? initMaxmindDb (config_->asnDbPath ().value (), ASN_DB_TYPE) : nullptr ;
136+ country_db_ = config_->countryDbPath ()
137+ ? initMaxmindDb (config_->countryDbPath ().value (), COUNTRY_DB_TYPE)
138+ : nullptr ;
129139 mmdb_reload_dispatcher_ = api.allocateDispatcher (" mmdb_reload_routine" );
130140 mmdb_watcher_ = dispatcher.createFilesystemWatcher ();
131141 mmdb_reload_thread_ = api.threadFactory ().createThread (
@@ -157,6 +167,13 @@ GeoipProvider::GeoipProvider(Event::Dispatcher& dispatcher, Api::Api& api,
157167 return onMaxmindDbUpdate (config_->asnDbPath ().value (), ASN_DB_TYPE);
158168 }));
159169 }
170+ if (config_->countryDbPath ()) {
171+ THROW_IF_NOT_OK (mmdb_watcher_->addWatch (
172+ config_->countryDbPath ().value (), Filesystem::Watcher::Events::MovedTo,
173+ [this ](uint32_t ) {
174+ return onMaxmindDbUpdate (config_->countryDbPath ().value (), COUNTRY_DB_TYPE);
175+ }));
176+ }
160177 mmdb_reload_dispatcher_->run (Event::Dispatcher::RunType::RunUntilExit);
161178 },
162179 Thread::Options{std::string (" mmdb_reload_routine" )});
@@ -177,6 +194,7 @@ void GeoipProvider::lookup(Geolocation::LookupRequest&& request,
177194 Geolocation::LookupGeoHeadersCallback&& cb) const {
178195 auto & remote_address = request.remoteAddress ();
179196 auto lookup_result = absl::flat_hash_map<std::string, std::string>{};
197+ lookupInCountryDb (remote_address, lookup_result);
180198 lookupInCityDb (remote_address, lookup_result);
181199 lookupInAsnDb (remote_address, lookup_result);
182200 lookupInAnonDb (remote_address, lookup_result);
@@ -187,9 +205,12 @@ void GeoipProvider::lookup(Geolocation::LookupRequest&& request,
187205void GeoipProvider::lookupInCityDb (
188206 const Network::Address::InstanceConstSharedPtr& remote_address,
189207 absl::flat_hash_map<std::string, std::string>& lookup_result) const {
208+ // Country lookup falls back to City DB only if Country DB is not configured.
209+ const bool should_lookup_country_from_city_db =
210+ !config_->isCountryDbPathSet () && config_->isLookupEnabledForHeader (config_->countryHeader ());
190211 if (config_->isLookupEnabledForHeader (config_->cityHeader ()) ||
191212 config_->isLookupEnabledForHeader (config_->regionHeader ()) ||
192- config_-> isLookupEnabledForHeader (config_-> countryHeader ()) ) {
213+ should_lookup_country_from_city_db ) {
193214 int mmdb_error;
194215 auto city_db_ptr = getCityDb ();
195216 // Used for testing.
@@ -217,7 +238,8 @@ void GeoipProvider::lookupInCityDb(
217238 config_->regionHeader ().value (), MMDB_REGION_LOOKUP_ARGS[0 ],
218239 MMDB_REGION_LOOKUP_ARGS[1 ], MMDB_REGION_LOOKUP_ARGS[2 ]);
219240 }
220- if (config_->isLookupEnabledForHeader (config_->countryHeader ())) {
241+ // Country lookup from City DB only when Country DB is not configured.
242+ if (should_lookup_country_from_city_db) {
221243 populateGeoLookupResult (mmdb_lookup_result, lookup_result,
222244 config_->countryHeader ().value (), MMDB_COUNTRY_LOOKUP_ARGS[0 ],
223245 MMDB_COUNTRY_LOOKUP_ARGS[1 ]);
@@ -384,6 +406,51 @@ void GeoipProvider::lookupInIspDb(
384406 }
385407}
386408
409+ void GeoipProvider::lookupInCountryDb (
410+ const Network::Address::InstanceConstSharedPtr& remote_address,
411+ absl::flat_hash_map<std::string, std::string>& lookup_result) const {
412+ if (config_->isLookupEnabledForHeader (config_->countryHeader ())) {
413+ // Country DB takes precedence if configured, otherwise fall back to City DB.
414+ if (!config_->isCountryDbPathSet ()) {
415+ // Country lookup will be handled by lookupInCityDb.
416+ return ;
417+ }
418+ int mmdb_error;
419+ auto country_db_ptr = getCountryDb ();
420+ // Used for testing.
421+ synchronizer_.syncPoint (std::string (COUNTRY_DB_TYPE).append (" _lookup_pre_complete" ));
422+ if (!country_db_ptr) {
423+ if (config_->isCityDbPathSet ()) {
424+ // Country information can be looked up from City database as well, so we don't need to
425+ // throw an error if it is not set.
426+ return ;
427+ }
428+ IS_ENVOY_BUG (" Maxmind country database must be initialised for performing lookups" );
429+ return ;
430+ }
431+ auto country_db = country_db_ptr.get ();
432+ MMDB_lookup_result_s mmdb_lookup_result = MMDB_lookup_sockaddr (
433+ country_db->mmdb (), reinterpret_cast <const sockaddr*>(remote_address->sockAddr ()),
434+ &mmdb_error);
435+ const uint32_t n_prev_hits = lookup_result.size ();
436+ if (!mmdb_error && mmdb_lookup_result.found_entry ) {
437+ MMDB_entry_data_list_s* entry_data_list;
438+ int status = MMDB_get_entry_data_list (&mmdb_lookup_result.entry , &entry_data_list);
439+ if (status == MMDB_SUCCESS) {
440+ populateGeoLookupResult (mmdb_lookup_result, lookup_result, config_->countryHeader ().value (),
441+ MMDB_COUNTRY_LOOKUP_ARGS[0 ], MMDB_COUNTRY_LOOKUP_ARGS[1 ]);
442+ if (lookup_result.size () > n_prev_hits) {
443+ config_->incHit (COUNTRY_DB_TYPE);
444+ }
445+ MMDB_free_entry_data_list (entry_data_list);
446+ } else {
447+ config_->incLookupError (COUNTRY_DB_TYPE);
448+ }
449+ }
450+ config_->incTotal (COUNTRY_DB_TYPE);
451+ }
452+ }
453+
387454MaxmindDbSharedPtr GeoipProvider::initMaxmindDb (const std::string& db_path,
388455 const absl::string_view& db_type, bool reload) {
389456 MMDB_s maxmind_db;
@@ -422,6 +489,9 @@ absl::Status GeoipProvider::mmdbReload(const MaxmindDbSharedPtr reloaded_db,
422489 } else if (db_type == ASN_DB_TYPE) {
423490 updateAsnDb (reloaded_db);
424491 config_->incDbReloadSuccess (db_type);
492+ } else if (db_type == COUNTRY_DB_TYPE) {
493+ updateCountryDb (reloaded_db);
494+ config_->incDbReloadSuccess (db_type);
425495 } else {
426496 ENVOY_LOG (error, " Unsupported maxmind db type {}" , db_type);
427497 return absl::InvalidArgumentError (fmt::format (" Unsupported maxmind db type {}" , db_type));
@@ -472,6 +542,17 @@ void GeoipProvider::updateAnonDb(MaxmindDbSharedPtr anon_db) ABSL_LOCKS_EXCLUDED
472542 anon_db_ = anon_db;
473543}
474544
545+ MaxmindDbSharedPtr GeoipProvider::getCountryDb () const ABSL_LOCKS_EXCLUDED(mmdb_mutex_) {
546+ absl::ReaderMutexLock lock (&mmdb_mutex_);
547+ return country_db_;
548+ }
549+
550+ void GeoipProvider::updateCountryDb (MaxmindDbSharedPtr country_db)
551+ ABSL_LOCKS_EXCLUDED(mmdb_mutex_) {
552+ absl::MutexLock lock (&mmdb_mutex_);
553+ country_db_ = country_db;
554+ }
555+
475556absl::Status GeoipProvider::onMaxmindDbUpdate (const std::string& db_path,
476557 const absl::string_view& db_type) {
477558 MaxmindDbSharedPtr reloaded_db = initMaxmindDb (db_path, db_type, true /* reload */ );
0 commit comments