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

#![cfg(unix)]
use mirror_tool::config::{Config, ReplicaPolicy, ServiceConfig, Team};
use std::{
    process::{Command, Stdio},
    time::Duration,
};
use tokio::{
    io::{AsyncReadExt, AsyncWriteExt},
    net::{TcpListener, TcpStream},
    time::{sleep, timeout},
};

#[tokio::test]
async fn partial_http_body_cannot_keep_process_or_proxy_listener_alive() {
    let temp = tempfile::tempdir().unwrap();
    let web = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let web_addr = web.local_addr().unwrap();
    drop(web);
    let data = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let data_addr = data.local_addr().unwrap();
    drop(data);
    let upstream = TcpListener::bind("127.0.0.1:0").await.unwrap();
    let mut cfg = Config::default();
    cfg.web.listen = web_addr;
    cfg.web.state_file = temp.path().join("state.json");
    cfg.network.drain_timeout_ms = 100;
    cfg.teams = vec![Team {
        id: "a".into(),
        name: "A".into(),
        address: "127.0.0.1".parse().unwrap(),
        hidden: false,
    }];
    cfg.services = vec![ServiceConfig {
        id: "notes".into(),
        name: "Notes".into(),
        listen: data_addr,
        upstream_port: upstream.local_addr().unwrap().port(),
        primary: Some("a".into()),
        replicas: vec![],
        enabled: true,
        auto_switch: false,
        replica_policy: ReplicaPolicy::Isolate,
        custom: false,
        custom_ip: None,
    }];
    let path = temp.path().join("mirror.toml");
    std::fs::write(&path, toml::to_string(&cfg).unwrap()).unwrap();
    let mut child = Command::new(env!("CARGO_BIN_EXE_mirror-tool"))
        .arg("--config")
        .arg(path)
        .env("MIRROR_PASSWORD", "loopback-fixture-passphrase")
        .stdout(Stdio::null())
        .stderr(Stdio::null())
        .spawn()
        .unwrap();
    let ready = timeout(Duration::from_secs(5), async {
        loop {
            if let Ok(s) = TcpStream::connect(web_addr).await {
                break s;
            }
            sleep(Duration::from_millis(10)).await;
        }
    })
    .await;
    if ready.is_err() {
        let _ = child.kill();
        panic!("server did not start");
    }
    let mut ready_http = ready.unwrap();
    ready_http
        .write_all(b"GET /api/state HTTP/1.1\r\nHost: localhost\r\nConnection: close\r\n\r\n")
        .await
        .unwrap();
    let mut reply = [0; 1024];
    assert!(
        timeout(Duration::from_secs(1), ready_http.read(&mut reply))
            .await
            .unwrap()
            .unwrap()
            > 0
    );
    let mut http = TcpStream::connect(web_addr).await.unwrap();
    http.write_all(b"POST /api/login HTTP/1.1\r\nHost: localhost\r\nContent-Type: application/json\r\nContent-Length: 1000\r\n\r\n{").await.unwrap();
    let _proxy = TcpStream::connect(data_addr).await.unwrap();
    let _upstream_peer = timeout(Duration::from_secs(1), upstream.accept())
        .await
        .unwrap()
        .unwrap();
    assert!(
        Command::new("kill")
            .args(["-TERM", &child.id().to_string()])
            .status()
            .unwrap()
            .success()
    );
    let closed = timeout(Duration::from_secs(1), async {
        while TcpStream::connect(data_addr).await.is_ok() {
            sleep(Duration::from_millis(5)).await;
        }
    })
    .await;
    if closed.is_err() {
        let _ = child.kill();
        panic!("proxy kept accepting during HTTP shutdown");
    }
    let result = timeout(Duration::from_secs(3), async {
        loop {
            if let Some(status) = child.try_wait().unwrap() {
                break status;
            }
            sleep(Duration::from_millis(10)).await;
        }
    })
    .await;
    if result.is_err() {
        let _ = child.kill();
        panic!("partial HTTP body prevented bounded process exit");
    }
    assert!(
        result.as_ref().unwrap().success(),
        "unexpected process status: {:?}",
        result
    );
    assert!(TcpStream::connect(data_addr).await.is_err());
    drop(http);
}