wgcloak/tests/supervisor.rs

use bytes::Bytes;
use std::{
    sync::{
        atomic::{AtomicUsize, Ordering},
        Arc,
    },
    time::Duration,
};
use tokio::{sync::Mutex, time::timeout};
use wgcloak::{
    client::supervisor::{self as supervisor, ProbeFuture, TransportProbe, TransportSupervisor},
    config::{self, Config, TransportMode},
    protocol,
    transport::{TransportError, TransportFuture, TransportSession},
};

struct Session;
impl TransportSession for Session {
    fn send_datagram<'a>(&'a self, _: Bytes) -> TransportFuture<'a, ()> {
        Box::pin(async { Ok(()) })
    }
    fn recv_datagram<'a>(&'a self) -> TransportFuture<'a, Bytes> {
        Box::pin(async { Err(TransportError::Closed) })
    }
    fn close<'a>(&'a self) -> TransportFuture<'a, ()> {
        Box::pin(async { Ok(()) })
    }
    fn max_datagram_size(&self) -> usize {
        1200
    }
}

struct EchoSession {
    payload: Mutex<Option<Bytes>>,
}

impl TransportSession for EchoSession {
    fn send_datagram<'a>(&'a self, data: Bytes) -> TransportFuture<'a, ()> {
        Box::pin(async move {
            *self.payload.lock().await = Some(data);
            Ok(())
        })
    }

    fn recv_datagram<'a>(&'a self) -> TransportFuture<'a, Bytes> {
        Box::pin(async move {
            self.payload
                .lock()
                .await
                .take()
                .ok_or(TransportError::Closed)
        })
    }

    fn close<'a>(&'a self) -> TransportFuture<'a, ()> {
        Box::pin(async { Ok(()) })
    }

    fn max_datagram_size(&self) -> usize {
        1200
    }
}
struct Probe {
    calls: Arc<Mutex<Vec<TransportMode>>>,
    auth: bool,
}
impl TransportProbe for Probe {
    fn connect(&self, mode: TransportMode) -> ProbeFuture {
        let calls = self.calls.clone();
        let auth = self.auth;
        Box::pin(async move {
            calls.lock().await.push(mode);
            if auth {
                Err(TransportError::AuthenticationRejected(
                    protocol::RejectReason::AuthenticationFailed,
                ))
            } else {
                Err(TransportError::HandshakeTimeout)
            }
        })
    }

    fn wait_for_session_failure<'a>(
        &'a self,
        _session: &'a dyn TransportSession,
    ) -> supervisor::SessionWatchFuture<'a> {
        Box::pin(async {
            Err(TransportError::Unsupported(
                "test probe has no session lifecycle".into(),
            ))
        })
    }
}
fn config(mode: &str) -> config::ClientConfig {
    config_with_order(mode, "\"direct_udp\",\"quic\"")
}

fn config_with_order(mode: &str, order: &str) -> config::ClientConfig {
    Config::from_str(&format!("[client]\nserver=\"127.0.0.1\"\nid=\"c\"\nsecret=\"{}\"\n[transport]\nmode=\"{}\"\n[transport.auto]\norder=[{}]\nretries=2\nprobe_timeout_ms=20\n", "A".repeat(64), mode, order)).unwrap().into_client().unwrap()
}

struct SelectingProbe {
    calls: Arc<Mutex<Vec<TransportMode>>>,
    selected: TransportMode,
}

impl TransportProbe for SelectingProbe {
    fn connect(&self, mode: TransportMode) -> ProbeFuture {
        let calls = self.calls.clone();
        let selected = self.selected;
        Box::pin(async move {
            calls.lock().await.push(mode);
            if mode == selected {
                Ok(Box::new(EchoSession {
                    payload: Mutex::new(None),
                }) as Box<dyn TransportSession>)
            } else {
                Err(TransportError::HandshakeTimeout)
            }
        })
    }

    fn wait_for_session_failure<'a>(
        &'a self,
        _session: &'a dyn TransportSession,
    ) -> supervisor::SessionWatchFuture<'a> {
        Box::pin(async { Err(TransportError::Unsupported("test probe".into())) })
    }
}

async fn assert_selected_with_roundtrip(
    config: config::ClientConfig,
    selected: TransportMode,
    expected_calls: Vec<TransportMode>,
) {
    let calls = Arc::new(Mutex::new(Vec::new()));
    let supervisor = TransportSupervisor::new(
        config,
        SelectingProbe {
            calls: calls.clone(),
            selected,
        },
    );
    let session = supervisor.connect_once().await.unwrap();
    assert_eq!(*calls.lock().await, expected_calls);
    let status = supervisor.status().await;
    assert_eq!(status.state, supervisor::SupervisorState::Connected);
    assert_eq!(status.selected_transport, Some(selected));
    session
        .send_datagram(Bytes::from_static(b"auto-payload"))
        .await
        .unwrap();
    assert_eq!(
        session.recv_datagram().await.unwrap(),
        Bytes::from_static(b"auto-payload")
    );
}

