Merge pull request #747 from telemt/flow

Restore active IP observability for users without unique-IP limits

Cargo.lock

@@ -2791,7 +2791,7 @@ checksum = "7b2093cf4c8eb1e67749a6762251bc9cd836b6fc171623bd0a9d324d37af2417"
 
 [[package]]
 name = "telemt"
-version = "3.4.7"
+version = "3.4.8"
 dependencies = [
  "aes",
  "anyhow",

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "telemt"
-version = "3.4.7"
+version = "3.4.8"
 edition = "2024"
 
 [features]

src/proxy/client.rs

@@ -1612,22 +1612,19 @@ impl RunningClientHandler {
             });
         }
 
-        let tracks_ip = ip_tracker.get_user_limit(user).await.is_some();
+        match ip_tracker.check_and_add(user, peer_addr.ip()).await {
-        if tracks_ip {
+            Ok(()) => {}
-            match ip_tracker.check_and_add(user, peer_addr.ip()).await {
+            Err(reason) => {
-                Ok(()) => {}
+                stats.decrement_user_curr_connects(user);
-                Err(reason) => {
+                warn!(
-                    stats.decrement_user_curr_connects(user);
+                    user = %user,
-                    warn!(
+                    ip = %peer_addr.ip(),
-                        user = %user,
+                    reason = %reason,
-                        ip = %peer_addr.ip(),
+                    "IP limit exceeded"
-                        reason = %reason,
+                );
-                        "IP limit exceeded"
+                return Err(ProxyError::ConnectionLimitExceeded {
-                    );
+                    user: user.to_string(),
-                    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 {
-            match ip_tracker.check_and_add(user, peer_addr.ip()).await {
+            Ok(()) => {
-                Ok(()) => {
+                ip_tracker.remove_ip(user, peer_addr.ip()).await;
-                    ip_tracker.remove_ip(user, peer_addr.ip()).await;
+            }
-                }
+            Err(reason) => {
-                Err(reason) => {
+                stats.decrement_user_curr_connects(user);
-                    stats.decrement_user_curr_connects(user);
+                warn!(
-                    warn!(
+                    user = %user,
-                        user = %user,
+                    ip = %peer_addr.ip(),
-                        ip = %peer_addr.ip(),
+                    reason = %reason,
-                        reason = %reason,
+                    "IP limit exceeded"
-                        "IP limit exceeded"
+                );
-                    );
+                return Err(ProxyError::ConnectionLimitExceeded {
-                    return Err(ProxyError::ConnectionLimitExceeded {
+                    user: user.to_string(),
-                        user: user.to_string(),
+                });
-                    });
-                }
             }
         }
 

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();