use std::fmt;
use std::slice;
use std::sync::{Arc, LazyLock};
use compact_str::CompactString;
use serde::de::{MapAccess, SeqAccess, Visitor};
use serde::ser::SerializeMap;
use serde::{Deserialize, Deserializer, Serialize, Serializer};
const TYPED_FIELDS: usize = 19;
const HINT_RESERVE: usize = 128;
static PRIVATE: LazyLock<Arc<GeoIpRecord>> = LazyLock::new(|| Arc::new(GeoIpRecord::private()));
#[derive(Clone, Copy, Debug, PartialEq)]
#[non_exhaustive]
pub enum FieldValue<'a> {
Str(&'a str),
Float(f64),
UInt(u64),
Int(i64),
UInt128(u128),
Bool(bool),
Bytes(&'a [u8]),
}
#[derive(Clone, Debug, PartialEq, Serialize)]
#[serde(untagged)]
#[non_exhaustive]
pub enum ExtraValue {
Str(CompactString),
Float(f64),
UInt(u64),
Int(i64),
UInt128(u128),
Bool(bool),
Bytes(Vec<u8>),
}
impl ExtraValue {
pub(super) fn from_i64(value: i64) -> Self {
u64::try_from(value).map_or(Self::Int(value), Self::UInt)
}
pub(super) fn from_u128(value: u128) -> Self {
u64::try_from(value).map_or(Self::UInt128(value), Self::UInt)
}
#[must_use]
pub fn as_field(&self) -> FieldValue<'_> {
match self {
Self::Str(value) => FieldValue::Str(value),
Self::Float(value) => FieldValue::Float(*value),
Self::UInt(value) => FieldValue::UInt(*value),
Self::Int(value) => FieldValue::Int(*value),
Self::UInt128(value) => FieldValue::UInt128(*value),
Self::Bool(value) => FieldValue::Bool(*value),
Self::Bytes(value) => FieldValue::Bytes(value),
}
}
#[must_use]
pub fn as_str(&self) -> Option<&str> {
match self {
Self::Str(value) => Some(value),
_ => None,
}
}
}
impl<'de> Deserialize<'de> for ExtraValue {
fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
deserializer.deserialize_any(ExtraValueVisitor)
}
}
struct ExtraValueVisitor;
impl<'de> Visitor<'de> for ExtraValueVisitor {
type Value = ExtraValue;
fn expecting(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str("a string, number, boolean or byte string")
}
fn visit_bool<E>(self, value: bool) -> Result<Self::Value, E> {
Ok(ExtraValue::Bool(value))
}
fn visit_i64<E>(self, value: i64) -> Result<Self::Value, E> {
Ok(ExtraValue::from_i64(value))
}
fn visit_u64<E>(self, value: u64) -> Result<Self::Value, E> {
Ok(ExtraValue::UInt(value))
}
fn visit_u128<E>(self, value: u128) -> Result<Self::Value, E> {
Ok(ExtraValue::from_u128(value))
}
fn visit_f64<E>(self, value: f64) -> Result<Self::Value, E> {
Ok(ExtraValue::Float(value))
}
fn visit_str<E>(self, value: &str) -> Result<Self::Value, E> {
Ok(ExtraValue::Str(CompactString::from(value)))
}
fn visit_bytes<E>(self, value: &[u8]) -> Result<Self::Value, E> {
Ok(ExtraValue::Bytes(value.to_vec()))
}
fn visit_seq<A: SeqAccess<'de>>(self, mut seq: A) -> Result<Self::Value, A::Error> {
let mut bytes = Vec::with_capacity(seq.size_hint().unwrap_or_default().min(HINT_RESERVE));
while let Some(byte) = seq.next_element::<u8>()? {
bytes.push(byte);
}
Ok(ExtraValue::Bytes(bytes))
}
}
#[derive(Clone, Debug, Default, PartialEq)]
pub struct ExtraFields {
fields: Vec<(CompactString, ExtraValue)>,
}
impl ExtraFields {
#[must_use]
pub const fn new() -> Self {
Self { fields: Vec::new() }
}
#[must_use]
pub fn get(&self, path: &str) -> Option<&ExtraValue> {
self.fields
.iter()
.find(|(key, _)| key == path)
.map(|(_, value)| value)
}
pub fn iter(&self) -> slice::Iter<'_, (CompactString, ExtraValue)> {
self.fields.iter()
}
#[must_use]
pub fn len(&self) -> usize {
self.fields.len()
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.fields.is_empty()
}
pub(super) fn push(&mut self, path: CompactString, value: ExtraValue) {
self.fields.push((path, value));
}
pub(super) fn reserve(&mut self, additional: usize) {
self.fields.reserve(additional);
}
}
impl<'a> IntoIterator for &'a ExtraFields {
type Item = &'a (CompactString, ExtraValue);
type IntoIter = slice::Iter<'a, (CompactString, ExtraValue)>;
fn into_iter(self) -> Self::IntoIter {
self.iter()
}
}
impl Serialize for ExtraFields {
fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
let mut map = serializer.serialize_map(Some(self.fields.len()))?;
for (path, value) in &self.fields {
map.serialize_entry(path.as_str(), value)?;
}
map.end()
}
}
impl<'de> Deserialize<'de> for ExtraFields {
fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
deserializer.deserialize_map(ExtraFieldsVisitor)
}
}
struct ExtraFieldsVisitor;
impl<'de> Visitor<'de> for ExtraFieldsVisitor {
type Value = ExtraFields;
fn expecting(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str("a map of source paths to values")
}
fn visit_map<A: MapAccess<'de>>(self, mut map: A) -> Result<Self::Value, A::Error> {
let mut fields = Vec::with_capacity(map.size_hint().unwrap_or_default().min(HINT_RESERVE));
while let Some(entry) = map.next_entry::<CompactString, ExtraValue>()? {
fields.push(entry);
}
Ok(ExtraFields { fields })
}
}
#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
pub struct GeoIpRecord {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub city_name: Option<CompactString>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub continent_code: Option<CompactString>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub continent_name: Option<CompactString>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub country_code: Option<CompactString>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub country_name: Option<CompactString>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub region_name: Option<CompactString>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub region_code: Option<CompactString>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub postal_code: Option<CompactString>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub timezone: Option<CompactString>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub latitude: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub longitude: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub metro_code: Option<u16>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub accuracy_radius: Option<u16>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub autonomous_system_number: Option<u32>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub autonomous_system_organization: Option<CompactString>,
pub is_private: bool,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub network: Option<CompactString>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub asn_network: Option<CompactString>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub as_domain: Option<CompactString>,
#[serde(default, skip_serializing_if = "ExtraFields::is_empty")]
pub extra: ExtraFields,
}
impl GeoIpRecord {
#[must_use]
pub const fn private() -> Self {
Self {
city_name: None,
continent_code: None,
continent_name: None,
country_code: None,
country_name: None,
region_name: None,
region_code: None,
postal_code: None,
timezone: None,
latitude: None,
longitude: None,
metro_code: None,
accuracy_radius: None,
autonomous_system_number: None,
autonomous_system_organization: None,
is_private: true,
network: None,
asn_network: None,
as_domain: None,
extra: ExtraFields::new(),
}
}
#[must_use]
pub fn private_shared() -> Arc<Self> {
Arc::clone(&PRIVATE)
}
#[must_use]
pub fn to_schema_map(&self) -> Vec<(&str, FieldValue<'_>)> {
let mut fields = Vec::with_capacity(TYPED_FIELDS + self.extra.len());
push(&mut fields, "city_name", self.city_name.as_deref());
push(
&mut fields,
"continent_code",
self.continent_code.as_deref(),
);
push(
&mut fields,
"continent_name",
self.continent_name.as_deref(),
);
push(&mut fields, "country_code", self.country_code.as_deref());
push(&mut fields, "country_name", self.country_name.as_deref());
push(&mut fields, "region_name", self.region_name.as_deref());
push(&mut fields, "region_code", self.region_code.as_deref());
push(&mut fields, "postal_code", self.postal_code.as_deref());
push(&mut fields, "timezone", self.timezone.as_deref());
push(&mut fields, "latitude", self.latitude);
push(&mut fields, "longitude", self.longitude);
push(&mut fields, "metro_code", self.metro_code.map(u32::from));
push(
&mut fields,
"accuracy_radius",
self.accuracy_radius.map(u32::from),
);
push(
&mut fields,
"autonomous_system_number",
self.autonomous_system_number,
);
push(
&mut fields,
"autonomous_system_organization",
self.autonomous_system_organization.as_deref(),
);
fields.push(("is_private", FieldValue::Bool(self.is_private)));
push(&mut fields, "network", self.network.as_deref());
push(&mut fields, "asn_network", self.asn_network.as_deref());
push(&mut fields, "as_domain", self.as_domain.as_deref());
for (path, value) in &self.extra {
fields.push((path.as_str(), value.as_field()));
}
fields
}
}
impl<'a> From<&'a str> for FieldValue<'a> {
fn from(value: &'a str) -> Self {
Self::Str(value)
}
}
impl From<f64> for FieldValue<'_> {
fn from(value: f64) -> Self {
Self::Float(value)
}
}
impl From<u32> for FieldValue<'_> {
fn from(value: u32) -> Self {
Self::UInt(u64::from(value))
}
}
impl From<u64> for FieldValue<'_> {
fn from(value: u64) -> Self {
Self::UInt(value)
}
}
fn push<'a, T: Into<FieldValue<'a>>>(
fields: &mut Vec<(&'a str, FieldValue<'a>)>,
key: &'static str,
value: Option<T>,
) {
if let Some(value) = value {
fields.push((key, value.into()));
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn a_default_record_carries_no_answer_and_is_not_private() {
let record = GeoIpRecord::default();
assert!(record.city_name.is_none());
assert!(record.network.is_none());
assert!(record.asn_network.is_none());
assert!(!record.is_private);
}
#[test]
fn the_private_record_is_flagged_and_otherwise_empty() {
let record = GeoIpRecord::private();
assert!(record.is_private);
assert!(record.country_code.is_none());
assert!(record.autonomous_system_number.is_none());
assert!(record.latitude.is_none());
}
#[test]
fn the_shared_private_record_is_one_allocation() {
let first = GeoIpRecord::private_shared();
let second = GeoIpRecord::private_shared();
assert!(Arc::ptr_eq(&first, &second));
assert!(first.is_private);
}
#[test]
fn the_schema_map_omits_absent_fields_but_always_states_privacy() {
let record = GeoIpRecord::default();
let fields = record.to_schema_map();
assert_eq!(fields, vec![("is_private", FieldValue::Bool(false))]);
}
#[test]
fn the_schema_map_keys_are_the_field_names_and_keep_their_types() {
let record = GeoIpRecord {
city_name: Some(CompactString::from("Boxford")),
country_code: Some(CompactString::from("GB")),
region_name: Some(CompactString::from("West Berkshire")),
latitude: Some(51.75),
accuracy_radius: Some(100),
autonomous_system_number: Some(15169),
network: Some(CompactString::from("2.125.160.216/29")),
asn_network: Some(CompactString::from("2.125.160.0/24")),
..GeoIpRecord::default()
};
let fields = record.to_schema_map();
assert!(fields.contains(&("city_name", FieldValue::Str("Boxford"))));
assert!(fields.contains(&("country_code", FieldValue::Str("GB"))));
assert!(fields.contains(&("region_name", FieldValue::Str("West Berkshire"))));
assert!(fields.contains(&("accuracy_radius", FieldValue::UInt(100))));
assert!(fields.contains(&("autonomous_system_number", FieldValue::UInt(15169))));
assert!(fields.contains(&("is_private", FieldValue::Bool(false))));
assert!(fields.contains(&("network", FieldValue::Str("2.125.160.216/29"))));
assert!(fields.contains(&("asn_network", FieldValue::Str("2.125.160.0/24"))));
let latitude = fields
.iter()
.find_map(|(key, value)| match (key, value) {
(&"latitude", &FieldValue::Float(value)) => Some(value),
_ => None,
})
.expect("latitude is in the map");
assert!((latitude - 51.75).abs() < f64::EPSILON);
}
#[test]
fn serialisation_drops_the_fields_that_resolved_to_nothing() {
let record = GeoIpRecord {
country_code: Some(CompactString::from("SE")),
..GeoIpRecord::default()
};
let json = serde_json::to_string(&record).unwrap();
assert!(json.contains(r#""country_code":"SE""#), "{json}");
assert!(json.contains(r#""is_private":false"#), "{json}");
assert!(!json.contains("city_name"), "{json}");
assert!(!json.contains("null"), "{json}");
}
fn with_extra() -> GeoIpRecord {
let mut record = GeoIpRecord {
city_name: Some(CompactString::from("Boxford")),
..GeoIpRecord::default()
};
record.extra.push(
CompactString::from("isp"),
ExtraValue::Str("Telstra".into()),
);
record.extra.push(
CompactString::from("city.geoname_id"),
ExtraValue::UInt(2_655_045),
);
record
}
#[test]
fn the_schema_map_carries_the_source_fields_beside_the_typed_ones() {
let record = with_extra();
let fields = record.to_schema_map();
assert!(fields.contains(&("city_name", FieldValue::Str("Boxford"))));
assert!(fields.contains(&("isp", FieldValue::Str("Telstra"))));
assert!(fields.contains(&("city.geoname_id", FieldValue::UInt(2_655_045))));
}
#[test]
fn the_source_fields_are_reached_by_path_and_kept_in_source_order() {
let record = with_extra();
assert_eq!(record.extra.len(), 2);
assert_eq!(
record.extra.get("isp"),
Some(&ExtraValue::Str("Telstra".into()))
);
assert!(record.extra.get("organization").is_none());
let paths: Vec<&str> = record.extra.iter().map(|(path, _)| path.as_str()).collect();
assert_eq!(paths, vec!["isp", "city.geoname_id"]);
}
#[test]
fn the_source_fields_serialise_as_an_object_and_round_trip() {
let record = with_extra();
let json = serde_json::to_string(&record).unwrap();
assert!(json.contains(r#""extra":{"isp":"Telstra""#), "{json}");
assert!(json.contains(r#""city.geoname_id":2655045"#), "{json}");
let decoded: GeoIpRecord = serde_json::from_str(&json).unwrap();
assert_eq!(decoded, record);
}
#[test]
fn a_record_with_no_source_fields_writes_no_map() {
let record = GeoIpRecord::default();
let json = serde_json::to_string(&record).unwrap();
assert!(!json.contains("extra"), "{json}");
}
#[test]
fn every_value_kind_the_database_format_holds_round_trips() {
for value in [
ExtraValue::Str(CompactString::from("residential")),
ExtraValue::Float(-1.25),
ExtraValue::UInt(2_655_045),
ExtraValue::Int(-42),
ExtraValue::Bool(true),
ExtraValue::Bytes(vec![0, 1, 254, 255]),
] {
let json = serde_json::to_string(&value).unwrap();
let decoded: ExtraValue = serde_json::from_str(&json).unwrap();
assert_eq!(decoded, value, "{json}");
}
}
#[test]
fn a_whole_number_reads_as_the_narrowest_kind_that_holds_it() {
assert_eq!(ExtraValue::from_i64(5), ExtraValue::UInt(5));
assert_eq!(ExtraValue::from_i64(-5), ExtraValue::Int(-5));
assert_eq!(ExtraValue::from_u128(5), ExtraValue::UInt(5));
let past_u64 = u128::from(u64::MAX) + 1;
assert_eq!(
ExtraValue::from_u128(past_u64),
ExtraValue::UInt128(past_u64)
);
}
#[test]
fn a_number_wider_than_json_carries_is_written_out_and_re_read_as_a_double() {
let value = ExtraValue::UInt128(u128::from(u64::MAX) + 1);
let json = serde_json::to_string(&value).unwrap();
assert_eq!(json, "18446744073709551616");
let decoded: ExtraValue = serde_json::from_str(&json).unwrap();
assert!(matches!(decoded, ExtraValue::Float(_)), "{decoded:?}");
}
#[test]
fn a_record_round_trips_through_serde() {
let record = GeoIpRecord {
city_name: Some(CompactString::from("Linkoping")),
country_code: Some(CompactString::from("SE")),
latitude: Some(58.4167),
metro_code: Some(500),
autonomous_system_organization: Some(CompactString::from("Bredband2 AB")),
network: Some(CompactString::from("89.160.20.112/28")),
asn_network: Some(CompactString::from("89.160.20.112/29")),
..GeoIpRecord::default()
};
let json = serde_json::to_string(&record).unwrap();
let decoded: GeoIpRecord = serde_json::from_str(&json).unwrap();
assert_eq!(decoded, record);
}
}