ctf-ad-mirror-tool-rust/tests/control_plane.rs

use mirror_tool::{
    auth::Auth,
    config::{Config, ReplicaPolicy, ServiceConfig, Team},
    control::Controller,
    web,
};
use serde_json::{Value, json};
use std::{net::SocketAddr, sync::Arc};
use tokio::{
    io::{AsyncReadExt, AsyncWriteExt},
    net::TcpListener,
};
use tokio_util::sync::CancellationToken;

async fn fixture() -> (
    Arc<Controller>,
    String,
    CancellationToken,
    tempfile::TempDir,
    tokio::task::JoinHandle<()>,
) {
    let temp = tempfile::tempdir().unwrap();
    let mut config = Config::default();
    config.web.state_file = temp.path().join("state.json");
    let controller = Controller::start(config.clone()).await.unwrap();
    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let addr = listener.local_addr().unwrap();
    let cancel = CancellationToken::new();
    let router = web::router(
        controller.clone(),
        Auth::new(config.web, "fixture-passphrase"),
        cancel.clone(),
    );
    let stop = cancel.clone();
    let task = tokio::spawn(async move {
        axum::serve(
            listener,
            router.into_make_service_with_connect_info::<SocketAddr>(),
        )
        .with_graceful_shutdown(stop.cancelled_owned())
        .await
        .unwrap();
    });
    (controller, format!("http://{addr}"), cancel, temp, task)
}
async fn login(client: &reqwest::Client, url: &str) -> (String, String) {
    let response = client
        .post(format!("{url}/api/login"))
        .json(&json!({"password":"fixture-passphrase"}))
        .send()
        .await
        .unwrap();
    assert_eq!(response.status(), 200);
    let cookie = response.headers()["set-cookie"]
        .to_str()
        .unwrap()
        .split(';')
        .next()
        .unwrap()
        .to_owned();
    let body: Value = response.json().await.unwrap();
    (cookie, body["csrf"].as_str().unwrap().into())
}
#[tokio::test]
async fn authenticated_api_csrf_persistence_and_sse() {
    let (controller, url, cancel, temp, task) = fixture().await;
    let client = reqwest::Client::new();
    assert_eq!(
        client
            .get(format!("{url}/api/snapshot"))
            .send()
            .await
            .unwrap()
            .status(),
        401
    );
    assert_eq!(
        client
            .post(format!("{url}/api/login"))
            .header("origin", "http://other.example")
            .json(&json!({"password":"fixture-passphrase"}))
            .send()
            .await
            .unwrap()
            .status(),
        403
    );
    let (cookie, csrf) = login(&client, &url).await;
    let team = json!({"id":"local","name":"Local","address":"127.0.0.1"});
    assert_eq!(
        client
            .put(format!("{url}/api/teams/local"))
            .header("cookie", &cookie)
            .json(&team)
            .send()
            .await
            .unwrap()
            .status(),
        403
    );
    assert_eq!(
        client
            .put(format!("{url}/api/teams/local"))
            .header("cookie", &cookie)
            .header("x-csrf-token", &csrf)
            .json(&team)
            .send()
            .await
            .unwrap()
            .status(),
        200
    );
    let service = json!({"id":"notes","name":"Notes","listen":"127.0.0.1:18081","upstream_port":18082,"primary":"local","replicas":[],"enabled":false,"auto_switch":true,"replica_policy":"isolate"});
    assert_eq!(
        client
            .put(format!("{url}/api/services/notes"))
            .header("cookie", &cookie)
            .header("x-csrf-token", &csrf)
            .json(&service)
            .send()
            .await
            .unwrap()
            .status(),
        200
    );
    let saved = std::fs::read_to_string(temp.path().join("state.json")).unwrap();
    assert!(saved.contains("notes"));
    assert!(!saved.contains("fixture-passphrase"));
    assert!(!saved.contains(&csrf));
    let snapshot: Value = client
        .get(format!("{url}/api/snapshot"))
        .header("cookie", &cookie)
        .send()
        .await
        .unwrap()
        .json()
        .await
        .unwrap();
    assert_eq!(snapshot["services"][0]["id"], "notes");
    let mut events = client
        .get(format!("{url}/api/events"))
        .header("cookie", &cookie)
        .send()
        .await
        .unwrap();
    let first = events.chunk().await.unwrap().unwrap();
    assert!(
        std::str::from_utf8(&first)
            .unwrap()
            .contains("event: snapshot")
    );
    drop(events);
    assert_eq!(
        client
            .post(format!("{url}/api/logout"))
            .header("cookie", &cookie)
            .header("x-csrf-token", &csrf)
            .send()
            .await
            .unwrap()
            .status(),
        200
    );
    assert_eq!(
        client
            .get(format!("{url}/api/session"))
            .header("cookie", &cookie)
            .send()
            .await
            .unwrap()
            .status(),
        401
    );
    cancel.cancel();
    task.await.unwrap();
    controller.engine.shutdown().await;
    let restarted = Controller::start(controller.config.clone()).await.unwrap();
    assert_eq!(restarted.topology().services[0].id, "notes");
    restarted.engine.shutdown().await;
}
#[tokio::test]
async fn state_failure_rolls_back_listener_and_topology() {
    let temp = tempfile::tempdir().unwrap();
    let bad_parent = temp.path().join("not-directory");
    std::fs::write(&bad_parent, "x").unwrap();
    let mut config = Config::default();
    config.web.state_file = bad_parent.join("state.json");
    let controller = Controller::start(config).await.unwrap();
    let team = Team {
        hidden: false,
        id: "a".into(),
        name: "A".into(),
        address: "127.0.0.1".parse().unwrap(),
    };
    assert!(controller.put_team(team).await.is_err());
    assert!(controller.topology().teams.is_empty());
    controller.engine.shutdown().await;
}
#[tokio::test]
async fn automatic_failover_is_durable_and_primary_requests_keep_working() {
    let temp = tempfile::tempdir().unwrap();
    let primary = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let port = primary.local_addr().unwrap().port();
    drop(primary);
    let replica = TcpListener::bind((std::net::Ipv6Addr::LOCALHOST, port))
        .await
        .unwrap();
    let upstream = tokio::spawn(async move {
        loop {
            let (mut s, _) = replica.accept().await.unwrap();
            tokio::spawn(async move {
                let mut data = vec![];
                s.read_to_end(&mut data).await.unwrap();
                s.write_all(&data).await.unwrap();
            });
        }
    });
    let reserve = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let listen = reserve.local_addr().unwrap();
    drop(reserve);
    let mut config = Config::default();
    config.web.state_file = temp.path().join("state.json");
    config.monitoring.cooldown_ms = 0;
    config.monitoring.timeout_ms = 100;
    config.teams = vec![
        Team {
            hidden: false,
            id: "a".into(),
            name: "A".into(),
            address: "127.0.0.1".parse().unwrap(),
        },
        Team {
            hidden: false,
            id: "b".into(),
            name: "B".into(),
            address: "::1".parse().unwrap(),
        },
    ];
    config.services = vec![ServiceConfig {
        custom: true,
        custom_ip: Some("127.0.0.1".parse().unwrap()),
        id: "notes".into(),
        name: "Notes".into(),
        listen,
        upstream_port: port,
        primary: Some("a".into()),
        replicas: vec!["b".into()],
        enabled: true,
        auto_switch: true,
        replica_policy: ReplicaPolicy::Isolate,
    }];
    let controller = Controller::start(config.clone()).await.unwrap();
    controller.check_cycle().await;
    controller.check_cycle().await;
    assert_eq!(
        controller.topology().services[0].primary.as_deref(),
        Some("b")
    );
    assert!(controller.topology().services[0].custom_ip.is_none());
    assert_eq!(controller.topology().services[0].replicas, vec!["a"]);
    let mut client = tokio::net::TcpStream::connect(listen).await.unwrap();
    client.write_all(b"hello failover").await.unwrap();
    client.shutdown().await.unwrap();
    let mut bytes = vec![];
    client.read_to_end(&mut bytes).await.unwrap();
    assert_eq!(bytes, b"hello failover");
    assert_eq!(
        mirror_tool::store::Topology::load(&config)
            .unwrap()
            .services[0]
            .primary
            .as_deref(),
        Some("b")
    );
    controller.engine.shutdown().await;
    upstream.abort();
    let _ = upstream.await;
}

