This commit is contained in:
Alexey
2026-03-21 15:45:29 +03:00
parent 7a8f946029
commit d7bbb376c9
154 changed files with 6194 additions and 3775 deletions

View File

@@ -5,8 +5,8 @@ use crate::stream::BufferPool;
use rand::rngs::StdRng;
use rand::{RngExt, SeedableRng};
use std::sync::Arc;
use tokio::io::{duplex, AsyncRead, AsyncReadExt, AsyncWriteExt};
use tokio::time::{timeout, Duration, Instant};
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWriteExt, duplex};
use tokio::time::{Duration, Instant, timeout};
async fn read_available<R: AsyncRead + Unpin>(reader: &mut R, budget: Duration) -> usize {
let start = Instant::now();
@@ -52,7 +52,10 @@ async fn integration_full_duplex_exact_budget_then_hard_cutoff() {
Arc::new(BufferPool::new()),
));
client_peer.write_all(&[0x10, 0x11, 0x12, 0x13]).await.unwrap();
client_peer
.write_all(&[0x10, 0x11, 0x12, 0x13])
.await
.unwrap();
let mut c2s = [0u8; 4];
server_peer.read_exact(&mut c2s).await.unwrap();
assert_eq!(c2s, [0x10, 0x11, 0x12, 0x13]);
@@ -70,8 +73,16 @@ async fn integration_full_duplex_exact_budget_then_hard_cutoff() {
let mut probe_server = [0u8; 1];
let mut probe_client = [0u8; 1];
let leaked_to_server = timeout(Duration::from_millis(120), server_peer.read(&mut probe_server)).await;
let leaked_to_client = timeout(Duration::from_millis(120), client_peer.read(&mut probe_client)).await;
let leaked_to_server = timeout(
Duration::from_millis(120),
server_peer.read(&mut probe_server),
)
.await;
let leaked_to_client = timeout(
Duration::from_millis(120),
client_peer.read(&mut probe_client),
)
.await;
assert!(
!matches!(leaked_to_server, Ok(Ok(n)) if n > 0),
@@ -126,14 +137,23 @@ async fn negative_preloaded_quota_blocks_both_directions_immediately() {
let leaked_to_server = read_available(&mut server_peer, Duration::from_millis(120)).await;
let leaked_to_client = read_available(&mut client_peer, Duration::from_millis(120)).await;
assert_eq!(leaked_to_server, 0, "preloaded limit must block C->S immediately");
assert_eq!(leaked_to_client, 0, "preloaded limit must block S->C immediately");
assert_eq!(
leaked_to_server, 0,
"preloaded limit must block C->S immediately"
);
assert_eq!(
leaked_to_client, 0,
"preloaded limit must block S->C immediately"
);
let relay_result = timeout(Duration::from_secs(2), relay)
.await
.expect("relay must terminate under preloaded cutoff")
.expect("relay task must not panic");
assert!(matches!(relay_result, Err(ProxyError::DataQuotaExceeded { .. })));
assert!(matches!(
relay_result,
Err(ProxyError::DataQuotaExceeded { .. })
));
assert!(stats.get_user_total_octets(user) <= 5);
}
@@ -160,19 +180,24 @@ async fn edge_quota_one_bidirectional_race_allows_at_most_one_forwarded_octet()
Arc::new(BufferPool::new()),
));
let _ = tokio::join!(client_peer.write_all(&[0xAA]), server_peer.write_all(&[0xBB]));
let _ = tokio::join!(
client_peer.write_all(&[0xAA]),
server_peer.write_all(&[0xBB])
);
let mut to_server = [0u8; 1];
let mut to_client = [0u8; 1];
let delivered_server = match timeout(Duration::from_millis(120), server_peer.read(&mut to_server)).await {
Ok(Ok(n)) => n,
_ => 0,
};
let delivered_client = match timeout(Duration::from_millis(120), client_peer.read(&mut to_client)).await {
Ok(Ok(n)) => n,
_ => 0,
};
let delivered_server =
match timeout(Duration::from_millis(120), server_peer.read(&mut to_server)).await {
Ok(Ok(n)) => n,
_ => 0,
};
let delivered_client =
match timeout(Duration::from_millis(120), client_peer.read(&mut to_client)).await {
Ok(Ok(n)) => n,
_ => 0,
};
assert!(
delivered_server + delivered_client <= 1,
@@ -183,7 +208,10 @@ async fn edge_quota_one_bidirectional_race_allows_at_most_one_forwarded_octet()
.await
.expect("relay must terminate under quota=1")
.expect("relay task must not panic");
assert!(matches!(relay_result, Err(ProxyError::DataQuotaExceeded { .. })));
assert!(matches!(
relay_result,
Err(ProxyError::DataQuotaExceeded { .. })
));
assert!(stats.get_user_total_octets(user) <= 1);
}
@@ -241,7 +269,10 @@ async fn adversarial_blackhat_alternating_fragmented_jitter_never_overshoots_glo
.expect("relay must terminate under black-hat jitter attack")
.expect("relay task must not panic");
assert!(matches!(relay_result, Err(ProxyError::DataQuotaExceeded { .. })));
assert!(matches!(
relay_result,
Err(ProxyError::DataQuotaExceeded { .. })
));
assert!(
delivered_to_server + delivered_to_client <= quota as usize,
"combined forwarded bytes must never exceed configured quota"
@@ -291,13 +322,17 @@ async fn light_fuzz_randomized_schedule_preserves_quota_and_forwarded_byte_invar
if rng.random::<bool>() {
let _ = client_peer.write_all(&[rng.random::<u8>()]).await;
let mut one = [0u8; 1];
if let Ok(Ok(n)) = timeout(Duration::from_millis(3), server_peer.read(&mut one)).await {
if let Ok(Ok(n)) =
timeout(Duration::from_millis(3), server_peer.read(&mut one)).await
{
delivered_total = delivered_total.saturating_add(n);
}
} else {
let _ = server_peer.write_all(&[rng.random::<u8>()]).await;
let mut one = [0u8; 1];
if let Ok(Ok(n)) = timeout(Duration::from_millis(3), client_peer.read(&mut one)).await {
if let Ok(Ok(n)) =
timeout(Duration::from_millis(3), client_peer.read(&mut one)).await
{
delivered_total = delivered_total.saturating_add(n);
}
}
@@ -312,7 +347,8 @@ async fn light_fuzz_randomized_schedule_preserves_quota_and_forwarded_byte_invar
.expect("fuzz relay task must not panic");
assert!(
relay_result.is_ok() || matches!(relay_result, Err(ProxyError::DataQuotaExceeded { .. })),
relay_result.is_ok()
|| matches!(relay_result, Err(ProxyError::DataQuotaExceeded { .. })),
"relay must either close cleanly or terminate via typed quota error"
);
assert!(
@@ -371,18 +407,25 @@ async fn stress_multi_relay_same_user_mixed_direction_jitter_respects_global_quo
if ((step as usize + worker_id as usize) & 1) == 0 {
let _ = client_peer.write_all(&[step ^ 0x3C]).await;
let mut one = [0u8; 1];
if let Ok(Ok(n)) = timeout(Duration::from_millis(3), server_peer.read(&mut one)).await {
if let Ok(Ok(n)) =
timeout(Duration::from_millis(3), server_peer.read(&mut one)).await
{
delivered = delivered.saturating_add(n);
}
} else {
let _ = server_peer.write_all(&[step ^ 0xC3]).await;
let mut one = [0u8; 1];
if let Ok(Ok(n)) = timeout(Duration::from_millis(3), client_peer.read(&mut one)).await {
if let Ok(Ok(n)) =
timeout(Duration::from_millis(3), client_peer.read(&mut one)).await
{
delivered = delivered.saturating_add(n);
}
}
tokio::time::sleep(Duration::from_millis((((worker_id as u64) + (step as u64)) % 3) + 1)).await;
tokio::time::sleep(Duration::from_millis(
(((worker_id as u64) + (step as u64)) % 3) + 1,
))
.await;
}
drop(client_peer);
@@ -393,7 +436,8 @@ async fn stress_multi_relay_same_user_mixed_direction_jitter_respects_global_quo
.expect("stress relay task must not panic");
assert!(
relay_result.is_ok() || matches!(relay_result, Err(ProxyError::DataQuotaExceeded { .. })),
relay_result.is_ok()
|| matches!(relay_result, Err(ProxyError::DataQuotaExceeded { .. })),
"stress relay must either close cleanly or terminate via typed quota error"
);
delivered
@@ -402,7 +446,8 @@ async fn stress_multi_relay_same_user_mixed_direction_jitter_respects_global_quo
let mut delivered_sum = 0usize;
for worker in workers {
delivered_sum = delivered_sum.saturating_add(worker.await.expect("stress worker must not panic"));
delivered_sum =
delivered_sum.saturating_add(worker.await.expect("stress worker must not panic"));
}
assert!(