use std::fmt;
use std::net::IpAddr;
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant, SystemTime};
use arc_swap::ArcSwap;
use compact_str::{CompactString, ToCompactString};
use maxminddb::{Mmap, Reader, geoip2};
use quick_cache::sync::{Cache, GuardResult};
use tracing::debug;
use super::DEFAULT_RESIDENT_MAX_BYTES;
#[cfg(feature = "geoip-download")]
use super::download::{Database, DatabaseFormat, Databases};
use super::extra;
use super::private::is_private;
use super::record::GeoIpRecord;
const DEFAULT_CAPACITY: usize = 100_000;
#[derive(Debug, thiserror::Error)]
#[non_exhaustive]
pub enum GeoIpLookupError {
#[error("{} is a {format} database, a format not supported by the lookup engine", .path.display())]
UnsupportedFormat {
path: PathBuf,
format: &'static str,
},
#[error("the {kind} database was provisioned with no files")]
NoFiles {
kind: &'static str,
},
#[error("could not open {}", .path.display())]
Open {
path: PathBuf,
#[source]
source: maxminddb::MaxMindDbError,
},
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct CacheConfig {
pub capacity: usize,
pub max_age: Option<Duration>,
pub resident_max_bytes: u64,
pub collect_extra_fields: bool,
}
impl Default for CacheConfig {
fn default() -> Self {
Self {
capacity: DEFAULT_CAPACITY,
max_age: None,
resident_max_bytes: DEFAULT_RESIDENT_MAX_BYTES,
collect_extra_fields: true,
}
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[non_exhaustive]
pub enum DatabaseBacking {
Resident,
Mapped,
}
impl DatabaseBacking {
#[must_use]
pub const fn label(self) -> &'static str {
match self {
Self::Resident => "resident",
Self::Mapped => "mapped",
}
}
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct DatabaseBackings {
pub city: Option<DatabaseBacking>,
pub asn: Option<DatabaseBacking>,
}
pub(super) enum Backing {
Resident(Reader<Vec<u8>>),
Mapped(Reader<Mmap>),
}
impl Backing {
const fn reported(&self) -> DatabaseBacking {
match self {
Self::Resident(_) => DatabaseBacking::Resident,
Self::Mapped(_) => DatabaseBacking::Mapped,
}
}
}
pub(super) struct Readers {
pub(super) city: Option<Arc<Backing>>,
pub(super) asn: Option<Arc<Backing>>,
pub(super) generation: u64,
}
impl Readers {
fn backings(&self) -> DatabaseBackings {
DatabaseBackings {
city: self.city.as_deref().map(Backing::reported),
asn: self.asn.as_deref().map(Backing::reported),
}
}
}
pub(super) struct Source {
pub(super) path: PathBuf,
pub(super) mtime: Option<SystemTime>,
}
pub(super) struct Sources {
pub(super) city: Option<Source>,
pub(super) asn: Option<Source>,
}
#[derive(Clone)]
pub(super) struct Entry {
loaded: Instant,
generation: u64,
record: Option<Arc<GeoIpRecord>>,
}
pub(super) struct Inner {
pub(super) readers: ArcSwap<Readers>,
generation: AtomicU64,
cache: Cache<IpAddr, Entry>,
pub(super) sources: Mutex<Sources>,
max_age: Option<Duration>,
pub(super) resident_max_bytes: u64,
collect_extra_fields: bool,
}
impl Inner {
pub(super) fn swap_readers(&self, mut readers: Readers) {
readers.generation = self.generation.load(Ordering::Acquire) + 1;
let generation = readers.generation;
let backings = readers.backings();
self.readers.store(Arc::new(readers));
self.generation.store(generation, Ordering::Release);
self.cache.clear();
telemetry::size(self.cache.len());
backing_telemetry::backings(backings);
}
}
#[derive(Clone)]
pub struct GeoIp {
pub(super) inner: Arc<Inner>,
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct DatabasePaths<'a> {
pub city: Option<&'a Path>,
pub asn: Option<&'a Path>,
}
impl<'a> DatabasePaths<'a> {
#[must_use]
pub const fn city_only(path: &'a Path) -> Self {
Self {
city: Some(path),
asn: None,
}
}
#[must_use]
pub const fn asn_only(path: &'a Path) -> Self {
Self {
city: None,
asn: Some(path),
}
}
}
impl GeoIp {
pub fn open(paths: DatabasePaths<'_>, cache: CacheConfig) -> Result<Self, GeoIpLookupError> {
let open = |path| open_source(path, cache.resident_max_bytes);
let (city_source, city_reader) = paths.city.map(open).transpose()?.unzip();
let (asn_source, asn_reader) = paths.asn.map(open).transpose()?.unzip();
let readers = Readers {
city: city_reader,
asn: asn_reader,
generation: 0,
};
backing_telemetry::backings(readers.backings());
Ok(Self {
inner: Arc::new(Inner {
readers: ArcSwap::from_pointee(readers),
generation: AtomicU64::new(0),
cache: Cache::new(cache.capacity),
sources: Mutex::new(Sources {
city: city_source,
asn: asn_source,
}),
max_age: cache.max_age,
resident_max_bytes: cache.resident_max_bytes,
collect_extra_fields: cache.collect_extra_fields,
}),
})
}
#[cfg(feature = "geoip-download")]
#[cfg_attr(docsrs, doc(cfg(feature = "geoip-download")))]
pub fn from_databases(
databases: &Databases,
cache: CacheConfig,
) -> Result<Self, GeoIpLookupError> {
let paths = DatabasePaths {
city: mmdb_path(databases.city.as_ref(), "city")?,
asn: mmdb_path(databases.asn.as_ref(), "asn")?,
};
Self::open(paths, cache)
}
#[must_use]
pub fn lookup(&self, ip: IpAddr) -> Option<Arc<GeoIpRecord>> {
if is_private(ip) {
return Some(GeoIpRecord::private_shared());
}
let started = telemetry::started();
let entry = match self.inner.cache.get_value_or_guard(&ip, None) {
GuardResult::Value(entry) if self.current(&entry) => {
telemetry::hit();
entry
}
GuardResult::Value(_) => {
let entry = self.miss(ip);
self.inner.cache.insert(ip, entry.clone());
telemetry::size(self.inner.cache.len());
entry
}
GuardResult::Guard(guard) => {
let entry = self.miss(ip);
drop(guard.insert(entry.clone()));
telemetry::size(self.inner.cache.len());
entry
}
GuardResult::Timeout => self.miss(ip),
};
telemetry::elapsed(started);
entry.record
}
pub async fn lookup_async(&self, ip: IpAddr) -> Option<Arc<GeoIpRecord>> {
if is_private(ip) {
return Some(GeoIpRecord::private_shared());
}
let started = telemetry::started();
let entry = match self.inner.cache.get_value_or_guard_async(&ip).await {
Ok(entry) if self.current(&entry) => {
telemetry::hit();
entry
}
Ok(_) => {
let entry = self.miss(ip);
self.inner.cache.insert(ip, entry.clone());
telemetry::size(self.inner.cache.len());
entry
}
Err(guard) => {
let entry = self.miss(ip);
drop(guard.insert(entry.clone()));
telemetry::size(self.inner.cache.len());
entry
}
};
telemetry::elapsed(started);
entry.record
}
#[must_use]
pub fn lookup_many(&self, ips: &[IpAddr]) -> Vec<Option<Arc<GeoIpRecord>>> {
ips.iter().map(|&ip| self.lookup(ip)).collect()
}
#[must_use]
pub fn cached_entries(&self) -> usize {
self.inner.cache.len()
}
#[must_use]
pub fn backings(&self) -> DatabaseBackings {
self.inner.readers.load().backings()
}
pub fn clear_cache(&self) {
self.inner.cache.clear();
telemetry::size(self.inner.cache.len());
}
fn miss(&self, ip: IpAddr) -> Entry {
telemetry::miss();
let readers = self.inner.readers.load_full();
let entry = Entry {
loaded: Instant::now(),
generation: readers.generation,
record: read(&readers, ip, self.inner.collect_extra_fields).map(Arc::new),
};
#[cfg(test)]
pause::point();
entry
}
fn current(&self, entry: &Entry) -> bool {
entry.generation == self.inner.generation.load(Ordering::Acquire)
&& self
.inner
.max_age
.is_none_or(|max_age| entry.loaded.elapsed() <= max_age)
}
}
#[cfg(test)]
pub(super) mod pause {
use std::sync::Mutex;
use std::sync::mpsc::{self, Receiver, Sender};
static HELD: Mutex<Option<(Sender<()>, Receiver<()>)>> = Mutex::new(None);
pub(in crate::geoip) fn hold() -> (Receiver<()>, Sender<()>) {
let (read, has_read) = mpsc::channel();
let (release, released) = mpsc::channel();
*HELD.lock().unwrap() = Some((read, released));
(has_read, release)
}
pub(super) fn point() {
let held = HELD.lock().unwrap().take();
if let Some((read, release)) = held {
let _ = read.send(());
let _ = release.recv();
}
}
}
impl fmt::Debug for GeoIp {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
let readers = self.inner.readers.load();
f.debug_struct("GeoIp")
.field("city", &readers.city.is_some())
.field("asn", &readers.asn.is_some())
.field("cached_entries", &self.inner.cache.len())
.finish_non_exhaustive()
}
}
fn read(readers: &Readers, ip: IpAddr, collect_extra: bool) -> Option<GeoIpRecord> {
let mut record = GeoIpRecord::default();
let mut found = false;
if let Some(backing) = readers.city.as_deref() {
found |= match backing {
Backing::Resident(reader) => read_city(reader, ip, &mut record, collect_extra),
Backing::Mapped(reader) => read_city(reader, ip, &mut record, collect_extra),
};
}
if let Some(backing) = readers.asn.as_deref() {
found |= match backing {
Backing::Resident(reader) => read_asn(reader, ip, &mut record, collect_extra),
Backing::Mapped(reader) => read_asn(reader, ip, &mut record, collect_extra),
};
}
found.then_some(record)
}
#[derive(serde::Deserialize)]
struct IpInfoRecord<'a> {
#[serde(borrow, default)]
country_code: Option<&'a str>,
#[serde(borrow, default)]
country: Option<&'a str>,
#[serde(borrow, default)]
continent_code: Option<&'a str>,
#[serde(borrow, default)]
continent: Option<&'a str>,
#[serde(borrow, default)]
asn: Option<&'a str>,
#[serde(borrow, default)]
as_name: Option<&'a str>,
#[serde(borrow, default)]
as_domain: Option<&'a str>,
}
fn is_ipinfo<S: AsRef<[u8]>>(reader: &Reader<S>) -> bool {
reader.metadata().database_type.starts_with("ipinfo")
}
pub(super) fn asn_number(text: &str) -> Option<u32> {
text.strip_prefix("AS").unwrap_or(text).parse().ok()
}
fn read_ipinfo<S: AsRef<[u8]>>(
reader: &Reader<S>,
ip: IpAddr,
record: &mut GeoIpRecord,
collect_extra: bool,
) -> bool {
let Some(result) = lookup_in(reader, ip, "ipinfo") else {
return false;
};
let found = match result.decode::<IpInfoRecord<'_>>() {
Ok(Some(found)) => found,
Ok(None) => return false,
Err(e) => {
debug!(%ip, error = %e, "IPinfo record did not decode");
return false;
}
};
record.country_code = found.country_code.map(CompactString::from);
record.country_name = found.country.map(CompactString::from);
record.continent_code = found.continent_code.map(CompactString::from);
record.continent_name = found.continent.map(CompactString::from);
record.autonomous_system_number = found.asn.and_then(asn_number);
record.autonomous_system_organization = found.as_name.map(CompactString::from);
record.as_domain = found.as_domain.map(CompactString::from);
let network = network_of(&result);
record.network.clone_from(&network);
record.asn_network = network;
if collect_extra {
extra::collect(&result, record);
}
true
}
fn read_city<S: AsRef<[u8]>>(
reader: &Reader<S>,
ip: IpAddr,
record: &mut GeoIpRecord,
collect_extra: bool,
) -> bool {
if is_ipinfo(reader) {
return read_ipinfo(reader, ip, record, collect_extra);
}
let Some(result) = lookup_in(reader, ip, "city") else {
return false;
};
let city = match result.decode::<geoip2::City>() {
Ok(Some(city)) => city,
Ok(None) => return false,
Err(e) => {
debug!(%ip, error = %e, "city record did not decode");
return false;
}
};
record.city_name = city.city.names.english.map(CompactString::from);
record.continent_code = city.continent.code.map(CompactString::from);
record.continent_name = city.continent.names.english.map(CompactString::from);
record.country_code = city.country.iso_code.map(CompactString::from);
record.country_name = city.country.names.english.map(CompactString::from);
if let Some(region) = city.subdivisions.last() {
record.region_name = region.names.english.map(CompactString::from);
record.region_code = region.iso_code.map(CompactString::from);
}
record.postal_code = city.postal.code.map(CompactString::from);
record.timezone = city.location.time_zone.map(CompactString::from);
record.latitude = city.location.latitude;
record.longitude = city.location.longitude;
record.metro_code = city.location.metro_code;
record.accuracy_radius = city.location.accuracy_radius;
record.network = network_of(&result);
if collect_extra {
extra::collect(&result, record);
}
true
}
fn read_asn<S: AsRef<[u8]>>(
reader: &Reader<S>,
ip: IpAddr,
record: &mut GeoIpRecord,
collect_extra: bool,
) -> bool {
if is_ipinfo(reader) {
return read_ipinfo(reader, ip, record, collect_extra);
}
let Some(result) = lookup_in(reader, ip, "asn") else {
return false;
};
let asn = match result.decode::<geoip2::Asn>() {
Ok(Some(asn)) => asn,
Ok(None) => return false,
Err(e) => {
debug!(%ip, error = %e, "ASN record did not decode");
return false;
}
};
record.autonomous_system_number = asn.autonomous_system_number;
record.autonomous_system_organization =
asn.autonomous_system_organization.map(CompactString::from);
record.asn_network = network_of(&result);
if collect_extra {
extra::collect(&result, record);
}
true
}
fn lookup_in<'a, S: AsRef<[u8]>>(
reader: &'a Reader<S>,
ip: IpAddr,
kind: &'static str,
) -> Option<maxminddb::LookupResult<'a, S>> {
match reader.lookup(ip) {
Ok(result) => Some(result),
Err(e) => {
debug!(%ip, kind, error = %e, "GeoIP database refused the address");
None
}
}
}
fn network_of<S: AsRef<[u8]>>(result: &maxminddb::LookupResult<'_, S>) -> Option<CompactString> {
result
.network()
.ok()
.map(|network| network.to_compact_string())
}
fn open_source(
path: &Path,
resident_max_bytes: u64,
) -> Result<(Source, Arc<Backing>), GeoIpLookupError> {
let mtime = file_mtime(path);
let backing = open_backing(path, resident_max_bytes)?;
Ok((
Source {
path: path.to_path_buf(),
mtime,
},
Arc::new(backing),
))
}
pub(super) fn open_backing(
path: &Path,
resident_max_bytes: u64,
) -> Result<Backing, GeoIpLookupError> {
if file_len(path).is_none_or(|len| len > resident_max_bytes) {
return open_reader(path).map(Backing::Mapped);
}
Reader::open_readfile(path)
.map(Backing::Resident)
.map_err(|source| GeoIpLookupError::Open {
path: path.to_path_buf(),
source,
})
}
#[expect(
unsafe_code,
reason = "the reader offers no safe mmap constructor, and the provisioning \
half never writes a database in place, which is what makes the \
mapping sound"
)]
pub(super) fn open_reader(path: &Path) -> Result<Reader<Mmap>, GeoIpLookupError> {
unsafe { Reader::open_mmap(path) }.map_err(|source| GeoIpLookupError::Open {
path: path.to_path_buf(),
source,
})
}
pub(super) fn file_mtime(path: &Path) -> Option<SystemTime> {
std::fs::metadata(path)
.and_then(|metadata| metadata.modified())
.ok()
}
fn file_len(path: &Path) -> Option<u64> {
std::fs::metadata(path).map(|metadata| metadata.len()).ok()
}
#[cfg(feature = "geoip-download")]
fn mmdb_path<'a>(
database: Option<&'a Database>,
kind: &'static str,
) -> Result<Option<&'a Path>, GeoIpLookupError> {
let Some(database) = database else {
return Ok(None);
};
let Some(path) = database.files.first() else {
return Err(GeoIpLookupError::NoFiles { kind });
};
if database.format != DatabaseFormat::Mmdb {
return Err(GeoIpLookupError::UnsupportedFormat {
path: path.clone(),
format: database.format.name(),
});
}
Ok(Some(path.as_path()))
}
#[cfg(feature = "metrics-lookup")]
mod telemetry {
use std::time::Instant;
const TYPE: &str = "geoip";
#[derive(Clone, Copy)]
pub(super) struct Started(Instant);
pub(super) fn started() -> Started {
Started(Instant::now())
}
pub(super) fn elapsed(started: Started) {
metrics::histogram!("enrichment_duration_seconds", "type" => TYPE)
.record(started.0.elapsed().as_secs_f64());
}
pub(super) fn hit() {
metrics::counter!("enrichment_cache_hits_total", "type" => TYPE).increment(1);
}
pub(super) fn miss() {
metrics::counter!("enrichment_cache_misses_total", "type" => TYPE).increment(1);
}
#[allow(clippy::cast_precision_loss)]
pub(super) fn size(entries: usize) {
metrics::gauge!("enrichment_cache_size", "type" => TYPE).set(entries as f64);
}
}
#[cfg(not(feature = "metrics-lookup"))]
mod telemetry {
#[derive(Clone, Copy)]
pub(super) struct Started;
pub(super) const fn started() -> Started {
Started
}
pub(super) const fn elapsed(_started: Started) {}
pub(super) const fn hit() {}
pub(super) const fn miss() {}
pub(super) const fn size(_entries: usize) {}
}
#[cfg(feature = "metrics")]
mod backing_telemetry {
use super::{DatabaseBacking, DatabaseBackings};
const TYPE: &str = "geoip";
pub(super) fn backings(chosen: DatabaseBackings) {
if let Some(city) = chosen.city {
report("city", city);
}
if let Some(asn) = chosen.asn {
report("asn", asn);
}
}
fn report(kind: &'static str, chosen: DatabaseBacking) {
for candidate in [DatabaseBacking::Resident, DatabaseBacking::Mapped] {
metrics::gauge!(
"enrichment_database_backing",
"type" => TYPE,
"kind" => kind,
"backing" => candidate.label(),
)
.set(if candidate == chosen { 1.0 } else { 0.0 });
}
}
}
#[cfg(not(feature = "metrics"))]
mod backing_telemetry {
use super::DatabaseBackings;
pub(super) const fn backings(_chosen: DatabaseBackings) {}
}
#[cfg(test)]
mod tests {
use std::sync::Barrier;
use std::thread;
use super::*;
use crate::geoip::{ExtraFields, ExtraValue};
const CITY_DB: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/tests/data/city-test.mmdb");
const ASN_DB: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/tests/data/asn-test.mmdb");
const BOXFORD: &str = "2.125.160.216";
const LINKOPING: &str = "89.160.20.112";
const LINKOPING_NAME: &str = "Link\u{f6}ping";
const ABSENT: &str = "8.8.8.8";
fn both() -> GeoIp {
GeoIp::open(
DatabasePaths {
city: Some(Path::new(CITY_DB)),
asn: Some(Path::new(ASN_DB)),
},
CacheConfig::default(),
)
.unwrap()
}
const IPINFO_DB: &str = concat!(
env!("CARGO_MANIFEST_DIR"),
"/tests/data/IPinfo-Lite-Test.mmdb"
);
const CITY_RICH_DB: &str = concat!(
env!("CARGO_MANIFEST_DIR"),
"/tests/data/city-rich-test.mmdb"
);
const ISP_DB: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/tests/data/isp-test.mmdb");
fn ip(literal: &str) -> IpAddr {
literal.parse().unwrap()
}
fn ipinfo() -> GeoIp {
GeoIp::open(
DatabasePaths::city_only(Path::new(IPINFO_DB)),
CacheConfig::default(),
)
.unwrap()
}
#[test]
fn an_ipv6_address_resolves_from_the_city_database() {
let record = both().lookup(ip("2001:218::1")).unwrap();
assert_eq!(record.country_code.as_deref(), Some("JP"));
assert_eq!(record.continent_code.as_deref(), Some("AS"));
assert_eq!(record.continent_name.as_deref(), Some("Asia"));
assert_eq!(record.network.as_deref(), Some("2001:218::/32"));
assert!(!record.is_private);
}
#[test]
fn an_ipv6_address_resolves_from_the_asn_database() {
let record = both().lookup(ip("2001:2000::1")).unwrap();
assert_eq!(record.autonomous_system_number, Some(1299));
assert_eq!(
record.autonomous_system_organization.as_deref(),
Some("TeliaSonera International Carrier")
);
assert_eq!(record.asn_network.as_deref(), Some("2001:2000::/20"));
}
#[test]
fn an_ipv6_address_only_one_database_holds_still_answers() {
let city_only = both().lookup(ip("2001:218::1")).unwrap();
assert!(city_only.country_code.is_some());
assert!(city_only.autonomous_system_number.is_none());
let asn_only = both().lookup(ip("2001:2000::1")).unwrap();
assert!(asn_only.country_code.is_none());
assert!(asn_only.autonomous_system_number.is_some());
}
#[test]
fn a_v6_literal_and_its_long_form_are_one_cache_entry() {
let geoip = both();
let short = geoip.lookup(ip("2001:218::1")).unwrap();
let long = geoip
.lookup(ip("2001:0218:0000:0000:0000:0000:0000:0001"))
.unwrap();
assert!(Arc::ptr_eq(&short, &long));
}
#[test]
fn an_ipinfo_database_resolves_location_and_network_together() {
let record = ipinfo().lookup(ip("8.8.8.8")).unwrap();
assert_eq!(record.country_code.as_deref(), Some("US"));
assert_eq!(record.country_name.as_deref(), Some("United States"));
assert_eq!(record.continent_code.as_deref(), Some("NA"));
assert_eq!(record.autonomous_system_number, Some(15169));
assert_eq!(
record.autonomous_system_organization.as_deref(),
Some("Google LLC")
);
assert_eq!(record.network.as_deref(), Some("8.8.8.0/24"));
assert_eq!(record.asn_network.as_deref(), Some("8.8.8.0/24"));
assert_eq!(record.as_domain.as_deref(), Some("google.com"));
}
#[test]
fn an_ipinfo_database_answers_ipv6() {
let record = ipinfo().lookup(ip("2606:4700:4700::1111")).unwrap();
assert_eq!(record.country_code.as_deref(), Some("US"));
assert_eq!(record.autonomous_system_number, Some(13335));
}
#[test]
fn an_ipinfo_record_without_network_fields_still_resolves() {
let record = ipinfo().lookup(ip("45.45.45.45")).unwrap();
assert_eq!(record.country_code.as_deref(), Some("AU"));
assert!(record.autonomous_system_number.is_none());
assert!(record.autonomous_system_organization.is_none());
}
#[test]
fn an_address_the_ipinfo_database_omits_reads_as_absent() {
assert!(ipinfo().lookup(ip("9.9.9.9")).is_none());
}
#[test]
fn an_ipinfo_database_serves_the_asn_slot_too() {
let geoip = GeoIp::open(
DatabasePaths::asn_only(Path::new(IPINFO_DB)),
CacheConfig::default(),
)
.unwrap();
let record = geoip.lookup(ip("1.1.1.1")).unwrap();
assert_eq!(record.autonomous_system_number, Some(13335));
assert_eq!(record.country_code.as_deref(), Some("AU"));
}
#[test]
fn an_as_prefixed_number_parses_and_a_malformed_one_does_not() {
assert_eq!(asn_number("AS15169"), Some(15169));
assert_eq!(asn_number("15169"), Some(15169));
assert_eq!(asn_number("AS"), None);
assert_eq!(asn_number(""), None);
assert_eq!(asn_number("ASN15169"), None);
}
#[test]
fn a_known_city_resolves_every_location_field() {
let record = both().lookup(ip(BOXFORD)).unwrap();
assert_eq!(record.city_name.as_deref(), Some("Boxford"));
assert_eq!(record.continent_code.as_deref(), Some("EU"));
assert_eq!(record.continent_name.as_deref(), Some("Europe"));
assert_eq!(record.country_code.as_deref(), Some("GB"));
assert_eq!(record.country_name.as_deref(), Some("United Kingdom"));
assert_eq!(record.postal_code.as_deref(), Some("OX1"));
assert_eq!(record.timezone.as_deref(), Some("Europe/London"));
assert_eq!(record.accuracy_radius, Some(100));
assert_eq!(record.network.as_deref(), Some("2.125.160.216/29"));
assert!(!record.is_private);
assert!(record.metro_code.is_none());
let latitude = record.latitude.unwrap();
let longitude = record.longitude.unwrap();
assert!((latitude - 51.75).abs() < 1e-6, "{latitude}");
assert!((longitude + 1.25).abs() < 1e-6, "{longitude}");
}
#[test]
fn the_region_is_the_last_subdivision_not_the_first() {
let record = both().lookup(ip(BOXFORD)).unwrap();
assert_eq!(record.region_name.as_deref(), Some("West Berkshire"));
assert_eq!(record.region_code.as_deref(), Some("WBK"));
assert_ne!(record.region_name.as_deref(), Some("England"));
assert_ne!(record.region_code.as_deref(), Some("ENG"));
}
#[test]
fn a_known_asn_resolves_its_number_and_organisation() {
let record = both().lookup(ip("1.0.0.1")).unwrap();
assert_eq!(record.autonomous_system_number, Some(15169));
assert_eq!(
record.autonomous_system_organization.as_deref(),
Some("Google Inc.")
);
assert_eq!(record.asn_network.as_deref(), Some("1.0.0.0/24"));
assert!(record.city_name.is_none());
assert!(record.network.is_none());
}
#[test]
fn the_two_databases_report_their_own_prefixes() {
let record = both().lookup(ip(LINKOPING)).unwrap();
assert_eq!(record.city_name.as_deref(), Some(LINKOPING_NAME));
assert_eq!(record.country_code.as_deref(), Some("SE"));
assert_eq!(
record.autonomous_system_organization.as_deref(),
Some("Bredband2 AB")
);
assert_eq!(record.network.as_deref(), Some("89.160.20.112/28"));
assert_eq!(record.asn_network.as_deref(), Some("89.160.0.0/17"));
assert_ne!(record.network, record.asn_network);
}
#[test]
fn a_reserved_address_answers_without_touching_the_cache() {
let geoip = both();
for reserved in ["10.0.0.1", "127.0.0.1", "::1", "100.64.0.1", "240.0.0.1"] {
let record = geoip.lookup(ip(reserved)).unwrap();
assert!(record.is_private, "{reserved}");
assert!(record.country_code.is_none(), "{reserved}");
}
assert_eq!(geoip.cached_entries(), 0);
}
#[test]
fn every_reserved_answer_is_the_same_allocation() {
let geoip = both();
let first = geoip.lookup(ip("10.0.0.1")).unwrap();
let second = geoip.lookup(ip("192.168.1.1")).unwrap();
assert!(Arc::ptr_eq(&first, &second));
}
#[test]
fn an_address_no_database_holds_is_cached_as_a_miss() {
let geoip = both();
assert!(geoip.lookup(ip(ABSENT)).is_none());
assert_eq!(geoip.cached_entries(), 1);
assert!(geoip.lookup(ip(ABSENT)).is_none());
assert_eq!(geoip.cached_entries(), 1);
}
#[test]
fn a_repeat_lookup_returns_the_first_answer_rather_than_a_copy() {
let geoip = both();
let first = geoip.lookup(ip(BOXFORD)).unwrap();
let second = geoip.lookup(ip(BOXFORD)).unwrap();
assert!(Arc::ptr_eq(&first, &second));
assert_eq!(geoip.cached_entries(), 1);
}
#[test]
fn concurrent_misses_on_one_address_share_a_single_read() {
const THREADS: usize = 16;
let geoip = both();
let barrier = Arc::new(Barrier::new(THREADS));
let records: Vec<Arc<GeoIpRecord>> = thread::scope(|scope| {
let handles: Vec<_> = (0..THREADS)
.map(|_| {
let geoip = geoip.clone();
let barrier = Arc::clone(&barrier);
scope.spawn(move || {
barrier.wait();
geoip.lookup(ip(BOXFORD)).unwrap()
})
})
.collect();
handles
.into_iter()
.map(|handle| handle.join().unwrap())
.collect::<Vec<_>>()
});
for record in &records {
assert!(Arc::ptr_eq(&records[0], record));
}
assert_eq!(geoip.cached_entries(), 1);
}
#[test]
fn a_batch_resolves_each_repeat_once_and_keeps_its_order() {
let geoip = both();
let batch = [
ip(BOXFORD),
ip(LINKOPING),
ip(BOXFORD),
ip("10.0.0.1"),
ip(BOXFORD),
ip(ABSENT),
];
let answers = geoip.lookup_many(&batch);
assert_eq!(answers.len(), batch.len());
assert_eq!(
answers[0].as_ref().unwrap().city_name.as_deref(),
Some("Boxford")
);
assert_eq!(
answers[1].as_ref().unwrap().city_name.as_deref(),
Some(LINKOPING_NAME)
);
assert!(answers[3].as_ref().unwrap().is_private);
assert!(answers[5].is_none());
assert!(Arc::ptr_eq(
answers[0].as_ref().unwrap(),
answers[2].as_ref().unwrap()
));
assert!(Arc::ptr_eq(
answers[0].as_ref().unwrap(),
answers[4].as_ref().unwrap()
));
assert_eq!(geoip.cached_entries(), 3);
}
#[tokio::test]
async fn the_async_entry_point_reads_the_same_cache() {
let geoip = both();
let synchronous = geoip.lookup(ip(BOXFORD)).unwrap();
let asynchronous = geoip.lookup_async(ip(BOXFORD)).await.unwrap();
assert!(Arc::ptr_eq(&synchronous, &asynchronous));
assert_eq!(geoip.cached_entries(), 1);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn concurrent_async_misses_on_one_address_share_a_single_read() {
const TASKS: usize = 16;
let geoip = both();
let gate = Arc::new(tokio::sync::Barrier::new(TASKS));
let tasks: Vec<_> = (0..TASKS)
.map(|_| {
let geoip = geoip.clone();
let gate = Arc::clone(&gate);
tokio::spawn(async move {
gate.wait().await;
geoip.lookup_async(ip(BOXFORD)).await.unwrap()
})
})
.collect();
let mut records = Vec::with_capacity(TASKS);
for task in tasks {
records.push(task.await.unwrap());
}
for record in &records {
assert!(Arc::ptr_eq(&records[0], record));
}
assert_eq!(geoip.cached_entries(), 1);
}
#[tokio::test]
async fn the_async_entry_point_short_circuits_reserved_space_too() {
let geoip = both();
let record = geoip.lookup_async(ip("192.168.0.1")).await.unwrap();
assert!(record.is_private);
assert_eq!(geoip.cached_entries(), 0);
}
#[test]
fn an_answer_inside_the_age_limit_is_served_from_the_cache() {
let geoip = GeoIp::open(
DatabasePaths::city_only(Path::new(CITY_DB)),
CacheConfig {
max_age: Some(Duration::from_secs(3600)),
..CacheConfig::default()
},
)
.unwrap();
let first = geoip.lookup(ip(BOXFORD)).unwrap();
let second = geoip.lookup(ip(BOXFORD)).unwrap();
assert!(Arc::ptr_eq(&first, &second));
}
#[test]
fn an_age_limit_re_reads_the_answer_it_expired() {
let geoip = GeoIp::open(
DatabasePaths::city_only(Path::new(CITY_DB)),
CacheConfig {
max_age: Some(Duration::ZERO),
..CacheConfig::default()
},
)
.unwrap();
let first = geoip.lookup(ip(BOXFORD)).unwrap();
thread::sleep(Duration::from_millis(1));
let second = geoip.lookup(ip(BOXFORD)).unwrap();
assert!(!Arc::ptr_eq(&first, &second));
assert_eq!(first.city_name, second.city_name);
assert_eq!(geoip.cached_entries(), 1);
}
#[test]
fn a_missing_database_leaves_the_other_half_answering() {
let city_only = GeoIp::open(
DatabasePaths::city_only(Path::new(CITY_DB)),
CacheConfig::default(),
)
.unwrap();
let record = city_only.lookup(ip(LINKOPING)).unwrap();
assert_eq!(record.city_name.as_deref(), Some(LINKOPING_NAME));
assert!(record.autonomous_system_number.is_none());
assert!(record.asn_network.is_none());
}
#[test]
fn an_enricher_over_no_database_answers_nothing_and_does_not_fail() {
let empty = GeoIp::open(DatabasePaths::default(), CacheConfig::default()).unwrap();
assert!(empty.lookup(ip(BOXFORD)).is_none());
assert!(empty.lookup(ip("10.0.0.1")).unwrap().is_private);
}
#[test]
fn opening_something_that_is_not_a_database_names_the_file() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("not-a-database.mmdb");
std::fs::write(&path, b"not a MaxMind DB").unwrap();
let err = GeoIp::open(DatabasePaths::city_only(&path), CacheConfig::default()).unwrap_err();
assert!(matches!(err, GeoIpLookupError::Open { .. }));
assert!(err.to_string().contains("not-a-database.mmdb"));
}
#[test]
fn clearing_the_cache_leaves_the_readers_answering() {
let geoip = both();
assert!(geoip.lookup(ip(BOXFORD)).is_some());
assert_eq!(geoip.cached_entries(), 1);
geoip.clear_cache();
assert_eq!(geoip.cached_entries(), 0);
assert!(geoip.lookup(ip(BOXFORD)).is_some());
}
#[test]
fn a_clone_shares_the_cache_rather_than_building_its_own() {
let geoip = both();
let worker = geoip.clone();
let first = geoip.lookup(ip(BOXFORD)).unwrap();
let second = worker.lookup(ip(BOXFORD)).unwrap();
assert!(Arc::ptr_eq(&first, &second));
assert_eq!(geoip.cached_entries(), 1);
assert_eq!(worker.cached_entries(), 1);
}
#[test]
fn debug_reports_the_cache_without_dumping_the_databases() {
let geoip = both();
let rendered = format!("{geoip:?}");
assert!(rendered.contains("cached_entries"), "{rendered}");
assert!(!rendered.contains("city-test"), "{rendered}");
}
fn size_of(path: &str) -> u64 {
std::fs::metadata(path).unwrap().len()
}
fn both_under(resident_max_bytes: u64) -> GeoIp {
GeoIp::open(
DatabasePaths {
city: Some(Path::new(CITY_DB)),
asn: Some(Path::new(ASN_DB)),
},
CacheConfig {
resident_max_bytes,
..CacheConfig::default()
},
)
.unwrap()
}
#[test]
fn a_database_under_the_ceiling_is_held_in_memory() {
assert_eq!(
both().backings(),
DatabaseBackings {
city: Some(DatabaseBacking::Resident),
asn: Some(DatabaseBacking::Resident),
}
);
}
#[test]
fn a_database_over_the_ceiling_is_mapped() {
let geoip = both_under(size_of(ASN_DB));
assert_eq!(
geoip.backings(),
DatabaseBackings {
city: Some(DatabaseBacking::Mapped),
asn: Some(DatabaseBacking::Resident),
}
);
}
#[test]
fn the_ceiling_admits_a_file_of_exactly_its_size() {
let exact = size_of(CITY_DB);
assert_eq!(
both_under(exact).backings().city,
Some(DatabaseBacking::Resident)
);
assert_eq!(
both_under(exact - 1).backings().city,
Some(DatabaseBacking::Mapped)
);
}
#[test]
fn a_zero_ceiling_maps_every_database() {
assert_eq!(
both_under(0).backings(),
DatabaseBackings {
city: Some(DatabaseBacking::Mapped),
asn: Some(DatabaseBacking::Mapped),
}
);
}
#[test]
fn the_two_backings_answer_the_same_record() {
let resident = both();
let mapped = both_under(0);
for literal in [BOXFORD, LINKOPING, "1.0.0.1", "2001:218::1", ABSENT] {
assert_eq!(
resident.lookup(ip(literal)).map(|record| (*record).clone()),
mapped.lookup(ip(literal)).map(|record| (*record).clone()),
"{literal}"
);
}
}
#[test]
fn an_enricher_over_no_database_reports_no_backing() {
let empty = GeoIp::open(DatabasePaths::default(), CacheConfig::default()).unwrap();
assert_eq!(empty.backings(), DatabaseBackings::default());
}
fn city_rich() -> GeoIp {
GeoIp::open(
DatabasePaths::city_only(Path::new(CITY_RICH_DB)),
CacheConfig::default(),
)
.unwrap()
}
fn extra(record: &GeoIpRecord, path: &str) -> ExtraValue {
record
.extra
.get(path)
.unwrap_or_else(|| panic!("{path} is in the map: {:?}", record.extra))
.clone()
}
fn extra_str(record: &GeoIpRecord, path: &str) -> CompactString {
match extra(record, path) {
ExtraValue::Str(text) => text,
other => panic!("{path} is text, not {other:?}"),
}
}
#[test]
fn a_field_no_typed_field_names_reaches_the_map() {
let record = city_rich().lookup(ip(BOXFORD)).unwrap();
assert_eq!(extra_str(&record, "city.names.de"), "Boxford");
assert_eq!(
extra(&record, "city.geoname_id"),
ExtraValue::UInt(2_655_045)
);
assert_eq!(extra(&record, "city.confidence"), ExtraValue::UInt(25));
assert_eq!(
extra(&record, "country.is_in_european_union"),
ExtraValue::Bool(false)
);
assert_eq!(
extra(&record, "location.average_income"),
ExtraValue::UInt(32_323)
);
assert_eq!(extra(&record, "traits.is_anycast"), ExtraValue::Bool(true));
assert_eq!(extra_str(&record, "traits.user_type"), "residential");
assert_eq!(extra_str(&record, "represented_country.type"), "military");
assert_eq!(extra_str(&record, "registered_country.iso_code"), "GB");
}
#[test]
fn a_nested_name_keeps_its_own_characters() {
let record = city_rich().lookup(ip(BOXFORD)).unwrap();
assert_eq!(
extra_str(&record, "city.names.ja"),
"\u{30dc}\u{30c3}\u{30af}\u{30b9}\u{30d5}\u{30a9}\u{30fc}\u{30c9}"
);
}
#[test]
fn the_typed_fields_still_resolve_beside_the_map() {
let record = city_rich().lookup(ip(BOXFORD)).unwrap();
assert_eq!(record.city_name.as_deref(), Some("Boxford"));
assert_eq!(record.continent_name.as_deref(), Some("Europe"));
assert_eq!(record.country_code.as_deref(), Some("GB"));
assert_eq!(record.region_name.as_deref(), Some("West Berkshire"));
assert_eq!(record.region_code.as_deref(), Some("WBK"));
assert_eq!(record.postal_code.as_deref(), Some("OX1"));
assert_eq!(record.timezone.as_deref(), Some("Europe/London"));
assert_eq!(record.accuracy_radius, Some(100));
assert_eq!(record.network.as_deref(), Some("2.125.160.216/29"));
}
#[test]
fn a_field_a_typed_field_took_is_not_in_the_map_as_well() {
let record = city_rich().lookup(ip(BOXFORD)).unwrap();
for path in [
"city.names.en",
"continent.code",
"continent.names.en",
"country.iso_code",
"country.names.en",
"postal.code",
"location.time_zone",
"location.latitude",
"location.longitude",
"location.accuracy_radius",
] {
assert!(record.extra.get(path).is_none(), "{path}");
}
}
#[test]
fn an_array_is_carried_under_its_indices() {
let record = city_rich().lookup(ip(BOXFORD)).unwrap();
assert_eq!(extra_str(&record, "subdivisions.0.names.en"), "England");
assert_eq!(extra_str(&record, "subdivisions.0.iso_code"), "ENG");
assert_eq!(
extra(&record, "subdivisions.0.geoname_id"),
ExtraValue::UInt(6_269_131)
);
assert!(record.extra.get("subdivisions.1.names.en").is_none());
assert!(record.extra.get("subdivisions.1.iso_code").is_none());
assert_eq!(
extra(&record, "subdivisions.1.confidence"),
ExtraValue::UInt(40)
);
}
#[test]
fn the_isp_edition_delivers_the_fields_it_is_bought_for() {
let geoip = GeoIp::open(
DatabasePaths::asn_only(Path::new(ISP_DB)),
CacheConfig::default(),
)
.unwrap();
let record = geoip.lookup(ip("1.0.0.1")).unwrap();
assert_eq!(record.autonomous_system_number, Some(15169));
assert_eq!(
record.autonomous_system_organization.as_deref(),
Some("Google Inc.")
);
assert_eq!(extra_str(&record, "isp"), "Telstra Mobile");
assert_eq!(extra_str(&record, "organization"), "Telstra Mobile Data");
assert_eq!(extra_str(&record, "mobile_country_code"), "505");
assert_eq!(extra_str(&record, "mobile_network_code"), "01");
assert!(record.extra.get("autonomous_system_number").is_none());
assert!(record.extra.get("autonomous_system_organization").is_none());
}
#[test]
fn the_lean_fixtures_were_discarding_fields_too() {
let record = both().lookup(ip(BOXFORD)).unwrap();
assert_eq!(extra_str(&record, "subdivisions.0.names.en"), "England");
assert_eq!(extra_str(&record, "subdivisions.0.iso_code"), "ENG");
}
#[test]
fn a_source_the_record_fully_names_carries_no_map() {
let record = ipinfo().lookup(ip("8.8.8.8")).unwrap();
assert_eq!(record.continent_name.as_deref(), Some("North America"));
assert!(record.extra.is_empty(), "{:?}", record.extra);
}
#[test]
fn a_reserved_address_carries_no_source_fields() {
let record = city_rich().lookup(ip("10.0.0.1")).unwrap();
assert!(record.is_private);
assert!(record.extra.is_empty());
}
#[test]
fn both_databases_contribute_to_one_map() {
let geoip = GeoIp::open(
DatabasePaths {
city: Some(Path::new(CITY_RICH_DB)),
asn: Some(Path::new(ISP_DB)),
},
CacheConfig::default(),
)
.unwrap();
let record = geoip.lookup(ip(LINKOPING)).unwrap();
assert_eq!(extra_str(&record, "isp"), "Bredband2 AB");
assert_eq!(extra_str(&record, "organization"), "Bredband2 Customer");
let boxford = geoip.lookup(ip(BOXFORD)).unwrap();
assert_eq!(extra_str(&boxford, "traits.user_type"), "residential");
}
fn without_extras() -> CacheConfig {
CacheConfig {
collect_extra_fields: false,
..CacheConfig::default()
}
}
fn lean(paths: DatabasePaths<'_>) -> GeoIp {
GeoIp::open(paths, without_extras()).unwrap()
}
#[test]
fn the_source_fields_are_kept_unless_the_setting_says_otherwise() {
assert!(CacheConfig::default().collect_extra_fields);
}
#[test]
fn the_source_fields_are_dropped_when_collection_is_off() {
let rich = city_rich().lookup(ip(BOXFORD)).unwrap();
let lean = lean(DatabasePaths::city_only(Path::new(CITY_RICH_DB)))
.lookup(ip(BOXFORD))
.unwrap();
assert!(!rich.extra.is_empty());
assert!(lean.extra.is_empty(), "{:?}", lean.extra);
let mut typed_only = (*rich).clone();
typed_only.extra = ExtraFields::new();
assert_eq!(typed_only, *lean);
}
#[test]
fn the_typed_fields_resolve_with_collection_off() {
let record = lean(DatabasePaths::city_only(Path::new(CITY_RICH_DB)))
.lookup(ip(BOXFORD))
.unwrap();
assert_eq!(record.city_name.as_deref(), Some("Boxford"));
assert_eq!(record.continent_name.as_deref(), Some("Europe"));
assert_eq!(record.country_code.as_deref(), Some("GB"));
assert_eq!(record.region_name.as_deref(), Some("West Berkshire"));
assert_eq!(record.region_code.as_deref(), Some("WBK"));
assert_eq!(record.postal_code.as_deref(), Some("OX1"));
assert_eq!(record.timezone.as_deref(), Some("Europe/London"));
assert_eq!(record.accuracy_radius, Some(100));
assert_eq!(record.network.as_deref(), Some("2.125.160.216/29"));
assert!(record.extra.is_empty());
}
#[test]
fn the_asn_half_drops_its_fields_and_keeps_its_typed_ones() {
let geoip = lean(DatabasePaths::asn_only(Path::new(ISP_DB)));
let record = geoip.lookup(ip("1.0.0.1")).unwrap();
assert!(record.extra.is_empty(), "{:?}", record.extra);
assert_eq!(record.autonomous_system_number, Some(15169));
assert_eq!(
record.autonomous_system_organization.as_deref(),
Some("Google Inc.")
);
assert_eq!(record.asn_network.as_deref(), Some("1.0.0.0/24"));
}
#[test]
fn an_ipinfo_database_answers_with_collection_off() {
let geoip = lean(DatabasePaths::city_only(Path::new(IPINFO_DB)));
let record = geoip.lookup(ip("8.8.8.8")).unwrap();
assert_eq!(record.country_code.as_deref(), Some("US"));
assert_eq!(record.autonomous_system_number, Some(15169));
assert_eq!(record.as_domain.as_deref(), Some("google.com"));
assert!(record.extra.is_empty());
}
#[test]
fn neither_database_contributes_a_field_when_collection_is_off() {
let geoip = lean(DatabasePaths {
city: Some(Path::new(CITY_RICH_DB)),
asn: Some(Path::new(ISP_DB)),
});
let linkoping = geoip.lookup(ip(LINKOPING)).unwrap();
assert_eq!(
linkoping.autonomous_system_organization.as_deref(),
Some("Bredband2 AB")
);
assert!(linkoping.extra.is_empty(), "{:?}", linkoping.extra);
let boxford = geoip.lookup(ip(BOXFORD)).unwrap();
assert_eq!(boxford.city_name.as_deref(), Some("Boxford"));
assert!(boxford.extra.is_empty(), "{:?}", boxford.extra);
}
#[cfg(feature = "geoip-download")]
#[test]
fn a_csv_database_is_refused_by_name() {
let databases = Databases {
city: Some(Database {
format: DatabaseFormat::Csv,
files: vec![
PathBuf::from("/var/lib/geoip/city-ipv4.csv"),
PathBuf::from("/var/lib/geoip/city-ipv6.csv"),
],
}),
asn: None,
};
let err = GeoIp::from_databases(&databases, CacheConfig::default()).unwrap_err();
assert!(matches!(err, GeoIpLookupError::UnsupportedFormat { .. }));
let message = err.to_string();
assert!(
message.contains("/var/lib/geoip/city-ipv4.csv"),
"{message}"
);
assert!(message.contains("CSV"), "{message}");
}
#[cfg(feature = "geoip-download")]
#[test]
fn a_provisioned_mmdb_database_opens_through_the_download_types() {
let databases = Databases {
city: Some(Database {
format: DatabaseFormat::Mmdb,
files: vec![PathBuf::from(CITY_DB)],
}),
asn: Some(Database {
format: DatabaseFormat::Mmdb,
files: vec![PathBuf::from(ASN_DB)],
}),
};
let geoip = GeoIp::from_databases(&databases, CacheConfig::default()).unwrap();
let record = geoip.lookup(ip(LINKOPING)).unwrap();
assert_eq!(record.country_code.as_deref(), Some("SE"));
assert_eq!(record.autonomous_system_number, Some(29518));
}
#[cfg(feature = "geoip-download")]
#[test]
fn a_database_resolved_to_no_files_is_refused_by_kind() {
let databases = Databases {
city: None,
asn: Some(Database {
format: DatabaseFormat::Mmdb,
files: Vec::new(),
}),
};
let err = GeoIp::from_databases(&databases, CacheConfig::default()).unwrap_err();
assert!(matches!(err, GeoIpLookupError::NoFiles { kind: "asn" }));
}
}