diff --git a/src/api/model.rs b/src/api/model.rs index efe8ebb..09eebdc 100644 --- a/src/api/model.rs +++ b/src/api/model.rs @@ -369,6 +369,7 @@ pub(super) struct UserInfo { pub(super) max_unique_ips: Option, pub(super) current_connections: u64, pub(super) active_unique_ips: usize, + pub(super) recent_unique_ips: usize, pub(super) total_octets: u64, pub(super) links: UserLinks, } diff --git a/src/api/users.rs b/src/api/users.rs index 8e90c7f..32823f1 100644 --- a/src/api/users.rs +++ b/src/api/users.rs @@ -112,6 +112,7 @@ pub(super) async fn create_user( max_unique_ips: updated_limit, current_connections: 0, active_unique_ips: 0, + recent_unique_ips: 0, total_octets: 0, links: build_user_links( &cfg, @@ -300,7 +301,7 @@ pub(super) async fn users_from_config( startup_detected_ip_v4: Option, startup_detected_ip_v6: Option, ) -> Vec { - let ip_counts = ip_tracker + let active_ip_counts = ip_tracker .get_stats() .await .into_iter() @@ -309,6 +310,7 @@ pub(super) async fn users_from_config( 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 mut users = Vec::with_capacity(names.len()); for username in names { @@ -340,7 +342,8 @@ 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: ip_counts.get(&username).copied().unwrap_or(0), + active_unique_ips: active_ip_counts.get(&username).copied().unwrap_or(0), + recent_unique_ips: recent_ip_counts.get(&username).copied().unwrap_or(0), total_octets: stats.get_user_total_octets(&username), links, username, diff --git a/src/ip_tracker.rs b/src/ip_tracker.rs index 626d591..3ff6042 100644 --- a/src/ip_tracker.rs +++ b/src/ip_tracker.rs @@ -67,22 +67,15 @@ impl UserIpTracker { let max_ips = self.max_ips.read().await; max_ips.get(username).copied() }; + let mode = *self.limit_mode.read().await; + let window = *self.limit_window.read().await; + let now = Instant::now(); let mut active_ips = self.active_ips.write().await; let user_active = active_ips .entry(username.to_string()) .or_insert_with(HashSet::new); - if limit.is_none() { - user_active.insert(ip); - return Ok(()); - } - - let limit = limit.unwrap_or_default(); - let mode = *self.limit_mode.read().await; - let window = *self.limit_window.read().await; - let now = Instant::now(); - let mut recent_ips = self.recent_ips.write().await; let user_recent = recent_ips .entry(username.to_string()) @@ -94,24 +87,26 @@ impl UserIpTracker { return Ok(()); } - let active_limit_reached = user_active.len() >= limit; - let recent_limit_reached = user_recent.len() >= limit; - let deny = match mode { - UserMaxUniqueIpsMode::ActiveWindow => active_limit_reached, - UserMaxUniqueIpsMode::TimeWindow => recent_limit_reached, - UserMaxUniqueIpsMode::Combined => active_limit_reached || recent_limit_reached, - }; + if let Some(limit) = limit { + let active_limit_reached = user_active.len() >= limit; + let recent_limit_reached = user_recent.len() >= limit; + let deny = match mode { + UserMaxUniqueIpsMode::ActiveWindow => active_limit_reached, + UserMaxUniqueIpsMode::TimeWindow => recent_limit_reached, + UserMaxUniqueIpsMode::Combined => active_limit_reached || recent_limit_reached, + }; - if deny { - return Err(format!( - "IP limit reached for user '{}': active={}/{} recent={}/{} mode={:?}", - username, - user_active.len(), - limit, - user_recent.len(), - limit, - mode - )); + if deny { + return Err(format!( + "IP limit reached for user '{}': active={}/{} recent={}/{} mode={:?}", + username, + user_active.len(), + limit, + user_recent.len(), + limit, + mode + )); + } } user_active.insert(ip); @@ -127,18 +122,26 @@ impl UserIpTracker { active_ips.remove(username); } } - drop(active_ips); + } - let mode = *self.limit_mode.read().await; - if matches!(mode, UserMaxUniqueIpsMode::ActiveWindow) { - let mut recent_ips = self.recent_ips.write().await; - if let Some(user_recent) = recent_ips.get_mut(username) { - user_recent.remove(&ip); - if user_recent.is_empty() { - recent_ips.remove(username); - } - } + pub async fn get_recent_counts_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 counts = HashMap::with_capacity(users.len()); + for user in users { + let count = if let Some(user_recent) = recent_ips.get_mut(user) { + Self::prune_recent(user_recent, now, window); + user_recent.len() + } else { + 0 + }; + counts.insert(user.clone(), count); } + + recent_ips.retain(|_, user_recent| !user_recent.is_empty()); + counts } pub async fn get_active_ip_count(&self, username: &str) -> usize { diff --git a/src/metrics.rs b/src/metrics.rs index eae69d1..1595445 100644 --- a/src/metrics.rs +++ b/src/metrics.rs @@ -1267,11 +1267,21 @@ async fn render_metrics(stats: &Stats, config: &ProxyConfig, ip_tracker: &UserIp .collect(); let mut unique_users = BTreeSet::new(); + unique_users.extend(config.access.users.keys().cloned()); unique_users.extend(config.access.user_max_unique_ips.keys().cloned()); unique_users.extend(ip_counts.keys().cloned()); + let unique_users_vec: Vec = unique_users.iter().cloned().collect(); + let recent_counts = ip_tracker + .get_recent_counts_for_users(&unique_users_vec) + .await; let _ = writeln!(out, "# HELP telemt_user_unique_ips_current Per-user current number of unique active IPs"); let _ = writeln!(out, "# TYPE telemt_user_unique_ips_current gauge"); + let _ = writeln!( + out, + "# HELP telemt_user_unique_ips_recent_window Per-user unique IPs seen in configured observation window" + ); + let _ = writeln!(out, "# TYPE telemt_user_unique_ips_recent_window gauge"); let _ = writeln!(out, "# HELP telemt_user_unique_ips_limit Per-user configured unique IP limit (0 means unlimited)"); let _ = writeln!(out, "# TYPE telemt_user_unique_ips_limit gauge"); let _ = writeln!(out, "# HELP telemt_user_unique_ips_utilization Per-user unique IP usage ratio (0 for unlimited)"); @@ -1286,6 +1296,12 @@ async fn render_metrics(stats: &Stats, config: &ProxyConfig, ip_tracker: &UserIp 0.0 }; let _ = writeln!(out, "telemt_user_unique_ips_current{{user=\"{}\"}} {}", user, current); + let _ = writeln!( + out, + "telemt_user_unique_ips_recent_window{{user=\"{}\"}} {}", + user, + recent_counts.get(&user).copied().unwrap_or(0) + ); let _ = writeln!(out, "telemt_user_unique_ips_limit{{user=\"{}\"}} {}", user, limit); let _ = writeln!( out, @@ -1378,6 +1394,7 @@ mod tests { assert!(output.contains("telemt_user_msgs_from_client{user=\"alice\"} 1")); assert!(output.contains("telemt_user_msgs_to_client{user=\"alice\"} 2")); assert!(output.contains("telemt_user_unique_ips_current{user=\"alice\"} 1")); + assert!(output.contains("telemt_user_unique_ips_recent_window{user=\"alice\"} 1")); assert!(output.contains("telemt_user_unique_ips_limit{user=\"alice\"} 4")); assert!(output.contains("telemt_user_unique_ips_utilization{user=\"alice\"} 0.250000")); } @@ -1391,7 +1408,8 @@ mod tests { assert!(output.contains("telemt_connections_total 0")); assert!(output.contains("telemt_connections_bad_total 0")); assert!(output.contains("telemt_handshake_timeouts_total 0")); - assert!(!output.contains("user=")); + assert!(output.contains("telemt_user_unique_ips_current{user=")); + assert!(output.contains("telemt_user_unique_ips_recent_window{user=")); } #[tokio::test] @@ -1412,6 +1430,7 @@ mod tests { "# TYPE telemt_me_writer_removed_unexpected_minus_restored_total gauge" )); assert!(output.contains("# TYPE telemt_user_unique_ips_current gauge")); + assert!(output.contains("# TYPE telemt_user_unique_ips_recent_window gauge")); assert!(output.contains("# TYPE telemt_user_unique_ips_limit gauge")); assert!(output.contains("# TYPE telemt_user_unique_ips_utilization gauge")); }