AWG2_harness/traffic_core/src/traffic.rs

use crate::{
    config::Config,
    docker,
    lock::FileLock,
    netns::{NsTransport, build_transport},
    wg,
};

use std::collections::BTreeSet;

#[derive(Debug, Clone)]
struct TrafficDrift {
    has_drift: bool,
    reason: String,
}

pub fn inspect_drift_summary(cfg: &Config) -> anyhow::Result<Option<String>> {
    if !cfg.traffic.enabled {
        return Ok(None);
    }

    let target = docker::discover_target(cfg)?;
    let transport = build_transport(target.container_pid, &target.container_name);
    let peers = wg::discover_peers(
        target.container_pid,
        &target.container_name,
        &cfg.amnezia.peer_ipv4_cidr,
    )?;

    if peers.is_empty() {
        return Ok(Some("no peers discovered".to_string()));
    }

    let drift = inspect_and_diff(
        transport.as_ref(),
        &target.interface_name,
        &cfg.traffic.ifb_name,
        &peers,
    )?;

    if drift.has_drift {
        Ok(Some(drift.reason))
    } else {
        Ok(None)
    }
}

pub fn reconcile_traffic(cfg: &Config) -> anyhow::Result<()> {
    if !cfg.traffic.enabled {
        return Ok(());
    }

    let _lock = FileLock::acquire("/run/awg2-harness/traffic.lock")?;
    let target = docker::discover_target(cfg)?;
    let transport = build_transport(target.container_pid, &target.container_name);
    let peers = wg::discover_peers(
        target.container_pid,
        &target.container_name,
        &cfg.amnezia.peer_ipv4_cidr,
    )?;

    if peers.is_empty() {
        println!("no peers found, skip tc apply");
        return Ok(());
    }

    let drift = inspect_and_diff(
        transport.as_ref(),
        &target.interface_name,
        &cfg.traffic.ifb_name,
        &peers,
    )?;

    if !drift.has_drift {
        println!("traffic rules up-to-date, skip apply");
        return Ok(());
    }

    println!("traffic drift detected: {}", drift.reason);

    if cfg.traffic.full_rebuild_on_drift {
        teardown_owned_topology(transport.as_ref(), &target.interface_name, &cfg.traffic.ifb_name)?;
        apply_full_topology(
            cfg,
            transport.as_ref(),
            &target.interface_name,
            &cfg.traffic.ifb_name,
            &peers,
        )?;
    }

    Ok(())
}

fn inspect_and_diff(
    transport: &dyn NsTransport,
    iface: &str,
    ifb: &str,
    peers: &[wg::Peer],
) -> anyhow::Result<TrafficDrift> {
    let root_qdisc = transport
        .run(&["sh", "-c", &format!("tc qdisc show dev {iface} 2>/dev/null")])
        .unwrap_or_default();
    if !root_qdisc.contains("htb") {
        return Ok(TrafficDrift {
            has_drift: true,
            reason: "missing egress root htb qdisc".to_string(),
        });
    }

    let ingress_qdisc = transport
        .run(&["sh", "-c", &format!("tc qdisc show dev {iface} 2>/dev/null")])
        .unwrap_or_default();
    if !ingress_qdisc.contains("ingress") {
        return Ok(TrafficDrift {
            has_drift: true,
            reason: "missing ingress qdisc".to_string(),
        });
    }

    let ifb_link = transport
        .run(&["sh", "-c", &format!("ip link show {ifb} 2>/dev/null")])
        .unwrap_or_default();
    if ifb_link.trim().is_empty() {
        return Ok(TrafficDrift {
            has_drift: true,
            reason: "missing ifb interface".to_string(),
        });
    }

    let filters = transport
        .run(&[
            "sh",
            "-c",
            &format!("tc filter show dev {iface} parent 1: 2>/dev/null"),
        ])
        .unwrap_or_default();

    let actual_ips = extract_filter_ips(&filters);
    let expected_ips: BTreeSet<String> = peers.iter().map(|p| p.ipv4.to_string()).collect();

    if actual_ips != expected_ips {
        return Ok(TrafficDrift {
            has_drift: true,
            reason: format!(
                "peer filter mismatch (actual={}, expected={})",
                actual_ips.len(),
                expected_ips.len()
            ),
        });
    }

    Ok(TrafficDrift {
        has_drift: false,
        reason: "no drift".to_string(),
    })
}

fn extract_filter_ips(filters: &str) -> BTreeSet<String> {
    let mut set = BTreeSet::new();
    for line in filters.lines() {
        if !line.contains("match") || !line.contains("dst") {
            continue;
        }
        let toks: Vec<&str> = line.split_whitespace().collect();
        for (idx, tok) in toks.iter().enumerate() {
            if *tok == "dst" {
                if let Some(ip) = toks.get(idx + 1) {
                    set.insert((*ip).to_string());
                }
                break;
            }
        }
    }
    set
}

fn teardown_owned_topology(transport: &dyn NsTransport, iface: &str, ifb: &str) -> anyhow::Result<()> {
    let cmd = format!(
        "tc qdisc del dev {iface} root 2>/dev/null || true; \
         tc qdisc del dev {iface} ingress 2>/dev/null || true; \
         ip link del {ifb} 2>/dev/null || true"
    );
    let _ = transport.run(&["sh", "-c", &cmd])?;
    Ok(())
}

fn apply_full_topology(
    cfg: &Config,
    transport: &dyn NsTransport,
    iface: &str,
    ifb: &str,
    peers: &[wg::Peer],
) -> anyhow::Result<()> {
    let setup_cmd = format!(
        "tc qdisc add dev {iface} root handle 1: htb; \
         tc class add dev {iface} parent 1: classid 1:1 htb rate {root} ceil {root}; \
         modprobe ifb 2>/dev/null || true; \
         ip link add name {ifb} type ifb 2>/dev/null || true; \
         ip link set dev {ifb} up; \
         tc qdisc add dev {iface} handle ffff: ingress; \
         tc filter add dev {iface} parent ffff: protocol all u32 match u32 0 0 action mirred egress redirect dev {ifb}; \
         tc qdisc add dev {ifb} root handle 2: htb; \
         tc class add dev {ifb} parent 2: classid 2:1 htb rate {root} ceil {root}",
        root = cfg.traffic.root_rate
    );
    let _ = transport.run(&["sh", "-c", &setup_cmd])?;

    for (idx, p) in peers.iter().enumerate() {
        let class_id = cfg.traffic.class_start_id + idx as u16;
        let ip = p.ipv4;
        let per_peer = format!(
            "tc class add dev {iface} parent 1:1 classid 1:{class_id} htb rate {limit} ceil {limit}; \
             tc qdisc add dev {iface} parent 1:{class_id} fq_codel; \
             tc filter add dev {iface} parent 1: protocol ip u32 match ip dst {ip} flowid 1:{class_id}; \
             tc class add dev {ifb} parent 2:1 classid 2:{class_id} htb rate {limit} ceil {limit}; \
             tc qdisc add dev {ifb} parent 2:{class_id} fq_codel; \
             tc filter add dev {ifb} parent 2: protocol ip u32 match ip src {ip} flowid 2:{class_id}",
            limit = cfg.traffic.peer_limit
        );
        let _ = transport.run(&["sh", "-c", &per_peer])?;
    }

    Ok(())
}