AWG2_harness/traffic_core/src/docker.rs

use std::process::Command;

use serde_json::Value;

use crate::{config::Config, error::AppError};

#[derive(Debug, Clone)]
pub struct RuntimeTarget {
    pub container_name: String,
    pub container_id: String,
    pub container_pid: u32,
    pub netns_path: String,
    pub interface_name: String,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TriState {
    True,
    False,
    Unknown,
}

impl TriState {
    pub fn as_str(self) -> &'static str {
        match self {
            Self::True => "true",
            Self::False => "false",
            Self::Unknown => "unknown",
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PublishedUdpPortPath {
    pub host_port: u16,
    pub container_ip: String,
    pub container_port: u16,
}

#[derive(Debug, Clone)]
pub struct PortPathDiagnostics {
    pub docker_userland_proxy: TriState,
    pub published_udp_port: Option<PublishedUdpPortPath>,
    pub host_forward_policy_drop: TriState,
    pub docker_forward_rules_present: TriState,
}

impl PortPathDiagnostics {
    pub fn warning(&self) -> Option<&'static str> {
        if self.docker_userland_proxy == TriState::False
            && self.host_forward_policy_drop == TriState::True
            && self.docker_forward_rules_present != TriState::True
        {
            return Some(
                "userland-proxy=false AND host FORWARD policy DROP AND Docker accept rules are missing/unknown",
            );
        }
        None
    }
}

pub fn discover_target(cfg: &Config) -> anyhow::Result<RuntimeTarget> {
    let container_name = find_running_container(&cfg.amnezia.container_names)?;
    let container_id = run_capture(&["docker", "inspect", "--format", "{{.Id}}", &container_name])?
        .trim()
        .to_string();
    let pid_str = run_capture(&[
        "docker",
        "inspect",
        "--format",
        "{{.State.Pid}}",
        &container_name,
    ])?;
    let container_pid = pid_str.trim().parse::<u32>()?;
    let netns_path = format!("/proc/{container_pid}/ns/net");
    let interface_name = crate::iface::detect_interface(cfg, container_pid, &container_name)?;

    Ok(RuntimeTarget {
        container_name,
        container_id,
        container_pid,
        netns_path,
        interface_name,
    })
}

fn find_running_container(names: &[String]) -> anyhow::Result<String> {
    let ps = run_capture(&["docker", "ps", "--format", "{{.Names}}"])?;
    for name in names {
        if ps.lines().any(|x| x.trim() == name) {
            return Ok(name.clone());
        }
    }
    Err(AppError::NotFound("amnezia container not running".to_string()).into())
}

pub fn run_capture(args: &[&str]) -> anyhow::Result<String> {
    let (program, rest) = args
        .split_first()
        .ok_or_else(|| anyhow::anyhow!("empty command"))?;
    let out = Command::new(program).args(rest).output()?;
    if !out.status.success() {
        return Err(AppError::Command(format!(
            "{} {:?}: {}",
            program,
            rest,
            String::from_utf8_lossy(&out.stderr)
        ))
        .into());
    }
    Ok(String::from_utf8_lossy(&out.stdout).to_string())
}

pub fn collect_port_path_diagnostics(
    container_name: &str,
    expected_udp_port: u16,
) -> PortPathDiagnostics {
    let docker_userland_proxy = detect_userland_proxy().unwrap_or(TriState::Unknown);
    let published_udp_port = discover_published_udp_port(container_name, expected_udp_port).ok();
    let (host_forward_policy_drop, docker_forward_rules_present) =
        inspect_iptables_forward_state().unwrap_or((TriState::Unknown, TriState::Unknown));

    PortPathDiagnostics {
        docker_userland_proxy,
        published_udp_port,
        host_forward_policy_drop,
        docker_forward_rules_present,
    }
}

fn detect_userland_proxy() -> anyhow::Result<TriState> {
    let raw = run_capture(&["docker", "info", "--format", "{{json .}}"])?;
    Ok(parse_userland_proxy_from_docker_info(&raw))
}

fn discover_published_udp_port(
    container_name: &str,
    expected_udp_port: u16,
) -> anyhow::Result<PublishedUdpPortPath> {
    let raw = run_capture(&["docker", "inspect", container_name])?;
    parse_udp_port_path_from_inspect(&raw, expected_udp_port)
}

fn inspect_iptables_forward_state() -> anyhow::Result<(TriState, TriState)> {
    let fwd = run_capture(&["iptables", "-S", "FORWARD"])?;
    Ok(parse_forward_policy_and_docker_rules(&fwd))
}

fn parse_userland_proxy_from_docker_info(raw: &str) -> TriState {
    let Ok(v) = serde_json::from_str::<Value>(raw) else {
        return TriState::Unknown;
    };
    match v.get("UserlandProxy").and_then(|x| x.as_bool()) {
        Some(true) => TriState::True,
        Some(false) => TriState::False,
        None => TriState::Unknown,
    }
}

fn parse_udp_port_path_from_inspect(
    raw: &str,
    expected_udp_port: u16,
) -> anyhow::Result<PublishedUdpPortPath> {
    let v: Value = serde_json::from_str(raw)?;
    let arr = v
        .as_array()
        .ok_or_else(|| anyhow::anyhow!("docker inspect output is not array"))?;
    let obj = arr
        .first()
        .ok_or_else(|| anyhow::anyhow!("docker inspect output is empty"))?;

    let networks = obj
        .get("NetworkSettings")
        .and_then(|x| x.get("Networks"))
        .and_then(|x| x.as_object());
    let container_ip = networks
        .and_then(|nets| {
            nets.values()
                .find_map(|n| n.get("IPAddress").and_then(|x| x.as_str()))
        })
        .unwrap_or("unknown")
        .to_string();

    let ports = obj
        .get("NetworkSettings")
        .and_then(|x| x.get("Ports"))
        .and_then(|x| x.as_object())
        .ok_or_else(|| anyhow::anyhow!("docker inspect has no NetworkSettings.Ports"))?;

    let mut udp_bindings: Vec<(u16, u16)> = Vec::new();
    for (k, v) in ports {
        if !k.ends_with("/udp") {
            continue;
        }
        let Some(host_port) = extract_host_port_from_binding(v) else {
            continue;
        };
        let container_port = k
            .split('/')
            .next()
            .and_then(|x| x.parse::<u16>().ok())
            .ok_or_else(|| anyhow::anyhow!("invalid udp port key in docker inspect: {k}"))?;
        udp_bindings.push((host_port, container_port));
    }

    if let Some((host_port, container_port)) = udp_bindings
        .iter()
        .copied()
        .find(|(host_port, _)| *host_port == expected_udp_port)
    {
        return Ok(PublishedUdpPortPath {
            host_port,
            container_ip,
            container_port,
        });
    }

    let expected_key = format!("{expected_udp_port}/udp");
    if let Some(host_port) = extract_host_port_for_udp_key(ports, &expected_key) {
        return Ok(PublishedUdpPortPath {
            host_port,
            container_ip,
            container_port: expected_udp_port,
        });
    }

    if udp_bindings.len() == 1 {
        let (host_port, container_port) = udp_bindings[0];
        return Ok(PublishedUdpPortPath {
            host_port,
            container_ip,
            container_port,
        });
    }

    if udp_bindings.len() > 1 {
        return Err(anyhow::anyhow!(
            "ambiguous udp published ports in docker inspect"
        ));
    }

    Err(anyhow::anyhow!(
        "no published udp port found in docker inspect"
    ))
}

fn extract_host_port_for_udp_key(ports: &serde_json::Map<String, Value>, key: &str) -> Option<u16> {
    let v = ports.get(key)?;
    extract_host_port_from_binding(v)
}

fn extract_host_port_from_binding(v: &Value) -> Option<u16> {
    let arr = v.as_array()?;
    let first = arr.first()?.as_object()?;
    first
        .get("HostPort")
        .and_then(|x| x.as_str())
        .and_then(|x| x.parse::<u16>().ok())
}

fn parse_forward_policy_and_docker_rules(raw: &str) -> (TriState, TriState) {
    let mut policy_drop = TriState::Unknown;
    let mut docker_rules = TriState::False;

    for line in raw.lines() {
        let t = line.trim();
        if t.starts_with("-P FORWARD ") {
            policy_drop = if t.eq_ignore_ascii_case("-P FORWARD DROP") {
                TriState::True
            } else {
                TriState::False
            };
        }
        if t.contains(" DOCKER") || t.contains(" DOCKER-") {
            docker_rules = TriState::True;
        }
    }

    (policy_drop, docker_rules)
}

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

