use std::fs;
use std::path::Path;
use super::{DatabaseFormat, GeoIpDownloadError, Kind};
use crate::geoip::config::AutoDownloadConfig;
use crate::table::{Schema, TableFormat};
#[cfg(feature = "geoip-lookup")]
use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
#[cfg(not(feature = "geoip-lookup"))]
use tracing::debug;
const PERCENT: u64 = 100;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) struct Guard {
probe: Option<Kind>,
parses: Option<TableFormat>,
declared: Option<usize>,
min_size_percent: u8,
}
impl Guard {
#[cfg(test)]
pub(super) const OFF: Self = Self {
probe: None,
parses: None,
declared: None,
min_size_percent: 0,
};
pub(super) const fn new(kind: Kind, auto: &AutoDownloadConfig) -> Self {
Self {
probe: if auto.verify_content {
Some(kind)
} else {
None
},
parses: None,
declared: None,
min_size_percent: auto.min_size_percent,
}
}
pub(crate) fn for_table(
format: TableFormat,
schema: &Schema,
auto: &AutoDownloadConfig,
) -> Self {
Self {
probe: None,
parses: if auto.verify_content {
Some(format)
} else {
None
},
declared: match schema {
Schema::Named(names) => Some(names.len()),
Schema::Auto => None,
},
min_size_percent: auto.min_size_percent,
}
}
pub(super) fn admit(
self,
staged: &Path,
dest: &Path,
format: DatabaseFormat,
) -> Result<(), GeoIpDownloadError> {
self.check_size(staged, dest)?;
if let Some(kind) = self.probe {
probe(staged, dest, kind, format)?;
}
if let Some(table) = self.parses {
parses(staged, dest, table, self.declared)?;
}
Ok(())
}
fn check_size(self, staged: &Path, dest: &Path) -> Result<(), GeoIpDownloadError> {
if self.min_size_percent == 0 {
return Ok(());
}
let Ok(existing) = fs::metadata(dest).map(|metadata| metadata.len()) else {
return Ok(());
};
if existing == 0 {
return Ok(());
}
let actual = fs::metadata(staged)?.len();
if actual.saturating_mul(PERCENT)
< existing.saturating_mul(u64::from(self.min_size_percent))
{
return Err(GeoIpDownloadError::Undersized {
path: dest.display().to_string(),
actual,
existing,
floor_percent: self.min_size_percent,
});
}
Ok(())
}
}
#[cfg(feature = "geoip-lookup")]
const PROBE_ADDRESSES: [IpAddr; 6] = [
IpAddr::V4(Ipv4Addr::new(8, 8, 8, 8)),
IpAddr::V4(Ipv4Addr::new(1, 1, 1, 1)),
IpAddr::V4(Ipv4Addr::new(9, 9, 9, 9)),
IpAddr::V6(Ipv6Addr::new(0x2001, 0x4860, 0x4860, 0, 0, 0, 0, 0x8888)),
IpAddr::V6(Ipv6Addr::new(0x2606, 0x4700, 0x4700, 0, 0, 0, 0, 0x1111)),
IpAddr::V6(Ipv6Addr::new(0x2620, 0x00fe, 0, 0, 0, 0, 0, 0x00fe)),
];
#[cfg(feature = "geoip-lookup")]
const COUNTRY_PATHS: [&[&str]; 3] = [
&["country", "iso_code"],
&["registered_country", "iso_code"],
&["country_code"],
];
#[cfg(feature = "geoip-lookup")]
const ORGANISATION_PATHS: [&[&str]; 3] = [
&["autonomous_system_organization"],
&["isp"],
&["as_name"],
];
#[cfg(feature = "geoip-lookup")]
fn probe(
staged: &Path,
dest: &Path,
kind: Kind,
format: DatabaseFormat,
) -> Result<(), GeoIpDownloadError> {
if format != DatabaseFormat::Mmdb {
return Ok(());
}
let reader = crate::geoip::enricher::open_reader(staged).map_err(|_| {
GeoIpDownloadError::NotADatabase {
url: dest.display().to_string(),
}
})?;
let answered = PROBE_ADDRESSES
.iter()
.any(|&address| answers(&reader, address, kind));
if answered {
Ok(())
} else {
Err(GeoIpDownloadError::Unpopulated {
path: dest.display().to_string(),
field: field_of(kind),
})
}
}
#[cfg(not(feature = "geoip-lookup"))]
#[allow(clippy::unnecessary_wraps)]
fn probe(
staged: &Path,
_dest: &Path,
_kind: Kind,
_format: DatabaseFormat,
) -> Result<(), GeoIpDownloadError> {
debug!(
staged = %staged.display(),
"no lookup engine is compiled in, so the staged database is not probed"
);
Ok(())
}
fn parses(
staged: &Path,
dest: &Path,
format: TableFormat,
declared: Option<usize>,
) -> Result<(), GeoIpDownloadError> {
crate::table::probe(staged, format, declared).map_err(|e| GeoIpDownloadError::Unparseable {
path: dest.display().to_string(),
detail: e.to_string(),
})
}
#[cfg(feature = "geoip-lookup")]
fn answers(reader: &maxminddb::Reader<maxminddb::Mmap>, address: IpAddr, kind: Kind) -> bool {
let Ok(result) = reader.lookup(address) else {
return false;
};
paths_of(kind).iter().any(|path| {
let elements: Vec<maxminddb::PathElement<'_>> = path
.iter()
.copied()
.map(maxminddb::PathElement::Key)
.collect();
matches!(
result.decode_path::<String>(&elements),
Ok(Some(value)) if !value.trim().is_empty()
)
})
}
#[cfg(feature = "geoip-lookup")]
const fn paths_of(kind: Kind) -> &'static [&'static [&'static str]] {
match kind {
Kind::City => &COUNTRY_PATHS,
Kind::Asn => &ORGANISATION_PATHS,
}
}
#[cfg(feature = "geoip-lookup")]
const fn field_of(kind: Kind) -> &'static str {
match kind {
Kind::City => "country",
Kind::Asn => "organisation name",
}
}
#[cfg(test)]
pub(super) mod fixtures;
#[cfg(test)]
mod tests {
use super::*;
fn defaults() -> AutoDownloadConfig {
AutoDownloadConfig::default()
}
fn sized(path: &Path, bytes: usize) {
fs::write(path, vec![b'x'; bytes]).unwrap();
}
#[test]
fn the_default_guard_checks_both_halves() {
let guard = Guard::new(Kind::Asn, &defaults());
assert_eq!(guard.probe, Some(Kind::Asn));
assert_eq!(guard.min_size_percent, 50);
assert_ne!(guard, Guard::OFF);
}
#[test]
fn each_check_is_switched_off_on_its_own() {
let no_content = AutoDownloadConfig {
verify_content: false,
..defaults()
};
let no_floor = AutoDownloadConfig {
min_size_percent: 0,
..defaults()
};
assert_eq!(Guard::new(Kind::City, &no_content).probe, None);
assert_eq!(Guard::new(Kind::City, &no_content).min_size_percent, 50);
assert_eq!(Guard::new(Kind::City, &no_floor).probe, Some(Kind::City));
assert_eq!(Guard::new(Kind::City, &no_floor).min_size_percent, 0);
}
#[test]
fn a_first_download_has_nothing_to_compare_against() {
let dir = tempfile::tempdir().unwrap();
let staged = dir.path().join("db.mmdb.staged");
let dest = dir.path().join("db.mmdb");
sized(&staged, 16);
let guard = Guard::new(Kind::Asn, &defaults());
guard.check_size(&staged, &dest).unwrap();
}
#[test]
fn a_replacement_under_the_floor_is_refused() {
let dir = tempfile::tempdir().unwrap();
let staged = dir.path().join("db.mmdb.staged");
let dest = dir.path().join("db.mmdb");
sized(&dest, 10_000);
sized(&staged, 4_999);
let err = Guard::new(Kind::Asn, &defaults())
.check_size(&staged, &dest)
.unwrap_err();
let message = err.to_string();
assert!(
matches!(
err,
GeoIpDownloadError::Undersized {
actual: 4_999,
existing: 10_000,
floor_percent: 50,
..
}
),
"{err:?}"
);
assert!(message.contains("4999"), "{message}");
assert!(message.contains("10000"), "{message}");
}
#[test]
fn a_replacement_on_the_floor_is_admitted() {
let dir = tempfile::tempdir().unwrap();
let staged = dir.path().join("db.mmdb.staged");
let dest = dir.path().join("db.mmdb");
sized(&dest, 10_000);
sized(&staged, 5_000);
Guard::new(Kind::Asn, &defaults())
.check_size(&staged, &dest)
.unwrap();
}
#[test]
fn ordinary_month_to_month_shrinkage_is_not_a_refusal() {
let dir = tempfile::tempdir().unwrap();
let staged = dir.path().join("db.mmdb.staged");
let dest = dir.path().join("db.mmdb");
sized(&dest, 10_000);
sized(&staged, 8_000);
Guard::new(Kind::City, &defaults())
.check_size(&staged, &dest)
.unwrap();
}
#[test]
fn a_zero_floor_admits_a_stub() {
let dir = tempfile::tempdir().unwrap();
let staged = dir.path().join("db.mmdb.staged");
let dest = dir.path().join("db.mmdb");
sized(&dest, 10_000);
sized(&staged, 1);
let off = AutoDownloadConfig {
min_size_percent: 0,
..defaults()
};
Guard::new(Kind::Asn, &off)
.check_size(&staged, &dest)
.unwrap();
}
#[test]
fn an_empty_destination_is_not_a_baseline() {
let dir = tempfile::tempdir().unwrap();
let staged = dir.path().join("db.mmdb.staged");
let dest = dir.path().join("db.mmdb");
fs::write(&dest, b"").unwrap();
sized(&staged, 1);
Guard::new(Kind::Asn, &defaults())
.check_size(&staged, &dest)
.unwrap();
}
#[cfg(feature = "geoip-lookup")]
mod probes {
use super::super::fixtures;
use super::*;
fn admit(body: &[u8], kind: Kind) -> Result<(), GeoIpDownloadError> {
let dir = tempfile::tempdir().unwrap();
let staged = dir.path().join("db.mmdb.staged");
let dest = dir.path().join("db.mmdb");
fs::write(&staged, body).unwrap();
Guard::new(kind, &defaults()).admit(&staged, &dest, DatabaseFormat::Mmdb)
}
#[test]
fn an_asn_database_carrying_operator_names_is_admitted() {
admit(&fixtures::asn_mmdb(13_335, "CLOUDFLARENET"), Kind::Asn).unwrap();
}
#[test]
fn an_asn_database_with_a_blank_operator_name_is_refused() {
let err = admit(&fixtures::asn_mmdb(13_335, ""), Kind::Asn).unwrap_err();
let message = err.to_string();
assert!(
matches!(
err,
GeoIpDownloadError::Unpopulated {
field: "organisation name",
..
}
),
"{err:?}"
);
assert!(message.contains("organisation name"), "{message}");
}
#[test]
fn a_whitespace_operator_name_counts_as_blank() {
let err = admit(&fixtures::asn_mmdb(13_335, " "), Kind::Asn).unwrap_err();
assert!(
matches!(err, GeoIpDownloadError::Unpopulated { .. }),
"{err:?}"
);
}
#[test]
fn a_city_database_resolving_a_country_is_admitted() {
admit(&fixtures::city_mmdb("US"), Kind::City).unwrap();
}
#[test]
fn a_city_database_resolving_no_country_is_refused() {
let err = admit(&fixtures::city_mmdb(""), Kind::City).unwrap_err();
assert!(
matches!(
err,
GeoIpDownloadError::Unpopulated {
field: "country",
..
}
),
"{err:?}"
);
}
#[test]
fn a_database_published_for_ipv6_is_probed_the_same_way() {
admit(&fixtures::asn_mmdb_v6(13_335, "CLOUDFLARENET"), Kind::Asn).unwrap();
admit(&fixtures::city_mmdb_v6("US"), Kind::City).unwrap();
let err = admit(&fixtures::asn_mmdb_v6(13_335, ""), Kind::Asn).unwrap_err();
assert!(
matches!(err, GeoIpDownloadError::Unpopulated { .. }),
"{err:?}"
);
}
#[test]
fn an_ipv6_address_answers_from_a_database_published_for_it() {
let dir = tempfile::tempdir().unwrap();
let staged = dir.path().join("db.mmdb.staged");
fs::write(&staged, fixtures::asn_mmdb_v6(13_335, "CLOUDFLARENET")).unwrap();
let reader = crate::geoip::enricher::open_reader(&staged).unwrap();
let resolver: IpAddr = "2606:4700:4700::1111".parse().unwrap();
assert!(answers(&reader, resolver, Kind::Asn));
assert!(!answers(&reader, resolver, Kind::City));
}
#[test]
fn a_top_level_country_code_is_read_too() {
admit(&fixtures::flat_city_mmdb("AU"), Kind::City).unwrap();
}
#[test]
fn an_asn_database_offered_as_a_city_one_is_refused() {
let err = admit(&fixtures::asn_mmdb(13_335, "CLOUDFLARENET"), Kind::City).unwrap_err();
assert!(
matches!(err, GeoIpDownloadError::Unpopulated { .. }),
"{err:?}"
);
}
#[test]
fn a_database_the_reader_will_not_open_is_refused() {
let mut body = b"not a database".to_vec();
body.extend_from_slice(b"\xab\xcd\xefMaxMind.com");
body.extend_from_slice(b"nor is this metadata");
let err = admit(&body, Kind::Asn).unwrap_err();
assert!(
matches!(err, GeoIpDownloadError::NotADatabase { .. }),
"{err:?}"
);
}
#[test]
fn content_verification_off_admits_a_blank_database() {
let dir = tempfile::tempdir().unwrap();
let staged = dir.path().join("db.mmdb.staged");
let dest = dir.path().join("db.mmdb");
fs::write(&staged, fixtures::asn_mmdb(13_335, "")).unwrap();
let off = AutoDownloadConfig {
verify_content: false,
..defaults()
};
Guard::new(Kind::Asn, &off)
.admit(&staged, &dest, DatabaseFormat::Mmdb)
.unwrap();
}
#[test]
fn a_text_database_is_admitted_on_its_structure_alone() {
let dir = tempfile::tempdir().unwrap();
let staged = dir.path().join("ranges.csv.staged");
let dest = dir.path().join("ranges.csv");
fs::write(&staged, b"1.1.1.0,1.1.1.255,13335,CLOUDFLARENET\n").unwrap();
Guard::new(Kind::Asn, &defaults())
.admit(&staged, &dest, DatabaseFormat::Csv)
.unwrap();
}
}
}