Runtime Ownership hardened

Co-Authored-By: brekotis <93345790+brekotis@users.noreply.github.com>

src/proxy/adaptive_buffers.rs

@@ -2,7 +2,7 @@
 
 use dashmap::DashMap;
 use std::cmp::max;
-use std::sync::OnceLock;
+use std::sync::{Mutex, OnceLock};
 use std::time::{Duration, Instant};
 
 const EMA_ALPHA: f64 = 0.2;
@@ -294,6 +294,11 @@ fn profiles() -> &'static DashMap<String, UserAdaptiveProfile> {
     USER_PROFILES.get_or_init(DashMap::new)
 }
 
+fn profile_insert_guard() -> &'static Mutex<()> {
+    static PROFILE_INSERT_GUARD: OnceLock<Mutex<()>> = OnceLock::new();
+    PROFILE_INSERT_GUARD.get_or_init(|| Mutex::new(()))
+}
+
 /// Returns a fresh user's recent successful Direct tier, or `Base` when stale.
 #[allow(dead_code)]
 pub fn seed_tier_for_user(user: &str) -> AdaptiveTier {
@@ -320,29 +325,58 @@ pub fn record_user_tier(user: &str, tier: AdaptiveTier) {
     if user.len() > MAX_USER_KEY_BYTES {
         return;
     }
+    record_user_tier_with_cap(
+        profiles(),
+        profile_insert_guard(),
+        user,
+        tier,
+        MAX_USER_PROFILES_ENTRIES,
+    );
+}
+
+fn record_user_tier_with_cap(
+    profiles: &DashMap<String, UserAdaptiveProfile>,
+    insert_guard: &Mutex<()>,
+    user: &str,
+    tier: AdaptiveTier,
+    max_entries: usize,
+) {
     let now = Instant::now();
-    let mut was_vacant = false;
-    match profiles().entry(user.to_string()) {
-        dashmap::mapref::entry::Entry::Occupied(mut entry) => {
-            let existing = *entry.get();
-            let effective = if now.saturating_duration_since(existing.seen_at) > PROFILE_TTL {
-                tier
-            } else {
-                max(existing.tier, tier)
-            };
-            entry.insert(UserAdaptiveProfile {
-                tier: effective,
-                seen_at: now,
-            });
-        }
-        dashmap::mapref::entry::Entry::Vacant(slot) => {
-            slot.insert(UserAdaptiveProfile { tier, seen_at: now });
-            was_vacant = true;
-        }
+    if let Some(mut entry) = profiles.get_mut(user) {
+        let effective = if now.saturating_duration_since(entry.seen_at) > PROFILE_TTL {
+            tier
+        } else {
+            max(entry.tier, tier)
+        };
+        *entry = UserAdaptiveProfile {
+            tier: effective,
+            seen_at: now,
+        };
+        return;
     }
-    if was_vacant && profiles().len() > MAX_USER_PROFILES_ENTRIES {
-        profiles().retain(|_, v| now.saturating_duration_since(v.seen_at) <= PROFILE_TTL);
+
+    let _guard = insert_guard
+        .lock()
+        .unwrap_or_else(|poisoned| poisoned.into_inner());
+    if let Some(mut entry) = profiles.get_mut(user) {
+        let effective = if now.saturating_duration_since(entry.seen_at) > PROFILE_TTL {
+            tier
+        } else {
+            max(entry.tier, tier)
+        };
+        *entry = UserAdaptiveProfile {
+            tier: effective,
+            seen_at: now,
+        };
+        return;
     }
+    if profiles.len() >= max_entries {
+        profiles.retain(|_, value| now.saturating_duration_since(value.seen_at) <= PROFILE_TTL);
+    }
+    if profiles.len() >= max_entries {
+        return;
+    }
+    profiles.insert(user.to_string(), UserAdaptiveProfile { tier, seen_at: now });
 }
 
 #[cfg(test)]
@@ -438,6 +472,17 @@ mod adaptive_direct_budget_policy_tests;
 mod tests {
     use super::*;
 
+    #[test]
+    fn fresh_profile_cardinality_is_hard_bounded() {
+        let profiles = DashMap::new();
+        let insert_guard = Mutex::new(());
+        for index in 0..64 {
+            let user = format!("user-{index}");
+            record_user_tier_with_cap(&profiles, &insert_guard, &user, AdaptiveTier::Base, 16);
+        }
+        assert_eq!(profiles.len(), 16);
+    }
+
     fn sample(
         c2s_bytes: u64,
         s2c_requested_bytes: u64,

src/proxy/client.rs

@@ -50,7 +50,6 @@ use crate::proxy::handshake::{
 };
 #[cfg(test)]
 use crate::proxy::handshake::{handle_mtproto_handshake, handle_tls_handshake};
-use crate::proxy::masking::handle_bad_client_with_shared;
 #[cfg(test)]
 use crate::proxy::route_mode::RelayRouteMode;
 use crate::proxy::route_mode::RouteRuntimeController;
@@ -247,6 +246,7 @@ fn masking_outcome<R, W>(
     peer: SocketAddr,
     local_addr: SocketAddr,
     config: Arc<ProxyConfig>,
+    upstream_manager: Arc<UpstreamManager>,
     beobachten: Arc<BeobachtenStore>,
     shared: Arc<ProxySharedState>,
 ) -> HandshakeOutcome
@@ -264,7 +264,7 @@ where
         )
         .await;
 
-        handle_bad_client_with_shared(
+        crate::proxy::masking::handle_bad_client_with_shared_resolver(
             reader,
             writer,
             &initial_data,
@@ -273,6 +273,7 @@ where
             &config,
             &beobachten,
             shared.as_ref(),
+            Some(upstream_manager.as_ref()),
         )
         .await;
         Ok(())
@@ -657,6 +658,7 @@ where
                     real_peer,
                     local_addr,
                     config.clone(),
+                    upstream_manager.clone(),
                     beobachten.clone(),
                     shared.clone(),
                 ));
@@ -679,6 +681,7 @@ where
                         real_peer,
                         local_addr,
                         config.clone(),
+                        upstream_manager.clone(),
                         beobachten.clone(),
                         shared.clone(),
                     ));
@@ -698,6 +701,7 @@ where
                     real_peer,
                     local_addr,
                     config.clone(),
+                    upstream_manager.clone(),
                     beobachten.clone(),
                     shared.clone(),
                 ));
