telemt/src/proxy/client_adversarial_tests.rs

use super::*;
use crate::config::ProxyConfig;
use crate::stats::Stats;
use crate::ip_tracker::UserIpTracker;
use crate::error::ProxyError;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::net::{IpAddr, Ipv4Addr, SocketAddr};

// ------------------------------------------------------------------
// Priority 3: Massive Concurrency Stress (OWASP ASVS 5.1.6)
// ------------------------------------------------------------------

#[tokio::test]
async fn client_stress_10k_connections_limit_strict() {
    let user = "stress-user";
    let limit = 512;
    
    let stats = Arc::new(Stats::new());
    let ip_tracker = Arc::new(UserIpTracker::new());
    
    let mut config = ProxyConfig::default();
    config.access.user_max_tcp_conns.insert(user.to_string(), limit);
    
    let iterations = 1000;
    let mut tasks = Vec::new();

    for i in 0..iterations {
        let stats = Arc::clone(&stats);
        let ip_tracker = Arc::clone(&ip_tracker);
        let config = config.clone();
        let user_str = user.to_string();
        
        tasks.push(tokio::spawn(async move {
            let peer = SocketAddr::new(
                IpAddr::V4(Ipv4Addr::new(127, 0, 0, (i % 254 + 1) as u8)),
                10000 + (i % 1000) as u16,
            );
            
            match RunningClientHandler::acquire_user_connection_reservation_static(
                &user_str,
                &config,
                stats,
                peer,
                ip_tracker,
            ).await {
                Ok(res) => Ok(res),
                Err(ProxyError::ConnectionLimitExceeded { .. }) => Err(()),
                Err(e) => panic!("Unexpected error: {:?}", e),
            }
        }));
    }

    let results = futures::future::join_all(tasks).await;
    let mut successes = 0;
    let mut failures = 0;
    let mut reservations = Vec::new();

    for res in results {
        match res.unwrap() {
            Ok(r) => {
                successes += 1;
                reservations.push(r);
            }
            Err(_) => failures += 1,
        }
    }

    assert_eq!(successes, limit, "Should allow exactly 'limit' connections");
    assert_eq!(failures, iterations - limit, "Should fail the rest with LimitExceeded");
    assert_eq!(stats.get_user_curr_connects(user), limit as u64);

    drop(reservations);
    
    ip_tracker.drain_cleanup_queue().await;
    
    assert_eq!(stats.get_user_curr_connects(user), 0, "Stats must converge to 0 after all drops");
    assert_eq!(ip_tracker.get_active_ip_count(user).await, 0, "IP tracker must converge to 0");
}

// ------------------------------------------------------------------
// Priority 3: IP Tracker Race Stress
// ------------------------------------------------------------------

#[tokio::test]
async fn client_ip_tracker_race_condition_stress() {
    let user = "race-user";
    let ip_tracker = Arc::new(UserIpTracker::new());
    ip_tracker.set_user_limit(user, 100).await;
    
    let iterations = 1000;
    let mut tasks = Vec::new();

    for i in 0..iterations {
        let ip_tracker = Arc::clone(&ip_tracker);
        let ip = IpAddr::V4(Ipv4Addr::new(10, 0, 0, (i % 254 + 1) as u8));
        
        tasks.push(tokio::spawn(async move {
            for _ in 0..10 {
                if let Ok(()) = ip_tracker.check_and_add("race-user", ip).await {
                    ip_tracker.remove_ip("race-user", ip).await;
                }
            }
        }));
    }

    futures::future::join_all(tasks).await;
    
    assert_eq!(ip_tracker.get_active_ip_count(user).await, 0, "IP count must be zero after balanced add/remove burst");
}

