Add comprehensive tests for relay quota management and adversarial scenarios

- Introduced `relay_quota_boundary_blackhat_tests.rs` to validate behavior under quota limits, including edge cases and adversarial conditions.
- Added `relay_quota_model_adversarial_tests.rs` to ensure quota management maintains integrity during bidirectional communication and various load scenarios.
- Created `relay_quota_overflow_regression_tests.rs` to address overflow issues and ensure that quota limits are respected during aggressive data transmission.
- Implemented `route_mode_coherence_adversarial_tests.rs` to verify the consistency of route mode transitions and timestamp management across different relay modes.

src/proxy/tests/relay_quota_boundary_blackhat_tests.rs

@@ -0,0 +1,416 @@
+use super::relay_bidirectional;
+use crate::error::ProxyError;
+use crate::stats::Stats;
+use crate::stream::BufferPool;
+use rand::rngs::StdRng;
+use rand::{Rng, SeedableRng};
+use std::sync::Arc;
+use tokio::io::{duplex, AsyncRead, AsyncReadExt, AsyncWriteExt};
+use tokio::time::{timeout, Duration, Instant};
+
+async fn read_available<R: AsyncRead + Unpin>(reader: &mut R, budget: Duration) -> usize {
+    let start = Instant::now();
+    let mut total = 0usize;
+    let mut buf = [0u8; 256];
+
+    loop {
+        let elapsed = start.elapsed();
+        if elapsed >= budget {
+            break;
+        }
+        let remaining = budget.saturating_sub(elapsed);
+        match timeout(remaining, reader.read(&mut buf)).await {
+            Ok(Ok(0)) => break,
+            Ok(Ok(n)) => total = total.saturating_add(n),
+            Ok(Err(_)) | Err(_) => break,
+        }
+    }
+
+    total
+}
+
+#[tokio::test]
+async fn integration_full_duplex_exact_budget_then_hard_cutoff() {
+    let stats = Arc::new(Stats::new());
+    let user = "quota-full-duplex-boundary-user";
+
+    let (mut client_peer, relay_client) = duplex(4096);
+    let (relay_server, mut server_peer) = duplex(4096);
+    let (client_reader, client_writer) = tokio::io::split(relay_client);
+    let (server_reader, server_writer) = tokio::io::split(relay_server);
+
+    let relay = tokio::spawn(relay_bidirectional(
+        client_reader,
+        client_writer,
+        server_reader,
+        server_writer,
+        1024,
+        1024,
+        user,
+        Arc::clone(&stats),
+        Some(10),
+        Arc::new(BufferPool::new()),
+    ));
+
+    client_peer.write_all(&[0x10, 0x11, 0x12, 0x13]).await.unwrap();
+    let mut c2s = [0u8; 4];
+    server_peer.read_exact(&mut c2s).await.unwrap();
+    assert_eq!(c2s, [0x10, 0x11, 0x12, 0x13]);
+
+    server_peer
+        .write_all(&[0x20, 0x21, 0x22, 0x23, 0x24, 0x25])
+        .await
+        .unwrap();
+    let mut s2c = [0u8; 6];
+    client_peer.read_exact(&mut s2c).await.unwrap();
+    assert_eq!(s2c, [0x20, 0x21, 0x22, 0x23, 0x24, 0x25]);
+
+    let _ = client_peer.write_all(&[0x99]).await;
+    let _ = server_peer.write_all(&[0x88]).await;
+
+    let mut probe_server = [0u8; 1];
+    let mut probe_client = [0u8; 1];
+    let leaked_to_server = timeout(Duration::from_millis(120), server_peer.read(&mut probe_server)).await;
+    let leaked_to_client = timeout(Duration::from_millis(120), client_peer.read(&mut probe_client)).await;
+
+    assert!(
+        !matches!(leaked_to_server, Ok(Ok(n)) if n > 0),
+        "once quota is exhausted, no extra client byte must be forwarded"
+    );
+    assert!(
+        !matches!(leaked_to_client, Ok(Ok(n)) if n > 0),
+        "once quota is exhausted, no extra server byte must be forwarded"
+    );
+
+    let relay_result = timeout(Duration::from_secs(2), relay)
+        .await
+        .expect("relay must terminate under quota cutoff")
+        .expect("relay task must not panic");
+
+    assert!(matches!(
+        relay_result,
+        Err(ProxyError::DataQuotaExceeded { ref user }) if user == "quota-full-duplex-boundary-user"
+    ));
+    assert!(stats.get_user_total_octets(user) <= 10);
+}
+
+#[tokio::test]
+async fn negative_preloaded_quota_blocks_both_directions_immediately() {
+    let stats = Arc::new(Stats::new());
+    let user = "quota-preloaded-cutoff-user";
+    stats.add_user_octets_from(user, 5);
+
+    let (mut client_peer, relay_client) = duplex(2048);
+    let (relay_server, mut server_peer) = duplex(2048);
+    let (client_reader, client_writer) = tokio::io::split(relay_client);
+    let (server_reader, server_writer) = tokio::io::split(relay_server);
+
+    let relay = tokio::spawn(relay_bidirectional(
+        client_reader,
+        client_writer,
+        server_reader,
+        server_writer,
+        512,
+        512,
+        user,
+        Arc::clone(&stats),
+        Some(5),
+        Arc::new(BufferPool::new()),
+    ));
+
+    let _ = tokio::join!(
+        client_peer.write_all(&[0x41, 0x42]),
+        server_peer.write_all(&[0x51, 0x52]),
+    );
+
+    let leaked_to_server = read_available(&mut server_peer, Duration::from_millis(120)).await;
+    let leaked_to_client = read_available(&mut client_peer, Duration::from_millis(120)).await;
+
+    assert_eq!(leaked_to_server, 0, "preloaded limit must block C->S immediately");
+    assert_eq!(leaked_to_client, 0, "preloaded limit must block S->C immediately");
+
+    let relay_result = timeout(Duration::from_secs(2), relay)
+        .await
+        .expect("relay must terminate under preloaded cutoff")
+        .expect("relay task must not panic");
+    assert!(matches!(relay_result, Err(ProxyError::DataQuotaExceeded { .. })));
+    assert!(stats.get_user_total_octets(user) <= 5);
+}
+
+#[tokio::test]
+async fn edge_quota_one_bidirectional_race_allows_at_most_one_forwarded_octet() {
+    let stats = Arc::new(Stats::new());
+    let user = "quota-one-race-user";
+
+    let (mut client_peer, relay_client) = duplex(1024);
+    let (relay_server, mut server_peer) = duplex(1024);
+    let (client_reader, client_writer) = tokio::io::split(relay_client);
+    let (server_reader, server_writer) = tokio::io::split(relay_server);
+
+    let relay = tokio::spawn(relay_bidirectional(
+        client_reader,
+        client_writer,
+        server_reader,
+        server_writer,
+        256,
+        256,
+        user,
+        Arc::clone(&stats),
+        Some(1),
+        Arc::new(BufferPool::new()),
+    ));
+
+    let _ = tokio::join!(client_peer.write_all(&[0xAA]), server_peer.write_all(&[0xBB]));
+
+    let mut to_server = [0u8; 1];
+    let mut to_client = [0u8; 1];
+
+    let delivered_server = match timeout(Duration::from_millis(120), server_peer.read(&mut to_server)).await {
+        Ok(Ok(n)) => n,
+        _ => 0,
+    };
+    let delivered_client = match timeout(Duration::from_millis(120), client_peer.read(&mut to_client)).await {
+        Ok(Ok(n)) => n,
+        _ => 0,
+    };
+
+    assert!(
+        delivered_server + delivered_client <= 1,
+        "quota=1 must not allow >1 forwarded byte across both directions"
+    );
+
+    let relay_result = timeout(Duration::from_secs(2), relay)
+        .await
+        .expect("relay must terminate under quota=1")
+        .expect("relay task must not panic");
+    assert!(matches!(relay_result, Err(ProxyError::DataQuotaExceeded { .. })));
+    assert!(stats.get_user_total_octets(user) <= 1);
+}
+
+#[tokio::test]
+async fn adversarial_blackhat_alternating_fragmented_jitter_never_overshoots_global_quota() {
+    let stats = Arc::new(Stats::new());
+    let user = "quota-blackhat-jitter-user";
+    let quota = 32u64;
+
+    let (mut client_peer, relay_client) = duplex(4096);
+    let (relay_server, mut server_peer) = duplex(4096);
+    let (client_reader, client_writer) = tokio::io::split(relay_client);
+    let (server_reader, server_writer) = tokio::io::split(relay_server);
+
+    let relay = tokio::spawn(relay_bidirectional(
+        client_reader,
+        client_writer,
+        server_reader,
+        server_writer,
+        256,
+        256,
+        user,
+        Arc::clone(&stats),
+        Some(quota),
+        Arc::new(BufferPool::new()),
+    ));
+
+    let mut delivered_to_server = 0usize;
+    let mut delivered_to_client = 0usize;
+
+    for i in 0..256usize {
+        if relay.is_finished() {
+            break;
+        }
+
+        if (i & 1) == 0 {
+            let _ = client_peer.write_all(&[(i as u8) ^ 0x5A]).await;
+            let mut one = [0u8; 1];
+            if let Ok(Ok(n)) = timeout(Duration::from_millis(4), server_peer.read(&mut one)).await {
+                delivered_to_server = delivered_to_server.saturating_add(n);
+            }
+        } else {
+            let _ = server_peer.write_all(&[(i as u8) ^ 0xA5]).await;
+            let mut one = [0u8; 1];
+            if let Ok(Ok(n)) = timeout(Duration::from_millis(4), client_peer.read(&mut one)).await {
+                delivered_to_client = delivered_to_client.saturating_add(n);
+            }
+        }
+
+        tokio::time::sleep(Duration::from_millis(((i % 3) + 1) as u64)).await;
+    }
+
+    let relay_result = timeout(Duration::from_secs(2), relay)
+        .await
+        .expect("relay must terminate under black-hat jitter attack")
+        .expect("relay task must not panic");
+
+    assert!(matches!(relay_result, Err(ProxyError::DataQuotaExceeded { .. })));
+    assert!(
+        delivered_to_server + delivered_to_client <= quota as usize,
+        "combined forwarded bytes must never exceed configured quota"
+    );
+    assert!(stats.get_user_total_octets(user) <= quota);
+}
+
+#[tokio::test]
+async fn light_fuzz_randomized_schedule_preserves_quota_and_forwarded_byte_invariants() {
+    let mut rng = StdRng::seed_from_u64(0xD15C_A11E_F00D_BAAD);
+
+    for case in 0..48u64 {
+        let stats = Arc::new(Stats::new());
+        let user = format!("quota-fuzz-schedule-{case}");
+        let quota = rng.random_range(1u64..=32u64);
+
+        let (mut client_peer, relay_client) = duplex(4096);
+        let (relay_server, mut server_peer) = duplex(4096);
+        let (client_reader, client_writer) = tokio::io::split(relay_client);
+        let (server_reader, server_writer) = tokio::io::split(relay_server);
+
+        let relay_user = user.clone();
+        let relay_stats = Arc::clone(&stats);
+        let relay = tokio::spawn(async move {
+            relay_bidirectional(
+                client_reader,
+                client_writer,
+                server_reader,
+                server_writer,
+                256,
+                256,
+                &relay_user,
+                Arc::clone(&relay_stats),
+                Some(quota),
+                Arc::new(BufferPool::new()),
+            )
+            .await
+        });
+
+        let mut delivered_total = 0usize;
+
+        for _ in 0..96usize {
+            if relay.is_finished() {
+                break;
+            }
+
+            if rng.random::<bool>() {
+                let _ = client_peer.write_all(&[rng.random::<u8>()]).await;
+                let mut one = [0u8; 1];
+                if let Ok(Ok(n)) = timeout(Duration::from_millis(3), server_peer.read(&mut one)).await {
+                    delivered_total = delivered_total.saturating_add(n);
+                }
+            } else {
+                let _ = server_peer.write_all(&[rng.random::<u8>()]).await;
+                let mut one = [0u8; 1];
+                if let Ok(Ok(n)) = timeout(Duration::from_millis(3), client_peer.read(&mut one)).await {
+                    delivered_total = delivered_total.saturating_add(n);
+                }
+            }
+        }
+
+        drop(client_peer);
+        drop(server_peer);
+
+        let relay_result = timeout(Duration::from_secs(2), relay)
+            .await
+            .expect("fuzz relay must terminate")
+            .expect("fuzz relay task must not panic");
+
+        assert!(
+            relay_result.is_ok() || matches!(relay_result, Err(ProxyError::DataQuotaExceeded { .. })),
+            "relay must either close cleanly or terminate via typed quota error"
+        );
+        assert!(
+            delivered_total <= quota as usize,
+            "fuzz case {case}: forwarded bytes must not exceed quota"
+        );
+        assert!(
+            stats.get_user_total_octets(&user) <= quota,
+            "fuzz case {case}: accounted bytes must not exceed quota"
+        );
+    }
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
+async fn stress_multi_relay_same_user_mixed_direction_jitter_respects_global_quota() {
+    let stats = Arc::new(Stats::new());
+    let user = "quota-stress-multi-relay-user";
+    let quota = 64u64;
+
+    let mut workers = Vec::new();
+
+    for worker_id in 0..4u8 {
+        let stats = Arc::clone(&stats);
+        let user = user.to_string();
+
+        workers.push(tokio::spawn(async move {
+            let (mut client_peer, relay_client) = duplex(4096);
+            let (relay_server, mut server_peer) = duplex(4096);
+            let (client_reader, client_writer) = tokio::io::split(relay_client);
+            let (server_reader, server_writer) = tokio::io::split(relay_server);
+
+            let relay_user = user.clone();
+            let relay = tokio::spawn(async move {
+                relay_bidirectional(
+                    client_reader,
+                    client_writer,
+                    server_reader,
+                    server_writer,
+                    256,
+                    256,
+                    &relay_user,
+                    Arc::clone(&stats),
+                    Some(quota),
+                    Arc::new(BufferPool::new()),
+                )
+                .await
+            });
+
+            let mut delivered = 0usize;
+
+            for step in 0..96u8 {
+                if relay.is_finished() {
+                    break;
+                }
+
+                if ((step as usize + worker_id as usize) & 1) == 0 {
+                    let _ = client_peer.write_all(&[step ^ 0x3C]).await;
+                    let mut one = [0u8; 1];
+                    if let Ok(Ok(n)) = timeout(Duration::from_millis(3), server_peer.read(&mut one)).await {
+                        delivered = delivered.saturating_add(n);
+                    }
+                } else {
+                    let _ = server_peer.write_all(&[step ^ 0xC3]).await;
+                    let mut one = [0u8; 1];
+                    if let Ok(Ok(n)) = timeout(Duration::from_millis(3), client_peer.read(&mut one)).await {
+                        delivered = delivered.saturating_add(n);
+                    }
+                }
+
+                tokio::time::sleep(Duration::from_millis((((worker_id as u64) + (step as u64)) % 3) + 1)).await;
+            }
+
+            drop(client_peer);
+            drop(server_peer);
+            let relay_result = timeout(Duration::from_secs(2), relay)
+                .await
+                .expect("stress relay must terminate")
+                .expect("stress relay task must not panic");
+
+            assert!(
+                relay_result.is_ok() || matches!(relay_result, Err(ProxyError::DataQuotaExceeded { .. })),
+                "stress relay must either close cleanly or terminate via typed quota error"
+            );
+            delivered
+        }));
+    }
+
+    let mut delivered_sum = 0usize;
+    for worker in workers {
+        delivered_sum = delivered_sum.saturating_add(worker.await.expect("stress worker must not panic"));
+    }
+
+    assert!(
+        stats.get_user_total_octets(user) <= quota,
+        "global per-user quota must hold under concurrent mixed-direction relay stress"
+    );
+    assert!(
+        delivered_sum <= quota as usize,
+        "combined delivered bytes across relays must stay within global quota"
+    );
+}