fn direct_service(listen: SocketAddr, port: u16) -> ServiceConfig {
    ServiceConfig {
        id: "custom_notes".into(),
        name: "Custom notes".into(),
        listen,
        upstream_port: port,
        primary: None,
        replicas: vec![],
        enabled: true,
        auto_switch: false,
        replica_policy: ReplicaPolicy::Isolate,
        custom: true,
        custom_ip: Some("127.0.0.1".parse().unwrap()),
    }
}

#[tokio::test]
async fn custom_ip_is_atomic_reused_persisted_and_clearable() {
    let temp = tempfile::tempdir().unwrap();
    let mut cfg = Config::default();
    cfg.web.state_file = temp.path().join("state.json");
    let upstream = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let reserve = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let listen = reserve.local_addr().unwrap();
    drop(reserve);
    let controller = Controller::start(cfg.clone()).await.unwrap();
    let mut service = direct_service(listen, upstream.local_addr().unwrap().port());
    controller.put_service(service.clone()).await.unwrap();
    let state = controller.topology();
    assert_eq!(state.teams.len(), 1);
    assert!(state.teams[0].hidden);
    assert_eq!(state.teams[0].name, "127.0.0.1");
    service.replicas = vec![state.teams[0].id.clone()];
    controller.put_service(service.clone()).await.unwrap();
    assert_eq!(controller.topology().teams.len(), 1);
    assert!(controller.topology().services[0].replicas.is_empty());
    let saved = mirror_tool::store::Topology::load(&cfg).unwrap();
    assert_eq!(saved.services, controller.topology().services);
    assert_eq!(saved.teams, controller.topology().teams);
    service.primary = Some(state.teams[0].id.clone());
    service.custom_ip = None;
    service.replicas.clear();
    controller.put_service(service).await.unwrap();
    assert!(controller.topology().services[0].custom_ip.is_none());
    controller.shutdown().await;
    let mut stopped = controller.topology().services[0].clone();
    stopped.enabled = false;
    assert!(controller.put_service(stopped).await.is_err());
    assert!(
        controller
            .put_team(Team {
                id: "after".into(),
                name: "after".into(),
                address: "::1".parse().unwrap(),
                hidden: false
            })
            .await
            .is_err()
    );
}

