use std::fmt::{self, Write as _};
use std::sync::Once;
use compact_str::CompactString;
use maxminddb::LookupResult;
use serde::de::{DeserializeSeed, MapAccess, SeqAccess, Visitor};
use serde::{Deserialize, Deserializer};
use tracing::{debug, warn};
use super::enricher::asn_number;
use super::record::{ExtraValue, GeoIpRecord};
const MAX_FIELDS: usize = 512;
const MAX_BYTES: usize = 64 * 1024;
static TRUNCATED: Once = Once::new();
pub(super) fn collect<S: AsRef<[u8]>>(result: &LookupResult<'_, S>, record: &mut GeoIpRecord) {
let kept = match result.decode::<Leaves>() {
Ok(Some(leaves)) => leaves.0,
Ok(None) => return,
Err(e) => {
debug!(error = %e, "record fields did not decode");
return;
}
};
if kept.full() {
truncated(result);
}
record.extra.reserve(kept.fields.len());
for (path, value) in kept.fields {
if !typed(&path, &value, record) {
record.extra.push(path, value);
}
}
}
fn truncated<S: AsRef<[u8]>>(result: &LookupResult<'_, S>) {
TRUNCATED.call_once(|| {
warn!(
network = ?result.network().ok(),
offset = ?result.offset(),
max_fields = MAX_FIELDS,
max_bytes = MAX_BYTES,
"database record carries more fields than the map holds; it was truncated"
);
});
}
struct Leaves(Kept);
struct Kept {
fields: Vec<(CompactString, ExtraValue)>,
bytes: usize,
full: bool,
}
impl Kept {
const fn new() -> Self {
Self {
fields: Vec::new(),
bytes: 0,
full: false,
}
}
const fn full(&self) -> bool {
self.full
}
fn push(&mut self, path: &str, value: ExtraValue) {
let cost = path.len() + payload(&value);
if self.bytes + cost > MAX_BYTES {
self.full = true;
return;
}
self.bytes += cost;
self.fields.push((CompactString::from(path), value));
self.full = self.fields.len() >= MAX_FIELDS;
}
}
fn payload(value: &ExtraValue) -> usize {
match value {
ExtraValue::Str(text) => text.len(),
ExtraValue::Bytes(bytes) => bytes.len(),
_ => 0,
}
}
impl<'de> Deserialize<'de> for Leaves {
fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
let mut kept = Kept::new();
let mut path = String::new();
Walk {
kept: &mut kept,
path: &mut path,
}
.deserialize(deserializer)?;
Ok(Self(kept))
}
}
struct Walk<'a> {
kept: &'a mut Kept,
path: &'a mut String,
}
impl Walk<'_> {
fn leaf(self, value: ExtraValue) {
self.kept.push(self.path, value);
}
}
impl<'de> DeserializeSeed<'de> for Walk<'_> {
type Value = ();
fn deserialize<D: Deserializer<'de>>(self, deserializer: D) -> Result<(), D::Error> {
deserializer.deserialize_any(self)
}
}
impl<'de> Visitor<'de> for Walk<'_> {
type Value = ();
fn expecting(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str("any value a database record can hold")
}
fn visit_bool<E>(self, value: bool) -> Result<(), E> {
self.leaf(ExtraValue::Bool(value));
Ok(())
}
fn visit_i64<E>(self, value: i64) -> Result<(), E> {
self.leaf(ExtraValue::from_i64(value));
Ok(())
}
fn visit_u64<E>(self, value: u64) -> Result<(), E> {
self.leaf(ExtraValue::UInt(value));
Ok(())
}
fn visit_u128<E>(self, value: u128) -> Result<(), E> {
self.leaf(ExtraValue::from_u128(value));
Ok(())
}
fn visit_f64<E>(self, value: f64) -> Result<(), E> {
self.leaf(ExtraValue::Float(value));
Ok(())
}
fn visit_str<E>(self, value: &str) -> Result<(), E> {
self.leaf(ExtraValue::Str(CompactString::from(value)));
Ok(())
}
fn visit_bytes<E>(self, value: &[u8]) -> Result<(), E> {
self.leaf(ExtraValue::Bytes(value.to_vec()));
Ok(())
}
fn visit_map<A: MapAccess<'de>>(self, mut map: A) -> Result<(), A::Error> {
while !self.kept.full() {
let Some(key) = map.next_key::<&str>()? else {
break;
};
let parent = self.path.len();
if parent != 0 {
self.path.push('.');
}
self.path.push_str(key);
map.next_value_seed(Walk {
kept: &mut *self.kept,
path: &mut *self.path,
})?;
self.path.truncate(parent);
}
Ok(())
}
fn visit_seq<A: SeqAccess<'de>>(self, mut seq: A) -> Result<(), A::Error> {
let mut index = 0usize;
while !self.kept.full() {
let parent = self.path.len();
if parent != 0 {
self.path.push('.');
}
let _ = write!(self.path, "{index}");
let element = seq.next_element_seed(Walk {
kept: &mut *self.kept,
path: &mut *self.path,
})?;
self.path.truncate(parent);
if element.is_none() {
break;
}
index += 1;
}
Ok(())
}
}
fn typed(path: &str, value: &ExtraValue, record: &GeoIpRecord) -> bool {
match path {
"city.names.en" => is_str(value, record.city_name.as_deref()),
"continent.code" | "continent_code" => is_str(value, record.continent_code.as_deref()),
"continent.names.en" | "continent" => is_str(value, record.continent_name.as_deref()),
"country.iso_code" | "country_code" => is_str(value, record.country_code.as_deref()),
"country.names.en" | "country" => is_str(value, record.country_name.as_deref()),
"postal.code" => is_str(value, record.postal_code.as_deref()),
"location.time_zone" => is_str(value, record.timezone.as_deref()),
"location.latitude" => is_float(value, record.latitude),
"location.longitude" => is_float(value, record.longitude),
"location.metro_code" => is_uint(value, record.metro_code.map(u64::from)),
"location.accuracy_radius" => is_uint(value, record.accuracy_radius.map(u64::from)),
"autonomous_system_number" => {
is_uint(value, record.autonomous_system_number.map(u64::from))
}
"autonomous_system_organization" | "as_name" => {
is_str(value, record.autonomous_system_organization.as_deref())
}
"as_domain" => is_str(value, record.as_domain.as_deref()),
"asn" => {
record.autonomous_system_number.is_some()
&& value.as_str().and_then(asn_number) == record.autonomous_system_number
}
_ => match subdivision(path) {
Some("names.en") => is_str(value, record.region_name.as_deref()),
Some("iso_code") => is_str(value, record.region_code.as_deref()),
_ => false,
},
}
}
fn subdivision(path: &str) -> Option<&str> {
let (index, field) = path.strip_prefix("subdivisions.")?.split_once('.')?;
index
.bytes()
.all(|byte| byte.is_ascii_digit())
.then_some(field)
}
fn is_str(value: &ExtraValue, held: Option<&str>) -> bool {
matches!((value.as_str(), held), (Some(value), Some(held)) if value == held)
}
fn is_uint(value: &ExtraValue, held: Option<u64>) -> bool {
matches!((value, held), (ExtraValue::UInt(value), Some(held)) if *value == held)
}
fn is_float(value: &ExtraValue, held: Option<f64>) -> bool {
matches!((value, held), (ExtraValue::Float(value), Some(held)) if value.to_bits() == held.to_bits())
}
#[cfg(test)]
mod tests {
use super::*;
fn located() -> GeoIpRecord {
GeoIpRecord {
city_name: Some(CompactString::from("Boxford")),
region_name: Some(CompactString::from("West Berkshire")),
latitude: Some(51.75),
accuracy_radius: Some(100),
autonomous_system_number: Some(15169),
..GeoIpRecord::default()
}
}
#[test]
fn a_path_the_typed_decode_took_is_not_repeated() {
let record = located();
assert!(typed(
"city.names.en",
&ExtraValue::Str(CompactString::from("Boxford")),
&record
));
assert!(typed(
"location.accuracy_radius",
&ExtraValue::UInt(100),
&record
));
assert!(typed(
"location.latitude",
&ExtraValue::Float(51.75),
&record
));
}
#[test]
fn a_path_no_typed_field_names_is_kept() {
let record = located();
assert!(!typed(
"city.names.de",
&ExtraValue::Str(CompactString::from("Boxford")),
&record
));
assert!(!typed(
"city.geoname_id",
&ExtraValue::UInt(2_655_045),
&record
));
}
#[test]
fn a_superseded_subdivision_is_kept_and_the_specific_one_is_not() {
let record = located();
assert!(!typed(
"subdivisions.0.names.en",
&ExtraValue::Str(CompactString::from("England")),
&record
));
assert!(typed(
"subdivisions.1.names.en",
&ExtraValue::Str(CompactString::from("West Berkshire")),
&record
));
}
#[test]
fn a_value_the_typed_decode_could_not_read_is_kept() {
let record = GeoIpRecord::default();
assert!(!typed(
"location.accuracy_radius",
&ExtraValue::UInt(100),
&record
));
}
#[test]
fn the_as_prefixed_asn_matches_the_number_it_resolved_to() {
let record = located();
assert!(typed(
"asn",
&ExtraValue::Str(CompactString::from("AS15169")),
&record
));
assert!(!typed(
"asn",
&ExtraValue::Str(CompactString::from("AS13335")),
&record
));
}
#[test]
fn a_subdivision_path_is_read_only_when_it_carries_an_index() {
assert_eq!(subdivision("subdivisions.0.iso_code"), Some("iso_code"));
assert_eq!(subdivision("subdivisions.12.names.en"), Some("names.en"));
assert_eq!(subdivision("subdivisions.x.iso_code"), None);
assert_eq!(subdivision("subdivisions.0"), None);
assert_eq!(subdivision("city.names.en"), None);
}
}
#[cfg(test)]
mod crafted {
use crate::geoip::mmdb_wire::{IPV4, MAP, database, header, string};
const DATABASE_TYPE: &str = "GeoLite2-City";
const POINTER: u8 = 1;
const POINTER_LIMIT: usize = 2048;
fn pointer(target: usize) -> Vec<u8> {
assert!(target < POINTER_LIMIT, "crafted records stay under 2 KiB");
vec![
(POINTER << 5) | u8::try_from(target >> 8).unwrap(),
u8::try_from(target & 0xff).unwrap(),
]
}
pub(super) fn pointer_bomb(levels: u32, leaf: &str) -> Vec<u8> {
let mut data = string(leaf);
let mut below = 0;
for _ in 0..levels {
let here = data.len();
let branch = pointer(below);
data.extend_from_slice(&header(MAP, 2));
data.extend_from_slice(&string("a"));
data.extend_from_slice(&branch);
data.extend_from_slice(&string("b"));
data.extend_from_slice(&branch);
below = here;
}
database(&data, below, DATABASE_TYPE, IPV4)
}
pub(super) fn wide_record(fields: usize, value: &str) -> Vec<u8> {
let mut data = header(MAP, fields);
for index in 0..fields {
data.extend_from_slice(&string(&format!("f{index}")));
data.extend_from_slice(&string(value));
}
database(&data, 0, DATABASE_TYPE, IPV4)
}
}
#[cfg(test)]
mod bounds {
use std::io::Write as _;
use std::net::IpAddr;
use std::path::Path;
use std::time::{Duration, Instant};
use tempfile::NamedTempFile;
use super::crafted;
use super::{ExtraValue, MAX_BYTES, MAX_FIELDS};
use crate::geoip::{CacheConfig, DatabasePaths, GeoIp, GeoIpRecord};
use crate::maxminddb::Reader;
const CITY_RICH_DB: &str = concat!(
env!("CARGO_MANIFEST_DIR"),
"/tests/data/city-rich-test.mmdb"
);
const BOXFORD: &str = "2.125.160.216";
const LONG_LEAF_BYTES: usize = 200;
const DEEP: u32 = 45;
const BUDGET: Duration = Duration::from_secs(2);
fn ip(literal: &str) -> IpAddr {
literal.parse().unwrap()
}
fn written(bytes: &[u8]) -> NamedTempFile {
let mut file = NamedTempFile::new().unwrap();
file.write_all(bytes).unwrap();
file.flush().unwrap();
file
}
fn collected(bytes: &[u8]) -> (GeoIpRecord, Duration) {
let file = written(bytes);
let reader = Reader::open_readfile(file.path()).unwrap();
let result = reader.lookup(ip("1.2.3.4")).unwrap();
let mut record = GeoIpRecord::default();
let started = Instant::now();
super::collect(&result, &mut record);
(record, started.elapsed())
}
fn retained(record: &GeoIpRecord) -> usize {
record
.extra
.iter()
.map(|(path, value)| path.len() + super::payload(value))
.sum()
}
#[test]
fn a_pointer_graph_cannot_grow_the_map_or_stall_the_walk() {
let bytes = crafted::pointer_bomb(DEEP, "x");
assert!(bytes.len() < 1024, "{}", bytes.len());
let (record, elapsed) = collected(&bytes);
assert_eq!(record.extra.len(), MAX_FIELDS);
assert!(elapsed < BUDGET, "{elapsed:?}");
}
#[test]
fn leaves_sharing_one_long_string_are_bounded_in_bytes_not_just_in_count() {
let leaf = "v".repeat(LONG_LEAF_BYTES);
let bytes = crafted::pointer_bomb(DEEP, &leaf);
assert!(bytes.len() < 1024, "{}", bytes.len());
let (record, elapsed) = collected(&bytes);
assert!(retained(&record) <= MAX_BYTES, "{}", retained(&record));
assert!(record.extra.len() < MAX_FIELDS, "{}", record.extra.len());
assert!(elapsed < BUDGET, "{elapsed:?}");
}
#[test]
fn a_crafted_record_cannot_stall_a_lookup() {
let bytes = crafted::pointer_bomb(20, "x");
let file = written(&bytes);
let geoip = GeoIp::open(
DatabasePaths::city_only(file.path()),
CacheConfig::default(),
)
.unwrap();
let started = Instant::now();
let record = geoip.lookup(ip("1.2.3.4")).unwrap();
let elapsed = started.elapsed();
assert_eq!(record.extra.len(), MAX_FIELDS);
assert!(elapsed < BUDGET, "{elapsed:?}");
}
#[test]
fn an_honest_record_over_the_cap_is_truncated_rather_than_refused() {
let bytes = crafted::wide_record(MAX_FIELDS + 100, "v");
let (record, _) = collected(&bytes);
assert_eq!(record.extra.len(), MAX_FIELDS);
assert_eq!(record.extra.get("f0").unwrap().as_str(), Some("v"));
}
#[test]
fn an_honest_record_under_the_cap_keeps_every_field() {
let fields = MAX_FIELDS - 1;
let bytes = crafted::wide_record(fields, "v");
let (record, _) = collected(&bytes);
assert_eq!(record.extra.len(), fields);
}
#[test]
fn the_richest_published_record_is_nowhere_near_the_bounds() {
let geoip = GeoIp::open(
DatabasePaths::city_only(Path::new(CITY_RICH_DB)),
CacheConfig::default(),
)
.unwrap();
let record = geoip.lookup(ip(BOXFORD)).unwrap();
assert!(
record.extra.len() * 8 < MAX_FIELDS,
"{}",
record.extra.len()
);
assert!(retained(&record) * 8 < MAX_BYTES, "{}", retained(&record));
assert_eq!(
record
.extra
.get("city.names.de")
.and_then(ExtraValue::as_str),
Some("Boxford")
);
assert_eq!(
record
.extra
.get("traits.user_type")
.and_then(ExtraValue::as_str),
Some("residential")
);
assert!(record.extra.get("city.geoname_id").is_some());
assert!(record.extra.get("subdivisions.1.confidence").is_some());
}
}