mirror of
https://github.com/telemt/telemt.git
synced 2026-06-17 16:38:30 +03:00
Format
This commit is contained in:
@@ -1,11 +1,11 @@
|
||||
use super::*;
|
||||
use crate::config::ProxyConfig;
|
||||
use crate::stats::Stats;
|
||||
use crate::ip_tracker::UserIpTracker;
|
||||
use crate::error::ProxyError;
|
||||
use crate::ip_tracker::UserIpTracker;
|
||||
use crate::stats::Stats;
|
||||
use std::net::{IpAddr, Ipv4Addr, SocketAddr};
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
use std::net::{IpAddr, Ipv4Addr, SocketAddr};
|
||||
|
||||
// ------------------------------------------------------------------
|
||||
// Priority 3: Massive Concurrency Stress (OWASP ASVS 5.1.6)
|
||||
@@ -15,13 +15,16 @@ use std::net::{IpAddr, Ipv4Addr, SocketAddr};
|
||||
async fn client_stress_10k_connections_limit_strict() {
|
||||
let user = "stress-user";
|
||||
let limit = 512;
|
||||
|
||||
|
||||
let stats = Arc::new(Stats::new());
|
||||
let ip_tracker = Arc::new(UserIpTracker::new());
|
||||
|
||||
|
||||
let mut config = ProxyConfig::default();
|
||||
config.access.user_max_tcp_conns.insert(user.to_string(), limit);
|
||||
|
||||
config
|
||||
.access
|
||||
.user_max_tcp_conns
|
||||
.insert(user.to_string(), limit);
|
||||
|
||||
let iterations = 1000;
|
||||
let mut tasks = Vec::new();
|
||||
|
||||
@@ -30,20 +33,18 @@ async fn client_stress_10k_connections_limit_strict() {
|
||||
let ip_tracker = Arc::clone(&ip_tracker);
|
||||
let config = config.clone();
|
||||
let user_str = user.to_string();
|
||||
|
||||
|
||||
tasks.push(tokio::spawn(async move {
|
||||
let peer = SocketAddr::new(
|
||||
IpAddr::V4(Ipv4Addr::new(127, 0, 0, (i % 254 + 1) as u8)),
|
||||
10000 + (i % 1000) as u16,
|
||||
);
|
||||
|
||||
|
||||
match RunningClientHandler::acquire_user_connection_reservation_static(
|
||||
&user_str,
|
||||
&config,
|
||||
stats,
|
||||
peer,
|
||||
ip_tracker,
|
||||
).await {
|
||||
&user_str, &config, stats, peer, ip_tracker,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(res) => Ok(res),
|
||||
Err(ProxyError::ConnectionLimitExceeded { .. }) => Err(()),
|
||||
Err(e) => panic!("Unexpected error: {:?}", e),
|
||||
@@ -67,15 +68,27 @@ async fn client_stress_10k_connections_limit_strict() {
|
||||
}
|
||||
|
||||
assert_eq!(successes, limit, "Should allow exactly 'limit' connections");
|
||||
assert_eq!(failures, iterations - limit, "Should fail the rest with LimitExceeded");
|
||||
assert_eq!(
|
||||
failures,
|
||||
iterations - limit,
|
||||
"Should fail the rest with LimitExceeded"
|
||||
);
|
||||
assert_eq!(stats.get_user_curr_connects(user), limit as u64);
|
||||
|
||||
drop(reservations);
|
||||
|
||||
|
||||
ip_tracker.drain_cleanup_queue().await;
|
||||
|
||||
assert_eq!(stats.get_user_curr_connects(user), 0, "Stats must converge to 0 after all drops");
|
||||
assert_eq!(ip_tracker.get_active_ip_count(user).await, 0, "IP tracker must converge to 0");
|
||||
|
||||
assert_eq!(
|
||||
stats.get_user_curr_connects(user),
|
||||
0,
|
||||
"Stats must converge to 0 after all drops"
|
||||
);
|
||||
assert_eq!(
|
||||
ip_tracker.get_active_ip_count(user).await,
|
||||
0,
|
||||
"IP tracker must converge to 0"
|
||||
);
|
||||
}
|
||||
|
||||
// ------------------------------------------------------------------
|
||||
@@ -87,14 +100,14 @@ async fn client_ip_tracker_race_condition_stress() {
|
||||
let user = "race-user";
|
||||
let ip_tracker = Arc::new(UserIpTracker::new());
|
||||
ip_tracker.set_user_limit(user, 100).await;
|
||||
|
||||
|
||||
let iterations = 1000;
|
||||
let mut tasks = Vec::new();
|
||||
|
||||
for i in 0..iterations {
|
||||
let ip_tracker = Arc::clone(&ip_tracker);
|
||||
let ip = IpAddr::V4(Ipv4Addr::new(10, 0, 0, (i % 254 + 1) as u8));
|
||||
|
||||
|
||||
tasks.push(tokio::spawn(async move {
|
||||
for _ in 0..10 {
|
||||
if let Ok(()) = ip_tracker.check_and_add("race-user", ip).await {
|
||||
@@ -105,8 +118,12 @@ async fn client_ip_tracker_race_condition_stress() {
|
||||
}
|
||||
|
||||
futures::future::join_all(tasks).await;
|
||||
|
||||
assert_eq!(ip_tracker.get_active_ip_count(user).await, 0, "IP count must be zero after balanced add/remove burst");
|
||||
|
||||
assert_eq!(
|
||||
ip_tracker.get_active_ip_count(user).await,
|
||||
0,
|
||||
"IP count must be zero after balanced add/remove burst"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -119,7 +136,10 @@ async fn client_limit_burst_peak_never_exceeds_cap() {
|
||||
let ip_tracker = Arc::new(UserIpTracker::new());
|
||||
|
||||
let mut config = ProxyConfig::default();
|
||||
config.access.user_max_tcp_conns.insert(user.to_string(), limit);
|
||||
config
|
||||
.access
|
||||
.user_max_tcp_conns
|
||||
.insert(user.to_string(), limit);
|
||||
|
||||
let peak = Arc::new(AtomicU64::new(0));
|
||||
let mut tasks = Vec::with_capacity(attempts);
|
||||
@@ -207,10 +227,10 @@ async fn client_expiration_rejection_never_mutates_live_counters() {
|
||||
let ip_tracker = Arc::new(UserIpTracker::new());
|
||||
|
||||
let mut config = ProxyConfig::default();
|
||||
config
|
||||
.access
|
||||
.user_expirations
|
||||
.insert(user.to_string(), chrono::Utc::now() - chrono::Duration::seconds(1));
|
||||
config.access.user_expirations.insert(
|
||||
user.to_string(),
|
||||
chrono::Utc::now() - chrono::Duration::seconds(1),
|
||||
);
|
||||
|
||||
let peer: SocketAddr = "198.51.100.202:31112".parse().unwrap();
|
||||
let res = RunningClientHandler::acquire_user_connection_reservation_static(
|
||||
@@ -235,7 +255,10 @@ async fn client_ip_limit_failure_rolls_back_counter_exactly() {
|
||||
ip_tracker.set_user_limit(user, 1).await;
|
||||
|
||||
let mut config = ProxyConfig::default();
|
||||
config.access.user_max_tcp_conns.insert(user.to_string(), 16);
|
||||
config
|
||||
.access
|
||||
.user_max_tcp_conns
|
||||
.insert(user.to_string(), 16);
|
||||
|
||||
let first_peer: SocketAddr = "198.51.100.203:31113".parse().unwrap();
|
||||
let first = RunningClientHandler::acquire_user_connection_reservation_static(
|
||||
@@ -258,7 +281,10 @@ async fn client_ip_limit_failure_rolls_back_counter_exactly() {
|
||||
)
|
||||
.await;
|
||||
|
||||
assert!(matches!(second, Err(ProxyError::ConnectionLimitExceeded { .. })));
|
||||
assert!(matches!(
|
||||
second,
|
||||
Err(ProxyError::ConnectionLimitExceeded { .. })
|
||||
));
|
||||
assert_eq!(stats.get_user_curr_connects(user), 1);
|
||||
|
||||
drop(first);
|
||||
@@ -276,7 +302,10 @@ async fn client_parallel_limit_checks_success_path_leaves_no_residue() {
|
||||
ip_tracker.set_user_limit(user, 128).await;
|
||||
|
||||
let mut config = ProxyConfig::default();
|
||||
config.access.user_max_tcp_conns.insert(user.to_string(), 128);
|
||||
config
|
||||
.access
|
||||
.user_max_tcp_conns
|
||||
.insert(user.to_string(), 128);
|
||||
|
||||
let mut tasks = Vec::new();
|
||||
for i in 0..128u16 {
|
||||
@@ -310,7 +339,10 @@ async fn client_parallel_limit_checks_failure_path_leaves_no_residue() {
|
||||
ip_tracker.set_user_limit(user, 0).await;
|
||||
|
||||
let mut config = ProxyConfig::default();
|
||||
config.access.user_max_tcp_conns.insert(user.to_string(), 512);
|
||||
config
|
||||
.access
|
||||
.user_max_tcp_conns
|
||||
.insert(user.to_string(), 512);
|
||||
|
||||
let mut tasks = Vec::new();
|
||||
for i in 0..64u16 {
|
||||
@@ -319,7 +351,10 @@ async fn client_parallel_limit_checks_failure_path_leaves_no_residue() {
|
||||
let config = config.clone();
|
||||
|
||||
tasks.push(tokio::spawn(async move {
|
||||
let peer = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(172, 16, 0, (i % 250 + 1) as u8)), 33000 + i);
|
||||
let peer = SocketAddr::new(
|
||||
IpAddr::V4(Ipv4Addr::new(172, 16, 0, (i % 250 + 1) as u8)),
|
||||
33000 + i,
|
||||
);
|
||||
RunningClientHandler::check_user_limits_static(user, &config, &stats, peer, &ip_tracker)
|
||||
.await
|
||||
}));
|
||||
@@ -360,11 +395,7 @@ async fn client_churn_mixed_success_failure_converges_to_zero_state() {
|
||||
34000 + (i % 32),
|
||||
);
|
||||
let maybe_res = RunningClientHandler::acquire_user_connection_reservation_static(
|
||||
user,
|
||||
&config,
|
||||
stats,
|
||||
peer,
|
||||
ip_tracker,
|
||||
user, &config, stats, peer, ip_tracker,
|
||||
)
|
||||
.await;
|
||||
|
||||
@@ -401,11 +432,7 @@ async fn client_same_ip_parallel_attempts_allow_at_most_one_when_limit_is_one()
|
||||
let config = config.clone();
|
||||
tasks.push(tokio::spawn(async move {
|
||||
RunningClientHandler::acquire_user_connection_reservation_static(
|
||||
user,
|
||||
&config,
|
||||
stats,
|
||||
peer,
|
||||
ip_tracker,
|
||||
user, &config, stats, peer, ip_tracker,
|
||||
)
|
||||
.await
|
||||
}));
|
||||
@@ -424,7 +451,10 @@ async fn client_same_ip_parallel_attempts_allow_at_most_one_when_limit_is_one()
|
||||
}
|
||||
}
|
||||
|
||||
assert_eq!(granted, 1, "only one reservation may be granted for same IP with limit=1");
|
||||
assert_eq!(
|
||||
granted, 1,
|
||||
"only one reservation may be granted for same IP with limit=1"
|
||||
);
|
||||
drop(reservations);
|
||||
ip_tracker.drain_cleanup_queue().await;
|
||||
assert_eq!(stats.get_user_curr_connects(user), 0);
|
||||
@@ -439,7 +469,10 @@ async fn client_repeat_acquire_release_cycles_never_accumulate_state() {
|
||||
ip_tracker.set_user_limit(user, 32).await;
|
||||
|
||||
let mut config = ProxyConfig::default();
|
||||
config.access.user_max_tcp_conns.insert(user.to_string(), 32);
|
||||
config
|
||||
.access
|
||||
.user_max_tcp_conns
|
||||
.insert(user.to_string(), 32);
|
||||
|
||||
for i in 0..500u16 {
|
||||
let peer = SocketAddr::new(
|
||||
@@ -484,11 +517,7 @@ async fn client_multi_user_isolation_under_parallel_limit_exhaustion() {
|
||||
37000 + i,
|
||||
);
|
||||
RunningClientHandler::acquire_user_connection_reservation_static(
|
||||
user,
|
||||
&config,
|
||||
stats,
|
||||
peer,
|
||||
ip_tracker,
|
||||
user, &config, stats, peer, ip_tracker,
|
||||
)
|
||||
.await
|
||||
}));
|
||||
@@ -497,7 +526,11 @@ async fn client_multi_user_isolation_under_parallel_limit_exhaustion() {
|
||||
let mut u1_success = 0usize;
|
||||
let mut u2_success = 0usize;
|
||||
let mut reservations = Vec::new();
|
||||
for (idx, result) in futures::future::join_all(tasks).await.into_iter().enumerate() {
|
||||
for (idx, result) in futures::future::join_all(tasks)
|
||||
.await
|
||||
.into_iter()
|
||||
.enumerate()
|
||||
{
|
||||
let user = if idx % 2 == 0 { "u1" } else { "u2" };
|
||||
match result.unwrap() {
|
||||
Ok(reservation) => {
|
||||
@@ -556,7 +589,10 @@ async fn client_limit_recovery_after_full_rejection_wave() {
|
||||
ip_tracker.clone(),
|
||||
)
|
||||
.await;
|
||||
assert!(matches!(denied, Err(ProxyError::ConnectionLimitExceeded { .. })));
|
||||
assert!(matches!(
|
||||
denied,
|
||||
Err(ProxyError::ConnectionLimitExceeded { .. })
|
||||
));
|
||||
}
|
||||
|
||||
drop(reservation);
|
||||
@@ -572,7 +608,10 @@ async fn client_limit_recovery_after_full_rejection_wave() {
|
||||
ip_tracker.clone(),
|
||||
)
|
||||
.await;
|
||||
assert!(recovered.is_ok(), "capacity must recover after prior holder drops");
|
||||
assert!(
|
||||
recovered.is_ok(),
|
||||
"capacity must recover after prior holder drops"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
@@ -619,7 +658,10 @@ async fn client_dual_limit_cross_product_never_leaks_on_reject() {
|
||||
ip_tracker.clone(),
|
||||
)
|
||||
.await;
|
||||
assert!(matches!(denied, Err(ProxyError::ConnectionLimitExceeded { .. })));
|
||||
assert!(matches!(
|
||||
denied,
|
||||
Err(ProxyError::ConnectionLimitExceeded { .. })
|
||||
));
|
||||
}
|
||||
|
||||
assert_eq!(stats.get_user_curr_connects(user), 2);
|
||||
@@ -637,7 +679,10 @@ async fn client_check_user_limits_concurrent_churn_no_counter_drift() {
|
||||
ip_tracker.set_user_limit(user, 64).await;
|
||||
|
||||
let mut config = ProxyConfig::default();
|
||||
config.access.user_max_tcp_conns.insert(user.to_string(), 64);
|
||||
config
|
||||
.access
|
||||
.user_max_tcp_conns
|
||||
.insert(user.to_string(), 64);
|
||||
|
||||
let mut tasks = Vec::new();
|
||||
for i in 0..512u16 {
|
||||
|
||||
Reference in New Issue
Block a user