mirror of
https://github.com/telemt/telemt.git
synced 2026-04-18 19:14:09 +03:00
Add comprehensive tests for relay quota management and adversarial scenarios
- Introduced `relay_quota_boundary_blackhat_tests.rs` to validate behavior under quota limits, including edge cases and adversarial conditions. - Added `relay_quota_model_adversarial_tests.rs` to ensure quota management maintains integrity during bidirectional communication and various load scenarios. - Created `relay_quota_overflow_regression_tests.rs` to address overflow issues and ensure that quota limits are respected during aggressive data transmission. - Implemented `route_mode_coherence_adversarial_tests.rs` to verify the consistency of route mode transitions and timestamp management across different relay modes.
This commit is contained in:
David Osipov
300
src/proxy/tests/relay_quota_model_adversarial_tests.rs
Normal file
300
src/proxy/tests/relay_quota_model_adversarial_tests.rs
Normal file
@@ -0,0 +1,300 @@
|
||||
use super::relay_bidirectional;
|
||||
use crate::error::ProxyError;
|
||||
use crate::stats::Stats;
|
||||
use crate::stream::BufferPool;
|
||||
use rand::rngs::StdRng;
|
||||
use rand::{Rng, SeedableRng};
|
||||
use std::sync::Arc;
|
||||
use tokio::io::{duplex, AsyncRead, AsyncReadExt, AsyncWriteExt};
|
||||
use tokio::sync::Barrier;
|
||||
use tokio::time::{timeout, Duration};
|
||||
|
||||
fn assert_is_prefix(received: &[u8], sent: &[u8], direction: &str) {
|
||||
assert!(
|
||||
sent.starts_with(received),
|
||||
"{direction} stream corruption: received={} sent={} (received must be prefix of sent)",
|
||||
received.len(),
|
||||
sent.len()
|
||||
);
|
||||
}
|
||||
|
||||
async fn drain_available<R: AsyncRead + Unpin>(reader: &mut R, out: &mut Vec<u8>, rounds: usize) {
|
||||
for _ in 0..rounds {
|
||||
let mut buf = [0u8; 64];
|
||||
match timeout(Duration::from_millis(2), reader.read(&mut buf)).await {
|
||||
Ok(Ok(0)) => break,
|
||||
Ok(Ok(n)) => out.extend_from_slice(&buf[..n]),
|
||||
Ok(Err(_)) | Err(_) => break,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn model_fuzz_bidirectional_schedule_preserves_prefixes_and_quota_budget() {
|
||||
let mut rng = StdRng::seed_from_u64(0xC0DE_CAFE_D15C_F00D);
|
||||
|
||||
for case in 0..64u64 {
|
||||
let stats = Arc::new(Stats::new());
|
||||
let user = format!("quota-model-fuzz-{case}");
|
||||
let quota = rng.random_range(1u64..=64u64);
|
||||
|
||||
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_user = user.clone();
|
||||
let relay_stats = Arc::clone(&stats);
|
||||
let relay = tokio::spawn(async move {
|
||||
relay_bidirectional(
|
||||
client_reader,
|
||||
client_writer,
|
||||
server_reader,
|
||||
server_writer,
|
||||
256,
|
||||
256,
|
||||
&relay_user,
|
||||
relay_stats,
|
||||
Some(quota),
|
||||
Arc::new(BufferPool::new()),
|
||||
)
|
||||
.await
|
||||
});
|
||||
|
||||
let mut sent_c2s = Vec::new();
|
||||
let mut sent_s2c = Vec::new();
|
||||
let mut recv_at_server = Vec::new();
|
||||
let mut recv_at_client = Vec::new();
|
||||
|
||||
for _ in 0..96usize {
|
||||
if relay.is_finished() {
|
||||
break;
|
||||
}
|
||||
|
||||
let do_c2s = rng.random::<bool>();
|
||||
let chunk_len = rng.random_range(1usize..=12usize);
|
||||
let mut chunk = vec![0u8; chunk_len];
|
||||
for b in &mut chunk {
|
||||
*b = rng.random::<u8>();
|
||||
}
|
||||
|
||||
if do_c2s {
|
||||
if client_peer.write_all(&chunk).await.is_ok() {
|
||||
sent_c2s.extend_from_slice(&chunk);
|
||||
}
|
||||
} else if server_peer.write_all(&chunk).await.is_ok() {
|
||||
sent_s2c.extend_from_slice(&chunk);
|
||||
}
|
||||
|
||||
drain_available(&mut server_peer, &mut recv_at_server, 2).await;
|
||||
drain_available(&mut client_peer, &mut recv_at_client, 2).await;
|
||||
|
||||
assert_is_prefix(&recv_at_server, &sent_c2s, "C->S");
|
||||
assert_is_prefix(&recv_at_client, &sent_s2c, "S->C");
|
||||
assert!(
|
||||
recv_at_server.len() + recv_at_client.len() <= quota as usize,
|
||||
"fuzz case {case}: delivered bytes exceed quota"
|
||||
);
|
||||
assert!(
|
||||
stats.get_user_total_octets(&user) <= quota,
|
||||
"fuzz case {case}: accounted bytes exceed quota"
|
||||
);
|
||||
}
|
||||
|
||||
drop(client_peer);
|
||||
drop(server_peer);
|
||||
|
||||
let relay_result = timeout(Duration::from_secs(2), relay)
|
||||
.await
|
||||
.expect("fuzz relay must terminate")
|
||||
.expect("fuzz relay task must not panic");
|
||||
|
||||
assert!(
|
||||
relay_result.is_ok() || matches!(relay_result, Err(ProxyError::DataQuotaExceeded { .. })),
|
||||
"fuzz case {case}: relay must end cleanly or with typed quota error"
|
||||
);
|
||||
|
||||
assert_is_prefix(&recv_at_server, &sent_c2s, "C->S final");
|
||||
assert_is_prefix(&recv_at_client, &sent_s2c, "S->C final");
|
||||
assert!(recv_at_server.len() + recv_at_client.len() <= quota as usize);
|
||||
assert!(stats.get_user_total_octets(&user) <= quota);
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
|
||||
async fn adversarial_dual_direction_cutoff_race_allows_at_most_one_forwarded_byte() {
|
||||
let stats = Arc::new(Stats::new());
|
||||
let user = "quota-dual-race-user";
|
||||
|
||||
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 = tokio::spawn(relay_bidirectional(
|
||||
client_reader,
|
||||
client_writer,
|
||||
server_reader,
|
||||
server_writer,
|
||||
128,
|
||||
128,
|
||||
user,
|
||||
Arc::clone(&stats),
|
||||
Some(1),
|
||||
Arc::new(BufferPool::new()),
|
||||
));
|
||||
|
||||
let gate = Arc::new(Barrier::new(3));
|
||||
|
||||
let writer_c2s = {
|
||||
let gate = Arc::clone(&gate);
|
||||
tokio::spawn(async move {
|
||||
gate.wait().await;
|
||||
let _ = client_peer.write_all(&[0xA1]).await;
|
||||
client_peer
|
||||
})
|
||||
};
|
||||
|
||||
let writer_s2c = {
|
||||
let gate = Arc::clone(&gate);
|
||||
tokio::spawn(async move {
|
||||
gate.wait().await;
|
||||
let _ = server_peer.write_all(&[0xB2]).await;
|
||||
server_peer
|
||||
})
|
||||
};
|
||||
|
||||
gate.wait().await;
|
||||
|
||||
let mut client_peer = writer_c2s.await.expect("c2s writer must not panic");
|
||||
let mut server_peer = writer_s2c.await.expect("s2c writer must not panic");
|
||||
|
||||
let mut got_at_server = [0u8; 1];
|
||||
let mut got_at_client = [0u8; 1];
|
||||
|
||||
let n_server = match timeout(Duration::from_millis(120), server_peer.read(&mut got_at_server)).await {
|
||||
Ok(Ok(n)) => n,
|
||||
_ => 0,
|
||||
};
|
||||
let n_client = match timeout(Duration::from_millis(120), client_peer.read(&mut got_at_client)).await {
|
||||
Ok(Ok(n)) => n,
|
||||
_ => 0,
|
||||
};
|
||||
|
||||
assert!(
|
||||
n_server + n_client <= 1,
|
||||
"quota=1 race must not forward both concurrent direction bytes"
|
||||
);
|
||||
|
||||
drop(client_peer);
|
||||
drop(server_peer);
|
||||
|
||||
let relay_result = timeout(Duration::from_secs(2), relay)
|
||||
.await
|
||||
.expect("quota race relay must terminate")
|
||||
.expect("quota race relay task must not panic");
|
||||
|
||||
assert!(matches!(relay_result, Err(ProxyError::DataQuotaExceeded { .. })));
|
||||
assert!(stats.get_user_total_octets(user) <= 1);
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
|
||||
async fn stress_shared_user_multi_relay_global_quota_never_overshoots_under_model_load() {
|
||||
let stats = Arc::new(Stats::new());
|
||||
let user = "quota-model-stress-user";
|
||||
let quota = 96u64;
|
||||
|
||||
let mut workers = Vec::new();
|
||||
for worker_id in 0..6u64 {
|
||||
let stats = Arc::clone(&stats);
|
||||
let user = user.to_string();
|
||||
|
||||
workers.push(tokio::spawn(async move {
|
||||
let mut rng = StdRng::seed_from_u64(0x9E37_79B9_7F4A_7C15 ^ worker_id);
|
||||
|
||||
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_user = user.clone();
|
||||
let relay_stats = Arc::clone(&stats);
|
||||
let relay = tokio::spawn(async move {
|
||||
relay_bidirectional(
|
||||
client_reader,
|
||||
client_writer,
|
||||
server_reader,
|
||||
server_writer,
|
||||
192,
|
||||
192,
|
||||
&relay_user,
|
||||
relay_stats,
|
||||
Some(quota),
|
||||
Arc::new(BufferPool::new()),
|
||||
)
|
||||
.await
|
||||
});
|
||||
|
||||
let mut sent_c2s = Vec::new();
|
||||
let mut sent_s2c = Vec::new();
|
||||
let mut recv_at_server = Vec::new();
|
||||
let mut recv_at_client = Vec::new();
|
||||
|
||||
for _ in 0..64usize {
|
||||
if relay.is_finished() {
|
||||
break;
|
||||
}
|
||||
|
||||
let choose_c2s = rng.random::<bool>();
|
||||
let len = rng.random_range(1usize..=10usize);
|
||||
let mut payload = vec![0u8; len];
|
||||
for b in &mut payload {
|
||||
*b = rng.random::<u8>();
|
||||
}
|
||||
|
||||
if choose_c2s {
|
||||
if client_peer.write_all(&payload).await.is_ok() {
|
||||
sent_c2s.extend_from_slice(&payload);
|
||||
}
|
||||
} else if server_peer.write_all(&payload).await.is_ok() {
|
||||
sent_s2c.extend_from_slice(&payload);
|
||||
}
|
||||
|
||||
drain_available(&mut server_peer, &mut recv_at_server, 2).await;
|
||||
drain_available(&mut client_peer, &mut recv_at_client, 2).await;
|
||||
|
||||
assert_is_prefix(&recv_at_server, &sent_c2s, "stress C->S");
|
||||
assert_is_prefix(&recv_at_client, &sent_s2c, "stress S->C");
|
||||
}
|
||||
|
||||
drop(client_peer);
|
||||
drop(server_peer);
|
||||
|
||||
let relay_result = timeout(Duration::from_secs(2), relay)
|
||||
.await
|
||||
.expect("stress relay must terminate")
|
||||
.expect("stress relay task must not panic");
|
||||
|
||||
assert!(
|
||||
relay_result.is_ok() || matches!(relay_result, Err(ProxyError::DataQuotaExceeded { .. })),
|
||||
"stress relay must end cleanly or with typed quota error"
|
||||
);
|
||||
|
||||
recv_at_server.len() + recv_at_client.len()
|
||||
}));
|
||||
}
|
||||
|
||||
let mut delivered_sum = 0usize;
|
||||
for worker in workers {
|
||||
delivered_sum = delivered_sum.saturating_add(worker.await.expect("worker must not panic"));
|
||||
}
|
||||
|
||||
assert!(
|
||||
stats.get_user_total_octets(user) <= quota,
|
||||
"global per-user quota must never overshoot under concurrent multi-relay model load"
|
||||
);
|
||||
assert!(
|
||||
delivered_sum <= quota as usize,
|
||||
"aggregate delivered bytes across relays must remain within global quota"
|
||||
);
|
||||
}
|
||||
Reference in New Issue
Block a user