Unique IP always in Metrics+API: merge pull request #321 from telemt/flow-iplimit

Unique IP always in Metrics+API

src/api/model.rs

@@ -1,3 +1,5 @@
+use std::net::IpAddr;
+
 use chrono::{DateTime, Utc};
 use hyper::StatusCode;
 use rand::Rng;
@@ -369,6 +371,9 @@ pub(super) struct UserInfo {
     pub(super) max_unique_ips: Option<usize>,
     pub(super) current_connections: u64,
     pub(super) active_unique_ips: usize,
+    pub(super) active_unique_ips_list: Vec<IpAddr>,
+    pub(super) recent_unique_ips: usize,
+    pub(super) recent_unique_ips_list: Vec<IpAddr>,
     pub(super) total_octets: u64,
     pub(super) links: UserLinks,
 }

src/api/users.rs

@@ -1,4 +1,3 @@
-use std::collections::HashMap;
 use std::net::IpAddr;
 
 use hyper::StatusCode;
@@ -112,6 +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,
@@ -300,18 +302,21 @@ pub(super) async fn users_from_config(
     startup_detected_ip_v4: Option<IpAddr>,
     startup_detected_ip_v6: Option<IpAddr>,
 ) -> Vec<UserInfo> {
-    let ip_counts = ip_tracker
-        .get_stats()
-        .await
-        .into_iter()
-        .map(|(user, count, _)| (user, count))
-        .collect::<HashMap<_, _>>();
-
     let mut names = cfg.access.users.keys().cloned().collect::<Vec<_>>();
     names.sort();
+    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
@@ -340,7 +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: 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,

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<RwLock<HashMap<String, HashSet<IpAddr>>>>,
+    active_ips: Arc<RwLock<HashMap<String, HashMap<IpAddr, usize>>>>,
     recent_ips: Arc<RwLock<HashMap<String, HashMap<IpAddr, Instant>>>>,
     max_ips: Arc<RwLock<HashMap<String, usize>>>,
     limit_mode: Arc<RwLock<UserMaxUniqueIpsMode>>,
@@ -67,21 +67,14 @@ 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();
+            .or_insert_with(HashMap::new);
 
         let mut recent_ips = self.recent_ips.write().await;
         let user_recent = recent_ips
@@ -89,32 +82,35 @@ 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(());
         }
 
-        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);
+        user_active.insert(ip, 1);
         user_recent.insert(ip, now);
         Ok(())
     }
@@ -122,23 +118,73 @@ 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);
             }
         }
-        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<String, usize> {
+        let window = *self.limit_window.read().await;
+        let now = Instant::now();
+        let recent_ips = self.recent_ips.read().await;
+
+        let mut counts = HashMap::with_capacity(users.len());
+        for user in users {
+            let count = if let Some(user_recent) = recent_ips.get(user) {
+                user_recent
+                    .values()
+                    .filter(|seen_at| now.duration_since(**seen_at) <= window)
+                    .count()
+            } else {
+                0
+            };
+            counts.insert(user.clone(), count);
         }
+        counts
+    }
+
+    pub async fn get_active_ips_for_users(&self, users: &[String]) -> HashMap<String, Vec<IpAddr>> {
+        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::<Vec<_>>())
+                .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<String, Vec<IpAddr>> {
+        let window = *self.limit_window.read().await;
+        let now = Instant::now();
+        let recent_ips = self.recent_ips.read().await;
+
+        let mut out = HashMap::with_capacity(users.len());
+        for user in users {
+            let mut ips = if let Some(user_recent) = recent_ips.get(user) {
+                user_recent
+                    .iter()
+                    .filter(|(_, seen_at)| now.duration_since(**seen_at) <= window)
+                    .map(|(ip, _)| *ip)
+                    .collect::<Vec<_>>()
+            } else {
+                Vec::new()
+            };
+            ips.sort();
+            out.insert(user.clone(), ips);
+        }
+        out
     }
 
     pub async fn get_active_ip_count(&self, username: &str) -> usize {
@@ -150,7 +196,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)
     }
 
@@ -190,7 +236,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)
     }
 
@@ -266,6 +312,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();
@@ -278,6 +344,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();

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<String> = 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"));
     }