RustProxy/src/geoip.rs

//! Bounded country GeoIP reader for direct CSV pairs and native MaxMind DB files.
//! Managed downloads use direct HTTPS sources; publisher checksum URLs are optional.
mod mmdb;
mod update;
mod update_http;
pub use update::{
    UpdateResult, UpdateSettings, DEFAULT_INTERVAL_SECONDS, DEFAULT_IPV4_URL, DEFAULT_IPV6_URL,
    DEFAULT_MAX_BYTES, DEFAULT_MMDB_URL,
};
pub(crate) use update::{poll_loop as update_poll_loop, update_now};

use anyhow::{bail, ensure, Context, Result};
use serde::{Deserialize, 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 {
    // Network byte order preserves the numeric ordering of IPv6 addresses
    // without the 16-byte alignment/padding of u128 fields.
    start: [u8; 16],
    end: [u8; 16],
    country: [u8; 2],
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum DatabaseFormat {
    CsvPair,
    Mmdb,
}

impl Default for DatabaseFormat {
    fn default() -> Self {
        Self::CsvPair
    }
}

#[derive(Debug)]
enum Backend {
    Csv { ipv4: Box<[V4Range]>, ipv6: Box<[V6Range]> },
    Mmdb(mmdb::Database),
}

#[derive(Debug)]
pub struct Database {
    backend: Backend,
}
#[derive(Debug, Clone, Serialize)]
pub struct View {
    pub loaded: bool,
    pub ipv4_ranges: usize,
    pub ipv6_ranges: usize,
    pub source: String,
    pub format: &'static str,
}
impl Database {
    pub fn empty() -> Self {
        Self { backend: Backend::Csv { ipv4: Box::new([]), ipv6: Box::new([]) } }
    }
    pub fn load(directory: &Path) -> Result<Self> {
        if directory.join("country.mmdb").is_file() {
            Self::load_format(directory, DatabaseFormat::Mmdb)
        } else {
            Self::load_format(directory, DatabaseFormat::CsvPair)
        }
    }
    pub fn load_format(directory: &Path, format: DatabaseFormat) -> Result<Self> {
        match format {
            DatabaseFormat::CsvPair => {
                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");
                // Convert to exactly sized immutable slices: the CSV parser may have
                // doubled Vec capacity many times during loading.
                Ok(Self { backend: Backend::Csv { ipv4: ipv4.into_boxed_slice(), ipv6: ipv6.into_boxed_slice() } })
            }
            DatabaseFormat::Mmdb => Ok(Self {
                backend: Backend::Mmdb(mmdb::Database::load(&directory.join("country.mmdb"), MAX_FILE_BYTES)?),
            }),
        }
    }
    pub fn loaded(&self) -> bool {
        match &self.backend {
            Backend::Csv { ipv4, ipv6 } => !ipv4.is_empty() || !ipv6.is_empty(),
            Backend::Mmdb(_) => true,
        }
    }
    pub fn view(&self) -> View {
        match &self.backend {
            Backend::Csv { ipv4, ipv6 } => View {
                loaded: self.loaded(),
                ipv4_ranges: ipv4.len(),
                ipv6_ranges: ipv6.len(),
                source: "CSV country ranges".into(),
                format: "csv_pair",
            },
            Backend::Mmdb(database) => View {
                loaded: true,
                ipv4_ranges: 0,
                ipv6_ranges: 0,
                source: format!("MaxMind DB ({})", database.database_type()),
                format: "mmdb",
            },
        }
    }
    pub fn country(&self, ip: IpAddr) -> Option<String> {
        let ip = crate::vhost::canonical_ip(ip);
        match &self.backend {
            Backend::Csv { ipv4, ipv6 } => {
                let code = match ip {
                    IpAddr::V4(ip) => find_v4(ipv4, u32::from(ip)),
                    IpAddr::V6(ip) => find_v6(ipv6, u128::from(ip).to_be_bytes()),
                }?;
                Some(String::from_utf8_lossy(&code).into_owned())
            }
            Backend::Mmdb(database) => database.country(ip),
        }
    }
}

pub(super) fn normalize_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) = normalize_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) = normalize_country_code(country) {
            result.push(V6Range {
                start: start.to_be_bytes(),
                end: end.to_be_bytes(),
                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: [u8; 16]) -> 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:#}"))
}

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

#[cfg(test)]
include!("geoip/tests.rs");