#[tokio::test]
async fn failed_service_persistence_closes_new_listener_and_restores_existing_route() {
    let temp = tempfile::tempdir().unwrap();
    let parent = temp.path().join("state");
    std::fs::create_dir(&parent).unwrap();
    let mut cfg = Config::default();
    cfg.network.drain_timeout_ms = 100;
    cfg.web.state_file = parent.join("state.json");
    let upstream = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let reserve = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let listen = reserve.local_addr().unwrap();
    drop(reserve);
    let controller = Controller::start(cfg).await.unwrap();
    let service = direct_service(listen, upstream.local_addr().unwrap().port());
    controller.put_service(service.clone()).await.unwrap();
    let before = controller.topology();
    std::fs::remove_file(parent.join("state.json")).unwrap();
    std::fs::remove_dir(&parent).unwrap();
    std::fs::write(&parent, "blocked").unwrap();
    let reserve = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let changed_listen = reserve.local_addr().unwrap();
    drop(reserve);
    let mut changed = service.clone();
    changed.listen = changed_listen;
    assert!(controller.put_service(changed).await.is_err());
    assert_eq!(controller.topology().services, before.services);
    assert_eq!(controller.topology().teams, before.teams);
    assert!(
        tokio::net::TcpStream::connect(changed_listen)
            .await
            .is_err()
    );
    let _client = tokio::net::TcpStream::connect(listen).await.unwrap();
    tokio::time::timeout(std::time::Duration::from_secs(1), upstream.accept())
        .await
        .unwrap()
        .unwrap();
    let reserve = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let new_listen = reserve.local_addr().unwrap();
    drop(reserve);
    let mut new = service;
    new.id = "new_service".into();
    new.listen = new_listen;
    new.custom_ip = Some("::1".parse().unwrap());
    assert!(controller.put_service(new).await.is_err());
    assert!(tokio::net::TcpStream::connect(new_listen).await.is_err());
    assert_eq!(controller.topology().teams, before.teams);
    controller.shutdown().await;
}

