irontide-session 1.2.1

//! M255/ER1: per-peer country resolution from a user-supplied
//! `MaxMind`/`DB-IP` `.mmdb` database. Session-owned — the engine crates
//! never resolve; `PeerInfo.country_code` is stamped at the
//! `SessionCommand::GetPeerInfo` boundary.

use std::net::IpAddr;
use std::path::Path;

use irontide_settings::Settings;

/// Thin wrapper over the `MaxMind` reader holding the whole DB in memory
/// (`open_readfile`, no mmap — a country DB is a few MB read once).
pub(crate) struct GeoIpResolver {
    reader: maxminddb::Reader<Vec<u8>>,
}

impl std::fmt::Debug for GeoIpResolver {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("GeoIpResolver").finish_non_exhaustive()
    }
}

impl GeoIpResolver {
    /// Open a `.mmdb` database from disk. Errors bubble so the caller
    /// (`build_geoip_resolver`) can warn-and-degrade per M255 D5.
    pub(crate) fn open(path: &Path) -> Result<Self, maxminddb::MaxMindDbError> {
        Ok(Self {
            reader: maxminddb::Reader::open_readfile(path)?,
        })
    }

    /// Resolve an IP to an ISO-3166-1 alpha-2 code. `None` for unmapped
    /// IPs, lookup errors, or records without a 2-ASCII-byte `iso_code` —
    /// resolution is cosmetic and must never error a caller.
    #[must_use]
    pub(crate) fn lookup(&self, ip: IpAddr) -> Option<[u8; 2]> {
        let result = self.reader.lookup(ip).ok()?;
        let country = result.decode::<maxminddb::geoip2::Country>().ok()??;
        let iso = country.country.iso_code?;
        let bytes = iso.as_bytes();
        if bytes.len() == 2 && bytes.iter().all(u8::is_ascii_uppercase) {
            Some([bytes[0], bytes[1]])
        } else {
            None
        }
    }
}

/// Build (or decline to build) the resolver from live settings — the
/// single constructor used at session start AND on settings apply.
/// Enabled-with-valid-DB ⇒ `Some`; disabled ⇒ `None` silently; enabled
/// but unopenable ⇒ `None` with a WARN (D5 graceful degrade — a
/// vanished DB file must not fail a running session).
#[must_use]
pub(crate) fn build_geoip_resolver(settings: &Settings) -> Option<GeoIpResolver> {
    if !settings.resolve_peer_countries {
        return None;
    }
    let path = settings.peer_country_db_path.as_deref()?;
    match GeoIpResolver::open(path) {
        Ok(r) => Some(r),
        Err(e) => {
            tracing::warn!(
                path = %path.display(),
                error = %e,
                "resolve_peer_countries enabled but the GeoIP DB failed to open; \
                 peer country codes will be None"
            );
            None
        }
    }
}

