wgcloak/tests/reconnect.rs

use std::{path::Path, time::Duration};

use tempfile::tempdir;
use tokio::{
    net::UdpSocket,
    time::{sleep, timeout},
};
use wgcloak::{
    client,
    config::{self, Config},
    control, server,
};

const SECRET: &str = "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA";

fn free_udp_port() -> u16 {
    let socket = std::net::UdpSocket::bind("127.0.0.1:0").unwrap();
    socket.local_addr().unwrap().port()
}

fn server_config(port: u16, backend: std::net::SocketAddr) -> config::ServerConfig {
    Config::from_str(&format!(
        "[server]\nlisten='127.0.0.1'\n[backend]\ntarget='{backend}'\n[transports.direct_udp]\nenabled=true\nport={port}\n[transports.quic]\nenabled=false\n[transports.http3]\nenabled=false\n[transports.tls_tcp]\nenabled=false\n[[auth.clients]]\nid='reconnect-test'\nsecret='{SECRET}'\n"
    ))
    .unwrap()
    .into_server()
    .unwrap()
}

fn client_config(port: u16, local_port: u16) -> config::ClientConfig {
    Config::from_str(&format!(
        "[client]\nserver='127.0.0.1'\nid='reconnect-test'\nsecret='{SECRET}'\n[local]\nlisten='127.0.0.1:{local_port}'\n[transport]\nmode='direct_udp'\n[transport.direct_udp]\nport={port}\n[transport.auto]\nprobe_timeout_ms=150\nretries=1\n"
    ))
    .unwrap()
    .into_client()
    .unwrap()
}

async fn wait_for_contains(path: &Path, command: &str, needle: &str) -> String {
    timeout(Duration::from_secs(10), async {
        loop {
            if let Ok(response) = control::send_command(path, command).await {
                if response.contains(needle) {
                    return response;
                }
            }
            sleep(Duration::from_millis(20)).await;
        }
    })
    .await
    .unwrap_or_else(|_| panic!("control {command} did not contain {needle:?}"))
}

async fn wait_for_reconnected_stats(path: &Path) -> String {
    timeout(Duration::from_secs(10), async {
        loop {
            if let Ok(response) = control::send_command(path, "stats").await {
                if response.contains("\"state\":\"connected\"")
                    && response.contains("\"reconnections\":1")
                {
                    return response;
                }
            }
            sleep(Duration::from_millis(20)).await;
        }
    })
    .await
    .expect("client did not publish its connected reconnect state")
}

async fn wait_for_echo(app: &UdpSocket, local_port: u16, payload: &[u8]) {
    let mut response = [0u8; 128];
    app.send_to(payload, ("127.0.0.1", local_port))
        .await
        .unwrap();
    let (len, _) = timeout(Duration::from_secs(5), app.recv_from(&mut response))
        .await
        .expect("application echo timed out")
        .unwrap();
    assert_eq!(&response[..len], payload);
}

#[tokio::test]
async fn client_recovers_through_backoff_after_server_restart() {
    let directory = tempdir().unwrap();
    let backend = UdpSocket::bind("127.0.0.1:0").await.unwrap();
    let backend_addr = backend.local_addr().unwrap();
    let backend_task = tokio::spawn(async move {
        let mut packet = vec![0u8; 2048];
        while let Ok((len, peer)) = backend.recv_from(&mut packet).await {
            let _ = backend.send_to(&packet[..len], peer).await;
        }
    });

    let port = free_udp_port();
    let local_port = free_udp_port();
    let runtime_dir = directory.path().join("runtime");
    let server_control = runtime_dir.join("server.sock");
    let client_control = runtime_dir.join("client.sock");
    let server_task = tokio::spawn(server::run_with_control_path(
        server_config(port, backend_addr),
        server_control.clone(),
    ));
    wait_for_contains(&server_control, "status", "transports=direct_udp").await;

    let client_task = tokio::spawn(client::run_with_http3_connector_and_control_path(
        client_config(port, local_port),
        client_control.clone(),
        None,
    ));
    wait_for_contains(&client_control, "status", "connected; server=127.0.0.1").await;
    let app = UdpSocket::bind("127.0.0.1:0").await.unwrap();
    wait_for_echo(&app, local_port, b"before-restart").await;

    let stopped = control::send_command(&server_control, "shutdown")
        .await
        .unwrap();
    assert!(stopped.contains("\"ok\":true"));
    timeout(Duration::from_secs(5), server_task)
        .await
        .expect("first server instance should stop")
        .unwrap()
        .unwrap();

    let requested = control::send_command(&client_control, "reconnect")
        .await
        .unwrap();
    assert!(requested.contains("\"ok\":true"));
    // Wait until the client has made a failed connection attempt and entered
    // jittered backoff before starting the replacement server. This exercises
    // the daemon loop without relying on a guessed sleep duration.
    wait_for_contains(&client_control, "status", "error=").await;

    let replacement = tokio::spawn(server::run_with_control_path(
        server_config(port, backend_addr),
        server_control.clone(),
    ));
    wait_for_contains(&server_control, "status", "transports=direct_udp").await;
    let _stats = wait_for_reconnected_stats(&client_control).await;
    wait_for_echo(&app, local_port, b"after-restart").await;

    let stopped = control::send_command(&server_control, "shutdown")
        .await
        .unwrap();
    assert!(stopped.contains("\"ok\":true"));
    timeout(Duration::from_secs(5), replacement)
        .await
        .expect("replacement server should stop")
        .unwrap()
        .unwrap();
    let stopped = control::send_command(&client_control, "shutdown")
        .await
        .unwrap();
    assert!(stopped.contains("\"ok\":true"));
    timeout(Duration::from_secs(5), client_task)
        .await
        .expect("client should stop")
        .unwrap()
        .unwrap();
    backend_task.abort();
}