Line data    Source code

       1             : #pragma once
       2             : 
       3             : #include "envoy/common/platform.h"
       4             : #include "envoy/extensions/geoip_providers/maxmind/v3/maxmind.pb.h"
       5             : #include "envoy/geoip/geoip_provider_driver.h"
       6             : 
       7             : #include "source/common/common/logger.h"
       8             : 
       9             : #include "maxminddb.h"
      10             : 
      11             : namespace Envoy {
      12             : namespace Extensions {
      13             : namespace GeoipProviders {
      14             : namespace Maxmind {
      15             : 
      16             : class GeoipProviderConfig {
      17             : public:
      18             :   GeoipProviderConfig(const envoy::extensions::geoip_providers::maxmind::v3::MaxMindConfig& config,
      19             :                       const std::string& stat_prefix, Stats::Scope& scope);
      20             : 
      21           0 :   const absl::optional<std::string>& cityDbPath() const { return city_db_path_; }
      22           0 :   const absl::optional<std::string>& ispDbPath() const { return isp_db_path_; }
      23           0 :   const absl::optional<std::string>& anonDbPath() const { return anon_db_path_; }
      24             : 
      25             :   bool isLookupEnabledForHeader(const absl::optional<std::string>& header);
      26             : 
      27           0 :   const absl::optional<std::string>& countryHeader() const { return country_header_; }
      28           0 :   const absl::optional<std::string>& cityHeader() const { return city_header_; }
      29           0 :   const absl::optional<std::string>& regionHeader() const { return region_header_; }
      30           0 :   const absl::optional<std::string>& asnHeader() const { return asn_header_; }
      31             : 
      32           0 :   const absl::optional<std::string>& anonHeader() const { return anon_header_; }
      33           0 :   const absl::optional<std::string>& anonVpnHeader() const { return anon_vpn_header_; }
      34           0 :   const absl::optional<std::string>& anonHostingHeader() const { return anon_hosting_header_; }
      35           0 :   const absl::optional<std::string>& anonTorHeader() const { return anon_tor_header_; }
      36           0 :   const absl::optional<std::string>& anonProxyHeader() const { return anon_proxy_header_; }
      37             : 
      38           0 :   void incLookupError(absl::string_view maxmind_db_type) {
      39           0 :     incCounter(
      40           0 :         stat_name_set_->getBuiltin(absl::StrCat(maxmind_db_type, ".lookup_error"), unknown_hit_));
      41           0 :   }
      42             : 
      43           0 :   void incTotal(absl::string_view maxmind_db_type) {
      44           0 :     incCounter(stat_name_set_->getBuiltin(absl::StrCat(maxmind_db_type, ".total"), unknown_hit_));
      45           0 :   }
      46             : 
      47           0 :   void incHit(absl::string_view maxmind_db_type) {
      48           0 :     incCounter(stat_name_set_->getBuiltin(absl::StrCat(maxmind_db_type, ".hit"), unknown_hit_));
      49           0 :   }
      50             : 
      51             :   void registerGeoDbStats(const std::string& db_type);
      52             : 
      53           0 :   Stats::Scope& getStatsScopeForTest() const { return *stats_scope_; }
      54             : 
      55             : private:
      56             :   absl::optional<std::string> city_db_path_;
      57             :   absl::optional<std::string> isp_db_path_;
      58             :   absl::optional<std::string> anon_db_path_;
      59             : 
      60             :   absl::optional<std::string> country_header_;
      61             :   absl::optional<std::string> city_header_;
      62             :   absl::optional<std::string> region_header_;
      63             :   absl::optional<std::string> asn_header_;
      64             : 
      65             :   absl::optional<std::string> anon_header_;
      66             :   absl::optional<std::string> anon_vpn_header_;
      67             :   absl::optional<std::string> anon_hosting_header_;
      68             :   absl::optional<std::string> anon_tor_header_;
      69             :   absl::optional<std::string> anon_proxy_header_;
      70             : 
      71             :   Stats::ScopeSharedPtr stats_scope_;
      72             :   Stats::StatNameSetPtr stat_name_set_;
      73             :   const Stats::StatName unknown_hit_;
      74             :   void incCounter(Stats::StatName name);
      75             : };
      76             : 
      77             : using GeoipProviderConfigSharedPtr = std::shared_ptr<GeoipProviderConfig>;
      78             : 
      79             : using MaxmindDbPtr = std::unique_ptr<MMDB_s>;
      80             : class GeoipProvider : public Envoy::Geolocation::Driver,
      81             :                       public Logger::Loggable<Logger::Id::geolocation> {
      82             : 
      83             : public:
      84             :   GeoipProvider(Singleton::InstanceSharedPtr owner, GeoipProviderConfigSharedPtr config)
      85           0 :       : config_(config), owner_(owner) {
      86           0 :     city_db_ = initMaxMindDb(config_->cityDbPath());
      87           0 :     isp_db_ = initMaxMindDb(config_->ispDbPath());
      88           0 :     anon_db_ = initMaxMindDb(config_->anonDbPath());
      89           0 :   };
      90             : 
      91             :   ~GeoipProvider();
      92             : 
      93             :   // Envoy::Geolocation::Driver
      94             :   void lookup(Geolocation::LookupRequest&&, Geolocation::LookupGeoHeadersCallback&&) const override;
      95             : 
      96             : private:
      97             :   // Allow the unit test to have access to private members.
      98             :   friend class GeoipProviderPeer;
      99             :   GeoipProviderConfigSharedPtr config_;
     100             :   MaxmindDbPtr city_db_;
     101             :   MaxmindDbPtr isp_db_;
     102             :   MaxmindDbPtr anon_db_;
     103             :   MaxmindDbPtr initMaxMindDb(const absl::optional<std::string>& db_path);
     104             :   void lookupInCityDb(const Network::Address::InstanceConstSharedPtr& remote_address,
     105             :                       absl::flat_hash_map<std::string, std::string>& lookup_result) const;
     106             :   void lookupInAsnDb(const Network::Address::InstanceConstSharedPtr& remote_address,
     107             :                      absl::flat_hash_map<std::string, std::string>& lookup_result) const;
     108             :   void lookupInAnonDb(const Network::Address::InstanceConstSharedPtr& remote_address,
     109             :                       absl::flat_hash_map<std::string, std::string>& lookup_result) const;
     110             :   template <typename... Params>
     111             :   void populateGeoLookupResult(MMDB_lookup_result_s& mmdb_lookup_result,
     112             :                                absl::flat_hash_map<std::string, std::string>& lookup_result,
     113             :                                const std::string& result_key, Params... lookup_params) const;
     114             :   // A shared_ptr to keep the provider singleton alive as long as any of its providers are in use.
     115             :   const Singleton::InstanceSharedPtr owner_;
     116             : };
     117             : 
     118             : using GeoipProviderSharedPtr = std::shared_ptr<GeoipProvider>;
     119             : 
     120             : } // namespace Maxmind
     121             : } // namespace GeoipProviders
     122             : } // namespace Extensions
     123             : } // namespace Envoy