/// Stamp `country_code` onto a `GetPeerInfo` reply. No-op when the
/// resolver is absent (disabled / failed open).
pub(crate) fn stamp_country_codes(
    peers: &mut [crate::types::PeerInfo],
    resolver: Option<&GeoIpResolver>,
) {
    let Some(resolver) = resolver else { return };
    for peer in peers {
        peer.country_code = resolver.lookup(peer.addr.ip());
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::io::Write as _;
    use std::str::FromStr as _;

    #[derive(serde::Serialize)]
    struct CountryRec {
        country: CountryInner,
    }
    #[derive(serde::Serialize)]
    struct CountryInner {
        iso_code: &'static str,
    }

    /// Generate a minimal country DB: 81.2.69.0/24 → GB (`MaxMind`'s
    /// canonical documentation range) + 127.0.0.0/8 → XL (ISO
    /// private-use — lets loopback test peers resolve).
    fn make_test_mmdb() -> Vec<u8> {
        let mut db = maxminddb_writer::Database::default();
        // OV F4: pin a v4 tree explicitly (also the 0.1.2 writer default —
        // verified in source) so v4 CIDR insertions and the reader's v4
        // lookups agree on tree width; all fixture data is v4.
        db.metadata.ip_version = maxminddb_writer::metadata::IpVersion::V4;
        let gb = db
            .insert_value(CountryRec {
                country: CountryInner { iso_code: "GB" },
            })
            .expect("insert GB");
        let xl = db
            .insert_value(CountryRec {
                country: CountryInner { iso_code: "XL" },
            })
            .expect("insert XL");
        db.insert_node(
            maxminddb_writer::paths::IpAddrWithMask::from_str("81.2.69.0/24").expect("cidr"),
            gb,
        );
        db.insert_node(
            maxminddb_writer::paths::IpAddrWithMask::from_str("127.0.0.0/8").expect("cidr"),
            xl,
        );
        db.write_to(Vec::new()).expect("serialize mmdb")
    }

    fn fixture_file() -> tempfile::NamedTempFile {
        let mut f = tempfile::NamedTempFile::new().expect("tmp");
        f.write_all(&make_test_mmdb()).expect("write fixture");
        f.flush().expect("flush");
        f
    }

    fn test_peer(addr: &str) -> crate::types::PeerInfo {
        crate::types::PeerInfo {
            addr: addr.parse().unwrap(),
            client: String::new(),
            peer_choking: false,
            peer_interested: false,
            am_choking: false,
            am_interested: false,
            download_rate: 0,
            upload_rate: 0,
            num_pieces: 0,
            source: irontide_session_types::PeerSource::Tracker,
            supports_fast: false,
            upload_only: false,
            snubbed: false,
            connected_duration_secs: 0,
            num_pending_requests: 0,
            num_incoming_requests: 0,
            is_optimistic: false,
            is_encrypted: false,
            uses_utp: false,
            uses_holepunch: false,
            in_flight_requests: 0,
            target_pipeline_depth: 0,
            relevance: 0.0,
            connection_kind: crate::types::PeerConnectionKind::Tcp,
            progress: 0.0,
            country_code: None,
        }
    }

    #[test]
    fn m255_lookup_known_ip_resolves_gb() {
        let f = fixture_file();
        let r = GeoIpResolver::open(f.path()).expect("open fixture");
        assert_eq!(r.lookup("81.2.69.142".parse().unwrap()), Some(*b"GB"));
    }

    #[test]
    fn m255_lookup_unmapped_ip_is_none() {
        let f = fixture_file();
        let r = GeoIpResolver::open(f.path()).expect("open fixture");
        assert_eq!(r.lookup("10.1.2.3".parse().unwrap()), None);
    }

    #[test]
    fn m255_open_missing_file_errors() {
        assert!(GeoIpResolver::open(Path::new("/nonexistent/geo.mmdb")).is_err());
    }

    #[test]
    fn m255_open_garbage_file_errors() {
        let mut f = tempfile::NamedTempFile::new().expect("tmp");
        f.write_all(b"not an mmdb").expect("write");
        f.flush().expect("flush");
        assert!(GeoIpResolver::open(f.path()).is_err());
    }

    #[test]
    fn m255_build_resolver_disabled_is_none() {
        let s = Settings::default();
        assert!(build_geoip_resolver(&s).is_none());
    }

    #[test]
    fn m255_build_resolver_enabled_bad_path_degrades_to_none() {
        let s = Settings {
            resolve_peer_countries: true,
            peer_country_db_path: Some("/nonexistent/geo.mmdb".into()),
            ..Settings::default()
        };
        assert!(build_geoip_resolver(&s).is_none());
    }

    #[test]
    fn m255_build_resolver_enabled_valid_is_some() {
        let f = fixture_file();
        let s = Settings {
            resolve_peer_countries: true,
            peer_country_db_path: Some(f.path().to_path_buf()),
            ..Settings::default()
        };
        assert!(build_geoip_resolver(&s).is_some());
    }

    #[test]
    fn m255_stamp_country_codes_stamps_and_skips() {
        let f = fixture_file();
        let r = GeoIpResolver::open(f.path()).expect("open fixture");
        let mut peers = vec![
            test_peer("81.2.69.10:6881"),
            test_peer("203.0.113.5:6881"), // TEST-NET-3, unmapped
        ];
        stamp_country_codes(&mut peers, Some(&r));
        assert_eq!(peers[0].country_code, Some(*b"GB"));
        assert_eq!(peers[1].country_code, None);
        // Resolver absent ⇒ untouched no-op.
        let mut more = vec![test_peer("81.2.69.10:6881")];
        stamp_country_codes(&mut more, None);
        assert_eq!(more[0].country_code, None);
    }
}