From 0b1a8cd3f82dec0b20c01b2fbdd24244264ec0b2 Mon Sep 17 00:00:00 2001 From: Alexey <247128645+axkurcom@users.noreply.github.com> Date: Thu, 5 Mar 2026 13:41:41 +0300 Subject: [PATCH] IP Limit fixes --- src/api/model.rs | 4 ++ src/api/users.rs | 27 +++++++------ src/ip_tracker.rs | 96 +++++++++++++++++++++++++++++++++++++++++++---- 3 files changed, 108 insertions(+), 19 deletions(-) diff --git a/src/api/model.rs b/src/api/model.rs index 09eebdc..2f6c58e 100644 --- a/src/api/model.rs +++ b/src/api/model.rs @@ -1,3 +1,5 @@ +use std::net::IpAddr; + use chrono::{DateTime, Utc}; use hyper::StatusCode; use rand::Rng; @@ -369,7 +371,9 @@ pub(super) struct UserInfo { pub(super) max_unique_ips: Option, pub(super) current_connections: u64, pub(super) active_unique_ips: usize, + pub(super) active_unique_ips_list: Vec, pub(super) recent_unique_ips: usize, + pub(super) recent_unique_ips_list: Vec, pub(super) total_octets: u64, pub(super) links: UserLinks, } diff --git a/src/api/users.rs b/src/api/users.rs index 32823f1..d156896 100644 --- a/src/api/users.rs +++ b/src/api/users.rs @@ -1,4 +1,3 @@ -use std::collections::HashMap; use std::net::IpAddr; use hyper::StatusCode; @@ -112,7 +111,9 @@ pub(super) async fn create_user( max_unique_ips: updated_limit, current_connections: 0, active_unique_ips: 0, + active_unique_ips_list: Vec::new(), recent_unique_ips: 0, + recent_unique_ips_list: Vec::new(), total_octets: 0, links: build_user_links( &cfg, @@ -301,19 +302,21 @@ pub(super) async fn users_from_config( startup_detected_ip_v4: Option, startup_detected_ip_v6: Option, ) -> Vec { - let active_ip_counts = ip_tracker - .get_stats() - .await - .into_iter() - .map(|(user, count, _)| (user, count)) - .collect::>(); - let mut names = cfg.access.users.keys().cloned().collect::>(); names.sort(); - let recent_ip_counts = ip_tracker.get_recent_counts_for_users(&names).await; + let active_ip_lists = ip_tracker.get_active_ips_for_users(&names).await; + let recent_ip_lists = ip_tracker.get_recent_ips_for_users(&names).await; let mut users = Vec::with_capacity(names.len()); for username in names { + let active_ip_list = active_ip_lists + .get(&username) + .cloned() + .unwrap_or_else(Vec::new); + let recent_ip_list = recent_ip_lists + .get(&username) + .cloned() + .unwrap_or_else(Vec::new); let links = cfg .access .users @@ -342,8 +345,10 @@ pub(super) async fn users_from_config( data_quota_bytes: cfg.access.user_data_quota.get(&username).copied(), max_unique_ips: cfg.access.user_max_unique_ips.get(&username).copied(), current_connections: stats.get_user_curr_connects(&username), - active_unique_ips: active_ip_counts.get(&username).copied().unwrap_or(0), - recent_unique_ips: recent_ip_counts.get(&username).copied().unwrap_or(0), + active_unique_ips: active_ip_list.len(), + active_unique_ips_list: active_ip_list, + recent_unique_ips: recent_ip_list.len(), + recent_unique_ips_list: recent_ip_list, total_octets: stats.get_user_total_octets(&username), links, username, diff --git a/src/ip_tracker.rs b/src/ip_tracker.rs index 3ff6042..492d642 100644 --- a/src/ip_tracker.rs +++ b/src/ip_tracker.rs @@ -2,7 +2,7 @@ #![allow(dead_code)] -use std::collections::{HashMap, HashSet}; +use std::collections::HashMap; use std::net::IpAddr; use std::sync::Arc; use std::time::{Duration, Instant}; @@ -13,7 +13,7 @@ use crate::config::UserMaxUniqueIpsMode; #[derive(Debug, Clone)] pub struct UserIpTracker { - active_ips: Arc>>>, + active_ips: Arc>>>, recent_ips: Arc>>>, max_ips: Arc>>, limit_mode: Arc>, @@ -74,7 +74,7 @@ impl UserIpTracker { let mut active_ips = self.active_ips.write().await; let user_active = active_ips .entry(username.to_string()) - .or_insert_with(HashSet::new); + .or_insert_with(HashMap::new); let mut recent_ips = self.recent_ips.write().await; let user_recent = recent_ips @@ -82,7 +82,8 @@ impl UserIpTracker { .or_insert_with(HashMap::new); Self::prune_recent(user_recent, now, window); - if user_active.contains(&ip) { + if let Some(count) = user_active.get_mut(&ip) { + *count = count.saturating_add(1); user_recent.insert(ip, now); return Ok(()); } @@ -109,7 +110,7 @@ impl UserIpTracker { } } - user_active.insert(ip); + user_active.insert(ip, 1); user_recent.insert(ip, now); Ok(()) } @@ -117,7 +118,13 @@ impl UserIpTracker { pub async fn remove_ip(&self, username: &str, ip: IpAddr) { let mut active_ips = self.active_ips.write().await; if let Some(user_ips) = active_ips.get_mut(username) { - user_ips.remove(&ip); + if let Some(count) = user_ips.get_mut(&ip) { + if *count > 1 { + *count -= 1; + } else { + user_ips.remove(&ip); + } + } if user_ips.is_empty() { active_ips.remove(username); } @@ -144,6 +151,41 @@ impl UserIpTracker { counts } + pub async fn get_active_ips_for_users(&self, users: &[String]) -> HashMap> { + let active_ips = self.active_ips.read().await; + let mut out = HashMap::with_capacity(users.len()); + for user in users { + let mut ips = active_ips + .get(user) + .map(|per_ip| per_ip.keys().copied().collect::>()) + .unwrap_or_else(Vec::new); + ips.sort(); + out.insert(user.clone(), ips); + } + out + } + + pub async fn get_recent_ips_for_users(&self, users: &[String]) -> HashMap> { + let window = *self.limit_window.read().await; + let now = Instant::now(); + let mut recent_ips = self.recent_ips.write().await; + + let mut out = HashMap::with_capacity(users.len()); + for user in users { + let mut ips = if let Some(user_recent) = recent_ips.get_mut(user) { + Self::prune_recent(user_recent, now, window); + user_recent.keys().copied().collect::>() + } else { + Vec::new() + }; + ips.sort(); + out.insert(user.clone(), ips); + } + + recent_ips.retain(|_, user_recent| !user_recent.is_empty()); + out + } + pub async fn get_active_ip_count(&self, username: &str) -> usize { let active_ips = self.active_ips.read().await; active_ips.get(username).map(|ips| ips.len()).unwrap_or(0) @@ -153,7 +195,7 @@ impl UserIpTracker { let active_ips = self.active_ips.read().await; active_ips .get(username) - .map(|ips| ips.iter().copied().collect()) + .map(|ips| ips.keys().copied().collect()) .unwrap_or_else(Vec::new) } @@ -193,7 +235,7 @@ impl UserIpTracker { let active_ips = self.active_ips.read().await; active_ips .get(username) - .map(|ips| ips.contains(&ip)) + .map(|ips| ips.contains_key(&ip)) .unwrap_or(false) } @@ -269,6 +311,26 @@ mod tests { assert_eq!(tracker.get_active_ip_count("test_user").await, 2); } + #[tokio::test] + async fn test_active_window_rejects_new_ip_and_keeps_existing_session() { + let tracker = UserIpTracker::new(); + tracker.set_user_limit("test_user", 1).await; + tracker + .set_limit_policy(UserMaxUniqueIpsMode::ActiveWindow, 30) + .await; + + let ip1 = test_ipv4(10, 10, 10, 1); + let ip2 = test_ipv4(10, 10, 10, 2); + + assert!(tracker.check_and_add("test_user", ip1).await.is_ok()); + assert!(tracker.is_ip_active("test_user", ip1).await); + assert!(tracker.check_and_add("test_user", ip2).await.is_err()); + + // Existing session remains active; only new unique IP is denied. + assert!(tracker.is_ip_active("test_user", ip1).await); + assert_eq!(tracker.get_active_ip_count("test_user").await, 1); + } + #[tokio::test] async fn test_reconnection_from_same_ip() { let tracker = UserIpTracker::new(); @@ -281,6 +343,24 @@ mod tests { assert_eq!(tracker.get_active_ip_count("test_user").await, 1); } + #[tokio::test] + async fn test_same_ip_disconnect_keeps_active_while_other_session_alive() { + let tracker = UserIpTracker::new(); + tracker.set_user_limit("test_user", 2).await; + + let ip1 = test_ipv4(192, 168, 1, 1); + + assert!(tracker.check_and_add("test_user", ip1).await.is_ok()); + assert!(tracker.check_and_add("test_user", ip1).await.is_ok()); + assert_eq!(tracker.get_active_ip_count("test_user").await, 1); + + tracker.remove_ip("test_user", ip1).await; + assert_eq!(tracker.get_active_ip_count("test_user").await, 1); + + tracker.remove_ip("test_user", ip1).await; + assert_eq!(tracker.get_active_ip_count("test_user").await, 0); + } + #[tokio::test] async fn test_ip_removal() { let tracker = UserIpTracker::new();