AWG2_harness/traffic_core/src/firewall.rs

use crate::{
    config::Config,
    docker,
    netns::{NsTransport, build_host_transport, build_transport},
};
use serde_json::Value;

#[derive(Debug, Clone)]
struct FirewallActual {
    table_exists: bool,
    counter_exists: bool,
    chain_exists: bool,
    rule_matches: bool,
}

pub fn apply_firewall(cfg: &Config) -> anyhow::Result<()> {
    let transport = firewall_transport(cfg)?;

    if !cfg.anti_torrent.enabled {
        run_allow_notfound(
            transport.as_ref(),
            &["nft", "delete", "table", "inet", "awg2_harness"],
        )?;
        return Ok(());
    }

    let actual = inspect_firewall(transport.as_ref(), cfg)?;
    if actual.table_exists && actual.counter_exists && actual.chain_exists && actual.rule_matches {
        return Ok(());
    }

    run_allow_notfound(
        transport.as_ref(),
        &["nft", "delete", "table", "inet", "awg2_harness"],
    )?;

    run_step(
        transport.as_ref(),
        &["nft", "add", "table", "inet", "awg2_harness"],
    )?;
    run_step(
        transport.as_ref(),
        &[
            "nft",
            "add",
            "counter",
            "inet",
            "awg2_harness",
            "antitorrent_dropped_v4",
        ],
    )?;
    run_step_chain(transport.as_ref())?;
    apply_rule(
        transport.as_ref(),
        &cfg.anti_torrent.peer_ipv4_cidr,
        cfg.anti_torrent.max_connections_per_peer,
    )?;

    let verify = inspect_firewall(transport.as_ref(), cfg)?;
    if !(verify.table_exists && verify.counter_exists && verify.chain_exists && verify.rule_matches)
    {
        anyhow::bail!("firewall verify-after-apply failed");
    }
    Ok(())
}

fn inspect_firewall(transport: &dyn NsTransport, cfg: &Config) -> anyhow::Result<FirewallActual> {
    let table = transport
        .run(&["nft", "list", "table", "inet", "awg2_harness"])
        .unwrap_or_default();

    let table_exists = !table.trim().is_empty();
    let counter_exists = table.contains("counter antitorrent_dropped_v4");
    let chain_exists = table.contains("chain forward_antitorrent");

    let needle_cidr = format!("ip saddr {}", cfg.anti_torrent.peer_ipv4_cidr);
    let needle_max = format!(
        "ct count over {}",
        cfg.anti_torrent.max_connections_per_peer
    );
    let rule_matches = table.contains(&needle_cidr)
        && table.contains(&needle_max)
        && table.contains("ct state new")
        && table.contains("meter awg2_peer_connlimit_v4")
        && table.contains("counter name antitorrent_dropped_v4")
        && table.contains("drop")
        && table.contains("hook forward")
        && table.contains("policy accept");

    Ok(FirewallActual {
        table_exists,
        counter_exists,
        chain_exists,
        rule_matches,
    })
}

fn firewall_transport(cfg: &Config) -> anyhow::Result<Box<dyn NsTransport>> {
    if cfg.anti_torrent.apply_in_amnezia_netns {
        let target = docker::discover_target(cfg)?;
        Ok(build_transport(
            target.container_pid,
            &target.container_name,
        ))
    } else {
        Ok(build_host_transport())
    }
}

fn apply_rule(transport: &dyn NsTransport, cidr: &str, max_conn: u32) -> anyhow::Result<()> {
    let _ = transport.run(&[
        "nft",
        "add",
        "rule",
        "inet",
        "awg2_harness",
        "forward_antitorrent",
        "ip",
        "saddr",
        cidr,
        "ct",
        "state",
        "new",
        "meter",
        "awg2_peer_connlimit_v4",
        "{",
        "ip",
        "saddr",
        "ct",
        "count",
        "over",
        &max_conn.to_string(),
        "}",
        "counter",
        "name",
        "antitorrent_dropped_v4",
        "drop",
    ])?;
    Ok(())
}

fn run_step(transport: &dyn NsTransport, args: &[&str]) -> anyhow::Result<()> {
    let _ = transport.run(args)?;
    Ok(())
}

fn run_step_chain(transport: &dyn NsTransport) -> anyhow::Result<()> {
    run_step(
        transport,
        &[
            "nft",
            "add",
            "chain",
            "inet",
            "awg2_harness",
            "forward_antitorrent",
            "{",
            "type",
            "filter",
            "hook",
            "forward",
            "priority",
            "filter;",
            "policy",
            "accept;",
            "}",
        ],
    )
}

fn run_allow_notfound(transport: &dyn NsTransport, args: &[&str]) -> anyhow::Result<()> {
    match transport.run(args) {
        Ok(_) => Ok(()),
        Err(e) => {
            let msg = e.to_string();
            if msg.contains("No such file")
                || msg.contains("No such")
                || msg.contains("not found")
                || msg.contains("not exist")
            {
                return Ok(());
            }
            Err(e)
        }
    }
}

pub fn antitorrent_drop_counter(cfg: &Config) -> anyhow::Result<Option<(u64, u64)>> {
    if !cfg.anti_torrent.enabled {
        return Ok(None);
    }

    let transport = firewall_transport(cfg)?;
    let json = transport.run(&[
        "nft",
        "-j",
        "list",
        "counter",
        "inet",
        "awg2_harness",
        "antitorrent_dropped_v4",
    ]);

    if let Ok(raw) = json {
        if let Some(v) = parse_counter_json(&raw)? {
            return Ok(Some(v));
        }
    }

    let table = transport
        .run(&["nft", "list", "table", "inet", "awg2_harness"])
        .unwrap_or_default();
    for line in table.lines() {
        let t = line.trim();
        if !t.contains("packets") {
            continue;
        }
        let toks: Vec<&str> = t.split_whitespace().collect();
        let packets = toks
            .iter()
            .position(|x| *x == "packets")
            .and_then(|i| toks.get(i + 1))
            .and_then(|s| s.parse::<u64>().ok());
        let bytes = toks
            .iter()
            .position(|x| *x == "bytes")
            .and_then(|i| toks.get(i + 1))
            .and_then(|s| s.parse::<u64>().ok())
            .unwrap_or(0);
        if let Some(p) = packets {
            return Ok(Some((p, bytes)));
        }
    }

    Ok(None)
}

fn parse_counter_json(raw: &str) -> anyhow::Result<Option<(u64, u64)>> {
    let v: Value = serde_json::from_str(raw)?;
    let Some(arr) = v.get("nftables").and_then(|x| x.as_array()) else {
        return Ok(None);
    };
    for item in arr {
        if let Some(counter) = item.get("counter") {
            let packets = counter.get("packets").and_then(|x| x.as_u64());
            let bytes = counter.get("bytes").and_then(|x| x.as_u64());
            if let (Some(p), Some(b)) = (packets, bytes) {
                return Ok(Some((p, b)));
            }
        }
    }
    Ok(None)
}

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

    #[test]
    fn parse_nft_counter_json() {
        let raw = r#"{"nftables":[{"metainfo":{"json_schema_version":1}},{"counter":{"family":"inet","table":"awg2_harness","name":"antitorrent_dropped_v4","handle":2,"packets":7,"bytes":448}}]}"#;
        let got = parse_counter_json(raw).expect("json parse");
        assert_eq!(got, Some((7, 448)));
    }
}