RustProxy/src/geoip.rs

//! Minimal reader for sapics/ip-location-db user-country CSV files.
//! Installer verification pins each downloaded file to the publisher-provided
//! SHA-256 checksum; runtime parsing stays dependency-free and bounded.
use anyhow::{bail, ensure, Context, Result};
use serde::Serialize;
use std::{
    fs::File,
    io::{BufRead, BufReader},
    net::{IpAddr, Ipv4Addr, Ipv6Addr},
    path::Path,
};

const MAX_FILE_BYTES: u64 = 256 * 1024 * 1024;
const MAX_RANGES: usize = 2_000_000;

#[derive(Debug, Clone, Copy)]
struct V4Range {
    start: u32,
    end: u32,
    country: [u8; 2],
}
#[derive(Debug, Clone, Copy)]
struct V6Range {
    start: u128,
    end: u128,
    country: [u8; 2],
}
#[derive(Debug, Default)]
pub struct Database {
    ipv4: Vec<V4Range>,
    ipv6: Vec<V6Range>,
}
#[derive(Debug, Clone, Serialize)]
pub struct View {
    pub loaded: bool,
    pub ipv4_ranges: usize,
    pub ipv6_ranges: usize,
    pub source: &'static str,
}
impl Database {
    pub fn empty() -> Self {
        Self::default()
    }
    pub fn load(directory: &Path) -> Result<Self> {
        let ipv4 = parse_v4(&directory.join("user-country-ipv4.csv"))?;
        let ipv6 = parse_v6(&directory.join("user-country-ipv6.csv"))?;
        ensure!(
            !ipv4.is_empty() && !ipv6.is_empty(),
            "GeoIP country database is empty"
        );
        Ok(Self { ipv4, ipv6 })
    }
    pub fn loaded(&self) -> bool {
        !self.ipv4.is_empty() || !self.ipv6.is_empty()
    }
    pub fn view(&self) -> View {
        View {
            loaded: self.loaded(),
            ipv4_ranges: self.ipv4.len(),
            ipv6_ranges: self.ipv6.len(),
            source: "sapics/ip-location-db user-country (PDDL)",
        }
    }
    pub fn country(&self, ip: IpAddr) -> Option<String> {
        let code = match crate::vhost::canonical_ip(ip) {
            IpAddr::V4(ip) => find_v4(&self.ipv4, u32::from(ip)),
            IpAddr::V6(ip) => find_v6(&self.ipv6, u128::from(ip)),
        }?;
        Some(String::from_utf8_lossy(&code).into_owned())
    }
}

fn country_code(raw: &str) -> Option<[u8; 2]> {
    let raw = raw.trim().trim_matches('"').as_bytes();
    if raw.len() == 2 && raw.iter().all(|b| b.is_ascii_alphabetic()) {
        Some([raw[0].to_ascii_uppercase(), raw[1].to_ascii_uppercase()])
    } else {
        None
    }
}
fn prepare(path: &Path) -> Result<BufReader<File>> {
    let file =
        File::open(path).with_context(|| format!("open GeoIP database {}", path.display()))?;
    ensure!(
        file.metadata()?.len() <= MAX_FILE_BYTES,
        "GeoIP database file is larger than 256 MiB"
    );
    Ok(BufReader::new(file))
}
fn fields(line: &str) -> Result<(&str, &str, &str)> {
    let mut parts = line.split(',');
    let start = parts
        .next()
        .context("GeoIP CSV missing start address")?
        .trim();
    let end = parts
        .next()
        .context("GeoIP CSV missing end address")?
        .trim();
    let country = parts.next().context("GeoIP CSV missing country")?.trim();
    ensure!(parts.next().is_none(), "GeoIP CSV has unexpected columns");
    Ok((start, end, country))
}
fn parse_v4(path: &Path) -> Result<Vec<V4Range>> {
    let mut result = Vec::new();
    let mut previous_end: Option<u32> = None;
    for (index, line) in prepare(path)?.lines().enumerate() {
        let line = line.with_context(|| format!("read GeoIP IPv4 row {}", index + 1))?;
        let line = line.trim();
        if line.is_empty() || line.starts_with('#') {
            continue;
        }
        let (start, end, country) = fields(line)?;
        let start: Ipv4Addr = start
            .trim_matches('"')
            .parse()
            .context("invalid GeoIP IPv4 start")?;
        let end: Ipv4Addr = end
            .trim_matches('"')
            .parse()
            .context("invalid GeoIP IPv4 end")?;
        let start = u32::from(start);
        let end = u32::from(end);
        ensure!(start <= end, "invalid GeoIP IPv4 range");
        if let Some(previous) = previous_end {
            ensure!(
                start > previous,
                "GeoIP IPv4 ranges overlap or are unsorted"
            );
        }
        previous_end = Some(end);
        if let Some(country) = country_code(country) {
            result.push(V4Range {
                start,
                end,
                country,
            });
            ensure!(result.len() <= MAX_RANGES, "too many GeoIP IPv4 ranges");
        }
    }
    Ok(result)
}
fn parse_v6(path: &Path) -> Result<Vec<V6Range>> {
    let mut result = Vec::new();
    let mut previous_end: Option<u128> = None;
    for (index, line) in prepare(path)?.lines().enumerate() {
        let line = line.with_context(|| format!("read GeoIP IPv6 row {}", index + 1))?;
        let line = line.trim();
        if line.is_empty() || line.starts_with('#') {
            continue;
        }
        let (start, end, country) = fields(line)?;
        let start: Ipv6Addr = start
            .trim_matches('"')
            .parse()
            .context("invalid GeoIP IPv6 start")?;
        let end: Ipv6Addr = end
            .trim_matches('"')
            .parse()
            .context("invalid GeoIP IPv6 end")?;
        let start = u128::from(start);
        let end = u128::from(end);
        ensure!(start <= end, "invalid GeoIP IPv6 range");
        if let Some(previous) = previous_end {
            ensure!(
                start > previous,
                "GeoIP IPv6 ranges overlap or are unsorted"
            );
        }
        previous_end = Some(end);
        if let Some(country) = country_code(country) {
            result.push(V6Range {
                start,
                end,
                country,
            });
            ensure!(result.len() <= MAX_RANGES, "too many GeoIP IPv6 ranges");
        }
    }
    Ok(result)
}
fn find_v4(ranges: &[V4Range], ip: u32) -> Option<[u8; 2]> {
    let index = ranges
        .partition_point(|item| item.start <= ip)
        .checked_sub(1)?;
    let item = ranges.get(index)?;
    (ip <= item.end).then_some(item.country)
}
fn find_v6(ranges: &[V6Range], ip: u128) -> Option<[u8; 2]> {
    let index = ranges
        .partition_point(|item| item.start <= ip)
        .checked_sub(1)?;
    let item = ranges.get(index)?;
    (ip <= item.end).then_some(item.country)
}

pub fn validate_directory(directory: &Path) -> Result<()> {
    Database::load(directory)
        .map(|_| ())
        .or_else(|error| bail!("GeoIP country database is invalid: {error:#}"))
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn range_lookup_handles_gaps() {
        let db = Database {
            ipv4: vec![V4Range {
                start: u32::from(Ipv4Addr::new(192, 0, 2, 0)),
                end: u32::from(Ipv4Addr::new(192, 0, 2, 255)),
                country: *b"PL",
            }],
            ipv6: vec![],
        };
        assert_eq!(
            db.country("192.0.2.7".parse().unwrap()).as_deref(),
            Some("PL")
        );
        assert!(db.country("198.51.100.1".parse().unwrap()).is_none());
    }
}