mirror of
https://github.com/telemt/telemt.git
synced 2026-05-01 17:34:09 +03:00
Format
This commit is contained in:
@@ -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!(
|
||||
|
||||
Reference in New Issue
Block a user