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(); }