#[test]
fn wildcard_loop_aliases_mapped_addresses_and_address_families() {
    let team = |ip: &str| Team {
        id: "a".into(),
        name: "A".into(),
        address: ip.parse().unwrap(),
        hidden: false,
    };
    let mut service = direct_service("0.0.0.0:38080".parse().unwrap(), 38080);
    service.custom_ip = None;
    service.primary = Some("a".into());
    for ip in ["127.0.0.1", "127.0.0.2", "::ffff:127.0.0.2"] {
        assert!(
            mirror_tool::config::validate_topology(&[team(ip)], &[service.clone()]).is_err(),
            "accepted self-loop {ip}"
        );
    }
    // An IPv4-only listener cannot receive native IPv6 loopback traffic.
    assert!(mirror_tool::config::validate_topology(&[team("::1")], &[service.clone()]).is_ok());
    let mut other = service.clone();
    other.id = "other".into();
    other.listen = "0.0.0.0:38081".parse().unwrap();
    other.upstream_port = 38082;
    service.upstream_port = 38081;
    assert!(
        mirror_tool::config::validate_topology(&[team("127.0.0.2")], &[service, other]).is_err()
    );
}

#[cfg(target_os = "linux")]
#[tokio::test]
async fn linux_unenumerated_loopback_alias_reaches_ipv4_wildcard() {
    let wildcard = TcpListener::bind("0.0.0.0:0").await.unwrap();
    let port = wildcard.local_addr().unwrap().port();
    let _client = tokio::net::TcpStream::connect((std::net::Ipv4Addr::new(127, 0, 0, 2), port))
        .await
        .unwrap();
    tokio::time::timeout(std::time::Duration::from_secs(1), wildcard.accept())
        .await
        .unwrap()
        .unwrap();
}

#[tokio::test]
async fn custom_ip_ids_avoid_collisions_and_reuse_operator_teams() {
    let temp = tempfile::tempdir().unwrap();
    let mut cfg = Config::default();
    cfg.web.state_file = temp.path().join("state.json");
    cfg.teams = vec![Team {
        id: "direct_127_0_0_1".into(),
        name: "Operator team".into(),
        address: "::1".parse().unwrap(),
        hidden: false,
    }];
    let controller = Controller::start(cfg.clone()).await.unwrap();
    let mut service = direct_service("127.0.0.1:38090".parse().unwrap(), 38091);
    service.enabled = false;
    controller.put_service(service.clone()).await.unwrap();
    let generated = controller.topology().services[0].primary.clone().unwrap();
    assert_eq!(generated, "direct_127_0_0_1_1");
    service.custom_ip = Some("::1".parse().unwrap());
    controller.put_service(service.clone()).await.unwrap();
    assert_eq!(
        controller.topology().services[0].primary.as_deref(),
        Some("direct_127_0_0_1")
    );
    assert!(!controller.topology().teams[0].hidden);
    service.custom_ip = Some("127.0.0.1".parse().unwrap());
    controller.put_service(service).await.unwrap();
    assert_eq!(
        controller.topology().services[0].primary.as_deref(),
        Some(generated.as_str())
    );
    controller.shutdown().await;
    let restored = Controller::start(cfg).await.unwrap();
    assert_eq!(restored.topology().teams.len(), 2);
    assert_eq!(
        restored.topology().services[0].primary.as_deref(),
        Some(generated.as_str())
    );
    assert!(restored.topology().teams[1].hidden);
    restored.shutdown().await;
}

#[tokio::test]
async fn shutdown_serializes_with_in_flight_service_transaction() {
    let temp = tempfile::tempdir().unwrap();
    let mut cfg = Config::default();
    cfg.web.state_file = temp.path().join("state.json");
    cfg.network.drain_timeout_ms = 100;
    let reserve = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let listen = reserve.local_addr().unwrap();
    drop(reserve);
    let controller = Controller::start(cfg.clone()).await.unwrap();
    let service = direct_service(listen, 39001);
    let (saved, ()) = tokio::join!(biased; controller.put_service(service), controller.shutdown());
    saved.unwrap();
    assert_eq!(
        mirror_tool::store::Topology::load(&cfg).unwrap().services,
        controller.topology().services
    );
    assert!(tokio::net::TcpStream::connect(listen).await.is_err());
    assert!(!controller.engine.snapshots()["custom_notes"].running);
}