telemt/src/proxy/tests/client_more_advanced_tests.rs

use super::*;
use crate::config::ProxyConfig;
use crate::stats::Stats;
use crate::transport::UpstreamManager;
use std::sync::Arc;
use std::time::Duration;
use tokio::io::{AsyncReadExt, AsyncWriteExt, duplex};

fn preload_user_quota(stats: &Stats, user: &str, bytes: u64) {
    let user_stats = stats.get_or_create_user_stats_handle(user);
    stats.quota_charge_post_write(user_stats.as_ref(), bytes);
}

#[tokio::test]
async fn edge_mask_delay_bypassed_if_max_is_zero() {
    let mut config = ProxyConfig::default();
    config.censorship.server_hello_delay_min_ms = 10_000;
    config.censorship.server_hello_delay_max_ms = 0;

    let start = std::time::Instant::now();
    maybe_apply_mask_reject_delay(&config).await;
    assert!(start.elapsed() < Duration::from_millis(50));
}

#[test]
fn edge_beobachten_ttl_clamps_exactly_to_24_hours() {
    let mut config = ProxyConfig::default();
    config.general.beobachten = true;
    config.general.beobachten_minutes = 100_000;

    let ttl = beobachten_ttl(&config);
    assert_eq!(ttl.as_secs(), 24 * 60 * 60);
}

#[test]
fn edge_wrap_tls_application_record_empty_payload() {
    let wrapped = wrap_tls_application_record(&[]);
    assert_eq!(wrapped.len(), 5);
    assert_eq!(wrapped[0], TLS_RECORD_APPLICATION);
    assert_eq!(&wrapped[3..5], &[0, 0]);
}

#[tokio::test]
async fn boundary_user_data_quota_exact_match_rejects() {
    let user = "quota-boundary-user";
    let mut config = ProxyConfig::default();
    config.access.user_data_quota.insert(user.to_string(), 1024);

    let stats = Arc::new(Stats::new());
    preload_user_quota(stats.as_ref(), user, 1024);

    let ip_tracker = Arc::new(UserIpTracker::new());
    let peer = "198.51.100.10:55000".parse().unwrap();

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

    assert!(matches!(result, Err(ProxyError::DataQuotaExceeded { .. })));
}

#[tokio::test]
async fn boundary_user_expiration_in_past_rejects() {
    let user = "expired-boundary-user";
    let mut config = ProxyConfig::default();
    let expired_time = chrono::Utc::now() - chrono::Duration::milliseconds(1);
    config
        .access
        .user_expirations
        .insert(user.to_string(), expired_time);

    let stats = Arc::new(Stats::new());
    let ip_tracker = Arc::new(UserIpTracker::new());
    let peer = "198.51.100.11:55000".parse().unwrap();

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

    assert!(matches!(result, Err(ProxyError::UserExpired { .. })));
}

#[tokio::test]
async fn blackhat_proxy_protocol_massive_garbage_rejected_quickly() {
    let mut cfg = ProxyConfig::default();
    cfg.server.proxy_protocol_header_timeout_ms = 300;
    let config = Arc::new(cfg);
    let stats = Arc::new(Stats::new());

    let (server_side, mut client_side) = duplex(4096);
    let handler = tokio::spawn(handle_client_stream(
        server_side,
        "198.51.100.12:55000".parse().unwrap(),
        config,
        stats.clone(),
        Arc::new(UpstreamManager::new(
            vec![],
            1,
            1,
            1,
            10,
            1,
            false,
            stats.clone(),
        )),
        Arc::new(ReplayChecker::new(128, Duration::from_secs(60))),
        Arc::new(BufferPool::new()),
        Arc::new(SecureRandom::new()),
        None,
        Arc::new(RouteRuntimeController::new(RelayRouteMode::Direct)),
        None,
        Arc::new(UserIpTracker::new()),
        Arc::new(BeobachtenStore::new()),
        true,
    ));

    client_side.write_all(&vec![b'A'; 2000]).await.unwrap();

    let result = tokio::time::timeout(Duration::from_secs(2), handler)
        .await
        .unwrap()
        .unwrap();
    assert!(matches!(result, Err(ProxyError::InvalidProxyProtocol)));
    assert_eq!(stats.get_connects_bad(), 1);
}

