wgcloak/tests/limits.rs

use std::{net::SocketAddr, sync::Arc, thread, time::Duration};

use wgcloak::{
    config::{self, LimitsConfig},
    transport::{AdmissionController, AdmissionError},
};

fn limits(
    max_sessions: usize,
    max_sessions_per_ip: usize,
    max_pending_handshakes: usize,
    max_auth_attempts_per_ip: usize,
    auth_attempt_window_secs: u64,
) -> LimitsConfig {
    LimitsConfig {
        max_sessions,
        max_sessions_per_ip,
        max_pending_handshakes,
        max_auth_attempts_per_ip,
        auth_attempt_window_secs,
        handshake_timeout_ms: 1_000,
    }
}

fn address(ip: &str) -> SocketAddr {
    format!("{ip}:40000").parse().unwrap()
}

#[test]
fn pending_limit_and_failed_handshake_release_reservation() {
    let controller = AdmissionController::new(&limits(8, 4, 1, 100, 60));
    let first = controller.begin(address("192.0.2.1")).unwrap();
    assert_eq!(controller.pending_handshakes(), 1);
    assert!(matches!(
        controller.begin(address("192.0.2.2")),
        Err(AdmissionError::Busy)
    ));

    // Dropping the pending guard models any failed or timed-out handshake.
    drop(first);
    assert_eq!(controller.pending_handshakes(), 0);
    let second = controller.begin(address("192.0.2.2")).unwrap();
    drop(second);
    assert_eq!(controller.pending_handshakes(), 0);
}

#[test]
fn active_and_per_ip_limits_release_on_close_or_drop() {
    let controller = AdmissionController::new(&limits(2, 1, 2, 100, 60));
    let first = controller.begin(address("192.0.2.1")).unwrap().promote();
    assert_eq!(controller.active_sessions(), 1);
    assert!(matches!(
        controller.begin(address("192.0.2.1")),
        Err(AdmissionError::Busy)
    ));

    let second = controller.begin(address("192.0.2.2")).unwrap().promote();
    assert_eq!(controller.active_sessions(), 2);
    assert!(matches!(
        controller.begin(address("192.0.2.3")),
        Err(AdmissionError::Busy)
    ));

    drop(first);
    assert_eq!(controller.active_sessions(), 1);
    drop(second);
    assert_eq!(controller.active_sessions(), 0);
    assert!(controller.begin(address("192.0.2.3")).is_ok());
}

#[test]
fn concurrent_admissions_never_exceed_global_or_per_ip_capacity() {
    let controller = AdmissionController::new(&limits(3, 3, 32, 100, 60));
    let barrier = Arc::new(std::sync::Barrier::new(33));
    let (sender, receiver) = std::sync::mpsc::channel();
    let mut workers = Vec::new();
    for _ in 0..32 {
        let controller = controller.clone();
        let barrier = barrier.clone();
        let sender = sender.clone();
        workers.push(thread::spawn(move || {
            barrier.wait();
            let permit = controller
                .begin(address("192.0.2.10"))
                .ok()
                .map(|pending| pending.promote());
            barrier.wait();
            sender.send(permit).unwrap();
        }));
    }
    barrier.wait();
    barrier.wait();
    drop(sender);
    let permits: Vec<_> = receiver.into_iter().flatten().collect();
    for worker in workers {
        worker.join().unwrap();
    }
    assert_eq!(permits.len(), 3);
    assert_eq!(controller.active_sessions(), 3);
}

#[tokio::test]
async fn authentication_attempt_limit_is_per_ip_and_expires() {
    let controller = AdmissionController::new(&limits(8, 4, 8, 2, 1));
    for _ in 0..2 {
        drop(controller.begin(address("192.0.2.1")).unwrap());
    }
    assert!(matches!(
        controller.begin(address("192.0.2.1")),
        Err(AdmissionError::RateLimited)
    ));
    assert!(controller.begin(address("192.0.2.2")).is_ok());

    tokio::time::sleep(Duration::from_millis(1_050)).await;
    assert!(controller.begin(address("192.0.2.1")).is_ok());
}

#[test]
fn server_limits_have_safe_defaults_and_reject_out_of_bounds_values() {
    let defaults = LimitsConfig::default();
    assert_eq!(defaults.max_sessions, 1024);
    assert_eq!(defaults.max_sessions_per_ip, 32);
    assert_eq!(defaults.max_pending_handshakes, 256);
    assert_eq!(defaults.max_auth_attempts_per_ip, 20);
    assert_eq!(defaults.auth_attempt_window_secs, 60);
    assert_eq!(defaults.handshake_timeout_ms, 3_000);

    let valid = "[server]\nlisten='127.0.0.1'\n[backend]\ntarget='127.0.0.1:51820'\n[[auth.clients]]\nid='one'\nsecret='AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA'\n";
    assert!(config::Config::from_str(valid).unwrap().validate().is_ok());

    let invalid = format!("{valid}\n[limits]\nmax_auth_attempts_per_ip=0\n");
    assert!(config::Config::from_str(&invalid)
        .unwrap()
        .validate()
        .is_err());
    let too_long = format!("{valid}\n[limits]\nhandshake_timeout_ms=300001\n");
    assert!(config::Config::from_str(&too_long)
        .unwrap()
        .validate()
        .is_err());
}