#[tokio::test]
async fn client_limit_burst_peak_never_exceeds_cap() {
    let user = "peak-cap-user";
    let limit = 32;
    let attempts = 256;

    let stats = Arc::new(Stats::new());
    let ip_tracker = Arc::new(UserIpTracker::new());

    let mut config = ProxyConfig::default();
    config.access.user_max_tcp_conns.insert(user.to_string(), limit);

    let peak = Arc::new(AtomicU64::new(0));
    let mut tasks = Vec::with_capacity(attempts);

    for i in 0..attempts {
        let stats = Arc::clone(&stats);
        let ip_tracker = Arc::clone(&ip_tracker);
        let config = config.clone();
        let peak = Arc::clone(&peak);

        tasks.push(tokio::spawn(async move {
            let peer = SocketAddr::new(
                IpAddr::V4(Ipv4Addr::new(203, 0, 113, (i % 250 + 1) as u8)),
                20000 + i as u16,
            );

            let acquired = RunningClientHandler::acquire_user_connection_reservation_static(
                user,
                &config,
                stats.clone(),
                peer,
                ip_tracker,
            )
            .await;

            if let Ok(reservation) = acquired {
                let now = stats.get_user_curr_connects(user);
                loop {
                    let prev = peak.load(Ordering::Relaxed);
                    if now <= prev {
                        break;
                    }
                    if peak
                        .compare_exchange(prev, now, Ordering::Relaxed, Ordering::Relaxed)
                        .is_ok()
                    {
                        break;
                    }
                }
                tokio::time::sleep(Duration::from_millis(2)).await;
                drop(reservation);
            }
        }));
    }

    futures::future::join_all(tasks).await;
    ip_tracker.drain_cleanup_queue().await;

    assert!(
        peak.load(Ordering::Relaxed) <= limit as u64,
        "peak concurrent reservations must not exceed configured cap"
    );
    assert_eq!(stats.get_user_curr_connects(user), 0);
    assert_eq!(ip_tracker.get_active_ip_count(user).await, 0);
}

#[tokio::test]
async fn client_quota_rejection_never_mutates_live_counters() {
    let user = "quota-reject-user";
    let stats = Arc::new(Stats::new());
    let ip_tracker = Arc::new(UserIpTracker::new());

    let mut config = ProxyConfig::default();
    config.access.user_data_quota.insert(user.to_string(), 0);

    let peer: SocketAddr = "198.51.100.201:31111".parse().unwrap();
    let res = RunningClientHandler::acquire_user_connection_reservation_static(
        user,
        &config,
        stats.clone(),
        peer,
        ip_tracker.clone(),
    )
    .await;

    assert!(matches!(res, Err(ProxyError::DataQuotaExceeded { .. })));
    assert_eq!(stats.get_user_curr_connects(user), 0);
    assert_eq!(ip_tracker.get_active_ip_count(user).await, 0);
}

#[tokio::test]
async fn client_expiration_rejection_never_mutates_live_counters() {
    let user = "expired-user";
    let stats = Arc::new(Stats::new());
    let ip_tracker = Arc::new(UserIpTracker::new());

    let mut config = ProxyConfig::default();
    config
        .access
        .user_expirations
        .insert(user.to_string(), chrono::Utc::now() - chrono::Duration::seconds(1));

    let peer: SocketAddr = "198.51.100.202:31112".parse().unwrap();
    let res = RunningClientHandler::acquire_user_connection_reservation_static(
        user,
        &config,
        stats.clone(),
        peer,
        ip_tracker.clone(),
    )
    .await;

    assert!(matches!(res, Err(ProxyError::UserExpired { .. })));
    assert_eq!(stats.get_user_curr_connects(user), 0);
    assert_eq!(ip_tracker.get_active_ip_count(user).await, 0);
}