@@ -729,6 +733,7 @@ where
                         real_peer,
                         local_addr,
                         config.clone(),
+                        upstream_manager.clone(),
                         beobachten.clone(),
                         shared.clone(),
                     ));
@@ -787,6 +792,7 @@ where
                         real_peer,
                         local_addr,
                         config.clone(),
+                        upstream_manager.clone(),
                         beobachten.clone(),
                         shared.clone(),
                     ));
@@ -817,6 +823,7 @@ where
                     real_peer,
                     local_addr,
                     config.clone(),
+                    upstream_manager.clone(),
                     beobachten.clone(),
                     shared.clone(),
                 ));
@@ -843,6 +850,7 @@ where
                         real_peer,
                         local_addr,
                         config.clone(),
+                        upstream_manager.clone(),
                         beobachten.clone(),
                         shared.clone(),
                     ));
@@ -1279,6 +1287,7 @@ impl RunningClientHandler {
                 peer,
                 local_addr,
                 self.config.clone(),
+                self.upstream_manager.clone(),
                 self.beobachten.clone(),
                 self.shared.clone(),
             ));
@@ -1301,6 +1310,7 @@ impl RunningClientHandler {
                     peer,
                     local_addr,
                     self.config.clone(),
+                    self.upstream_manager.clone(),
                     self.beobachten.clone(),
                     self.shared.clone(),
                 ));
