From 7d543aeb676ed0011e2ef1be1673ad43d0fc3969 Mon Sep 17 00:00:00 2001
From: Alexey <247128645+axkurcom@users.noreply.github.com>
Date: Fri, 5 Jun 2026 12:59:37 +0300
Subject: [PATCH] Fixes for Adversarial Timing Profile Latency-flake by #761

Co-Authored-By: brekotis <93345790+brekotis@users.noreply.github.com>
---
 ...client_timing_profile_adversarial_tests.rs | 26 ++++++++++---------
 1 file changed, 14 insertions(+), 12 deletions(-)

diff --git a/src/proxy/tests/client_timing_profile_adversarial_tests.rs b/src/proxy/tests/client_timing_profile_adversarial_tests.rs
index 9f61e3c..44e194b 100644
--- a/src/proxy/tests/client_timing_profile_adversarial_tests.rs
+++ b/src/proxy/tests/client_timing_profile_adversarial_tests.rs
@@ -75,12 +75,17 @@ async fn run_generic_once(class: ProbeClass) -> u128 {
     let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
     let backend_addr = listener.local_addr().unwrap();
     let backend_reply = REPLY_404.to_vec();
+    let probe = match class {
+        ProbeClass::MalformedTlsTruncation => malformed_tls_probe(),
+        ProbeClass::PlainWebBaseline => plain_web_probe(),
+    };
 
     let accept_task = tokio::spawn({
         let backend_reply = backend_reply.clone();
+        let expected_probe_len = probe.len();
         async move {
             let (mut stream, _) = listener.accept().await.unwrap();
-            let mut buf = [0u8; 5];
+            let mut buf = vec![0u8; expected_probe_len];
             stream.read_exact(&mut buf).await.unwrap();
             stream.write_all(&backend_reply).await.unwrap();
         }
@@ -94,6 +99,7 @@ async fn run_generic_once(class: ProbeClass) -> u128 {
     cfg.censorship.mask_host = Some("127.0.0.1".to_string());
     cfg.censorship.mask_port = backend_addr.port();
     cfg.censorship.mask_proxy_protocol = 0;
+    cfg.censorship.mask_shape_hardening = false;
 
     if matches!(class, ProbeClass::PlainWebBaseline) {
         cfg.general.modes.classic = false;
@@ -130,11 +136,6 @@ async fn run_generic_once(class: ProbeClass) -> u128 {
         false,
     ));
 
-    let probe = match class {
-        ProbeClass::MalformedTlsTruncation => malformed_tls_probe(),
-        ProbeClass::PlainWebBaseline => plain_web_probe(),
-    };
-
     let started = Instant::now();
     client_side.write_all(&probe).await.unwrap();
     client_side.shutdown().await.unwrap();
@@ -170,11 +171,16 @@ async fn run_client_handler_once(class: ProbeClass) -> u128 {
     let front_addr = front_listener.local_addr().unwrap();
 
     let backend_reply = REPLY_404.to_vec();
+    let probe = match class {
+        ProbeClass::MalformedTlsTruncation => malformed_tls_probe(),
+        ProbeClass::PlainWebBaseline => plain_web_probe(),
+    };
     let mask_accept_task = tokio::spawn({
         let backend_reply = backend_reply.clone();
+        let expected_probe_len = probe.len();
         async move {
             let (mut stream, _) = mask_listener.accept().await.unwrap();
-            let mut buf = [0u8; 5];
+            let mut buf = vec![0u8; expected_probe_len];
             stream.read_exact(&mut buf).await.unwrap();
             stream.write_all(&backend_reply).await.unwrap();
         }
@@ -188,6 +194,7 @@ async fn run_client_handler_once(class: ProbeClass) -> u128 {
     cfg.censorship.mask_host = Some("127.0.0.1".to_string());
     cfg.censorship.mask_port = backend_addr.port();
     cfg.censorship.mask_proxy_protocol = 0;
+    cfg.censorship.mask_shape_hardening = false;
 
     if matches!(class, ProbeClass::PlainWebBaseline) {
         cfg.general.modes.classic = false;
@@ -240,11 +247,6 @@ async fn run_client_handler_once(class: ProbeClass) -> u128 {
         })
     };
 
-    let probe = match class {
-        ProbeClass::MalformedTlsTruncation => malformed_tls_probe(),
-        ProbeClass::PlainWebBaseline => plain_web_probe(),
-    };
-
     let mut client = TcpStream::connect(front_addr).await.unwrap();
     let started = Instant::now();
     client.write_all(&probe).await.unwrap();