From db114f09c377a6f0916d76cead6da01489fdf7cc Mon Sep 17 00:00:00 2001
From: Alexey <247128645+axkurcom@users.noreply.github.com>
Date: Tue, 21 Apr 2026 13:30:11 +0300
Subject: [PATCH] Sync tests with code

---
 .../tests/load_mask_shape_security_tests.rs   | 11 +++---
 ...asking_consume_stress_adversarial_tests.rs | 19 +++++++---
 ...roduction_cap_regression_security_tests.rs | 35 ++++++++++++++++---
 3 files changed, 49 insertions(+), 16 deletions(-)

diff --git a/src/config/tests/load_mask_shape_security_tests.rs b/src/config/tests/load_mask_shape_security_tests.rs
index bccd36f..41b1f94 100644
--- a/src/config/tests/load_mask_shape_security_tests.rs
+++ b/src/config/tests/load_mask_shape_security_tests.rs
@@ -238,7 +238,7 @@ mask_shape_above_cap_blur_max_bytes = 8
 }
 
 #[test]
-fn load_rejects_zero_mask_relay_max_bytes() {
+fn load_accepts_zero_mask_relay_max_bytes_as_unlimited() {
     let path = write_temp_config(
         r#"
 [censorship]
@@ -246,12 +246,9 @@ mask_relay_max_bytes = 0
 "#,
     );
 
-    let err = ProxyConfig::load(&path).expect_err("mask_relay_max_bytes must be > 0");
-    let msg = err.to_string();
-    assert!(
-        msg.contains("censorship.mask_relay_max_bytes must be > 0"),
-        "error must explain non-zero relay cap invariant, got: {msg}"
-    );
+    let cfg = ProxyConfig::load(&path)
+        .expect("mask_relay_max_bytes=0 must be accepted as unlimited relay cap");
+    assert_eq!(cfg.censorship.mask_relay_max_bytes, 0);
 
     remove_temp_config(&path);
 }
diff --git a/src/proxy/tests/masking_consume_stress_adversarial_tests.rs b/src/proxy/tests/masking_consume_stress_adversarial_tests.rs
index 7579a9c..efe9f49 100644
--- a/src/proxy/tests/masking_consume_stress_adversarial_tests.rs
+++ b/src/proxy/tests/masking_consume_stress_adversarial_tests.rs
@@ -58,11 +58,22 @@ async fn consume_stall_stress_finishes_within_idle_budget() {
 }
 
 #[tokio::test]
-async fn consume_zero_cap_returns_immediately() {
+async fn consume_zero_cap_is_idle_bounded_on_stall() {
     let started = Instant::now();
-    consume_client_data(tokio::io::empty(), 0, MASK_RELAY_IDLE_TIMEOUT).await;
+    tokio::time::timeout(
+        MASK_RELAY_TIMEOUT,
+        consume_client_data(OneByteThenStall { sent: false }, 0, MASK_RELAY_IDLE_TIMEOUT),
+    )
+    .await
+    .expect("zero-cap consume path must remain bounded by timeout guards");
+
+    let elapsed = started.elapsed();
     assert!(
-        started.elapsed() < MASK_RELAY_IDLE_TIMEOUT,
-        "zero byte cap must return immediately"
+        elapsed >= (MASK_RELAY_IDLE_TIMEOUT / 2),
+        "zero cap must not short-circuit before idle timeout path, got {elapsed:?}"
+    );
+    assert!(
+        elapsed < MASK_RELAY_TIMEOUT,
+        "zero-cap consume path must complete before relay timeout, got {elapsed:?}"
     );
 }
diff --git a/src/proxy/tests/masking_production_cap_regression_security_tests.rs b/src/proxy/tests/masking_production_cap_regression_security_tests.rs
index c5d542e..46abcff 100644
--- a/src/proxy/tests/masking_production_cap_regression_security_tests.rs
+++ b/src/proxy/tests/masking_production_cap_regression_security_tests.rs
@@ -148,9 +148,10 @@ async fn positive_copy_with_production_cap_stops_exactly_at_budget() {
 }
 
 #[tokio::test]
-async fn negative_consume_with_zero_cap_performs_no_reads() {
-    let read_calls = Arc::new(AtomicUsize::new(0));
-    let reader = FinitePatternReader::new(1024, 64, Arc::clone(&read_calls));
+async fn consume_with_zero_cap_drains_until_eof() {
+    let payload = 256 * 1024;
+    let total_read = Arc::new(AtomicUsize::new(0));
+    let reader = BudgetProbeReader::new(payload, Arc::clone(&total_read));
 
     consume_client_data_with_timeout_and_cap(
         reader,
@@ -161,9 +162,33 @@ async fn negative_consume_with_zero_cap_performs_no_reads() {
     .await;
 
     assert_eq!(
-        read_calls.load(Ordering::Relaxed),
+        total_read.load(Ordering::Relaxed),
+        payload,
+        "zero cap must disable byte budget and drain finite payload to EOF"
+    );
+}
+
+#[tokio::test]
+async fn copy_with_zero_cap_drains_until_eof() {
+    let read_calls = Arc::new(AtomicUsize::new(0));
+    let payload = 73 * 1024;
+    let mut reader = FinitePatternReader::new(payload, 3072, read_calls);
+    let mut writer = CountingWriter::default();
+
+    let outcome = copy_with_idle_timeout(
+        &mut reader,
+        &mut writer,
         0,
-        "zero cap must return before reading attacker-controlled bytes"
+        true,
+        MASK_RELAY_IDLE_TIMEOUT,
+    )
+    .await;
+
+    assert_eq!(outcome.total, payload);
+    assert_eq!(writer.written, payload);
+    assert!(
+        outcome.ended_by_eof,
+        "zero cap must not terminate relay early on byte budget"
     );
 }