Runtime Ownership hardened

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

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),
                     )