diff --git a/Cargo.lock b/Cargo.lock index 2bbc6c8..b191314 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3696,7 +3696,9 @@ dependencies = [ "ipnetwork", "log", "memchr", + "memmap2", "serde", + "simdutf8", "thiserror 2.0.20", ] @@ -3706,6 +3708,15 @@ version = "2.8.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" +[[package]] +name = "memmap2" +version = "0.9.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d1219ed1b7f229ee7104d281dd01d6802fe28bb6e95d292942c4daacdeb798c0" +dependencies = [ + "libc", +] + [[package]] name = "merlin" version = "3.0.0" diff --git a/Cargo.toml b/Cargo.toml index 19a2cf7..d631a7a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -82,7 +82,7 @@ mail-parser = { version = "0.11.3", features = ["full_encoding"] } mail-send = { version = "0.6.0", default-features = false, features = [ "builder", ] } -maxminddb = "0.29" +maxminddb = { version = "0.29", features = ["mmap", "simdutf8"] } mockall = "0.15" mock-io = { version = "0.3.2", features = ["full"] } moka = { version = "0.12.15", features = ["sync"] } diff --git a/ic-bn-lib/src/geoip.rs b/ic-bn-lib/src/geoip.rs new file mode 100644 index 0000000..c1b9c1b --- /dev/null +++ b/ic-bn-lib/src/geoip.rs @@ -0,0 +1,333 @@ +use std::{fmt::Display, net::IpAddr, ops::Deref, path::PathBuf}; + +use anyhow::Context; +use arrayvec::ArrayString; +use maxminddb::geoip2; +use serde::{Deserialize, Serialize}; + +use crate::{Error, TruncatesString}; + +const CITY_NAME_MAX_LENGTH: usize = 32; + +/// Two-letter country code. +/// See https://en.wikipedia.org/wiki/ISO_3166-1_alpha-2 +#[derive(Clone, Copy, Debug, Eq, PartialEq, Ord, PartialOrd, Hash, Deserialize, Serialize)] +pub struct CountryCode(pub ArrayString<2>); + +impl Deref for CountryCode { + type Target = str; + + fn deref(&self) -> &Self::Target { + self.0.as_str() + } +} + +impl Display for CountryCode { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{}", self.0) + } +} + +/// Location representation +pub struct Location { + pub lat: f64, + pub lon: f64, +} + +/// GeoIP lookup city representation. +pub struct City { + pub name: Option>, + pub country_code: Option, + pub location: Option, +} + +impl From> for City { + fn from(city: geoip2::City<'_>) -> Self { + Self { + // Try English, then German, otherwise None + name: city + .city + .names + .english + .or(city.city.names.german) + // SAFETY: truncate_bytes makes the string *no longer* than CITY_NAME_MAX_LENGTH + // so it will always fit into Arraystring + .map(|x| x.truncate_bytes(CITY_NAME_MAX_LENGTH).try_into().unwrap()), + + country_code: city + .country + .iso_code + .and_then(|x| x.try_into().ok()) + .map(CountryCode), + + // Location is Some only when both lat & lon are available + location: city + .location + .latitude + .zip(city.location.longitude) + .map(|(lat, lon)| Location { lat, lon }), + } + } +} + +/// Looks up the client's country using his IP address +pub struct GeoIp { + db: maxminddb::Reader>, +} + +impl GeoIp { + /// Creates a new GeoIp instance from a provided database + pub fn new(db_path: &PathBuf) -> Result { + Ok(Self { + db: maxminddb::Reader::open_readfile(db_path).context("unable to load GeoIP DB")?, + }) + } + + /// Looks up the country code from an IP + pub fn lookup_country(&self, ip: IpAddr) -> Option { + let country: Option = self.db.lookup(ip).ok()?.decode().ok()?; + // Country code should always fit into 2-letter ArrayString. + // If for whatever reason it does not - return None. + Some(CountryCode(country?.country.iso_code?.try_into().ok()?)) + } + + /// Looks up the city from an IP + pub fn lookup_city(&self, ip: IpAddr) -> Option { + let city: geoip2::City = self.db.lookup(ip).ok()?.decode().ok()??; + Some(city.into()) + } +} + +#[cfg(test)] +mod test { + use std::net::{Ipv4Addr, Ipv6Addr}; + + use super::*; + + // Known entries in the MaxMind test DBs (present in both Country & City DBs) + const IP_KNOWN: Ipv4Addr = Ipv4Addr::new(89, 160, 20, 112); + const COUNTRY_KNOWN: &str = "SE"; + const CITY_KNOWN: &str = "Linköping"; + const LAT_KNOWN: f64 = 58.4167; + const LON_KNOWN: f64 = 15.6167; + const IP_UNKNOWN: Ipv4Addr = Ipv4Addr::new(10, 10, 10, 10); + // City DB: country & location, but no city name + const IP_NO_CITY_NAME: Ipv4Addr = Ipv4Addr::new(149, 101, 100, 1); + // City DB: location only, no country and no city name + const IP_LOCATION_ONLY: Ipv6Addr = Ipv6Addr::new(0x2a02, 0xd500, 0, 0, 0, 0, 0, 1); + // City DB: record exists but is empty + const IP_EMPTY_RECORD: Ipv4Addr = Ipv4Addr::new(2, 3, 3, 1); + + /// MaxMind GeoIP2-Country test DB + fn test_db_path() -> PathBuf { + PathBuf::from(concat!( + env!("CARGO_MANIFEST_DIR"), + "/test-data/geoip-test-db.mmdb" + )) + } + + /// MaxMind GeoIP2-City test DB + fn test_city_db_path() -> PathBuf { + PathBuf::from(concat!( + env!("CARGO_MANIFEST_DIR"), + "/test-data/geoip-city-test-db.mmdb" + )) + } + + /// Builds a raw MaxMind city record with only the fields we care about + fn geoip2_city<'a>( + english: Option<&'a str>, + german: Option<&'a str>, + iso_code: Option<&'a str>, + latitude: Option, + longitude: Option, + ) -> geoip2::City<'a> { + geoip2::City { + city: geoip2::city::City { + names: geoip2::Names { + english, + german, + ..Default::default() + }, + ..Default::default() + }, + country: geoip2::city::Country { + iso_code, + ..Default::default() + }, + location: geoip2::city::Location { + latitude, + longitude, + ..Default::default() + }, + ..Default::default() + } + } + + #[test] + fn lookup_known_ip_returns_country_code() { + let geoip = GeoIp::new(&test_db_path()).unwrap(); + assert_eq!( + geoip + .lookup_country(IpAddr::V4(IP_KNOWN)) + .unwrap() + .0 + .as_str(), + COUNTRY_KNOWN + ); + } + + #[test] + fn lookup_unknown_ip_returns_none() { + let geoip = GeoIp::new(&test_db_path()).unwrap(); + assert!(geoip.lookup_country(IpAddr::V4(IP_UNKNOWN)).is_none()); + } + + #[test] + fn lookup_city_known_ip_returns_name_country_and_location() { + let geoip = GeoIp::new(&test_city_db_path()).unwrap(); + let city = geoip.lookup_city(IpAddr::V4(IP_KNOWN)).unwrap(); + + assert_eq!(city.name.unwrap().as_str(), CITY_KNOWN); + assert_eq!(city.country_code.unwrap().0.as_str(), COUNTRY_KNOWN); + + let location = city.location.unwrap(); + assert_eq!(location.lat, LAT_KNOWN); + assert_eq!(location.lon, LON_KNOWN); + } + + #[test] + fn lookup_city_unknown_ip_returns_none() { + let geoip = GeoIp::new(&test_city_db_path()).unwrap(); + assert!(geoip.lookup_city(IpAddr::V4(IP_UNKNOWN)).is_none()); + } + + #[test] + fn lookup_city_without_name_returns_country_and_location() { + let geoip = GeoIp::new(&test_city_db_path()).unwrap(); + let city = geoip.lookup_city(IpAddr::V4(IP_NO_CITY_NAME)).unwrap(); + + assert!(city.name.is_none()); + assert_eq!(city.country_code.unwrap().0.as_str(), "US"); + + let location = city.location.unwrap(); + assert_eq!(location.lat, 37.751); + assert_eq!(location.lon, -97.822); + } + + #[test] + fn lookup_city_ipv6_with_location_only() { + let geoip = GeoIp::new(&test_city_db_path()).unwrap(); + let city = geoip.lookup_city(IpAddr::V6(IP_LOCATION_ONLY)).unwrap(); + + assert!(city.name.is_none()); + assert!(city.country_code.is_none()); + + let location = city.location.unwrap(); + assert_eq!(location.lat, 48.69096); + assert_eq!(location.lon, 9.14062); + } + + #[test] + fn lookup_city_empty_record_returns_city_without_fields() { + let geoip = GeoIp::new(&test_city_db_path()).unwrap(); + let city = geoip.lookup_city(IpAddr::V4(IP_EMPTY_RECORD)).unwrap(); + + assert!(city.name.is_none()); + assert!(city.country_code.is_none()); + assert!(city.location.is_none()); + } + + #[test] + fn lookup_city_with_country_db_returns_country_only() { + let geoip = GeoIp::new(&test_db_path()).unwrap(); + let city = geoip.lookup_city(IpAddr::V4(IP_KNOWN)).unwrap(); + + assert!(city.name.is_none()); + assert_eq!(city.country_code.unwrap().0.as_str(), COUNTRY_KNOWN); + assert!(city.location.is_none()); + } + + #[test] + fn city_from_geoip2_prefers_english_name() { + let city = City::from(geoip2_city( + Some("Singapore"), + Some("Singapur"), + None, + None, + None, + )); + assert_eq!(city.name.unwrap().as_str(), "Singapore"); + } + + #[test] + fn city_from_geoip2_falls_back_to_german_name() { + let city = City::from(geoip2_city(None, Some("Singapur"), None, None, None)); + assert_eq!(city.name.unwrap().as_str(), "Singapur"); + } + + #[test] + fn city_from_geoip2_without_names_has_no_name() { + let city = City::from(geoip2_city(None, None, None, None, None)); + assert!(city.name.is_none()); + } + + #[test] + fn city_from_geoip2_truncates_long_name() { + // 58 ASCII bytes, truncated to exactly CITY_NAME_MAX_LENGTH bytes + let long = "Llanfairpwllgwyngyllgogerychwyrndrobwllllantysiliogogogoch"; + assert!(long.len() > CITY_NAME_MAX_LENGTH); + let city = City::from(geoip2_city(Some(long), None, None, None, None)); + assert_eq!(city.name.unwrap().as_str(), &long[..CITY_NAME_MAX_LENGTH]); + + // 17 two-byte chars = 34 bytes, byte 32 is a char boundary -> 16 chars kept + let name = "ä".repeat(17); + let city = City::from(geoip2_city(Some(&name), None, None, None, None)); + assert_eq!(city.name.unwrap().as_str(), "ä".repeat(16).as_str()); + + // 11 three-byte chars = 33 bytes, byte 32 is mid-char -> 10 chars (30 bytes) kept + let name = "€".repeat(11); + let city = City::from(geoip2_city(Some(&name), None, None, None, None)); + assert_eq!(city.name.unwrap().as_str(), "€".repeat(10).as_str()); + + // Truncation also applies to the German fallback + let city = City::from(geoip2_city(None, Some(long), None, None, None)); + assert_eq!(city.name.unwrap().as_str(), &long[..CITY_NAME_MAX_LENGTH]); + } + + #[test] + fn city_from_geoip2_country_code() { + let city = City::from(geoip2_city(None, None, Some("SE"), None, None)); + assert_eq!(city.country_code.unwrap().0.as_str(), "SE"); + + // ISO code that doesn't fit into two letters is dropped + let city = City::from(geoip2_city(None, None, Some("SWE"), None, None)); + assert!(city.country_code.is_none()); + + let city = City::from(geoip2_city(None, None, None, None, None)); + assert!(city.country_code.is_none()); + } + + #[test] + fn city_from_geoip2_location_requires_both_coordinates() { + let city = City::from(geoip2_city( + None, + None, + None, + Some(LAT_KNOWN), + Some(LON_KNOWN), + )); + let location = city.location.unwrap(); + assert_eq!(location.lat, LAT_KNOWN); + assert_eq!(location.lon, LON_KNOWN); + + let city = City::from(geoip2_city(None, None, None, Some(LAT_KNOWN), None)); + assert!(city.location.is_none()); + + let city = City::from(geoip2_city(None, None, None, None, Some(LON_KNOWN))); + assert!(city.location.is_none()); + + let city = City::from(geoip2_city(None, None, None, None, None)); + assert!(city.location.is_none()); + } +} diff --git a/ic-bn-lib/src/http/middleware/request_meta.rs b/ic-bn-lib/src/http/middleware/request_meta.rs index 852d775..ef05b8a 100644 --- a/ic-bn-lib/src/http/middleware/request_meta.rs +++ b/ic-bn-lib/src/http/middleware/request_meta.rs @@ -8,7 +8,6 @@ use std::{ }; use anyhow::Context; -use arrayvec::ArrayString; use axum::{ extract::{Request, State}, middleware::Next, @@ -17,11 +16,11 @@ use axum::{ use bytes::Bytes; use http::{HeaderMap, header::HeaderValue}; use ipnet::IpNet; -use maxminddb::geoip2; use serde::{Deserialize, Serialize}; use crate::{ Error, + geoip::GeoIp, http::{ headers::{X_REAL_IP, X_REQUEST_ID}, server::conn::ConnInfo, @@ -71,51 +70,10 @@ impl Display for RequestId { } } -/// Two-letter country code. -/// See https://en.wikipedia.org/wiki/ISO_3166-1_alpha-2 -#[derive(Clone, Copy, Debug, Eq, PartialEq, Ord, PartialOrd, Hash, Deserialize, Serialize)] -pub struct CountryCode(pub ArrayString<2>); - -impl Deref for CountryCode { - type Target = str; - - fn deref(&self) -> &Self::Target { - self.0.as_str() - } -} - -impl Display for CountryCode { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!(f, "{}", self.0) - } -} - -/// Looks up the client's country using his IP address -pub struct GeoIp { - db: maxminddb::Reader>, -} - -impl GeoIp { - /// Creates a new GeoIp instance from a provided database - pub fn new(db_path: &PathBuf) -> Result { - Ok(Self { - db: maxminddb::Reader::open_readfile(db_path).context("unable to load GeoIP DB")?, - }) - } - - /// Looks up the country code from an IP - pub fn lookup(&self, ip: IpAddr) -> Option { - let country: Option = self.db.lookup(ip).ok()?.decode().ok()?; - // Country code should always fit into 2-letter ArrayString. - // If for whatever reason it does not - return None. - Some(CountryCode(country?.country.iso_code?.try_into().ok()?)) - } -} - /// State for [`middleware`] pub struct RequestMetaState { /// Optional GeoIP database - geoip: Option, + geoip: Option>, /// Trust incoming headers from these subnets for the purpose of IP address extraction. /// If not set - headers will not be used. @@ -136,23 +94,36 @@ impl RequestMetaState { } } - /// Creates a new [`RequestMetaState`] with a GeoIP DB - pub fn new_with_geoip( + /// Creates a new [`RequestMetaState`] with a GeoIP provided + pub const fn new_with_geoip( + trust_ip_from: Vec, + trust_request_id_from: Vec, + geoip: Option>, + ) -> Self { + Self { + geoip, + trust_ip_from, + trust_request_id_from, + } + } + + /// Creates a new [`RequestMetaState`] with a GeoIP DB path + pub fn new_with_geoip_db( trust_ip_from: Vec, trust_request_id_from: Vec, geoip_db_path: Option, ) -> Result { let geoip = if let Some(v) = geoip_db_path { - Some(GeoIp::new(&v).context("unable to init GeoIP")?) + Some(Arc::new(GeoIp::new(&v).context("unable to init GeoIP")?)) } else { None }; - Ok(Self { - geoip, + Ok(Self::new_with_geoip( trust_ip_from, trust_request_id_from, - }) + geoip, + )) } /// Extracts remote IP address from the `x-real-ip` header if remote is trusted & header exists @@ -217,7 +188,7 @@ pub async fn middleware( request.extensions_mut().insert(v); // Look up country code if GeoIP is enabled - state.geoip.as_ref().and_then(|x| x.lookup(v.0)) + state.geoip.as_ref().and_then(|x| x.lookup_country(v.0)) }); if let Some(v) = country_code { @@ -265,7 +236,7 @@ mod test { use tower::Service; use super::*; - use crate::{hname, http::server::conn::ConnInfo, hval, network::Addr}; + use crate::{geoip::CountryCode, hname, http::server::conn::ConnInfo, hval, network::Addr}; const X_TEST_REQUEST_ID: &str = "x-test-request-id"; const X_TEST_REMOTE_ADDR: &str = "x-test-remote-addr"; @@ -284,21 +255,6 @@ mod test { )) } - #[test] - fn lookup_known_ip_returns_country_code() { - let geoip = GeoIp::new(&test_db_path()).unwrap(); - assert_eq!( - geoip.lookup(IpAddr::V4(IP_KNOWN)).unwrap().0.as_str(), - COUNTRY_KNOWN - ); - } - - #[test] - fn lookup_unknown_ip_returns_none() { - let geoip = GeoIp::new(&test_db_path()).unwrap(); - assert!(geoip.lookup(IpAddr::V4(IP_UNKNOWN)).is_none()); - } - #[test] fn extract_ip_ignores_header_from_untrusted_source() { let state = RequestMetaState::new(vec![], vec![]); @@ -657,7 +613,8 @@ mod test { #[tokio::test] async fn middleware_geoip_unknown_ip_no_country_code() { - let state = RequestMetaState::new_with_geoip(vec![], vec![], Some(test_db_path())).unwrap(); + let state = + RequestMetaState::new_with_geoip_db(vec![], vec![], Some(test_db_path())).unwrap(); let mut app = app(state); let req = Request::builder().body(Body::empty()).unwrap(); @@ -671,7 +628,8 @@ mod test { #[tokio::test] async fn middleware_geoip_known_ip_attaches_country_code_to_request_and_response() { - let state = RequestMetaState::new_with_geoip(vec![], vec![], Some(test_db_path())).unwrap(); + let state = + RequestMetaState::new_with_geoip_db(vec![], vec![], Some(test_db_path())).unwrap(); let mut app = app(state); let req = Request::builder().body(Body::empty()).unwrap(); diff --git a/ic-bn-lib/src/http/middleware/waf.rs b/ic-bn-lib/src/http/middleware/waf.rs index 716c824..08e4165 100644 --- a/ic-bn-lib/src/http/middleware/waf.rs +++ b/ic-bn-lib/src/http/middleware/waf.rs @@ -42,11 +42,8 @@ use tracing::warn; use url::Url; use crate::{ - http::{ - Error, - client::Client, - middleware::{RemoteAddr, request_meta::CountryCode}, - }, + geoip::CountryCode, + http::{Error, client::Client, middleware::RemoteAddr}, tasks::Run, }; diff --git a/ic-bn-lib/src/lib.rs b/ic-bn-lib/src/lib.rs index 7294fd2..d96eb05 100644 --- a/ic-bn-lib/src/lib.rs +++ b/ic-bn-lib/src/lib.rs @@ -9,6 +9,7 @@ #[cfg(feature = "custom-domains")] pub mod custom_domains; pub mod dns; +pub mod geoip; pub mod health; pub mod http; pub mod ic; diff --git a/ic-bn-lib/test-data/geoip-city-test-db.mmdb b/ic-bn-lib/test-data/geoip-city-test-db.mmdb new file mode 100644 index 0000000..e20846a Binary files /dev/null and b/ic-bn-lib/test-data/geoip-city-test-db.mmdb differ