#[tokio::test]
async fn client_ip_limit_failure_rolls_back_counter_exactly() {
    let user = "ip-limit-rollback-user";
    let stats = Arc::new(Stats::new());
    let ip_tracker = Arc::new(UserIpTracker::new());
    ip_tracker.set_user_limit(user, 1).await;

    let mut config = ProxyConfig::default();
    config.access.user_max_tcp_conns.insert(user.to_string(), 16);

    let first_peer: SocketAddr = "198.51.100.203:31113".parse().unwrap();
    let first = RunningClientHandler::acquire_user_connection_reservation_static(
        user,
        &config,
        stats.clone(),
        first_peer,
        ip_tracker.clone(),
    )
    .await
    .unwrap();

    let second_peer: SocketAddr = "198.51.100.204:31114".parse().unwrap();
    let second = RunningClientHandler::acquire_user_connection_reservation_static(
        user,
        &config,
        stats.clone(),
        second_peer,
        ip_tracker.clone(),
    )
    .await;

    assert!(matches!(second, Err(ProxyError::ConnectionLimitExceeded { .. })));
    assert_eq!(stats.get_user_curr_connects(user), 1);

    drop(first);
    ip_tracker.drain_cleanup_queue().await;

    assert_eq!(stats.get_user_curr_connects(user), 0);
    assert_eq!(ip_tracker.get_active_ip_count(user).await, 0);
}

#[tokio::test]
async fn client_parallel_limit_checks_success_path_leaves_no_residue() {
    let user = "parallel-check-success-user";
    let stats = Arc::new(Stats::new());
    let ip_tracker = Arc::new(UserIpTracker::new());
    ip_tracker.set_user_limit(user, 128).await;

    let mut config = ProxyConfig::default();
    config.access.user_max_tcp_conns.insert(user.to_string(), 128);

    let mut tasks = Vec::new();
    for i in 0..128u16 {
        let stats = Arc::clone(&stats);
        let ip_tracker = Arc::clone(&ip_tracker);
        let config = config.clone();

        tasks.push(tokio::spawn(async move {
            let peer = SocketAddr::new(
                IpAddr::V4(Ipv4Addr::new(10, 10, (i / 255) as u8, (i % 255 + 1) as u8)),
                32000 + i,
            );
            RunningClientHandler::check_user_limits_static(user, &config, &stats, peer, &ip_tracker)
                .await
        }));
    }

    for result in futures::future::join_all(tasks).await {
        assert!(result.unwrap().is_ok());
    }

    assert_eq!(stats.get_user_curr_connects(user), 0);
    assert_eq!(ip_tracker.get_active_ip_count(user).await, 0);
}

#[tokio::test]
async fn client_parallel_limit_checks_failure_path_leaves_no_residue() {
    let user = "parallel-check-failure-user";
    let stats = Arc::new(Stats::new());
    let ip_tracker = Arc::new(UserIpTracker::new());
    ip_tracker.set_user_limit(user, 0).await;

    let mut config = ProxyConfig::default();
    config.access.user_max_tcp_conns.insert(user.to_string(), 512);

    let mut tasks = Vec::new();
    for i in 0..64u16 {
        let stats = Arc::clone(&stats);
        let ip_tracker = Arc::clone(&ip_tracker);
        let config = config.clone();

        tasks.push(tokio::spawn(async move {
            let peer = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(172, 16, 0, (i % 250 + 1) as u8)), 33000 + i);
            RunningClientHandler::check_user_limits_static(user, &config, &stats, peer, &ip_tracker)
                .await
        }));
    }

    let mut _denied = 0usize;
    for result in futures::future::join_all(tasks).await {
        match result.unwrap() {
            Ok(()) => {}
            Err(ProxyError::ConnectionLimitExceeded { .. }) => _denied += 1,
            Err(other) => panic!("unexpected error: {other}"),
        }
    }

    assert_eq!(stats.get_user_curr_connects(user), 0);
    assert_eq!(ip_tracker.get_active_ip_count(user).await, 0);
}