#[tokio::test]
async fn auto_selects_quic_after_direct_udp_is_unavailable() {
    assert_selected_with_roundtrip(
        config_with_order("auto", "\"direct_udp\",\"quic\""),
        TransportMode::Quic,
        vec![
            TransportMode::DirectUdp,
            TransportMode::DirectUdp,
            TransportMode::Quic,
        ],
    )
    .await;
}

#[tokio::test]
async fn auto_selects_tls_tcp_after_udp_carriers_are_unavailable() {
    assert_selected_with_roundtrip(
        config_with_order("auto", "\"direct_udp\",\"quic\",\"http3\",\"tls_tcp\""),
        TransportMode::TlsTcp,
        vec![
            TransportMode::DirectUdp,
            TransportMode::DirectUdp,
            TransportMode::Quic,
            TransportMode::Quic,
            TransportMode::Http3,
            TransportMode::Http3,
            TransportMode::TlsTcp,
        ],
    )
    .await;
}

#[tokio::test]
async fn selection_order_retries_and_auth_stop() {
    let calls = Arc::new(Mutex::new(Vec::new()));
    let supervisor = TransportSupervisor::new(
        config("auto"),
        Probe {
            calls: calls.clone(),
            auth: false,
        },
    );
    assert!(matches!(
        supervisor.connect_once().await,
        Err(TransportError::HandshakeTimeout)
    ));
    assert_eq!(
        *calls.lock().await,
        vec![
            TransportMode::DirectUdp,
            TransportMode::DirectUdp,
            TransportMode::Quic,
            TransportMode::Quic
        ]
    );
    let calls = Arc::new(Mutex::new(Vec::new()));
    let supervisor = TransportSupervisor::new(
        config("auto"),
        Probe {
            calls: calls.clone(),
            auth: true,
        },
    );
    assert!(matches!(
        supervisor.connect_once().await,
        Err(TransportError::AuthenticationRejected(_))
    ));
    assert_eq!(*calls.lock().await, vec![TransportMode::DirectUdp]);
    assert_eq!(
        supervisor.status().await.state,
        supervisor::SupervisorState::AuthenticationRejected
    );
}

#[test]
fn backoff_is_bounded_and_resettable() {
    let mut supervisor = TransportSupervisor::new(
        config("direct_udp"),
        Probe {
            calls: Arc::new(Mutex::new(Vec::new())),
            auth: false,
        },
    );
    for _ in 0..30 {
        let delay = supervisor.backoff_delay();
        assert!((250..=30_000).contains(&delay.as_millis()));
    }
    supervisor.reset_backoff();
    assert_eq!(supervisor.backoff_delay(), Duration::from_millis(250));
}

struct LifecycleProbe {
    calls: Arc<AtomicUsize>,
}

impl TransportProbe for LifecycleProbe {
    fn connect(&self, _mode: TransportMode) -> ProbeFuture {
        let calls = self.calls.clone();
        Box::pin(async move {
            if calls.fetch_add(1, Ordering::SeqCst) == 0 {
                Ok(Box::new(Session) as Box<dyn TransportSession>)
            } else {
                Err(TransportError::HandshakeTimeout)
            }
        })
    }

    fn wait_for_session_failure<'a>(
        &'a self,
        _session: &'a dyn TransportSession,
    ) -> supervisor::SessionWatchFuture<'a> {
        Box::pin(async { Err(TransportError::Closed) })
    }
}

#[tokio::test]
async fn established_failure_reprobes_without_reading_datagrams() {
    let calls = Arc::new(AtomicUsize::new(0));
    let supervisor = TransportSupervisor::new(
        config("direct_udp"),
        LifecycleProbe {
            calls: calls.clone(),
        },
    );
    let shutdown = supervisor.shutdown_token();
    let task = tokio::spawn(supervisor.run());
    // The failed established session uses the normal minimum reconnect
    // backoff before the second probe.
    tokio::time::sleep(Duration::from_millis(320)).await;
    shutdown.cancel();
    let result = timeout(Duration::from_secs(2), task)
        .await
        .expect("supervisor should stop")
        .unwrap();
    assert!(result.is_ok());
    assert!(calls.load(Ordering::SeqCst) >= 2);
}

struct UnsupportedLifecycleProbe;

impl TransportProbe for UnsupportedLifecycleProbe {
    fn connect(&self, _mode: TransportMode) -> ProbeFuture {
        Box::pin(async { Ok(Box::new(Session) as Box<dyn TransportSession>) })
    }

    fn wait_for_session_failure<'a>(
        &'a self,
        _session: &'a dyn TransportSession,
    ) -> supervisor::SessionWatchFuture<'a> {
        Box::pin(async { Err(TransportError::Unsupported("lifecycle unavailable".into())) })
    }
}

#[tokio::test]
async fn unsupported_session_lifecycle_is_explicitly_terminal() {
    let supervisor = TransportSupervisor::new(config("direct_udp"), UnsupportedLifecycleProbe);
    let result = timeout(Duration::from_secs(1), supervisor.run())
        .await
        .expect("unsupported lifecycle should not hang")
        .unwrap_err();
    assert!(matches!(result, TransportError::Unsupported(_)));
}