    #[test]
    fn parse_userland_proxy_from_docker_info_works() {
        assert_eq!(
            parse_userland_proxy_from_docker_info(r#"{"UserlandProxy":true}"#),
            TriState::True
        );
        assert_eq!(
            parse_userland_proxy_from_docker_info(r#"{"UserlandProxy":false}"#),
            TriState::False
        );
        assert_eq!(
            parse_userland_proxy_from_docker_info("{}"),
            TriState::Unknown
        );
    }

    #[test]
    fn parse_udp_port_path_from_inspect_works() {
        let raw = r#"[
  {
    "NetworkSettings": {
      "Networks": {
        "bridge": { "IPAddress": "172.29.172.2" }
      },
      "Ports": {
        "443/udp": [{ "HostIp": "0.0.0.0", "HostPort": "443" }]
      }
    }
  }
]"#;
        let got = parse_udp_port_path_from_inspect(raw, 443).expect("parse inspect");
        assert_eq!(
            got,
            PublishedUdpPortPath {
                host_port: 443,
                container_ip: "172.29.172.2".to_string(),
                container_port: 443,
            }
        );
    }

    #[test]
    fn parse_udp_port_prefers_host_port_match() {
        let raw = r#"[
  {
    "NetworkSettings": {
      "Networks": {
        "bridge": { "IPAddress": "172.29.172.2" }
      },
      "Ports": {
        "44344/udp": [{ "HostIp": "0.0.0.0", "HostPort": "443" }],
        "443/udp": [{ "HostIp": "0.0.0.0", "HostPort": "5443" }]
      }
    }
  }
]"#;
        let got = parse_udp_port_path_from_inspect(raw, 443).expect("parse inspect");
        assert_eq!(
            got,
            PublishedUdpPortPath {
                host_port: 443,
                container_ip: "172.29.172.2".to_string(),
                container_port: 44344,
            }
        );
    }

