astelm/telemt
1
0
Fork 0
mirror of https://github.com/telemt/telemt.git synced 2026-04-25 14:34:10 +03:00
Code Issues Packages Projects Releases Wiki Activity

Avoid IP tracking when unique-IP limits are disabled and cap beobachten memory

Browse Source 
Co-Authored-By: brekotis <93345790+brekotis@users.noreply.github.com>
Signed-off-by: Alexey <247128645+axkurcom@users.noreply.github.com>
This commit is contained in:
Alexey
2026-04-25 12:00:46 +03:00
parent 8874396ba5
commit e78592ef9b
3 changed files with 91 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files

22
src/proxy/client.rs
View File

@@ -31,16 +31,24 @@ struct UserConnectionReservation {
ip_tracker: Arc<UserIpTracker>,
user: String,
ip: IpAddr,
tracks_ip: bool,
active: bool,
}
impl UserConnectionReservation {
fn new(stats: Arc<Stats>, ip_tracker: Arc<UserIpTracker>, user: String, ip: IpAddr) -> Self {
fn new(
stats: Arc<Stats>,
ip_tracker: Arc<UserIpTracker>,
user: String,
ip: IpAddr,
tracks_ip: bool,
) -> Self {
Self {
stats,
ip_tracker,
user,
ip,
tracks_ip,
active: true,
}
}
@@ -49,7 +57,9 @@ impl UserConnectionReservation {
if !self.active {
return;
}
if self.tracks_ip {
self.ip_tracker.remove_ip(&self.user, self.ip).await;
}
self.active = false;
self.stats.decrement_user_curr_connects(&self.user);
}
@@ -62,9 +72,11 @@ impl Drop for UserConnectionReservation {
}
self.active = false;
self.stats.decrement_user_curr_connects(&self.user);
if self.tracks_ip {
self.ip_tracker.enqueue_cleanup(self.user.clone(), self.ip);
}
}
}
use crate::config::ProxyConfig;
use crate::crypto::SecureRandom;
@@ -1600,6 +1612,8 @@ impl RunningClientHandler {
});
}
let tracks_ip = ip_tracker.get_user_limit(user).await.is_some();
if tracks_ip {
match ip_tracker.check_and_add(user, peer_addr.ip()).await {
Ok(()) => {}
Err(reason) => {
@@ -1615,12 +1629,14 @@ impl RunningClientHandler {
});
}
}
}
Ok(UserConnectionReservation::new(
stats,
ip_tracker,
user.to_string(),
peer_addr.ip(),
tracks_ip,
))
}
@@ -1663,10 +1679,10 @@ impl RunningClientHandler {
});
}
if ip_tracker.get_user_limit(user).await.is_some() {
match ip_tracker.check_and_add(user, peer_addr.ip()).await {
Ok(()) => {
ip_tracker.remove_ip(user, peer_addr.ip()).await;
stats.decrement_user_curr_connects(user);
}
Err(reason) => {
stats.decrement_user_curr_connects(user);
@@ -1681,7 +1697,9 @@ impl RunningClientHandler {
});
}
}
}
stats.decrement_user_curr_connects(user);
Ok(())
}
}

25
src/proxy/tests/client_security_tests.rs
View File