#[tokio::test]
async fn edge_tls_body_immediate_eof_triggers_masking_and_bad_connect() {
    let mut cfg = ProxyConfig::default();
    cfg.general.beobachten = true;
    cfg.general.beobachten_minutes = 1;

    let config = Arc::new(cfg);
    let stats = Arc::new(Stats::new());
    let beobachten = Arc::new(BeobachtenStore::new());

    let (server_side, mut client_side) = duplex(4096);
    let handler = tokio::spawn(handle_client_stream(
        server_side,
        "198.51.100.13:55000".parse().unwrap(),
        config,
        stats.clone(),
        Arc::new(UpstreamManager::new(
            vec![],
            1,
            1,
            1,
            10,
            1,
            false,
            stats.clone(),
        )),
        Arc::new(ReplayChecker::new(128, Duration::from_secs(60))),
        Arc::new(BufferPool::new()),
        Arc::new(SecureRandom::new()),
        None,
        Arc::new(RouteRuntimeController::new(RelayRouteMode::Direct)),
        None,
        Arc::new(UserIpTracker::new()),
        beobachten.clone(),
        false,
    ));

    client_side
        .write_all(&[0x16, 0x03, 0x01, 0x00, 100])
        .await
        .unwrap();
    client_side.shutdown().await.unwrap();

    let _ = tokio::time::timeout(Duration::from_secs(2), handler)
        .await
        .unwrap();

    assert_eq!(stats.get_connects_bad(), 1);
}

#[tokio::test]
async fn security_classic_mode_disabled_masks_valid_length_payload() {
    let mut cfg = ProxyConfig::default();
    cfg.general.modes.classic = false;
    cfg.general.modes.secure = false;
    cfg.censorship.mask = true;

    let config = Arc::new(cfg);
    let stats = Arc::new(Stats::new());

    let (server_side, mut client_side) = duplex(4096);
    let handler = tokio::spawn(handle_client_stream(
        server_side,
        "198.51.100.15:55000".parse().unwrap(),
        config,
        stats.clone(),
        Arc::new(UpstreamManager::new(
            vec![],
            1,
            1,
            1,
            10,
            1,
            false,
            stats.clone(),
        )),
        Arc::new(ReplayChecker::new(128, Duration::from_secs(60))),
        Arc::new(BufferPool::new()),
        Arc::new(SecureRandom::new()),
        None,
        Arc::new(RouteRuntimeController::new(RelayRouteMode::Direct)),
        None,
        Arc::new(UserIpTracker::new()),
        Arc::new(BeobachtenStore::new()),
        false,
    ));

    client_side.write_all(&vec![0xEF; 64]).await.unwrap();
    client_side.shutdown().await.unwrap();

    let _ = tokio::time::timeout(Duration::from_secs(2), handler)
        .await
        .unwrap();
    assert_eq!(stats.get_connects_bad(), 1);
}

#[tokio::test]
async fn concurrency_ip_tracker_strict_limit_one_rapid_churn() {
    let user = "rapid-churn-user";
    let mut config = ProxyConfig::default();
    config
        .access
        .user_max_tcp_conns
        .insert(user.to_string(), 10);

    let stats = Arc::new(Stats::new());
    let ip_tracker = Arc::new(UserIpTracker::new());
    ip_tracker.set_user_limit(user, 1).await;

    let peer = "198.51.100.16:55000".parse().unwrap();

    for _ in 0..500 {
        let reservation = RunningClientHandler::acquire_user_connection_reservation_static(
            user,
            &config,
            stats.clone(),
            peer,
            ip_tracker.clone(),
        )
        .await
        .unwrap();
        reservation.release().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 quirk_read_with_progress_zero_length_buffer_returns_zero_immediately() {
    let (mut server_side, _client_side) = duplex(4096);
    let mut empty_buf = &mut [][..];

    let result = tokio::time::timeout(
        Duration::from_millis(50),
        read_with_progress(&mut server_side, &mut empty_buf),
    )
    .await;

    assert!(result.is_ok());
    assert_eq!(result.unwrap().unwrap(), 0);
}

#[tokio::test]
async fn stress_read_with_progress_cancellation_safety() {
    let (mut server_side, mut client_side) = duplex(4096);

    client_side.write_all(b"12345").await.unwrap();

    let mut buf = [0u8; 10];
    let result = tokio::time::timeout(
        Duration::from_millis(50),
        read_with_progress(&mut server_side, &mut buf),
    )
    .await;

    assert!(result.is_err());

    client_side.write_all(b"67890").await.unwrap();
    let mut buf2 = [0u8; 5];
    server_side.read_exact(&mut buf2).await.unwrap();
    assert_eq!(&buf2, b"67890");
}