diff --git a/src/proxy/tests/middle_relay_blackhat_campaign_integration_tests.rs b/src/proxy/tests/middle_relay_blackhat_campaign_integration_tests.rs deleted file mode 100644 index 6f0e91a..0000000 --- a/src/proxy/tests/middle_relay_blackhat_campaign_integration_tests.rs +++ /dev/null @@ -1,113 +0,0 @@ -use super::*; -use crate::stats::Stats; -use dashmap::DashMap; -use std::sync::Arc; -use std::sync::atomic::{AtomicUsize, Ordering}; -use tokio::sync::Barrier; -use tokio::time::{Duration, timeout}; - -#[tokio::test(flavor = "multi_thread", worker_threads = 4)] -async fn blackhat_campaign_saturation_quota_race_with_queue_pressure_stays_fail_closed() { - let _guard = super::quota_user_lock_test_scope(); - let _pressure_guard = super::relay_idle_pressure_test_scope(); - let map = QUOTA_USER_LOCKS.get_or_init(DashMap::new); - map.clear(); - - let mut retained = Vec::with_capacity(QUOTA_USER_LOCKS_MAX); - for idx in 0..QUOTA_USER_LOCKS_MAX { - retained.push(quota_user_lock(&format!( - "middle-blackhat-held-{}-{idx}", - std::process::id() - ))); - } - - assert_eq!( - map.len(), - QUOTA_USER_LOCKS_MAX, - "precondition: bounded lock cache must be saturated" - ); - - let (tx, _rx) = mpsc::channel::(1); - tx.send(C2MeCommand::Close) - .await - .expect("queue prefill should succeed"); - - let pressure_seq_before = relay_pressure_event_seq(); - let pressure_errors = Arc::new(AtomicUsize::new(0)); - let mut pressure_workers = Vec::new(); - for _ in 0..16 { - let tx = tx.clone(); - let pressure_errors = Arc::clone(&pressure_errors); - pressure_workers.push(tokio::spawn(async move { - if enqueue_c2me_command(&tx, C2MeCommand::Close).await.is_err() { - pressure_errors.fetch_add(1, Ordering::Relaxed); - } - })); - } - - let stats = Arc::new(Stats::new()); - let user = format!("middle-blackhat-quota-race-{}", std::process::id()); - let gate = Arc::new(Barrier::new(16)); - - let mut quota_workers = Vec::new(); - for _ in 0..16u8 { - let stats = Arc::clone(&stats); - let user = user.clone(); - let gate = Arc::clone(&gate); - quota_workers.push(tokio::spawn(async move { - gate.wait().await; - let user_lock = quota_user_lock(&user); - let _quota_guard = user_lock.lock().await; - - if quota_would_be_exceeded_for_user(&stats, &user, Some(1), 1) { - return false; - } - stats.add_user_octets_to(&user, 1); - true - })); - } - - let mut ok_count = 0usize; - let mut denied_count = 0usize; - for worker in quota_workers { - let result = timeout(Duration::from_secs(2), worker) - .await - .expect("quota worker must finish") - .expect("quota worker must not panic"); - if result { - ok_count += 1; - } else { - denied_count += 1; - } - } - - for worker in pressure_workers { - timeout(Duration::from_secs(2), worker) - .await - .expect("pressure worker must finish") - .expect("pressure worker must not panic"); - } - - assert_eq!( - stats.get_user_total_octets(&user), - 1, - "black-hat campaign must not overshoot same-user quota under saturation" - ); - assert!(ok_count <= 1, "at most one quota contender may succeed"); - assert!( - denied_count >= 15, - "all remaining contenders must be quota-denied" - ); - - let pressure_seq_after = relay_pressure_event_seq(); - assert!( - pressure_seq_after > pressure_seq_before, - "queue pressure leg must trigger pressure accounting" - ); - assert!( - pressure_errors.load(Ordering::Relaxed) >= 1, - "at least one pressure worker should fail from persistent backpressure" - ); - - drop(retained); -} diff --git a/src/proxy/tests/middle_relay_coverage_high_risk_security_tests.rs b/src/proxy/tests/middle_relay_coverage_high_risk_security_tests.rs deleted file mode 100644 index 44c201f..0000000 --- a/src/proxy/tests/middle_relay_coverage_high_risk_security_tests.rs +++ /dev/null @@ -1,777 +0,0 @@ -use super::*; -use crate::crypto::AesCtr; -use crate::crypto::SecureRandom; -use crate::stats::Stats; -use crate::stream::{BufferPool, PooledBuffer}; -use std::sync::Arc; -use tokio::io::AsyncReadExt; -use tokio::io::duplex; -use tokio::sync::mpsc; -use tokio::time::{Duration as TokioDuration, timeout}; - -fn make_pooled_payload(data: &[u8]) -> PooledBuffer { - let pool = Arc::new(BufferPool::with_config(data.len().max(1), 4)); - let mut payload = pool.get(); - payload.resize(data.len(), 0); - payload[..data.len()].copy_from_slice(data); - payload -} - -#[tokio::test] -async fn write_client_payload_abridged_short_quickack_sets_flag_and_preserves_payload() { - let (mut read_side, write_side) = duplex(4096); - let key = [0u8; 32]; - let iv = 0u128; - - let mut writer = CryptoWriter::new(write_side, AesCtr::new(&key, iv), 8 * 1024); - let mut decryptor = AesCtr::new(&key, iv); - let rng = SecureRandom::new(); - let mut frame_buf = Vec::new(); - let payload = vec![0xA1, 0xB2, 0xC3, 0xD4, 0x10, 0x20, 0x30, 0x40]; - - write_client_payload( - &mut writer, - ProtoTag::Abridged, - RPC_FLAG_QUICKACK, - &payload, - &rng, - &mut frame_buf, - ) - .await - .expect("abridged quickack payload should serialize"); - writer.flush().await.expect("flush must succeed"); - - let mut encrypted = vec![0u8; 1 + payload.len()]; - read_side - .read_exact(&mut encrypted) - .await - .expect("must read serialized abridged frame"); - let plaintext = decryptor.decrypt(&encrypted); - - assert_eq!(plaintext[0], 0x80 | ((payload.len() / 4) as u8)); - assert_eq!(&plaintext[1..], payload.as_slice()); -} - -#[tokio::test] -async fn write_client_payload_abridged_extended_header_is_encoded_correctly() { - let (mut read_side, write_side) = duplex(16 * 1024); - let key = [0u8; 32]; - let iv = 0u128; - - let mut writer = CryptoWriter::new(write_side, AesCtr::new(&key, iv), 8 * 1024); - let mut decryptor = AesCtr::new(&key, iv); - let rng = SecureRandom::new(); - let mut frame_buf = Vec::new(); - - // Boundary where abridged switches to extended length encoding. - let payload = vec![0x5Au8; 0x7f * 4]; - - write_client_payload( - &mut writer, - ProtoTag::Abridged, - RPC_FLAG_QUICKACK, - &payload, - &rng, - &mut frame_buf, - ) - .await - .expect("extended abridged payload should serialize"); - writer.flush().await.expect("flush must succeed"); - - let mut encrypted = vec![0u8; 4 + payload.len()]; - read_side - .read_exact(&mut encrypted) - .await - .expect("must read serialized extended abridged frame"); - let plaintext = decryptor.decrypt(&encrypted); - - assert_eq!(plaintext[0], 0xff, "0x7f with quickack bit must be set"); - assert_eq!(&plaintext[1..4], &[0x7f, 0x00, 0x00]); - assert_eq!(&plaintext[4..], payload.as_slice()); -} - -#[tokio::test] -async fn write_client_payload_abridged_misaligned_is_rejected_fail_closed() { - let (_read_side, write_side) = duplex(1024); - let key = [0u8; 32]; - let iv = 0u128; - - let mut writer = CryptoWriter::new(write_side, AesCtr::new(&key, iv), 8 * 1024); - let rng = SecureRandom::new(); - let mut frame_buf = Vec::new(); - - let err = write_client_payload( - &mut writer, - ProtoTag::Abridged, - 0, - &[1, 2, 3], - &rng, - &mut frame_buf, - ) - .await - .expect_err("misaligned abridged payload must be rejected"); - - let msg = format!("{err}"); - assert!( - msg.contains("4-byte aligned"), - "error should explain alignment contract, got: {msg}" - ); -} - -#[tokio::test] -async fn write_client_payload_secure_misaligned_is_rejected_fail_closed() { - let (_read_side, write_side) = duplex(1024); - let key = [0u8; 32]; - let iv = 0u128; - - let mut writer = CryptoWriter::new(write_side, AesCtr::new(&key, iv), 8 * 1024); - let rng = SecureRandom::new(); - let mut frame_buf = Vec::new(); - - let err = write_client_payload( - &mut writer, - ProtoTag::Secure, - 0, - &[9, 8, 7, 6, 5], - &rng, - &mut frame_buf, - ) - .await - .expect_err("misaligned secure payload must be rejected"); - - let msg = format!("{err}"); - assert!( - msg.contains("Secure payload must be 4-byte aligned"), - "error should be explicit for fail-closed triage, got: {msg}" - ); -} - -#[tokio::test] -async fn write_client_payload_intermediate_quickack_sets_length_msb() { - let (mut read_side, write_side) = duplex(4096); - let key = [0u8; 32]; - let iv = 0u128; - - let mut writer = CryptoWriter::new(write_side, AesCtr::new(&key, iv), 8 * 1024); - let mut decryptor = AesCtr::new(&key, iv); - let rng = SecureRandom::new(); - let mut frame_buf = Vec::new(); - let payload = b"hello-middle-relay"; - - write_client_payload( - &mut writer, - ProtoTag::Intermediate, - RPC_FLAG_QUICKACK, - payload, - &rng, - &mut frame_buf, - ) - .await - .expect("intermediate quickack payload should serialize"); - writer.flush().await.expect("flush must succeed"); - - let mut encrypted = vec![0u8; 4 + payload.len()]; - read_side - .read_exact(&mut encrypted) - .await - .expect("must read intermediate frame"); - let plaintext = decryptor.decrypt(&encrypted); - - let mut len_bytes = [0u8; 4]; - len_bytes.copy_from_slice(&plaintext[..4]); - let len_with_flags = u32::from_le_bytes(len_bytes); - assert_ne!(len_with_flags & 0x8000_0000, 0, "quickack bit must be set"); - assert_eq!((len_with_flags & 0x7fff_ffff) as usize, payload.len()); - assert_eq!(&plaintext[4..], payload); -} - -#[tokio::test] -async fn write_client_payload_secure_quickack_prefix_and_padding_bounds_hold() { - let (mut read_side, write_side) = duplex(4096); - let key = [0u8; 32]; - let iv = 0u128; - - let mut writer = CryptoWriter::new(write_side, AesCtr::new(&key, iv), 8 * 1024); - let mut decryptor = AesCtr::new(&key, iv); - let rng = SecureRandom::new(); - let mut frame_buf = Vec::new(); - let payload = vec![0x33u8; 100]; // 4-byte aligned as required by secure mode. - - write_client_payload( - &mut writer, - ProtoTag::Secure, - RPC_FLAG_QUICKACK, - &payload, - &rng, - &mut frame_buf, - ) - .await - .expect("secure quickack payload should serialize"); - writer.flush().await.expect("flush must succeed"); - - // Secure mode adds 1..=3 bytes of randomized tail padding. - let mut encrypted_header = [0u8; 4]; - read_side - .read_exact(&mut encrypted_header) - .await - .expect("must read secure header"); - let decrypted_header = decryptor.decrypt(&encrypted_header); - let header: [u8; 4] = decrypted_header - .try_into() - .expect("decrypted secure header must be 4 bytes"); - let wire_len_raw = u32::from_le_bytes(header); - - assert_ne!( - wire_len_raw & 0x8000_0000, - 0, - "secure quickack bit must be set" - ); - - let wire_len = (wire_len_raw & 0x7fff_ffff) as usize; - assert!(wire_len >= payload.len()); - let padding_len = wire_len - payload.len(); - assert!( - (1..=3).contains(&padding_len), - "secure writer must add bounded random tail padding, got {padding_len}" - ); - - let mut encrypted_body = vec![0u8; wire_len]; - read_side - .read_exact(&mut encrypted_body) - .await - .expect("must read secure body"); - let decrypted_body = decryptor.decrypt(&encrypted_body); - assert_eq!(&decrypted_body[..payload.len()], payload.as_slice()); -} - -#[tokio::test] -#[ignore = "heavy: allocates >64MiB to validate abridged too-large fail-closed branch"] -async fn write_client_payload_abridged_too_large_is_rejected_fail_closed() { - let (_read_side, write_side) = duplex(1024); - let key = [0u8; 32]; - let iv = 0u128; - - let mut writer = CryptoWriter::new(write_side, AesCtr::new(&key, iv), 8 * 1024); - let rng = SecureRandom::new(); - let mut frame_buf = Vec::new(); - - // Exactly one 4-byte word above the encodable 24-bit abridged length range. - let payload = vec![0x00u8; (1 << 24) * 4]; - let err = write_client_payload( - &mut writer, - ProtoTag::Abridged, - 0, - &payload, - &rng, - &mut frame_buf, - ) - .await - .expect_err("oversized abridged payload must be rejected"); - - let msg = format!("{err}"); - assert!( - msg.contains("Abridged frame too large"), - "error must clearly indicate oversize fail-close path, got: {msg}" - ); -} - -#[tokio::test] -async fn write_client_ack_intermediate_is_little_endian() { - let (mut read_side, write_side) = duplex(1024); - let key = [0u8; 32]; - let iv = 0u128; - let mut writer = CryptoWriter::new(write_side, AesCtr::new(&key, iv), 8 * 1024); - let mut decryptor = AesCtr::new(&key, iv); - - write_client_ack(&mut writer, ProtoTag::Intermediate, 0x11_22_33_44) - .await - .expect("ack serialization should succeed"); - writer.flush().await.expect("flush must succeed"); - - let mut encrypted = [0u8; 4]; - read_side - .read_exact(&mut encrypted) - .await - .expect("must read ack bytes"); - let plain = decryptor.decrypt(&encrypted); - assert_eq!(plain.as_slice(), &0x11_22_33_44u32.to_le_bytes()); -} - -#[tokio::test] -async fn write_client_ack_abridged_is_big_endian() { - let (mut read_side, write_side) = duplex(1024); - let key = [0u8; 32]; - let iv = 0u128; - let mut writer = CryptoWriter::new(write_side, AesCtr::new(&key, iv), 8 * 1024); - let mut decryptor = AesCtr::new(&key, iv); - - write_client_ack(&mut writer, ProtoTag::Abridged, 0xDE_AD_BE_EF) - .await - .expect("ack serialization should succeed"); - writer.flush().await.expect("flush must succeed"); - - let mut encrypted = [0u8; 4]; - read_side - .read_exact(&mut encrypted) - .await - .expect("must read ack bytes"); - let plain = decryptor.decrypt(&encrypted); - assert_eq!(plain.as_slice(), &0xDE_AD_BE_EFu32.to_be_bytes()); -} - -#[tokio::test] -async fn write_client_payload_abridged_short_boundary_0x7e_is_single_byte_header() { - let (mut read_side, write_side) = duplex(1024 * 1024); - let key = [0u8; 32]; - let iv = 0u128; - let mut writer = CryptoWriter::new(write_side, AesCtr::new(&key, iv), 8 * 1024); - let mut decryptor = AesCtr::new(&key, iv); - let rng = SecureRandom::new(); - let mut frame_buf = Vec::new(); - let payload = vec![0xABu8; 0x7e * 4]; - - write_client_payload( - &mut writer, - ProtoTag::Abridged, - 0, - &payload, - &rng, - &mut frame_buf, - ) - .await - .expect("boundary payload should serialize"); - writer.flush().await.expect("flush must succeed"); - - let mut encrypted = vec![0u8; 1 + payload.len()]; - read_side.read_exact(&mut encrypted).await.unwrap(); - let plain = decryptor.decrypt(&encrypted); - assert_eq!(plain[0], 0x7e); - assert_eq!(&plain[1..], payload.as_slice()); -} - -#[tokio::test] -async fn write_client_payload_abridged_extended_without_quickack_has_clean_prefix() { - let (mut read_side, write_side) = duplex(16 * 1024); - let key = [0u8; 32]; - let iv = 0u128; - let mut writer = CryptoWriter::new(write_side, AesCtr::new(&key, iv), 8 * 1024); - let mut decryptor = AesCtr::new(&key, iv); - let rng = SecureRandom::new(); - let mut frame_buf = Vec::new(); - let payload = vec![0x42u8; 0x80 * 4]; - - write_client_payload( - &mut writer, - ProtoTag::Abridged, - 0, - &payload, - &rng, - &mut frame_buf, - ) - .await - .expect("extended payload should serialize"); - writer.flush().await.expect("flush must succeed"); - - let mut encrypted = vec![0u8; 4 + payload.len()]; - read_side.read_exact(&mut encrypted).await.unwrap(); - let plain = decryptor.decrypt(&encrypted); - assert_eq!(plain[0], 0x7f); - assert_eq!(&plain[1..4], &[0x80, 0x00, 0x00]); - assert_eq!(&plain[4..], payload.as_slice()); -} - -#[tokio::test] -async fn write_client_payload_intermediate_zero_length_emits_header_only() { - let (mut read_side, write_side) = duplex(1024); - let key = [0u8; 32]; - let iv = 0u128; - let mut writer = CryptoWriter::new(write_side, AesCtr::new(&key, iv), 8 * 1024); - let mut decryptor = AesCtr::new(&key, iv); - let rng = SecureRandom::new(); - let mut frame_buf = Vec::new(); - - write_client_payload( - &mut writer, - ProtoTag::Intermediate, - 0, - &[], - &rng, - &mut frame_buf, - ) - .await - .expect("zero-length intermediate payload should serialize"); - writer.flush().await.expect("flush must succeed"); - - let mut encrypted = [0u8; 4]; - read_side.read_exact(&mut encrypted).await.unwrap(); - let plain = decryptor.decrypt(&encrypted); - assert_eq!(plain.as_slice(), &[0, 0, 0, 0]); -} - -#[tokio::test] -async fn write_client_payload_intermediate_ignores_unrelated_flags() { - let (mut read_side, write_side) = duplex(1024); - let key = [0u8; 32]; - let iv = 0u128; - let mut writer = CryptoWriter::new(write_side, AesCtr::new(&key, iv), 8 * 1024); - let mut decryptor = AesCtr::new(&key, iv); - let rng = SecureRandom::new(); - let mut frame_buf = Vec::new(); - let payload = [7u8; 12]; - - write_client_payload( - &mut writer, - ProtoTag::Intermediate, - 0x4000_0000, - &payload, - &rng, - &mut frame_buf, - ) - .await - .expect("payload should serialize"); - writer.flush().await.expect("flush must succeed"); - - let mut encrypted = [0u8; 16]; - read_side.read_exact(&mut encrypted).await.unwrap(); - let plain = decryptor.decrypt(&encrypted); - let len = u32::from_le_bytes(plain[0..4].try_into().unwrap()); - assert_eq!(len, payload.len() as u32, "only quickack bit may affect header"); - assert_eq!(&plain[4..], payload.as_slice()); -} - -#[tokio::test] -async fn write_client_payload_secure_without_quickack_keeps_msb_clear() { - let (mut read_side, write_side) = duplex(4096); - let key = [0u8; 32]; - let iv = 0u128; - let mut writer = CryptoWriter::new(write_side, AesCtr::new(&key, iv), 8 * 1024); - let mut decryptor = AesCtr::new(&key, iv); - let rng = SecureRandom::new(); - let mut frame_buf = Vec::new(); - let payload = [0x1Du8; 64]; - - write_client_payload( - &mut writer, - ProtoTag::Secure, - 0, - &payload, - &rng, - &mut frame_buf, - ) - .await - .expect("payload should serialize"); - writer.flush().await.expect("flush must succeed"); - - let mut encrypted_header = [0u8; 4]; - read_side.read_exact(&mut encrypted_header).await.unwrap(); - let plain_header = decryptor.decrypt(&encrypted_header); - let h: [u8; 4] = plain_header.as_slice().try_into().unwrap(); - let wire_len_raw = u32::from_le_bytes(h); - assert_eq!(wire_len_raw & 0x8000_0000, 0, "quickack bit must stay clear"); -} - -#[tokio::test] -async fn secure_padding_light_fuzz_distribution_has_multiple_outcomes() { - let (mut read_side, write_side) = duplex(256 * 1024); - let key = [0u8; 32]; - let iv = 0u128; - let mut writer = CryptoWriter::new(write_side, AesCtr::new(&key, iv), 8 * 1024); - let mut decryptor = AesCtr::new(&key, iv); - let rng = SecureRandom::new(); - let mut frame_buf = Vec::new(); - let payload = [0x55u8; 100]; - let mut seen = [false; 4]; - - for _ in 0..96 { - write_client_payload( - &mut writer, - ProtoTag::Secure, - 0, - &payload, - &rng, - &mut frame_buf, - ) - .await - .expect("secure payload should serialize"); - writer.flush().await.expect("flush must succeed"); - - let mut encrypted_header = [0u8; 4]; - read_side.read_exact(&mut encrypted_header).await.unwrap(); - let plain_header = decryptor.decrypt(&encrypted_header); - let h: [u8; 4] = plain_header.as_slice().try_into().unwrap(); - let wire_len = (u32::from_le_bytes(h) & 0x7fff_ffff) as usize; - let padding_len = wire_len - payload.len(); - assert!((1..=3).contains(&padding_len)); - seen[padding_len] = true; - - let mut encrypted_body = vec![0u8; wire_len]; - read_side.read_exact(&mut encrypted_body).await.unwrap(); - let _ = decryptor.decrypt(&encrypted_body); - } - - let distinct = (1..=3).filter(|idx| seen[*idx]).count(); - assert!( - distinct >= 2, - "padding generator should not collapse to a single outcome under campaign" - ); -} - -#[tokio::test] -async fn write_client_payload_mixed_proto_sequence_preserves_stream_sync() { - let (mut read_side, write_side) = duplex(128 * 1024); - let key = [0u8; 32]; - let iv = 0u128; - let mut writer = CryptoWriter::new(write_side, AesCtr::new(&key, iv), 8 * 1024); - let mut decryptor = AesCtr::new(&key, iv); - let rng = SecureRandom::new(); - let mut frame_buf = Vec::new(); - - let p1 = vec![1u8; 8]; - let p2 = vec![2u8; 16]; - let p3 = vec![3u8; 20]; - - write_client_payload(&mut writer, ProtoTag::Abridged, 0, &p1, &rng, &mut frame_buf) - .await - .unwrap(); - write_client_payload( - &mut writer, - ProtoTag::Intermediate, - RPC_FLAG_QUICKACK, - &p2, - &rng, - &mut frame_buf, - ) - .await - .unwrap(); - write_client_payload(&mut writer, ProtoTag::Secure, 0, &p3, &rng, &mut frame_buf) - .await - .unwrap(); - writer.flush().await.unwrap(); - - // Frame 1: abridged short. - let mut e1 = vec![0u8; 1 + p1.len()]; - read_side.read_exact(&mut e1).await.unwrap(); - let d1 = decryptor.decrypt(&e1); - assert_eq!(d1[0], (p1.len() / 4) as u8); - assert_eq!(&d1[1..], p1.as_slice()); - - // Frame 2: intermediate with quickack. - let mut e2 = vec![0u8; 4 + p2.len()]; - read_side.read_exact(&mut e2).await.unwrap(); - let d2 = decryptor.decrypt(&e2); - let l2 = u32::from_le_bytes(d2[0..4].try_into().unwrap()); - assert_ne!(l2 & 0x8000_0000, 0); - assert_eq!((l2 & 0x7fff_ffff) as usize, p2.len()); - assert_eq!(&d2[4..], p2.as_slice()); - - // Frame 3: secure with bounded tail. - let mut e3h = [0u8; 4]; - read_side.read_exact(&mut e3h).await.unwrap(); - let d3h = decryptor.decrypt(&e3h); - let l3 = (u32::from_le_bytes(d3h.as_slice().try_into().unwrap()) & 0x7fff_ffff) as usize; - assert!(l3 >= p3.len()); - assert!((1..=3).contains(&(l3 - p3.len()))); - let mut e3b = vec![0u8; l3]; - read_side.read_exact(&mut e3b).await.unwrap(); - let d3b = decryptor.decrypt(&e3b); - assert_eq!(&d3b[..p3.len()], p3.as_slice()); -} - -#[test] -fn should_yield_sender_boundary_matrix_blackhat() { - assert!(!should_yield_c2me_sender(0, false)); - assert!(!should_yield_c2me_sender(0, true)); - assert!(!should_yield_c2me_sender(C2ME_SENDER_FAIRNESS_BUDGET - 1, true)); - assert!(!should_yield_c2me_sender(C2ME_SENDER_FAIRNESS_BUDGET, false)); - assert!(should_yield_c2me_sender(C2ME_SENDER_FAIRNESS_BUDGET, true)); - assert!(should_yield_c2me_sender( - C2ME_SENDER_FAIRNESS_BUDGET.saturating_add(1024), - true - )); -} - -#[test] -fn should_yield_sender_light_fuzz_matches_oracle() { - let mut s: u64 = 0xD00D_BAAD_F00D_CAFE; - for _ in 0..5000 { - s ^= s << 7; - s ^= s >> 9; - s ^= s << 8; - let sent = (s as usize) & 0x1fff; - let backlog = (s & 1) != 0; - - let expected = backlog && sent >= C2ME_SENDER_FAIRNESS_BUDGET; - assert_eq!(should_yield_c2me_sender(sent, backlog), expected); - } -} - -#[test] -fn quota_would_be_exceeded_exact_remaining_one_byte() { - let stats = Stats::new(); - let user = "quota-edge"; - let quota = 100u64; - stats.add_user_octets_to(user, 99); - - assert!( - !quota_would_be_exceeded_for_user(&stats, user, Some(quota), 1), - "exactly remaining budget should be allowed" - ); - assert!( - quota_would_be_exceeded_for_user(&stats, user, Some(quota), 2), - "one byte beyond remaining budget must be rejected" - ); -} - -#[test] -fn quota_would_be_exceeded_saturating_edge_remains_fail_closed() { - let stats = Stats::new(); - let user = "quota-saturating-edge"; - let quota = u64::MAX - 3; - stats.add_user_octets_to(user, u64::MAX - 4); - - assert!( - quota_would_be_exceeded_for_user(&stats, user, Some(quota), 2), - "saturating arithmetic edge must stay fail-closed" - ); -} - -#[test] -fn quota_exceeded_boundary_is_inclusive() { - let stats = Stats::new(); - let user = "quota-inclusive-boundary"; - stats.add_user_octets_to(user, 50); - - assert!(quota_exceeded_for_user(&stats, user, Some(50))); - assert!(!quota_exceeded_for_user(&stats, user, Some(51))); -} - -#[test] -fn quota_soft_helper_matches_capped_generic_helper_matrix() { - let stats = Stats::new(); - let user = "quota-soft-parity"; - - for used in [0u64, 1, 7, 63, 127, 255] { - stats.sub_user_octets_to(user, stats.get_user_total_octets(user)); - stats.add_user_octets_to(user, used); - - for quota in [8u64, 64, 128, 256] { - for overshoot in [0u64, 1, 5, 32] { - for bytes in [0u64, 1, 2, 7, 31, 64] { - let soft = quota_would_be_exceeded_for_user_soft( - &stats, - user, - Some(quota), - bytes, - overshoot, - ); - let capped = quota_would_be_exceeded_for_user( - &stats, - user, - Some(quota_soft_cap(quota, overshoot)), - bytes, - ); - assert_eq!( - soft, capped, - "soft helper parity mismatch: used={used} quota={quota} overshoot={overshoot} bytes={bytes}" - ); - } - } - } - } -} - -#[test] -fn quota_soft_helper_none_limit_never_rejects() { - let stats = Stats::new(); - let user = "quota-soft-none"; - stats.add_user_octets_to(user, u64::MAX); - - assert!(!quota_would_be_exceeded_for_user_soft( - &stats, - user, - None, - u64::MAX, - u64::MAX, - )); -} - -#[test] -fn quota_soft_cap_saturates_and_stays_fail_closed() { - let stats = Stats::new(); - let user = "quota-soft-saturating"; - let quota = u64::MAX - 2; - let overshoot = 100; - - assert_eq!(quota_soft_cap(quota, overshoot), u64::MAX); - - stats.add_user_octets_to(user, u64::MAX - 1); - assert!(quota_would_be_exceeded_for_user_soft( - &stats, - user, - Some(quota), - 2, - overshoot, - )); -} - -#[tokio::test] -async fn enqueue_c2me_close_fast_path_succeeds_without_backpressure() { - let (tx, mut rx) = mpsc::channel::(4); - enqueue_c2me_command(&tx, C2MeCommand::Close) - .await - .expect("close should enqueue on fast path"); - - let recv = timeout(TokioDuration::from_millis(50), rx.recv()) - .await - .expect("must receive close command") - .expect("close command should be present"); - assert!(matches!(recv, C2MeCommand::Close)); -} - -#[tokio::test] -async fn enqueue_c2me_data_full_then_drain_preserves_order() { - let (tx, mut rx) = mpsc::channel::(1); - tx.send(C2MeCommand::Data { - payload: make_pooled_payload(&[1]), - flags: 10, - }) - .await - .unwrap(); - - let tx2 = tx.clone(); - let producer = tokio::spawn(async move { - enqueue_c2me_command( - &tx2, - C2MeCommand::Data { - payload: make_pooled_payload(&[2, 2]), - flags: 20, - }, - ) - .await - }); - - tokio::time::sleep(TokioDuration::from_millis(10)).await; - - let first = rx.recv().await.expect("first item should exist"); - match first { - C2MeCommand::Data { payload, flags } => { - assert_eq!(payload.as_ref(), &[1]); - assert_eq!(flags, 10); - } - C2MeCommand::Close => panic!("unexpected close as first item"), - } - - producer.await.unwrap().expect("producer should complete"); - - let second = timeout(TokioDuration::from_millis(100), rx.recv()) - .await - .unwrap() - .expect("second item should exist"); - match second { - C2MeCommand::Data { payload, flags } => { - assert_eq!(payload.as_ref(), &[2, 2]); - assert_eq!(flags, 20); - } - C2MeCommand::Close => panic!("unexpected close as second item"), - } -} diff --git a/src/proxy/tests/middle_relay_cross_mode_lock_release_regression_security_tests.rs b/src/proxy/tests/middle_relay_cross_mode_lock_release_regression_security_tests.rs deleted file mode 100644 index a787aa6..0000000 --- a/src/proxy/tests/middle_relay_cross_mode_lock_release_regression_security_tests.rs +++ /dev/null @@ -1,295 +0,0 @@ -use super::*; -use crate::crypto::{AesCtr, SecureRandom}; -use crate::stats::Stats; -use crate::stream::CryptoWriter; -use bytes::Bytes; -use std::pin::Pin; -use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; -use std::sync::{Arc, Mutex}; -use std::task::{Context, Poll, Waker}; -use tokio::io::AsyncWrite; -use tokio::sync::Notify; -use tokio::task::JoinSet; -use tokio::time::{Duration, timeout}; - -fn make_crypto_writer(writer: W) -> CryptoWriter -where - W: tokio::io::AsyncWrite + Unpin, -{ - let key = [0u8; 32]; - let iv = 0u128; - CryptoWriter::new(writer, AesCtr::new(&key, iv), 8 * 1024) -} - -#[derive(Default)] -struct BlockingWriteState { - write_entered: AtomicBool, - released: AtomicBool, - write_waker: Mutex>, - write_entered_notify: Notify, -} - -struct BlockingWrite { - state: Arc, -} - -impl BlockingWrite { - fn new(state: Arc) -> Self { - Self { state } - } -} - -impl AsyncWrite for BlockingWrite { - fn poll_write( - self: Pin<&mut Self>, - cx: &mut Context<'_>, - buf: &[u8], - ) -> Poll> { - self.state.write_entered.store(true, Ordering::Release); - self.state.write_entered_notify.notify_waiters(); - - if self.state.released.load(Ordering::Acquire) { - return Poll::Ready(Ok(buf.len())); - } - - if let Ok(mut slot) = self.state.write_waker.lock() { - *slot = Some(cx.waker().clone()); - } - - Poll::Pending - } - - fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { - Poll::Ready(Ok(())) - } - - fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { - Poll::Ready(Ok(())) - } -} - -async fn wait_until_blocking_write_entered(state: &Arc) { - for _ in 0..8 { - if state.write_entered.load(Ordering::Acquire) { - return; - } - let _ = timeout(Duration::from_millis(25), state.write_entered_notify.notified()).await; - } - - panic!("blocking writer did not enter poll_write in bounded time"); -} - -fn release_blocking_write(state: &Arc) { - state.released.store(true, Ordering::Release); - if let Ok(mut slot) = state.write_waker.lock() - && let Some(waker) = slot.take() - { - waker.wake(); - } -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn adversarial_blocked_write_releases_cross_mode_lock_and_preserves_fail_closed_quota() { - let stats = Arc::new(Stats::new()); - let user = format!("middle-cross-release-regression-{}", std::process::id()); - let cross_mode_lock = Arc::new(cross_mode_quota_user_lock_for_tests(&user)); - let bytes_me2c = Arc::new(AtomicU64::new(0)); - let writer_state = Arc::new(BlockingWriteState::default()); - - let first = { - let stats = Arc::clone(&stats); - let user = user.clone(); - let cross_mode_lock = Arc::clone(&cross_mode_lock); - let bytes_me2c = Arc::clone(&bytes_me2c); - let writer_state = Arc::clone(&writer_state); - tokio::spawn(async move { - let mut writer = make_crypto_writer(BlockingWrite::new(writer_state)); - let mut frame_buf = Vec::new(); - process_me_writer_response_with_cross_mode_lock( - MeResponse::Data { - flags: 0, - data: Bytes::from_static(&[0xAA, 0xBB, 0xCC, 0xDD]), - }, - &mut writer, - ProtoTag::Intermediate, - &SecureRandom::new(), - &mut frame_buf, - stats.as_ref(), - &user, - Some(4), - 0, - Some(&cross_mode_lock), - bytes_me2c.as_ref(), - 41_000, - false, - false, - ) - .await - }) - }; - - wait_until_blocking_write_entered(&writer_state).await; - - let guard = timeout(Duration::from_millis(40), cross_mode_lock.lock()) - .await - .expect("cross-mode lock must be released while first write is pending"); - drop(guard); - - let second = { - let stats = Arc::clone(&stats); - let user = user.clone(); - let cross_mode_lock = Arc::clone(&cross_mode_lock); - let bytes_me2c = Arc::clone(&bytes_me2c); - tokio::spawn(async move { - let mut writer = make_crypto_writer(tokio::io::sink()); - let mut frame_buf = Vec::new(); - timeout( - Duration::from_millis(150), - process_me_writer_response_with_cross_mode_lock( - MeResponse::Data { - flags: 0, - data: Bytes::from_static(&[0xEE]), - }, - &mut writer, - ProtoTag::Intermediate, - &SecureRandom::new(), - &mut frame_buf, - stats.as_ref(), - &user, - Some(4), - 0, - Some(&cross_mode_lock), - bytes_me2c.as_ref(), - 41_001, - false, - false, - ), - ) - .await - }) - }; - - let second_result = second - .await - .expect("second task must not panic") - .expect("second write must not block on cross-mode lock"); - assert!( - matches!(second_result, Err(ProxyError::DataQuotaExceeded { .. })), - "second write must fail closed due to first write reservation" - ); - - release_blocking_write(&writer_state); - - let first_result = timeout(Duration::from_millis(300), first) - .await - .expect("first task timed out") - .expect("first task must not panic"); - assert!(first_result.is_ok()); - - assert_eq!(stats.get_user_total_octets(&user), 4); - assert_eq!(bytes_me2c.load(Ordering::Relaxed), 4); -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 4)] -async fn stress_pending_write_does_not_starve_same_user_waiters_after_quota_boundary() { - let stats = Arc::new(Stats::new()); - let user = format!("middle-cross-release-stress-{}", std::process::id()); - let cross_mode_lock = Arc::new(cross_mode_quota_user_lock_for_tests(&user)); - let bytes_me2c = Arc::new(AtomicU64::new(0)); - let writer_state = Arc::new(BlockingWriteState::default()); - - let first = { - let stats = Arc::clone(&stats); - let user = user.clone(); - let cross_mode_lock = Arc::clone(&cross_mode_lock); - let bytes_me2c = Arc::clone(&bytes_me2c); - let writer_state = Arc::clone(&writer_state); - tokio::spawn(async move { - let mut writer = make_crypto_writer(BlockingWrite::new(writer_state)); - let mut frame_buf = Vec::new(); - process_me_writer_response_with_cross_mode_lock( - MeResponse::Data { - flags: 0, - data: Bytes::from_static(&[0x01, 0x02]), - }, - &mut writer, - ProtoTag::Intermediate, - &SecureRandom::new(), - &mut frame_buf, - stats.as_ref(), - &user, - Some(3), - 0, - Some(&cross_mode_lock), - bytes_me2c.as_ref(), - 41_100, - false, - false, - ) - .await - }) - }; - - wait_until_blocking_write_entered(&writer_state).await; - - let mut set = JoinSet::new(); - for idx in 0..48u64 { - let stats = Arc::clone(&stats); - let user = user.clone(); - let cross_mode_lock = Arc::clone(&cross_mode_lock); - let bytes_me2c = Arc::clone(&bytes_me2c); - set.spawn(async move { - let mut writer = make_crypto_writer(tokio::io::sink()); - let mut frame_buf = Vec::new(); - timeout( - Duration::from_millis(200), - process_me_writer_response_with_cross_mode_lock( - MeResponse::Data { - flags: 0, - data: Bytes::from_static(&[0x10]), - }, - &mut writer, - ProtoTag::Intermediate, - &SecureRandom::new(), - &mut frame_buf, - stats.as_ref(), - &user, - Some(3), - 0, - Some(&cross_mode_lock), - bytes_me2c.as_ref(), - 41_200 + idx, - false, - false, - ), - ) - .await - }); - } - - let mut ok = 0usize; - let mut quota_exceeded = 0usize; - while let Some(done) = set.join_next().await { - let timed = done.expect("waiter task must not panic"); - let result = timed.expect("waiter must not block behind pending first write"); - match result { - Ok(_) => ok += 1, - Err(ProxyError::DataQuotaExceeded { .. }) => quota_exceeded += 1, - Err(other) => panic!("unexpected error in waiter: {other:?}"), - } - } - - assert_eq!(ok, 1, "exactly one waiter should consume remaining one-byte quota"); - assert_eq!(quota_exceeded, 47); - - release_blocking_write(&writer_state); - - let first_result = timeout(Duration::from_millis(300), first) - .await - .expect("first task timed out") - .expect("first task must not panic"); - assert!(first_result.is_ok()); - - assert_eq!(stats.get_user_total_octets(&user), 3); - assert_eq!(bytes_me2c.load(Ordering::Relaxed), 3); -} diff --git a/src/proxy/tests/middle_relay_cross_mode_lookup_efficiency_security_tests.rs b/src/proxy/tests/middle_relay_cross_mode_lookup_efficiency_security_tests.rs deleted file mode 100644 index 37e1b87..0000000 --- a/src/proxy/tests/middle_relay_cross_mode_lookup_efficiency_security_tests.rs +++ /dev/null @@ -1,116 +0,0 @@ -use super::*; -use crate::crypto::{AesCtr, SecureRandom}; -use crate::stats::Stats; -use crate::stream::CryptoWriter; -use bytes::Bytes; -use std::sync::Arc; -use std::sync::atomic::{AtomicU64, Ordering}; -use std::sync::{Mutex, OnceLock}; - -fn make_crypto_writer(writer: W) -> CryptoWriter -where - W: tokio::io::AsyncWrite + Unpin, -{ - let key = [0u8; 32]; - let iv = 0u128; - CryptoWriter::new(writer, AesCtr::new(&key, iv), 8 * 1024) -} - -fn lookup_counter_test_lock() -> &'static Mutex<()> { - static LOCK: OnceLock> = OnceLock::new(); - LOCK.get_or_init(|| Mutex::new(())) -} - -#[tokio::test] -async fn tdd_prefetched_cross_mode_lock_avoids_per_frame_registry_lookup_in_me_to_client_writer() { - let _guard = lookup_counter_test_lock() - .lock() - .unwrap_or_else(|poison| poison.into_inner()); - - let stats = Stats::new(); - let user = format!("middle-cross-mode-lookup-{}", std::process::id()); - let cross_mode_lock = crate::proxy::quota_lock_registry::cross_mode_quota_user_lock(&user); - - crate::proxy::quota_lock_registry::reset_cross_mode_quota_user_lock_lookup_count_for_tests(); - - let mut writer = make_crypto_writer(tokio::io::sink()); - let mut frame_buf = Vec::new(); - let bytes_me2c = AtomicU64::new(0); - - for idx in 0..8u64 { - let outcome = process_me_writer_response_with_cross_mode_lock( - MeResponse::Data { - flags: 0, - data: Bytes::from_static(&[0xAB]), - }, - &mut writer, - ProtoTag::Intermediate, - &SecureRandom::new(), - &mut frame_buf, - &stats, - &user, - Some(1024), - 0, - Some(&cross_mode_lock), - &bytes_me2c, - 20_000 + idx, - false, - false, - ) - .await; - - assert!(outcome.is_ok()); - } - - assert_eq!( - crate::proxy::quota_lock_registry::cross_mode_quota_user_lock_lookup_count_for_user_for_tests(&user), - 0, - "prefetched lock path must not re-query lock registry per frame" - ); - assert_eq!(stats.get_user_total_octets(&user), 8); - assert_eq!(bytes_me2c.load(Ordering::Relaxed), 8); -} - -#[tokio::test] -async fn control_without_prefetched_lock_still_uses_registry_lookup_path() { - let _guard = lookup_counter_test_lock() - .lock() - .unwrap_or_else(|poison| poison.into_inner()); - - let stats = Stats::new(); - let user = format!("middle-cross-mode-lookup-control-{}", std::process::id()); - - crate::proxy::quota_lock_registry::reset_cross_mode_quota_user_lock_lookup_count_for_tests(); - - let mut writer = make_crypto_writer(tokio::io::sink()); - let mut frame_buf = Vec::new(); - let bytes_me2c = AtomicU64::new(0); - - let outcome = process_me_writer_response_with_cross_mode_lock( - MeResponse::Data { - flags: 0, - data: Bytes::from_static(&[0xCD]), - }, - &mut writer, - ProtoTag::Intermediate, - &SecureRandom::new(), - &mut frame_buf, - &stats, - &user, - Some(1024), - 0, - None, - &bytes_me2c, - 20_100, - false, - false, - ) - .await; - - assert!(outcome.is_ok()); - assert_eq!( - crate::proxy::quota_lock_registry::cross_mode_quota_user_lock_lookup_count_for_user_for_tests(&user), - 1, - "fallback path without prefetched lock should perform a registry lookup" - ); -} diff --git a/src/proxy/tests/middle_relay_cross_mode_quota_lock_matrix_security_tests.rs b/src/proxy/tests/middle_relay_cross_mode_quota_lock_matrix_security_tests.rs deleted file mode 100644 index bc7c857..0000000 --- a/src/proxy/tests/middle_relay_cross_mode_quota_lock_matrix_security_tests.rs +++ /dev/null @@ -1,376 +0,0 @@ -use super::*; -use crate::crypto::{AesCtr, SecureRandom}; -use crate::stats::Stats; -use crate::stream::CryptoWriter; -use bytes::Bytes; -use std::sync::Arc; -use std::sync::atomic::{AtomicU64, Ordering}; -use tokio::time::{Duration, timeout}; - -fn make_crypto_writer(writer: W) -> CryptoWriter -where - W: tokio::io::AsyncWrite + Unpin, -{ - let key = [0u8; 32]; - let iv = 0u128; - CryptoWriter::new(writer, AesCtr::new(&key, iv), 8 * 1024) -} - -#[tokio::test] -async fn positive_quota_limited_me_to_client_write_updates_counters_exactly_once() { - let stats = Stats::new(); - let user = format!("middle-cross-matrix-positive-{}", std::process::id()); - let mut writer = make_crypto_writer(tokio::io::sink()); - let mut frame_buf = Vec::new(); - let bytes_me2c = AtomicU64::new(0); - - let result = process_me_writer_response( - MeResponse::Data { - flags: 0, - data: Bytes::from_static(&[1, 2, 3, 4]), - }, - &mut writer, - ProtoTag::Intermediate, - &SecureRandom::new(), - &mut frame_buf, - &stats, - &user, - Some(128), - 0, - &bytes_me2c, - 10_001, - false, - false, - ) - .await; - - assert!(result.is_ok()); - assert_eq!(stats.get_user_total_octets(&user), 4); - assert_eq!(bytes_me2c.load(Ordering::Relaxed), 4); -} - -#[tokio::test] -async fn negative_held_cross_mode_lock_blocks_quota_limited_me_to_client_path() { - let stats = Stats::new(); - let user = format!("middle-cross-matrix-negative-{}", std::process::id()); - let held = cross_mode_quota_user_lock_for_tests(&user); - let held_guard = held - .try_lock() - .expect("test must hold lock before ME->C call"); - - let mut writer = make_crypto_writer(tokio::io::sink()); - let mut frame_buf = Vec::new(); - let bytes_me2c = AtomicU64::new(0); - - let blocked = timeout( - Duration::from_millis(25), - process_me_writer_response( - MeResponse::Data { - flags: 0, - data: Bytes::from_static(&[0x41]), - }, - &mut writer, - ProtoTag::Intermediate, - &SecureRandom::new(), - &mut frame_buf, - &stats, - &user, - Some(256), - 0, - &bytes_me2c, - 10_002, - false, - false, - ), - ) - .await; - - assert!(blocked.is_err()); - drop(held_guard); -} - -#[tokio::test] -async fn edge_quota_none_bypasses_cross_mode_lock_guard_in_me_to_client_path() { - let stats = Stats::new(); - let user = format!("middle-cross-matrix-edge-none-{}", std::process::id()); - let held = cross_mode_quota_user_lock_for_tests(&user); - let held_guard = held - .try_lock() - .expect("test must hold lock while quota is disabled"); - - let mut writer = make_crypto_writer(tokio::io::sink()); - let mut frame_buf = Vec::new(); - let bytes_me2c = AtomicU64::new(0); - - let outcome = timeout( - Duration::from_millis(80), - process_me_writer_response( - MeResponse::Data { - flags: 0, - data: Bytes::from_static(&[0x11, 0x22]), - }, - &mut writer, - ProtoTag::Intermediate, - &SecureRandom::new(), - &mut frame_buf, - &stats, - &user, - None, - 0, - &bytes_me2c, - 10_003, - false, - false, - ), - ) - .await - .expect("quota-none path must not wait on cross-mode lock"); - - assert!(outcome.is_ok()); - drop(held_guard); -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 4)] -async fn adversarial_same_user_parallel_quota_limited_writes_stay_hard_capped() { - let stats = Arc::new(Stats::new()); - let user = format!("middle-cross-matrix-adversarial-{}", std::process::id()); - let limit = 64u64; - let bytes_me2c = Arc::new(AtomicU64::new(0)); - let mut tasks = Vec::new(); - - for idx in 0..256u64 { - let stats = Arc::clone(&stats); - let bytes_me2c = Arc::clone(&bytes_me2c); - let user = user.clone(); - tasks.push(tokio::spawn(async move { - let mut writer = make_crypto_writer(tokio::io::sink()); - let mut frame_buf = Vec::new(); - process_me_writer_response( - MeResponse::Data { - flags: 0, - data: Bytes::from_static(&[0xEE]), - }, - &mut writer, - ProtoTag::Intermediate, - &SecureRandom::new(), - &mut frame_buf, - stats.as_ref(), - &user, - Some(limit), - 0, - bytes_me2c.as_ref(), - 11_000 + idx, - false, - false, - ) - .await - })); - } - - let mut ok = 0usize; - for task in tasks { - match task.await.expect("task must not panic") { - Ok(_) => ok += 1, - Err(ProxyError::DataQuotaExceeded { .. }) => {} - Err(other) => panic!("unexpected error in adversarial parallel case: {other:?}"), - } - } - - assert_eq!(ok, limit as usize); - assert_eq!(stats.get_user_total_octets(&user), limit); - assert_eq!(bytes_me2c.load(Ordering::Relaxed), limit); -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 4)] -async fn integration_shared_lock_blocks_direct_relay_and_middle_relay_for_same_user() { - let user = format!("middle-cross-matrix-integration-{}", std::process::id()); - let relay_lock = crate::proxy::relay::cross_mode_quota_user_lock_for_tests(&user); - let middle_lock = cross_mode_quota_user_lock_for_tests(&user); - assert!( - Arc::ptr_eq(&relay_lock, &middle_lock), - "relay and middle-relay must share the same cross-mode lock identity" - ); - - let held_guard = relay_lock - .try_lock() - .expect("test must hold shared cross-mode lock"); - - let stats = Stats::new(); - let mut writer = make_crypto_writer(tokio::io::sink()); - let mut frame_buf = Vec::new(); - let bytes_me2c = AtomicU64::new(0); - - let middle_blocked = timeout( - Duration::from_millis(25), - process_me_writer_response( - MeResponse::Data { - flags: 0, - data: Bytes::from_static(&[0x92]), - }, - &mut writer, - ProtoTag::Intermediate, - &SecureRandom::new(), - &mut frame_buf, - &stats, - &user, - Some(1024), - 0, - &bytes_me2c, - 12_001, - false, - false, - ), - ) - .await; - assert!(middle_blocked.is_err()); - - drop(held_guard); - - let middle_ready = timeout( - Duration::from_millis(250), - process_me_writer_response( - MeResponse::Data { - flags: 0, - data: Bytes::from_static(&[0x94]), - }, - &mut writer, - ProtoTag::Intermediate, - &SecureRandom::new(), - &mut frame_buf, - &stats, - &user, - Some(1024), - 0, - &bytes_me2c, - 12_002, - false, - false, - ), - ) - .await - .expect("middle path must complete after release"); - - assert!(middle_ready.is_ok()); -} - -#[tokio::test] -async fn light_fuzz_mixed_payload_sizes_with_periodic_lock_holds_keeps_accounting_consistent() { - let stats = Stats::new(); - let user = format!("middle-cross-matrix-fuzz-{}", std::process::id()); - let bytes_me2c = AtomicU64::new(0); - let mut seed = 0xC0DE_1234_55AA_9988u64; - - for case in 0..96u32 { - seed ^= seed << 7; - seed ^= seed >> 9; - seed ^= seed << 8; - - let hold = (seed & 0x03) == 0; - let mut held_lock = None; - let maybe_guard = if hold { - held_lock = Some(cross_mode_quota_user_lock_for_tests(&user)); - Some( - held_lock - .as_ref() - .expect("held lock should be present") - .try_lock() - .expect("cross-mode lock should be acquirable in fuzz round"), - ) - } else { - None - }; - - let payload_len = ((seed >> 8) as usize % 8) + 1; - let payload = vec![(seed & 0xff) as u8; payload_len]; - let before = stats.get_user_total_octets(&user); - let mut writer = make_crypto_writer(tokio::io::sink()); - let mut frame_buf = Vec::new(); - - let timed = timeout( - Duration::from_millis(20), - process_me_writer_response( - MeResponse::Data { - flags: 0, - data: Bytes::from(payload), - }, - &mut writer, - ProtoTag::Intermediate, - &SecureRandom::new(), - &mut frame_buf, - &stats, - &user, - Some(1024), - 0, - &bytes_me2c, - 13_000 + case as u64, - false, - false, - ), - ) - .await; - - if hold { - assert!(timed.is_err(), "held-lock fuzz round must block within timeout"); - assert_eq!(stats.get_user_total_octets(&user), before); - } else { - let done = timed.expect("unheld fuzz round must complete in time"); - assert!(done.is_ok()); - } - - drop(maybe_guard); - drop(held_lock); - assert_eq!(bytes_me2c.load(Ordering::Relaxed), stats.get_user_total_octets(&user)); - } -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 4)] -async fn stress_held_user_lock_does_not_block_other_users_me_to_client_writes() { - let held_user = format!("middle-cross-matrix-stress-held-{}", std::process::id()); - let free_user = format!("middle-cross-matrix-stress-free-{}", std::process::id()); - - let held = cross_mode_quota_user_lock_for_tests(&held_user); - let held_guard = held - .try_lock() - .expect("test must hold lock for blocked user"); - - let mut tasks = Vec::new(); - for idx in 0..64u64 { - let user = free_user.clone(); - tasks.push(tokio::spawn(async move { - let stats = Stats::new(); - let mut writer = make_crypto_writer(tokio::io::sink()); - let mut frame_buf = Vec::new(); - let bytes_me2c = AtomicU64::new(0); - process_me_writer_response( - MeResponse::Data { - flags: 0, - data: Bytes::from_static(&[0xA0]), - }, - &mut writer, - ProtoTag::Intermediate, - &SecureRandom::new(), - &mut frame_buf, - &stats, - &user, - Some(1), - 0, - &bytes_me2c, - 14_000 + idx, - false, - false, - ) - .await - })); - } - - timeout(Duration::from_secs(2), async { - for task in tasks { - let done = task.await.expect("free-user task must not panic"); - assert!(done.is_ok()); - } - }) - .await - .expect("free-user tasks should complete without waiting for held user's lock"); - - drop(held_guard); -} diff --git a/src/proxy/tests/middle_relay_cross_mode_quota_reservation_security_tests.rs b/src/proxy/tests/middle_relay_cross_mode_quota_reservation_security_tests.rs deleted file mode 100644 index 51092bd..0000000 --- a/src/proxy/tests/middle_relay_cross_mode_quota_reservation_security_tests.rs +++ /dev/null @@ -1,254 +0,0 @@ -use super::*; -use crate::crypto::{AesCtr, SecureRandom}; -use crate::stats::Stats; -use crate::stream::CryptoWriter; -use bytes::Bytes; -use std::pin::Pin; -use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; -use std::sync::{Arc, Mutex}; -use std::task::{Context, Poll, Waker}; -use tokio::io::AsyncWrite; -use tokio::sync::Notify; -use tokio::time::{Duration, timeout}; - -fn make_crypto_writer(writer: W) -> CryptoWriter -where - W: tokio::io::AsyncWrite + Unpin, -{ - let key = [0u8; 32]; - let iv = 0u128; - CryptoWriter::new(writer, AesCtr::new(&key, iv), 8 * 1024) -} - -#[derive(Default)] -struct BlockingWriteState { - write_entered: AtomicBool, - released: AtomicBool, - write_waker: Mutex>, - write_entered_notify: Notify, -} - -struct BlockingWrite { - state: Arc, -} - -impl BlockingWrite { - fn new(state: Arc) -> Self { - Self { state } - } -} - -impl AsyncWrite for BlockingWrite { - fn poll_write( - self: Pin<&mut Self>, - cx: &mut Context<'_>, - buf: &[u8], - ) -> Poll> { - self.state.write_entered.store(true, Ordering::Release); - self.state.write_entered_notify.notify_waiters(); - - if self.state.released.load(Ordering::Acquire) { - return Poll::Ready(Ok(buf.len())); - } - - if let Ok(mut slot) = self.state.write_waker.lock() { - *slot = Some(cx.waker().clone()); - } - Poll::Pending - } - - fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { - Poll::Ready(Ok(())) - } - - fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { - Poll::Ready(Ok(())) - } -} - -async fn wait_until_blocking_write_entered(state: &Arc) { - for _ in 0..8 { - if state.write_entered.load(Ordering::Acquire) { - return; - } - let _ = timeout(Duration::from_millis(25), state.write_entered_notify.notified()).await; - } - - panic!("blocking writer did not enter poll_write in bounded time"); -} - -fn release_blocking_write(state: &Arc) { - state.released.store(true, Ordering::Release); - if let Ok(mut slot) = state.write_waker.lock() - && let Some(waker) = slot.take() - { - waker.wake(); - } -} - -#[tokio::test] -async fn adversarial_held_cross_mode_lock_blocks_me_to_client_quota_reservation_path() { - let stats = Stats::new(); - let user = format!("middle-me2c-cross-mode-held-{}", std::process::id()); - let held = cross_mode_quota_user_lock_for_tests(&user); - let held_guard = held - .try_lock() - .expect("test must hold shared cross-mode lock before ME->C write path"); - - let mut writer = make_crypto_writer(tokio::io::sink()); - let mut frame_buf = Vec::new(); - let bytes_me2c = AtomicU64::new(0); - - let blocked = timeout( - Duration::from_millis(25), - process_me_writer_response( - MeResponse::Data { - flags: 0, - data: Bytes::from_static(&[0x41]), - }, - &mut writer, - ProtoTag::Intermediate, - &SecureRandom::new(), - &mut frame_buf, - &stats, - &user, - Some(1024), - 0, - &bytes_me2c, - 9901, - false, - false, - ), - ) - .await; - - assert!( - blocked.is_err(), - "ME->C quota reservation path must be serialized by held shared cross-mode lock" - ); - - drop(held_guard); - - let released = timeout( - Duration::from_millis(250), - process_me_writer_response( - MeResponse::Data { - flags: 0, - data: Bytes::from_static(&[0x42]), - }, - &mut writer, - ProtoTag::Intermediate, - &SecureRandom::new(), - &mut frame_buf, - &stats, - &user, - Some(1024), - 0, - &bytes_me2c, - 9902, - false, - false, - ), - ) - .await - .expect("ME->C write must complete after cross-mode lock release"); - - assert!(released.is_ok()); -} - -#[tokio::test] -async fn business_uncontended_cross_mode_lock_allows_me_to_client_quota_reservation() { - let stats = Stats::new(); - let user = format!("middle-me2c-cross-mode-free-{}", std::process::id()); - let mut writer = make_crypto_writer(tokio::io::sink()); - let mut frame_buf = Vec::new(); - let bytes_me2c = AtomicU64::new(0); - - let outcome = timeout( - Duration::from_millis(250), - process_me_writer_response( - MeResponse::Data { - flags: 0, - data: Bytes::from_static(&[0x55, 0x66]), - }, - &mut writer, - ProtoTag::Intermediate, - &SecureRandom::new(), - &mut frame_buf, - &stats, - &user, - Some(1024), - 0, - &bytes_me2c, - 9903, - false, - false, - ), - ) - .await - .expect("uncontended ME->C path should not stall"); - - assert!(outcome.is_ok()); - assert_eq!(stats.get_user_total_octets(&user), 2); - assert_eq!(bytes_me2c.load(std::sync::atomic::Ordering::Relaxed), 2); -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn adversarial_cross_mode_lock_is_released_before_me_to_client_write_await() { - let stats = Arc::new(Stats::new()); - let user = format!("middle-me2c-lock-drop-before-write-{}", std::process::id()); - let cross_mode_lock = cross_mode_quota_user_lock_for_tests(&user); - let bytes_me2c = Arc::new(AtomicU64::new(0)); - let writer_state = Arc::new(BlockingWriteState::default()); - - let worker = { - let stats = Arc::clone(&stats); - let user = user.clone(); - let cross_mode_lock = Arc::clone(&cross_mode_lock); - let bytes_me2c = Arc::clone(&bytes_me2c); - let writer_state = Arc::clone(&writer_state); - tokio::spawn(async move { - let mut writer = make_crypto_writer(BlockingWrite::new(writer_state)); - let mut frame_buf = Vec::new(); - let rng = SecureRandom::new(); - process_me_writer_response_with_cross_mode_lock( - MeResponse::Data { - flags: 0, - data: Bytes::from_static(&[0xDE, 0xAD, 0xBE, 0xEF]), - }, - &mut writer, - ProtoTag::Intermediate, - &rng, - &mut frame_buf, - stats.as_ref(), - &user, - Some(1024), - 0, - Some(&cross_mode_lock), - bytes_me2c.as_ref(), - 9910, - false, - false, - ) - .await - }) - }; - - wait_until_blocking_write_entered(&writer_state).await; - - let acquired_guard = timeout(Duration::from_millis(40), cross_mode_lock.lock()) - .await - .expect("cross-mode lock must be free while ME->C write is pending"); - drop(acquired_guard); - - release_blocking_write(&writer_state); - - let result = timeout(Duration::from_millis(300), worker) - .await - .expect("ME->C worker timed out after releasing blocking writer") - .expect("ME->C worker must not panic"); - - assert!(result.is_ok()); - assert_eq!(stats.get_user_total_octets(&user), 4); - assert_eq!(bytes_me2c.load(Ordering::Relaxed), 4); -} diff --git a/src/proxy/tests/middle_relay_hol_quota_security_tests.rs b/src/proxy/tests/middle_relay_hol_quota_security_tests.rs deleted file mode 100644 index 3ce0235..0000000 --- a/src/proxy/tests/middle_relay_hol_quota_security_tests.rs +++ /dev/null @@ -1,232 +0,0 @@ -use super::*; -use crate::crypto::{AesCtr, SecureRandom}; -use crate::stats::Stats; -use crate::stream::CryptoWriter; -use bytes::Bytes; -use std::io; -use std::pin::Pin; -use std::sync::Arc; -use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; -use std::task::{Context, Poll, Waker}; -use tokio::io::AsyncWrite; -use tokio::time::{Duration, timeout}; - -fn make_crypto_writer(writer: W) -> CryptoWriter -where - W: tokio::io::AsyncWrite + Unpin, -{ - let key = [0u8; 32]; - let iv = 0u128; - CryptoWriter::new(writer, AesCtr::new(&key, iv), 8 * 1024) -} - -#[derive(Default)] -struct GateState { - open: AtomicBool, - parked_waker: std::sync::Mutex>, -} - -impl GateState { - fn open(&self) { - self.open.store(true, Ordering::Relaxed); - if let Ok(mut guard) = self.parked_waker.lock() - && let Some(w) = guard.take() - { - w.wake(); - } - } - - fn has_waiter(&self) -> bool { - self.parked_waker - .lock() - .map(|guard| guard.is_some()) - .unwrap_or(false) - } -} - -#[derive(Default)] -struct GateWriter { - gate: Arc, -} - -impl GateWriter { - fn new(gate: Arc) -> Self { - Self { gate } - } -} - -impl AsyncWrite for GateWriter { - fn poll_write( - self: Pin<&mut Self>, - cx: &mut Context<'_>, - buf: &[u8], - ) -> Poll> { - if self.gate.open.load(Ordering::Relaxed) { - return Poll::Ready(Ok(buf.len())); - } - - if let Ok(mut guard) = self.gate.parked_waker.lock() { - *guard = Some(cx.waker().clone()); - } - Poll::Pending - } - - fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { - Poll::Ready(Ok(())) - } - - fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { - Poll::Ready(Ok(())) - } -} - -struct FailingWriter; - -impl AsyncWrite for FailingWriter { - fn poll_write( - self: Pin<&mut Self>, - _cx: &mut Context<'_>, - _buf: &[u8], - ) -> Poll> { - Poll::Ready(Err(io::Error::new( - io::ErrorKind::BrokenPipe, - "injected writer failure", - ))) - } - - fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { - Poll::Ready(Ok(())) - } - - fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { - Poll::Ready(Ok(())) - } -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn adversarial_same_user_slow_writer_must_not_hol_block_peer_connection() { - let stats = Stats::new(); - let bytes_me2c = AtomicU64::new(0); - let rng = SecureRandom::new(); - let quota_limit = Some(1024); - let user = "hol-quota-user"; - - let gate = Arc::new(GateState::default()); - - let mut blocked_writer = make_crypto_writer(GateWriter::new(Arc::clone(&gate))); - let slow_task = tokio::spawn(async move { - let mut frame_buf = Vec::new(); - process_me_writer_response( - MeResponse::Data { - flags: 0, - data: Bytes::from_static(&[0x10, 0x20, 0x30, 0x40]), - }, - &mut blocked_writer, - ProtoTag::Intermediate, - &rng, - &mut frame_buf, - &stats, - user, - quota_limit, - 0, - &bytes_me2c, - 7001, - false, - false, - ) - .await - }); - - timeout(Duration::from_millis(100), async { - loop { - if gate.has_waiter() { - break; - } - tokio::task::yield_now().await; - } - }) - .await - .expect("first writer must reach backpressure and park"); - - let stats_fast = Stats::new(); - let bytes_fast = AtomicU64::new(0); - let rng_fast = SecureRandom::new(); - let mut fast_writer = make_crypto_writer(tokio::io::sink()); - let mut frame_buf_fast = Vec::new(); - - timeout( - Duration::from_millis(50), - process_me_writer_response( - MeResponse::Data { - flags: 0, - data: Bytes::from_static(&[0x41]), - }, - &mut fast_writer, - ProtoTag::Intermediate, - &rng_fast, - &mut frame_buf_fast, - &stats_fast, - user, - quota_limit, - 0, - &bytes_fast, - 7002, - false, - false, - ), - ) - .await - .expect("peer connection must not be blocked by same-user stalled write") - .expect("fast peer write must succeed"); - - gate.open(); - let slow_result = timeout(Duration::from_secs(1), slow_task) - .await - .expect("stalled task must complete once gate opens") - .expect("stalled task must not panic"); - assert!(slow_result.is_ok()); -} - -#[tokio::test] -async fn negative_write_failure_rolls_back_pre_accounted_quota_and_forensics_bytes() { - let stats = Stats::new(); - let user = "rollback-user"; - stats.add_user_octets_from(user, 7); - - let bytes_me2c = AtomicU64::new(0); - let rng = SecureRandom::new(); - let mut writer = make_crypto_writer(FailingWriter); - let mut frame_buf = Vec::new(); - - let result = process_me_writer_response( - MeResponse::Data { - flags: 0, - data: Bytes::from_static(&[1, 2, 3, 4]), - }, - &mut writer, - ProtoTag::Intermediate, - &rng, - &mut frame_buf, - &stats, - user, - Some(64), - 0, - &bytes_me2c, - 7003, - false, - false, - ) - .await; - - assert!(matches!(result, Err(ProxyError::Io(_)))); - assert_eq!( - stats.get_user_total_octets(user), - 7, - "failed client write must not overcharge user quota accounting" - ); - assert_eq!( - bytes_me2c.load(Ordering::Relaxed), - 0, - "failed client write must not inflate ME->C forensic byte counter" - ); -} \ No newline at end of file diff --git a/src/proxy/tests/middle_relay_quota_extended_attack_surface_security_tests.rs b/src/proxy/tests/middle_relay_quota_extended_attack_surface_security_tests.rs deleted file mode 100644 index 29384e0..0000000 --- a/src/proxy/tests/middle_relay_quota_extended_attack_surface_security_tests.rs +++ /dev/null @@ -1,372 +0,0 @@ -use super::*; -use crate::crypto::{AesCtr, SecureRandom}; -use crate::error::ProxyError; -use crate::stats::Stats; -use crate::stream::CryptoWriter; -use bytes::Bytes; -use std::sync::atomic::{AtomicU64, Ordering}; -use std::sync::{Arc, OnceLock, Mutex}; -use tokio::sync::Mutex as AsyncMutex; -use tokio::task::JoinSet; -use tokio::time::{Duration, timeout}; - -fn make_crypto_writer(writer: W) -> CryptoWriter -where - W: tokio::io::AsyncWrite + Unpin, -{ - let key = [0u8; 32]; - let iv = 0u128; - CryptoWriter::new(writer, AesCtr::new(&key, iv), 8 * 1024) -} - -fn lookup_test_lock() -> &'static Mutex<()> { - static LOCK: OnceLock> = OnceLock::new(); - LOCK.get_or_init(|| Mutex::new(())) -} - -#[tokio::test] -async fn positive_me2c_quota_counts_bytes_exactly_once() { - let _guard = lookup_test_lock().lock().unwrap(); - let stats = Stats::new(); - let user = format!("quota-middle-ext-positive-{}", std::process::id()); - let lock = Arc::new(AsyncMutex::new(())); - - let mut writer = make_crypto_writer(tokio::io::sink()); - let mut frame_buf = Vec::new(); - let bytes_me2c = AtomicU64::new(0); - - let result = process_me_writer_response_with_cross_mode_lock( - MeResponse::Data { - flags: 0, - data: Bytes::from_static(&[1, 2, 3, 4, 5]), - }, - &mut writer, - ProtoTag::Intermediate, - &SecureRandom::new(), - &mut frame_buf, - &stats, - &user, - Some(64), - 0, - Some(&lock), - &bytes_me2c, - 70_001, - false, - false, - ) - .await; - - assert!(result.is_ok()); - assert_eq!(stats.get_user_total_octets(&user), 5); - assert_eq!(bytes_me2c.load(Ordering::Relaxed), 5); -} - -#[tokio::test] -async fn negative_held_crossmode_lock_blocks_me2c_write() { - let _guard = lookup_test_lock().lock().unwrap(); - let stats = Stats::new(); - let user = format!("quota-middle-ext-negative-{}", std::process::id()); - - let lock = Arc::new(AsyncMutex::new(())); - let _held = lock.try_lock().expect("lock must be held"); - - let mut writer = make_crypto_writer(tokio::io::sink()); - let mut frame_buf = Vec::new(); - let bytes_me2c = AtomicU64::new(0); - - let blocked = timeout( - Duration::from_millis(25), - process_me_writer_response_with_cross_mode_lock( - MeResponse::Data { - flags: 0, - data: Bytes::from_static(&[0xFE]), - }, - &mut writer, - ProtoTag::Intermediate, - &SecureRandom::new(), - &mut frame_buf, - &stats, - &user, - Some(16), - 0, - Some(&lock), - &bytes_me2c, - 70_101, - false, - false, - ), - ) - .await; - - assert!(blocked.is_err()); - assert_eq!(stats.get_user_total_octets(&user), 0); - assert_eq!(bytes_me2c.load(Ordering::Relaxed), 0); -} - -#[tokio::test] -async fn edge_zero_quota_zero_payload_is_fail_closed() { - let _guard = lookup_test_lock().lock().unwrap(); - let stats = Stats::new(); - let user = format!("quota-middle-ext-edge-{}", std::process::id()); - - let lock = Arc::new(AsyncMutex::new(())); - let mut writer = make_crypto_writer(tokio::io::sink()); - let mut frame_buf = Vec::new(); - let bytes_me2c = AtomicU64::new(0); - - let result = process_me_writer_response_with_cross_mode_lock( - MeResponse::Data { - flags: 0, - data: Bytes::new(), - }, - &mut writer, - ProtoTag::Intermediate, - &SecureRandom::new(), - &mut frame_buf, - &stats, - &user, - Some(0), - 0, - Some(&lock), - &bytes_me2c, - 70_201, - false, - false, - ) - .await; - - assert!(matches!(result, Err(ProxyError::DataQuotaExceeded { .. }))); - assert_eq!(stats.get_user_total_octets(&user), 0); -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 4)] -async fn adversarial_parallel_me2c_race_falls_back_to_quota_error() { - let _guard = lookup_test_lock().lock().unwrap(); - let stats = Arc::new(Stats::new()); - let user = format!("quota-middle-ext-blackhat-{}", std::process::id()); - let quota = 64u64; - let lock = Arc::new(AsyncMutex::new(())); - let bytes_me2c = Arc::new(AtomicU64::new(0)); - - let mut set = JoinSet::new(); - for i in 0..256u64 { - let stats = Arc::clone(&stats); - let user = user.clone(); - let lock = Arc::clone(&lock); - let bytes_me2c = Arc::clone(&bytes_me2c); - - set.spawn(async move { - let mut writer = make_crypto_writer(tokio::io::sink()); - let mut frame_buf = Vec::new(); - let payload = vec![((i & 0xFF) as u8); (i % 4 + 1) as usize]; - - process_me_writer_response_with_cross_mode_lock( - MeResponse::Data { - flags: 0, - data: Bytes::from(payload), - }, - &mut writer, - ProtoTag::Intermediate, - &SecureRandom::new(), - &mut frame_buf, - stats.as_ref(), - &user, - Some(quota), - 0, - Some(&lock), - bytes_me2c.as_ref(), - 70_301 + i, - false, - false, - ) - .await - }); - } - - let mut succeeded = 0usize; - while let Some(done) = set.join_next().await { - match done.expect("task must not panic") { - Ok(_) => succeeded += 1, - Err(ProxyError::DataQuotaExceeded { .. }) => {} - Err(other) => panic!("unexpected error {other:?}"), - } - } - - assert_eq!(stats.get_user_total_octets(&user), bytes_me2c.load(Ordering::Relaxed)); - assert!(stats.get_user_total_octets(&user) <= quota); - assert!(succeeded <= quota as usize); -} - -#[tokio::test] -async fn integration_shared_prefetched_lock_blocks_then_releases_writer() { - let stats = Stats::new(); - let user = format!("quota-middle-ext-integration-{}", std::process::id()); - let lock = Arc::new(AsyncMutex::new(())); - let held = lock - .try_lock() - .expect("integration test must hold prefetched lock first"); - - let mut writer = make_crypto_writer(tokio::io::sink()); - let mut frame_buf = Vec::new(); - let bytes_me2c = AtomicU64::new(0); - - let blocked = timeout( - Duration::from_millis(25), - process_me_writer_response_with_cross_mode_lock( - MeResponse::Data { - flags: 0, - data: Bytes::from_static(&[0xA1]), - }, - &mut writer, - ProtoTag::Intermediate, - &SecureRandom::new(), - &mut frame_buf, - &stats, - &user, - Some(8), - 0, - Some(&lock), - &bytes_me2c, - 70_360, - false, - false, - ), - ) - .await; - assert!(blocked.is_err()); - - drop(held); - - let after_release = timeout( - Duration::from_millis(150), - process_me_writer_response_with_cross_mode_lock( - MeResponse::Data { - flags: 0, - data: Bytes::from_static(&[0xA2]), - }, - &mut writer, - ProtoTag::Intermediate, - &SecureRandom::new(), - &mut frame_buf, - &stats, - &user, - Some(8), - 0, - Some(&lock), - &bytes_me2c, - 70_361, - false, - false, - ), - ) - .await - .expect("writer should progress once the shared lock is released"); - - assert!(after_release.is_ok()); -} - -#[tokio::test] -async fn light_fuzz_small_payloads_toggle_lock_state_stays_consistent() { - let _guard = lookup_test_lock().lock().unwrap(); - let stats = Stats::new(); - let user = format!("quota-middle-ext-fuzz-{}", std::process::id()); - let mut seed = 0xCAFE_BABE_1234u64; - let bytes_me2c = AtomicU64::new(0); - - for case in 0..48u32 { - seed ^= seed << 5; - seed ^= seed >> 12; - seed ^= seed << 13; - let hold = (seed & 0x1) == 0; - - let lock = Arc::new(AsyncMutex::new(())); - let maybe_guard = if hold { - Some(lock.try_lock().unwrap()) - } else { - None - }; - - let mut writer = make_crypto_writer(tokio::io::sink()); - let mut frame_buf = Vec::new(); - - let result = timeout( - Duration::from_millis(30), - process_me_writer_response_with_cross_mode_lock( - MeResponse::Data { - flags: 0, - data: Bytes::from(vec![(seed & 0xFF) as u8; ((seed as usize % 5) + 1)]), - }, - &mut writer, - ProtoTag::Intermediate, - &SecureRandom::new(), - &mut frame_buf, - &stats, - &user, - Some(128), - 0, - Some(&lock), - &bytes_me2c, - 70_401 + case as u64, - false, - false, - ), - ) - .await; - - if hold { - assert!(result.is_err()); - } else { - assert!(result.unwrap().is_ok()); - } - - drop(maybe_guard); - } -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 4)] -async fn stress_parallel_free_users_during_held_user_lock_maintains_liveness() { - let _guard = lookup_test_lock().lock().unwrap(); - let held = Arc::new(AsyncMutex::new(())); - let _held_guard = held.try_lock().unwrap(); - - let mut set = JoinSet::new(); - for i in 0..48u64 { - set.spawn(async move { - let stats = Stats::new(); - let user = format!("quota-middle-ext-stress-free-{i}"); - let mut writer = make_crypto_writer(tokio::io::sink()); - let mut frame_buf = Vec::new(); - let bytes_me2c = AtomicU64::new(0); - let free_lock = Arc::new(AsyncMutex::new(())); - - process_me_writer_response_with_cross_mode_lock( - MeResponse::Data { - flags: 0, - data: Bytes::from_static(&[0xEE]), - }, - &mut writer, - ProtoTag::Intermediate, - &SecureRandom::new(), - &mut frame_buf, - &stats, - &user, - Some(1), - 0, - Some(&free_lock), - &bytes_me2c, - 70_500 + i, - false, - false, - ) - .await - }); - } - - timeout(Duration::from_secs(2), async { - while let Some(task) = set.join_next().await { - task.unwrap().unwrap(); - } - }) - .await - .unwrap(); -} diff --git a/src/proxy/tests/middle_relay_quota_overflow_lock_security_tests.rs b/src/proxy/tests/middle_relay_quota_overflow_lock_security_tests.rs deleted file mode 100644 index d06e103..0000000 --- a/src/proxy/tests/middle_relay_quota_overflow_lock_security_tests.rs +++ /dev/null @@ -1,131 +0,0 @@ -use super::*; -use dashmap::DashMap; -use std::sync::Arc; - -#[test] -fn saturation_uses_stable_overflow_lock_without_cache_growth() { - let _guard = super::quota_user_lock_test_scope(); - let map = QUOTA_USER_LOCKS.get_or_init(DashMap::new); - map.clear(); - - let prefix = format!("middle-quota-held-{}", std::process::id()); - let mut retained = Vec::with_capacity(QUOTA_USER_LOCKS_MAX); - for idx in 0..QUOTA_USER_LOCKS_MAX { - retained.push(quota_user_lock(&format!("{prefix}-{idx}"))); - } - - assert_eq!(map.len(), QUOTA_USER_LOCKS_MAX); - - let user = format!("middle-quota-overflow-{}", std::process::id()); - let first = quota_user_lock(&user); - let second = quota_user_lock(&user); - - assert!( - Arc::ptr_eq(&first, &second), - "overflow user must get deterministic same lock while cache is saturated" - ); - assert_eq!( - map.len(), - QUOTA_USER_LOCKS_MAX, - "overflow path must not grow bounded lock map" - ); - assert!( - map.get(&user).is_none(), - "overflow user should stay outside bounded lock map under saturation" - ); - - drop(retained); -} - -#[test] -fn overflow_striping_keeps_different_users_distributed() { - let _guard = super::quota_user_lock_test_scope(); - let map = QUOTA_USER_LOCKS.get_or_init(DashMap::new); - map.clear(); - - let prefix = format!("middle-quota-dist-held-{}", std::process::id()); - let mut retained = Vec::with_capacity(QUOTA_USER_LOCKS_MAX); - for idx in 0..QUOTA_USER_LOCKS_MAX { - retained.push(quota_user_lock(&format!("{prefix}-{idx}"))); - } - - let a = quota_user_lock("middle-overflow-user-a"); - let b = quota_user_lock("middle-overflow-user-b"); - let c = quota_user_lock("middle-overflow-user-c"); - - let distinct = [ - Arc::as_ptr(&a) as usize, - Arc::as_ptr(&b) as usize, - Arc::as_ptr(&c) as usize, - ] - .iter() - .copied() - .collect::>() - .len(); - - assert!( - distinct >= 2, - "striped overflow lock set should avoid collapsing all users to one lock" - ); - - drop(retained); -} - -#[test] -fn reclaim_path_caches_new_user_after_stale_entries_drop() { - let _guard = super::quota_user_lock_test_scope(); - let map = QUOTA_USER_LOCKS.get_or_init(DashMap::new); - map.clear(); - - let prefix = format!("middle-quota-reclaim-held-{}", std::process::id()); - let mut retained = Vec::with_capacity(QUOTA_USER_LOCKS_MAX); - for idx in 0..QUOTA_USER_LOCKS_MAX { - retained.push(quota_user_lock(&format!("{prefix}-{idx}"))); - } - - drop(retained); - - let user = format!("middle-quota-reclaim-user-{}", std::process::id()); - let got = quota_user_lock(&user); - assert!(map.get(&user).is_some()); - assert!( - Arc::strong_count(&got) >= 2, - "after reclaim, lock should be held both by caller and map" - ); -} - -#[test] -fn overflow_path_same_user_is_stable_across_parallel_threads() { - let _guard = super::quota_user_lock_test_scope(); - let map = QUOTA_USER_LOCKS.get_or_init(DashMap::new); - map.clear(); - - let mut retained = Vec::with_capacity(QUOTA_USER_LOCKS_MAX); - for idx in 0..QUOTA_USER_LOCKS_MAX { - retained.push(quota_user_lock(&format!( - "middle-quota-thread-held-{}-{idx}", - std::process::id() - ))); - } - - let user = format!("middle-quota-overflow-thread-user-{}", std::process::id()); - let mut workers = Vec::new(); - for _ in 0..32 { - let user = user.clone(); - workers.push(std::thread::spawn(move || quota_user_lock(&user))); - } - - let first = workers - .remove(0) - .join() - .expect("thread must return lock handle"); - for worker in workers { - let got = worker.join().expect("thread must return lock handle"); - assert!( - Arc::ptr_eq(&first, &got), - "same overflow user should resolve to one striped lock even under contention" - ); - } - - drop(retained); -} diff --git a/src/proxy/tests/middle_relay_quota_reservation_adversarial_tests.rs b/src/proxy/tests/middle_relay_quota_reservation_adversarial_tests.rs deleted file mode 100644 index 963b3e0..0000000 --- a/src/proxy/tests/middle_relay_quota_reservation_adversarial_tests.rs +++ /dev/null @@ -1,1066 +0,0 @@ -use super::*; -use crate::crypto::{AesCtr, SecureRandom}; -use crate::stats::Stats; -use crate::stream::CryptoWriter; -use bytes::Bytes; -use std::sync::Arc; -use std::sync::atomic::{AtomicU64, Ordering}; -use std::task::{Context, Poll}; -use tokio::io::AsyncWrite; -use tokio::task::JoinSet; - -fn make_crypto_writer(writer: W) -> CryptoWriter -where - W: tokio::io::AsyncWrite + Unpin, -{ - let key = [0u8; 32]; - let iv = 0u128; - CryptoWriter::new(writer, AesCtr::new(&key, iv), 8 * 1024) -} - -struct FailingWriter; - -impl AsyncWrite for FailingWriter { - fn poll_write( - self: std::pin::Pin<&mut Self>, - _cx: &mut Context<'_>, - _buf: &[u8], - ) -> Poll> { - Poll::Ready(Err(std::io::Error::other("forced writer failure"))) - } - - fn poll_flush( - self: std::pin::Pin<&mut Self>, - _cx: &mut Context<'_>, - ) -> Poll> { - Poll::Ready(Ok(())) - } - - fn poll_shutdown( - self: std::pin::Pin<&mut Self>, - _cx: &mut Context<'_>, - ) -> Poll> { - Poll::Ready(Ok(())) - } -} - -struct FailAfterBudgetWriter { - remaining: usize, - written: usize, -} - -impl FailAfterBudgetWriter { - fn new(remaining: usize) -> Self { - Self { - remaining, - written: 0, - } - } -} - -impl AsyncWrite for FailAfterBudgetWriter { - fn poll_write( - mut self: std::pin::Pin<&mut Self>, - _cx: &mut Context<'_>, - buf: &[u8], - ) -> Poll> { - if self.remaining == 0 { - return Poll::Ready(Err(std::io::Error::other("forced short-write exhaustion"))); - } - - let n = self.remaining.min(buf.len()); - self.remaining -= n; - self.written += n; - Poll::Ready(Ok(n)) - } - - fn poll_flush( - self: std::pin::Pin<&mut Self>, - _cx: &mut Context<'_>, - ) -> Poll> { - Poll::Ready(Ok(())) - } - - fn poll_shutdown( - self: std::pin::Pin<&mut Self>, - _cx: &mut Context<'_>, - ) -> Poll> { - Poll::Ready(Ok(())) - } -} - -#[tokio::test] -async fn positive_exact_quota_boundary_allows_last_frame_and_blocks_next() { - let stats = Stats::new(); - let user = "quota-boundary-user"; - let bytes_me2c = AtomicU64::new(0); - - stats.add_user_octets_from(user, 5); - - let mut writer_one = make_crypto_writer(tokio::io::sink()); - let mut frame_buf_one = Vec::new(); - let first = process_me_writer_response( - MeResponse::Data { - flags: 0, - data: Bytes::from_static(&[1, 2, 3]), - }, - &mut writer_one, - ProtoTag::Intermediate, - &SecureRandom::new(), - &mut frame_buf_one, - &stats, - user, - Some(8), - 0, - &bytes_me2c, - 7101, - false, - false, - ) - .await; - - assert!(first.is_ok(), "frame that reaches boundary must be allowed"); - assert_eq!(stats.get_user_total_octets(user), 8); - - let mut writer_two = make_crypto_writer(tokio::io::sink()); - let mut frame_buf_two = Vec::new(); - let second = process_me_writer_response( - MeResponse::Data { - flags: 0, - data: Bytes::from_static(&[9]), - }, - &mut writer_two, - ProtoTag::Intermediate, - &SecureRandom::new(), - &mut frame_buf_two, - &stats, - user, - Some(8), - 0, - &bytes_me2c, - 7102, - false, - false, - ) - .await; - - assert!( - matches!(second, Err(ProxyError::DataQuotaExceeded { .. })), - "frame after boundary must be rejected" - ); - assert_eq!(stats.get_user_total_octets(user), 8); - assert_eq!(bytes_me2c.load(Ordering::Relaxed), 3); -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 4)] -async fn adversarial_parallel_reservation_stress_never_overshoots_quota_or_counters() { - let stats = Arc::new(Stats::new()); - let user = "reservation-stress-user"; - let quota_limit = 64u64; - let bytes_me2c = Arc::new(AtomicU64::new(0)); - let mut tasks = JoinSet::new(); - - for idx in 0..256u64 { - let user_owned = user.to_string(); - let stats_ref = Arc::clone(&stats); - let bytes_ref = Arc::clone(&bytes_me2c); - - tasks.spawn(async move { - let mut writer = make_crypto_writer(tokio::io::sink()); - let mut frame_buf = Vec::new(); - process_me_writer_response( - MeResponse::Data { - flags: 0, - data: Bytes::from_static(&[0xAB]), - }, - &mut writer, - ProtoTag::Intermediate, - &SecureRandom::new(), - &mut frame_buf, - stats_ref.as_ref(), - &user_owned, - Some(quota_limit), - 0, - bytes_ref.as_ref(), - 7200 + idx, - false, - false, - ) - .await - }); - } - - let mut ok = 0usize; - let mut denied = 0usize; - while let Some(joined) = tasks.join_next().await { - match joined.expect("reservation stress task must not panic") { - Ok(_) => ok += 1, - Err(ProxyError::DataQuotaExceeded { .. }) => denied += 1, - Err(other) => panic!("unexpected error in stress case: {other:?}"), - } - } - - let total = stats.get_user_total_octets(user); - assert_eq!( - total, quota_limit, - "quota must be exactly exhausted without overshoot" - ); - assert_eq!( - bytes_me2c.load(Ordering::Relaxed), - total, - "ME->C forensic bytes must track committed quota usage" - ); - assert_eq!(ok, quota_limit as usize, "exactly quota_limit tasks must succeed"); - assert_eq!( - denied, - 256usize - (quota_limit as usize), - "remaining tasks must be exactly denied without silently swallowing state" - ); -} - -#[tokio::test] -async fn light_fuzz_random_frame_sizes_preserve_quota_and_counter_consistency() { - let stats = Stats::new(); - let user = "reservation-fuzz-user"; - let quota_limit = 128u64; - let bytes_me2c = AtomicU64::new(0); - let mut seed = 0xC0FE_EE11_8899_2211u64; - - for conn in 0..512u64 { - seed ^= seed << 7; - seed ^= seed >> 9; - seed ^= seed << 8; - let len = ((seed & 0x0f) + 1) as usize; - let payload = vec![0x5A; len]; - - let mut writer = make_crypto_writer(tokio::io::sink()); - let mut frame_buf = Vec::new(); - let result = process_me_writer_response( - MeResponse::Data { - flags: 0, - data: Bytes::from(payload), - }, - &mut writer, - ProtoTag::Intermediate, - &SecureRandom::new(), - &mut frame_buf, - &stats, - user, - Some(quota_limit), - 0, - &bytes_me2c, - 7300 + conn, - false, - false, - ) - .await; - - if let Err(err) = result { - assert!( - matches!(err, ProxyError::DataQuotaExceeded { .. }), - "fuzz run produced unexpected error variant: {err:?}" - ); - } - } - - let total = stats.get_user_total_octets(user); - assert!(total <= quota_limit); - assert_eq!(bytes_me2c.load(Ordering::Relaxed), total); -} - -#[tokio::test] -async fn positive_soft_overshoot_allows_burst_inside_soft_cap_then_blocks() { - let stats = Stats::new(); - let user = "soft-cap-boundary-user"; - let bytes_me2c = AtomicU64::new(0); - let quota_limit = 10u64; - let overshoot = 3u64; - - stats.add_user_octets_from(user, 10); - - let mut writer_one = make_crypto_writer(tokio::io::sink()); - let mut frame_buf_one = Vec::new(); - let first = process_me_writer_response( - MeResponse::Data { - flags: 0, - data: Bytes::from_static(&[1, 2, 3]), - }, - &mut writer_one, - ProtoTag::Intermediate, - &SecureRandom::new(), - &mut frame_buf_one, - &stats, - user, - Some(quota_limit), - overshoot, - &bytes_me2c, - 7401, - false, - false, - ) - .await; - assert!(first.is_ok(), "soft-cap buffer should allow reaching limit+overshoot"); - assert_eq!(stats.get_user_total_octets(user), 13); - - let mut writer_two = make_crypto_writer(tokio::io::sink()); - let mut frame_buf_two = Vec::new(); - let second = process_me_writer_response( - MeResponse::Data { - flags: 0, - data: Bytes::from_static(&[9]), - }, - &mut writer_two, - ProtoTag::Intermediate, - &SecureRandom::new(), - &mut frame_buf_two, - &stats, - user, - Some(quota_limit), - overshoot, - &bytes_me2c, - 7402, - false, - false, - ) - .await; - assert!(matches!(second, Err(ProxyError::DataQuotaExceeded { .. }))); - assert_eq!(stats.get_user_total_octets(user), 13); - assert_eq!(bytes_me2c.load(Ordering::Relaxed), 3); -} - -#[tokio::test] -async fn negative_soft_overshoot_rejects_when_payload_exceeds_remaining_soft_budget() { - let stats = Stats::new(); - let user = "soft-cap-remaining-user"; - let bytes_me2c = AtomicU64::new(0); - let quota_limit = 10u64; - let overshoot = 4u64; - - stats.add_user_octets_from(user, 12); - - let mut writer = make_crypto_writer(tokio::io::sink()); - let mut frame_buf = Vec::new(); - let result = process_me_writer_response( - MeResponse::Data { - flags: 0, - data: Bytes::from_static(&[1, 2, 3]), - }, - &mut writer, - ProtoTag::Intermediate, - &SecureRandom::new(), - &mut frame_buf, - &stats, - user, - Some(quota_limit), - overshoot, - &bytes_me2c, - 7501, - false, - false, - ) - .await; - - assert!(matches!(result, Err(ProxyError::DataQuotaExceeded { .. }))); - assert_eq!(stats.get_user_total_octets(user), 12); - assert_eq!(bytes_me2c.load(Ordering::Relaxed), 0); -} - -#[tokio::test] -async fn negative_write_failure_rolls_back_reservation_under_soft_cap_mode() { - let stats = Stats::new(); - let user = "soft-cap-rollback-user"; - let bytes_me2c = AtomicU64::new(0); - let mut writer = make_crypto_writer(FailingWriter); - let mut frame_buf = Vec::new(); - - stats.add_user_octets_from(user, 9); - - let result = process_me_writer_response( - MeResponse::Data { - flags: 0, - data: Bytes::from_static(&[1, 2, 3]), - }, - &mut writer, - ProtoTag::Intermediate, - &SecureRandom::new(), - &mut frame_buf, - &stats, - user, - Some(10), - 8, - &bytes_me2c, - 7601, - false, - false, - ) - .await; - - assert!(matches!(result, Err(ProxyError::Io(_)))); - assert_eq!(stats.get_user_total_octets(user), 9); - assert_eq!(bytes_me2c.load(Ordering::Relaxed), 0); -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 4)] -async fn adversarial_parallel_soft_cap_stress_never_exceeds_soft_limit() { - let stats = Arc::new(Stats::new()); - let user = "soft-cap-stress-user"; - let quota_limit = 40u64; - let overshoot = 5u64; - let soft_limit = quota_limit + overshoot; - let bytes_me2c = Arc::new(AtomicU64::new(0)); - let mut tasks = JoinSet::new(); - - for idx in 0..256u64 { - let user_owned = user.to_string(); - let stats_ref = Arc::clone(&stats); - let bytes_ref = Arc::clone(&bytes_me2c); - tasks.spawn(async move { - let mut writer = make_crypto_writer(tokio::io::sink()); - let mut frame_buf = Vec::new(); - process_me_writer_response( - MeResponse::Data { - flags: 0, - data: Bytes::from_static(&[0x42]), - }, - &mut writer, - ProtoTag::Intermediate, - &SecureRandom::new(), - &mut frame_buf, - stats_ref.as_ref(), - &user_owned, - Some(quota_limit), - overshoot, - bytes_ref.as_ref(), - 7700 + idx, - false, - false, - ) - .await - }); - } - - while let Some(joined) = tasks.join_next().await { - match joined.expect("soft-cap stress task must not panic") { - Ok(_) | Err(ProxyError::DataQuotaExceeded { .. }) => {} - Err(other) => panic!("unexpected error in soft-cap stress case: {other:?}"), - } - } - - let total = stats.get_user_total_octets(user); - assert!(total <= soft_limit, "soft-cap stress must never overshoot soft limit"); - assert_eq!(bytes_me2c.load(Ordering::Relaxed), total); -} - -#[tokio::test] -async fn light_fuzz_soft_cap_matrix_keeps_counters_and_limits_consistent() { - let stats = Stats::new(); - let user = "soft-cap-fuzz-user"; - let bytes_me2c = AtomicU64::new(0); - let mut seed = 0x9E37_79B9_7F4A_7C15u64; - - for conn in 0..1024u64 { - seed ^= seed << 7; - seed ^= seed >> 9; - seed ^= seed << 8; - - let quota_limit = 32 + (seed & 0x3f); - let overshoot = seed.rotate_left(13) & 0x0f; - let len = ((seed >> 3) & 0x07) + 1; - let payload = vec![0xA5; len as usize]; - let before = stats.get_user_total_octets(user); - - let mut writer = make_crypto_writer(tokio::io::sink()); - let mut frame_buf = Vec::new(); - let result = process_me_writer_response( - MeResponse::Data { - flags: 0, - data: Bytes::from(payload), - }, - &mut writer, - ProtoTag::Intermediate, - &SecureRandom::new(), - &mut frame_buf, - &stats, - user, - Some(quota_limit), - overshoot, - &bytes_me2c, - 7800 + conn, - false, - false, - ) - .await; - - if let Err(ref err) = result { - assert!( - matches!(err, ProxyError::DataQuotaExceeded { .. }), - "soft-cap fuzz produced unexpected error variant: {err:?}" - ); - } - - let after = stats.get_user_total_octets(user); - let soft_limit = quota_limit.saturating_add(overshoot); - match result { - Ok(_) => { - assert_eq!(after, before.saturating_add(len)); - assert!(after <= soft_limit, "accepted write must stay within active soft cap"); - } - Err(_) => { - assert_eq!(after, before, "rejected write must not mutate quota state"); - } - } - assert_eq!( - bytes_me2c.load(Ordering::Relaxed), - after, - "soft-cap fuzz must keep counters synchronized" - ); - } -} - -#[tokio::test] -async fn positive_no_quota_limit_accumulates_data_octets_exactly() { - let stats = Stats::new(); - let user = "no-quota-user"; - let bytes_me2c = AtomicU64::new(0); - let mut expected = 0u64; - - for (idx, len) in [1usize, 2, 3, 5, 8, 13, 21].iter().copied().enumerate() { - let mut writer = make_crypto_writer(tokio::io::sink()); - let mut frame_buf = Vec::new(); - let payload = vec![0x41; len]; - let result = process_me_writer_response( - MeResponse::Data { - flags: 0, - data: Bytes::from(payload), - }, - &mut writer, - ProtoTag::Intermediate, - &SecureRandom::new(), - &mut frame_buf, - &stats, - user, - None, - 0, - &bytes_me2c, - 7900 + idx as u64, - false, - false, - ) - .await; - - assert!(result.is_ok()); - expected += len as u64; - } - - assert_eq!(stats.get_user_total_octets(user), expected); - assert_eq!(bytes_me2c.load(Ordering::Relaxed), expected); -} - -#[tokio::test] -async fn negative_zero_quota_rejects_non_empty_payload() { - let stats = Stats::new(); - let user = "zero-quota-user"; - let bytes_me2c = AtomicU64::new(0); - - let mut writer = make_crypto_writer(tokio::io::sink()); - let mut frame_buf = Vec::new(); - let result = process_me_writer_response( - MeResponse::Data { - flags: 0, - data: Bytes::from_static(&[0xAA]), - }, - &mut writer, - ProtoTag::Intermediate, - &SecureRandom::new(), - &mut frame_buf, - &stats, - user, - Some(0), - 0, - &bytes_me2c, - 8001, - false, - false, - ) - .await; - - assert!(matches!(result, Err(ProxyError::DataQuotaExceeded { .. }))); - assert_eq!(stats.get_user_total_octets(user), 0); - assert_eq!(bytes_me2c.load(Ordering::Relaxed), 0); -} - -#[tokio::test] -async fn edge_zero_length_payload_with_zero_quota_is_fail_closed() { - let stats = Stats::new(); - let user = "zero-len-zero-quota-user"; - let bytes_me2c = AtomicU64::new(0); - - let mut writer = make_crypto_writer(tokio::io::sink()); - let mut frame_buf = Vec::new(); - let result = process_me_writer_response( - MeResponse::Data { - flags: 0, - data: Bytes::new(), - }, - &mut writer, - ProtoTag::Intermediate, - &SecureRandom::new(), - &mut frame_buf, - &stats, - user, - Some(0), - 0, - &bytes_me2c, - 8002, - false, - false, - ) - .await; - - assert!(matches!(result, Err(ProxyError::DataQuotaExceeded { .. }))); - assert_eq!(stats.get_user_total_octets(user), 0); - assert_eq!(bytes_me2c.load(Ordering::Relaxed), 0); -} - -#[tokio::test] -async fn positive_ack_response_does_not_touch_quota_counters() { - let stats = Stats::new(); - let user = "ack-accounting-user"; - let bytes_me2c = AtomicU64::new(11); - stats.add_user_octets_to(user, 23); - - let mut writer = make_crypto_writer(tokio::io::sink()); - let mut frame_buf = Vec::new(); - let result = process_me_writer_response( - MeResponse::Ack(0x33445566), - &mut writer, - ProtoTag::Intermediate, - &SecureRandom::new(), - &mut frame_buf, - &stats, - user, - Some(24), - 0, - &bytes_me2c, - 8003, - true, - true, - ) - .await; - - assert!(result.is_ok()); - assert_eq!(stats.get_user_total_octets(user), 23); - assert_eq!(bytes_me2c.load(Ordering::Relaxed), 11); -} - -#[tokio::test] -async fn edge_close_response_is_accounting_noop() { - let stats = Stats::new(); - let user = "close-accounting-user"; - let bytes_me2c = AtomicU64::new(19); - stats.add_user_octets_to(user, 31); - - let mut writer = make_crypto_writer(tokio::io::sink()); - let mut frame_buf = Vec::new(); - let result = process_me_writer_response( - MeResponse::Close, - &mut writer, - ProtoTag::Intermediate, - &SecureRandom::new(), - &mut frame_buf, - &stats, - user, - Some(40), - 3, - &bytes_me2c, - 8004, - false, - true, - ) - .await; - - assert!(result.is_ok()); - assert_eq!(stats.get_user_total_octets(user), 31); - assert_eq!(bytes_me2c.load(Ordering::Relaxed), 19); -} - -#[tokio::test] -async fn negative_preloaded_above_soft_cap_rejects_even_single_byte() { - let stats = Stats::new(); - let user = "preloaded-over-soft-cap-user"; - let bytes_me2c = AtomicU64::new(0); - let quota_limit = 20u64; - let overshoot = 2u64; - stats.add_user_octets_to(user, quota_limit + overshoot + 1); - - let mut writer = make_crypto_writer(tokio::io::sink()); - let mut frame_buf = Vec::new(); - let result = process_me_writer_response( - MeResponse::Data { - flags: 0, - data: Bytes::from_static(&[1]), - }, - &mut writer, - ProtoTag::Intermediate, - &SecureRandom::new(), - &mut frame_buf, - &stats, - user, - Some(quota_limit), - overshoot, - &bytes_me2c, - 8005, - false, - false, - ) - .await; - - assert!(matches!(result, Err(ProxyError::DataQuotaExceeded { .. }))); - assert_eq!(bytes_me2c.load(Ordering::Relaxed), 0); - assert_eq!(stats.get_user_total_octets(user), quota_limit + overshoot + 1); -} - -#[tokio::test] -async fn adversarial_fail_writer_path_never_desynchronizes_quota_accounting() { - let stats = Stats::new(); - let user = "partial-write-rollback-user"; - let bytes_me2c = AtomicU64::new(0); - let mut writer = make_crypto_writer(FailAfterBudgetWriter::new(7)); - let mut frame_buf = Vec::new(); - let payload_len = 16 * 1024u64; - - let result = process_me_writer_response( - MeResponse::Data { - flags: 0, - data: Bytes::from(vec![0x42; 16 * 1024]), - }, - &mut writer, - ProtoTag::Intermediate, - &SecureRandom::new(), - &mut frame_buf, - &stats, - user, - Some(payload_len), - 0, - &bytes_me2c, - 8006, - false, - false, - ) - .await; - - let total_after = stats.get_user_total_octets(user); - let forensic_after = bytes_me2c.load(Ordering::Relaxed); - assert_eq!(forensic_after, total_after); - assert!( - total_after == 0 || total_after == payload_len, - "writer failure path must either roll back fully or commit exactly one payload" - ); - - // Regardless of whether I/O failure surfaced immediately or was deferred, - // accounting must remain fail-closed and prevent silent overshoot. - let mut writer_two = make_crypto_writer(tokio::io::sink()); - let mut frame_buf_two = Vec::new(); - let second = process_me_writer_response( - MeResponse::Data { - flags: 0, - data: Bytes::from_static(&[0x99]), - }, - &mut writer_two, - ProtoTag::Intermediate, - &SecureRandom::new(), - &mut frame_buf_two, - &stats, - user, - Some(payload_len), - 0, - &bytes_me2c, - 8007, - false, - false, - ) - .await; - - if total_after == payload_len { - assert!(matches!(second, Err(ProxyError::DataQuotaExceeded { .. }))); - } else { - assert!(second.is_ok()); - } -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 4)] -async fn stress_parallel_oversized_frames_fail_closed_without_counter_leak() { - let stats = Arc::new(Stats::new()); - let user = "parallel-fail-rollback-user"; - let bytes_me2c = Arc::new(AtomicU64::new(0)); - let mut tasks = JoinSet::new(); - - for idx in 0..256u64 { - let user_owned = user.to_string(); - let stats_ref = Arc::clone(&stats); - let bytes_ref = Arc::clone(&bytes_me2c); - tasks.spawn(async move { - let mut writer = make_crypto_writer(tokio::io::sink()); - let mut frame_buf = Vec::new(); - process_me_writer_response( - MeResponse::Data { - flags: 0, - data: Bytes::from(vec![0xEE; 12 * 1024]), - }, - &mut writer, - ProtoTag::Intermediate, - &SecureRandom::new(), - &mut frame_buf, - stats_ref.as_ref(), - &user_owned, - Some(512), - 0, - bytes_ref.as_ref(), - 8100 + idx, - false, - false, - ) - .await - }); - } - - while let Some(joined) = tasks.join_next().await { - let result = joined.expect("parallel fail writer task must not panic"); - assert!(matches!(result, Err(ProxyError::DataQuotaExceeded { .. }))); - } - - assert_eq!(stats.get_user_total_octets(user), 0); - assert_eq!(bytes_me2c.load(Ordering::Relaxed), 0); -} - -#[tokio::test] -async fn integration_mixed_data_ack_close_sequence_preserves_data_only_accounting() { - let stats = Stats::new(); - let user = "mixed-sequence-user"; - let bytes_me2c = AtomicU64::new(0); - - let mut writer = make_crypto_writer(tokio::io::sink()); - let mut frame_buf = Vec::new(); - - let data_one = process_me_writer_response( - MeResponse::Data { - flags: 0, - data: Bytes::from_static(&[1, 2, 3]), - }, - &mut writer, - ProtoTag::Intermediate, - &SecureRandom::new(), - &mut frame_buf, - &stats, - user, - Some(32), - 0, - &bytes_me2c, - 8201, - false, - false, - ) - .await; - assert!(data_one.is_ok()); - - let ack = process_me_writer_response( - MeResponse::Ack(0x0102_0304), - &mut writer, - ProtoTag::Intermediate, - &SecureRandom::new(), - &mut frame_buf, - &stats, - user, - Some(32), - 0, - &bytes_me2c, - 8202, - true, - true, - ) - .await; - assert!(ack.is_ok()); - - let data_two = process_me_writer_response( - MeResponse::Data { - flags: 0, - data: Bytes::from_static(&[4, 5]), - }, - &mut writer, - ProtoTag::Intermediate, - &SecureRandom::new(), - &mut frame_buf, - &stats, - user, - Some(32), - 0, - &bytes_me2c, - 8203, - false, - true, - ) - .await; - assert!(data_two.is_ok()); - - let close = process_me_writer_response( - MeResponse::Close, - &mut writer, - ProtoTag::Intermediate, - &SecureRandom::new(), - &mut frame_buf, - &stats, - user, - Some(32), - 0, - &bytes_me2c, - 8204, - false, - true, - ) - .await; - assert!(close.is_ok()); - - assert_eq!(stats.get_user_total_octets(user), 5); - assert_eq!(bytes_me2c.load(Ordering::Relaxed), 5); -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 4)] -async fn stress_parallel_multi_user_quota_isolation_no_cross_user_leakage() { - let stats = Arc::new(Stats::new()); - let user_a = "quota-isolation-a"; - let user_b = "quota-isolation-b"; - let limit_a = 50u64; - let limit_b = 80u64; - let bytes_a = Arc::new(AtomicU64::new(0)); - let bytes_b = Arc::new(AtomicU64::new(0)); - - let mut tasks = JoinSet::new(); - for idx in 0..200u64 { - let stats_ref = Arc::clone(&stats); - let bytes_ref = Arc::clone(&bytes_a); - tasks.spawn(async move { - let mut writer = make_crypto_writer(tokio::io::sink()); - let mut frame_buf = Vec::new(); - process_me_writer_response( - MeResponse::Data { - flags: 0, - data: Bytes::from_static(&[0xA1]), - }, - &mut writer, - ProtoTag::Intermediate, - &SecureRandom::new(), - &mut frame_buf, - stats_ref.as_ref(), - user_a, - Some(limit_a), - 0, - bytes_ref.as_ref(), - 8300 + idx, - false, - false, - ) - .await - }); - } - - for idx in 0..220u64 { - let stats_ref = Arc::clone(&stats); - let bytes_ref = Arc::clone(&bytes_b); - tasks.spawn(async move { - let mut writer = make_crypto_writer(tokio::io::sink()); - let mut frame_buf = Vec::new(); - process_me_writer_response( - MeResponse::Data { - flags: 0, - data: Bytes::from_static(&[0xB2]), - }, - &mut writer, - ProtoTag::Intermediate, - &SecureRandom::new(), - &mut frame_buf, - stats_ref.as_ref(), - user_b, - Some(limit_b), - 0, - bytes_ref.as_ref(), - 8500 + idx, - false, - false, - ) - .await - }); - } - - while let Some(joined) = tasks.join_next().await { - let result = joined.expect("quota isolation task must not panic"); - assert!(result.is_ok() || matches!(result, Err(ProxyError::DataQuotaExceeded { .. }))); - } - - assert_eq!(stats.get_user_total_octets(user_a), limit_a); - assert_eq!(stats.get_user_total_octets(user_b), limit_b); - assert_eq!(bytes_a.load(Ordering::Relaxed), limit_a); - assert_eq!(bytes_b.load(Ordering::Relaxed), limit_b); -} - -#[tokio::test] -async fn light_fuzz_mixed_me_responses_preserve_quota_and_counter_invariants() { - let stats = Stats::new(); - let user = "mixed-fuzz-user"; - let bytes_me2c = AtomicU64::new(0); - let quota_limit = 96u64; - let mut seed = 0xDEAD_BEEF_2026_0323u64; - - for idx in 0..2048u64 { - seed ^= seed << 7; - seed ^= seed >> 9; - seed ^= seed << 8; - - let choice = (seed & 0x03) as u8; - let response = if choice == 0 { - MeResponse::Ack((seed >> 8) as u32) - } else if choice == 1 { - MeResponse::Close - } else { - let len = ((seed >> 16) & 0x07) as usize; - let mut payload = vec![0u8; len]; - payload.fill((seed & 0xff) as u8); - MeResponse::Data { - flags: 0, - data: Bytes::from(payload), - } - }; - - let mut writer = make_crypto_writer(tokio::io::sink()); - let mut frame_buf = Vec::new(); - let result = process_me_writer_response( - response, - &mut writer, - ProtoTag::Intermediate, - &SecureRandom::new(), - &mut frame_buf, - &stats, - user, - Some(quota_limit), - 0, - &bytes_me2c, - 8800 + idx, - (idx & 1) == 0, - (idx & 2) == 0, - ) - .await; - - if let Err(err) = result { - assert!( - matches!(err, ProxyError::DataQuotaExceeded { .. }), - "mixed fuzz produced unexpected error variant: {err:?}" - ); - } - - let total = stats.get_user_total_octets(user); - assert!( - total <= quota_limit, - "mixed fuzz must keep usage at or below quota limit" - ); - assert_eq!(bytes_me2c.load(Ordering::Relaxed), total); - } -} \ No newline at end of file diff --git a/src/proxy/tests/middle_relay_quota_reservation_extreme_security_tests.rs b/src/proxy/tests/middle_relay_quota_reservation_extreme_security_tests.rs deleted file mode 100644 index e4d0c6e..0000000 --- a/src/proxy/tests/middle_relay_quota_reservation_extreme_security_tests.rs +++ /dev/null @@ -1,399 +0,0 @@ -use super::*; -use crate::crypto::{AesCtr, SecureRandom}; -use crate::stats::Stats; -use crate::stream::CryptoWriter; -use bytes::Bytes; -use std::sync::atomic::{AtomicU64, Ordering}; -use std::sync::{Arc, Mutex, OnceLock}; -use tokio::sync::Mutex as AsyncMutex; -use tokio::task::JoinSet; -use tokio::time::{Duration, timeout}; - -fn make_crypto_writer(writer: W) -> CryptoWriter -where - W: tokio::io::AsyncWrite + Unpin, -{ - let key = [0u8; 32]; - let iv = 0u128; - CryptoWriter::new(writer, AesCtr::new(&key, iv), 8 * 1024) -} - -fn lookup_counter_test_lock() -> &'static Mutex<()> { - static LOCK: OnceLock> = OnceLock::new(); - LOCK.get_or_init(|| Mutex::new(())) -} - -#[tokio::test] -async fn positive_prefetched_cross_mode_lock_multi_frame_accounting_is_exact() { - let _guard = lookup_counter_test_lock() - .lock() - .unwrap_or_else(|poison| poison.into_inner()); - - let stats = Stats::new(); - let user = format!("quota-extreme-positive-{}", std::process::id()); - let lock = crate::proxy::quota_lock_registry::cross_mode_quota_user_lock(&user); - crate::proxy::quota_lock_registry::reset_cross_mode_quota_user_lock_lookup_count_for_tests(); - - let mut writer = make_crypto_writer(tokio::io::sink()); - let mut frame_buf = Vec::new(); - let bytes_me2c = AtomicU64::new(0); - - for idx in 0..12u64 { - let payload = vec![0x5A; ((idx % 4) + 1) as usize]; - let result = process_me_writer_response_with_cross_mode_lock( - MeResponse::Data { - flags: 0, - data: Bytes::from(payload), - }, - &mut writer, - ProtoTag::Intermediate, - &SecureRandom::new(), - &mut frame_buf, - &stats, - &user, - Some(512), - 0, - Some(&lock), - &bytes_me2c, - 31_000 + idx, - false, - false, - ) - .await; - - assert!(result.is_ok()); - } - - assert_eq!( - crate::proxy::quota_lock_registry::cross_mode_quota_user_lock_lookup_count_for_user_for_tests(&user), - 0, - "prefetched lock path must avoid hot-path registry lookups" - ); - assert_eq!( - stats.get_user_total_octets(&user), - bytes_me2c.load(Ordering::Relaxed), - "forensics and quota accounting must remain synchronized" - ); -} - -#[tokio::test] -async fn negative_held_prefetched_lock_blocks_writer_without_accounting_mutation() { - let _guard = lookup_counter_test_lock() - .lock() - .unwrap_or_else(|poison| poison.into_inner()); - - let stats = Stats::new(); - let user = format!("quota-extreme-negative-{}", std::process::id()); - let lock = crate::proxy::quota_lock_registry::cross_mode_quota_user_lock(&user); - let held_guard = lock - .try_lock() - .expect("test must hold lock before calling ME->C writer"); - - let mut writer = make_crypto_writer(tokio::io::sink()); - let mut frame_buf = Vec::new(); - let bytes_me2c = AtomicU64::new(0); - - let blocked = timeout( - Duration::from_millis(25), - process_me_writer_response_with_cross_mode_lock( - MeResponse::Data { - flags: 0, - data: Bytes::from_static(&[1, 2, 3]), - }, - &mut writer, - ProtoTag::Intermediate, - &SecureRandom::new(), - &mut frame_buf, - &stats, - &user, - Some(64), - 0, - Some(&lock), - &bytes_me2c, - 31_100, - false, - false, - ), - ) - .await; - - assert!(blocked.is_err()); - assert_eq!(stats.get_user_total_octets(&user), 0); - assert_eq!(bytes_me2c.load(Ordering::Relaxed), 0); - - drop(held_guard); -} - -#[tokio::test] -async fn edge_zero_quota_and_zero_payload_is_fail_closed() { - let _guard = lookup_counter_test_lock() - .lock() - .unwrap_or_else(|poison| poison.into_inner()); - - let stats = Stats::new(); - let user = format!("quota-extreme-edge-{}", std::process::id()); - let lock = crate::proxy::quota_lock_registry::cross_mode_quota_user_lock(&user); - - let mut writer = make_crypto_writer(tokio::io::sink()); - let mut frame_buf = Vec::new(); - let bytes_me2c = AtomicU64::new(0); - - let result = process_me_writer_response_with_cross_mode_lock( - MeResponse::Data { - flags: 0, - data: Bytes::new(), - }, - &mut writer, - ProtoTag::Intermediate, - &SecureRandom::new(), - &mut frame_buf, - &stats, - &user, - Some(0), - 0, - Some(&lock), - &bytes_me2c, - 31_200, - false, - false, - ) - .await; - - assert!(matches!(result, Err(ProxyError::DataQuotaExceeded { .. }))); - assert_eq!(stats.get_user_total_octets(&user), 0); - assert_eq!(bytes_me2c.load(Ordering::Relaxed), 0); -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 4)] -async fn adversarial_blackhat_parallel_quota_race_never_overshoots_soft_cap() { - let _guard = lookup_counter_test_lock() - .lock() - .unwrap_or_else(|poison| poison.into_inner()); - - let stats = Arc::new(Stats::new()); - let user = format!("quota-extreme-blackhat-{}", std::process::id()); - let quota = 80u64; - let overshoot = 7u64; - let soft_limit = quota + overshoot; - let lock = Arc::new(crate::proxy::quota_lock_registry::cross_mode_quota_user_lock(&user)); - let bytes_me2c = Arc::new(AtomicU64::new(0)); - - let mut set = JoinSet::new(); - for idx in 0..256u64 { - let stats = Arc::clone(&stats); - let user = user.clone(); - let lock = Arc::clone(&lock); - let bytes_me2c = Arc::clone(&bytes_me2c); - - set.spawn(async move { - let mut writer = make_crypto_writer(tokio::io::sink()); - let mut frame_buf = Vec::new(); - let len = ((idx % 5) + 1) as usize; - let payload = vec![0xAA; len]; - - process_me_writer_response_with_cross_mode_lock( - MeResponse::Data { - flags: 0, - data: Bytes::from(payload), - }, - &mut writer, - ProtoTag::Intermediate, - &SecureRandom::new(), - &mut frame_buf, - stats.as_ref(), - &user, - Some(quota), - overshoot, - Some(&lock), - bytes_me2c.as_ref(), - 31_300 + idx, - false, - false, - ) - .await - }); - } - - while let Some(done) = set.join_next().await { - match done.expect("task must not panic") { - Ok(_) | Err(ProxyError::DataQuotaExceeded { .. }) => {} - Err(other) => panic!("unexpected error variant under black-hat race: {other:?}"), - } - } - - let total = stats.get_user_total_octets(&user); - assert!( - total <= soft_limit, - "parallel adversarial race must stay under soft cap" - ); - assert_eq!(bytes_me2c.load(Ordering::Relaxed), total); -} - -#[tokio::test] -async fn integration_without_prefetched_lock_uses_registry_lookup_path() { - let _guard = lookup_counter_test_lock() - .lock() - .unwrap_or_else(|poison| poison.into_inner()); - - let stats = Stats::new(); - let user = format!("quota-extreme-integration-{}", std::process::id()); - crate::proxy::quota_lock_registry::reset_cross_mode_quota_user_lock_lookup_count_for_tests(); - - let mut writer = make_crypto_writer(tokio::io::sink()); - let mut frame_buf = Vec::new(); - let bytes_me2c = AtomicU64::new(0); - - for idx in 0..3u64 { - let result = process_me_writer_response_with_cross_mode_lock( - MeResponse::Data { - flags: 0, - data: Bytes::from_static(&[0x41]), - }, - &mut writer, - ProtoTag::Intermediate, - &SecureRandom::new(), - &mut frame_buf, - &stats, - &user, - Some(16), - 0, - None, - &bytes_me2c, - 31_400 + idx, - false, - false, - ) - .await; - - assert!(result.is_ok()); - } - - assert_eq!( - crate::proxy::quota_lock_registry::cross_mode_quota_user_lock_lookup_count_for_user_for_tests(&user), - 3, - "control path should perform one lock-registry lookup per call" - ); -} - -#[tokio::test] -async fn light_fuzz_quota_matrix_preserves_fail_closed_accounting() { - let _guard = lookup_counter_test_lock() - .lock() - .unwrap_or_else(|poison| poison.into_inner()); - - let stats = Stats::new(); - let user = format!("quota-extreme-fuzz-{}", std::process::id()); - let lock = crate::proxy::quota_lock_registry::cross_mode_quota_user_lock(&user); - let bytes_me2c = AtomicU64::new(0); - let mut seed = 0xA11C_55EE_2026_0323u64; - - for idx in 0..512u64 { - seed ^= seed << 7; - seed ^= seed >> 9; - seed ^= seed << 8; - - let quota = 24 + (seed & 0x3f); - let overshoot = (seed >> 13) & 0x0f; - let len = ((seed >> 19) & 0x07) + 1; - - let mut writer = make_crypto_writer(tokio::io::sink()); - let mut frame_buf = Vec::new(); - let before = stats.get_user_total_octets(&user); - - let result = process_me_writer_response_with_cross_mode_lock( - MeResponse::Data { - flags: 0, - data: Bytes::from(vec![0x11; len as usize]), - }, - &mut writer, - ProtoTag::Intermediate, - &SecureRandom::new(), - &mut frame_buf, - &stats, - &user, - Some(quota), - overshoot, - Some(&lock), - &bytes_me2c, - 31_500 + idx, - false, - false, - ) - .await; - - let after = stats.get_user_total_octets(&user); - if result.is_ok() { - assert!(after >= before); - } else { - assert!(matches!(result, Err(ProxyError::DataQuotaExceeded { .. }))); - assert_eq!(after, before); - } - assert_eq!(bytes_me2c.load(Ordering::Relaxed), after); - } -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 4)] -async fn stress_prefetched_lock_high_fanout_exact_quota_success_count() { - let _guard = lookup_counter_test_lock() - .lock() - .unwrap_or_else(|poison| poison.into_inner()); - - let stats = Arc::new(Stats::new()); - let user = format!("quota-extreme-stress-{}", std::process::id()); - let quota = 96u64; - let lock: Arc> = crate::proxy::quota_lock_registry::cross_mode_quota_user_lock(&user); - let bytes_me2c = Arc::new(AtomicU64::new(0)); - - crate::proxy::quota_lock_registry::reset_cross_mode_quota_user_lock_lookup_count_for_tests(); - - let mut set = JoinSet::new(); - for idx in 0..384u64 { - let stats = Arc::clone(&stats); - let user = user.clone(); - let lock = Arc::clone(&lock); - let bytes_me2c = Arc::clone(&bytes_me2c); - - set.spawn(async move { - let mut writer = make_crypto_writer(tokio::io::sink()); - let mut frame_buf = Vec::new(); - process_me_writer_response_with_cross_mode_lock( - MeResponse::Data { - flags: 0, - data: Bytes::from_static(&[0xFF]), - }, - &mut writer, - ProtoTag::Intermediate, - &SecureRandom::new(), - &mut frame_buf, - stats.as_ref(), - &user, - Some(quota), - 0, - Some(&lock), - bytes_me2c.as_ref(), - 31_600 + idx, - false, - false, - ) - .await - }); - } - - let mut success = 0usize; - while let Some(done) = set.join_next().await { - match done.expect("task must not panic") { - Ok(_) => success += 1, - Err(ProxyError::DataQuotaExceeded { .. }) => {} - Err(other) => panic!("unexpected error variant in stress fanout: {other:?}"), - } - } - - assert_eq!(success, quota as usize); - assert_eq!(stats.get_user_total_octets(&user), quota); - assert_eq!(bytes_me2c.load(Ordering::Relaxed), quota); - assert_eq!( - crate::proxy::quota_lock_registry::cross_mode_quota_user_lock_lookup_count_for_user_for_tests(&user), - 0, - "stress prefetched path must not use lock registry lookups" - ); -} diff --git a/src/proxy/tests/middle_relay_security_tests.rs b/src/proxy/tests/middle_relay_security_tests.rs deleted file mode 100644 index 1d3b736..0000000 --- a/src/proxy/tests/middle_relay_security_tests.rs +++ /dev/null @@ -1,2517 +0,0 @@ -use super::*; -use crate::config::{GeneralConfig, MeRouteNoWriterMode, MeSocksKdfPolicy, MeWriterPickMode}; -use crate::crypto::AesCtr; -use crate::crypto::SecureRandom; -use crate::network::probe::NetworkDecision; -use crate::proxy::handshake::HandshakeSuccess; -use crate::proxy::route_mode::{RelayRouteMode, RouteRuntimeController}; -use crate::stats::Stats; -use crate::stream::{BufferPool, CryptoReader, CryptoWriter, PooledBuffer}; -use crate::transport::middle_proxy::MePool; -use bytes::Bytes; -use rand::rngs::StdRng; -use rand::{RngExt, SeedableRng}; -use std::collections::{HashMap, HashSet}; -use std::net::SocketAddr; -use std::sync::Arc; -use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; -use std::sync::Mutex; -use std::thread; -use tokio::io::AsyncReadExt; -use tokio::io::AsyncWriteExt; -use tokio::io::duplex; -use tokio::sync::Barrier; -use tokio::time::{Duration as TokioDuration, timeout}; - -fn make_pooled_payload(data: &[u8]) -> PooledBuffer { - let pool = Arc::new(BufferPool::with_config(data.len().max(1), 4)); - let mut payload = pool.get(); - payload.resize(data.len(), 0); - payload[..data.len()].copy_from_slice(data); - payload -} - -fn make_pooled_payload_from(pool: &Arc, data: &[u8]) -> PooledBuffer { - let mut payload = pool.get(); - payload.resize(data.len(), 0); - payload[..data.len()].copy_from_slice(data); - payload -} - -#[test] -fn should_yield_sender_only_on_budget_with_backlog() { - assert!(!should_yield_c2me_sender(0, true)); - assert!(!should_yield_c2me_sender( - C2ME_SENDER_FAIRNESS_BUDGET - 1, - true - )); - assert!(!should_yield_c2me_sender( - C2ME_SENDER_FAIRNESS_BUDGET, - false - )); - assert!(should_yield_c2me_sender(C2ME_SENDER_FAIRNESS_BUDGET, true)); -} - -#[tokio::test] -async fn enqueue_c2me_command_uses_try_send_fast_path() { - let (tx, mut rx) = mpsc::channel::(2); - enqueue_c2me_command( - &tx, - C2MeCommand::Data { - payload: make_pooled_payload(&[1, 2, 3]), - flags: 0, - }, - ) - .await - .unwrap(); - - let recv = timeout(TokioDuration::from_millis(50), rx.recv()) - .await - .unwrap() - .unwrap(); - match recv { - C2MeCommand::Data { payload, flags } => { - assert_eq!(payload.as_ref(), &[1, 2, 3]); - assert_eq!(flags, 0); - } - C2MeCommand::Close => panic!("unexpected close command"), - } -} - -#[tokio::test] -async fn enqueue_c2me_command_falls_back_to_send_when_queue_is_full() { - let (tx, mut rx) = mpsc::channel::(1); - tx.send(C2MeCommand::Data { - payload: make_pooled_payload(&[9]), - flags: 9, - }) - .await - .unwrap(); - - let tx2 = tx.clone(); - let producer = tokio::spawn(async move { - enqueue_c2me_command( - &tx2, - C2MeCommand::Data { - payload: make_pooled_payload(&[7, 7]), - flags: 7, - }, - ) - .await - .unwrap(); - }); - - let _ = timeout(TokioDuration::from_millis(100), rx.recv()) - .await - .unwrap(); - producer.await.unwrap(); - - let recv = timeout(TokioDuration::from_millis(100), rx.recv()) - .await - .unwrap() - .unwrap(); - match recv { - C2MeCommand::Data { payload, flags } => { - assert_eq!(payload.as_ref(), &[7, 7]); - assert_eq!(flags, 7); - } - C2MeCommand::Close => panic!("unexpected close command"), - } -} - -#[tokio::test] -async fn enqueue_c2me_command_closed_channel_recycles_payload() { - let pool = Arc::new(BufferPool::with_config(64, 4)); - let payload = make_pooled_payload_from(&pool, &[1, 2, 3, 4]); - let (tx, rx) = mpsc::channel::(1); - drop(rx); - - let result = enqueue_c2me_command(&tx, C2MeCommand::Data { payload, flags: 0 }).await; - - assert!(result.is_err(), "closed queue must fail enqueue"); - drop(result); - assert!( - pool.stats().pooled >= 1, - "payload must return to pool when enqueue fails on closed channel" - ); -} - -#[tokio::test] -async fn enqueue_c2me_command_full_then_closed_recycles_waiting_payload() { - let pool = Arc::new(BufferPool::with_config(64, 4)); - let (tx, rx) = mpsc::channel::(1); - - tx.send(C2MeCommand::Data { - payload: make_pooled_payload_from(&pool, &[9]), - flags: 1, - }) - .await - .unwrap(); - - let tx2 = tx.clone(); - let pool2 = pool.clone(); - let blocked_send = tokio::spawn(async move { - enqueue_c2me_command( - &tx2, - C2MeCommand::Data { - payload: make_pooled_payload_from(&pool2, &[7, 7, 7]), - flags: 2, - }, - ) - .await - }); - - tokio::time::sleep(TokioDuration::from_millis(10)).await; - drop(rx); - - let result = timeout(TokioDuration::from_secs(1), blocked_send) - .await - .expect("blocked send task must finish") - .expect("blocked send task must not panic"); - - assert!( - result.is_err(), - "closing receiver while sender is blocked must fail enqueue" - ); - drop(result); - assert!( - pool.stats().pooled >= 2, - "both queued and blocked payloads must return to pool after channel close" - ); -} - -#[tokio::test] -async fn enqueue_c2me_command_full_queue_times_out_without_receiver_progress() { - let (tx, _rx) = mpsc::channel::(1); - tx.send(C2MeCommand::Data { - payload: make_pooled_payload(&[1]), - flags: 0, - }) - .await - .unwrap(); - - let started = Instant::now(); - let result = enqueue_c2me_command( - &tx, - C2MeCommand::Data { - payload: make_pooled_payload(&[2, 2]), - flags: 1, - }, - ) - .await; - - assert!( - result.is_err(), - "enqueue must fail when queue stays full beyond bounded timeout" - ); - assert!( - started.elapsed() < TokioDuration::from_millis(400), - "full-queue timeout must resolve promptly" - ); -} - -#[test] -fn desync_dedup_cache_is_bounded() { - let _guard = desync_dedup_test_lock() - .lock() - .expect("desync dedup test lock must be available"); - clear_desync_dedup_for_testing(); - - let now = Instant::now(); - for key in 0..DESYNC_DEDUP_MAX_ENTRIES as u64 { - assert!( - should_emit_full_desync(key, false, now), - "unique keys up to cap must be tracked" - ); - } - - assert!( - should_emit_full_desync(u64::MAX, false, now), - "new key above cap must emit once after bounded eviction for forensic visibility" - ); - - assert!( - !should_emit_full_desync(u64::MAX, false, now), - "already tracked key inside dedup window must stay suppressed" - ); -} - -#[test] -fn quota_user_lock_cache_reuses_entry_for_same_user() { - let _guard = super::quota_user_lock_test_scope(); - - let map = QUOTA_USER_LOCKS.get_or_init(DashMap::new); - map.clear(); - - let a = quota_user_lock("quota-user-a"); - let b = quota_user_lock("quota-user-a"); - assert!(Arc::ptr_eq(&a, &b), "same user must reuse same quota lock"); -} - -#[test] -fn quota_user_lock_cache_is_bounded_under_unique_churn() { - let _guard = super::quota_user_lock_test_scope(); - - let map = QUOTA_USER_LOCKS.get_or_init(DashMap::new); - map.clear(); - - for idx in 0..(QUOTA_USER_LOCKS_MAX + 128) { - let user = format!("quota-user-{idx}"); - let lock = quota_user_lock(&user); - drop(lock); - } - - assert!( - map.len() <= QUOTA_USER_LOCKS_MAX, - "quota lock cache must stay within configured bound" - ); -} - -#[test] -fn quota_user_lock_cache_saturation_returns_stable_overflow_lock_without_growth() { - let _guard = super::quota_user_lock_test_scope(); - - let map = QUOTA_USER_LOCKS.get_or_init(DashMap::new); - for attempt in 0..8u32 { - map.clear(); - - let prefix = format!("quota-held-user-{}-{attempt}", std::process::id()); - let mut retained = Vec::with_capacity(QUOTA_USER_LOCKS_MAX); - for idx in 0..QUOTA_USER_LOCKS_MAX { - let user = format!("{prefix}-{idx}"); - retained.push(quota_user_lock(&user)); - } - - if map.len() != QUOTA_USER_LOCKS_MAX { - drop(retained); - continue; - } - - let overflow_user = format!("quota-overflow-user-{}-{attempt}", std::process::id()); - let overflow_a = quota_user_lock(&overflow_user); - let overflow_b = quota_user_lock(&overflow_user); - - assert_eq!( - map.len(), - QUOTA_USER_LOCKS_MAX, - "overflow acquisition must not grow cache past hard limit" - ); - assert!( - map.get(&overflow_user).is_none(), - "overflow path should not cache new user lock when map is saturated and all entries are retained" - ); - assert!( - Arc::ptr_eq(&overflow_a, &overflow_b), - "overflow user lock should use deterministic striping under saturation" - ); - - drop(retained); - return; - } - - panic!("unable to observe stable saturated lock-cache precondition after bounded retries"); -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 4)] -async fn adversarial_quota_race_under_lock_cache_saturation_still_allows_only_one_winner() { - let map = QUOTA_USER_LOCKS.get_or_init(DashMap::new); - map.clear(); - - let mut retained = Vec::with_capacity(QUOTA_USER_LOCKS_MAX); - for idx in 0..QUOTA_USER_LOCKS_MAX { - let user = format!("quota-saturated-user-{idx}"); - retained.push(quota_user_lock(&user)); - } - - assert_eq!( - map.len(), - QUOTA_USER_LOCKS_MAX, - "precondition: cache must be saturated for overflow-user race test" - ); - - let stats = Stats::new(); - let bytes_me2c = AtomicU64::new(0); - let user = "gap-t04-saturated-lock-race-user"; - let barrier = Arc::new(Barrier::new(2)); - - let one = run_quota_race_attempt(&stats, &bytes_me2c, user, 0x55, 9101, barrier.clone()); - let two = run_quota_race_attempt(&stats, &bytes_me2c, user, 0x66, 9102, barrier); - let (r1, r2) = tokio::join!(one, two); - - assert!( - matches!(r1, Ok(_) | Err(ProxyError::DataQuotaExceeded { .. })) - && matches!(r2, Ok(_) | Err(ProxyError::DataQuotaExceeded { .. })), - "both racers must resolve cleanly without unexpected errors" - ); - assert!( - matches!(r1, Err(ProxyError::DataQuotaExceeded { .. })) - || matches!(r2, Err(ProxyError::DataQuotaExceeded { .. })), - "at least one racer must be quota-rejected even when lock cache is saturated" - ); - assert_eq!( - stats.get_user_total_octets(user), - 1, - "saturated lock cache must not permit double-success quota overshoot" - ); - - drop(retained); -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 4)] -async fn stress_quota_race_under_lock_cache_saturation_never_allows_double_success() { - let map = QUOTA_USER_LOCKS.get_or_init(DashMap::new); - map.clear(); - - let mut retained = Vec::with_capacity(QUOTA_USER_LOCKS_MAX); - for idx in 0..QUOTA_USER_LOCKS_MAX { - let user = format!("quota-saturated-stress-holder-{idx}"); - retained.push(quota_user_lock(&user)); - } - - let stats = Stats::new(); - let bytes_me2c = AtomicU64::new(0); - - for round in 0..128u64 { - let user = format!("gap-t04-saturated-race-round-{round}"); - let barrier = Arc::new(Barrier::new(2)); - - let one = run_quota_race_attempt( - &stats, - &bytes_me2c, - &user, - 0x71, - 12_000 + round, - barrier.clone(), - ); - let two = run_quota_race_attempt(&stats, &bytes_me2c, &user, 0x72, 13_000 + round, barrier); - - let (r1, r2) = tokio::join!(one, two); - assert!( - matches!(r1, Ok(_) | Err(ProxyError::DataQuotaExceeded { .. })) - && matches!(r2, Ok(_) | Err(ProxyError::DataQuotaExceeded { .. })), - "round {round}: racers must resolve cleanly" - ); - assert!( - matches!(r1, Err(ProxyError::DataQuotaExceeded { .. })) - || matches!(r2, Err(ProxyError::DataQuotaExceeded { .. })), - "round {round}: at least one racer must be quota-rejected" - ); - assert_eq!( - stats.get_user_total_octets(&user), - 1, - "round {round}: saturated cache must still enforce exactly one forwarded byte" - ); - } - - drop(retained); -} - -#[test] -fn adversarial_forensics_trace_id_should_not_alias_conn_id() { - let now = Instant::now(); - let trace_id = 0x1122_3344_5566_7788; - let conn_id = 0x8877_6655_4433_2211; - let state = RelayForensicsState { - trace_id, - conn_id, - user: "trace-user".to_string(), - peer: "198.51.100.17:443".parse().unwrap(), - peer_hash: 0x8877_6655_4433_2211, - started_at: now, - bytes_c2me: 0, - bytes_me2c: Arc::new(AtomicU64::new(0)), - desync_all_full: false, - }; - - assert_ne!( - state.trace_id, state.conn_id, - "security expectation: trace correlation should be independent of connection identity" - ); - assert_eq!(state.trace_id, trace_id); - assert_eq!(state.conn_id, conn_id); -} - -#[tokio::test] -async fn abridged_ack_uses_big_endian_confirm_bytes_after_decryption() { - let (mut writer_side, reader_side) = duplex(8); - let key = [0u8; 32]; - let iv = 0u128; - let mut writer = CryptoWriter::new(reader_side, AesCtr::new(&key, iv), 8 * 1024); - - write_client_ack(&mut writer, ProtoTag::Abridged, 0x11_22_33_44) - .await - .expect("ack write must succeed"); - - let mut observed = [0u8; 4]; - writer_side - .read_exact(&mut observed) - .await - .expect("ack bytes must be readable"); - let mut decryptor = AesCtr::new(&key, iv); - let decrypted = decryptor.decrypt(&observed); - - assert_eq!( - decrypted, - 0x11_22_33_44u32.to_be_bytes(), - "abridged ACK should encode confirm bytes in big-endian order" - ); -} - -#[test] -fn desync_dedup_full_cache_churn_stays_suppressed() { - let _guard = desync_dedup_test_lock() - .lock() - .expect("desync dedup test lock must be available"); - clear_desync_dedup_for_testing(); - - let now = Instant::now(); - for key in 0..DESYNC_DEDUP_MAX_ENTRIES as u64 { - assert!(should_emit_full_desync(key, false, now)); - } - - for offset in 0..2048u64 { - let emitted = should_emit_full_desync(u64::MAX - offset, false, now); - if offset == 0 { - assert!( - emitted, - "first full-cache newcomer should emit for forensic visibility" - ); - } else { - assert!( - !emitted, - "full-cache newcomer churn inside emit interval must stay suppressed" - ); - } - } -} - -#[test] -fn dedup_hash_is_stable_for_same_input_within_process() { - let sample = ( - "scope_user", - hash_ip("198.51.100.7".parse().unwrap()), - ProtoTag::Secure, - ); - let first = hash_value(&sample); - let second = hash_value(&sample); - assert_eq!( - first, second, - "dedup hash must be stable within a process for cache lookups" - ); -} - -#[test] -fn dedup_hash_resists_simple_collision_bursts_for_peer_ip_space() { - let mut seen = HashSet::new(); - - for octet in 1u16..=2048 { - let third = ((octet / 256) & 0xff) as u8; - let fourth = (octet & 0xff) as u8; - let ip = IpAddr::V4(std::net::Ipv4Addr::new(198, 51, third, fourth)); - let key = hash_value(&( - "scope_user", - hash_ip(ip), - ProtoTag::Secure, - DESYNC_ERROR_CLASS, - )); - seen.insert(key); - } - - assert_eq!( - seen.len(), - 2048, - "adversarial peer-IP burst should not collapse dedup keys via trivial collisions" - ); -} - -#[test] -fn light_fuzz_dedup_hash_collision_rate_stays_negligible() { - let mut rng = StdRng::seed_from_u64(0x9E37_79B9_A1B2_C3D4); - let mut seen = HashSet::new(); - let samples = 8192usize; - - for _ in 0..samples { - let user_seed: u64 = rng.random(); - let peer_seed: u64 = rng.random(); - let proto = if (peer_seed & 1) == 0 { - ProtoTag::Secure - } else { - ProtoTag::Intermediate - }; - let key = hash_value(&(user_seed, peer_seed, proto, DESYNC_ERROR_CLASS)); - seen.insert(key); - } - - let collisions = samples - seen.len(); - assert!( - collisions <= 1, - "light fuzz collision count should remain negligible for 64-bit dedup keys" - ); -} - -#[test] -fn stress_desync_dedup_churn_keeps_cache_hard_bounded() { - let _guard = desync_dedup_test_lock() - .lock() - .expect("desync dedup test lock must be available"); - clear_desync_dedup_for_testing(); - - let now = Instant::now(); - let total = DESYNC_DEDUP_MAX_ENTRIES + 8192; - - let mut emitted_count = 0usize; - for key in 0..total as u64 { - let emitted = should_emit_full_desync(key, false, now); - if emitted { - emitted_count += 1; - } - } - - assert_eq!( - emitted_count, - DESYNC_DEDUP_MAX_ENTRIES + 1, - "after capacity is reached, same-tick newcomer churn must be rate-limited" - ); - - let len = DESYNC_DEDUP - .get() - .expect("dedup cache must be initialized by stress run") - .len(); - assert!( - len <= DESYNC_DEDUP_MAX_ENTRIES, - "dedup cache must stay bounded under stress churn" - ); -} - -#[test] -fn full_cache_newcomer_emission_is_rate_limited_but_periodic() { - let _guard = desync_dedup_test_lock() - .lock() - .expect("desync dedup test lock must be available"); - clear_desync_dedup_for_testing(); - - let dedup = DESYNC_DEDUP.get_or_init(DashMap::new); - let base_now = Instant::now(); - - for key in 0..DESYNC_DEDUP_MAX_ENTRIES as u64 { - dedup.insert(key, base_now - TokioDuration::from_millis(10)); - } - - // Same-tick newcomer storm: only the first should emit full forensic record. - let mut burst_emits = 0usize; - for i in 0..1024u64 { - if should_emit_full_desync(10_000_000 + i, false, base_now) { - burst_emits += 1; - } - } - assert_eq!( - burst_emits, 1, - "full-cache newcomer burst must be bounded to a single full emit per interval" - ); - - // After each interval elapses, one newcomer may emit again. - for step in 1..=6u64 { - let t = base_now + DESYNC_FULL_CACHE_EMIT_MIN_INTERVAL * step as u32; - assert!( - should_emit_full_desync(20_000_000 + step, false, t), - "full-cache newcomer should re-emit once interval has elapsed" - ); - assert!( - !should_emit_full_desync(30_000_000 + step, false, t), - "additional newcomers in the same interval tick must remain suppressed" - ); - } -} - -#[test] -fn full_cache_mode_override_emits_every_event() { - let _guard = desync_dedup_test_lock() - .lock() - .expect("desync dedup test lock must be available"); - clear_desync_dedup_for_testing(); - - let now = Instant::now(); - for i in 0..10_000u64 { - assert!( - should_emit_full_desync(100_000_000 + i, true, now), - "desync_all_full override must bypass dedup and rate-limit suppression" - ); - } -} - -#[test] -fn report_desync_stats_follow_rate_limited_full_cache_policy() { - let _guard = desync_dedup_test_lock() - .lock() - .expect("desync dedup test lock must be available"); - clear_desync_dedup_for_testing(); - - let dedup = DESYNC_DEDUP.get_or_init(DashMap::new); - let base_now = Instant::now(); - for key in 0..DESYNC_DEDUP_MAX_ENTRIES as u64 { - dedup.insert(key, base_now - TokioDuration::from_millis(10)); - } - - let stats = Stats::new(); - let mut state = make_forensics_state(); - state.started_at = base_now; - - for i in 0..128u64 { - state.peer_hash = 0xABC0_0000_0000_0000u64 ^ i; - let _ = report_desync_frame_too_large( - &state, - ProtoTag::Secure, - 3, - 1024, - 4096, - Some([0x16, 0x03, 0x03, 0x00]), - &stats, - ); - } - - assert_eq!( - stats.get_desync_total(), - 128, - "every detected desync must increment total counter" - ); - assert_eq!( - stats.get_desync_full_logged(), - 1, - "same-interval full-cache newcomer storm must allow only one full forensic emit" - ); - assert_eq!( - stats.get_desync_suppressed(), - 127, - "remaining same-interval full-cache newcomer events must be suppressed" - ); - - // After one full interval in real wall clock, a newcomer should emit again. - thread::sleep(DESYNC_FULL_CACHE_EMIT_MIN_INTERVAL + TokioDuration::from_millis(20)); - state.peer_hash = 0xDEAD_BEEF_DEAD_BEEFu64; - let _ = report_desync_frame_too_large( - &state, - ProtoTag::Secure, - 4, - 1024, - 4097, - Some([0x16, 0x03, 0x03, 0x01]), - &stats, - ); - - assert_eq!( - stats.get_desync_full_logged(), - 2, - "full forensic emission must recover after rate-limit interval" - ); -} - -#[test] -fn concurrent_full_cache_newcomer_storm_is_single_emit_per_interval() { - let _guard = desync_dedup_test_lock() - .lock() - .expect("desync dedup test lock must be available"); - clear_desync_dedup_for_testing(); - - let dedup = DESYNC_DEDUP.get_or_init(DashMap::new); - let base_now = Instant::now(); - for key in 0..DESYNC_DEDUP_MAX_ENTRIES as u64 { - dedup.insert(key, base_now - TokioDuration::from_millis(10)); - } - - let emits = Arc::new(AtomicUsize::new(0)); - let mut workers = Vec::new(); - for worker_id in 0..32u64 { - let emits = Arc::clone(&emits); - workers.push(thread::spawn(move || { - for i in 0..512u64 { - let key = 0x7000_0000_0000_0000u64 ^ (worker_id << 20) ^ i; - if should_emit_full_desync(key, false, base_now) { - emits.fetch_add(1, Ordering::Relaxed); - } - } - })); - } - - for worker in workers { - worker.join().expect("worker thread must not panic"); - } - - assert_eq!( - emits.load(Ordering::Relaxed), - 1, - "concurrent same-interval full-cache storm must allow only one full forensic emit" - ); -} - -#[test] -fn light_fuzz_full_cache_rate_limit_oracle_matches_model() { - let _guard = desync_dedup_test_lock() - .lock() - .expect("desync dedup test lock must be available"); - clear_desync_dedup_for_testing(); - - let dedup = DESYNC_DEDUP.get_or_init(DashMap::new); - let base_now = Instant::now(); - for key in 0..DESYNC_DEDUP_MAX_ENTRIES as u64 { - dedup.insert(key, base_now - TokioDuration::from_millis(10)); - } - - let mut rng = StdRng::seed_from_u64(0xD15EA5E5_F00DBAAD); - let mut model_last_emit: Option = None; - - for i in 0..4096u64 { - let jitter_ms: u64 = rng.random_range(0..=3000); - let t = base_now + TokioDuration::from_millis(jitter_ms); - let key = 0x55AA_0000_0000_0000u64 ^ i ^ rng.random::(); - let actual = should_emit_full_desync(key, false, t); - - let expected = match model_last_emit { - None => { - model_last_emit = Some(t); - true - } - Some(last) => { - match t.checked_duration_since(last) { - Some(elapsed) if elapsed >= DESYNC_FULL_CACHE_EMIT_MIN_INTERVAL => { - model_last_emit = Some(t); - true - } - Some(_) => false, - None => { - // Match production fail-open behavior for non-monotonic synthetic input. - model_last_emit = Some(t); - true - } - } - } - }; - - assert_eq!( - actual, expected, - "full-cache rate-limit gate diverged from reference model under light fuzz" - ); - } -} - -#[test] -fn full_cache_gate_lock_poison_is_fail_closed_without_panic() { - let _guard = desync_dedup_test_lock() - .lock() - .expect("desync dedup test lock must be available"); - clear_desync_dedup_for_testing(); - - let dedup = DESYNC_DEDUP.get_or_init(DashMap::new); - let base_now = Instant::now(); - for key in 0..DESYNC_DEDUP_MAX_ENTRIES as u64 { - dedup.insert(key, base_now - TokioDuration::from_millis(10)); - } - - // Poison the full-cache gate lock intentionally. - let gate = DESYNC_FULL_CACHE_LAST_EMIT_AT.get_or_init(|| Mutex::new(None)); - let _ = std::panic::catch_unwind(|| { - let _lock = gate - .lock() - .expect("gate lock must be lockable before poison"); - panic!("intentional gate poison for fail-closed regression"); - }); - - let emitted = should_emit_full_desync(0xFACE_0000_0000_0001, false, base_now); - assert!( - !emitted, - "poisoned full-cache gate must fail-closed (suppress) instead of panic or fail-open" - ); - assert!( - dedup.len() <= DESYNC_DEDUP_MAX_ENTRIES, - "dedup cache must remain bounded even when gate lock is poisoned" - ); -} - -#[test] -fn full_cache_non_monotonic_time_emits_and_resets_gate_safely() { - let _guard = desync_dedup_test_lock() - .lock() - .expect("desync dedup test lock must be available"); - clear_desync_dedup_for_testing(); - - let dedup = DESYNC_DEDUP.get_or_init(DashMap::new); - let base_now = Instant::now(); - for key in 0..DESYNC_DEDUP_MAX_ENTRIES as u64 { - dedup.insert(key, base_now - TokioDuration::from_millis(10)); - } - - // First event seeds the gate. - assert!(should_emit_full_desync( - 0xABCD_0000_0000_0001, - false, - base_now + TokioDuration::from_millis(900) - )); - - // Synthetic earlier timestamp must not panic; it should fail-open and reset gate. - assert!(should_emit_full_desync( - 0xABCD_0000_0000_0002, - false, - base_now + TokioDuration::from_millis(100) - )); - - // Same instant again remains suppressed after reset. - assert!(!should_emit_full_desync( - 0xABCD_0000_0000_0003, - false, - base_now + TokioDuration::from_millis(100) - )); -} - -#[test] -fn desync_dedup_full_cache_inserts_new_key_with_bounded_single_key_churn() { - let _guard = desync_dedup_test_lock() - .lock() - .expect("desync dedup test lock must be available"); - clear_desync_dedup_for_testing(); - - let dedup = DESYNC_DEDUP.get_or_init(DashMap::new); - let base_now = Instant::now(); - - // Fill with fresh entries so stale-pruning does not apply. - for key in 0..DESYNC_DEDUP_MAX_ENTRIES as u64 { - dedup.insert(key, base_now - TokioDuration::from_millis(10)); - } - - let before_keys: std::collections::HashSet = dedup.iter().map(|e| *e.key()).collect(); - - let newcomer_key = u64::MAX; - let emitted = should_emit_full_desync(newcomer_key, false, base_now); - assert!( - emitted, - "new entry under full fresh cache must emit after bounded eviction" - ); - assert!( - dedup.get(&newcomer_key).is_some(), - "new key must be inserted after bounded eviction" - ); - - let after_keys: std::collections::HashSet = dedup.iter().map(|e| *e.key()).collect(); - let removed_count = before_keys.difference(&after_keys).count(); - let added_count = after_keys.difference(&before_keys).count(); - - assert_eq!( - removed_count, 1, - "full-cache insertion must evict exactly one prior key" - ); - assert_eq!( - added_count, 1, - "full-cache insertion must add exactly one newcomer key" - ); - assert!( - dedup.len() <= DESYNC_DEDUP_MAX_ENTRIES, - "dedup cache must remain hard-bounded after full-cache churn" - ); -} - -#[test] -fn light_fuzz_desync_dedup_temporal_gate_behavior_is_stable() { - let _guard = desync_dedup_test_lock() - .lock() - .expect("desync dedup test lock must be available"); - clear_desync_dedup_for_testing(); - - let key = 0xC0DE_CAFE_u64; - let start = Instant::now(); - - assert!( - should_emit_full_desync(key, false, start), - "first event for key must emit full forensic record" - ); - - // Deterministic pseudo-random time deltas around dedup window edge. - let mut s: u64 = 0x1234_5678_9ABC_DEF0; - for _ in 0..2048 { - s ^= s << 7; - s ^= s >> 9; - s ^= s << 8; - - let delta_ms = s % (DESYNC_DEDUP_WINDOW.as_millis() as u64 * 2 + 1); - let now = start + TokioDuration::from_millis(delta_ms); - let emitted = should_emit_full_desync(key, false, now); - - if delta_ms < DESYNC_DEDUP_WINDOW.as_millis() as u64 { - assert!( - !emitted, - "events inside dedup window must remain suppressed" - ); - } else { - // Once window elapsed for this key, at least one sample should re-emit and refresh. - if emitted { - return; - } - } - } - - panic!("expected at least one post-window sample to re-emit forensic record"); -} - -fn make_forensics_state() -> RelayForensicsState { - RelayForensicsState { - trace_id: 1, - conn_id: 2, - user: "test-user".to_string(), - peer: "127.0.0.1:50000".parse::().unwrap(), - peer_hash: 3, - started_at: Instant::now(), - bytes_c2me: 0, - bytes_me2c: Arc::new(AtomicU64::new(0)), - desync_all_full: false, - } -} - -fn make_crypto_reader(reader: R) -> CryptoReader -where - R: tokio::io::AsyncRead + Unpin, -{ - let key = [0u8; 32]; - let iv = 0u128; - CryptoReader::new(reader, AesCtr::new(&key, iv)) -} - -fn make_crypto_writer(writer: W) -> CryptoWriter -where - W: tokio::io::AsyncWrite + Unpin, -{ - let key = [0u8; 32]; - let iv = 0u128; - CryptoWriter::new(writer, AesCtr::new(&key, iv), 8 * 1024) -} - -async fn make_me_pool_for_abort_test(stats: Arc) -> Arc { - let general = GeneralConfig::default(); - - MePool::new( - None, - vec![1u8; 32], - None, - false, - None, - Vec::new(), - 1, - None, - 12, - 1200, - HashMap::new(), - HashMap::new(), - None, - NetworkDecision::default(), - None, - Arc::new(SecureRandom::new()), - stats, - general.me_keepalive_enabled, - general.me_keepalive_interval_secs, - general.me_keepalive_jitter_secs, - general.me_keepalive_payload_random, - general.rpc_proxy_req_every, - general.me_warmup_stagger_enabled, - general.me_warmup_step_delay_ms, - general.me_warmup_step_jitter_ms, - general.me_reconnect_max_concurrent_per_dc, - general.me_reconnect_backoff_base_ms, - general.me_reconnect_backoff_cap_ms, - general.me_reconnect_fast_retry_count, - general.me_single_endpoint_shadow_writers, - general.me_single_endpoint_outage_mode_enabled, - general.me_single_endpoint_outage_disable_quarantine, - general.me_single_endpoint_outage_backoff_min_ms, - general.me_single_endpoint_outage_backoff_max_ms, - general.me_single_endpoint_shadow_rotate_every_secs, - general.me_floor_mode, - general.me_adaptive_floor_idle_secs, - general.me_adaptive_floor_min_writers_single_endpoint, - general.me_adaptive_floor_min_writers_multi_endpoint, - general.me_adaptive_floor_recover_grace_secs, - general.me_adaptive_floor_writers_per_core_total, - general.me_adaptive_floor_cpu_cores_override, - general.me_adaptive_floor_max_extra_writers_single_per_core, - general.me_adaptive_floor_max_extra_writers_multi_per_core, - general.me_adaptive_floor_max_active_writers_per_core, - general.me_adaptive_floor_max_warm_writers_per_core, - general.me_adaptive_floor_max_active_writers_global, - general.me_adaptive_floor_max_warm_writers_global, - general.hardswap, - general.me_pool_drain_ttl_secs, - general.me_instadrain, - general.me_pool_drain_threshold, - general.me_pool_drain_soft_evict_enabled, - general.me_pool_drain_soft_evict_grace_secs, - general.me_pool_drain_soft_evict_per_writer, - general.me_pool_drain_soft_evict_budget_per_core, - general.me_pool_drain_soft_evict_cooldown_ms, - general.effective_me_pool_force_close_secs(), - general.me_pool_min_fresh_ratio, - general.me_hardswap_warmup_delay_min_ms, - general.me_hardswap_warmup_delay_max_ms, - general.me_hardswap_warmup_extra_passes, - general.me_hardswap_warmup_pass_backoff_base_ms, - general.me_bind_stale_mode, - general.me_bind_stale_ttl_secs, - general.me_secret_atomic_snapshot, - general.me_deterministic_writer_sort, - MeWriterPickMode::default(), - general.me_writer_pick_sample_size, - MeSocksKdfPolicy::default(), - general.me_writer_cmd_channel_capacity, - general.me_route_channel_capacity, - general.me_route_backpressure_base_timeout_ms, - general.me_route_backpressure_high_timeout_ms, - general.me_route_backpressure_high_watermark_pct, - general.me_reader_route_data_wait_ms, - general.me_health_interval_ms_unhealthy, - general.me_health_interval_ms_healthy, - general.me_warn_rate_limit_ms, - MeRouteNoWriterMode::default(), - general.me_route_no_writer_wait_ms, - general.me_route_inline_recovery_attempts, - general.me_route_inline_recovery_wait_ms, - ) -} - -fn encrypt_for_reader(plaintext: &[u8]) -> Vec { - let key = [0u8; 32]; - let iv = 0u128; - let mut cipher = AesCtr::new(&key, iv); - cipher.encrypt(plaintext) -} - -#[tokio::test] -async fn read_client_payload_times_out_on_header_stall() { - let _guard = desync_dedup_test_lock() - .lock() - .expect("middle relay test lock must be available"); - let (reader, _writer) = duplex(1024); - let mut crypto_reader = make_crypto_reader(reader); - let buffer_pool = Arc::new(BufferPool::new()); - let stats = Stats::new(); - let forensics = make_forensics_state(); - let mut frame_counter = 0; - - let result = read_client_payload( - &mut crypto_reader, - ProtoTag::Intermediate, - 1024, - TokioDuration::from_millis(25), - &buffer_pool, - &forensics, - &mut frame_counter, - &stats, - ) - .await; - - assert!( - matches!(result, Err(ProxyError::Io(ref e)) if e.kind() == std::io::ErrorKind::TimedOut), - "stalled header read must time out" - ); -} - -#[tokio::test] -async fn read_client_payload_times_out_on_payload_stall() { - let _guard = desync_dedup_test_lock() - .lock() - .expect("middle relay test lock must be available"); - let (reader, mut writer) = duplex(1024); - let encrypted_len = encrypt_for_reader(&[8, 0, 0, 0]); - writer.write_all(&encrypted_len).await.unwrap(); - - let mut crypto_reader = make_crypto_reader(reader); - let buffer_pool = Arc::new(BufferPool::new()); - let stats = Stats::new(); - let forensics = make_forensics_state(); - let mut frame_counter = 0; - - let result = read_client_payload( - &mut crypto_reader, - ProtoTag::Intermediate, - 1024, - TokioDuration::from_millis(25), - &buffer_pool, - &forensics, - &mut frame_counter, - &stats, - ) - .await; - - assert!( - matches!(result, Err(ProxyError::Io(ref e)) if e.kind() == std::io::ErrorKind::TimedOut), - "stalled payload body read must time out" - ); -} - -#[tokio::test] -async fn read_client_payload_large_intermediate_frame_is_exact() { - let _guard = desync_dedup_test_lock() - .lock() - .expect("middle relay test lock must be available"); - - let (reader, mut writer) = duplex(262_144); - let mut crypto_reader = make_crypto_reader(reader); - let buffer_pool = Arc::new(BufferPool::new()); - let stats = Stats::new(); - let forensics = make_forensics_state(); - let mut frame_counter = 0; - - let payload_len = buffer_pool.buffer_size().saturating_mul(3).max(65_537); - let mut plaintext = Vec::with_capacity(4 + payload_len); - plaintext.extend_from_slice(&(payload_len as u32).to_le_bytes()); - plaintext.extend((0..payload_len).map(|idx| (idx as u8).wrapping_mul(31))); - - let encrypted = encrypt_for_reader(&plaintext); - writer.write_all(&encrypted).await.unwrap(); - - let read = read_client_payload( - &mut crypto_reader, - ProtoTag::Intermediate, - payload_len + 16, - TokioDuration::from_secs(1), - &buffer_pool, - &forensics, - &mut frame_counter, - &stats, - ) - .await - .expect("payload read must succeed") - .expect("frame must be present"); - - let (frame, quickack) = read; - assert!(!quickack, "quickack flag must be unset"); - assert_eq!( - frame.len(), - payload_len, - "payload size must match wire length" - ); - for (idx, byte) in frame.iter().enumerate() { - assert_eq!(*byte, (idx as u8).wrapping_mul(31)); - } - assert_eq!(frame_counter, 1, "exactly one frame must be counted"); -} - -#[tokio::test] -async fn read_client_payload_secure_strips_tail_padding_bytes() { - let _guard = desync_dedup_test_lock() - .lock() - .expect("middle relay test lock must be available"); - - let (reader, mut writer) = duplex(1024); - let mut crypto_reader = make_crypto_reader(reader); - let buffer_pool = Arc::new(BufferPool::new()); - let stats = Stats::new(); - let forensics = make_forensics_state(); - let mut frame_counter = 0; - - let payload = [0x11u8, 0x22, 0x33, 0x44, 0xaa, 0xbb, 0xcc, 0xdd]; - let tail = [0xeeu8, 0xff, 0x99]; - let wire_len = payload.len() + tail.len(); - - let mut plaintext = Vec::with_capacity(4 + wire_len); - plaintext.extend_from_slice(&(wire_len as u32).to_le_bytes()); - plaintext.extend_from_slice(&payload); - plaintext.extend_from_slice(&tail); - let encrypted = encrypt_for_reader(&plaintext); - writer.write_all(&encrypted).await.unwrap(); - - let read = read_client_payload( - &mut crypto_reader, - ProtoTag::Secure, - 1024, - TokioDuration::from_secs(1), - &buffer_pool, - &forensics, - &mut frame_counter, - &stats, - ) - .await - .expect("secure payload read must succeed") - .expect("secure frame must be present"); - - let (frame, quickack) = read; - assert!(!quickack, "quickack flag must be unset"); - assert_eq!(frame.as_ref(), &payload); - assert_eq!(frame_counter, 1, "one secure frame must be counted"); -} - -#[tokio::test] -async fn read_client_payload_secure_rejects_wire_len_below_4() { - let _guard = desync_dedup_test_lock() - .lock() - .expect("middle relay test lock must be available"); - - let (reader, mut writer) = duplex(1024); - let mut crypto_reader = make_crypto_reader(reader); - let buffer_pool = Arc::new(BufferPool::new()); - let stats = Stats::new(); - let forensics = make_forensics_state(); - let mut frame_counter = 0; - - let mut plaintext = Vec::with_capacity(7); - plaintext.extend_from_slice(&3u32.to_le_bytes()); - plaintext.extend_from_slice(&[1u8, 2, 3]); - let encrypted = encrypt_for_reader(&plaintext); - writer.write_all(&encrypted).await.unwrap(); - - let result = read_client_payload( - &mut crypto_reader, - ProtoTag::Secure, - 1024, - TokioDuration::from_secs(1), - &buffer_pool, - &forensics, - &mut frame_counter, - &stats, - ) - .await; - - assert!( - matches!(result, Err(ProxyError::Proxy(ref msg)) if msg.contains("Frame too small: 3")), - "secure wire length below 4 must be fail-closed by the frame-too-small guard" - ); -} - -#[tokio::test] -async fn read_client_payload_intermediate_skips_zero_len_frame() { - let _guard = desync_dedup_test_lock() - .lock() - .expect("middle relay test lock must be available"); - - let (reader, mut writer) = duplex(1024); - let mut crypto_reader = make_crypto_reader(reader); - let buffer_pool = Arc::new(BufferPool::new()); - let stats = Stats::new(); - let forensics = make_forensics_state(); - let mut frame_counter = 0; - - let payload = [7u8, 6, 5, 4, 3, 2, 1, 0]; - let mut plaintext = Vec::with_capacity(4 + 4 + payload.len()); - plaintext.extend_from_slice(&0u32.to_le_bytes()); - plaintext.extend_from_slice(&(payload.len() as u32).to_le_bytes()); - plaintext.extend_from_slice(&payload); - let encrypted = encrypt_for_reader(&plaintext); - writer.write_all(&encrypted).await.unwrap(); - - let read = read_client_payload( - &mut crypto_reader, - ProtoTag::Intermediate, - 1024, - TokioDuration::from_secs(1), - &buffer_pool, - &forensics, - &mut frame_counter, - &stats, - ) - .await - .expect("intermediate payload read must succeed") - .expect("frame must be present"); - - let (frame, quickack) = read; - assert!(!quickack, "quickack flag must be unset"); - assert_eq!(frame.as_ref(), &payload); - assert_eq!(frame_counter, 1, "zero-length frame must be skipped"); -} - -#[tokio::test] -async fn read_client_payload_abridged_extended_len_sets_quickack() { - let _guard = desync_dedup_test_lock() - .lock() - .expect("middle relay test lock must be available"); - - let (reader, mut writer) = duplex(4096); - let mut crypto_reader = make_crypto_reader(reader); - let buffer_pool = Arc::new(BufferPool::new()); - let stats = Stats::new(); - let forensics = make_forensics_state(); - let mut frame_counter = 0; - - let payload_len = 4 * 130; - let len_words = (payload_len / 4) as u32; - let mut plaintext = Vec::with_capacity(1 + 3 + payload_len); - plaintext.push(0xff | 0x80); - let lw = len_words.to_le_bytes(); - plaintext.extend_from_slice(&lw[..3]); - plaintext.extend((0..payload_len).map(|idx| (idx as u8).wrapping_add(17))); - - let encrypted = encrypt_for_reader(&plaintext); - writer.write_all(&encrypted).await.unwrap(); - - let read = read_client_payload( - &mut crypto_reader, - ProtoTag::Abridged, - payload_len + 16, - TokioDuration::from_secs(1), - &buffer_pool, - &forensics, - &mut frame_counter, - &stats, - ) - .await - .expect("abridged payload read must succeed") - .expect("frame must be present"); - - let (frame, quickack) = read; - assert!( - quickack, - "quickack bit must be propagated from abridged header" - ); - assert_eq!(frame.len(), payload_len); - assert_eq!(frame_counter, 1, "one abridged frame must be counted"); -} - -#[tokio::test] -async fn read_client_payload_returns_buffer_to_pool_after_emit() { - let _guard = desync_dedup_test_lock() - .lock() - .expect("middle relay test lock must be available"); - - let pool = Arc::new(BufferPool::with_config(64, 8)); - pool.preallocate(1); - assert_eq!(pool.stats().pooled, 1, "precondition: one pooled buffer"); - - let (reader, mut writer) = duplex(4096); - let mut crypto_reader = make_crypto_reader(reader); - let stats = Stats::new(); - let forensics = make_forensics_state(); - let mut frame_counter = 0; - - // Force growth beyond default pool buffer size to catch ownership-take regressions. - let payload_len = 257usize; - let mut plaintext = Vec::with_capacity(4 + payload_len); - plaintext.extend_from_slice(&(payload_len as u32).to_le_bytes()); - plaintext.extend((0..payload_len).map(|idx| (idx as u8).wrapping_mul(13))); - - let encrypted = encrypt_for_reader(&plaintext); - writer.write_all(&encrypted).await.unwrap(); - - let _ = read_client_payload( - &mut crypto_reader, - ProtoTag::Intermediate, - payload_len + 8, - TokioDuration::from_secs(1), - &pool, - &forensics, - &mut frame_counter, - &stats, - ) - .await - .expect("payload read must succeed") - .expect("frame must be present"); - - assert_eq!(frame_counter, 1); - let pool_stats = pool.stats(); - assert!( - pool_stats.pooled >= 1, - "emitted payload buffer must be returned to pool to avoid pool drain" - ); -} - -#[tokio::test] -async fn read_client_payload_keeps_pool_buffer_checked_out_until_frame_drop() { - let _guard = desync_dedup_test_lock() - .lock() - .expect("middle relay test lock must be available"); - - let pool = Arc::new(BufferPool::with_config(64, 2)); - pool.preallocate(1); - assert_eq!( - pool.stats().pooled, - 1, - "one pooled buffer must be available" - ); - - let (reader, mut writer) = duplex(1024); - let mut crypto_reader = make_crypto_reader(reader); - let stats = Stats::new(); - let forensics = make_forensics_state(); - let mut frame_counter = 0; - - let payload = [0x41u8, 0x42, 0x43, 0x44, 0x45, 0x46, 0x47, 0x48]; - let mut plaintext = Vec::with_capacity(4 + payload.len()); - plaintext.extend_from_slice(&(payload.len() as u32).to_le_bytes()); - plaintext.extend_from_slice(&payload); - let encrypted = encrypt_for_reader(&plaintext); - writer.write_all(&encrypted).await.unwrap(); - - let (frame, quickack) = read_client_payload( - &mut crypto_reader, - ProtoTag::Intermediate, - 1024, - TokioDuration::from_secs(1), - &pool, - &forensics, - &mut frame_counter, - &stats, - ) - .await - .expect("payload read must succeed") - .expect("frame must be present"); - - assert!(!quickack); - assert_eq!(frame.as_ref(), &payload); - assert_eq!( - pool.stats().pooled, - 0, - "buffer must stay checked out while frame payload is alive" - ); - - drop(frame); - assert!( - pool.stats().pooled >= 1, - "buffer must return to pool only after frame drop" - ); -} - -#[tokio::test] -async fn enqueue_c2me_close_unblocks_after_queue_drain() { - let (tx, mut rx) = mpsc::channel::(1); - tx.send(C2MeCommand::Data { - payload: make_pooled_payload(&[0x41]), - flags: 0, - }) - .await - .unwrap(); - - let tx2 = tx.clone(); - let close_task = - tokio::spawn(async move { enqueue_c2me_command(&tx2, C2MeCommand::Close).await }); - - tokio::time::sleep(TokioDuration::from_millis(10)).await; - - let first = timeout(TokioDuration::from_millis(100), rx.recv()) - .await - .unwrap() - .expect("first queued item must be present"); - assert!(matches!(first, C2MeCommand::Data { .. })); - - close_task - .await - .unwrap() - .expect("close enqueue must succeed after drain"); - - let second = timeout(TokioDuration::from_millis(100), rx.recv()) - .await - .unwrap() - .expect("close command must follow after queue drain"); - assert!(matches!(second, C2MeCommand::Close)); -} - -#[tokio::test] -async fn enqueue_c2me_close_full_then_receiver_drop_fails_cleanly() { - let (tx, rx) = mpsc::channel::(1); - tx.send(C2MeCommand::Data { - payload: make_pooled_payload(&[0x42]), - flags: 0, - }) - .await - .unwrap(); - - let tx2 = tx.clone(); - let close_task = - tokio::spawn(async move { enqueue_c2me_command(&tx2, C2MeCommand::Close).await }); - - tokio::time::sleep(TokioDuration::from_millis(10)).await; - drop(rx); - - let result = timeout(TokioDuration::from_secs(1), close_task) - .await - .expect("close task must finish") - .expect("close task must not panic"); - assert!( - result.is_err(), - "close enqueue must fail cleanly when receiver is dropped under pressure" - ); -} - -#[tokio::test] -async fn process_me_writer_response_ack_obeys_flush_policy() { - let (writer_side, _reader_side) = duplex(1024); - let mut writer = make_crypto_writer(writer_side); - let rng = SecureRandom::new(); - let mut frame_buf = Vec::new(); - let stats = Stats::new(); - let bytes_me2c = AtomicU64::new(0); - - let immediate = process_me_writer_response( - MeResponse::Ack(0x11223344), - &mut writer, - ProtoTag::Intermediate, - &rng, - &mut frame_buf, - &stats, - "user", - None, - 0, - &bytes_me2c, - 77, - true, - false, - ) - .await - .expect("ack response must be processed"); - - assert!(matches!( - immediate, - MeWriterResponseOutcome::Continue { - frames: 1, - bytes: 4, - flush_immediately: true, - } - )); - - let delayed = process_me_writer_response( - MeResponse::Ack(0x55667788), - &mut writer, - ProtoTag::Intermediate, - &rng, - &mut frame_buf, - &stats, - "user", - None, - 0, - &bytes_me2c, - 77, - false, - false, - ) - .await - .expect("ack response must be processed"); - - assert!(matches!( - delayed, - MeWriterResponseOutcome::Continue { - frames: 1, - bytes: 4, - flush_immediately: false, - } - )); -} - -#[tokio::test] -async fn process_me_writer_response_data_updates_byte_accounting() { - let (writer_side, _reader_side) = duplex(1024); - let mut writer = make_crypto_writer(writer_side); - let rng = SecureRandom::new(); - let mut frame_buf = Vec::new(); - let stats = Stats::new(); - let bytes_me2c = AtomicU64::new(0); - - let payload = vec![1u8, 2, 3, 4, 5, 6, 7, 8, 9]; - let outcome = process_me_writer_response( - MeResponse::Data { - flags: 0, - data: Bytes::from(payload.clone()), - }, - &mut writer, - ProtoTag::Intermediate, - &rng, - &mut frame_buf, - &stats, - "user", - None, - 0, - &bytes_me2c, - 88, - false, - false, - ) - .await - .expect("data response must be processed"); - - assert!(matches!( - outcome, - MeWriterResponseOutcome::Continue { - frames: 1, - bytes, - flush_immediately: false, - } if bytes == payload.len() - )); - assert_eq!( - bytes_me2c.load(std::sync::atomic::Ordering::Relaxed), - payload.len() as u64, - "ME->C byte accounting must increase by emitted payload size" - ); -} - -#[tokio::test] -async fn process_me_writer_response_data_enforces_live_user_quota() { - let (writer_side, mut reader_side) = duplex(1024); - let mut writer = make_crypto_writer(writer_side); - let rng = SecureRandom::new(); - let mut frame_buf = Vec::new(); - let stats = Stats::new(); - let bytes_me2c = AtomicU64::new(0); - - stats.add_user_octets_from("quota-user", 10); - - let result = process_me_writer_response( - MeResponse::Data { - flags: 0, - data: Bytes::from(vec![1u8, 2, 3, 4]), - }, - &mut writer, - ProtoTag::Intermediate, - &rng, - &mut frame_buf, - &stats, - "quota-user", - Some(12), - 0, - &bytes_me2c, - 89, - false, - false, - ) - .await; - - assert!( - matches!(result, Err(ProxyError::DataQuotaExceeded { user }) if user == "quota-user"), - "ME->client runtime path must terminate when live user quota is crossed" - ); - - let mut raw = [0u8; 1]; - assert!( - timeout(TokioDuration::from_millis(100), reader_side.read(&mut raw)) - .await - .is_err(), - "quota exhaustion must not write any ciphertext to the client stream" - ); -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn process_me_writer_response_concurrent_same_user_quota_does_not_overshoot_limit() { - let stats = Stats::new(); - let bytes_me2c = AtomicU64::new(0); - let user = "quota-race-user"; - - let (writer_side_a, _reader_side_a) = duplex(1024); - let (writer_side_b, _reader_side_b) = duplex(1024); - let mut writer_a = make_crypto_writer(writer_side_a); - let mut writer_b = make_crypto_writer(writer_side_b); - let mut frame_buf_a = Vec::new(); - let mut frame_buf_b = Vec::new(); - let rng_a = SecureRandom::new(); - let rng_b = SecureRandom::new(); - - let fut_a = process_me_writer_response( - MeResponse::Data { - flags: 0, - data: Bytes::from_static(&[0x11]), - }, - &mut writer_a, - ProtoTag::Intermediate, - &rng_a, - &mut frame_buf_a, - &stats, - user, - Some(1), - 0, - &bytes_me2c, - 91, - false, - false, - ); - let fut_b = process_me_writer_response( - MeResponse::Data { - flags: 0, - data: Bytes::from_static(&[0x22]), - }, - &mut writer_b, - ProtoTag::Intermediate, - &rng_b, - &mut frame_buf_b, - &stats, - user, - Some(1), - 0, - &bytes_me2c, - 92, - false, - false, - ); - - let (result_a, result_b) = tokio::join!(fut_a, fut_b); - - assert!( - matches!(result_a, Err(ProxyError::DataQuotaExceeded { ref user }) if user == "quota-race-user") - || matches!(result_a, Ok(_)), - "concurrent quota test must complete without panicking" - ); - assert!( - matches!(result_b, Err(ProxyError::DataQuotaExceeded { ref user }) if user == "quota-race-user") - || matches!(result_b, Ok(_)), - "concurrent quota test must complete without panicking" - ); - assert!( - stats.get_user_total_octets(user) <= 1, - "same-user concurrent middle-relay responses must not overshoot the configured quota" - ); -} - -#[tokio::test] -async fn process_me_writer_response_data_does_not_forward_partial_payload_when_remaining_quota_is_smaller_than_message() - { - let (writer_side, mut reader_side) = duplex(1024); - let mut writer = make_crypto_writer(writer_side); - let rng = SecureRandom::new(); - let mut frame_buf = Vec::new(); - let stats = Stats::new(); - let bytes_me2c = AtomicU64::new(0); - - stats.add_user_octets_to("partial-quota-user", 3); - - let result = process_me_writer_response( - MeResponse::Data { - flags: 0, - data: Bytes::from(vec![1u8, 2, 3, 4]), - }, - &mut writer, - ProtoTag::Intermediate, - &rng, - &mut frame_buf, - &stats, - "partial-quota-user", - Some(4), - 0, - &bytes_me2c, - 90, - false, - false, - ) - .await; - - assert!( - matches!(result, Err(ProxyError::DataQuotaExceeded { user }) if user == "partial-quota-user"), - "ME->client runtime path must reject oversized payloads before writing" - ); - - let mut raw = [0u8; 1]; - assert!( - timeout(TokioDuration::from_millis(100), reader_side.read(&mut raw)) - .await - .is_err(), - "oversized payloads must not leak any partial ciphertext to the client stream" - ); -} - -#[tokio::test] -async fn middle_relay_abort_midflight_releases_route_gauge() { - let stats = Arc::new(Stats::new()); - let me_pool = make_me_pool_for_abort_test(stats.clone()).await; - let config = Arc::new(ProxyConfig::default()); - let buffer_pool = Arc::new(BufferPool::new()); - let rng = Arc::new(SecureRandom::new()); - - let route_runtime = Arc::new(RouteRuntimeController::new(RelayRouteMode::Middle)); - let route_snapshot = route_runtime.snapshot(); - - let (server_side, client_side) = duplex(64 * 1024); - let (server_reader, server_writer) = tokio::io::split(server_side); - let crypto_reader = make_crypto_reader(server_reader); - let crypto_writer = make_crypto_writer(server_writer); - - let success = HandshakeSuccess { - user: "abort-middle-user".to_string(), - dc_idx: 2, - proto_tag: ProtoTag::Intermediate, - dec_key: [0u8; 32], - dec_iv: 0, - enc_key: [0u8; 32], - enc_iv: 0, - peer: "127.0.0.1:50001".parse().unwrap(), - is_tls: false, - }; - - let relay_task = tokio::spawn(handle_via_middle_proxy( - crypto_reader, - crypto_writer, - success, - me_pool, - stats.clone(), - config, - buffer_pool, - "127.0.0.1:443".parse().unwrap(), - rng, - route_runtime.subscribe(), - route_snapshot, - 0xdecafbad, - )); - - let started = tokio::time::timeout(TokioDuration::from_secs(2), async { - loop { - if stats.get_current_connections_me() == 1 { - break; - } - tokio::time::sleep(TokioDuration::from_millis(10)).await; - } - }) - .await; - assert!( - started.is_ok(), - "middle relay must increment route gauge before abort" - ); - - relay_task.abort(); - let joined = relay_task.await; - assert!( - joined.is_err(), - "aborted middle relay task must return join error" - ); - - tokio::time::sleep(TokioDuration::from_millis(20)).await; - assert_eq!( - stats.get_current_connections_me(), - 0, - "route gauge must be released when middle relay task is aborted mid-flight" - ); - - drop(client_side); -} - -#[tokio::test] -async fn middle_relay_cutover_midflight_releases_route_gauge() { - let stats = Arc::new(Stats::new()); - let me_pool = make_me_pool_for_abort_test(stats.clone()).await; - let config = Arc::new(ProxyConfig::default()); - let buffer_pool = Arc::new(BufferPool::new()); - let rng = Arc::new(SecureRandom::new()); - - let route_runtime = Arc::new(RouteRuntimeController::new(RelayRouteMode::Middle)); - let route_snapshot = route_runtime.snapshot(); - - let (server_side, client_side) = duplex(64 * 1024); - let (server_reader, server_writer) = tokio::io::split(server_side); - let crypto_reader = make_crypto_reader(server_reader); - let crypto_writer = make_crypto_writer(server_writer); - - let success = HandshakeSuccess { - user: "cutover-middle-user".to_string(), - dc_idx: 2, - proto_tag: ProtoTag::Intermediate, - dec_key: [0u8; 32], - dec_iv: 0, - enc_key: [0u8; 32], - enc_iv: 0, - peer: "127.0.0.1:50003".parse().unwrap(), - is_tls: false, - }; - - let relay_task = tokio::spawn(handle_via_middle_proxy( - crypto_reader, - crypto_writer, - success, - me_pool, - stats.clone(), - config, - buffer_pool, - "127.0.0.1:443".parse().unwrap(), - rng, - route_runtime.subscribe(), - route_snapshot, - 0xfeed_beef, - )); - - tokio::time::timeout(TokioDuration::from_secs(2), async { - loop { - if stats.get_current_connections_me() == 1 { - break; - } - tokio::time::sleep(TokioDuration::from_millis(10)).await; - } - }) - .await - .expect("middle relay must increment route gauge before cutover"); - - assert!( - route_runtime.set_mode(RelayRouteMode::Direct).is_some(), - "cutover must advance route generation" - ); - - let relay_result = tokio::time::timeout(TokioDuration::from_secs(6), relay_task) - .await - .expect("middle relay must terminate after cutover") - .expect("middle relay task must not panic"); - assert!( - relay_result.is_err(), - "cutover should terminate middle relay session" - ); - assert!( - matches!( - relay_result, - Err(ProxyError::Proxy(ref msg)) if msg == ROUTE_SWITCH_ERROR_MSG - ), - "client-visible cutover error must stay generic and avoid route-internal metadata" - ); - - assert_eq!( - stats.get_current_connections_me(), - 0, - "route gauge must be released when middle relay exits on cutover" - ); - - drop(client_side); -} - -async fn run_quota_race_attempt( - stats: &Stats, - bytes_me2c: &AtomicU64, - user: &str, - payload: u8, - conn_id: u64, - barrier: Arc, -) -> Result { - let (writer_side, _reader_side) = duplex(1024); - let mut writer = make_crypto_writer(writer_side); - let rng = SecureRandom::new(); - let mut frame_buf = Vec::new(); - - barrier.wait().await; - process_me_writer_response( - MeResponse::Data { - flags: 0, - data: Bytes::from(vec![payload]), - }, - &mut writer, - ProtoTag::Intermediate, - &rng, - &mut frame_buf, - stats, - user, - Some(1), - 0, - bytes_me2c, - conn_id, - false, - false, - ) - .await -} - -#[tokio::test] -async fn abridged_max_extended_length_fails_closed_without_panic_or_partial_read() { - let _guard = desync_dedup_test_lock() - .lock() - .expect("middle relay test lock must be available"); - - let (reader, mut writer) = duplex(256); - let mut crypto_reader = make_crypto_reader(reader); - let buffer_pool = Arc::new(BufferPool::new()); - let stats = Stats::new(); - let forensics = make_forensics_state(); - let mut frame_counter = 0; - - let plaintext = vec![0x7f, 0xff, 0xff, 0xff]; - let encrypted = encrypt_for_reader(&plaintext); - writer.write_all(&encrypted).await.unwrap(); - - let result = read_client_payload( - &mut crypto_reader, - ProtoTag::Abridged, - 4096, - TokioDuration::from_secs(1), - &buffer_pool, - &forensics, - &mut frame_counter, - &stats, - ) - .await; - - assert!( - result.is_err(), - "oversized abridged length must fail closed" - ); - assert_eq!( - frame_counter, 0, - "oversized frame must not be counted as accepted" - ); -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 4)] -async fn deterministic_quota_race_exactly_one_succeeds_and_one_is_rejected() { - let stats = Stats::new(); - let bytes_me2c = AtomicU64::new(0); - let user = "gap-t04-race-user"; - let barrier = Arc::new(Barrier::new(2)); - - let f1 = run_quota_race_attempt(&stats, &bytes_me2c, user, 0x11, 5001, barrier.clone()); - let f2 = run_quota_race_attempt(&stats, &bytes_me2c, user, 0x22, 5002, barrier); - - let (r1, r2) = tokio::join!(f1, f2); - - assert!( - matches!(r1, Ok(_) | Err(ProxyError::DataQuotaExceeded { .. })), - "first racer must either finish or fail closed on quota" - ); - assert!( - matches!(r2, Ok(_) | Err(ProxyError::DataQuotaExceeded { .. })), - "second racer must either finish or fail closed on quota" - ); - assert!( - matches!(r1, Err(ProxyError::DataQuotaExceeded { .. })) - || matches!(r2, Err(ProxyError::DataQuotaExceeded { .. })), - "at least one racer must be quota-rejected" - ); - assert_eq!( - stats.get_user_total_octets(user), - 1, - "same-user race must forward/account exactly one payload byte" - ); -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 4)] -async fn stress_quota_race_bursts_never_allow_double_success_per_round() { - let stats = Stats::new(); - let bytes_me2c = AtomicU64::new(0); - - for round in 0..128u64 { - let user = format!("gap-t04-race-burst-{round}"); - let barrier = Arc::new(Barrier::new(2)); - - let one = run_quota_race_attempt( - &stats, - &bytes_me2c, - &user, - 0x33, - 6000 + round, - barrier.clone(), - ); - let two = run_quota_race_attempt(&stats, &bytes_me2c, &user, 0x44, 7000 + round, barrier); - - let (r1, r2) = tokio::join!(one, two); - assert!( - matches!(r1, Ok(_) | Err(ProxyError::DataQuotaExceeded { .. })) - && matches!(r2, Ok(_) | Err(ProxyError::DataQuotaExceeded { .. })), - "round {round}: racers must resolve cleanly without unexpected errors" - ); - assert!( - matches!(r1, Err(ProxyError::DataQuotaExceeded { .. })) - || matches!(r2, Err(ProxyError::DataQuotaExceeded { .. })), - "round {round}: at least one racer must be quota-rejected" - ); - assert_eq!( - stats.get_user_total_octets(&user), - 1, - "round {round}: same-user total octets must remain exactly 1 (single forwarded winner)" - ); - } -} - -#[tokio::test] -async fn middle_relay_cutover_storm_multi_session_keeps_generic_errors_and_releases_gauge() { - let session_count = 6usize; - let stats = Arc::new(Stats::new()); - let me_pool = make_me_pool_for_abort_test(stats.clone()).await; - let config = Arc::new(ProxyConfig::default()); - let buffer_pool = Arc::new(BufferPool::new()); - let rng = Arc::new(SecureRandom::new()); - - let route_runtime = Arc::new(RouteRuntimeController::new(RelayRouteMode::Middle)); - let route_snapshot = route_runtime.snapshot(); - - let mut relay_tasks = Vec::with_capacity(session_count); - let mut client_sides = Vec::with_capacity(session_count); - - for idx in 0..session_count { - let (server_side, client_side) = duplex(64 * 1024); - client_sides.push(client_side); - let (server_reader, server_writer) = tokio::io::split(server_side); - let crypto_reader = make_crypto_reader(server_reader); - let crypto_writer = make_crypto_writer(server_writer); - - let success = HandshakeSuccess { - user: format!("cutover-storm-middle-user-{idx}"), - dc_idx: 2, - proto_tag: ProtoTag::Intermediate, - dec_key: [0u8; 32], - dec_iv: 0, - enc_key: [0u8; 32], - enc_iv: 0, - peer: SocketAddr::new( - std::net::IpAddr::V4(std::net::Ipv4Addr::new(127, 0, 0, 1)), - 52000 + idx as u16, - ), - is_tls: false, - }; - - relay_tasks.push(tokio::spawn(handle_via_middle_proxy( - crypto_reader, - crypto_writer, - success, - me_pool.clone(), - stats.clone(), - config.clone(), - buffer_pool.clone(), - "127.0.0.1:443".parse().unwrap(), - rng.clone(), - route_runtime.subscribe(), - route_snapshot, - 0xB000_0000 + idx as u64, - ))); - } - - tokio::time::timeout(TokioDuration::from_secs(4), async { - loop { - if stats.get_current_connections_me() == session_count as u64 { - break; - } - tokio::time::sleep(TokioDuration::from_millis(10)).await; - } - }) - .await - .expect("all middle sessions must become active before cutover storm"); - - let route_runtime_flipper = route_runtime.clone(); - let flipper = tokio::spawn(async move { - for step in 0..64u32 { - let mode = if (step & 1) == 0 { - RelayRouteMode::Direct - } else { - RelayRouteMode::Middle - }; - let _ = route_runtime_flipper.set_mode(mode); - tokio::time::sleep(TokioDuration::from_millis(15)).await; - } - }); - - for relay_task in relay_tasks { - let relay_result = tokio::time::timeout(TokioDuration::from_secs(10), relay_task) - .await - .expect("middle relay task must finish under cutover storm") - .expect("middle relay task must not panic"); - - assert!( - matches!( - relay_result, - Err(ProxyError::Proxy(ref msg)) if msg == ROUTE_SWITCH_ERROR_MSG - ), - "storm-cutover termination must remain generic for all middle sessions" - ); - } - - flipper.abort(); - let _ = flipper.await; - - assert_eq!( - stats.get_current_connections_me(), - 0, - "middle route gauge must return to zero after cutover storm" - ); - - drop(client_sides); -} - -#[tokio::test] -async fn secure_padding_distribution_in_relay_writer() { - timeout(TokioDuration::from_secs(10), async { - let (mut client_side, relay_side) = duplex(512 * 1024); - let key = [0u8; 32]; - let iv = 0u128; - let mut writer = CryptoWriter::new(relay_side, AesCtr::new(&key, iv), 8 * 1024); - let rng = Arc::new(SecureRandom::new()); - let mut frame_buf = Vec::new(); - let mut decryptor = AesCtr::new(&key, iv); - - let mut padding_counts = [0usize; 4]; - let iterations = 180usize; - let payload = vec![0xAAu8; 100]; // 4-byte aligned - - for _ in 0..iterations { - write_client_payload( - &mut writer, - ProtoTag::Secure, - 0, - &payload, - &rng, - &mut frame_buf, - ) - .await - .expect("payload write must succeed"); - writer - .flush() - .await - .expect("writer flush must complete so encrypted frame becomes readable"); - - let mut len_buf = [0u8; 4]; - client_side - .read_exact(&mut len_buf) - .await - .expect("must read encrypted secure length"); - let decrypted_len_bytes = decryptor.decrypt(&len_buf); - let decrypted_len_bytes: [u8; 4] = decrypted_len_bytes - .try_into() - .expect("decrypted length must be 4 bytes"); - let wire_len = (u32::from_le_bytes(decrypted_len_bytes) & 0x7fff_ffff) as usize; - - assert!( - wire_len >= payload.len(), - "wire length must include at least payload bytes" - ); - let padding_len = wire_len - payload.len(); - assert!(padding_len >= 1 && padding_len <= 3); - padding_counts[padding_len] += 1; - - // Drain and decrypt frame bytes so CTR state stays aligned across writes. - let mut trash = vec![0u8; wire_len]; - client_side - .read_exact(&mut trash) - .await - .expect("must read encrypted secure frame body"); - let _ = decryptor.decrypt(&trash); - } - - for p in 1..=3 { - let count = padding_counts[p]; - assert!( - count > iterations / 8, - "padding length {p} is under-represented ({count}/{iterations})" - ); - } - }) - .await - .expect("secure padding distribution test exceeded runtime budget"); -} - -#[tokio::test] -async fn negative_middle_end_connection_lost_during_relay_exits_on_client_eof() { - let (client_reader_side, client_writer_side) = duplex(1024); - let (_relay_reader_side, relay_writer_side) = duplex(1024); - - let key = [0u8; 32]; - let iv = 0u128; - let crypto_reader = CryptoReader::new(client_reader_side, AesCtr::new(&key, iv)); - let crypto_writer = CryptoWriter::new(relay_writer_side, AesCtr::new(&key, iv), 1024); - - let stats = Arc::new(Stats::new()); - let config = Arc::new(ProxyConfig::default()); - let buffer_pool = Arc::new(BufferPool::with_config(1024, 1)); - let rng = Arc::new(SecureRandom::new()); - let route_runtime = RouteRuntimeController::new(RelayRouteMode::Middle); - - // Create an ME pool. - let me_pool = make_me_pool_for_abort_test(stats.clone()).await; - - // ConnRegistry ids are monotonic; reserve one id so we can predict the - // next session conn_id and close it deterministically without relying on - // writer-bound views such as active_conn_ids(). - let (probe_conn_id, probe_rx) = me_pool.registry().register().await; - drop(probe_rx); - me_pool.registry().unregister(probe_conn_id).await; - let target_conn_id = probe_conn_id.wrapping_add(1); - - let success = HandshakeSuccess { - user: "test-user".to_string(), - peer: "127.0.0.1:12345".parse().unwrap(), - dc_idx: 1, - proto_tag: ProtoTag::Intermediate, - enc_key: key, - enc_iv: iv, - dec_key: key, - dec_iv: iv, - is_tls: false, - }; - - let session_task = tokio::spawn(handle_via_middle_proxy( - crypto_reader, - crypto_writer, - success, - me_pool.clone(), - stats.clone(), - config.clone(), - buffer_pool.clone(), - "127.0.0.1:443".parse().unwrap(), - rng.clone(), - route_runtime.subscribe(), - route_runtime.snapshot(), - 0x1234_5678, - )); - - // Wait until session startup is visible, then unregister the predicted - // conn_id to close the per-session ME response channel. - timeout(TokioDuration::from_millis(500), async { - loop { - if stats.get_current_connections_me() >= 1 { - break; - } - tokio::time::sleep(TokioDuration::from_millis(10)).await; - } - }) - .await - .expect("ME session must start before channel close simulation"); - - me_pool.registry().unregister(target_conn_id).await; - - drop(client_writer_side); - - let result = timeout(TokioDuration::from_secs(2), session_task) - .await - .expect("Session task must terminate after ME drop and client EOF") - .expect("Session task must not panic"); - - assert!( - result.is_ok(), - "Session should complete cleanly after ME drop when client closes, got: {:?}", - result - ); -} - -#[tokio::test] -async fn adversarial_middle_end_drop_plus_cutover_returns_generic_route_switch() { - let (client_reader_side, _client_writer_side) = duplex(1024); - let (_relay_reader_side, relay_writer_side) = duplex(1024); - - let key = [0u8; 32]; - let iv = 0u128; - let crypto_reader = CryptoReader::new(client_reader_side, AesCtr::new(&key, iv)); - let crypto_writer = CryptoWriter::new(relay_writer_side, AesCtr::new(&key, iv), 1024); - - let stats = Arc::new(Stats::new()); - let config = Arc::new(ProxyConfig::default()); - let buffer_pool = Arc::new(BufferPool::with_config(1024, 1)); - let rng = Arc::new(SecureRandom::new()); - let route_runtime = Arc::new(RouteRuntimeController::new(RelayRouteMode::Middle)); - - let me_pool = make_me_pool_for_abort_test(stats.clone()).await; - - // Predict the next conn_id so we can force-drop its ME channel deterministically. - let (probe_conn_id, probe_rx) = me_pool.registry().register().await; - drop(probe_rx); - me_pool.registry().unregister(probe_conn_id).await; - let target_conn_id = probe_conn_id.wrapping_add(1); - - let success = HandshakeSuccess { - user: "test-user-cutover".to_string(), - peer: "127.0.0.1:12345".parse().unwrap(), - dc_idx: 1, - proto_tag: ProtoTag::Intermediate, - enc_key: key, - enc_iv: iv, - dec_key: key, - dec_iv: iv, - is_tls: false, - }; - - let runtime_clone = route_runtime.clone(); - let session_task = tokio::spawn(handle_via_middle_proxy( - crypto_reader, - crypto_writer, - success, - me_pool.clone(), - stats.clone(), - config, - buffer_pool, - "127.0.0.1:443".parse().unwrap(), - rng, - runtime_clone.subscribe(), - runtime_clone.snapshot(), - 0xC001_CAFE, - )); - - timeout(TokioDuration::from_millis(500), async { - loop { - if stats.get_current_connections_me() >= 1 { - break; - } - tokio::time::sleep(TokioDuration::from_millis(10)).await; - } - }) - .await - .expect("ME session must start before race trigger"); - - // Race ME channel drop with route cutover and assert generic client-visible outcome. - me_pool.registry().unregister(target_conn_id).await; - assert!( - route_runtime.set_mode(RelayRouteMode::Direct).is_some(), - "cutover must advance generation" - ); - - let relay_result = timeout(TokioDuration::from_secs(6), session_task) - .await - .expect("session must terminate under ME-drop + cutover race") - .expect("session task must not panic"); - - assert!( - matches!( - relay_result, - Err(ProxyError::Proxy(ref msg)) if msg == ROUTE_SWITCH_ERROR_MSG - ), - "race outcome must remain generic and not leak ME internals, got: {:?}", - relay_result - ); -} - -#[tokio::test] -async fn stress_middle_end_drop_with_client_eof_never_hangs_across_burst() { - let stats = Arc::new(Stats::new()); - let me_pool = make_me_pool_for_abort_test(stats.clone()).await; - - for round in 0..32u64 { - let (client_reader_side, client_writer_side) = duplex(1024); - let (_relay_reader_side, relay_writer_side) = duplex(1024); - - let key = [0u8; 32]; - let iv = 0u128; - let crypto_reader = CryptoReader::new(client_reader_side, AesCtr::new(&key, iv)); - let crypto_writer = CryptoWriter::new(relay_writer_side, AesCtr::new(&key, iv), 1024); - - let config = Arc::new(ProxyConfig::default()); - let buffer_pool = Arc::new(BufferPool::with_config(1024, 1)); - let rng = Arc::new(SecureRandom::new()); - let route_runtime = RouteRuntimeController::new(RelayRouteMode::Middle); - - let (probe_conn_id, probe_rx) = me_pool.registry().register().await; - drop(probe_rx); - me_pool.registry().unregister(probe_conn_id).await; - let target_conn_id = probe_conn_id.wrapping_add(1); - - let success = HandshakeSuccess { - user: format!("stress-me-drop-eof-{round}"), - peer: "127.0.0.1:12345".parse().unwrap(), - dc_idx: 1, - proto_tag: ProtoTag::Intermediate, - enc_key: key, - enc_iv: iv, - dec_key: key, - dec_iv: iv, - is_tls: false, - }; - - let session_task = tokio::spawn(handle_via_middle_proxy( - crypto_reader, - crypto_writer, - success, - me_pool.clone(), - stats.clone(), - config, - buffer_pool, - "127.0.0.1:443".parse().unwrap(), - rng, - route_runtime.subscribe(), - route_runtime.snapshot(), - 0xD00D_0000 + round, - )); - - timeout(TokioDuration::from_millis(500), async { - loop { - if stats.get_current_connections_me() >= 1 { - break; - } - tokio::time::sleep(TokioDuration::from_millis(10)).await; - } - }) - .await - .expect("session must start before forced drop in burst round"); - - me_pool.registry().unregister(target_conn_id).await; - drop(client_writer_side); - - let result = timeout(TokioDuration::from_secs(2), session_task) - .await - .expect("burst round session must terminate quickly") - .expect("burst round session must not panic"); - - assert!( - result.is_ok(), - "burst round {round}: expected clean shutdown after ME drop + EOF, got: {:?}", - result - ); - } -} diff --git a/src/proxy/tests/quota_lock_registry_cross_mode_adversarial_tests.rs b/src/proxy/tests/quota_lock_registry_cross_mode_adversarial_tests.rs deleted file mode 100644 index fb0cf93..0000000 --- a/src/proxy/tests/quota_lock_registry_cross_mode_adversarial_tests.rs +++ /dev/null @@ -1,108 +0,0 @@ -use super::*; -use std::sync::Arc; -use std::sync::{Mutex, OnceLock}; - -fn cross_mode_lock_test_guard() -> std::sync::MutexGuard<'static, ()> { - static TEST_LOCK: OnceLock> = OnceLock::new(); - TEST_LOCK - .get_or_init(|| Mutex::new(())) - .lock() - .unwrap_or_else(|poisoned| poisoned.into_inner()) -} - -#[test] -fn same_user_returns_same_lock_identity() { - let _guard = cross_mode_lock_test_guard(); - let locks = CROSS_MODE_QUOTA_USER_LOCKS.get_or_init(DashMap::new); - locks.clear(); - - let a = cross_mode_quota_user_lock("cross-mode-same-user"); - let b = cross_mode_quota_user_lock("cross-mode-same-user"); - - assert!( - Arc::ptr_eq(&a, &b), - "same user must reuse a stable lock identity" - ); -} - -#[test] -fn saturation_overflow_path_returns_stable_striped_lock_without_cache_growth() { - let _guard = cross_mode_lock_test_guard(); - let locks = CROSS_MODE_QUOTA_USER_LOCKS.get_or_init(DashMap::new); - locks.clear(); - - let prefix = format!("cross-mode-saturated-{}", std::process::id()); - let mut retained = Vec::with_capacity(CROSS_MODE_QUOTA_USER_LOCKS_MAX); - for idx in 0..CROSS_MODE_QUOTA_USER_LOCKS_MAX { - retained.push(cross_mode_quota_user_lock(&format!("{prefix}-{idx}"))); - } - - assert_eq!( - locks.len(), - CROSS_MODE_QUOTA_USER_LOCKS_MAX, - "lock cache must be saturated for overflow check" - ); - - let overflow_user = format!("cross-mode-overflow-{}", std::process::id()); - let overflow_a = cross_mode_quota_user_lock(&overflow_user); - let overflow_b = cross_mode_quota_user_lock(&overflow_user); - - assert_eq!( - locks.len(), - CROSS_MODE_QUOTA_USER_LOCKS_MAX, - "overflow path must not grow bounded lock cache" - ); - assert!( - locks.get(&overflow_user).is_none(), - "overflow user must stay on striped fallback while cache is saturated" - ); - assert!( - Arc::ptr_eq(&overflow_a, &overflow_b), - "overflow user must receive a stable striped lock across repeated lookups" - ); - - drop(retained); -} - -#[test] -fn reclaim_drops_stale_entries_but_preserves_active_user_lock_identity() { - let _guard = cross_mode_lock_test_guard(); - let locks = CROSS_MODE_QUOTA_USER_LOCKS.get_or_init(DashMap::new); - locks.clear(); - - let prefix = format!("cross-mode-reclaim-{}", std::process::id()); - let protected_user = format!("{prefix}-protected"); - - let protected_lock = cross_mode_quota_user_lock(&protected_user); - let mut retained = Vec::with_capacity(CROSS_MODE_QUOTA_USER_LOCKS_MAX.saturating_sub(1)); - for idx in 0..(CROSS_MODE_QUOTA_USER_LOCKS_MAX.saturating_sub(1)) { - retained.push(cross_mode_quota_user_lock(&format!("{prefix}-{idx}"))); - } - - assert_eq!( - locks.len(), - CROSS_MODE_QUOTA_USER_LOCKS_MAX, - "fixture must saturate lock cache before reclaim path is exercised" - ); - - drop(retained); - - let newcomer_user = format!("{prefix}-newcomer"); - let _newcomer = cross_mode_quota_user_lock(&newcomer_user); - - assert!( - locks.get(&protected_user).is_some(), - "active protected user must remain cache-resident after reclaim" - ); - let locked = locks - .get(&protected_user) - .expect("protected user must remain in map after reclaim"); - assert!( - Arc::ptr_eq(locked.value(), &protected_lock), - "reclaim must not swap active user lock identity" - ); - assert!( - locks.get(&newcomer_user).is_some(), - "newcomer should become cacheable after stale entries are reclaimed" - ); -} diff --git a/src/proxy/tests/relay_cross_mode_pipeline_hol_integration_security_tests.rs b/src/proxy/tests/relay_cross_mode_pipeline_hol_integration_security_tests.rs deleted file mode 100644 index 9ea921c..0000000 --- a/src/proxy/tests/relay_cross_mode_pipeline_hol_integration_security_tests.rs +++ /dev/null @@ -1,267 +0,0 @@ -use super::relay_bidirectional; -use crate::stats::Stats; -use crate::stream::BufferPool; -use std::sync::Arc; -use tokio::io::{AsyncReadExt, AsyncWriteExt, duplex}; -use tokio::time::{Duration, timeout}; - -fn quota_test_guard() -> impl Drop { - super::quota_user_lock_test_scope() -} - -#[tokio::test] -async fn negative_same_user_pipeline_stalls_while_middle_lock_is_held() { - let _guard = quota_test_guard(); - - let user = format!("relay-pipeline-stall-{}", std::process::id()); - let held = crate::proxy::middle_relay::cross_mode_quota_user_lock_for_tests(&user); - let held_guard = held - .try_lock() - .expect("test must hold shared cross-mode lock"); - - let stats = Arc::new(Stats::new()); - 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_user = user.clone(); - let relay_stats = Arc::clone(&stats); - let relay_task = tokio::spawn(async move { - relay_bidirectional( - client_reader, - client_writer, - server_reader, - server_writer, - 256, - 256, - &relay_user, - relay_stats, - Some(1024), - Arc::new(BufferPool::new()), - ) - .await - }); - - server_peer - .write_all(&[0xA1]) - .await - .expect("server write should enqueue while relay is stalled"); - - let mut one = [0u8; 1]; - let blocked_read = timeout(Duration::from_millis(40), client_peer.read_exact(&mut one)).await; - assert!( - blocked_read.is_err(), - "same-user relay must remain blocked while cross-mode lock is held" - ); - - drop(held_guard); - - timeout(Duration::from_millis(400), client_peer.read_exact(&mut one)) - .await - .expect("blocked relay must resume after cross-mode lock release") - .expect("resumed relay must deliver queued byte"); - assert_eq!(one, [0xA1]); - - drop(client_peer); - drop(server_peer); - - let relay_result = timeout(Duration::from_secs(1), relay_task) - .await - .expect("relay task must complete") - .expect("relay task must not panic"); - assert!(relay_result.is_ok()); -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 4)] -async fn integration_other_user_pipeline_progresses_while_blocked_user_is_stalled() { - let _guard = quota_test_guard(); - - let blocked_user = format!("relay-pipeline-blocked-{}", std::process::id()); - let free_user = format!("relay-pipeline-free-{}", std::process::id()); - let held = crate::proxy::middle_relay::cross_mode_quota_user_lock_for_tests(&blocked_user); - let held_guard = held - .try_lock() - .expect("test must hold blocked user's shared cross-mode lock"); - - let stats_blocked = Arc::new(Stats::new()); - let stats_free = Arc::new(Stats::new()); - - let (mut blocked_client, blocked_relay_client) = duplex(1024); - let (blocked_relay_server, mut blocked_server) = duplex(1024); - let (blocked_client_reader, blocked_client_writer) = tokio::io::split(blocked_relay_client); - let (blocked_server_reader, blocked_server_writer) = tokio::io::split(blocked_relay_server); - - let (mut free_client, free_relay_client) = duplex(1024); - let (free_relay_server, mut free_server) = duplex(1024); - let (free_client_reader, free_client_writer) = tokio::io::split(free_relay_client); - let (free_server_reader, free_server_writer) = tokio::io::split(free_relay_server); - - let blocked_task = { - let user = blocked_user.clone(); - let stats = Arc::clone(&stats_blocked); - tokio::spawn(async move { - relay_bidirectional( - blocked_client_reader, - blocked_client_writer, - blocked_server_reader, - blocked_server_writer, - 256, - 256, - &user, - stats, - Some(1024), - Arc::new(BufferPool::new()), - ) - .await - }) - }; - - let free_task = { - let user = free_user.clone(); - let stats = Arc::clone(&stats_free); - tokio::spawn(async move { - relay_bidirectional( - free_client_reader, - free_client_writer, - free_server_reader, - free_server_writer, - 256, - 256, - &user, - stats, - Some(1024), - Arc::new(BufferPool::new()), - ) - .await - }) - }; - - blocked_server - .write_all(&[0xB1]) - .await - .expect("blocked user server write should queue"); - free_server - .write_all(&[0xC1]) - .await - .expect("free user server write should queue"); - - let mut blocked_buf = [0u8; 1]; - let mut free_buf = [0u8; 1]; - - let blocked_stalled = timeout( - Duration::from_millis(40), - blocked_client.read_exact(&mut blocked_buf), - ) - .await; - assert!( - blocked_stalled.is_err(), - "blocked user must remain stalled while its lock is held" - ); - - timeout(Duration::from_millis(250), free_client.read_exact(&mut free_buf)) - .await - .expect("free user must make progress while other user is blocked") - .expect("free user read must succeed"); - assert_eq!(free_buf, [0xC1]); - - drop(held_guard); - - timeout(Duration::from_millis(400), blocked_client.read_exact(&mut blocked_buf)) - .await - .expect("blocked user must resume after release") - .expect("blocked user resumed read must succeed"); - assert_eq!(blocked_buf, [0xB1]); - - drop(blocked_client); - drop(blocked_server); - drop(free_client); - drop(free_server); - - assert!( - timeout(Duration::from_secs(1), blocked_task) - .await - .expect("blocked relay task must complete") - .expect("blocked relay task must not panic") - .is_ok() - ); - assert!( - timeout(Duration::from_secs(1), free_task) - .await - .expect("free relay task must complete") - .expect("free relay task must not panic") - .is_ok() - ); -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 4)] -async fn stress_jittered_hold_release_cycles_preserve_pipeline_liveness() { - let _guard = quota_test_guard(); - - let mut seed = 0x5EED_C0DE_2026_0323u64; - for round in 0..24u32 { - seed ^= seed << 7; - seed ^= seed >> 9; - seed ^= seed << 8; - - let hold_ms = 2 + (seed % 10); - let user = format!("relay-pipeline-fuzz-{}-{round}", std::process::id()); - let held = crate::proxy::middle_relay::cross_mode_quota_user_lock_for_tests(&user); - let held_guard = held - .try_lock() - .expect("test must hold lock during fuzz round"); - - let stats = Arc::new(Stats::new()); - 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_user = user.clone(); - let relay_stats = Arc::clone(&stats); - let relay_task = tokio::spawn(async move { - relay_bidirectional( - client_reader, - client_writer, - server_reader, - server_writer, - 256, - 256, - &relay_user, - relay_stats, - Some(1024), - Arc::new(BufferPool::new()), - ) - .await - }); - - server_peer - .write_all(&[0xD1]) - .await - .expect("server write should queue in fuzz round"); - - let mut one = [0u8; 1]; - let stalled = timeout(Duration::from_millis(30), client_peer.read_exact(&mut one)).await; - assert!(stalled.is_err(), "held phase must stall same-user relay"); - - tokio::time::sleep(Duration::from_millis(hold_ms)).await; - drop(held_guard); - - timeout(Duration::from_millis(400), client_peer.read_exact(&mut one)) - .await - .expect("released phase must resume same-user relay") - .expect("released phase read must succeed"); - assert_eq!(one, [0xD1]); - - drop(client_peer); - drop(server_peer); - - assert!( - timeout(Duration::from_secs(1), relay_task) - .await - .expect("fuzz relay task must complete") - .expect("fuzz relay task must not panic") - .is_ok() - ); - } -} \ No newline at end of file diff --git a/src/proxy/tests/relay_cross_mode_pipeline_latency_benchmark_security_tests.rs b/src/proxy/tests/relay_cross_mode_pipeline_latency_benchmark_security_tests.rs deleted file mode 100644 index c967861..0000000 --- a/src/proxy/tests/relay_cross_mode_pipeline_latency_benchmark_security_tests.rs +++ /dev/null @@ -1,213 +0,0 @@ -use super::relay_bidirectional; -use crate::stats::Stats; -use crate::stream::BufferPool; -use std::sync::{Arc, Mutex}; -use tokio::io::{AsyncReadExt, AsyncWriteExt, duplex}; -use tokio::sync::{Barrier, watch}; -use tokio::time::{Duration, Instant, timeout}; - -fn quota_test_guard() -> impl Drop { - super::quota_user_lock_test_scope() -} - -fn percentile_index(len: usize, percentile: usize) -> usize { - ((len * percentile) / 100).min(len.saturating_sub(1)) -} - -#[tokio::test] -async fn micro_benchmark_pipeline_release_to_delivery_latency_stays_bounded() { - let _guard = quota_test_guard(); - - let rounds = 64usize; - let user = format!("relay-pipeline-latency-single-{}", std::process::id()); - let mut samples_ms = Vec::with_capacity(rounds); - - for round in 0..rounds { - let held = crate::proxy::middle_relay::cross_mode_quota_user_lock_for_tests(&user); - let held_guard = held - .try_lock() - .expect("test must hold shared cross-mode lock before round"); - - let stats = Arc::new(Stats::new()); - 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_user = user.clone(); - let relay_stats = Arc::clone(&stats); - let relay_task = tokio::spawn(async move { - relay_bidirectional( - client_reader, - client_writer, - server_reader, - server_writer, - 256, - 256, - &relay_user, - relay_stats, - Some(2048), - Arc::new(BufferPool::new()), - ) - .await - }); - - server_peer - .write_all(&[(round as u8) ^ 0xA5]) - .await - .expect("server write should queue before release"); - - let release_at = Instant::now(); - drop(held_guard); - - let mut one = [0u8; 1]; - timeout(Duration::from_millis(450), client_peer.read_exact(&mut one)) - .await - .expect("client must receive queued byte after release") - .expect("queued byte read must succeed"); - samples_ms.push(release_at.elapsed().as_millis() as u64); - - drop(client_peer); - drop(server_peer); - - let relay_result = timeout(Duration::from_secs(1), relay_task) - .await - .expect("relay task must complete") - .expect("relay task must not panic"); - assert!(relay_result.is_ok()); - } - - samples_ms.sort_unstable(); - let p50_ms = samples_ms[percentile_index(samples_ms.len(), 50)]; - let p95_ms = samples_ms[percentile_index(samples_ms.len(), 95)]; - - assert!( - p50_ms <= 45, - "single-flow release latency p50 must stay bounded; p50_ms={p50_ms}, samples={samples_ms:?}" - ); - assert!( - p95_ms <= 130, - "single-flow release latency p95 must stay bounded; p95_ms={p95_ms}, samples={samples_ms:?}" - ); -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 4)] -async fn stress_128_waiter_pipeline_release_latency_p95_stays_bounded() { - let _guard = quota_test_guard(); - - let waiters = 128usize; - let user = format!("relay-pipeline-latency-fanout-{}", std::process::id()); - let held = crate::proxy::middle_relay::cross_mode_quota_user_lock_for_tests(&user); - let held_guard = held - .try_lock() - .expect("test must hold shared lock before fanout release benchmark"); - - let ready_barrier = Arc::new(Barrier::new(waiters + 1)); - let release_at = Arc::new(Mutex::new(None::)); - let (release_tx, release_rx) = watch::channel(false); - let mut tasks = Vec::with_capacity(waiters); - - for idx in 0..waiters { - let user = user.clone(); - let barrier = Arc::clone(&ready_barrier); - let release_at = Arc::clone(&release_at); - let mut release_rx = release_rx.clone(); - - tasks.push(tokio::spawn(async move { - let stats = Arc::new(Stats::new()); - let (mut client_peer, relay_client) = duplex(512); - let (relay_server, mut server_peer) = duplex(512); - let (client_reader, client_writer) = tokio::io::split(relay_client); - let (server_reader, server_writer) = tokio::io::split(relay_server); - - let relay_user = user; - let relay_stats = Arc::clone(&stats); - let relay_task = tokio::spawn(async move { - relay_bidirectional( - client_reader, - client_writer, - server_reader, - server_writer, - 256, - 256, - &relay_user, - relay_stats, - Some(2048), - Arc::new(BufferPool::new()), - ) - .await - }); - - server_peer - .write_all(&[(idx as u8) ^ 0x5A]) - .await - .expect("fanout server write should queue before release"); - - barrier.wait().await; - release_rx - .changed() - .await - .expect("release signal should remain available"); - - let started = { - let guard = release_at.lock().unwrap_or_else(|poison| poison.into_inner()); - guard.expect("release timestamp must be populated before signal") - }; - - let mut one = [0u8; 1]; - timeout(Duration::from_millis(900), client_peer.read_exact(&mut one)) - .await - .expect("fanout waiter must receive queued byte after release") - .expect("fanout waiter read must succeed"); - - drop(client_peer); - drop(server_peer); - - let relay_result = timeout(Duration::from_secs(2), relay_task) - .await - .expect("fanout relay task must complete") - .expect("fanout relay task must not panic"); - assert!(relay_result.is_ok()); - - started.elapsed().as_millis() as u64 - })); - } - - ready_barrier.wait().await; - { - let mut guard = release_at.lock().unwrap_or_else(|poison| poison.into_inner()); - *guard = Some(Instant::now()); - } - drop(held_guard); - release_tx - .send(true) - .expect("release broadcast must succeed"); - - let mut samples_ms = Vec::with_capacity(waiters); - timeout(Duration::from_secs(8), async { - for task in tasks { - let elapsed = task.await.expect("fanout waiter must not panic"); - samples_ms.push(elapsed); - } - }) - .await - .expect("fanout benchmark must complete in bounded time"); - - samples_ms.sort_unstable(); - let p50_ms = samples_ms[percentile_index(samples_ms.len(), 50)]; - let p95_ms = samples_ms[percentile_index(samples_ms.len(), 95)]; - let max_ms = *samples_ms.last().unwrap_or(&0); - - assert!( - p50_ms <= 120, - "fanout release latency p50 must stay bounded; p50_ms={p50_ms}, p95_ms={p95_ms}, max_ms={max_ms}" - ); - assert!( - p95_ms <= 260, - "fanout release latency p95 must stay bounded; p50_ms={p50_ms}, p95_ms={p95_ms}, max_ms={max_ms}" - ); - assert!( - max_ms <= 700, - "fanout release latency max must stay bounded; p50_ms={p50_ms}, p95_ms={p95_ms}, max_ms={max_ms}" - ); -} \ No newline at end of file diff --git a/src/proxy/tests/relay_cross_mode_quota_fairness_tdd_tests.rs b/src/proxy/tests/relay_cross_mode_quota_fairness_tdd_tests.rs deleted file mode 100644 index adbdb22..0000000 --- a/src/proxy/tests/relay_cross_mode_quota_fairness_tdd_tests.rs +++ /dev/null @@ -1,604 +0,0 @@ -use super::*; -use crate::stats::Stats; -use std::pin::Pin; -use std::sync::Arc; -use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; -use std::task::{Context, Poll, Waker}; -use tokio::io::{AsyncReadExt, AsyncWriteExt, ReadBuf}; -use tokio::sync::Barrier; -use tokio::time::{Duration, timeout}; - -#[derive(Default)] -struct WakeCounter { - wakes: AtomicUsize, -} - -impl std::task::Wake for WakeCounter { - fn wake(self: Arc) { - self.wakes.fetch_add(1, Ordering::Relaxed); - } - - fn wake_by_ref(self: &Arc) { - self.wakes.fetch_add(1, Ordering::Relaxed); - } -} - -fn quota_test_guard() -> impl Drop { - super::quota_user_lock_test_scope() -} - -fn build_context() -> (Arc, Context<'static>) { - let wake_counter = Arc::new(WakeCounter::default()); - let waker = Waker::from(Arc::clone(&wake_counter)); - let leaked_waker: &'static Waker = Box::leak(Box::new(waker)); - (wake_counter, Context::from_waker(leaked_waker)) -} - -#[tokio::test] -async fn positive_cross_mode_uncontended_writer_progresses() { - let _guard = quota_test_guard(); - - let mut io = StatsIo::new( - tokio::io::sink(), - Arc::new(SharedCounters::new()), - Arc::new(Stats::new()), - "cross-mode-tdd-uncontended".to_string(), - Some(4096), - Arc::new(AtomicBool::new(false)), - tokio::time::Instant::now(), - ); - - let result = io.write_all(&[0x11, 0x22]).await; - assert!(result.is_ok(), "uncontended writer must progress"); -} - -#[tokio::test] -async fn adversarial_held_cross_mode_lock_blocks_writer_even_if_local_lock_free() { - let _guard = quota_test_guard(); - - let user = format!("cross-mode-tdd-held-{}", std::process::id()); - let held = crate::proxy::quota_lock_registry::cross_mode_quota_user_lock(&user); - let _held_guard = held - .try_lock() - .expect("test must hold cross-mode lock before polling writer"); - - let mut io = StatsIo::new( - tokio::io::sink(), - Arc::new(SharedCounters::new()), - Arc::new(Stats::new()), - user, - Some(4096), - Arc::new(AtomicBool::new(false)), - tokio::time::Instant::now(), - ); - - let wake_counter = Arc::new(WakeCounter::default()); - let waker = Waker::from(Arc::clone(&wake_counter)); - let mut cx = Context::from_waker(&waker); - - let poll = Pin::new(&mut io).poll_write(&mut cx, &[0xAA]); - assert!(poll.is_pending(), "writer must not bypass held cross-mode lock"); -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 4)] -async fn integration_parallel_waiters_resume_after_cross_mode_release() { - let _guard = quota_test_guard(); - - let user = format!("cross-mode-tdd-resume-{}", std::process::id()); - let held = crate::proxy::quota_lock_registry::cross_mode_quota_user_lock(&user); - let held_guard = held - .try_lock() - .expect("test must hold cross-mode lock before launching waiters"); - - let stats = Arc::new(Stats::new()); - let mut waiters = Vec::new(); - for _ in 0..16 { - let stats = Arc::clone(&stats); - let user = user.clone(); - waiters.push(tokio::spawn(async move { - let mut io = StatsIo::new( - tokio::io::sink(), - Arc::new(SharedCounters::new()), - stats, - user, - Some(4096), - Arc::new(AtomicBool::new(false)), - tokio::time::Instant::now(), - ); - io.write_all(&[0x7F]).await - })); - } - - tokio::time::sleep(Duration::from_millis(5)).await; - drop(held_guard); - - timeout(Duration::from_secs(1), async { - for waiter in waiters { - let result = waiter.await.expect("waiter task must not panic"); - assert!(result.is_ok(), "waiter must complete after cross-mode release"); - } - }) - .await - .expect("all waiters must complete in bounded time"); -} - -#[tokio::test] -async fn adversarial_cross_mode_contention_wake_budget_stays_bounded() { - let _guard = quota_test_guard(); - - let user = format!("cross-mode-tdd-wakes-{}", std::process::id()); - let held = crate::proxy::quota_lock_registry::cross_mode_quota_user_lock(&user); - let _held_guard = held - .try_lock() - .expect("test must hold cross-mode lock before polling"); - - let stats = Arc::new(Stats::new()); - let mut ios = Vec::new(); - let mut counters = Vec::new(); - for _ in 0..20 { - ios.push(StatsIo::new( - tokio::io::sink(), - Arc::new(SharedCounters::new()), - Arc::clone(&stats), - user.clone(), - Some(2048), - Arc::new(AtomicBool::new(false)), - tokio::time::Instant::now(), - )); - } - - for io in &mut ios { - let wake_counter = Arc::new(WakeCounter::default()); - let waker = Waker::from(Arc::clone(&wake_counter)); - let mut cx = Context::from_waker(&waker); - let poll = Pin::new(io).poll_write(&mut cx, &[0x33]); - assert!(poll.is_pending()); - counters.push(wake_counter); - } - - tokio::time::sleep(Duration::from_millis(25)).await; - let total_wakes: usize = counters - .iter() - .map(|counter| counter.wakes.load(Ordering::Relaxed)) - .sum(); - - assert!( - total_wakes <= 20 * 4, - "cross-mode contention should not create wake storms; wakes={total_wakes}" - ); -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 4)] -async fn light_fuzz_cross_mode_release_timing_preserves_read_write_liveness() { - let _guard = quota_test_guard(); - - let mut seed = 0xC0DE_BAAD_2026_0322u64; - for round in 0..16u32 { - seed ^= seed << 7; - seed ^= seed >> 9; - seed ^= seed << 8; - - let sleep_ms = 2 + (seed as u64 % 8); - let user = format!("cross-mode-tdd-fuzz-{}-{round}", std::process::id()); - let held = crate::proxy::quota_lock_registry::cross_mode_quota_user_lock(&user); - let held_guard = held - .try_lock() - .expect("test must hold cross-mode lock in fuzz round"); - - let stats = Arc::new(Stats::new()); - let user_reader = user.clone(); - let reader_task = tokio::spawn(async move { - let mut io = StatsIo::new( - tokio::io::empty(), - Arc::new(SharedCounters::new()), - Arc::clone(&stats), - user_reader, - Some(4096), - Arc::new(AtomicBool::new(false)), - tokio::time::Instant::now(), - ); - let mut one = [0u8; 1]; - io.read(&mut one).await - }); - - let user_writer = user.clone(); - let writer_task = tokio::spawn(async move { - let mut io = StatsIo::new( - tokio::io::sink(), - Arc::new(SharedCounters::new()), - Arc::new(Stats::new()), - user_writer, - Some(4096), - Arc::new(AtomicBool::new(false)), - tokio::time::Instant::now(), - ); - io.write_all(&[0x44]).await - }); - - tokio::time::sleep(Duration::from_millis(sleep_ms)).await; - drop(held_guard); - - let read_done = timeout(Duration::from_millis(350), reader_task) - .await - .expect("reader task must complete after release") - .expect("reader task must not panic"); - assert!(read_done.is_ok()); - - let write_done = timeout(Duration::from_millis(350), writer_task) - .await - .expect("writer task must complete after release") - .expect("writer task must not panic"); - assert!(write_done.is_ok()); - } -} - -#[tokio::test] -async fn integration_middle_lock_blocks_relay_reader_for_same_user() { - let _guard = quota_test_guard(); - - let user = format!("cross-mode-middle-reader-block-{}", std::process::id()); - let held = crate::proxy::middle_relay::cross_mode_quota_user_lock_for_tests(&user); - let _held_guard = held - .try_lock() - .expect("test must hold middle-relay shared lock"); - - let mut io = StatsIo::new( - tokio::io::empty(), - Arc::new(SharedCounters::new()), - Arc::new(Stats::new()), - user, - Some(1024), - Arc::new(AtomicBool::new(false)), - tokio::time::Instant::now(), - ); - - let (_wake_counter, mut cx) = build_context(); - let mut one = [0u8; 1]; - let mut buf = ReadBuf::new(&mut one); - let poll = Pin::new(&mut io).poll_read(&mut cx, &mut buf); - assert!(poll.is_pending()); -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn integration_middle_lock_release_unblocks_relay_reader() { - let _guard = quota_test_guard(); - - let user = format!("cross-mode-middle-reader-release-{}", std::process::id()); - let held = crate::proxy::middle_relay::cross_mode_quota_user_lock_for_tests(&user); - let held_guard = held - .try_lock() - .expect("test must hold middle-relay shared lock"); - - let task = tokio::spawn({ - let user = user.clone(); - async move { - let mut io = StatsIo::new( - tokio::io::empty(), - Arc::new(SharedCounters::new()), - Arc::new(Stats::new()), - user, - Some(1024), - Arc::new(AtomicBool::new(false)), - tokio::time::Instant::now(), - ); - let mut one = [0u8; 1]; - io.read(&mut one).await - } - }); - - tokio::time::sleep(Duration::from_millis(5)).await; - drop(held_guard); - - let done = timeout(Duration::from_millis(300), task) - .await - .expect("reader task must complete after release") - .expect("reader task must not panic"); - assert!(done.is_ok()); -} - -#[tokio::test] -async fn business_different_user_middle_lock_does_not_block_relay_writer() { - let _guard = quota_test_guard(); - - let held_user = format!("cross-mode-middle-held-{}", std::process::id()); - let active_user = format!("cross-mode-middle-active-{}", std::process::id()); - let held = crate::proxy::middle_relay::cross_mode_quota_user_lock_for_tests(&held_user); - let _held_guard = held - .try_lock() - .expect("test must hold middle-relay lock for other user"); - - let mut io = StatsIo::new( - tokio::io::sink(), - Arc::new(SharedCounters::new()), - Arc::new(Stats::new()), - active_user, - Some(1024), - Arc::new(AtomicBool::new(false)), - tokio::time::Instant::now(), - ); - - let (_wake_counter, mut cx) = build_context(); - let poll = Pin::new(&mut io).poll_write(&mut cx, &[0x61]); - assert!(matches!(poll, Poll::Ready(Ok(1)))); -} - -#[tokio::test] -async fn edge_quota_none_bypasses_cross_mode_lock_even_when_held() { - let _guard = quota_test_guard(); - - let user = format!("cross-mode-none-limit-{}", std::process::id()); - let held = crate::proxy::middle_relay::cross_mode_quota_user_lock_for_tests(&user); - let _held_guard = held - .try_lock() - .expect("test must hold lock while quota is disabled"); - - let mut io = StatsIo::new( - tokio::io::sink(), - Arc::new(SharedCounters::new()), - Arc::new(Stats::new()), - user, - None, - Arc::new(AtomicBool::new(false)), - tokio::time::Instant::now(), - ); - - let (_wake_counter, mut cx) = build_context(); - let poll = Pin::new(&mut io).poll_write(&mut cx, &[0x62, 0x63]); - assert!(matches!(poll, Poll::Ready(Ok(2)))); -} - -#[tokio::test] -async fn edge_quota_exceeded_flag_short_circuits_before_lock_path() { - let _guard = quota_test_guard(); - - let user = format!("cross-mode-pre-exceeded-{}", std::process::id()); - let held = crate::proxy::middle_relay::cross_mode_quota_user_lock_for_tests(&user); - let _held_guard = held - .try_lock() - .expect("test must hold shared lock before poll"); - - let quota_exceeded = Arc::new(AtomicBool::new(true)); - let mut io = StatsIo::new( - tokio::io::sink(), - Arc::new(SharedCounters::new()), - Arc::new(Stats::new()), - user, - Some(1024), - Arc::clone("a_exceeded), - tokio::time::Instant::now(), - ); - - let (_wake_counter, mut cx) = build_context(); - let poll = Pin::new(&mut io).poll_write(&mut cx, &[0x64]); - assert!(matches!(poll, Poll::Ready(Err(ref e)) if is_quota_io_error(e))); -} - -#[tokio::test] -async fn adversarial_repoll_while_middle_lock_held_keeps_pending_without_usage_leak() { - let _guard = quota_test_guard(); - - let user = format!("cross-mode-repoll-held-{}", std::process::id()); - let stats = Arc::new(Stats::new()); - let held = crate::proxy::middle_relay::cross_mode_quota_user_lock_for_tests(&user); - let _held_guard = held - .try_lock() - .expect("test must hold lock for repoll sequence"); - - let mut io = StatsIo::new( - tokio::io::sink(), - Arc::new(SharedCounters::new()), - Arc::clone(&stats), - user.clone(), - Some(1024), - Arc::new(AtomicBool::new(false)), - tokio::time::Instant::now(), - ); - - let (_wake_counter, mut cx) = build_context(); - for _ in 0..8 { - let poll = Pin::new(&mut io).poll_write(&mut cx, &[0x65]); - assert!(poll.is_pending()); - } - - assert_eq!(stats.get_user_total_octets(&user), 0); -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 4)] -async fn integration_same_user_mixed_read_write_waiters_resume_after_release() { - let _guard = quota_test_guard(); - - let user = format!("cross-mode-mixed-resume-{}", std::process::id()); - let held = crate::proxy::middle_relay::cross_mode_quota_user_lock_for_tests(&user); - let held_guard = held - .try_lock() - .expect("test must hold lock before spawning mixed waiters"); - - let mut tasks = Vec::new(); - for i in 0..12usize { - let user = user.clone(); - tasks.push(tokio::spawn(async move { - if i % 2 == 0 { - let mut io = StatsIo::new( - tokio::io::empty(), - Arc::new(SharedCounters::new()), - Arc::new(Stats::new()), - user, - Some(1024), - Arc::new(AtomicBool::new(false)), - tokio::time::Instant::now(), - ); - let mut b = [0u8; 1]; - io.read(&mut b).await.map(|_| ()) - } else { - let mut io = StatsIo::new( - tokio::io::sink(), - Arc::new(SharedCounters::new()), - Arc::new(Stats::new()), - user, - Some(1024), - Arc::new(AtomicBool::new(false)), - tokio::time::Instant::now(), - ); - io.write_all(&[0x66]).await - } - })); - } - - tokio::time::sleep(Duration::from_millis(8)).await; - drop(held_guard); - - timeout(Duration::from_secs(1), async { - for task in tasks { - let result = task.await.expect("mixed waiter task must not panic"); - assert!(result.is_ok()); - } - }) - .await - .expect("all mixed waiters must finish after release"); -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 4)] -async fn integration_one_user_blocked_other_user_progresses_under_middle_lock() { - let _guard = quota_test_guard(); - - let blocked_user = format!("cross-mode-blocked-{}", std::process::id()); - let free_user = format!("cross-mode-free-{}", std::process::id()); - let held = crate::proxy::middle_relay::cross_mode_quota_user_lock_for_tests(&blocked_user); - let held_guard = held - .try_lock() - .expect("test must hold blocked user lock"); - - let blocked_task = tokio::spawn({ - let blocked_user = blocked_user.clone(); - async move { - let mut io = StatsIo::new( - tokio::io::sink(), - Arc::new(SharedCounters::new()), - Arc::new(Stats::new()), - blocked_user, - Some(1024), - Arc::new(AtomicBool::new(false)), - tokio::time::Instant::now(), - ); - io.write_all(&[0x77]).await - } - }); - - let free_task = tokio::spawn({ - let free_user = free_user.clone(); - async move { - let mut io = StatsIo::new( - tokio::io::sink(), - Arc::new(SharedCounters::new()), - Arc::new(Stats::new()), - free_user, - Some(1024), - Arc::new(AtomicBool::new(false)), - tokio::time::Instant::now(), - ); - io.write_all(&[0x78]).await - } - }); - - let free_done = timeout(Duration::from_millis(250), free_task) - .await - .expect("free user must not be blocked") - .expect("free user task must not panic"); - assert!(free_done.is_ok()); - - drop(held_guard); - let blocked_done = timeout(Duration::from_secs(1), blocked_task) - .await - .expect("blocked user must resume after release") - .expect("blocked user task must not panic"); - assert!(blocked_done.is_ok()); -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 4)] -async fn stress_middle_lock_release_allows_high_waiter_fanout_completion() { - let _guard = quota_test_guard(); - - let user = format!("cross-mode-fanout-{}", std::process::id()); - let held = crate::proxy::middle_relay::cross_mode_quota_user_lock_for_tests(&user); - let held_guard = held - .try_lock() - .expect("test must hold lock before fanout"); - - let waiters = 48usize; - let gate = Arc::new(Barrier::new(waiters + 1)); - let mut tasks = Vec::new(); - for _ in 0..waiters { - let user = user.clone(); - let gate = Arc::clone(&gate); - tasks.push(tokio::spawn(async move { - let mut io = StatsIo::new( - tokio::io::sink(), - Arc::new(SharedCounters::new()), - Arc::new(Stats::new()), - user, - Some(1024), - Arc::new(AtomicBool::new(false)), - tokio::time::Instant::now(), - ); - gate.wait().await; - io.write_all(&[0x79]).await - })); - } - - gate.wait().await; - tokio::time::sleep(Duration::from_millis(10)).await; - drop(held_guard); - - timeout(Duration::from_secs(2), async { - for task in tasks { - let result = task.await.expect("fanout task must not panic"); - assert!(result.is_ok()); - } - }) - .await - .expect("fanout waiters must complete after release"); -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 4)] -async fn light_fuzz_middle_lock_hold_release_cycles_preserve_same_user_liveness() { - let _guard = quota_test_guard(); - - let mut seed = 0xA11C_EE55_2026_0323u64; - for round in 0..20u32 { - seed ^= seed << 7; - seed ^= seed >> 9; - seed ^= seed << 8; - - let hold_ms = 2 + (seed % 10); - let user = format!("cross-mode-middle-fuzz-{}-{round}", std::process::id()); - let held = crate::proxy::middle_relay::cross_mode_quota_user_lock_for_tests(&user); - let held_guard = held - .try_lock() - .expect("test must hold lock in fuzz round"); - - let writer = tokio::spawn({ - let user = user.clone(); - async move { - let mut io = StatsIo::new( - tokio::io::sink(), - Arc::new(SharedCounters::new()), - Arc::new(Stats::new()), - user, - Some(1024), - Arc::new(AtomicBool::new(false)), - tokio::time::Instant::now(), - ); - io.write_all(&[0x7A]).await - } - }); - - tokio::time::sleep(Duration::from_millis(hold_ms)).await; - drop(held_guard); - - let done = timeout(Duration::from_millis(400), writer) - .await - .expect("writer must complete after lock release") - .expect("writer task must not panic"); - assert!(done.is_ok()); - } -} diff --git a/src/proxy/tests/relay_cross_mode_quota_lock_security_tests.rs b/src/proxy/tests/relay_cross_mode_quota_lock_security_tests.rs deleted file mode 100644 index 5ea806a..0000000 --- a/src/proxy/tests/relay_cross_mode_quota_lock_security_tests.rs +++ /dev/null @@ -1,81 +0,0 @@ -use super::*; -use std::pin::Pin; -use std::sync::Arc; -use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; -use std::task::Waker; -use std::task::{Context, Poll}; - -#[derive(Default)] -struct WakeCounter { - wakes: AtomicUsize, -} - -impl std::task::Wake for WakeCounter { - fn wake(self: Arc) { - self.wakes.fetch_add(1, Ordering::Relaxed); - } - - fn wake_by_ref(self: &Arc) { - self.wakes.fetch_add(1, Ordering::Relaxed); - } -} - -fn build_context() -> (Arc, Context<'static>) { - let wake_counter = Arc::new(WakeCounter::default()); - let waker = Waker::from(Arc::clone(&wake_counter)); - let leaked_waker: &'static Waker = Box::leak(Box::new(waker)); - (wake_counter, Context::from_waker(leaked_waker)) -} - -#[tokio::test] -async fn adversarial_middle_held_cross_mode_lock_blocks_relay_writer() { - let _guard = quota_user_lock_test_scope(); - - let user = "cross-mode-lock-shared-user"; - let held = crate::proxy::middle_relay::cross_mode_quota_user_lock_for_tests(user); - let _held_guard = held - .try_lock() - .expect("test must hold shared cross-mode lock before relay poll"); - - let mut io = StatsIo::new( - tokio::io::sink(), - Arc::new(SharedCounters::new()), - Arc::new(crate::stats::Stats::new()), - user.to_string(), - Some(1024), - Arc::new(AtomicBool::new(false)), - tokio::time::Instant::now(), - ); - - let (_wake_counter, mut cx) = build_context(); - let poll = Pin::new(&mut io).poll_write(&mut cx, &[0x41, 0x42, 0x43]); - - assert!( - matches!(poll, Poll::Pending), - "relay writer must not bypass cross-mode lock held by middle-relay path" - ); -} - -#[tokio::test] -async fn business_cross_mode_lock_uncontended_allows_relay_writer_progress() { - let _guard = quota_user_lock_test_scope(); - - let user = "cross-mode-lock-progress-user"; - let mut io = StatsIo::new( - tokio::io::sink(), - Arc::new(SharedCounters::new()), - Arc::new(crate::stats::Stats::new()), - user.to_string(), - Some(1024), - Arc::new(AtomicBool::new(false)), - tokio::time::Instant::now(), - ); - - let (_wake_counter, mut cx) = build_context(); - let poll = Pin::new(&mut io).poll_write(&mut cx, &[0x51, 0x52]); - - assert!( - matches!(poll, Poll::Ready(Ok(2))), - "relay writer should progress when shared cross-mode lock is uncontended" - ); -} diff --git a/src/proxy/tests/relay_dual_lock_alternating_contention_security_tests.rs b/src/proxy/tests/relay_dual_lock_alternating_contention_security_tests.rs deleted file mode 100644 index 9ac4621..0000000 --- a/src/proxy/tests/relay_dual_lock_alternating_contention_security_tests.rs +++ /dev/null @@ -1,340 +0,0 @@ -use super::*; -use crate::stats::Stats; -use std::pin::Pin; -use std::sync::Arc; -use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; -use std::task::{Context, Waker}; -use tokio::io::AsyncWriteExt; -use tokio::time::{Duration, Instant, timeout}; - -#[derive(Default)] -struct WakeCounter { - wakes: AtomicUsize, -} - -impl std::task::Wake for WakeCounter { - fn wake(self: Arc) { - self.wakes.fetch_add(1, Ordering::Relaxed); - } - - fn wake_by_ref(self: &Arc) { - self.wakes.fetch_add(1, Ordering::Relaxed); - } -} - -fn quota_test_guard() -> impl Drop { - super::quota_user_lock_test_scope() -} - -#[tokio::test] -async fn positive_uncontended_dual_lock_writer_has_zero_retry_attempt() { - let _guard = quota_test_guard(); - - let mut io = StatsIo::new( - tokio::io::sink(), - Arc::new(SharedCounters::new()), - Arc::new(Stats::new()), - format!("dual-lock-alt-positive-{}", std::process::id()), - Some(2048), - Arc::new(AtomicBool::new(false)), - Instant::now(), - ); - - let write = io.write_all(&[0xAA, 0xBB]).await; - assert!(write.is_ok(), "uncontended write must complete"); - assert_eq!( - io.quota_write_retry_attempt, 0, - "uncontended write must not advance retry backoff" - ); -} - -#[tokio::test] -async fn adversarial_alternating_local_and_cross_mode_contention_preserves_backoff_growth() { - let _guard = quota_test_guard(); - - let user = format!("dual-lock-alt-adversarial-{}", std::process::id()); - let local_lock = quota_user_lock(&user); - let cross_mode_lock = crate::proxy::quota_lock_registry::cross_mode_quota_user_lock(&user); - - let mut local_guard = Some( - local_lock - .try_lock() - .expect("test must hold local quota lock initially"), - ); - let mut cross_guard = None; - - let mut io = StatsIo::new( - tokio::io::sink(), - Arc::new(SharedCounters::new()), - Arc::new(Stats::new()), - user, - Some(2048), - Arc::new(AtomicBool::new(false)), - Instant::now(), - ); - - let wake_counter = Arc::new(WakeCounter::default()); - let waker = Waker::from(Arc::clone(&wake_counter)); - let mut cx = Context::from_waker(&waker); - - let first = Pin::new(&mut io).poll_write(&mut cx, &[0x11]); - assert!(first.is_pending(), "held local lock must block first poll"); - - let mut observed_wakes = 0usize; - for idx in 0..18usize { - tokio::time::sleep(Duration::from_millis(6)).await; - - if idx % 2 == 0 { - drop(local_guard.take()); - cross_guard = Some( - cross_mode_lock - .try_lock() - .expect("cross-mode lock should be acquirable while local lock released"), - ); - } else { - drop(cross_guard.take()); - local_guard = Some( - local_lock - .try_lock() - .expect("local lock should be acquirable while cross lock released"), - ); - } - - let wakes = wake_counter.wakes.load(Ordering::Relaxed); - if wakes > observed_wakes { - observed_wakes = wakes; - let pending = Pin::new(&mut io).poll_write(&mut cx, &[0x12]); - assert!( - pending.is_pending(), - "alternating contention must keep write pending while one lock is held" - ); - } - } - - assert!( - io.quota_write_retry_attempt >= 2, - "alternating contention must still ramp retry backoff; got {}", - io.quota_write_retry_attempt - ); - assert!( - wake_counter.wakes.load(Ordering::Relaxed) <= 32, - "alternating contention must stay wake-rate-limited" - ); - - drop(local_guard); - drop(cross_guard); - let ready = Pin::new(&mut io).poll_write(&mut cx, &[0x13]); - assert!(ready.is_ready(), "writer must resume after both locks released"); -} - -#[tokio::test] -async fn edge_retry_scheduler_resets_after_alternating_contention_clears() { - let _guard = quota_test_guard(); - - let user = format!("dual-lock-alt-edge-reset-{}", std::process::id()); - let local_lock = quota_user_lock(&user); - let local_guard = local_lock - .try_lock() - .expect("test must hold local lock for edge scenario"); - - let mut io = StatsIo::new( - tokio::io::sink(), - Arc::new(SharedCounters::new()), - Arc::new(Stats::new()), - user, - Some(2048), - Arc::new(AtomicBool::new(false)), - Instant::now(), - ); - - let wake_counter = Arc::new(WakeCounter::default()); - let waker = Waker::from(Arc::clone(&wake_counter)); - let mut cx = Context::from_waker(&waker); - - let first = Pin::new(&mut io).poll_write(&mut cx, &[0x21]); - assert!(first.is_pending()); - tokio::time::sleep(Duration::from_millis(15)).await; - if wake_counter.wakes.load(Ordering::Relaxed) > 0 { - let next = Pin::new(&mut io).poll_write(&mut cx, &[0x22]); - assert!(next.is_pending()); - } - - drop(local_guard); - - let ready = Pin::new(&mut io).poll_write(&mut cx, &[0x23]); - assert!(ready.is_ready()); - assert_eq!( - io.quota_write_retry_attempt, 0, - "successful dual-lock acquisition must reset retry scheduler" - ); - assert!(!io.quota_write_wake_scheduled); - assert!(io.quota_write_retry_sleep.is_none()); -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 4)] -async fn integration_cross_mode_waiters_remain_live_under_alternating_contention_then_resume() { - let _guard = quota_test_guard(); - - let user = format!("dual-lock-alt-integration-{}", std::process::id()); - let local_lock = quota_user_lock(&user); - let cross_mode_lock = crate::proxy::quota_lock_registry::cross_mode_quota_user_lock(&user); - - let mut waiters = Vec::new(); - for _ in 0..16usize { - let user = user.clone(); - waiters.push(tokio::spawn(async move { - let mut io = StatsIo::new( - tokio::io::sink(), - Arc::new(SharedCounters::new()), - Arc::new(Stats::new()), - user, - Some(2048), - Arc::new(AtomicBool::new(false)), - Instant::now(), - ); - timeout(Duration::from_secs(2), io.write_all(&[0x31])).await - })); - } - - let mut local_guard = Some( - local_lock - .try_lock() - .expect("integration toggle must acquire local lock first"), - ); - let mut cross_guard = None; - - for idx in 0..24usize { - tokio::time::sleep(Duration::from_millis(4)).await; - if idx % 2 == 0 { - drop(local_guard.take()); - cross_guard = cross_mode_lock.try_lock().ok(); - } else { - drop(cross_guard.take()); - local_guard = local_lock.try_lock().ok(); - } - } - - drop(local_guard); - drop(cross_guard); - - for waiter in waiters { - let done = waiter.await.expect("waiter task must not panic"); - assert!( - done.is_ok(), - "waiter must finish once alternating contention window ends" - ); - assert!(done.expect("waiter timeout must not fire").is_ok()); - } -} - -#[tokio::test] -async fn light_fuzz_alternating_contention_matrix_preserves_lock_gating() { - let _guard = quota_test_guard(); - - let user = format!("dual-lock-alt-fuzz-{}", std::process::id()); - let local_lock = quota_user_lock(&user); - let cross_mode_lock = crate::proxy::quota_lock_registry::cross_mode_quota_user_lock(&user); - let mut seed = 0xD00D_BAAD_F00D_2026u64; - - for _round in 0..64u32 { - seed ^= seed << 7; - seed ^= seed >> 9; - seed ^= seed << 8; - - let hold_mode = (seed % 3) as u8; - let local_guard = if hold_mode == 0 { - Some( - local_lock - .try_lock() - .expect("fuzz local lock should be acquirable"), - ) - } else { - None - }; - let cross_guard = if hold_mode == 1 { - Some( - cross_mode_lock - .try_lock() - .expect("fuzz cross lock should be acquirable"), - ) - } else { - None - }; - - let mut io = StatsIo::new( - tokio::io::sink(), - Arc::new(SharedCounters::new()), - Arc::new(Stats::new()), - user.clone(), - Some(1024), - Arc::new(AtomicBool::new(false)), - Instant::now(), - ); - - let write = timeout(Duration::from_millis(35), io.write_all(&[0x51])).await; - if hold_mode == 2 { - assert!(write.is_ok(), "unheld fuzz round must make progress"); - assert!(write.expect("unheld round timeout").is_ok()); - } else { - assert!( - write.is_err(), - "held-lock fuzz round must remain pending inside bounded window" - ); - } - - drop(local_guard); - drop(cross_guard); - } -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 4)] -async fn stress_fanout_alternating_contention_recovers_without_hanging() { - let _guard = quota_test_guard(); - - let user = format!("dual-lock-alt-stress-{}", std::process::id()); - let local_lock = quota_user_lock(&user); - let cross_mode_lock = crate::proxy::quota_lock_registry::cross_mode_quota_user_lock(&user); - - let mut waiters = Vec::new(); - for _ in 0..48usize { - let user = user.clone(); - waiters.push(tokio::spawn(async move { - let mut io = StatsIo::new( - tokio::io::sink(), - Arc::new(SharedCounters::new()), - Arc::new(Stats::new()), - user, - Some(4096), - Arc::new(AtomicBool::new(false)), - Instant::now(), - ); - timeout(Duration::from_secs(3), io.write_all(&[0xA0, 0xA1])).await - })); - } - - let mut local_guard = Some( - local_lock - .try_lock() - .expect("stress toggle must acquire local lock first"), - ); - let mut cross_guard = None; - for idx in 0..40usize { - tokio::time::sleep(Duration::from_millis(3)).await; - if idx % 2 == 0 { - drop(local_guard.take()); - cross_guard = cross_mode_lock.try_lock().ok(); - } else { - drop(cross_guard.take()); - local_guard = local_lock.try_lock().ok(); - } - } - - drop(local_guard); - drop(cross_guard); - - for waiter in waiters { - let done = waiter.await.expect("stress waiter task must not panic"); - assert!(done.is_ok(), "stress waiter timed out under alternating contention"); - assert!(done.expect("stress waiter timeout should not fire").is_ok()); - } -} diff --git a/src/proxy/tests/relay_dual_lock_backoff_regression_security_tests.rs b/src/proxy/tests/relay_dual_lock_backoff_regression_security_tests.rs deleted file mode 100644 index ce26941..0000000 --- a/src/proxy/tests/relay_dual_lock_backoff_regression_security_tests.rs +++ /dev/null @@ -1,74 +0,0 @@ -use super::*; -use crate::stats::Stats; -use std::pin::Pin; -use std::sync::Arc; -use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; -use std::task::{Context, Waker}; -use tokio::time::{Duration, Instant}; - -#[derive(Default)] -struct WakeCounter { - wakes: AtomicUsize, -} - -impl std::task::Wake for WakeCounter { - fn wake(self: Arc) { - self.wakes.fetch_add(1, Ordering::Relaxed); - } - - fn wake_by_ref(self: &Arc) { - self.wakes.fetch_add(1, Ordering::Relaxed); - } -} - -fn quota_test_guard() -> impl Drop { - super::quota_user_lock_test_scope() -} - -#[tokio::test] -async fn adversarial_cross_mode_only_contention_backoff_attempt_must_ramp() { - let _guard = quota_test_guard(); - - let user = format!("dual-lock-backoff-{}", std::process::id()); - let cross_mode_lock = crate::proxy::quota_lock_registry::cross_mode_quota_user_lock(&user); - let held_cross_mode_guard = cross_mode_lock - .try_lock() - .expect("test must hold cross-mode lock before polling"); - - let mut io = StatsIo::new( - tokio::io::sink(), - Arc::new(SharedCounters::new()), - Arc::new(Stats::new()), - user, - Some(2048), - Arc::new(AtomicBool::new(false)), - Instant::now(), - ); - - let wake_counter = Arc::new(WakeCounter::default()); - let waker = Waker::from(Arc::clone(&wake_counter)); - let mut cx = Context::from_waker(&waker); - - let first = Pin::new(&mut io).poll_write(&mut cx, &[0xAA]); - assert!(first.is_pending(), "held cross-mode lock must block writer"); - - let started = Instant::now(); - let mut last_wakes = 0usize; - while started.elapsed() < Duration::from_millis(120) { - let wakes = wake_counter.wakes.load(Ordering::Relaxed); - if wakes > last_wakes { - last_wakes = wakes; - let next = Pin::new(&mut io).poll_write(&mut cx, &[0xAB]); - assert!(next.is_pending(), "writer must remain blocked while lock is held"); - } - tokio::time::sleep(Duration::from_millis(1)).await; - } - - assert!( - io.quota_write_retry_attempt >= 2, - "retry attempt must ramp under sustained second-lock contention; got {}", - io.quota_write_retry_attempt - ); - - drop(held_cross_mode_guard); -} diff --git a/src/proxy/tests/relay_dual_lock_contention_matrix_security_tests.rs b/src/proxy/tests/relay_dual_lock_contention_matrix_security_tests.rs deleted file mode 100644 index 513d92b..0000000 --- a/src/proxy/tests/relay_dual_lock_contention_matrix_security_tests.rs +++ /dev/null @@ -1,325 +0,0 @@ -use super::*; -use crate::stats::Stats; -use std::pin::Pin; -use std::sync::Arc; -use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; -use std::task::{Context, Waker}; -use tokio::io::{AsyncReadExt, AsyncWriteExt, ReadBuf}; -use tokio::time::{Duration, Instant, timeout}; - -#[derive(Default)] -struct WakeCounter { - wakes: AtomicUsize, -} - -impl std::task::Wake for WakeCounter { - fn wake(self: Arc) { - self.wakes.fetch_add(1, Ordering::Relaxed); - } - - fn wake_by_ref(self: &Arc) { - self.wakes.fetch_add(1, Ordering::Relaxed); - } -} - -fn quota_test_guard() -> impl Drop { - super::quota_user_lock_test_scope() -} - -fn build_context() -> (Arc, Context<'static>) { - let wake_counter = Arc::new(WakeCounter::default()); - let waker = Waker::from(Arc::clone(&wake_counter)); - let leaked_waker: &'static Waker = Box::leak(Box::new(waker)); - (wake_counter, Context::from_waker(leaked_waker)) -} - -#[tokio::test] -async fn positive_uncontended_dual_locks_writer_completes_without_retry_state() { - let _guard = quota_test_guard(); - - let mut io = StatsIo::new( - tokio::io::sink(), - Arc::new(SharedCounters::new()), - Arc::new(Stats::new()), - format!("dual-lock-positive-{}", std::process::id()), - Some(4096), - Arc::new(AtomicBool::new(false)), - Instant::now(), - ); - - let (_wake_counter, mut cx) = build_context(); - let poll = Pin::new(&mut io).poll_write(&mut cx, &[0x01, 0x02, 0x03]); - assert!(poll.is_ready()); - assert_eq!(io.quota_write_retry_attempt, 0); - assert!(!io.quota_write_wake_scheduled); - assert!(io.quota_write_retry_sleep.is_none()); -} - -#[tokio::test] -async fn negative_local_lock_contention_read_retry_attempt_ramps() { - let _guard = quota_test_guard(); - - let user = format!("dual-lock-local-contention-{}", std::process::id()); - let held = quota_user_lock(&user); - let held_guard = held - .try_lock() - .expect("test must hold local quota lock before polling"); - - let mut io = StatsIo::new( - tokio::io::empty(), - Arc::new(SharedCounters::new()), - Arc::new(Stats::new()), - user, - Some(2048), - Arc::new(AtomicBool::new(false)), - Instant::now(), - ); - - let (wake_counter, mut cx) = build_context(); - let mut one = [0u8; 1]; - let mut buf = ReadBuf::new(&mut one); - let first = Pin::new(&mut io).poll_read(&mut cx, &mut buf); - assert!(first.is_pending()); - - let started = Instant::now(); - let mut observed = 0usize; - while started.elapsed() < Duration::from_millis(120) { - let wakes = wake_counter.wakes.load(Ordering::Relaxed); - if wakes > observed { - observed = wakes; - let mut step_buf = ReadBuf::new(&mut one); - let next = Pin::new(&mut io).poll_read(&mut cx, &mut step_buf); - assert!(next.is_pending()); - } - tokio::time::sleep(Duration::from_millis(1)).await; - } - - assert!( - io.quota_read_retry_attempt >= 2, - "retry attempt must ramp under sustained local-lock contention; got {}", - io.quota_read_retry_attempt - ); - - drop(held_guard); -} - -#[tokio::test] -async fn edge_cross_mode_contention_release_resets_retry_scheduler_on_success() { - let _guard = quota_test_guard(); - - let user = format!("dual-lock-reset-{}", std::process::id()); - let cross_mode = crate::proxy::quota_lock_registry::cross_mode_quota_user_lock(&user); - let held_guard = cross_mode - .try_lock() - .expect("test must hold cross-mode lock before polling"); - - let mut io = StatsIo::new( - tokio::io::sink(), - Arc::new(SharedCounters::new()), - Arc::new(Stats::new()), - user, - Some(2048), - Arc::new(AtomicBool::new(false)), - Instant::now(), - ); - - let (wake_counter, mut cx) = build_context(); - let first = Pin::new(&mut io).poll_write(&mut cx, &[0x10]); - assert!(first.is_pending()); - - tokio::time::sleep(Duration::from_millis(20)).await; - if wake_counter.wakes.load(Ordering::Relaxed) > 0 { - let next = Pin::new(&mut io).poll_write(&mut cx, &[0x11]); - assert!(next.is_pending()); - } - - drop(held_guard); - let ready = Pin::new(&mut io).poll_write(&mut cx, &[0x12]); - assert!(ready.is_ready()); - assert_eq!(io.quota_write_retry_attempt, 0); - assert!(!io.quota_write_wake_scheduled); - assert!(io.quota_write_retry_sleep.is_none()); -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 4)] -async fn adversarial_cross_mode_hold_blocks_many_waiters_without_usage_leak() { - let _guard = quota_test_guard(); - - let user = format!("dual-lock-adversarial-{}", std::process::id()); - let stats = Arc::new(Stats::new()); - let held = crate::proxy::quota_lock_registry::cross_mode_quota_user_lock(&user); - let held_guard = held - .try_lock() - .expect("test must hold cross-mode lock before launching waiters"); - - let mut tasks = Vec::new(); - for _ in 0..24usize { - let stats = Arc::clone(&stats); - let user = user.clone(); - tasks.push(tokio::spawn(async move { - let mut io = StatsIo::new( - tokio::io::sink(), - Arc::new(SharedCounters::new()), - stats, - user, - Some(1024), - Arc::new(AtomicBool::new(false)), - Instant::now(), - ); - timeout(Duration::from_millis(40), io.write_all(&[0x33])).await - })); - } - - for task in tasks { - let timed = task.await.expect("waiter task must not panic"); - assert!(timed.is_err(), "held cross-mode lock must keep waiter pending"); - } - - assert_eq!(stats.get_user_total_octets(&user), 0); - drop(held_guard); -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 4)] -async fn integration_waiters_resume_after_cross_mode_release() { - let _guard = quota_test_guard(); - - let user = format!("dual-lock-integration-{}", std::process::id()); - let held = crate::proxy::quota_lock_registry::cross_mode_quota_user_lock(&user); - let held_guard = held - .try_lock() - .expect("test must hold cross-mode lock before starting waiter"); - - let task = tokio::spawn({ - let user = user.clone(); - async move { - let mut io = StatsIo::new( - tokio::io::sink(), - Arc::new(SharedCounters::new()), - Arc::new(Stats::new()), - user, - Some(1024), - Arc::new(AtomicBool::new(false)), - Instant::now(), - ); - io.write_all(&[0x44]).await - } - }); - - tokio::time::sleep(Duration::from_millis(10)).await; - drop(held_guard); - - let done = timeout(Duration::from_secs(1), task) - .await - .expect("waiter task must complete after release") - .expect("waiter task must not panic"); - assert!(done.is_ok()); -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 4)] -async fn light_fuzz_randomized_lock_holds_preserve_liveness_and_quota_bounds() { - let _guard = quota_test_guard(); - - let user = format!("dual-lock-fuzz-{}", std::process::id()); - let stats = Arc::new(Stats::new()); - let mut seed = 0xA55A_55AA_C3D2_E1F0u64; - - for _round in 0..48u32 { - seed ^= seed << 7; - seed ^= seed >> 9; - seed ^= seed << 8; - - let hold_mode = (seed % 3) as u8; - let mut local_lock = None; - let mut cross_lock = None; - let mut local_guard = None; - let mut cross_guard = None; - - if hold_mode == 0 { - local_lock = Some(quota_user_lock(&user)); - local_guard = Some( - local_lock - .as_ref() - .expect("local lock should be present") - .try_lock() - .expect("local lock should be acquirable in fuzz round"), - ); - } else if hold_mode == 1 { - cross_lock = Some(crate::proxy::quota_lock_registry::cross_mode_quota_user_lock( - &user, - )); - cross_guard = Some( - cross_lock - .as_ref() - .expect("cross lock should be present") - .try_lock() - .expect("cross lock should be acquirable in fuzz round"), - ); - } - - let mut io = StatsIo::new( - tokio::io::sink(), - Arc::new(SharedCounters::new()), - Arc::clone(&stats), - user.clone(), - Some(4096), - Arc::new(AtomicBool::new(false)), - Instant::now(), - ); - - let write = timeout(Duration::from_millis(25), io.write_all(&[0x7A])).await; - if hold_mode == 2 { - assert!(write.is_ok(), "unheld round must make progress"); - } else { - assert!(write.is_err(), "held-lock round must stay blocked within timeout"); - } - - drop(local_guard); - drop(cross_guard); - drop(local_lock); - drop(cross_lock); - } - - assert!(stats.get_user_total_octets(&user) <= 4096); -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 4)] -async fn stress_fanout_waiters_complete_after_release_without_panics() { - let _guard = quota_test_guard(); - - let user = format!("dual-lock-stress-{}", std::process::id()); - let held = crate::proxy::quota_lock_registry::cross_mode_quota_user_lock(&user); - let held_guard = held - .try_lock() - .expect("test must hold cross-mode lock before stress fanout"); - - let waiters = 64usize; - let mut tasks = Vec::new(); - for _ in 0..waiters { - let user = user.clone(); - tasks.push(tokio::spawn(async move { - let mut io = StatsIo::new( - tokio::io::empty(), - Arc::new(SharedCounters::new()), - Arc::new(Stats::new()), - user, - Some(1024), - Arc::new(AtomicBool::new(false)), - Instant::now(), - ); - let mut one = [0u8; 1]; - io.read(&mut one).await - })); - } - - tokio::time::sleep(Duration::from_millis(12)).await; - drop(held_guard); - - timeout(Duration::from_secs(2), async { - for task in tasks { - let result = task.await.expect("stress waiter task must not panic"); - assert!(result.is_ok()); - } - }) - .await - .expect("all stress waiters must complete after release"); -} diff --git a/src/proxy/tests/relay_dual_lock_race_harness_security_tests.rs b/src/proxy/tests/relay_dual_lock_race_harness_security_tests.rs deleted file mode 100644 index ec180e8..0000000 --- a/src/proxy/tests/relay_dual_lock_race_harness_security_tests.rs +++ /dev/null @@ -1,128 +0,0 @@ -use super::*; -use crate::stats::Stats; -use std::sync::Arc; -use std::sync::atomic::AtomicBool; -use tokio::io::AsyncWriteExt; -use tokio::time::{Duration, timeout}; - -fn quota_test_guard() -> impl Drop { - super::quota_user_lock_test_scope() -} - -fn make_stats_io(user: String) -> StatsIo { - StatsIo::new( - tokio::io::sink(), - Arc::new(SharedCounters::new()), - Arc::new(Stats::new()), - user, - Some(4096), - Arc::new(AtomicBool::new(false)), - tokio::time::Instant::now(), - ) -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 4)] -async fn light_fuzz_1024_round_hold_release_cycles_preserve_same_user_liveness() { - let _guard = quota_test_guard(); - - let user = format!("dual-lock-race-fuzz-{}", std::process::id()); - let mut seed = 0xD1CE_BAAD_5EED_1234u64; - - for round in 0..1024u32 { - seed ^= seed << 7; - seed ^= seed >> 9; - seed ^= seed << 8; - - let hold = (seed & 1) == 0; - let hold_ms = (seed % 3) as u64; - - let maybe_lock = if hold { - Some(crate::proxy::quota_lock_registry::cross_mode_quota_user_lock( - &user, - )) - } else { - None - }; - - let maybe_guard = maybe_lock.as_ref().map(|lock| { - lock.try_lock() - .expect("cross-mode lock must be acquirable in fuzz round") - }); - - if hold { - let mut blocked_io = make_stats_io(user.clone()); - let blocked = timeout(Duration::from_millis(5), blocked_io.write_all(&[0xA5])).await; - assert!( - blocked.is_err(), - "held round must block waiter before lock release (round={round})" - ); - - if hold_ms > 0 { - tokio::time::sleep(Duration::from_millis(hold_ms)).await; - } - } else { - let mut free_io = make_stats_io(user.clone()); - let free = timeout(Duration::from_millis(120), free_io.write_all(&[0xA5])).await; - assert!( - free.is_ok(), - "unheld round must complete promptly (round={round})" - ); - assert!(free.expect("unheld round should complete").is_ok()); - } - - drop(maybe_guard); - - let done = timeout(Duration::from_millis(350), async { - let user = user.clone(); - let mut io = make_stats_io(user); - io.write_all(&[0xA6]).await - }) - .await - .expect("post-release write must complete in bounded time"); - assert!(done.is_ok()); - } -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 4)] -async fn stress_jittered_three_waiter_rounds_do_not_starve_after_release() { - let _guard = quota_test_guard(); - - let user = format!("dual-lock-race-stress-{}", std::process::id()); - let mut seed = 0xC0FF_EE77_4444_9999u64; - - for round in 0..256u32 { - seed ^= seed << 7; - seed ^= seed >> 9; - seed ^= seed << 8; - - let hold_ms = (seed % 4) as u64; - let lock = crate::proxy::quota_lock_registry::cross_mode_quota_user_lock(&user); - let guard = lock - .try_lock() - .expect("cross-mode lock must be acquirable at round start"); - - let mut waiters = Vec::new(); - for _ in 0..3usize { - let user = user.clone(); - waiters.push(tokio::spawn(async move { - let mut io = make_stats_io(user); - io.write_all(&[0x55]).await - })); - } - - tokio::time::sleep(Duration::from_millis(hold_ms)).await; - drop(guard); - - timeout(Duration::from_secs(1), async { - for waiter in waiters { - let done = waiter.await.expect("waiter task must not panic"); - assert!( - done.is_ok(), - "waiter must complete after release (round={round})" - ); - } - }) - .await - .expect("all waiters must complete in bounded time after release"); - } -} diff --git a/src/proxy/tests/relay_quota_lock_eviction_lifecycle_tdd_tests.rs b/src/proxy/tests/relay_quota_lock_eviction_lifecycle_tdd_tests.rs deleted file mode 100644 index 806efb6..0000000 --- a/src/proxy/tests/relay_quota_lock_eviction_lifecycle_tdd_tests.rs +++ /dev/null @@ -1,79 +0,0 @@ -use super::*; -use dashmap::DashMap; -use std::sync::Arc; -use tokio::time::{Duration, timeout}; - -#[test] -fn tdd_explicit_quota_lock_evict_reclaims_only_unheld_entries() { - let _guard = quota_user_lock_test_scope(); - let map = QUOTA_USER_LOCKS.get_or_init(DashMap::new); - map.clear(); - - let held_user = format!("quota-evict-held-{}", std::process::id()); - let stale_a_user = format!("quota-evict-stale-a-{}", std::process::id()); - let stale_b_user = format!("quota-evict-stale-b-{}", std::process::id()); - - let held = quota_user_lock(&held_user); - let stale_a = quota_user_lock(&stale_a_user); - let stale_b = quota_user_lock(&stale_b_user); - - assert!(map.get(&held_user).is_some()); - assert!(map.get(&stale_a_user).is_some()); - assert!(map.get(&stale_b_user).is_some()); - - drop(stale_a); - drop(stale_b); - - quota_user_lock_evict(); - - assert!( - map.get(&held_user).is_some(), - "held entry must survive eviction" - ); - assert!( - map.get(&stale_a_user).is_none(), - "unheld stale entry must be reclaimed" - ); - assert!( - map.get(&stale_b_user).is_none(), - "unheld stale entry must be reclaimed" - ); - - drop(held); -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn tdd_periodic_quota_lock_evictor_reclaims_stale_entries_off_hot_path() { - let _guard = quota_user_lock_test_scope(); - let map = QUOTA_USER_LOCKS.get_or_init(DashMap::new); - map.clear(); - - let held_user = format!("quota-evict-loop-held-{}", std::process::id()); - let stale_user = format!("quota-evict-loop-stale-{}", std::process::id()); - - let held = quota_user_lock(&held_user); - let stale = quota_user_lock(&stale_user); - - assert_eq!(map.len(), 2); - drop(stale); - - let evictor = spawn_quota_user_lock_evictor_for_tests(Duration::from_millis(5)); - - timeout(Duration::from_millis(200), async { - loop { - if map.get(&stale_user).is_none() { - break; - } - tokio::time::sleep(Duration::from_millis(5)).await; - } - }) - .await - .expect("periodic quota lock evictor must reclaim stale entry"); - - evictor.abort(); - - assert!(map.get(&held_user).is_some()); - assert!(map.get(&stale_user).is_none()); - - drop(held); -} diff --git a/src/proxy/tests/relay_quota_lock_eviction_stress_security_tests.rs b/src/proxy/tests/relay_quota_lock_eviction_stress_security_tests.rs deleted file mode 100644 index 251582a..0000000 --- a/src/proxy/tests/relay_quota_lock_eviction_stress_security_tests.rs +++ /dev/null @@ -1,153 +0,0 @@ -use super::*; -use dashmap::DashMap; -use std::sync::Arc; -use tokio::task::JoinSet; -use tokio::time::{Duration, timeout}; - -#[tokio::test(flavor = "multi_thread", worker_threads = 4)] -async fn stress_background_evictor_with_high_churn_keeps_cache_bounded_and_live() { - let _guard = quota_user_lock_test_scope(); - let map = QUOTA_USER_LOCKS.get_or_init(DashMap::new); - map.clear(); - - let evictor = spawn_quota_user_lock_evictor_for_tests(Duration::from_millis(5)); - - let mut tasks = JoinSet::new(); - for worker in 0..24u32 { - tasks.spawn(async move { - for round in 0..320u32 { - let user = format!( - "quota-evict-stress-user-{}-{}-{}", - std::process::id(), - worker, - round - ); - let lock = quota_user_lock(&user); - if round % 19 == 0 { - tokio::task::yield_now().await; - } - drop(lock); - } - }); - } - - while let Some(done) = tasks.join_next().await { - done.expect("stress worker must not panic"); - } - - quota_user_lock_evict(); - tokio::time::sleep(Duration::from_millis(20)).await; - - assert!( - map.len() <= QUOTA_USER_LOCKS_MAX, - "quota lock map must remain bounded after churn + eviction" - ); - - let sanity_user = format!("quota-evict-stress-sanity-{}", std::process::id()); - let sanity_lock = quota_user_lock(&sanity_user); - assert!( - map.get(&sanity_user).is_some(), - "sanity user should be cacheable after eviction reclaimed stale entries" - ); - - drop(sanity_lock); - evictor.abort(); -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 4)] -async fn adversarial_held_lock_survives_repeated_eviction_then_reclaims_after_release() { - let _guard = quota_user_lock_test_scope(); - let map = QUOTA_USER_LOCKS.get_or_init(DashMap::new); - map.clear(); - - let held_user = format!("quota-evict-held-survive-{}", std::process::id()); - let held = quota_user_lock(&held_user); - - let evictor = spawn_quota_user_lock_evictor_for_tests(Duration::from_millis(3)); - - for idx in 0..512u32 { - let user = format!("quota-evict-held-churn-{}-{}", std::process::id(), idx); - let temp = quota_user_lock(&user); - drop(temp); - if idx % 32 == 0 { - tokio::task::yield_now().await; - } - } - - let reacquired = quota_user_lock(&held_user); - assert!( - Arc::ptr_eq(&held, &reacquired), - "held user lock identity must remain stable across repeated evictions" - ); - assert!( - map.get(&held_user).is_some(), - "held user entry must not be reclaimed while externally referenced" - ); - - drop(reacquired); - drop(held); - - timeout(Duration::from_millis(300), async { - loop { - if map.get(&held_user).is_none() { - break; - } - tokio::time::sleep(Duration::from_millis(5)).await; - } - }) - .await - .expect("released held lock must be reclaimed by periodic evictor"); - - evictor.abort(); -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 4)] -async fn stress_saturation_then_periodic_eviction_recovers_cacheability_without_inline_retain() { - let _guard = quota_user_lock_test_scope(); - let map = QUOTA_USER_LOCKS.get_or_init(DashMap::new); - map.clear(); - - let mut retained = Vec::with_capacity(QUOTA_USER_LOCKS_MAX); - let prefix = format!("quota-evict-saturated-{}", std::process::id()); - for idx in 0..QUOTA_USER_LOCKS_MAX { - retained.push(quota_user_lock(&format!("{prefix}-{idx}"))); - } - - assert_eq!(map.len(), QUOTA_USER_LOCKS_MAX); - - let overflow_user = format!("quota-evict-overflow-user-{}", std::process::id()); - let overflow_before = quota_user_lock(&overflow_user); - assert!( - map.get(&overflow_user).is_none(), - "saturated map must initially route new user to overflow stripe" - ); - - drop(retained); - - let evictor = spawn_quota_user_lock_evictor_for_tests(Duration::from_millis(4)); - - timeout(Duration::from_millis(400), async { - loop { - if map.len() < QUOTA_USER_LOCKS_MAX { - break; - } - tokio::time::sleep(Duration::from_millis(5)).await; - } - }) - .await - .expect("periodic evictor must reclaim stale saturated entries"); - - let overflow_after = quota_user_lock(&overflow_user); - assert!( - map.get(&overflow_user).is_some(), - "after eviction, overflow user should become cacheable again" - ); - assert!( - Arc::strong_count(&overflow_after) >= 2, - "cacheable lock should be held by map and caller" - ); - - drop(overflow_before); - drop(overflow_after); - evictor.abort(); -} diff --git a/src/proxy/tests/relay_quota_lock_identity_security_tests.rs b/src/proxy/tests/relay_quota_lock_identity_security_tests.rs deleted file mode 100644 index f717f54..0000000 --- a/src/proxy/tests/relay_quota_lock_identity_security_tests.rs +++ /dev/null @@ -1,135 +0,0 @@ -use super::*; -use crate::stats::Stats; -use std::pin::Pin; -use std::sync::Arc; -use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; -use std::task::Waker; -use std::task::{Context, Poll}; -use tokio::io::{AsyncRead, AsyncWrite, ReadBuf}; - -#[derive(Default)] -struct WakeCounter { - wakes: AtomicUsize, -} - -impl std::task::Wake for WakeCounter { - fn wake(self: Arc) { - self.wakes.fetch_add(1, Ordering::Relaxed); - } - - fn wake_by_ref(self: &Arc) { - self.wakes.fetch_add(1, Ordering::Relaxed); - } -} - -fn build_context() -> (Arc, Context<'static>) { - let wake_counter = Arc::new(WakeCounter::default()); - let waker = Waker::from(Arc::clone(&wake_counter)); - // Context stores a reference; leak one Waker for deterministic test scope. - let leaked_waker: &'static Waker = Box::leak(Box::new(waker)); - (wake_counter, Context::from_waker(leaked_waker)) -} - -#[tokio::test] -async fn adversarial_map_churn_cannot_bypass_held_writer_lock() { - let _guard = quota_user_lock_test_scope(); - let map = QUOTA_USER_LOCKS.get_or_init(DashMap::new); - map.clear(); - - let user = "quota-identity-writer-user"; - let held_lock = quota_user_lock(user); - let _held_guard = held_lock - .try_lock() - .expect("test must hold initial user lock before StatsIo poll"); - - let mut io = StatsIo::new( - tokio::io::sink(), - Arc::new(SharedCounters::new()), - Arc::new(Stats::new()), - user.to_string(), - Some(1024), - Arc::new(AtomicBool::new(false)), - tokio::time::Instant::now(), - ); - - map.clear(); - let churned_lock = quota_user_lock(user); - assert!( - !Arc::ptr_eq(&held_lock, &churned_lock), - "precondition: map churn should produce a distinct lock identity" - ); - - let (_wake_counter, mut cx) = build_context(); - let poll = Pin::new(&mut io).poll_write(&mut cx, &[0x11, 0x22, 0x33, 0x44]); - - assert!( - matches!(poll, Poll::Pending), - "writer must remain pending on the originally-held lock identity" - ); -} - -#[tokio::test] -async fn adversarial_map_churn_cannot_bypass_held_reader_lock() { - let _guard = quota_user_lock_test_scope(); - let map = QUOTA_USER_LOCKS.get_or_init(DashMap::new); - map.clear(); - - let user = "quota-identity-reader-user"; - let held_lock = quota_user_lock(user); - let _held_guard = held_lock - .try_lock() - .expect("test must hold initial user lock before StatsIo poll"); - - let mut io = StatsIo::new( - tokio::io::empty(), - Arc::new(SharedCounters::new()), - Arc::new(Stats::new()), - user.to_string(), - Some(1024), - Arc::new(AtomicBool::new(false)), - tokio::time::Instant::now(), - ); - - map.clear(); - let churned_lock = quota_user_lock(user); - assert!( - !Arc::ptr_eq(&held_lock, &churned_lock), - "precondition: map churn should produce a distinct lock identity" - ); - - let (_wake_counter, mut cx) = build_context(); - let mut storage = [0u8; 8]; - let mut read_buf = ReadBuf::new(&mut storage); - let poll = Pin::new(&mut io).poll_read(&mut cx, &mut read_buf); - - assert!( - matches!(poll, Poll::Pending), - "reader must remain pending on the originally-held lock identity" - ); -} - -#[tokio::test] -async fn business_no_lock_contention_keeps_writer_progress() { - let _guard = quota_user_lock_test_scope(); - let map = QUOTA_USER_LOCKS.get_or_init(DashMap::new); - map.clear(); - - let user = "quota-identity-progress-user"; - let mut io = StatsIo::new( - tokio::io::sink(), - Arc::new(SharedCounters::new()), - Arc::new(Stats::new()), - user.to_string(), - Some(1024), - Arc::new(AtomicBool::new(false)), - tokio::time::Instant::now(), - ); - - let (_wake_counter, mut cx) = build_context(); - let poll = Pin::new(&mut io).poll_write(&mut cx, &[0xAA, 0xBB]); - - assert!( - matches!(poll, Poll::Ready(Ok(2))), - "writer should progress immediately without contention" - ); -} diff --git a/src/proxy/tests/relay_quota_lock_pressure_adversarial_tests.rs b/src/proxy/tests/relay_quota_lock_pressure_adversarial_tests.rs deleted file mode 100644 index 5687965..0000000 --- a/src/proxy/tests/relay_quota_lock_pressure_adversarial_tests.rs +++ /dev/null @@ -1,440 +0,0 @@ -use super::*; -use crate::error::ProxyError; -use crate::stats::Stats; -use crate::stream::BufferPool; -use dashmap::DashMap; -use std::sync::Arc; -use std::sync::atomic::AtomicBool; -use std::time::Duration; -use tokio::io::{AsyncReadExt, AsyncWriteExt, duplex}; -use tokio::sync::Barrier; -use tokio::time::Instant; - -#[test] -fn quota_lock_same_user_returns_same_arc_instance() { - let _guard = super::quota_user_lock_test_scope(); - let map = QUOTA_USER_LOCKS.get_or_init(DashMap::new); - map.clear(); - - let a = quota_user_lock("quota-lock-same-user"); - let b = quota_user_lock("quota-lock-same-user"); - assert!(Arc::ptr_eq(&a, &b)); -} - -#[test] -fn quota_lock_parallel_same_user_reuses_single_lock() { - let _guard = super::quota_user_lock_test_scope(); - let map = QUOTA_USER_LOCKS.get_or_init(DashMap::new); - map.clear(); - - let user = "quota-lock-parallel-same"; - let mut handles = Vec::new(); - - for _ in 0..64 { - handles.push(std::thread::spawn(move || quota_user_lock(user))); - } - - let first = handles - .remove(0) - .join() - .expect("thread must return lock handle"); - - for handle in handles { - let got = handle.join().expect("thread must return lock handle"); - assert!(Arc::ptr_eq(&first, &got)); - } -} - -#[test] -fn quota_lock_unique_users_materialize_distinct_entries() { - let _guard = super::quota_user_lock_test_scope(); - let map = QUOTA_USER_LOCKS.get_or_init(DashMap::new); - - map.clear(); - - let base = format!("quota-lock-distinct-{}", std::process::id()); - let users: Vec = (0..(QUOTA_USER_LOCKS_MAX / 2)) - .map(|idx| format!("{base}-{idx}")) - .collect(); - - for user in &users { - let _ = quota_user_lock(user); - } - - for user in &users { - assert!( - map.get(user).is_some(), - "lock cache must contain entry for {user}" - ); - } -} - -#[test] -fn quota_lock_unique_churn_stress_keeps_all_inserted_keys_addressable() { - let _guard = super::quota_user_lock_test_scope(); - let map = QUOTA_USER_LOCKS.get_or_init(DashMap::new); - - map.clear(); - - let base = format!("quota-lock-churn-{}", std::process::id()); - for idx in 0..(QUOTA_USER_LOCKS_MAX + 256) { - let _ = quota_user_lock(&format!("{base}-{idx}")); - } - - assert!( - map.len() <= QUOTA_USER_LOCKS_MAX, - "quota lock cache must stay bounded under unique-user churn" - ); -} - -#[test] -fn quota_lock_saturation_returns_stable_overflow_lock_without_cache_growth() { - let _guard = super::quota_user_lock_test_scope(); - let map = QUOTA_USER_LOCKS.get_or_init(DashMap::new); - map.clear(); - - let prefix = format!("quota-held-{}", std::process::id()); - let mut retained = Vec::with_capacity(QUOTA_USER_LOCKS_MAX); - for idx in 0..QUOTA_USER_LOCKS_MAX { - retained.push(quota_user_lock(&format!("{prefix}-{idx}"))); - } - - assert_eq!( - map.len(), - QUOTA_USER_LOCKS_MAX, - "cache must be saturated for overflow check" - ); - - let overflow_user = format!("quota-overflow-{}", std::process::id()); - let overflow_a = quota_user_lock(&overflow_user); - let overflow_b = quota_user_lock(&overflow_user); - - assert_eq!( - map.len(), - QUOTA_USER_LOCKS_MAX, - "overflow path must not grow lock cache" - ); - assert!( - map.get(&overflow_user).is_none(), - "overflow user lock must stay outside bounded cache under saturation" - ); - assert!( - Arc::ptr_eq(&overflow_a, &overflow_b), - "overflow user must receive stable striped overflow lock while saturated" - ); - - drop(retained); -} - -#[test] -fn quota_lock_reclaims_unreferenced_entries_after_explicit_eviction_pass() { - let _guard = super::quota_user_lock_test_scope(); - let map = QUOTA_USER_LOCKS.get_or_init(DashMap::new); - map.clear(); - - // Saturate with retained strong references first so parallel tests cannot - // reclaim our fixture entries before we validate the reclaim path. - let prefix = format!("quota-reclaim-drop-{}", std::process::id()); - let mut retained = Vec::with_capacity(QUOTA_USER_LOCKS_MAX); - for idx in 0..QUOTA_USER_LOCKS_MAX { - retained.push(quota_user_lock(&format!("{prefix}-{idx}"))); - } - - drop(retained); - - quota_user_lock_evict(); - - let overflow_user = format!("quota-reclaim-overflow-{}", std::process::id()); - let overflow = quota_user_lock(&overflow_user); - - assert!( - map.get(&overflow_user).is_some(), - "after reclaiming stale entries, overflow user should become cacheable" - ); - assert!( - Arc::strong_count(&overflow) >= 2, - "cacheable overflow lock should be held by both map and caller" - ); -} - -#[test] -fn quota_lock_saturated_same_user_must_not_return_distinct_locks() { - let _guard = super::quota_user_lock_test_scope(); - let map = QUOTA_USER_LOCKS.get_or_init(DashMap::new); - map.clear(); - - let mut retained = Vec::with_capacity(QUOTA_USER_LOCKS_MAX); - for idx in 0..QUOTA_USER_LOCKS_MAX { - retained.push(quota_user_lock(&format!( - "quota-saturated-held-{}-{idx}", - std::process::id() - ))); - } - - let overflow_user = format!("quota-saturated-same-user-{}", std::process::id()); - let a = quota_user_lock(&overflow_user); - let b = quota_user_lock(&overflow_user); - - assert!( - Arc::ptr_eq(&a, &b), - "same user must not receive distinct locks under saturation because that enables quota race bypass" - ); - - drop(retained); -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 4)] -async fn quota_lock_saturation_concurrent_same_user_never_overshoots_quota() { - let _guard = super::quota_user_lock_test_scope(); - let map = QUOTA_USER_LOCKS.get_or_init(DashMap::new); - map.clear(); - - let mut retained = Vec::with_capacity(QUOTA_USER_LOCKS_MAX); - for idx in 0..QUOTA_USER_LOCKS_MAX { - retained.push(quota_user_lock(&format!( - "quota-saturated-race-held-{}-{idx}", - std::process::id() - ))); - } - - let stats = Arc::new(Stats::new()); - let user = format!("quota-saturated-race-user-{}", std::process::id()); - let gate = Arc::new(Barrier::new(2)); - - let worker = |label: u8, stats: Arc, user: String, gate: Arc| { - tokio::spawn(async move { - let counters = Arc::new(SharedCounters::new()); - let quota_exceeded = Arc::new(AtomicBool::new(false)); - let mut io = StatsIo::new( - tokio::io::sink(), - counters, - Arc::clone(&stats), - user, - Some(1), - quota_exceeded, - Instant::now(), - ); - gate.wait().await; - io.write_all(&[label]).await - }) - }; - - let one = worker(0x11, Arc::clone(&stats), user.clone(), Arc::clone(&gate)); - let two = worker(0x22, Arc::clone(&stats), user.clone(), Arc::clone(&gate)); - - let _ = tokio::time::timeout(Duration::from_secs(2), async { - let _ = one.await.expect("task one must not panic"); - let _ = two.await.expect("task two must not panic"); - }) - .await - .expect("quota race workers must complete"); - - assert!( - stats.get_user_total_octets(&user) <= 1, - "saturated lock path must never overshoot quota for same user" - ); - - drop(retained); -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 4)] -async fn quota_lock_saturation_stress_same_user_never_overshoots_quota() { - let _guard = super::quota_user_lock_test_scope(); - let map = QUOTA_USER_LOCKS.get_or_init(DashMap::new); - map.clear(); - - let mut retained = Vec::with_capacity(QUOTA_USER_LOCKS_MAX); - for idx in 0..QUOTA_USER_LOCKS_MAX { - retained.push(quota_user_lock(&format!( - "quota-saturated-stress-held-{}-{idx}", - std::process::id() - ))); - } - - for round in 0..128u32 { - let stats = Arc::new(Stats::new()); - let user = format!("quota-saturated-stress-user-{}-{round}", std::process::id()); - let gate = Arc::new(Barrier::new(2)); - - let one = { - let stats = Arc::clone(&stats); - let user = user.clone(); - let gate = Arc::clone(&gate); - tokio::spawn(async move { - let counters = Arc::new(SharedCounters::new()); - let quota_exceeded = Arc::new(AtomicBool::new(false)); - let mut io = StatsIo::new( - tokio::io::sink(), - counters, - Arc::clone(&stats), - user, - Some(1), - quota_exceeded, - Instant::now(), - ); - gate.wait().await; - io.write_all(&[0x31]).await - }) - }; - - let two = { - let stats = Arc::clone(&stats); - let user = user.clone(); - let gate = Arc::clone(&gate); - tokio::spawn(async move { - let counters = Arc::new(SharedCounters::new()); - let quota_exceeded = Arc::new(AtomicBool::new(false)); - let mut io = StatsIo::new( - tokio::io::sink(), - counters, - Arc::clone(&stats), - user, - Some(1), - quota_exceeded, - Instant::now(), - ); - gate.wait().await; - io.write_all(&[0x32]).await - }) - }; - - let _ = one.await.expect("stress task one must not panic"); - let _ = two.await.expect("stress task two must not panic"); - - assert!( - stats.get_user_total_octets(&user) <= 1, - "round {round}: saturated path must not overshoot quota" - ); - } - - drop(retained); -} - -#[test] -fn quota_error_classifier_accepts_internal_quota_sentinel_only() { - let err = quota_io_error(); - assert!(is_quota_io_error(&err)); -} - -#[test] -fn quota_error_classifier_rejects_plain_permission_denied() { - let err = std::io::Error::new(std::io::ErrorKind::PermissionDenied, "permission denied"); - assert!(!is_quota_io_error(&err)); -} - -#[test] -fn quota_lock_test_scope_recovers_after_guard_poison() { - let poison_result = std::thread::spawn(|| { - let _guard = super::quota_user_lock_test_scope(); - panic!("intentional test-only guard poison"); - }) - .join(); - assert!(poison_result.is_err(), "poison setup thread must panic"); - - let _guard = super::quota_user_lock_test_scope(); - let map = QUOTA_USER_LOCKS.get_or_init(DashMap::new); - map.clear(); - - let a = quota_user_lock("quota-lock-poison-recovery-user"); - let b = quota_user_lock("quota-lock-poison-recovery-user"); - assert!(Arc::ptr_eq(&a, &b)); -} - -#[tokio::test] -async fn quota_lock_integration_zero_quota_cuts_off_without_forwarding() { - let stats = Arc::new(Stats::new()); - let user = "quota-zero-user"; - - 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(0), - Arc::new(BufferPool::new()), - )); - - client_peer - .write_all(b"x") - .await - .expect("client write must succeed"); - - let mut probe = [0u8; 1]; - let forwarded = - tokio::time::timeout(Duration::from_millis(80), server_peer.read(&mut probe)).await; - if let Ok(Ok(n)) = forwarded { - assert_eq!(n, 0, "zero quota path must not forward payload bytes"); - } - - let result = tokio::time::timeout(Duration::from_secs(2), relay) - .await - .expect("relay must terminate under zero quota") - .expect("relay task must not panic"); - assert!(matches!(result, Err(ProxyError::DataQuotaExceeded { .. }))); -} - -#[tokio::test] -async fn quota_lock_integration_no_quota_relays_both_directions_under_burst() { - let stats = Arc::new(Stats::new()); - - let (mut client_peer, relay_client) = duplex(8192); - let (relay_server, mut server_peer) = duplex(8192); - 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, - "quota-none-burst-user", - Arc::clone(&stats), - None, - Arc::new(BufferPool::new()), - )); - - let c2s = vec![0xA5; 2048]; - let s2c = vec![0x5A; 1536]; - - client_peer - .write_all(&c2s) - .await - .expect("client burst write must succeed"); - let mut got_c2s = vec![0u8; c2s.len()]; - server_peer - .read_exact(&mut got_c2s) - .await - .expect("server must receive c2s burst"); - assert_eq!(got_c2s, c2s); - - server_peer - .write_all(&s2c) - .await - .expect("server burst write must succeed"); - let mut got_s2c = vec![0u8; s2c.len()]; - client_peer - .read_exact(&mut got_s2c) - .await - .expect("client must receive s2c burst"); - assert_eq!(got_s2c, s2c); - - drop(client_peer); - drop(server_peer); - - let done = tokio::time::timeout(Duration::from_secs(2), relay) - .await - .expect("relay must terminate after peers close") - .expect("relay task must not panic"); - assert!(done.is_ok()); -} diff --git a/src/proxy/tests/relay_quota_retry_allocation_latency_security_tests.rs b/src/proxy/tests/relay_quota_retry_allocation_latency_security_tests.rs deleted file mode 100644 index 447a090..0000000 --- a/src/proxy/tests/relay_quota_retry_allocation_latency_security_tests.rs +++ /dev/null @@ -1,249 +0,0 @@ -use super::*; -use crate::stats::Stats; -use std::pin::Pin; -use std::sync::Arc; -use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; -use std::task::{Context, Waker}; -use tokio::io::AsyncWriteExt; -use tokio::time::{Duration, Instant, timeout}; - -#[derive(Default)] -struct WakeCounter { - wakes: AtomicUsize, -} - -impl std::task::Wake for WakeCounter { - fn wake(self: Arc) { - self.wakes.fetch_add(1, Ordering::Relaxed); - } - - fn wake_by_ref(self: &Arc) { - self.wakes.fetch_add(1, Ordering::Relaxed); - } -} - -fn quota_test_guard() -> impl Drop { - super::quota_user_lock_test_scope() -} - -fn build_context() -> (Arc, Context<'static>) { - let wake_counter = Arc::new(WakeCounter::default()); - let waker = Waker::from(Arc::clone(&wake_counter)); - let leaked_waker: &'static Waker = Box::leak(Box::new(waker)); - (wake_counter, Context::from_waker(leaked_waker)) -} - -fn sleep_slot_ptr(slot: &Option>>) -> usize { - slot.as_ref() - .map(|sleep| (&**sleep) as *const tokio::time::Sleep as usize) - .unwrap_or(0) -} - -#[tokio::test] -async fn tdd_single_pending_timer_does_not_allocate_on_each_repoll() { - let _guard = quota_test_guard(); - - let user = format!("retry-alloc-single-pending-{}", std::process::id()); - let lock = quota_user_lock(&user); - let held_guard = lock - .try_lock() - .expect("test must hold local lock to force retry scheduling"); - - reset_quota_retry_sleep_allocs_for_tests(); - - let mut io = StatsIo::new( - tokio::io::sink(), - Arc::new(SharedCounters::new()), - Arc::new(Stats::new()), - user, - Some(2048), - Arc::new(AtomicBool::new(false)), - Instant::now(), - ); - - let (_wake_counter, mut cx) = build_context(); - - let first = Pin::new(&mut io).poll_write(&mut cx, &[0xA1]); - assert!(first.is_pending()); - let allocs_after_first = quota_retry_sleep_allocs_for_tests(); - let ptr_after_first = sleep_slot_ptr(&io.quota_write_retry_sleep); - - let second = Pin::new(&mut io).poll_write(&mut cx, &[0xA2]); - assert!(second.is_pending()); - let allocs_after_second = quota_retry_sleep_allocs_for_tests(); - let ptr_after_second = sleep_slot_ptr(&io.quota_write_retry_sleep); - - assert_eq!(allocs_after_first, 1, "first pending poll must allocate one timer"); - assert_eq!( - allocs_after_second, 1, - "repoll while the same timer is pending must not allocate again" - ); - assert_eq!( - ptr_after_first, ptr_after_second, - "repoll while pending should retain the same timer allocation" - ); - - drop(held_guard); -} - -#[tokio::test] -async fn tdd_retry_cycle_allocates_once_per_fired_timer_cycle_not_per_poll() { - let _guard = quota_test_guard(); - - let user = format!("retry-alloc-per-cycle-{}", std::process::id()); - let lock = quota_user_lock(&user); - let held_guard = lock - .try_lock() - .expect("test must hold local lock to keep write path pending"); - - reset_quota_retry_sleep_allocs_for_tests(); - - let mut io = StatsIo::new( - tokio::io::sink(), - Arc::new(SharedCounters::new()), - Arc::new(Stats::new()), - user, - Some(2048), - Arc::new(AtomicBool::new(false)), - Instant::now(), - ); - - let (wake_counter, mut cx) = build_context(); - - let mut polls = 0u64; - let mut observed_wakes = 0usize; - let started = Instant::now(); - while started.elapsed() < Duration::from_millis(70) { - let poll = Pin::new(&mut io).poll_write(&mut cx, &[0xB1]); - polls = polls.saturating_add(1); - assert!(poll.is_pending()); - - let wakes = wake_counter.wakes.load(Ordering::Relaxed); - if wakes > observed_wakes { - observed_wakes = wakes; - } - - tokio::time::sleep(Duration::from_millis(1)).await; - } - - let allocs = quota_retry_sleep_allocs_for_tests(); - assert!(allocs >= 2, "multiple fired cycles should allocate multiple timers"); - assert!( - allocs < polls, - "timer allocations must be bounded by cycles, not by every repoll (allocs={allocs}, polls={polls})" - ); - - drop(held_guard); -} - -#[tokio::test] -async fn adversarial_backoff_latency_envelope_stays_bounded_under_contention() { - let _guard = quota_test_guard(); - - let user = format!("retry-latency-envelope-{}", std::process::id()); - let lock = quota_user_lock(&user); - let held_guard = lock - .try_lock() - .expect("test must hold local lock for sustained contention"); - - reset_quota_retry_sleep_allocs_for_tests(); - - let mut io = StatsIo::new( - tokio::io::sink(), - Arc::new(SharedCounters::new()), - Arc::new(Stats::new()), - user, - Some(2048), - Arc::new(AtomicBool::new(false)), - Instant::now(), - ); - - let (wake_counter, mut cx) = build_context(); - - let first = Pin::new(&mut io).poll_write(&mut cx, &[0xC1]); - assert!(first.is_pending()); - - let started = Instant::now(); - let mut last_wakes = 0usize; - let mut wake_instants = Vec::new(); - - while started.elapsed() < Duration::from_millis(120) { - let wakes = wake_counter.wakes.load(Ordering::Relaxed); - if wakes > last_wakes { - last_wakes = wakes; - wake_instants.push(Instant::now()); - let pending = Pin::new(&mut io).poll_write(&mut cx, &[0xC2]); - assert!(pending.is_pending()); - } - tokio::time::sleep(Duration::from_millis(1)).await; - } - - let mut max_gap = Duration::from_millis(0); - for idx in 1..wake_instants.len() { - let gap = wake_instants[idx].saturating_duration_since(wake_instants[idx - 1]); - if gap > max_gap { - max_gap = gap; - } - } - - assert!( - max_gap <= Duration::from_millis(35), - "retry wake gap must remain bounded in test profile; observed max gap={max_gap:?}" - ); - assert!( - quota_retry_sleep_allocs_for_tests() <= 16, - "allocation cycles must remain bounded during a short contention window" - ); - - drop(held_guard); -} - -#[tokio::test] -async fn micro_benchmark_release_to_completion_latency_stays_bounded() { - let _guard = quota_test_guard(); - - let rounds = 96usize; - let mut samples_ms = Vec::with_capacity(rounds); - - for round in 0..rounds { - let user = format!("retry-release-latency-{}-{round}", std::process::id()); - let lock = quota_user_lock(&user); - let held_guard = lock - .try_lock() - .expect("test must hold local lock before spawning blocked writer"); - - let writer = tokio::spawn(async move { - let mut io = StatsIo::new( - tokio::io::sink(), - Arc::new(SharedCounters::new()), - Arc::new(Stats::new()), - user, - Some(2048), - Arc::new(AtomicBool::new(false)), - Instant::now(), - ); - io.write_all(&[0xD1]).await - }); - - tokio::time::sleep(Duration::from_millis(2)).await; - let release_at = Instant::now(); - drop(held_guard); - - let done = timeout(Duration::from_millis(120), writer) - .await - .expect("blocked writer must complete after release") - .expect("writer task must not panic"); - assert!(done.is_ok()); - - samples_ms.push(release_at.elapsed().as_millis() as u64); - } - - samples_ms.sort_unstable(); - let p95_idx = ((samples_ms.len() * 95) / 100).min(samples_ms.len().saturating_sub(1)); - let p95_ms = samples_ms[p95_idx]; - - assert!( - p95_ms <= 40, - "contention release->completion p95 must stay bounded; p95_ms={p95_ms}, samples={samples_ms:?}" - ); -} diff --git a/src/proxy/tests/relay_quota_retry_backoff_benchmark_security_tests.rs b/src/proxy/tests/relay_quota_retry_backoff_benchmark_security_tests.rs deleted file mode 100644 index 7083eb2..0000000 --- a/src/proxy/tests/relay_quota_retry_backoff_benchmark_security_tests.rs +++ /dev/null @@ -1,241 +0,0 @@ -use super::*; -use crate::stats::Stats; -use dashmap::DashMap; -use std::pin::Pin; -use std::sync::Arc; -use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; -use std::task::{Context, Waker}; -use tokio::io::ReadBuf; -use tokio::time::{Duration, Instant}; - -#[derive(Default)] -struct WakeCounter { - wakes: AtomicUsize, -} - -impl std::task::Wake for WakeCounter { - fn wake(self: Arc) { - self.wakes.fetch_add(1, Ordering::Relaxed); - } - - fn wake_by_ref(self: &Arc) { - self.wakes.fetch_add(1, Ordering::Relaxed); - } -} - -fn quota_test_guard() -> impl Drop { - super::quota_user_lock_test_scope() -} - -fn saturate_quota_user_locks() -> Vec>> { - let map = QUOTA_USER_LOCKS.get_or_init(DashMap::new); - map.clear(); - - let mut retained = Vec::with_capacity(QUOTA_USER_LOCKS_MAX); - for idx in 0..QUOTA_USER_LOCKS_MAX { - retained.push(quota_user_lock(&format!("quota-retry-bench-saturate-{idx}"))); - } - retained -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 4)] -async fn stress_contention_wake_rate_decays_with_backoff_curve() { - let _guard = quota_test_guard(); - - let _retained = saturate_quota_user_locks(); - let user = format!("quota-backoff-bench-{}", std::process::id()); - let stats = Arc::new(Stats::new()); - - let lock = quota_user_lock(&user); - let held_guard = lock - .try_lock() - .expect("test must hold quota lock before benchmark run"); - - let waiters = 64usize; - let mut ios = Vec::with_capacity(waiters); - let mut wake_counters = Vec::with_capacity(waiters); - - for _ in 0..waiters { - ios.push(StatsIo::new( - tokio::io::sink(), - Arc::new(SharedCounters::new()), - Arc::clone(&stats), - user.clone(), - Some(4096), - Arc::new(AtomicBool::new(false)), - tokio::time::Instant::now(), - )); - } - - for io in &mut ios { - let counter = Arc::new(WakeCounter::default()); - let waker = Waker::from(Arc::clone(&counter)); - let mut cx = Context::from_waker(&waker); - let pending = Pin::new(io).poll_write(&mut cx, &[0x71]); - assert!(pending.is_pending()); - wake_counters.push(counter); - } - - let mut observed = vec![0usize; waiters]; - let start = Instant::now(); - let mut wakes_at_40ms = 0usize; - let mut wakes_at_160ms = 0usize; - - while start.elapsed() < Duration::from_millis(200) { - for (idx, counter) in wake_counters.iter().enumerate() { - let wakes = counter.wakes.load(Ordering::Relaxed); - if wakes > observed[idx] { - observed[idx] = wakes; - let waker = Waker::from(Arc::clone(counter)); - let mut cx = Context::from_waker(&waker); - let pending = Pin::new(&mut ios[idx]).poll_write(&mut cx, &[0x72]); - assert!(pending.is_pending()); - } - } - - let elapsed = start.elapsed(); - if elapsed >= Duration::from_millis(40) && wakes_at_40ms == 0 { - wakes_at_40ms = wake_counters - .iter() - .map(|counter| counter.wakes.load(Ordering::Relaxed)) - .sum(); - } - if elapsed >= Duration::from_millis(160) && wakes_at_160ms == 0 { - wakes_at_160ms = wake_counters - .iter() - .map(|counter| counter.wakes.load(Ordering::Relaxed)) - .sum(); - } - - tokio::time::sleep(Duration::from_millis(1)).await; - } - - let total_wakes: usize = wake_counters - .iter() - .map(|counter| counter.wakes.load(Ordering::Relaxed)) - .sum(); - - let wakes_at_200ms = total_wakes; - let early_window_wakes = wakes_at_40ms; - let late_window_wakes = wakes_at_200ms.saturating_sub(wakes_at_160ms); - - assert!( - total_wakes <= waiters * 28, - "backoff benchmark exceeded wake budget; waiters={waiters}, wakes={total_wakes}" - ); - - assert!( - early_window_wakes > 0, - "benchmark failed to observe early contention wakes" - ); - - assert!( - late_window_wakes * 4 <= early_window_wakes * 3, - "wake-rate decay invariant violated; early_0_40ms={early_window_wakes}, late_160_200ms={late_window_wakes}, total={total_wakes}" - ); - - drop(held_guard); -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 4)] -async fn stress_read_contention_wake_rate_decays_with_backoff_curve() { - let _guard = quota_test_guard(); - - let _retained = saturate_quota_user_locks(); - let user = format!("quota-backoff-read-bench-{}", std::process::id()); - let stats = Arc::new(Stats::new()); - - let lock = quota_user_lock(&user); - let held_guard = lock - .try_lock() - .expect("test must hold quota lock before read benchmark run"); - - let waiters = 64usize; - let mut ios = Vec::with_capacity(waiters); - let mut wake_counters = Vec::with_capacity(waiters); - - for _ in 0..waiters { - ios.push(StatsIo::new( - tokio::io::empty(), - Arc::new(SharedCounters::new()), - Arc::clone(&stats), - user.clone(), - Some(4096), - Arc::new(AtomicBool::new(false)), - tokio::time::Instant::now(), - )); - } - - for io in &mut ios { - let counter = Arc::new(WakeCounter::default()); - let waker = Waker::from(Arc::clone(&counter)); - let mut cx = Context::from_waker(&waker); - let mut storage = [0u8; 1]; - let mut buf = ReadBuf::new(&mut storage); - let pending = Pin::new(io).poll_read(&mut cx, &mut buf); - assert!(pending.is_pending()); - wake_counters.push(counter); - } - - let mut observed = vec![0usize; waiters]; - let start = Instant::now(); - let mut wakes_at_40ms = 0usize; - let mut wakes_at_160ms = 0usize; - - while start.elapsed() < Duration::from_millis(200) { - for (idx, counter) in wake_counters.iter().enumerate() { - let wakes = counter.wakes.load(Ordering::Relaxed); - if wakes > observed[idx] { - observed[idx] = wakes; - let waker = Waker::from(Arc::clone(counter)); - let mut cx = Context::from_waker(&waker); - let mut storage = [0u8; 1]; - let mut buf = ReadBuf::new(&mut storage); - let pending = Pin::new(&mut ios[idx]).poll_read(&mut cx, &mut buf); - assert!(pending.is_pending()); - } - } - - let elapsed = start.elapsed(); - if elapsed >= Duration::from_millis(40) && wakes_at_40ms == 0 { - wakes_at_40ms = wake_counters - .iter() - .map(|counter| counter.wakes.load(Ordering::Relaxed)) - .sum(); - } - if elapsed >= Duration::from_millis(160) && wakes_at_160ms == 0 { - wakes_at_160ms = wake_counters - .iter() - .map(|counter| counter.wakes.load(Ordering::Relaxed)) - .sum(); - } - - tokio::time::sleep(Duration::from_millis(1)).await; - } - - let total_wakes: usize = wake_counters - .iter() - .map(|counter| counter.wakes.load(Ordering::Relaxed)) - .sum(); - - let wakes_at_200ms = total_wakes; - let early_window_wakes = wakes_at_40ms; - let late_window_wakes = wakes_at_200ms.saturating_sub(wakes_at_160ms); - - assert!( - total_wakes <= waiters * 28, - "read backoff benchmark exceeded wake budget; waiters={waiters}, wakes={total_wakes}" - ); - - assert!( - early_window_wakes > 0, - "read benchmark failed to observe early contention wakes" - ); - - assert!( - late_window_wakes * 4 <= early_window_wakes * 3, - "read wake-rate decay invariant violated; early_0_40ms={early_window_wakes}, late_160_200ms={late_window_wakes}, total={total_wakes}" - ); - - drop(held_guard); -} diff --git a/src/proxy/tests/relay_quota_retry_backoff_security_tests.rs b/src/proxy/tests/relay_quota_retry_backoff_security_tests.rs deleted file mode 100644 index 7f1e451..0000000 --- a/src/proxy/tests/relay_quota_retry_backoff_security_tests.rs +++ /dev/null @@ -1,339 +0,0 @@ -use super::*; -use crate::stats::Stats; -use dashmap::DashMap; -use std::pin::Pin; -use std::sync::Arc; -use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; -use std::task::{Context, Waker}; -use tokio::io::ReadBuf; -use tokio::time::{Duration, Instant}; - -#[derive(Default)] -struct WakeCounter { - wakes: AtomicUsize, -} - -impl std::task::Wake for WakeCounter { - fn wake(self: Arc) { - self.wakes.fetch_add(1, Ordering::Relaxed); - } - - fn wake_by_ref(self: &Arc) { - self.wakes.fetch_add(1, Ordering::Relaxed); - } -} - -fn quota_test_guard() -> impl Drop { - super::quota_user_lock_test_scope() -} - -fn saturate_quota_user_locks() -> Vec>> { - let map = QUOTA_USER_LOCKS.get_or_init(DashMap::new); - map.clear(); - - let mut retained = Vec::with_capacity(QUOTA_USER_LOCKS_MAX); - for idx in 0..QUOTA_USER_LOCKS_MAX { - retained.push(quota_user_lock(&format!("quota-retry-backoff-saturate-{idx}"))); - } - retained -} - -#[tokio::test] -async fn positive_uncontended_writer_keeps_retry_wakes_zero() { - let _guard = quota_test_guard(); - - let stats = Arc::new(Stats::new()); - let mut io = StatsIo::new( - tokio::io::sink(), - Arc::new(SharedCounters::new()), - Arc::clone(&stats), - "quota-backoff-positive".to_string(), - Some(2048), - Arc::new(AtomicBool::new(false)), - tokio::time::Instant::now(), - ); - - let wake_counter = Arc::new(WakeCounter::default()); - let waker = Waker::from(Arc::clone(&wake_counter)); - let mut cx = Context::from_waker(&waker); - - let poll = Pin::new(&mut io).poll_write(&mut cx, &[0x41, 0x42]); - assert!(poll.is_ready(), "uncontended writer must complete immediately"); - assert_eq!( - wake_counter.wakes.load(Ordering::Relaxed), - 0, - "uncontended path must not schedule deferred contention wakes" - ); -} - -#[tokio::test] -async fn adversarial_writer_sustained_contention_executor_repoll_is_rate_limited() { - let _guard = quota_test_guard(); - - let _retained = saturate_quota_user_locks(); - let user = "quota-backoff-adversarial-writer"; - let stats = Arc::new(Stats::new()); - - let lock = quota_user_lock(user); - let held_guard = lock - .try_lock() - .expect("test must hold quota lock before polling writer"); - - let mut io = StatsIo::new( - tokio::io::sink(), - Arc::new(SharedCounters::new()), - Arc::clone(&stats), - user.to_string(), - Some(2048), - Arc::new(AtomicBool::new(false)), - tokio::time::Instant::now(), - ); - - let wake_counter = Arc::new(WakeCounter::default()); - let waker = Waker::from(Arc::clone(&wake_counter)); - let mut cx = Context::from_waker(&waker); - - let first = Pin::new(&mut io).poll_write(&mut cx, &[0xAA]); - assert!(first.is_pending()); - - let start = Instant::now(); - let mut observed = 0usize; - while start.elapsed() < Duration::from_millis(80) { - let wakes = wake_counter.wakes.load(Ordering::Relaxed); - if wakes > observed { - observed = wakes; - let pending = Pin::new(&mut io).poll_write(&mut cx, &[0xAB]); - assert!(pending.is_pending()); - } - tokio::time::sleep(Duration::from_millis(1)).await; - } - - assert!( - wake_counter.wakes.load(Ordering::Relaxed) <= 16, - "sustained contention must be rate limited; observed wakes={} in 80ms", - wake_counter.wakes.load(Ordering::Relaxed) - ); - - drop(held_guard); - let ready = Pin::new(&mut io).poll_write(&mut cx, &[0xAC]); - assert!(ready.is_ready()); -} - -#[tokio::test] -async fn adversarial_reader_sustained_contention_executor_repoll_is_rate_limited() { - let _guard = quota_test_guard(); - - let _retained = saturate_quota_user_locks(); - let user = "quota-backoff-adversarial-reader"; - let stats = Arc::new(Stats::new()); - - let lock = quota_user_lock(user); - let held_guard = lock - .try_lock() - .expect("test must hold quota lock before polling reader"); - - let mut io = StatsIo::new( - tokio::io::empty(), - Arc::new(SharedCounters::new()), - Arc::clone(&stats), - user.to_string(), - Some(2048), - Arc::new(AtomicBool::new(false)), - tokio::time::Instant::now(), - ); - - let wake_counter = Arc::new(WakeCounter::default()); - let waker = Waker::from(Arc::clone(&wake_counter)); - let mut cx = Context::from_waker(&waker); - let mut storage = [0u8; 1]; - - let mut buf = ReadBuf::new(&mut storage); - let first = Pin::new(&mut io).poll_read(&mut cx, &mut buf); - assert!(first.is_pending()); - - let start = Instant::now(); - let mut observed = 0usize; - while start.elapsed() < Duration::from_millis(80) { - let wakes = wake_counter.wakes.load(Ordering::Relaxed); - if wakes > observed { - observed = wakes; - let mut next = ReadBuf::new(&mut storage); - let pending = Pin::new(&mut io).poll_read(&mut cx, &mut next); - assert!(pending.is_pending()); - } - tokio::time::sleep(Duration::from_millis(1)).await; - } - - assert!( - wake_counter.wakes.load(Ordering::Relaxed) <= 16, - "sustained contention must be rate limited; observed wakes={} in 80ms", - wake_counter.wakes.load(Ordering::Relaxed) - ); - - drop(held_guard); - let mut done = ReadBuf::new(&mut storage); - let ready = Pin::new(&mut io).poll_read(&mut cx, &mut done); - assert!(ready.is_ready()); -} - -#[tokio::test] -async fn edge_backoff_attempt_resets_after_contention_release() { - let _guard = quota_test_guard(); - - let _retained = saturate_quota_user_locks(); - let user = "quota-backoff-edge-reset"; - let stats = Arc::new(Stats::new()); - - let lock = quota_user_lock(user); - let held_guard = lock - .try_lock() - .expect("test must hold quota lock before polling writer"); - - let mut io = StatsIo::new( - tokio::io::sink(), - Arc::new(SharedCounters::new()), - Arc::clone(&stats), - user.to_string(), - Some(2048), - Arc::new(AtomicBool::new(false)), - tokio::time::Instant::now(), - ); - - let wake_counter = Arc::new(WakeCounter::default()); - let waker = Waker::from(Arc::clone(&wake_counter)); - let mut cx = Context::from_waker(&waker); - - let initial = Pin::new(&mut io).poll_write(&mut cx, &[0x31]); - assert!(initial.is_pending()); - - tokio::time::sleep(Duration::from_millis(15)).await; - let wakes = wake_counter.wakes.load(Ordering::Relaxed); - if wakes > 0 { - let pending = Pin::new(&mut io).poll_write(&mut cx, &[0x32]); - assert!(pending.is_pending()); - } - - drop(held_guard); - let ready = Pin::new(&mut io).poll_write(&mut cx, &[0x33]); - assert!(ready.is_ready()); - assert!( - !io.quota_write_wake_scheduled, - "successful write must clear deferred wake scheduling flag" - ); - assert!( - io.quota_write_retry_sleep.is_none(), - "successful write must clear deferred sleep slot" - ); -} - -#[tokio::test] -async fn light_fuzz_writer_repoll_schedule_keeps_wake_budget_bounded() { - let _guard = quota_test_guard(); - - let _retained = saturate_quota_user_locks(); - let user = "quota-backoff-fuzz-writer"; - let stats = Arc::new(Stats::new()); - - let lock = quota_user_lock(user); - let held_guard = lock - .try_lock() - .expect("test must hold quota lock before fuzz loop"); - - let mut io = StatsIo::new( - tokio::io::sink(), - Arc::new(SharedCounters::new()), - Arc::clone(&stats), - user.to_string(), - Some(2048), - Arc::new(AtomicBool::new(false)), - tokio::time::Instant::now(), - ); - - let wake_counter = Arc::new(WakeCounter::default()); - let waker = Waker::from(Arc::clone(&wake_counter)); - let mut cx = Context::from_waker(&waker); - - let mut seed = 0x5EED_CAFE_7788_9900u64; - for _ in 0..64 { - let poll = Pin::new(&mut io).poll_write(&mut cx, &[0x51]); - assert!(poll.is_pending()); - - seed ^= seed << 7; - seed ^= seed >> 9; - seed ^= seed << 8; - let sleep_ms = (seed % 4) as u64; - tokio::time::sleep(Duration::from_millis(sleep_ms)).await; - } - - assert!( - wake_counter.wakes.load(Ordering::Relaxed) <= 24, - "fuzzed repoll schedule must keep wake budget bounded; observed wakes={}", - wake_counter.wakes.load(Ordering::Relaxed) - ); - - drop(held_guard); -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 4)] -async fn stress_multi_waiter_contention_keeps_global_wake_budget_bounded() { - let _guard = quota_test_guard(); - - let _retained = saturate_quota_user_locks(); - let user = format!("quota-backoff-stress-{}", std::process::id()); - let stats = Arc::new(Stats::new()); - - let lock = quota_user_lock(&user); - let held_guard = lock - .try_lock() - .expect("test must hold quota lock before launching stress waiters"); - - let waiters = 48usize; - let mut ios = Vec::with_capacity(waiters); - let mut wake_counters = Vec::with_capacity(waiters); - - for _ in 0..waiters { - ios.push(StatsIo::new( - tokio::io::sink(), - Arc::new(SharedCounters::new()), - Arc::clone(&stats), - user.clone(), - Some(4096), - Arc::new(AtomicBool::new(false)), - tokio::time::Instant::now(), - )); - } - - for io in &mut ios { - let counter = Arc::new(WakeCounter::default()); - let waker = Waker::from(Arc::clone(&counter)); - let mut cx = Context::from_waker(&waker); - let pending = Pin::new(io).poll_write(&mut cx, &[0x61]); - assert!(pending.is_pending()); - wake_counters.push(counter); - } - - let start = Instant::now(); - while start.elapsed() < Duration::from_millis(120) { - for (idx, counter) in wake_counters.iter().enumerate() { - if counter.wakes.load(Ordering::Relaxed) > 0 { - let waker = Waker::from(Arc::clone(counter)); - let mut cx = Context::from_waker(&waker); - let pending = Pin::new(&mut ios[idx]).poll_write(&mut cx, &[0x62]); - assert!(pending.is_pending()); - } - } - tokio::time::sleep(Duration::from_millis(1)).await; - } - - let total_wakes: usize = wake_counters - .iter() - .map(|counter| counter.wakes.load(Ordering::Relaxed)) - .sum(); - - assert!( - total_wakes <= waiters * 20, - "stress contention must keep aggregate wake budget bounded; waiters={waiters}, wakes={total_wakes}" - ); - - drop(held_guard); -} diff --git a/src/proxy/tests/relay_quota_retry_scheduler_tdd_tests.rs b/src/proxy/tests/relay_quota_retry_scheduler_tdd_tests.rs deleted file mode 100644 index 35a6b6e..0000000 --- a/src/proxy/tests/relay_quota_retry_scheduler_tdd_tests.rs +++ /dev/null @@ -1,246 +0,0 @@ -use super::*; -use crate::stats::Stats; -use std::pin::Pin; -use std::sync::Arc; -use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; -use std::task::{Context, Poll, Waker}; -use tokio::io::{AsyncReadExt, AsyncWriteExt, ReadBuf}; -use tokio::time::{Duration, timeout}; - -#[derive(Default)] -struct WakeCounter { - wakes: AtomicUsize, -} - -impl std::task::Wake for WakeCounter { - fn wake(self: Arc) { - self.wakes.fetch_add(1, Ordering::Relaxed); - } - - fn wake_by_ref(self: &Arc) { - self.wakes.fetch_add(1, Ordering::Relaxed); - } -} - -fn quota_test_guard() -> impl Drop { - super::quota_user_lock_test_scope() -} - -#[tokio::test] -async fn positive_uncontended_quota_limited_writer_completes() { - let _guard = quota_test_guard(); - - let stats = Arc::new(Stats::new()); - let mut io = StatsIo::new( - tokio::io::sink(), - Arc::new(SharedCounters::new()), - Arc::clone(&stats), - "tdd-uncontended".to_string(), - Some(1024), - Arc::new(AtomicBool::new(false)), - tokio::time::Instant::now(), - ); - - let result = io.write_all(&[0x41, 0x42, 0x43]).await; - assert!(result.is_ok(), "uncontended writer must complete"); -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 4)] -async fn adversarial_contended_writers_without_repoll_must_not_wake_storm() { - let _guard = quota_test_guard(); - - let user = format!("tdd-writer-storm-{}", std::process::id()); - let held = quota_user_lock(&user); - let _held_guard = held - .try_lock() - .expect("test must hold quota lock before polling writers"); - - let stats = Arc::new(Stats::new()); - let writers = 24usize; - let mut ios = Vec::with_capacity(writers); - let mut wake_counters = Vec::with_capacity(writers); - - for _ in 0..writers { - ios.push(StatsIo::new( - tokio::io::sink(), - Arc::new(SharedCounters::new()), - Arc::clone(&stats), - user.clone(), - Some(1024), - Arc::new(AtomicBool::new(false)), - tokio::time::Instant::now(), - )); - } - - for io in &mut ios { - let counter = Arc::new(WakeCounter::default()); - let waker = Waker::from(Arc::clone(&counter)); - let mut cx = Context::from_waker(&waker); - let poll = Pin::new(io).poll_write(&mut cx, &[0xAA]); - assert!(poll.is_pending(), "writer must be pending under held lock"); - wake_counters.push(counter); - } - - tokio::time::sleep(Duration::from_millis(25)).await; - - let total_wakes: usize = wake_counters - .iter() - .map(|counter| counter.wakes.load(Ordering::Relaxed)) - .sum(); - - assert!( - total_wakes <= writers * 4, - "retry scheduler must remain bounded without repoll; observed wakes={total_wakes}, writers={writers}" - ); -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 4)] -async fn adversarial_contended_readers_without_repoll_must_not_wake_storm() { - let _guard = quota_test_guard(); - - let user = format!("tdd-reader-storm-{}", std::process::id()); - let held = quota_user_lock(&user); - let _held_guard = held - .try_lock() - .expect("test must hold quota lock before polling readers"); - - let stats = Arc::new(Stats::new()); - let readers = 24usize; - let mut ios = Vec::with_capacity(readers); - let mut wake_counters = Vec::with_capacity(readers); - - for _ in 0..readers { - ios.push(StatsIo::new( - tokio::io::empty(), - Arc::new(SharedCounters::new()), - Arc::clone(&stats), - user.clone(), - Some(1024), - Arc::new(AtomicBool::new(false)), - tokio::time::Instant::now(), - )); - } - - for io in &mut ios { - let counter = Arc::new(WakeCounter::default()); - let waker = Waker::from(Arc::clone(&counter)); - let mut cx = Context::from_waker(&waker); - let mut storage = [0u8; 1]; - let mut buf = ReadBuf::new(&mut storage); - let poll = Pin::new(io).poll_read(&mut cx, &mut buf); - assert!(poll.is_pending(), "reader must be pending under held lock"); - wake_counters.push(counter); - } - - tokio::time::sleep(Duration::from_millis(25)).await; - - let total_wakes: usize = wake_counters - .iter() - .map(|counter| counter.wakes.load(Ordering::Relaxed)) - .sum(); - - assert!( - total_wakes <= readers * 4, - "retry scheduler must remain bounded without repoll; observed wakes={total_wakes}, readers={readers}" - ); -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 4)] -async fn integration_contended_waiters_resume_after_lock_release() { - let _guard = quota_test_guard(); - - let user = format!("tdd-resume-{}", std::process::id()); - let held = quota_user_lock(&user); - let held_guard = held - .try_lock() - .expect("test must hold quota lock before launching waiters"); - - let stats = Arc::new(Stats::new()); - let mut waiters = Vec::new(); - for _ in 0..12 { - let stats = Arc::clone(&stats); - let user = user.clone(); - waiters.push(tokio::spawn(async move { - let mut io = StatsIo::new( - tokio::io::sink(), - Arc::new(SharedCounters::new()), - stats, - user, - Some(2048), - Arc::new(AtomicBool::new(false)), - tokio::time::Instant::now(), - ); - io.write_all(&[0x5A]).await - })); - } - - tokio::time::sleep(Duration::from_millis(5)).await; - drop(held_guard); - - timeout(Duration::from_secs(1), async { - for waiter in waiters { - let result = waiter.await.expect("waiter task must not panic"); - assert!(result.is_ok(), "waiter must complete after release"); - } - }) - .await - .expect("all waiters must complete in bounded time"); -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 4)] -async fn light_fuzz_contention_rounds_keep_retry_wakes_bounded() { - let _guard = quota_test_guard(); - - let mut seed = 0x9E37_79B9_AA55_1234u64; - for round in 0..20u32 { - seed ^= seed << 7; - seed ^= seed >> 9; - seed ^= seed << 8; - - let writers = 8 + (seed as usize % 12); - let sleep_ms = 10 + (seed as u64 % 15); - let user = format!("tdd-fuzz-{}-{round}", std::process::id()); - - let held = quota_user_lock(&user); - let _held_guard = held - .try_lock() - .expect("test must hold quota lock in fuzz round"); - - let stats = Arc::new(Stats::new()); - let mut ios = Vec::with_capacity(writers); - let mut wake_counters = Vec::with_capacity(writers); - - for _ in 0..writers { - ios.push(StatsIo::new( - tokio::io::sink(), - Arc::new(SharedCounters::new()), - Arc::clone(&stats), - user.clone(), - Some(2048), - Arc::new(AtomicBool::new(false)), - tokio::time::Instant::now(), - )); - } - - for io in &mut ios { - let counter = Arc::new(WakeCounter::default()); - let waker = Waker::from(Arc::clone(&counter)); - let mut cx = Context::from_waker(&waker); - let poll = Pin::new(io).poll_write(&mut cx, &[0x7A]); - assert!(matches!(poll, Poll::Pending)); - wake_counters.push(counter); - } - - tokio::time::sleep(Duration::from_millis(sleep_ms)).await; - - let total_wakes: usize = wake_counters - .iter() - .map(|counter| counter.wakes.load(Ordering::Relaxed)) - .sum(); - - assert!( - total_wakes <= writers * 4, - "fuzz round must keep wakes bounded; round={round}, writers={writers}, wakes={total_wakes}, sleep_ms={sleep_ms}" - ); - } -} diff --git a/src/proxy/tests/relay_quota_wake_liveness_regression_tests.rs b/src/proxy/tests/relay_quota_wake_liveness_regression_tests.rs deleted file mode 100644 index 9f68258..0000000 --- a/src/proxy/tests/relay_quota_wake_liveness_regression_tests.rs +++ /dev/null @@ -1,294 +0,0 @@ -use super::*; -use crate::stats::Stats; -use dashmap::DashMap; -use std::sync::Arc; -use std::sync::atomic::AtomicBool; -use tokio::io::{AsyncReadExt, AsyncWriteExt}; -use tokio::sync::Barrier; -use tokio::time::{Duration, timeout}; - -fn saturate_lock_cache() -> Vec>> { - let map = QUOTA_USER_LOCKS.get_or_init(DashMap::new); - map.clear(); - - let mut retained = Vec::with_capacity(QUOTA_USER_LOCKS_MAX); - for idx in 0..QUOTA_USER_LOCKS_MAX { - retained.push(quota_user_lock(&format!("quota-liveness-saturated-{idx}"))); - } - retained -} - -fn quota_test_guard() -> impl Drop { - super::quota_user_lock_test_scope() -} - -#[tokio::test] -async fn positive_writer_progresses_after_contention_release_without_external_wake() { - let _guard = quota_test_guard(); - - let _retained = saturate_lock_cache(); - let user = "quota-liveness-writer-positive"; - let stats = Arc::new(Stats::new()); - - let lock = quota_user_lock(user); - let held_guard = lock - .try_lock() - .expect("test must hold user quota lock before write"); - - let counters = Arc::new(SharedCounters::new()); - let quota_exceeded = Arc::new(AtomicBool::new(false)); - let mut io = StatsIo::new( - tokio::io::sink(), - counters, - Arc::clone(&stats), - user.to_string(), - Some(1024), - quota_exceeded, - tokio::time::Instant::now(), - ); - - let writer = tokio::spawn(async move { io.write_all(&[0x11]).await }); - - // Let the initial deferred wake fire while contention is still active. - tokio::time::sleep(Duration::from_millis(4)).await; - - drop(held_guard); - - let completed = timeout(Duration::from_millis(250), writer) - .await - .expect("writer must be re-polled and complete after lock release") - .expect("writer task must not panic"); - assert!(completed.is_ok(), "writer must complete after lock release"); -} - -#[tokio::test] -async fn edge_reader_progresses_after_contention_release_without_external_wake() { - let _guard = quota_test_guard(); - - let _retained = saturate_lock_cache(); - let user = "quota-liveness-reader-edge"; - let stats = Arc::new(Stats::new()); - - let lock = quota_user_lock(user); - let held_guard = lock - .try_lock() - .expect("test must hold user quota lock before read"); - - let counters = Arc::new(SharedCounters::new()); - let quota_exceeded = Arc::new(AtomicBool::new(false)); - let mut io = StatsIo::new( - tokio::io::empty(), - counters, - Arc::clone(&stats), - user.to_string(), - Some(1024), - quota_exceeded, - tokio::time::Instant::now(), - ); - - let reader = tokio::spawn(async move { - let mut one = [0u8; 1]; - io.read(&mut one).await - }); - - tokio::time::sleep(Duration::from_millis(4)).await; - drop(held_guard); - - let completed = timeout(Duration::from_millis(250), reader) - .await - .expect("reader must be re-polled and complete after lock release") - .expect("reader task must not panic"); - assert!(completed.is_ok(), "reader must complete after lock release"); -} - -#[tokio::test] -async fn adversarial_early_deferred_wake_consumption_does_not_deadlock_writer() { - let _guard = quota_test_guard(); - - let _retained = saturate_lock_cache(); - let user = "quota-liveness-adversarial"; - let stats = Arc::new(Stats::new()); - - let lock = quota_user_lock(user); - let held_guard = lock - .try_lock() - .expect("test must hold user quota lock before adversarial write"); - - let counters = Arc::new(SharedCounters::new()); - let quota_exceeded = Arc::new(AtomicBool::new(false)); - let mut io = StatsIo::new( - tokio::io::sink(), - counters, - Arc::clone(&stats), - user.to_string(), - Some(1024), - quota_exceeded, - tokio::time::Instant::now(), - ); - - let writer = tokio::spawn(async move { io.write_all(&[0x22]).await }); - - // Force multiple scheduler rounds while lock remains held so the first - // deferred wake has already been consumed under contention. - for _ in 0..32 { - tokio::task::yield_now().await; - } - - drop(held_guard); - - let completed = timeout(Duration::from_millis(300), writer) - .await - .expect("writer must not stay parked forever after release") - .expect("writer task must not panic"); - assert!(completed.is_ok()); -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 4)] -async fn integration_parallel_waiters_resume_after_single_release_event() { - let _guard = quota_test_guard(); - - let _retained = saturate_lock_cache(); - let user = format!("quota-liveness-integration-{}", std::process::id()); - let stats = Arc::new(Stats::new()); - let barrier = Arc::new(Barrier::new(13)); - - let lock = quota_user_lock(&user); - let held_guard = lock - .try_lock() - .expect("test must hold user quota lock before launching waiters"); - - let mut waiters = Vec::new(); - for _ in 0..12 { - let stats = Arc::clone(&stats); - let user = user.clone(); - let barrier = Arc::clone(&barrier); - waiters.push(tokio::spawn(async move { - let counters = Arc::new(SharedCounters::new()); - let quota_exceeded = Arc::new(AtomicBool::new(false)); - let mut io = StatsIo::new( - tokio::io::sink(), - counters, - stats, - user, - Some(4096), - quota_exceeded, - tokio::time::Instant::now(), - ); - barrier.wait().await; - io.write_all(&[0x33]).await - })); - } - - barrier.wait().await; - tokio::time::sleep(Duration::from_millis(4)).await; - drop(held_guard); - - timeout(Duration::from_secs(1), async { - for waiter in waiters { - let outcome = waiter.await.expect("waiter must not panic"); - assert!( - outcome.is_ok(), - "waiter must resume and complete after release" - ); - } - }) - .await - .expect("all waiters must complete in bounded time"); -} - -#[tokio::test] -async fn light_fuzz_release_timing_matrix_preserves_liveness() { - let _guard = quota_test_guard(); - - let _retained = saturate_lock_cache(); - let stats = Arc::new(Stats::new()); - - let mut seed = 0xD1CE_F00D_0123_4567u64; - for round in 0..64u32 { - seed ^= seed << 7; - seed ^= seed >> 9; - seed ^= seed << 8; - - let delay_ms = 1 + (seed & 0x7) as u64; - let user = format!("quota-liveness-fuzz-{}-{round}", std::process::id()); - - let lock = quota_user_lock(&user); - let held_guard = lock - .try_lock() - .expect("test must hold user quota lock in fuzz round"); - - let counters = Arc::new(SharedCounters::new()); - let quota_exceeded = Arc::new(AtomicBool::new(false)); - let mut io = StatsIo::new( - tokio::io::sink(), - counters, - Arc::clone(&stats), - user, - Some(2048), - quota_exceeded, - tokio::time::Instant::now(), - ); - - let writer = tokio::spawn(async move { io.write_all(&[0x44]).await }); - - tokio::time::sleep(Duration::from_millis(delay_ms)).await; - drop(held_guard); - - let done = timeout(Duration::from_millis(300), writer) - .await - .expect("fuzz round writer must complete") - .expect("fuzz writer task must not panic"); - assert!( - done.is_ok(), - "fuzz round writer must not stall after release" - ); - } -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 4)] -async fn stress_repeated_contention_cycles_remain_live() { - let _guard = quota_test_guard(); - - let _retained = saturate_lock_cache(); - let stats = Arc::new(Stats::new()); - - for cycle in 0..40u32 { - let user = format!("quota-liveness-stress-{}-{cycle}", std::process::id()); - let lock = quota_user_lock(&user); - let held_guard = lock - .try_lock() - .expect("test must hold lock before stress cycle"); - - let mut tasks = Vec::new(); - for _ in 0..6 { - let stats = Arc::clone(&stats); - let user = user.clone(); - tasks.push(tokio::spawn(async move { - let counters = Arc::new(SharedCounters::new()); - let quota_exceeded = Arc::new(AtomicBool::new(false)); - let mut io = StatsIo::new( - tokio::io::sink(), - counters, - stats, - user, - Some(2048), - quota_exceeded, - tokio::time::Instant::now(), - ); - io.write_all(&[0x55]).await - })); - } - - tokio::task::yield_now().await; - drop(held_guard); - - timeout(Duration::from_millis(700), async { - for task in tasks { - let outcome = task.await.expect("stress task must not panic"); - assert!(outcome.is_ok(), "stress writer must complete"); - } - }) - .await - .expect("stress cycle must finish in bounded time"); - } -} diff --git a/src/proxy/tests/relay_quota_waker_storm_adversarial_tests.rs b/src/proxy/tests/relay_quota_waker_storm_adversarial_tests.rs deleted file mode 100644 index fa4878a..0000000 --- a/src/proxy/tests/relay_quota_waker_storm_adversarial_tests.rs +++ /dev/null @@ -1,310 +0,0 @@ -use super::*; -use crate::stats::Stats; -use dashmap::DashMap; -use std::pin::Pin; -use std::sync::Arc; -use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; -use std::task::{Context, Waker}; -use tokio::io::{AsyncWriteExt, ReadBuf}; -use tokio::time::{Duration, timeout}; - -#[derive(Default)] -struct WakeCounter { - wakes: AtomicUsize, -} - -impl std::task::Wake for WakeCounter { - fn wake(self: Arc) { - self.wakes.fetch_add(1, Ordering::Relaxed); - } - - fn wake_by_ref(self: &Arc) { - self.wakes.fetch_add(1, Ordering::Relaxed); - } -} - -fn quota_test_guard() -> impl Drop { - super::quota_user_lock_test_scope() -} - -fn saturate_quota_user_locks() -> Vec>> { - let map = QUOTA_USER_LOCKS.get_or_init(DashMap::new); - map.clear(); - - let mut retained = Vec::with_capacity(QUOTA_USER_LOCKS_MAX); - for idx in 0..QUOTA_USER_LOCKS_MAX { - retained.push(quota_user_lock(&format!("quota-waker-saturate-{idx}"))); - } - retained -} - -#[tokio::test] -async fn positive_contended_writer_emits_deferred_wake_for_liveness() { - let _guard = quota_test_guard(); - - let _retained = saturate_quota_user_locks(); - let stats = Arc::new(Stats::new()); - let user = "quota-waker-positive-user"; - - let lock = quota_user_lock(user); - let held_guard = lock - .try_lock() - .expect("test must hold overflow lock before polling writer"); - - let counters = Arc::new(SharedCounters::new()); - let quota_exceeded = Arc::new(AtomicBool::new(false)); - let mut io = StatsIo::new( - tokio::io::sink(), - counters, - Arc::clone(&stats), - user.to_string(), - Some(1024), - quota_exceeded, - tokio::time::Instant::now(), - ); - - let wake_counter = Arc::new(WakeCounter::default()); - let waker = Waker::from(Arc::clone(&wake_counter)); - let mut cx = Context::from_waker(&waker); - - let pending = Pin::new(&mut io).poll_write(&mut cx, &[0xA1]); - assert!(pending.is_pending()); - - timeout(Duration::from_millis(100), async { - loop { - if wake_counter.wakes.load(Ordering::Relaxed) >= 1 { - break; - } - tokio::task::yield_now().await; - } - }) - .await - .expect("contended writer must receive deferred wake"); - - drop(held_guard); - let ready = Pin::new(&mut io).poll_write(&mut cx, &[0xA2]); - assert!( - ready.is_ready(), - "writer must progress after contention release" - ); -} - -#[tokio::test] -async fn adversarial_blackhat_writer_contention_does_not_create_waker_storm() { - let _guard = quota_test_guard(); - - let _retained = saturate_quota_user_locks(); - let stats = Arc::new(Stats::new()); - let user = "quota-waker-blackhat-writer"; - - let lock = quota_user_lock(user); - let held_guard = lock - .try_lock() - .expect("test must hold overflow lock before polling writer"); - - let counters = Arc::new(SharedCounters::new()); - let quota_exceeded = Arc::new(AtomicBool::new(false)); - let mut io = StatsIo::new( - tokio::io::sink(), - counters, - Arc::clone(&stats), - user.to_string(), - Some(1024), - quota_exceeded, - tokio::time::Instant::now(), - ); - - let wake_counter = Arc::new(WakeCounter::default()); - let waker = Waker::from(Arc::clone(&wake_counter)); - let mut cx = Context::from_waker(&waker); - - for _ in 0..512 { - let poll = Pin::new(&mut io).poll_write(&mut cx, &[0xBE]); - assert!( - poll.is_pending(), - "writer must stay pending while lock is held" - ); - tokio::task::yield_now().await; - } - - let wakes = wake_counter.wakes.load(Ordering::Relaxed); - assert!( - wakes <= 128, - "pending writer retries must not trigger wake storm; observed wakes={wakes}" - ); - - drop(held_guard); - let ready = Pin::new(&mut io).poll_write(&mut cx, &[0xEF]); - assert!(ready.is_ready()); -} - -#[tokio::test] -async fn edge_read_path_contention_keeps_wake_budget_bounded() { - let _guard = quota_test_guard(); - - let _retained = saturate_quota_user_locks(); - let stats = Arc::new(Stats::new()); - let user = "quota-waker-read-edge"; - - let lock = quota_user_lock(user); - let held_guard = lock - .try_lock() - .expect("test must hold overflow lock before polling reader"); - - let counters = Arc::new(SharedCounters::new()); - let quota_exceeded = Arc::new(AtomicBool::new(false)); - let mut io = StatsIo::new( - tokio::io::empty(), - counters, - Arc::clone(&stats), - user.to_string(), - Some(1024), - quota_exceeded, - tokio::time::Instant::now(), - ); - - let wake_counter = Arc::new(WakeCounter::default()); - let waker = Waker::from(Arc::clone(&wake_counter)); - let mut cx = Context::from_waker(&waker); - let mut storage = [0u8; 1]; - - for _ in 0..512 { - let mut buf = ReadBuf::new(&mut storage); - let poll = Pin::new(&mut io).poll_read(&mut cx, &mut buf); - assert!(poll.is_pending()); - tokio::task::yield_now().await; - } - - let wakes = wake_counter.wakes.load(Ordering::Relaxed); - assert!( - wakes <= 128, - "pending reader retries must not trigger wake storm; observed wakes={wakes}" - ); - - drop(held_guard); - let mut buf = ReadBuf::new(&mut storage); - let ready = Pin::new(&mut io).poll_read(&mut cx, &mut buf); - assert!(ready.is_ready()); -} - -#[tokio::test] -async fn light_fuzz_mixed_poll_schedule_under_contention_stays_bounded() { - let _guard = quota_test_guard(); - - let _retained = saturate_quota_user_locks(); - let stats = Arc::new(Stats::new()); - let user = "quota-waker-fuzz-user"; - - let lock = quota_user_lock(user); - let held_guard = lock - .try_lock() - .expect("test must hold overflow lock before fuzz polling"); - - let counters_w = Arc::new(SharedCounters::new()); - let mut writer_io = StatsIo::new( - tokio::io::sink(), - counters_w, - Arc::clone(&stats), - user.to_string(), - Some(1024), - Arc::new(AtomicBool::new(false)), - tokio::time::Instant::now(), - ); - - let counters_r = Arc::new(SharedCounters::new()); - let mut reader_io = StatsIo::new( - tokio::io::empty(), - counters_r, - Arc::clone(&stats), - user.to_string(), - Some(1024), - Arc::new(AtomicBool::new(false)), - tokio::time::Instant::now(), - ); - - let wake_counter = Arc::new(WakeCounter::default()); - let waker = Waker::from(Arc::clone(&wake_counter)); - let mut cx = Context::from_waker(&waker); - let mut seed = 0xBADC_0FFE_EE11_2211u64; - let mut storage = [0u8; 1]; - - for _ in 0..1024 { - seed ^= seed << 7; - seed ^= seed >> 9; - seed ^= seed << 8; - - if (seed & 1) == 0 { - let poll = Pin::new(&mut writer_io).poll_write(&mut cx, &[0x44]); - assert!(poll.is_pending()); - } else { - let mut buf = ReadBuf::new(&mut storage); - let poll = Pin::new(&mut reader_io).poll_read(&mut cx, &mut buf); - assert!(poll.is_pending()); - } - tokio::task::yield_now().await; - } - - assert!( - wake_counter.wakes.load(Ordering::Relaxed) <= 192, - "mixed contention fuzz must keep deferred wake count tightly bounded" - ); - - drop(held_guard); - let ready_w = Pin::new(&mut writer_io).poll_write(&mut cx, &[0x55]); - assert!(ready_w.is_ready()); - - let mut buf = ReadBuf::new(&mut storage); - let ready_r = Pin::new(&mut reader_io).poll_read(&mut cx, &mut buf); - assert!(ready_r.is_ready()); -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 4)] -#[ignore = "red-team detector: reveals possible starvation if deferred wake fires before contention release"] -async fn stress_many_contended_writers_complete_after_release() { - let _guard = quota_test_guard(); - - let _retained = saturate_quota_user_locks(); - let user = "quota-waker-stress-user".to_string(); - let stats = Arc::new(Stats::new()); - - let lock = quota_user_lock(&user); - let held_guard = lock - .try_lock() - .expect("test must hold overflow lock before launching contended tasks"); - - let mut tasks = Vec::new(); - for _ in 0..32 { - let stats = Arc::clone(&stats); - let user = user.clone(); - tasks.push(tokio::spawn(async move { - let counters = Arc::new(SharedCounters::new()); - let quota_exceeded = Arc::new(AtomicBool::new(false)); - let mut io = StatsIo::new( - tokio::io::sink(), - counters, - stats, - user, - Some(2048), - quota_exceeded, - tokio::time::Instant::now(), - ); - - io.write_all(&[0xAA]).await - })); - } - - for _ in 0..8 { - tokio::task::yield_now().await; - } - - drop(held_guard); - - timeout(Duration::from_secs(2), async { - for task in tasks { - let result = task.await.expect("stress task must not panic"); - assert!(result.is_ok(), "task must complete after lock release"); - } - }) - .await - .expect("all contended writer tasks must finish in bounded time after release"); -} diff --git a/src/proxy/tests/relay_security_tests.rs b/src/proxy/tests/relay_security_tests.rs deleted file mode 100644 index 7375192..0000000 --- a/src/proxy/tests/relay_security_tests.rs +++ /dev/null @@ -1,1284 +0,0 @@ -use super::relay_bidirectional; -use crate::error::ProxyError; -use crate::stats::Stats; -use crate::stream::BufferPool; -use std::future::poll_fn; -use std::io; -use std::pin::Pin; -use std::sync::Arc; -use std::sync::Mutex; -use std::sync::atomic::{AtomicUsize, Ordering}; -use std::task::Waker; -use std::task::{Context, Poll}; -use tokio::io::{AsyncRead, ReadBuf}; -use tokio::io::{AsyncReadExt, AsyncWrite, AsyncWriteExt, duplex}; -use tokio::time::{Duration, timeout}; - -#[derive(Default)] -struct WakeCounter { - wakes: AtomicUsize, -} - -impl std::task::Wake for WakeCounter { - fn wake(self: Arc) { - self.wakes.fetch_add(1, Ordering::Relaxed); - } - - fn wake_by_ref(self: &Arc) { - self.wakes.fetch_add(1, Ordering::Relaxed); - } -} - -#[tokio::test] -async fn quota_lock_contention_does_not_self_wake_pending_writer() { - let _guard = super::quota_user_lock_test_scope(); - let map = super::QUOTA_USER_LOCKS.get_or_init(dashmap::DashMap::new); - map.clear(); - - let stats = Arc::new(Stats::new()); - let user = "quota-lock-contention-user"; - - let lock = super::quota_user_lock(user); - let _held_lock = lock - .try_lock() - .expect("test must hold the per-user quota lock before polling writer"); - - let counters = Arc::new(super::SharedCounters::new()); - let quota_exceeded = Arc::new(std::sync::atomic::AtomicBool::new(false)); - let mut io = super::StatsIo::new( - tokio::io::sink(), - counters, - Arc::clone(&stats), - user.to_string(), - Some(1024), - quota_exceeded, - tokio::time::Instant::now(), - ); - - let wake_counter = Arc::new(WakeCounter::default()); - let waker = Waker::from(Arc::clone(&wake_counter)); - let mut cx = Context::from_waker(&waker); - - let poll = Pin::new(&mut io).poll_write(&mut cx, &[0x11]); - assert!( - poll.is_pending(), - "writer must remain pending while lock is contended" - ); - assert_eq!( - wake_counter.wakes.load(Ordering::Relaxed), - 0, - "contended quota lock must not self-wake immediately and spin the executor" - ); -} - -#[tokio::test] -async fn quota_lock_contention_writer_schedules_single_deferred_wake_until_lock_acquired() { - let _guard = super::quota_user_lock_test_scope(); - let map = super::QUOTA_USER_LOCKS.get_or_init(dashmap::DashMap::new); - map.clear(); - - let stats = Arc::new(Stats::new()); - let user = "quota-lock-writer-liveness-user"; - - let lock = super::quota_user_lock(user); - let held_lock = lock - .try_lock() - .expect("test must hold the per-user quota lock before polling writer"); - - let counters = Arc::new(super::SharedCounters::new()); - let quota_exceeded = Arc::new(std::sync::atomic::AtomicBool::new(false)); - let mut io = super::StatsIo::new( - tokio::io::sink(), - counters, - Arc::clone(&stats), - user.to_string(), - Some(1024), - quota_exceeded, - tokio::time::Instant::now(), - ); - - let wake_counter = Arc::new(WakeCounter::default()); - let waker = Waker::from(Arc::clone(&wake_counter)); - let mut cx = Context::from_waker(&waker); - - let first = Pin::new(&mut io).poll_write(&mut cx, &[0x11]); - assert!( - first.is_pending(), - "writer must remain pending while lock is contended" - ); - assert_eq!( - wake_counter.wakes.load(Ordering::Relaxed), - 0, - "deferred wake must not fire synchronously" - ); - - timeout(Duration::from_millis(50), async { - loop { - if wake_counter.wakes.load(Ordering::Relaxed) >= 1 { - break; - } - tokio::task::yield_now().await; - } - }) - .await - .expect("contended writer must schedule a deferred wake in bounded time"); - let wakes_after_first_yield = wake_counter.wakes.load(Ordering::Relaxed); - assert!( - wakes_after_first_yield >= 1, - "contended writer must schedule at least one deferred wake for liveness" - ); - - let second = Pin::new(&mut io).poll_write(&mut cx, &[0x22]); - assert!( - second.is_pending(), - "writer remains pending while lock is still held" - ); - - for _ in 0..8 { - tokio::task::yield_now().await; - } - let wakes_after_second_window = wake_counter.wakes.load(Ordering::Relaxed); - assert!( - wakes_after_second_window <= wakes_after_first_yield.saturating_add(2), - "writer contention should keep retry wakes bounded before lock acquisition: before={wakes_after_first_yield}, after={wakes_after_second_window}" - ); - - drop(held_lock); - let released = Pin::new(&mut io).poll_write(&mut cx, &[0x33]); - assert!( - released.is_ready(), - "writer must make progress once quota lock is released" - ); -} - -#[tokio::test] -async fn quota_lock_contention_read_path_schedules_deferred_wake_for_liveness() { - let _guard = super::quota_user_lock_test_scope(); - let map = super::QUOTA_USER_LOCKS.get_or_init(dashmap::DashMap::new); - map.clear(); - - let stats = Arc::new(Stats::new()); - let user = "quota-lock-read-liveness-user"; - - let lock = super::quota_user_lock(user); - let held_lock = lock - .try_lock() - .expect("test must hold the per-user quota lock before polling reader"); - - let counters = Arc::new(super::SharedCounters::new()); - let quota_exceeded = Arc::new(std::sync::atomic::AtomicBool::new(false)); - let mut io = super::StatsIo::new( - tokio::io::empty(), - counters, - Arc::clone(&stats), - user.to_string(), - Some(1024), - quota_exceeded, - tokio::time::Instant::now(), - ); - - let wake_counter = Arc::new(WakeCounter::default()); - let waker = Waker::from(Arc::clone(&wake_counter)); - let mut cx = Context::from_waker(&waker); - let mut storage = [0u8; 1]; - let mut buf = ReadBuf::new(&mut storage); - - let first = Pin::new(&mut io).poll_read(&mut cx, &mut buf); - assert!( - first.is_pending(), - "reader must remain pending while lock is contended" - ); - assert_eq!( - wake_counter.wakes.load(Ordering::Relaxed), - 0, - "read contention wake must not fire synchronously" - ); - - timeout(Duration::from_millis(50), async { - loop { - if wake_counter.wakes.load(Ordering::Relaxed) >= 1 { - break; - } - tokio::task::yield_now().await; - } - }) - .await - .expect("read contention must schedule a deferred wake in bounded time"); - - drop(held_lock); - let mut buf_after_release = ReadBuf::new(&mut storage); - let released = Pin::new(&mut io).poll_read(&mut cx, &mut buf_after_release); - assert!( - released.is_ready(), - "reader must make progress once quota lock is released" - ); -} - -#[tokio::test] -async fn relay_bidirectional_enforces_live_user_quota() { - let stats = Arc::new(Stats::new()); - let user = "quota-user"; - stats.add_user_octets_from(user, 6); - - 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_task = tokio::spawn(relay_bidirectional( - client_reader, - client_writer, - server_reader, - server_writer, - 1024, - 1024, - user, - Arc::clone(&stats), - Some(8), - Arc::new(BufferPool::new()), - )); - - client_peer - .write_all(&[0x10, 0x20, 0x30, 0x40]) - .await - .expect("client write must succeed"); - - let mut forwarded = [0u8; 4]; - let _ = timeout( - Duration::from_millis(200), - server_peer.read_exact(&mut forwarded), - ) - .await; - - let relay_result = timeout(Duration::from_secs(2), relay_task) - .await - .expect("relay task must finish under quota cutoff") - .expect("relay task must not panic"); - - assert!( - matches!(relay_result, Err(ProxyError::DataQuotaExceeded { ref user }) if user == "quota-user"), - "relay must surface a typed quota error once live quota is exceeded" - ); -} - -#[tokio::test] -async fn relay_bidirectional_does_not_forward_server_bytes_after_quota_is_exhausted() { - let stats = Arc::new(Stats::new()); - let quota_user = "quota-exhausted-user"; - stats.add_user_octets_from(quota_user, 1); - - 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_task = tokio::spawn(relay_bidirectional( - client_reader, - client_writer, - server_reader, - server_writer, - 1024, - 1024, - quota_user, - Arc::clone(&stats), - Some(1), - Arc::new(BufferPool::new()), - )); - - server_peer - .write_all(&[0xde, 0xad, 0xbe, 0xef]) - .await - .expect("server write must succeed"); - - let mut observed = [0u8; 4]; - let forwarded = timeout( - Duration::from_millis(200), - client_peer.read_exact(&mut observed), - ) - .await; - - let relay_result = timeout(Duration::from_secs(2), relay_task) - .await - .expect("relay task must finish under quota cutoff") - .expect("relay task must not panic"); - - assert!( - !matches!(forwarded, Ok(Ok(n)) if n == observed.len()), - "no full server payload should be forwarded once quota is already exhausted" - ); - assert!( - matches!(relay_result, Err(ProxyError::DataQuotaExceeded { ref user }) if user == quota_user), - "relay must still terminate with a typed quota error" - ); -} - -#[tokio::test] -async fn relay_bidirectional_does_not_leak_partial_server_payload_when_remaining_quota_is_smaller_than_write() - { - let stats = Arc::new(Stats::new()); - let quota_user = "partial-leak-user"; - stats.add_user_octets_from(quota_user, 3); - - 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_task = tokio::spawn(relay_bidirectional( - client_reader, - client_writer, - server_reader, - server_writer, - 1024, - 1024, - quota_user, - Arc::clone(&stats), - Some(4), - Arc::new(BufferPool::new()), - )); - - server_peer - .write_all(&[0x11, 0x22, 0x33, 0x44]) - .await - .expect("server write must succeed"); - - let mut observed = [0u8; 8]; - let forwarded = timeout(Duration::from_millis(200), client_peer.read(&mut observed)).await; - - let relay_result = timeout(Duration::from_secs(2), relay_task) - .await - .expect("relay task must finish under quota cutoff") - .expect("relay task must not panic"); - - assert!( - !matches!(forwarded, Ok(Ok(n)) if n > 0), - "quota exhaustion must not leak any partial server payload when remaining quota is smaller than the write" - ); - assert!( - matches!(relay_result, Err(ProxyError::DataQuotaExceeded { ref user }) if user == quota_user), - "relay must still terminate with a typed quota error" - ); -} - -#[tokio::test] -async fn relay_bidirectional_zero_quota_remains_fail_closed_for_server_payloads_under_stress() { - let stats = Arc::new(Stats::new()); - let quota_user = "zero-quota-user"; - - for payload_len in [1usize, 16, 512, 4096] { - 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_task = tokio::spawn(relay_bidirectional( - client_reader, - client_writer, - server_reader, - server_writer, - 1024, - 1024, - quota_user, - Arc::clone(&stats), - Some(0), - Arc::new(BufferPool::new()), - )); - - let payload = vec![0x7f; payload_len]; - let _ = server_peer.write_all(&payload).await; - - let mut observed = vec![0u8; payload_len]; - let forwarded = timeout(Duration::from_millis(200), client_peer.read(&mut observed)).await; - - let relay_result = timeout(Duration::from_secs(2), relay_task) - .await - .expect("relay task must finish under zero-quota cutoff") - .expect("relay task must not panic"); - - assert!( - !matches!(forwarded, Ok(Ok(n)) if n > 0), - "zero quota must not forward any server bytes for payload_len={payload_len}" - ); - assert!( - matches!(relay_result, Err(ProxyError::DataQuotaExceeded { ref user }) if user == quota_user), - "zero quota must terminate with the typed quota error for payload_len={payload_len}" - ); - } -} - -#[tokio::test] -async fn relay_bidirectional_allows_exact_server_payload_at_quota_boundary() { - let stats = Arc::new(Stats::new()); - let quota_user = "exact-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_task = tokio::spawn(relay_bidirectional( - client_reader, - client_writer, - server_reader, - server_writer, - 1024, - 1024, - quota_user, - Arc::clone(&stats), - Some(4), - Arc::new(BufferPool::new()), - )); - - server_peer - .write_all(&[0x91, 0x92, 0x93, 0x94]) - .await - .expect("server write must succeed at exact quota boundary"); - - let mut observed = [0u8; 4]; - client_peer - .read_exact(&mut observed) - .await - .expect("client must receive the full payload at the exact quota boundary"); - assert_eq!(observed, [0x91, 0x92, 0x93, 0x94]); - - let relay_result = timeout(Duration::from_secs(2), relay_task) - .await - .expect("relay task must finish after exact boundary delivery") - .expect("relay task must not panic"); - - assert!( - matches!(relay_result, Err(ProxyError::DataQuotaExceeded { ref user }) if user == quota_user), - "relay must close with a typed quota error after reaching the exact boundary" - ); -} - -#[tokio::test] -async fn relay_bidirectional_does_not_forward_client_bytes_after_quota_is_exhausted() { - let stats = Arc::new(Stats::new()); - let quota_user = "client-exhausted-user"; - stats.add_user_octets_from(quota_user, 1); - - 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_task = tokio::spawn(relay_bidirectional( - client_reader, - client_writer, - server_reader, - server_writer, - 1024, - 1024, - quota_user, - Arc::clone(&stats), - Some(1), - Arc::new(BufferPool::new()), - )); - - client_peer - .write_all(&[0x51, 0x52, 0x53, 0x54]) - .await - .expect("client write must succeed even when quota is already exhausted"); - - let mut observed = [0u8; 4]; - let forwarded = timeout( - Duration::from_millis(200), - server_peer.read_exact(&mut observed), - ) - .await; - - let relay_result = timeout(Duration::from_secs(2), relay_task) - .await - .expect("relay task must finish under quota cutoff") - .expect("relay task must not panic"); - - assert!( - !matches!(forwarded, Ok(Ok(n)) if n == observed.len()), - "client payload must not be fully forwarded once quota is already exhausted" - ); - assert!( - matches!(relay_result, Err(ProxyError::DataQuotaExceeded { ref user }) if user == quota_user), - "relay must still terminate with a typed quota error" - ); -} - -#[tokio::test] -async fn relay_bidirectional_server_bytes_remain_blocked_even_under_multiple_payload_sizes() { - let stats = Arc::new(Stats::new()); - let quota_user = "quota-fuzz-user"; - stats.add_user_octets_from(quota_user, 2); - - for payload_len in [1usize, 32, 1024, 8192] { - 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_task = tokio::spawn(relay_bidirectional( - client_reader, - client_writer, - server_reader, - server_writer, - 1024, - 1024, - quota_user, - Arc::clone(&stats), - Some(2), - Arc::new(BufferPool::new()), - )); - - let payload = vec![0xaa; payload_len]; - let _ = server_peer.write_all(&payload).await; - - let mut observed = vec![0u8; payload_len]; - let forwarded = timeout( - Duration::from_millis(200), - client_peer.read_exact(&mut observed), - ) - .await; - - let relay_result = timeout(Duration::from_secs(2), relay_task) - .await - .expect("relay task must finish under quota cutoff") - .expect("relay task must not panic"); - - assert!( - !matches!(forwarded, Ok(Ok(n)) if n == payload_len), - "quota exhaustion must block full server-to-client forwarding for payload_len={payload_len}" - ); - assert!( - matches!(relay_result, Err(ProxyError::DataQuotaExceeded { ref user }) if user == quota_user), - "relay must keep returning the typed quota error for payload_len={payload_len}" - ); - } -} - -#[tokio::test] -async fn relay_bidirectional_terminates_on_activity_timeout() { - tokio::time::pause(); - let stats = Arc::new(Stats::new()); - let user = "timeout-user"; - - let (client_peer, relay_client) = duplex(4096); - let (relay_server, 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_task = tokio::spawn(relay_bidirectional( - client_reader, - client_writer, - server_reader, - server_writer, - 1024, - 1024, - user, - Arc::clone(&stats), - None, // No quota - Arc::new(BufferPool::new()), - )); - - // Wait past the activity timeout threshold (1800 seconds) + buffer - tokio::time::sleep(Duration::from_secs(1805)).await; - - // Resume time to process timeouts - tokio::time::resume(); - - let relay_result = timeout(Duration::from_secs(1), relay_task) - .await - .expect("relay task must finish inside bounded timeout due to inactivity cutoff") - .expect("relay task must not panic"); - - assert!( - relay_result.is_ok(), - "relay should complete successfully on scheduled inactivity timeout" - ); - - // Verify client/server sockets are closed - drop(client_peer); - drop(server_peer); -} - -#[tokio::test] -async fn relay_bidirectional_watchdog_resists_premature_execution() { - tokio::time::pause(); - let stats = Arc::new(Stats::new()); - let user = "activity-user"; - - let (mut client_peer, relay_client) = duplex(4096); - let (relay_server, 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 mut relay_task = tokio::spawn(relay_bidirectional( - client_reader, - client_writer, - server_reader, - server_writer, - 1024, - 1024, - user, - Arc::clone(&stats), - None, - Arc::new(BufferPool::new()), - )); - - // Advance by half the timeout - tokio::time::sleep(Duration::from_secs(900)).await; - - // Provide activity - client_peer - .write_all(&[0xaa, 0xbb]) - .await - .expect("client write must succeed"); - client_peer.flush().await.unwrap(); - - // Advance by another half (total time since start is 1800, but since last activity is 900) - tokio::time::sleep(Duration::from_secs(900)).await; - - tokio::time::resume(); - - // Re-evaluating the task, it should NOT have timed out and still be pending - let relay_result = timeout(Duration::from_millis(100), &mut relay_task).await; - assert!( - relay_result.is_err(), - "Relay must not exit prematurely as long as activity was received before timeout" - ); - - // Explicitly drop sockets to cleanly shut down relay loop - drop(client_peer); - drop(server_peer); - - let completion = timeout(Duration::from_secs(1), relay_task) - .await - .expect("relay task must complete securely after client disconnection") - .expect("relay task must not panic"); - assert!(completion.is_ok(), "relay exits clean"); -} - -#[tokio::test] -async fn relay_bidirectional_half_closure_terminates_cleanly() { - let stats = Arc::new(Stats::new()); - let (client_peer, relay_client) = duplex(4096); - let (relay_server, 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_task = tokio::spawn(relay_bidirectional( - client_reader, - client_writer, - server_reader, - server_writer, - 1024, - 1024, - "half-close", - stats, - None, - Arc::new(BufferPool::new()), - )); - - // Half closure: drop the client completely but leave the server active. - drop(client_peer); - - // Check that we don't immediately crash. Bidirectional relay stays open for the server -> client flush. - // Eventually dropping the server cleanly closes the task. - drop(server_peer); - timeout(Duration::from_secs(1), relay_task) - .await - .unwrap() - .unwrap() - .unwrap(); -} - -#[tokio::test] -async fn relay_bidirectional_zero_length_noise_fuzzing() { - let stats = Arc::new(Stats::new()); - 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_task = tokio::spawn(relay_bidirectional( - client_reader, - client_writer, - server_reader, - server_writer, - 1024, - 1024, - "fuzz", - stats, - None, - Arc::new(BufferPool::new()), - )); - - // Flood with zero-length payloads (edge cases in stream framing logic sometimes loop) - for _ in 0..100 { - client_peer.write_all(&[]).await.unwrap(); - } - client_peer.write_all(&[1, 2, 3]).await.unwrap(); - client_peer.flush().await.unwrap(); - - let mut buf = [0u8; 3]; - server_peer.read_exact(&mut buf).await.unwrap(); - assert_eq!(&buf, &[1, 2, 3]); - - drop(client_peer); - drop(server_peer); - timeout(Duration::from_secs(1), relay_task) - .await - .unwrap() - .unwrap() - .unwrap(); -} - -#[tokio::test] -async fn relay_bidirectional_asymmetric_backpressure() { - let stats = Arc::new(Stats::new()); - // Give the client stream an extremely narrow throughput limit explicitly - let (client_peer, relay_client) = duplex(1024); - 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_task = tokio::spawn(relay_bidirectional( - client_reader, - client_writer, - server_reader, - server_writer, - 1024, - 1024, - "slowloris", - stats, - None, - Arc::new(BufferPool::new()), - )); - - let payload = vec![0xba; 65536]; // 64k payload - - // Server attempts to shove 64KB into a relay whose client pipe only holds 1KB! - let write_res = - tokio::time::timeout(Duration::from_millis(50), server_peer.write_all(&payload)).await; - - assert!( - write_res.is_err(), - "Relay backpressure MUST halt the server writer from unbounded buffering when client stream is full!" - ); - - drop(client_peer); - drop(server_peer); - - let completion = timeout(Duration::from_secs(1), relay_task) - .await - .unwrap() - .unwrap(); - assert!( - completion.is_ok() || completion.is_err(), - "Task must unwind reliably (either Ok or BrokenPipe Err) when dropped despite active backpressure locks" - ); -} - -use rand::{RngExt, SeedableRng, rngs::StdRng}; - -#[tokio::test] -async fn relay_bidirectional_light_fuzzing_temporal_jitter() { - tokio::time::pause(); - let stats = Arc::new(Stats::new()); - let (mut client_peer, relay_client) = duplex(4096); - let (relay_server, 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 mut relay_task = tokio::spawn(relay_bidirectional( - client_reader, - client_writer, - server_reader, - server_writer, - 1024, - 1024, - "fuzz-user", - stats, - None, - Arc::new(BufferPool::new()), - )); - - let mut rng = StdRng::seed_from_u64(0xDEADBEEF); - - for _ in 0..10 { - // Vary timing significantly up to 1600 seconds (limit is 1800s) - let jitter = rng.random_range(100..1600); - tokio::time::sleep(Duration::from_secs(jitter)).await; - - client_peer.write_all(&[0x11]).await.unwrap(); - client_peer.flush().await.unwrap(); - - // Ensure task has not died - let res = timeout(Duration::from_millis(10), &mut relay_task).await; - assert!( - res.is_err(), - "Relay must remain open indefinitely under light temporal fuzzing with active jitter pulses" - ); - } - - drop(client_peer); - drop(server_peer); - timeout(Duration::from_secs(1), relay_task) - .await - .unwrap() - .unwrap() - .unwrap(); -} - -struct FaultyReader { - error_once: Option, -} - -struct TwoPartyGate { - arrivals: AtomicUsize, - total_bytes: AtomicUsize, - wakers: Mutex>, -} - -impl TwoPartyGate { - fn new() -> Self { - Self { - arrivals: AtomicUsize::new(0), - total_bytes: AtomicUsize::new(0), - wakers: Mutex::new(Vec::new()), - } - } - - fn arrive_or_park(&self, cx: &mut Context<'_>) -> bool { - if self.arrivals.load(Ordering::Relaxed) >= 2 { - return true; - } - - let prev = self.arrivals.fetch_add(1, Ordering::AcqRel); - if prev + 1 >= 2 { - let mut wakers = self.wakers.lock().unwrap_or_else(|p| p.into_inner()); - for waker in wakers.drain(..) { - waker.wake(); - } - true - } else { - let mut wakers = self.wakers.lock().unwrap_or_else(|p| p.into_inner()); - wakers.push(cx.waker().clone()); - false - } - } - - fn total_bytes(&self) -> usize { - self.total_bytes.load(Ordering::Relaxed) - } -} - -struct GateWriter { - gate: Arc, - entered: bool, -} - -impl GateWriter { - fn new(gate: Arc) -> Self { - Self { - gate, - entered: false, - } - } -} - -impl AsyncWrite for GateWriter { - fn poll_write( - mut self: Pin<&mut Self>, - cx: &mut Context<'_>, - buf: &[u8], - ) -> Poll> { - if !self.entered { - self.entered = true; - } - - if !self.gate.arrive_or_park(cx) { - return Poll::Pending; - } - - self.gate - .total_bytes - .fetch_add(buf.len(), Ordering::Relaxed); - Poll::Ready(Ok(buf.len())) - } - - fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { - Poll::Ready(Ok(())) - } - - fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { - Poll::Ready(Ok(())) - } -} - -struct GateReader { - gate: Arc, - entered: bool, - emitted: bool, -} - -impl GateReader { - fn new(gate: Arc) -> Self { - Self { - gate, - entered: false, - emitted: false, - } - } -} - -impl AsyncRead for GateReader { - fn poll_read( - mut self: Pin<&mut Self>, - cx: &mut Context<'_>, - buf: &mut ReadBuf<'_>, - ) -> Poll> { - if self.emitted { - return Poll::Ready(Ok(())); - } - - if !self.entered { - self.entered = true; - } - - if !self.gate.arrive_or_park(cx) { - return Poll::Pending; - } - - buf.put_slice(&[0x42]); - self.gate.total_bytes.fetch_add(1, Ordering::Relaxed); - self.emitted = true; - Poll::Ready(Ok(())) - } -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn adversarial_concurrent_quota_write_race_does_not_overshoot_limit() { - let stats = Arc::new(Stats::new()); - let gate = Arc::new(TwoPartyGate::new()); - let user = "concurrent-quota-write".to_string(); - - let writer_a = super::StatsIo::new( - GateWriter::new(Arc::clone(&gate)), - Arc::new(super::SharedCounters::new()), - Arc::clone(&stats), - user.clone(), - Some(1), - Arc::new(std::sync::atomic::AtomicBool::new(false)), - tokio::time::Instant::now(), - ); - - let writer_b = super::StatsIo::new( - GateWriter::new(Arc::clone(&gate)), - Arc::new(super::SharedCounters::new()), - Arc::clone(&stats), - user.clone(), - Some(1), - Arc::new(std::sync::atomic::AtomicBool::new(false)), - tokio::time::Instant::now(), - ); - - let task_a = tokio::spawn(async move { - let mut w = writer_a; - AsyncWriteExt::write_all(&mut w, &[0x01]).await - }); - let task_b = tokio::spawn(async move { - let mut w = writer_b; - AsyncWriteExt::write_all(&mut w, &[0x02]).await - }); - - let (res_a, res_b) = tokio::join!(task_a, task_b); - let _ = res_a.expect("task a must join"); - let _ = res_b.expect("task b must join"); - - assert!( - gate.total_bytes() <= 1, - "concurrent same-user writes must not forward more than one byte under quota=1" - ); - assert!( - stats.get_user_total_octets(&user) <= 1, - "concurrent same-user writes must not account over limit" - ); -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn adversarial_concurrent_quota_read_race_does_not_overshoot_limit() { - let stats = Arc::new(Stats::new()); - let gate = Arc::new(TwoPartyGate::new()); - let user = "concurrent-quota-read".to_string(); - - let reader_a = super::StatsIo::new( - GateReader::new(Arc::clone(&gate)), - Arc::new(super::SharedCounters::new()), - Arc::clone(&stats), - user.clone(), - Some(1), - Arc::new(std::sync::atomic::AtomicBool::new(false)), - tokio::time::Instant::now(), - ); - - let reader_b = super::StatsIo::new( - GateReader::new(Arc::clone(&gate)), - Arc::new(super::SharedCounters::new()), - Arc::clone(&stats), - user.clone(), - Some(1), - Arc::new(std::sync::atomic::AtomicBool::new(false)), - tokio::time::Instant::now(), - ); - - let task_a = tokio::spawn(async move { - let mut r = reader_a; - let mut one = [0u8; 1]; - AsyncReadExt::read_exact(&mut r, &mut one).await - }); - let task_b = tokio::spawn(async move { - let mut r = reader_b; - let mut one = [0u8; 1]; - AsyncReadExt::read_exact(&mut r, &mut one).await - }); - - let (res_a, res_b) = tokio::join!(task_a, task_b); - let _ = res_a.expect("task a must join"); - let _ = res_b.expect("task b must join"); - - assert!( - gate.total_bytes() <= 1, - "concurrent same-user reads must not consume more than one byte under quota=1" - ); - assert!( - stats.get_user_total_octets(&user) <= 1, - "concurrent same-user reads must not account over limit" - ); -} - -#[tokio::test] -async fn stress_same_user_quota_parallel_relays_never_exceed_limit() { - let stats = Arc::new(Stats::new()); - let user = "parallel-quota-user"; - - for _ in 0..128 { - let (mut client_peer_a, relay_client_a) = duplex(256); - let (relay_server_a, mut server_peer_a) = duplex(256); - let (mut client_peer_b, relay_client_b) = duplex(256); - let (relay_server_b, mut server_peer_b) = duplex(256); - - let (client_reader_a, client_writer_a) = tokio::io::split(relay_client_a); - let (server_reader_a, server_writer_a) = tokio::io::split(relay_server_a); - let (client_reader_b, client_writer_b) = tokio::io::split(relay_client_b); - let (server_reader_b, server_writer_b) = tokio::io::split(relay_server_b); - - let relay_a = tokio::spawn(relay_bidirectional( - client_reader_a, - client_writer_a, - server_reader_a, - server_writer_a, - 64, - 64, - user, - Arc::clone(&stats), - Some(1), - Arc::new(BufferPool::new()), - )); - - let relay_b = tokio::spawn(relay_bidirectional( - client_reader_b, - client_writer_b, - server_reader_b, - server_writer_b, - 64, - 64, - user, - Arc::clone(&stats), - Some(1), - Arc::new(BufferPool::new()), - )); - - let _ = tokio::join!( - client_peer_a.write_all(&[0x01]), - server_peer_a.write_all(&[0x02]), - client_peer_b.write_all(&[0x03]), - server_peer_b.write_all(&[0x04]), - ); - - let _ = timeout( - Duration::from_millis(50), - poll_fn(|cx| { - let mut one = [0u8; 1]; - let _ = Pin::new(&mut client_peer_a).poll_read(cx, &mut ReadBuf::new(&mut one)); - Poll::Ready(()) - }), - ) - .await; - - drop(client_peer_a); - drop(server_peer_a); - drop(client_peer_b); - drop(server_peer_b); - - let _ = timeout(Duration::from_secs(1), relay_a).await; - let _ = timeout(Duration::from_secs(1), relay_b).await; - - assert!( - stats.get_user_total_octets(user) <= 1, - "parallel relays must not exceed configured quota" - ); - } -} - -impl FaultyReader { - fn permission_denied_with_message(message: impl Into) -> Self { - Self { - error_once: Some(io::Error::new( - io::ErrorKind::PermissionDenied, - message.into(), - )), - } - } -} - -impl AsyncRead for FaultyReader { - fn poll_read( - mut self: Pin<&mut Self>, - _cx: &mut Context<'_>, - _buf: &mut ReadBuf<'_>, - ) -> Poll> { - if let Some(err) = self.error_once.take() { - return Poll::Ready(Err(err)); - } - Poll::Ready(Ok(())) - } -} - -#[tokio::test] -async fn relay_bidirectional_does_not_misclassify_transport_permission_denied_as_quota() { - let stats = Arc::new(Stats::new()); - let (client_peer, relay_client) = duplex(4096); - let (client_reader, client_writer) = tokio::io::split(relay_client); - - let relay_result = relay_bidirectional( - client_reader, - client_writer, - FaultyReader::permission_denied_with_message("user data quota exceeded"), - tokio::io::sink(), - 1024, - 1024, - "non-quota-permission-denied", - Arc::clone(&stats), - None, - Arc::new(BufferPool::new()), - ) - .await; - - drop(client_peer); - - assert!( - matches!(relay_result, Err(ProxyError::Io(ref err)) if err.kind() == io::ErrorKind::PermissionDenied), - "non-quota transport PermissionDenied errors must remain IO errors" - ); -} - -#[tokio::test] -async fn relay_bidirectional_light_fuzz_permission_denied_messages_remain_io_errors() { - let mut rng = StdRng::seed_from_u64(0xA11CE0B5); - - for i in 0..128u64 { - let stats = Arc::new(Stats::new()); - let (client_peer, relay_client) = duplex(1024); - let (client_reader, client_writer) = tokio::io::split(relay_client); - - let random_len = rng.random_range(1..=48); - let mut msg = String::with_capacity(random_len); - for _ in 0..random_len { - let ch = (b'a' + (rng.random::() % 26)) as char; - msg.push(ch); - } - // Include the legacy quota string in a subset of fuzz cases to validate - // collision resistance against message-based classification. - if i % 7 == 0 { - msg = "user data quota exceeded".to_string(); - } - - let relay_result = relay_bidirectional( - client_reader, - client_writer, - FaultyReader::permission_denied_with_message(msg), - tokio::io::sink(), - 1024, - 1024, - "fuzz-perm-denied", - Arc::clone(&stats), - None, - Arc::new(BufferPool::new()), - ) - .await; - - drop(client_peer); - - assert!( - matches!(relay_result, Err(ProxyError::Io(ref err)) if err.kind() == io::ErrorKind::PermissionDenied), - "transport PermissionDenied case must stay typed as IO regardless of message content" - ); - } -} - -#[tokio::test] -async fn relay_half_close_keeps_reverse_direction_progressing() { - let stats = Arc::new(Stats::new()); - let user = "half-close-user"; - - let (client_peer, relay_client) = duplex(1024); - let (relay_server, 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 (mut cp_reader, mut cp_writer) = tokio::io::split(client_peer); - let (mut sp_reader, mut sp_writer) = tokio::io::split(server_peer); - - let relay_task = tokio::spawn(relay_bidirectional( - client_reader, - client_writer, - server_reader, - server_writer, - 8192, - 8192, - user, - Arc::clone(&stats), - None, - Arc::new(BufferPool::new()), - )); - - sp_writer - .write_all(&[0x10, 0x20, 0x30, 0x40]) - .await - .unwrap(); - sp_writer.shutdown().await.unwrap(); - - let mut inbound = [0u8; 4]; - cp_reader.read_exact(&mut inbound).await.unwrap(); - assert_eq!(inbound, [0x10, 0x20, 0x30, 0x40]); - - cp_writer - .write_all(&[0xaa, 0xbb, 0xcc, 0xdd]) - .await - .unwrap(); - let mut outbound = [0u8; 4]; - sp_reader.read_exact(&mut outbound).await.unwrap(); - assert_eq!(outbound, [0xaa, 0xbb, 0xcc, 0xdd]); - - relay_task.abort(); - let joined = relay_task.await; - assert!(joined.is_err(), "aborted relay task must return join error"); -}