@@ -281,8 +281,13 @@ async fn user_connection_reservation_drop_enqueues_cleanup_synchronously() {
assert_eq!(ip_tracker.get_active_ip_count(&user).await, 1);
assert_eq!(stats.get_user_curr_connects(&user), 1);
let reservation =
UserConnectionReservation::new(stats.clone(), ip_tracker.clone(), user.clone(), ip);
let reservation = UserConnectionReservation::new(
stats.clone(),
ip_tracker.clone(),
user.clone(),
ip,
true,
);
// Drop the reservation synchronously without any tokio::spawn/await yielding!
drop(reservation);
@@ -320,6 +325,7 @@ async fn relay_task_abort_releases_user_gate_and_ip_reservation() {
let stats = Arc::new(Stats::new());
let ip_tracker = Arc::new(UserIpTracker::new());
ip_tracker.set_user_limit(user, 8).await;
let mut cfg = ProxyConfig::default();
cfg.access.user_max_tcp_conns.insert(user.to_string(), 8);
@@ -437,6 +443,7 @@ async fn relay_cutover_releases_user_gate_and_ip_reservation() {
let stats = Arc::new(Stats::new());
let ip_tracker = Arc::new(UserIpTracker::new());
ip_tracker.set_user_limit(user, 8).await;
let mut cfg = ProxyConfig::default();
cfg.access.user_max_tcp_conns.insert(user.to_string(), 8);
@@ -2879,6 +2886,7 @@ async fn explicit_reservation_release_cleans_user_and_ip_immediately() {
let stats = Arc::new(Stats::new());
let ip_tracker = Arc::new(UserIpTracker::new());
ip_tracker.set_user_limit(user, 4).await;
let reservation = RunningClientHandler::acquire_user_connection_reservation_static(
user,
@@ -2917,6 +2925,7 @@ async fn explicit_reservation_release_does_not_double_decrement_on_drop() {
let stats = Arc::new(Stats::new());
let ip_tracker = Arc::new(UserIpTracker::new());
ip_tracker.set_user_limit(user, 4).await;
let reservation = RunningClientHandler::acquire_user_connection_reservation_static(
user,
@@ -2947,6 +2956,7 @@ async fn drop_fallback_eventually_cleans_user_and_ip_reservation() {
let stats = Arc::new(Stats::new());
let ip_tracker = Arc::new(UserIpTracker::new());
ip_tracker.set_user_limit(user, 1).await;
let reservation = RunningClientHandler::acquire_user_connection_reservation_static(
user,
@@ -3029,6 +3039,7 @@ async fn release_abort_storm_does_not_leak_user_or_ip_reservations() {
let stats = Arc::new(Stats::new());
let ip_tracker = Arc::new(UserIpTracker::new());
ip_tracker.set_user_limit(user, ATTEMPTS + 16).await;
for idx in 0..ATTEMPTS {
let peer = SocketAddr::new(
@@ -3079,6 +3090,7 @@ async fn release_abort_loop_preserves_immediate_same_ip_reacquire() {
let stats = Arc::new(Stats::new());
let ip_tracker = Arc::new(UserIpTracker::new());
ip_tracker.set_user_limit(user, 1).await;
for _ in 0..ITERATIONS {
let reservation = RunningClientHandler::acquire_user_connection_reservation_static(
@@ -3137,6 +3149,7 @@ async fn adversarial_mixed_release_drop_abort_wave_converges_to_zero() {
let stats = Arc::new(Stats::new());
let ip_tracker = Arc::new(UserIpTracker::new());
ip_tracker.set_user_limit(user, RESERVATIONS + 8).await;
let mut reservations = Vec::with_capacity(RESERVATIONS);
for idx in 0..RESERVATIONS {
@@ -3217,6 +3230,8 @@ async fn parallel_users_abort_release_isolation_preserves_independent_cleanup()
let stats = Arc::new(Stats::new());
let ip_tracker = Arc::new(UserIpTracker::new());
ip_tracker.set_user_limit(user_a, 64).await;
ip_tracker.set_user_limit(user_b, 64).await;
let mut tasks = tokio::task::JoinSet::new();
for idx in 0..64usize {
@@ -3278,6 +3293,7 @@ async fn concurrent_release_storm_leaves_zero_user_and_ip_footprint() {
let stats = Arc::new(Stats::new());
let ip_tracker = Arc::new(UserIpTracker::new());
ip_tracker.set_user_limit(user, RESERVATIONS + 8).await;
let mut reservations = Vec::with_capacity(RESERVATIONS);
for idx in 0..RESERVATIONS {
@@ -3332,6 +3348,7 @@ async fn relay_connect_error_releases_user_and_ip_before_return() {
let stats = Arc::new(Stats::new());
let ip_tracker = Arc::new(UserIpTracker::new());
ip_tracker.set_user_limit(user, 8).await;
let mut config = ProxyConfig::default();
config.access.user_max_tcp_conns.insert(user.to_string(), 1);
@@ -3427,6 +3444,7 @@ async fn mixed_release_and_drop_same_ip_preserves_counter_correctness() {
let stats = Arc::new(Stats::new());
let ip_tracker = Arc::new(UserIpTracker::new());
ip_tracker.set_user_limit(user, 1).await;
let reservation_a = RunningClientHandler::acquire_user_connection_reservation_static(
user,
@@ -3487,6 +3505,7 @@ async fn drop_one_of_two_same_ip_reservations_keeps_ip_active() {
let stats = Arc::new(Stats::new());
let ip_tracker = Arc::new(UserIpTracker::new());
ip_tracker.set_user_limit(user, 1).await;
let reservation_a = RunningClientHandler::acquire_user_connection_reservation_static(
user,
@@ -3696,6 +3715,7 @@ async fn cross_thread_drop_uses_captured_runtime_for_ip_cleanup() {
let stats = Arc::new(Stats::new());
let ip_tracker = Arc::new(UserIpTracker::new());
ip_tracker.set_user_limit(user, 8).await;
let reservation = RunningClientHandler::acquire_user_connection_reservation_static(
user,
@@ -3740,6 +3760,7 @@ async fn immediate_reacquire_after_cross_thread_drop_succeeds() {
let stats = Arc::new(Stats::new());
let ip_tracker = Arc::new(UserIpTracker::new());
ip_tracker.set_user_limit(user, 1).await;
let reservation = RunningClientHandler::acquire_user_connection_reservation_static(
user,

20
src/stats/beobachten.rs
View File

@@ -7,6 +7,7 @@ use std::time::{Duration, Instant};
use parking_lot::Mutex;
const CLEANUP_INTERVAL: Duration = Duration::from_secs(30);
const MAX_BEOBACHTEN_ENTRIES: usize = 65_536;
#[derive(Default)]
struct BeobachtenInner {
@@ -48,12 +49,23 @@ impl BeobachtenStore {
Self::cleanup_if_needed(&mut guard, now, ttl);
let key = (class.to_string(), ip);
let entry = guard.entries.entry(key).or_insert(BeobachtenEntry {
tries: 0,
last_seen: now,
});
if let Some(entry) = guard.entries.get_mut(&key) {
entry.tries = entry.tries.saturating_add(1);
entry.last_seen = now;
return;
}
if guard.entries.len() >= MAX_BEOBACHTEN_ENTRIES {
return;
}
guard.entries.insert(
key,
BeobachtenEntry {
tries: 1,
last_seen: now,
},
);
}
pub fn snapshot_text(&self, ttl: Duration) -> String {