#[tokio::test]
async fn client_churn_mixed_success_failure_converges_to_zero_state() {
    let user = "mixed-churn-user";
    let stats = Arc::new(Stats::new());
    let ip_tracker = Arc::new(UserIpTracker::new());
    ip_tracker.set_user_limit(user, 4).await;

    let mut config = ProxyConfig::default();
    config.access.user_max_tcp_conns.insert(user.to_string(), 8);

    let mut tasks = Vec::new();
    for i in 0..200u16 {
        let stats = Arc::clone(&stats);
        let ip_tracker = Arc::clone(&ip_tracker);
        let config = config.clone();

        tasks.push(tokio::spawn(async move {
            let peer = SocketAddr::new(
                IpAddr::V4(Ipv4Addr::new(192, 0, 2, (i % 16 + 1) as u8)),
                34000 + (i % 32),
            );
            let maybe_res = RunningClientHandler::acquire_user_connection_reservation_static(
                user,
                &config,
                stats,
                peer,
                ip_tracker,
            )
            .await;

            if let Ok(reservation) = maybe_res {
                tokio::time::sleep(Duration::from_millis((i % 3) as u64)).await;
                drop(reservation);
            }
        }));
    }

    futures::future::join_all(tasks).await;
    ip_tracker.drain_cleanup_queue().await;

    assert_eq!(stats.get_user_curr_connects(user), 0);
    assert_eq!(ip_tracker.get_active_ip_count(user).await, 0);
}

#[tokio::test]
async fn client_same_ip_parallel_attempts_allow_at_most_one_when_limit_is_one() {
    let user = "same-ip-user";
    let stats = Arc::new(Stats::new());
    let ip_tracker = Arc::new(UserIpTracker::new());
    ip_tracker.set_user_limit(user, 1).await;

    let mut config = ProxyConfig::default();
    config.access.user_max_tcp_conns.insert(user.to_string(), 1);

    let peer: SocketAddr = "203.0.113.44:35555".parse().unwrap();
    let mut tasks = Vec::new();

    for _ in 0..64 {
        let stats = Arc::clone(&stats);
        let ip_tracker = Arc::clone(&ip_tracker);
        let config = config.clone();
        tasks.push(tokio::spawn(async move {
            RunningClientHandler::acquire_user_connection_reservation_static(
                user,
                &config,
                stats,
                peer,
                ip_tracker,
            )
            .await
        }));
    }

    let mut granted = 0usize;
    let mut reservations = Vec::new();
    for result in futures::future::join_all(tasks).await {
        match result.unwrap() {
            Ok(reservation) => {
                granted += 1;
                reservations.push(reservation);
            }
            Err(ProxyError::ConnectionLimitExceeded { .. }) => {}
            Err(other) => panic!("unexpected error: {other}"),
        }
    }

    assert_eq!(granted, 1, "only one reservation may be granted for same IP with limit=1");
    drop(reservations);
    ip_tracker.drain_cleanup_queue().await;
    assert_eq!(stats.get_user_curr_connects(user), 0);
    assert_eq!(ip_tracker.get_active_ip_count(user).await, 0);
}

#[tokio::test]
async fn client_repeat_acquire_release_cycles_never_accumulate_state() {
    let user = "repeat-cycle-user";
    let stats = Arc::new(Stats::new());
    let ip_tracker = Arc::new(UserIpTracker::new());
    ip_tracker.set_user_limit(user, 32).await;

    let mut config = ProxyConfig::default();
    config.access.user_max_tcp_conns.insert(user.to_string(), 32);

    for i in 0..500u16 {
        let peer = SocketAddr::new(
            IpAddr::V4(Ipv4Addr::new(198, 18, (i / 250) as u8, (i % 250 + 1) as u8)),
            36000 + (i % 128),
        );
        let reservation = RunningClientHandler::acquire_user_connection_reservation_static(
            user,
            &config,
            stats.clone(),
            peer,
            ip_tracker.clone(),
        )
        .await
        .unwrap();
        drop(reservation);
    }

    ip_tracker.drain_cleanup_queue().await;
    assert_eq!(stats.get_user_curr_connects(user), 0);
    assert_eq!(ip_tracker.get_active_ip_count(user).await, 0);
}