wgcloak/tests/runtime_lifecycle.rs

use std::{os::unix::fs::PermissionsExt, path::Path};
use tempfile::tempdir;
use tokio::time::Duration;
use tokio::{net::UdpSocket, time::timeout};
use wgcloak::{
    config::Config,
    control, server,
    transport::{direct_udp::DirectUdpSession, TransportSession},
};

const DIRECT_SECRET: &str = "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA";

fn private_runtime_dir(root: &Path) -> std::path::PathBuf {
    let path = root.join("runtime");
    std::fs::create_dir(&path).unwrap();
    std::fs::set_permissions(&path, std::fs::Permissions::from_mode(0o700)).unwrap();
    path
}

#[tokio::test]
async fn control_commands_drive_only_supported_runtime_actions() {
    let dir = tempdir().unwrap();
    let path = private_runtime_dir(dir.path()).join("runtime.sock");
    let state = control::ControlState::default();
    state.enable_reconnect();
    let server = control::ControlServer::bind(&path, state.clone())
        .await
        .unwrap();
    let task = tokio::spawn(server.run());

    state.set_status("connected; transport=direct_udp").await;
    state
        .set_stats_json("{\"datagrams_sent\":7,\"datagrams_dropped\":1}")
        .await;
    let status = control::send_command(&path, "status").await.unwrap();
    assert!(status.contains("connected; transport=direct_udp"));
    let stats = control::send_command(&path, "stats").await.unwrap();
    assert!(stats.contains("\"datagrams_sent\":7"));

    let reconnect = state.reconnect.notified();
    tokio::pin!(reconnect);
    let response = control::send_command(&path, "reconnect").await.unwrap();
    assert!(response.contains("\"ok\":true"));
    timeout(Duration::from_secs(1), &mut reconnect)
        .await
        .expect("reconnect should reach the running client loop");

    let response = control::send_command(&path, "reload").await.unwrap();
    assert!(response.contains("\"ok\":false"));
    assert!(response.contains("not supported"));

    let response = control::send_command(&path, "shutdown").await.unwrap();
    assert!(response.contains("\"ok\":true"));
    timeout(Duration::from_secs(1), task)
        .await
        .expect("control server should stop on shutdown")
        .unwrap()
        .unwrap();
}

#[tokio::test]
async fn server_control_rejects_client_only_reconnect() {
    let dir = tempdir().unwrap();
    let path = private_runtime_dir(dir.path()).join("server.sock");
    let state = control::ControlState::default();
    let server = control::ControlServer::bind(&path, state.clone())
        .await
        .unwrap();
    let task = tokio::spawn(server.run());

    let response = control::send_command(&path, "reconnect").await.unwrap();
    assert!(response.contains("\"ok\":false"));
    assert!(response.contains("not supported"));

    state.shutdown.cancel();
    timeout(Duration::from_secs(1), task)
        .await
        .expect("control server should stop")
        .unwrap()
        .unwrap();
}

#[tokio::test]
async fn direct_server_control_stats_track_authenticated_session_lifecycle() {
    let probe = UdpSocket::bind("127.0.0.1:0").await.unwrap();
    let port = probe.local_addr().unwrap().port();
    drop(probe);

    let config = Config::from_str(&format!(
        "[server]\nlisten='127.0.0.1'\n[backend]\ntarget='127.0.0.1:9'\n[transports.direct_udp]\nport={port}\n[[auth.clients]]\nid='runtime-test'\nsecret='{DIRECT_SECRET}'\n"
    ))
    .unwrap()
    .into_server()
    .unwrap();
    let client_config = Config::from_str(&format!(
        "[client]\nserver='127.0.0.1'\nid='runtime-test'\nsecret='{DIRECT_SECRET}'\n[transport.direct_udp]\nport={port}\n"
    ))
    .unwrap()
    .into_client()
    .unwrap();

    let dir = tempdir().unwrap();
    let control_path = private_runtime_dir(dir.path()).join("direct-runtime.sock");
    let server_task = tokio::spawn(server::run_with_control_path(config, control_path.clone()));

    let session = timeout(
        Duration::from_secs(3),
        DirectUdpSession::connect(&client_config),
    )
    .await
    .expect("direct handshake should complete")
    .unwrap();

    wait_for_active_count(&control_path, 1).await;
    let status = control::send_command(&control_path, "status")
        .await
        .unwrap();
    assert!(status.contains("active_sessions=1"));
    session.close().await.unwrap();
    wait_for_active_count(&control_path, 0).await;

    let response = control::send_command(&control_path, "shutdown")
        .await
        .unwrap();
    assert!(response.contains("\"ok\":true"));
    timeout(Duration::from_secs(3), server_task)
        .await
        .expect("server runtime should stop")
        .unwrap()
        .unwrap();
}

async fn wait_for_active_count(path: &std::path::Path, expected: usize) {
    timeout(Duration::from_secs(4), async {
        loop {
            if let Ok(response) = control::send_command(path, "stats").await {
                if response.contains(&format!("\"active_sessions\":{expected}")) {
                    return;
                }
            }
            tokio::time::sleep(Duration::from_millis(20)).await;
        }
    })
    .await
    .unwrap_or_else(|_| panic!("server control stats did not report {expected} active sessions"));
}