@@ -1321,6 +1331,7 @@ impl RunningClientHandler {
                 peer,
                 local_addr,
                 self.config.clone(),
+                self.upstream_manager.clone(),
                 self.beobachten.clone(),
                 self.shared.clone(),
             ));
@@ -1377,6 +1388,7 @@ impl RunningClientHandler {
                         peer,
                         local_addr,
                         config.clone(),
+                        self.upstream_manager.clone(),
                         self.beobachten.clone(),
                         self.shared.clone(),
                     ));
@@ -1445,6 +1457,7 @@ impl RunningClientHandler {
                     peer,
                     local_addr,
                     config.clone(),
+                    self.upstream_manager.clone(),
                     self.beobachten.clone(),
                     self.shared.clone(),
                 ));
@@ -1493,6 +1506,7 @@ impl RunningClientHandler {
                 peer,
                 local_addr,
                 self.config.clone(),
+                self.upstream_manager.clone(),
                 self.beobachten.clone(),
                 self.shared.clone(),
             ));
@@ -1532,6 +1546,7 @@ impl RunningClientHandler {
                     peer,
                     local_addr,
                     config.clone(),
+                    self.upstream_manager.clone(),
                     self.beobachten.clone(),
                     self.shared.clone(),
                 ));

src/proxy/handshake/auth_probe.rs

@@ -98,7 +98,7 @@ pub(super) fn auth_probe_is_throttled_in(
     };
     if auth_probe_state_expired(&entry, now) {
         drop(entry);
-        state.remove(&peer_ip);
+        state.remove_if(&peer_ip, |_, current| auth_probe_state_expired(current, now));
         return false;
     }
     now < entry.blocked_until
@@ -116,7 +116,7 @@ pub(super) fn auth_probe_saturation_grace_exhausted_in(
     };
     if auth_probe_state_expired(&entry, now) {
         drop(entry);
-        state.remove(&peer_ip);
+        state.remove_if(&peer_ip, |_, current| auth_probe_state_expired(current, now));
         return false;
     }
 
@@ -264,11 +264,19 @@ pub(super) fn auth_probe_record_failure_with_state_in(
                     }
                 }
 
-                let Some((evict_key, _, _)) = eviction_candidate else {
+                let Some((evict_key, evict_fail_streak, evict_last_seen)) = eviction_candidate else {
                     return;
                 };
-                state.remove(&evict_key);
-                break;
+                if state
+                    .remove_if(&evict_key, |_, current| {
+                        current.fail_streak == evict_fail_streak
+                            && current.last_seen == evict_last_seen
+                    })
+                    .is_some()
+                {
+                    break;
+                }
+                continue;
             }
 
             let mut stale_keys = Vec::new();
@@ -334,18 +342,22 @@ pub(super) fn auth_probe_record_failure_with_state_in(
             }
 
             for stale_key in stale_keys {
-                state.remove(&stale_key);
+                state.remove_if(&stale_key, |_, current| {
+                    auth_probe_state_expired(current, now)
+                });
             }
 
             if state.len() < AUTH_PROBE_TRACK_MAX_ENTRIES {
                 break;
             }
 
-            let Some((evict_key, _, _)) = eviction_candidate else {
+            let Some((evict_key, evict_fail_streak, evict_last_seen)) = eviction_candidate else {
                 auth_probe_note_saturation_in(shared, now);
                 return;
             };
-            state.remove(&evict_key);
+            state.remove_if(&evict_key, |_, current| {
+                current.fail_streak == evict_fail_streak && current.last_seen == evict_last_seen
+            });
             auth_probe_note_saturation_in(shared, now);
         }
     }

src/proxy/handshake/tls_handshake.rs

@@ -266,11 +266,11 @@ where
         return HandshakeResult::BadClient { reader, writer };
     }
 
