wgcloak/tests/supervisor.rs

#[path = "../src/config.rs"]
mod config;
#[path = "../src/datagram.rs"]
mod datagram;
#[path = "../src/protocol/mod.rs"]
mod protocol;
#[path = "../src/stats.rs"]
mod stats;
#[path = "../src/client/supervisor.rs"]
mod supervisor;
#[path = "../src/transport/mod.rs"]
mod transport;

use bytes::Bytes;
use config::{Config, TransportMode};
use std::{
    sync::{
        atomic::{AtomicUsize, Ordering},
        Arc,
    },
    time::Duration,
};
use supervisor::{ProbeFuture, TransportProbe, TransportSupervisor};
use tokio::{sync::Mutex, time::timeout};
use 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 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 config(mode: &str) -> config::ClientConfig {
    Config::from_str(&format!("[client]\nserver=\"127.0.0.1\"\nid=\"c\"\nsecret=\"{}\"\n[transport]\nmode=\"{}\"\n[transport.auto]\norder=[\"direct_udp\",\"quic\"]\nretries=2\nprobe_timeout_ms=20\n", "A".repeat(64), mode)).unwrap().into_client().unwrap()
}

#[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]);
}

#[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);
}