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.6"
+version = "3.4.8"
 dependencies = [
  "aes",
  "anyhow",

Cargo.toml

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

src/proxy/client.rs

@@ -1612,8 +1612,6 @@ 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) => {
@@ -1629,14 +1627,13 @@ impl RunningClientHandler {
                 });
             }
         }
-        }
 
         Ok(UserConnectionReservation::new(
             stats,
             ip_tracker,
             user.to_string(),
             peer_addr.ip(),
-            tracks_ip,
+            true,
         ))
     }
 
@@ -1679,7 +1676,6 @@ 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;
@@ -1697,7 +1693,6 @@ impl RunningClientHandler {
                 });
             }
         }
-        }
 
         stats.decrement_user_curr_connects(user);
         Ok(())

src/proxy/tests/client_security_tests.rs

@@ -281,13 +281,8 @@ async fn user_connection_reservation_drop_enqueues_cleanup_synchronously() {
     assert_eq!(ip_tracker.get_active_ip_count(&user).await, 1);
     assert_eq!(stats.get_user_curr_connects(&user), 1);
 
-    let reservation = UserConnectionReservation::new(
+    let reservation =
-        stats.clone(),
+        UserConnectionReservation::new(stats.clone(), ip_tracker.clone(), user.clone(), ip, true);
-        ip_tracker.clone(),
-        user.clone(),
-        ip,
-        true,
-    );
 
     // Drop the reservation synchronously without any tokio::spawn/await yielding!
     drop(reservation);
@@ -965,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();

src/transport/middle_proxy/registry.rs

@@ -8,8 +8,8 @@ use dashmap::DashMap;
 use tokio::sync::mpsc::error::TrySendError;
 use tokio::sync::{Mutex, Semaphore, mpsc};
 
-use super::{MeResponse, RouteBytePermit};
 use super::codec::WriterCommand;
+use super::{MeResponse, RouteBytePermit};
 
 const ROUTE_BACKPRESSURE_BASE_TIMEOUT_MS: u64 = 25;
 const ROUTE_BACKPRESSURE_HIGH_TIMEOUT_MS: u64 = 120;
@@ -129,7 +129,10 @@ impl ConnRegistry {
         )
     }
 
-    fn with_route_limits(route_channel_capacity: usize, route_byte_permits_per_conn: usize) -> Self {
+    fn with_route_limits(
+        route_channel_capacity: usize,
+        route_byte_permits_per_conn: usize,
+    ) -> Self {
         let start = rand::random::<u64>() | 1;
         let route_channel_capacity = route_channel_capacity.max(1);
         Self {
@@ -209,9 +212,10 @@ impl ConnRegistry {
         let id = self.next_id.fetch_add(1, Ordering::Relaxed);
         let (tx, rx) = mpsc::channel(self.route_channel_capacity);
         self.routing.map.insert(id, tx);
-        self.routing
+        self.routing.byte_budget.insert(
-            .byte_budget
+            id,
-            .insert(id, Arc::new(Semaphore::new(self.route_byte_permits_per_conn)));
+            Arc::new(Semaphore::new(self.route_byte_permits_per_conn)),
+        );
         (id, rx)
     }
 
@@ -287,8 +291,7 @@ impl ConnRegistry {
                 .map_err(|_| RouteResult::QueueFullHigh)?,
             Some(timeout_ms) => {
                 let acquire = semaphore.acquire_many_owned(permits);
-                match tokio::time::timeout(Duration::from_millis(timeout_ms.max(1)), acquire)
+                match tokio::time::timeout(Duration::from_millis(timeout_ms.max(1)), acquire).await
-                    .await
                 {
                     Ok(Ok(permit)) => permit,
                     Ok(Err(_)) => return Err(RouteResult::ChannelClosed),