checkpoint: implement wgcloak v1 foundations

tests/supervisor.rs

@@ -0,0 +1,166 @@
+#[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);
+}