AWG2_harness/traffic_core/src/config.rs

use std::{fs, net::Ipv4Addr, path::Path, process::Command};

use serde::Deserialize;

#[derive(Debug, Clone, Deserialize)]
pub struct Config {
    pub amnezia: AmneziaConfig,
    pub traffic: TrafficConfig,
    pub anti_torrent: AntiTorrentConfig,
    pub reboot: RebootConfig,
    pub swap: SwapConfig,
}

#[derive(Debug, Clone, Deserialize)]
pub struct AmneziaConfig {
    pub container_names: Vec<String>,
    pub interface_candidates: Vec<String>,
    pub interface_patterns: Vec<String>,
    pub peer_ipv4_cidr: String,
}

#[derive(Debug, Clone, Deserialize)]
pub struct TrafficConfig {
    pub enabled: bool,
    pub peer_limit: String,
    pub root_rate: String,
    pub reconcile_interval_secs: u64,
    pub container_ready_timeout_secs: u64,
    pub class_start_id: u16,
    pub ifb_name: String,
}

#[derive(Debug, Clone, Deserialize)]
pub struct AntiTorrentConfig {
    pub enabled: bool,
    pub max_connections_per_peer: u32,
    pub peer_ipv4_cidr: String,
    pub apply_in_amnezia_netns: bool,
}

#[derive(Debug, Clone, Deserialize)]
pub struct RebootConfig {
    pub enabled: bool,
    pub window_start_utc: String,
    pub window_end_utc: String,
}

#[derive(Debug, Clone, Deserialize)]
pub struct SwapConfig {
    pub enabled: bool,
    pub path: String,
    pub size_mb: u64,
    pub priority: i32,
    pub swappiness: u8,
}

impl Config {
    pub fn load(path: &Path) -> anyhow::Result<Self> {
        let raw = fs::read_to_string(path)?;
        let cfg: Self = toml::from_str(&raw)?;
        Ok(cfg)
    }

    pub fn validate(&self) -> anyhow::Result<()> {
        if self.amnezia.container_names.is_empty() {
            anyhow::bail!("amnezia.container_names must not be empty");
        }
        if self.traffic.class_start_id < 10 {
            anyhow::bail!("traffic.class_start_id must be >= 10");
        }
        if self.swap.swappiness > 100 {
            anyhow::bail!("swap.swappiness must be <= 100");
        }
        if self.traffic.container_ready_timeout_secs == 0 {
            anyhow::bail!("traffic.container_ready_timeout_secs must be > 0");
        }
        if self.traffic.reconcile_interval_secs == 0 {
            anyhow::bail!("traffic.reconcile_interval_secs must be > 0");
        }
        if self.reboot.enabled && self.reboot.window_start_utc == self.reboot.window_end_utc {
            anyhow::bail!("reboot window start and end must differ");
        }
        parse_cidr_v4(&self.amnezia.peer_ipv4_cidr)?;
        parse_cidr_v4(&self.anti_torrent.peer_ipv4_cidr)?;
        self.check_host_tools()?;
        Ok(())
    }

    fn check_host_tools(&self) -> anyhow::Result<()> {
        let mut required = vec!["tc", "ip", "wg", "docker", "modprobe"];
        if !crate::netns::nsenter_available() {
            eprintln!("nsenter is not available, will use docker exec fallback");
        } else {
            required.push("nsenter");
        }
        if self.anti_torrent.enabled {
            required.push("nft");
        }

        for tool in required {
            if !tool_exists(tool) {
                anyhow::bail!("required host tool not found in PATH: {tool}");
            }
        }
        Ok(())
    }
}

fn tool_exists(tool: &str) -> bool {
    let probes: &[&[&str]] = match tool {
        "tc" => &[&["-V"]],
        "ip" => &[&["-V"], &["-Version"]],
        "wg" => &[&["--version"]],
        "docker" => &[&["--version"]],
        "modprobe" => &[&["--version"]],
        "nsenter" => &[&["--help"]],
        "nft" => &[&["--version"]],
        _ => &[&["--version"], &["--help"], &["-V"]],
    };

    for args in probes {
        if Command::new(tool).args(*args).output().is_ok() {
            return true;
        }
    }
    false
}

pub fn parse_cidr_v4(cidr: &str) -> anyhow::Result<(Ipv4Addr, u8)> {
    let mut parts = cidr.split('/');
    let ip = parts
        .next()
        .ok_or_else(|| anyhow::anyhow!("invalid cidr: {cidr}"))?
        .parse::<Ipv4Addr>()?;
    let prefix = parts
        .next()
        .ok_or_else(|| anyhow::anyhow!("invalid cidr: {cidr}"))?
        .parse::<u8>()?;
    if prefix > 32 {
        anyhow::bail!("invalid prefix in cidr: {cidr}");
    }
    Ok((ip, prefix))
}

pub fn ip_in_cidr(ip: Ipv4Addr, cidr: &str) -> anyhow::Result<bool> {
    let (base, prefix) = parse_cidr_v4(cidr)?;
    let mask = if prefix == 0 {
        0u32
    } else {
        u32::MAX << (32 - prefix)
    };
    let ip_u = u32::from(ip);
    let base_u = u32::from(base);
    Ok((ip_u & mask) == (base_u & mask))
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn parse_cidr_invalid_prefix() {
        let err = parse_cidr_v4("10.8.1.0/33").err();
        assert!(err.is_some());
    }

    #[test]
    fn ip_membership() {
        let ok = ip_in_cidr("10.8.1.10".parse().expect("ip"), "10.8.1.0/24").expect("cidr");
        assert!(ok);
        let no = ip_in_cidr("10.9.1.10".parse().expect("ip"), "10.8.1.0/24").expect("cidr");
        assert!(!no);
    }
}