wgcloak/tests/control.rs

use wgcloak::control;

use std::os::unix::fs::PermissionsExt;
use tempfile::tempdir;
use tokio::io::AsyncReadExt;
use tokio::net::UnixStream;
use tokio::time::{sleep, timeout, Duration};

#[tokio::test]
async fn control_commands_and_socket_permissions() {
    let dir = tempdir().unwrap();
    let runtime_dir = dir.path().join("runtime");
    let path = runtime_dir.join("control.sock");
    let state = control::ControlState::default();
    state.enable_reconnect();
    state.set_status("ready").await;
    state.set_stats_json("{\"packets\":3}").await;
    let server = control::ControlServer::bind(&path, state.clone())
        .await
        .unwrap();
    assert_eq!(
        std::fs::metadata(&runtime_dir)
            .unwrap()
            .permissions()
            .mode()
            & 0o777,
        0o700
    );
    assert_eq!(
        std::fs::metadata(&path).unwrap().permissions().mode() & 0o777,
        0o600
    );
    let task = tokio::spawn(server.run());
    let second_bind = control::ControlServer::bind(&path, control::ControlState::default()).await;
    assert!(matches!(
        second_bind,
        Err(error) if error.kind() == std::io::ErrorKind::AddrInUse
    ));
    assert!(control::send_command(&path, "status")
        .await
        .unwrap()
        .contains("ready"));
    assert!(control::send_command(&path, "stats")
        .await
        .unwrap()
        .contains("packets"));
    assert!(control::send_command(&path, "reconnect")
        .await
        .unwrap()
        .contains("true"));
    let reload = control::send_command(&path, "reload").await.unwrap();
    assert!(reload.contains("\"ok\":false"));
    assert!(reload.contains("not supported"));
    assert!(control::send_command(&path, "shutdown")
        .await
        .unwrap()
        .contains("true"));
    task.await.unwrap().unwrap();
}

#[tokio::test]
async fn replaces_only_a_verified_stale_socket() {
    let dir = tempdir().unwrap();
    let private_dir = dir.path().join("runtime");
    std::fs::create_dir(&private_dir).unwrap();
    std::fs::set_permissions(&private_dir, std::fs::Permissions::from_mode(0o700)).unwrap();
    let path = private_dir.join("stale.sock");
    let stale = std::os::unix::net::UnixListener::bind(&path).unwrap();
    drop(stale);

    let server = control::ControlServer::bind(&path, control::ControlState::default())
        .await
        .expect("verified stale socket should be replaced");
    assert_eq!(
        std::fs::metadata(&path).unwrap().permissions().mode() & 0o777,
        0o600
    );
    drop(server);
    assert!(!path.exists());
}

#[tokio::test]
async fn refuses_to_replace_a_non_socket_path() {
    let dir = tempdir().unwrap();
    let private_dir = dir.path().join("runtime");
    std::fs::create_dir(&private_dir).unwrap();
    std::fs::set_permissions(&private_dir, std::fs::Permissions::from_mode(0o700)).unwrap();
    let path = private_dir.join("control.sock");
    std::fs::write(&path, b"keep this file").unwrap();
    let result = control::ControlServer::bind(&path, control::ControlState::default()).await;
    assert!(matches!(
        result,
        Err(error) if error.kind() == std::io::ErrorKind::AlreadyExists
    ));
    assert_eq!(std::fs::read(&path).unwrap(), b"keep this file");
}

#[test]
fn default_socket_path_never_targets_the_shared_tmp_socket() {
    let path = control::default_socket_path();
    assert_eq!(path.file_name().unwrap(), "wgcloak.sock");
    assert_ne!(path.parent().unwrap(), std::path::Path::new("/tmp"));
}

#[tokio::test]
async fn control_serializes_unusual_status_and_stats_keys_as_json() {
    let dir = tempdir().unwrap();
    let path = dir.path().join("runtime/control.sock");
    let state = control::ControlState::default();
    state.set_status("ready\"\n\t\0").await;
    let mut clients = std::collections::BTreeMap::new();
    clients.insert("client\"\n\t\0".to_owned(), 1);
    state
        .set_transport_stats_by_client_id("transport\"\n\t\0", 1, Default::default(), clients)
        .await;
    let server = control::ControlServer::bind(&path, state.clone())
        .await
        .unwrap();
    let task = tokio::spawn(server.run());

    let status: serde_json::Value =
        serde_json::from_str(&control::send_command(&path, "status").await.unwrap()).unwrap();
    assert_eq!(status["status"], "ready\"\n\t\0");
    let stats: serde_json::Value =
        serde_json::from_str(&control::send_command(&path, "stats").await.unwrap()).unwrap();
    assert_eq!(
        stats["stats"]["active_sessions_by_client_id"]["client\"\n\t\0"],
        1
    );
    assert_eq!(
        stats["stats"]["active_sessions_by_transport"]["transport\"\n\t\0"],
        1
    );

    state.shutdown.cancel();
    task.await.unwrap().unwrap();
}

#[tokio::test]
async fn idle_control_clients_are_closed_and_do_not_pin_shutdown() {
    let dir = tempdir().unwrap();
    let path = dir.path().join("runtime/control.sock");
    let state = control::ControlState::default();
    let server = control::ControlServer::bind(&path, state.clone())
        .await
        .unwrap();
    let task = tokio::spawn(server.run());
    let mut idle = UnixStream::connect(&path).await.unwrap();
    let mut response = Vec::new();
    timeout(
        control::CONTROL_REQUEST_TIMEOUT + Duration::from_secs(1),
        idle.read_to_end(&mut response),
    )
    .await
    .expect("idle control client must be closed")
    .unwrap();
    assert!(response.is_empty());

    state.shutdown.cancel();
    timeout(Duration::from_secs(1), task)
        .await
        .unwrap()
        .unwrap()
        .unwrap();
}

#[tokio::test]
async fn control_rejects_clients_beyond_the_concurrency_bound() {
    let dir = tempdir().unwrap();
    let path = dir.path().join("runtime/control.sock");
    let state = control::ControlState::default();
    let server = control::ControlServer::bind(&path, state.clone())
        .await
        .unwrap();
    let task = tokio::spawn(server.run());
    let mut held = Vec::new();
    for _ in 0..control::CONTROL_MAX_CLIENTS {
        held.push(UnixStream::connect(&path).await.unwrap());
    }
    sleep(Duration::from_millis(25)).await;
    let mut extra = UnixStream::connect(&path).await.unwrap();
    let mut response = Vec::new();
    timeout(Duration::from_secs(1), extra.read_to_end(&mut response))
        .await
        .expect("connection over the bound must be closed")
        .unwrap();
    assert!(response.is_empty());
    drop(held);
    state.shutdown.cancel();
    task.await.unwrap().unwrap();
}