+    let selected_tls_domain =
+        matched_tls_domain.unwrap_or(config.censorship.tls_domain.as_str());
     let cached_entry = if config.censorship.tls_emulation {
         if let Some(cache) = tls_cache.as_ref() {
-            let selected_domain =
-                matched_tls_domain.unwrap_or(config.censorship.tls_domain.as_str());
-            let cached_entry = cache.get(selected_domain).await;
+            let cached_entry = cache.get(selected_tls_domain).await;
             Some(cached_entry)
         } else {
             None
@@ -322,6 +322,7 @@ where
             if let Some(cache) = tls_cache.as_ref() {
                 cache
                     .take_full_cert_budget_for_ip(
+                        selected_tls_domain,
                         peer.ip(),
                         Duration::from_secs(config.censorship.tls_full_cert_ttl_secs),
                     )

src/proxy/masking.rs

@@ -1,7 +1,6 @@
 //! Masking - forward unrecognized traffic to mask host
 
 use crate::config::ProxyConfig;
-use crate::network::dns_overrides::resolve_socket_addr;
 use crate::protocol::tls;
 use crate::proxy::shared_state::ProxySharedState;
 use crate::stats::beobachten::BeobachtenStore;
@@ -41,6 +40,7 @@ const MASK_RELAY_IDLE_TIMEOUT: Duration = Duration::from_millis(100);
 const MASK_BUFFER_SIZE: usize = 8192;
 const MASK_BUFFER_GROW_AFTER_BYTES: usize = 256 * 1024;
 const MASK_BUFFER_MAX_SIZE: usize = 64 * 1024;
+const MASK_DNS_RESULT_MAX_ADDRESSES: usize = 64;
 #[cfg(unix)]
 #[cfg(not(test))]
 const LOCAL_INTERFACE_CACHE_TTL: Duration = Duration::from_secs(300);
@@ -532,19 +532,25 @@ fn parse_mask_host_ip_literal(host: &str) -> Option<IpAddr> {
 async fn resolve_mask_target_addrs(
     mask_host: &str,
     mask_port: u16,
+    upstream_manager: Option<&crate::transport::UpstreamManager>,
 ) -> std::io::Result<Vec<SocketAddr>> {
-    if let Some(addr) = resolve_socket_addr(mask_host, mask_port) {
-        return Ok(vec![addr]);
-    }
-
     if let Some(ip) = parse_mask_host_ip_literal(mask_host) {
         return Ok(vec![SocketAddr::new(ip, mask_port)]);
     }
 
+    if let Some(upstream_manager) = upstream_manager {
+        return upstream_manager
+            .resolve_all(mask_host, mask_port)
+            .await
+            .map_err(|error| IoError::new(ErrorKind::NotFound, error.to_string()));
+    }
+
     let addrs = timeout(MASK_TIMEOUT, lookup_host((mask_host, mask_port)))
         .await
         .map_err(|_| IoError::new(ErrorKind::TimedOut, "mask target DNS lookup timed out"))??;
-    let addrs = addrs.collect::<Vec<_>>();
+    let addrs = addrs
+        .take(MASK_DNS_RESULT_MAX_ADDRESSES)
+        .collect::<Vec<_>>();
     if addrs.is_empty() {
         return Err(IoError::new(
             ErrorKind::NotFound,
@@ -999,6 +1005,34 @@ pub(crate) async fn handle_bad_client_with_shared<R, W>(
 ) where
     R: AsyncRead + Unpin + Send + 'static,
     W: AsyncWrite + Unpin + Send + 'static,
+{
+    handle_bad_client_with_shared_resolver(
+        reader,
+        writer,
+        initial_data,
+        peer,
+        local_addr,
+        config,
+        beobachten,
+        shared,
+        None,
+    )
+    .await;
+}
+
+pub(super) async fn handle_bad_client_with_shared_resolver<R, W>(
+    reader: R,
+    writer: W,
+    initial_data: &[u8],
+    peer: SocketAddr,
+    local_addr: SocketAddr,
+    config: &ProxyConfig,
+    beobachten: &BeobachtenStore,
+    shared: &ProxySharedState,
+    upstream_manager: Option<&crate::transport::UpstreamManager>,
+) where
+    R: AsyncRead + Unpin + Send + 'static,
+    W: AsyncWrite + Unpin + Send + 'static,
 {
     let client_type = detect_client_type(initial_data);
     if config.general.beobachten {
@@ -1112,7 +1146,8 @@ pub(crate) async fn handle_bad_client_with_shared<R, W>(
     let mask_host = mask_target.host;
     let mask_port = mask_target.port;
 
-    let resolved_mask_addrs = match resolve_mask_target_addrs(mask_host, mask_port).await {
+    let resolved_mask_addrs =
+        match resolve_mask_target_addrs(mask_host, mask_port, upstream_manager).await {
         Ok(addrs) => addrs,
         Err(e) => {
             let outcome_started = Instant::now();

src/proxy/middle_relay.rs

@@ -34,7 +34,7 @@ use crate::stats::{
     MeD2cFlushReason, MeD2cQuotaRejectStage, MeD2cWriteMode, QuotaReserveError, Stats, UserStats,
 };
 use crate::stream::{BufferPool, CryptoReader, CryptoWriter, PooledBuffer};
-use crate::transport::middle_proxy::{MePool, MeResponse, proto_flags_for_tag};
+use crate::transport::middle_proxy::{ConnLease, MePool, MeResponse, proto_flags_for_tag};
 
 mod c2me;
 mod d2c;

src/proxy/middle_relay/session.rs

@@ -1,5 +1,42 @@
 use super::*;
 
+struct RelayConnLease {
+    connection: Option<ConnLease>,
+    conn_id: u64,
+    shared: Arc<ProxySharedState>,
+}
+
+impl RelayConnLease {
+    fn new(connection: ConnLease, shared: Arc<ProxySharedState>) -> Self {
+        let conn_id = connection.conn_id();
+        Self {
+            connection: Some(connection),
+            conn_id,
+            shared,
+        }
+    }
+
+    fn conn_id(&self) -> u64 {
+        self.conn_id
+    }
+
+    async fn unregister(mut self) {
+        let Some(connection) = self.connection.take() else {
+            return;
+        };
+        clear_relay_idle_candidate_in(self.shared.as_ref(), connection.conn_id());
+        connection.unregister().await;
+    }
+}
+
+impl Drop for RelayConnLease {
+    fn drop(&mut self) {
+        if let Some(connection) = self.connection.as_ref() {
+            clear_relay_idle_candidate_in(self.shared.as_ref(), connection.conn_id());
+        }
+    }
+}
+
 /// Runs Middle-End relay with explicit kernel-conntrack close publication policy.
 pub(crate) async fn handle_via_middle_proxy_with_conntrack<R, W>(
     mut crypto_reader: CryptoReader<R>,
@@ -44,7 +81,11 @@ where
         "Routing via Middle-End"
     );
 
-    let (conn_id, me_rx) = me_pool.registry().register().await;
+    let Some((connection, me_rx)) = me_pool.register_connection().await else {
+        return Err(ProxyError::MiddleConnectionLost);
+    };
+    let relay_connection = RelayConnLease::new(connection, Arc::clone(&shared));
+    let conn_id = relay_connection.conn_id();
     let trace_id = session_id;
     let bytes_me2c = Arc::new(AtomicU64::new(0));
     let mut forensics = RelayForensicsState {
@@ -76,7 +117,7 @@ where
         let _cutover_park_lease = stats.acquire_middle_cutover_park_lease();
         tokio::time::sleep(delay).await;
         let _ = me_pool.send_close(conn_id).await;
-        me_pool.registry().unregister(conn_id).await;
+        relay_connection.unregister().await;
         return Err(ProxyError::RouteSwitched);
     }
 
@@ -825,8 +866,7 @@ where
         }
     }
 
-    clear_relay_idle_candidate_in(shared.as_ref(), conn_id);
-    me_pool.registry().unregister(conn_id).await;
+    relay_connection.unregister().await;
     let pool_snapshot = buffer_pool.stats();
     stats.set_buffer_pool_gauges(
         pool_snapshot.pooled,

src/proxy/tests/masking_adversarial_tests.rs

@@ -229,7 +229,7 @@ async fn masking_fallback_down_mimics_timeout() {
 
 #[tokio::test]
 async fn masking_ssrf_resolve_internal_ranges_blocked() {
-    use crate::network::dns_overrides::resolve_socket_addr;
+    use crate::network::dns_overrides::DnsOverrides;
 
     let blocked_ips = [
         "127.0.0.1",
@@ -238,10 +238,11 @@ async fn masking_ssrf_resolve_internal_ranges_blocked() {
         "192.168.1.1",
         "0.0.0.0",
     ];
+    let resolver = DnsOverrides::default();
 
     for ip in blocked_ips {
         assert!(
-            resolve_socket_addr(ip, 80).is_none(),
+            resolver.resolve_socket_addr(ip, 80).is_none(),
             "runtime DNS overrides must not resolve unconfigured literal host targets"
         );
     }

src/proxy/tests/masking_connect_failure_close_matrix_security_tests.rs

@@ -1,5 +1,5 @@
 use super::*;
-use crate::network::dns_overrides::install_entries;
+use std::sync::Arc;
 use tokio::io::{AsyncReadExt, AsyncWriteExt, duplex};
 use tokio::time::{Duration, Instant, timeout};
 
@@ -8,6 +8,7 @@ async fn run_connect_failure_case(
     port: u16,
     timing_normalization_enabled: bool,
     peer: SocketAddr,
+    dns_overrides: Vec<String>,
 ) -> Duration {
     let mut config = ProxyConfig::default();
     config.general.beobachten = false;
@@ -21,6 +22,19 @@ async fn run_connect_failure_case(
 
     let local_addr: SocketAddr = "127.0.0.1:443".parse().unwrap();
     let beobachten = BeobachtenStore::new();
+    let upstream_manager = crate::transport::UpstreamManager::new(
+        Vec::new(),
+        1,
+        1,
+        100,
+        1,
+        1,
+        true,
+        Arc::new(crate::stats::Stats::new()),
+    )
+    .with_dns_overrides(&dns_overrides)
+    .unwrap();
+    let shared = ProxySharedState::new();
     let probe = b"CONNECT example.org:443 HTTP/1.1\r\nHost: example.org\r\n\r\n";
 
     let (mut client_writer, client_reader) = duplex(1024);
@@ -28,7 +42,7 @@ async fn run_connect_failure_case(
 
     let started = Instant::now();
     let task = tokio::spawn(async move {
-        handle_bad_client(
+        handle_bad_client_with_shared_resolver(
             client_reader,
             client_visible_writer,
             probe,
@@ -36,6 +50,8 @@ async fn run_connect_failure_case(
             local_addr,
             &config,
             &beobachten,
+            shared.as_ref(),
+            Some(&upstream_manager),
         )
         .await;
     });
@@ -71,8 +87,14 @@ async fn connect_failure_refusal_close_behavior_matrix() {
             .parse()
             .unwrap();
         let elapsed =
-            run_connect_failure_case("127.0.0.1", unused_port, timing_normalization_enabled, peer)
-                .await;
+            run_connect_failure_case(
+                "127.0.0.1",
+                unused_port,
+                timing_normalization_enabled,
+                peer,
+                Vec::new(),
+            )
+            .await;
 
         if timing_normalization_enabled {
             assert!(
@@ -94,9 +116,6 @@ async fn connect_failure_overridden_hostname_close_behavior_matrix() {
     let unused_port = temp_listener.local_addr().unwrap().port();
     drop(temp_listener);
 
-    // Make hostname resolution deterministic in tests so timing ceilings are meaningful.
-    install_entries(&[format!("mask.invalid:{}:127.0.0.1", unused_port)]).unwrap();
-
     for (idx, timing_normalization_enabled) in [false, true].into_iter().enumerate() {
         let peer: SocketAddr = format!("203.0.113.220:{}", 54200 + idx as u16)
             .parse()
@@ -106,6 +125,7 @@ async fn connect_failure_overridden_hostname_close_behavior_matrix() {
             unused_port,
             timing_normalization_enabled,
             peer,
+            vec![format!("mask.invalid:{}:127.0.0.1", unused_port)],
         )
         .await;
 
@@ -121,6 +141,4 @@ async fn connect_failure_overridden_hostname_close_behavior_matrix() {
             );
         }
     }
-
-    install_entries(&[]).unwrap();
 }

src/proxy/traffic_limiter.rs

@@ -355,9 +355,10 @@ impl CidrBucket {
     }
 
     fn acquire_user_share(&self, user: &str) -> Arc<CidrUserShare> {
-        let share = self.users.get_or_insert_with(user, CidrUserShare::new);
-        share.active_conns.fetch_add(1, Ordering::Relaxed);
-        share
+        self.users
+            .get_or_insert_with(user, CidrUserShare::new, |share| {
+                share.active_conns.fetch_add(1, Ordering::Relaxed);
+            })
     }
 
     fn release_user_share(&self, user: &str, share: &Arc<CidrUserShare>) {
@@ -487,15 +488,20 @@ impl<T> ShardedRegistry<T> {
         (hasher.finish() as usize) & self.mask
     }
 
-    fn get_or_insert_with<F>(&self, key: &str, make: F) -> Arc<T>
+    fn get_or_insert_with<F, A>(&self, key: &str, make: F, activate: A) -> Arc<T>
     where
         F: FnOnce() -> T,
+        A: FnOnce(&Arc<T>),
     {
         let shard = &self.shards[self.shard_index(key)];
         match shard.entry(key.to_string()) {
-            dashmap::mapref::entry::Entry::Occupied(entry) => Arc::clone(entry.get()),
+            dashmap::mapref::entry::Entry::Occupied(entry) => {
+                activate(entry.get());
+                Arc::clone(entry.get())
+            }
             dashmap::mapref::entry::Entry::Vacant(slot) => {
                 let value = Arc::new(make());
+                activate(&value);
                 slot.insert(Arc::clone(&value));
                 value
             }
@@ -516,14 +522,7 @@ impl<T> ShardedRegistry<T> {
         F: Fn(&Arc<T>) -> bool,
     {
         let shard = &self.shards[self.shard_index(key)];
-        let should_remove = match shard.get(key) {
-            Some(entry) => predicate(entry.value()),
-            None => false,
-        };
-        if !should_remove {
-            return false;
-        }
-        shard.remove(key).is_some()
+        shard.remove_if(key, |_, value| predicate(value)).is_some()
     }
 }
 
@@ -743,9 +742,10 @@ impl TrafficLimiter {
         if let Some(limit) = policy.user_limits.get(user).copied() {
             let bucket = self
                 .user_buckets
-                .get_or_insert_with(user, || UserBucket::new(limit));
+                .get_or_insert_with(user, || UserBucket::new(limit), |bucket| {
+                    bucket.active_leases.fetch_add(1, Ordering::Relaxed);
+                });
             bucket.set_rates(limit);
-            bucket.active_leases.fetch_add(1, Ordering::Relaxed);
             self.user_scope
                 .active_leases
                 .fetch_add(1, Ordering::Relaxed);
@@ -762,9 +762,10 @@ impl TrafficLimiter {
             };
             let bucket = self
                 .cidr_buckets
-                .get_or_insert_with(key, || CidrBucket::new(limits));
+                .get_or_insert_with(key, || CidrBucket::new(limits), |bucket| {
+                    bucket.active_leases.fetch_add(1, Ordering::Relaxed);
+                });
             bucket.set_rates(limits);
-            bucket.active_leases.fetch_add(1, Ordering::Relaxed);
             self.cidr_scope
                 .active_leases
                 .fetch_add(1, Ordering::Relaxed);