From 9412f089c08230a93d58a6c34e8ea89e8cf13a07 Mon Sep 17 00:00:00 2001
From: Alexey <247128645+axkurcom@users.noreply.github.com>
Date: Sat, 25 Apr 2026 15:49:28 +0300
Subject: [PATCH] Restore active IP observability for users without unique-IP
 limits

---
 src/proxy/client.rs                      | 63 +++++++++++-------------
 src/proxy/tests/client_security_tests.rs | 30 +++++++++++
 2 files changed, 59 insertions(+), 34 deletions(-)

diff --git a/src/proxy/client.rs b/src/proxy/client.rs
index 2ab02ce..e28d0b6 100644
--- a/src/proxy/client.rs
+++ b/src/proxy/client.rs
@@ -1612,22 +1612,19 @@ 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) => {
-                    stats.decrement_user_curr_connects(user);
-                    warn!(
-                        user = %user,
-                        ip = %peer_addr.ip(),
-                        reason = %reason,
-                        "IP limit exceeded"
-                    );
-                    return Err(ProxyError::ConnectionLimitExceeded {
-                        user: user.to_string(),
-                    });
-                }
+        match ip_tracker.check_and_add(user, peer_addr.ip()).await {
+            Ok(()) => {}
+            Err(reason) => {
+                stats.decrement_user_curr_connects(user);
+                warn!(
+                    user = %user,
+                    ip = %peer_addr.ip(),
+                    reason = %reason,
+                    "IP limit exceeded"
+                );
+                return Err(ProxyError::ConnectionLimitExceeded {
+                    user: user.to_string(),
+                });
             }
         }
 
@@ -1636,7 +1633,7 @@ impl RunningClientHandler {
             ip_tracker,
             user.to_string(),
             peer_addr.ip(),
-            tracks_ip,
+            true,
         ))
     }
 
@@ -1679,23 +1676,21 @@ 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;
-                }
-                Err(reason) => {
-                    stats.decrement_user_curr_connects(user);
-                    warn!(
-                        user = %user,
-                        ip = %peer_addr.ip(),
-                        reason = %reason,
-                        "IP limit exceeded"
-                    );
-                    return Err(ProxyError::ConnectionLimitExceeded {
-                        user: user.to_string(),
-                    });
-                }
+        match ip_tracker.check_and_add(user, peer_addr.ip()).await {
+            Ok(()) => {
+                ip_tracker.remove_ip(user, peer_addr.ip()).await;
+            }
+            Err(reason) => {
+                stats.decrement_user_curr_connects(user);
+                warn!(
+                    user = %user,
+                    ip = %peer_addr.ip(),
+                    reason = %reason,
+                    "IP limit exceeded"
+                );
+                return Err(ProxyError::ConnectionLimitExceeded {
+                    user: user.to_string(),
+                });
             }
         }
 
diff --git a/src/proxy/tests/client_security_tests.rs b/src/proxy/tests/client_security_tests.rs
index abd4213..2a60b57 100644
--- a/src/proxy/tests/client_security_tests.rs
+++ b/src/proxy/tests/client_security_tests.rs
@@ -960,6 +960,36 @@ async fn reservation_limit_failure_does_not_leak_curr_connects_counter() {
     assert_eq!(ip_tracker.get_active_ip_count(user).await, 0);
 }
 
+#[tokio::test]
+async fn unlimited_unique_ip_user_is_still_visible_in_active_ip_tracker() {
+    let user = "active-ip-observed-user";
+    let config = crate::config::ProxyConfig::default();
+    let stats = Arc::new(crate::stats::Stats::new());
+    let ip_tracker = Arc::new(crate::ip_tracker::UserIpTracker::new());
+    let peer = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(198, 51, 200, 17)), 50017);
+
+    let reservation = RunningClientHandler::acquire_user_connection_reservation_static(
+        user,
+        &config,
+        stats.clone(),
+        peer,
+        ip_tracker.clone(),
+    )
+    .await
+    .expect("reservation without unique-IP limit must succeed");
+
+    assert_eq!(stats.get_user_curr_connects(user), 1);
+    assert_eq!(
+        ip_tracker.get_active_ip_count(user).await,
+        1,
+        "active IP observability must not depend on unique-IP limit enforcement"
+    );
+
+    reservation.release().await;
+    assert_eq!(stats.get_user_curr_connects(user), 0);
+    assert_eq!(ip_tracker.get_active_ip_count(user).await, 0);
+}
+
 #[tokio::test]
 async fn short_tls_probe_is_masked_through_client_pipeline() {
     let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();