    #[test]
    fn parse_udp_port_fallbacks_to_expected_key() {
        let raw = r#"[
  {
    "NetworkSettings": {
      "Networks": {
        "bridge": { "IPAddress": "172.29.172.2" }
      },
      "Ports": {
        "443/udp": [{ "HostIp": "0.0.0.0", "HostPort": "5443" }],
        "44344/udp": [{ "HostIp": "0.0.0.0", "HostPort": "15443" }]
      }
    }
  }
]"#;
        let got = parse_udp_port_path_from_inspect(raw, 443).expect("parse inspect");
        assert_eq!(
            got,
            PublishedUdpPortPath {
                host_port: 5443,
                container_ip: "172.29.172.2".to_string(),
                container_port: 443,
            }
        );
    }

    #[test]
    fn parse_udp_port_uses_single_udp_binding() {
        let raw = r#"[
  {
    "NetworkSettings": {
      "Networks": {
        "bridge": { "IPAddress": "172.29.172.2" }
      },
      "Ports": {
        "44344/udp": [{ "HostIp": "0.0.0.0", "HostPort": "44344" }]
      }
    }
  }
]"#;
        let got = parse_udp_port_path_from_inspect(raw, 443).expect("parse inspect");
        assert_eq!(
            got,
            PublishedUdpPortPath {
                host_port: 44344,
                container_ip: "172.29.172.2".to_string(),
                container_port: 44344,
            }
        );
    }

    #[test]
    fn parse_udp_port_returns_ambiguous_on_multiple_without_match() {
        let raw = r#"[
  {
    "NetworkSettings": {
      "Networks": {
        "bridge": { "IPAddress": "172.29.172.2" }
      },
      "Ports": {
        "10000/udp": [{ "HostIp": "0.0.0.0", "HostPort": "10000" }],
        "20000/udp": [{ "HostIp": "0.0.0.0", "HostPort": "20000" }]
      }
    }
  }
]"#;
        let err = parse_udp_port_path_from_inspect(raw, 443).expect_err("must be ambiguous");
        assert!(err.to_string().contains("ambiguous udp published ports"));
    }
    #[test]
    fn parse_forward_policy_and_docker_rules_works() {
        let raw = r#"
-P FORWARD DROP
-A FORWARD -j DOCKER-USER
-A FORWARD -j DOCKER-FORWARD
"#;
        assert_eq!(
            parse_forward_policy_and_docker_rules(raw),
            (TriState::True, TriState::True)
        );
    }

    #[test]
    fn warning_is_emitted_for_risky_forward_state() {
        let diag = PortPathDiagnostics {
            docker_userland_proxy: TriState::False,
            published_udp_port: None,
            host_forward_policy_drop: TriState::True,
            docker_forward_rules_present: TriState::Unknown,
        };
        assert!(diag.warning().is_some());
    }
}