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 {
        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 {
            id: "a".into(),
            name: "A".into(),
            address: "127.0.0.1".parse().unwrap(),
        },
        Team {
            id: "b".into(),
            name: "B".into(),
            address: "::1".parse().unwrap(),
        },
    ];
    config.services = vec![ServiceConfig {
        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")
    );
    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;
}