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

@@ -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 {