Runtime Ownership hardened

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

src/stats/beobachten.rs

@@ -1,17 +1,21 @@
 //! Per-IP forensic buckets for scanner and handshake failure observation.
 
+use std::collections::hash_map::RandomState;
 use std::collections::{BTreeMap, HashMap};
+use std::hash::{BuildHasher, Hash, Hasher};
 use std::net::IpAddr;
 use std::time::{Duration, Instant};
 
 use parking_lot::Mutex;
 
 const CLEANUP_INTERVAL: Duration = Duration::from_secs(30);
-const MAX_BEOBACHTEN_ENTRIES: usize = 65_536;
+const BEOBACHTEN_SHARDS: usize = 64;
+const BEOBACHTEN_ENTRIES_PER_SHARD: usize = 1024;
 
 #[derive(Default)]
-struct BeobachtenInner {
-    entries: HashMap<(String, IpAddr), BeobachtenEntry>,
+struct BeobachtenShard {
+    classes: HashMap<String, HashMap<IpAddr, BeobachtenEntry>>,
+    entries: usize,
     last_cleanup: Option<Instant>,
 }
 
@@ -23,7 +27,8 @@ struct BeobachtenEntry {
 
 /// In-memory, TTL-scoped per-IP counters keyed by source class.
 pub struct BeobachtenStore {
-    inner: Mutex<BeobachtenInner>,
+    shards: Vec<Mutex<BeobachtenShard>>,
+    hash_builder: RandomState,
 }
 
 impl Default for BeobachtenStore {
@@ -35,7 +40,10 @@ impl Default for BeobachtenStore {
 impl BeobachtenStore {
     pub fn new() -> Self {
         Self {
-            inner: Mutex::new(BeobachtenInner::default()),
+            shards: (0..BEOBACHTEN_SHARDS)
+                .map(|_| Mutex::new(BeobachtenShard::default()))
+                .collect(),
+            hash_builder: RandomState::new(),
         }
     }
 
@@ -45,27 +53,31 @@ impl BeobachtenStore {
         }
 
         let now = Instant::now();
-        let mut guard = self.inner.lock();
-        Self::cleanup_if_needed(&mut guard, now, ttl);
+        let shard_index = self.shard_index(class, ip);
+        let mut shard = self.shards[shard_index].lock();
+        Self::cleanup_if_needed(&mut shard, now, ttl);
 
-        let key = (class.to_string(), ip);
-        if let Some(entry) = guard.entries.get_mut(&key) {
+        if let Some(entry) = shard
+            .classes
+            .get_mut(class)
+            .and_then(|entries| entries.get_mut(&ip))
+        {
             entry.tries = entry.tries.saturating_add(1);
             entry.last_seen = now;
             return;
         }
 
-        if guard.entries.len() >= MAX_BEOBACHTEN_ENTRIES {
+        if shard.entries >= BEOBACHTEN_ENTRIES_PER_SHARD {
             return;
         }
-
-        guard.entries.insert(
-            key,
+        shard.classes.entry(class.to_string()).or_default().insert(
+            ip,
             BeobachtenEntry {
                 tries: 1,
                 last_seen: now,
             },
         );
+        shard.entries = shard.entries.saturating_add(1);
     }
 
     pub fn snapshot_text(&self, ttl: Duration) -> String {
@@ -74,21 +86,15 @@ impl BeobachtenStore {
         }
 
         let now = Instant::now();
-        let entries = {
-            let mut guard = self.inner.lock();
-            Self::cleanup(&mut guard, now, ttl);
-            guard.last_cleanup = Some(now);
-
-            guard
-                .entries
-                .iter()
-                .map(|((class, ip), entry)| (class.clone(), *ip, entry.tries))
-                .collect::<Vec<_>>()
-        };
-
         let mut grouped = BTreeMap::<String, Vec<(IpAddr, u64)>>::new();
-        for (class, ip, tries) in entries {
-            grouped.entry(class).or_default().push((ip, tries));
+        for shard in &self.shards {
+            let mut shard = shard.lock();
+            Self::cleanup(&mut shard, now, ttl);
+            shard.last_cleanup = Some(now);
+            for (class, entries) in &shard.classes {
+                let output = grouped.entry(class.clone()).or_default();
+                output.extend(entries.iter().map(|(ip, entry)| (*ip, entry.tries)));
+            }
         }
 
         if grouped.is_empty() {
@@ -111,24 +117,61 @@ impl BeobachtenStore {
                 out.push_str(&format!("{ip}-{tries}\n"));
             }
         }
-
         out
     }
 
-    fn cleanup_if_needed(inner: &mut BeobachtenInner, now: Instant, ttl: Duration) {
-        let should_cleanup = match inner.last_cleanup {
+    fn shard_index(&self, class: &str, ip: IpAddr) -> usize {
+        let mut hasher = self.hash_builder.build_hasher();
+        class.hash(&mut hasher);
+        ip.hash(&mut hasher);
+        (hasher.finish() as usize) % BEOBACHTEN_SHARDS
+    }
+
+    fn cleanup_if_needed(shard: &mut BeobachtenShard, now: Instant, ttl: Duration) {
+        let should_cleanup = match shard.last_cleanup {
             Some(last) => now.saturating_duration_since(last) >= CLEANUP_INTERVAL,
             None => true,
         };
         if should_cleanup {
-            Self::cleanup(inner, now, ttl);
-            inner.last_cleanup = Some(now);
+            Self::cleanup(shard, now, ttl);
+            shard.last_cleanup = Some(now);
         }
     }
 
-    fn cleanup(inner: &mut BeobachtenInner, now: Instant, ttl: Duration) {
-        inner
-            .entries
-            .retain(|_, entry| now.saturating_duration_since(entry.last_seen) <= ttl);
+    fn cleanup(shard: &mut BeobachtenShard, now: Instant, ttl: Duration) {
+        for entries in shard.classes.values_mut() {
+            entries.retain(|_, entry| now.saturating_duration_since(entry.last_seen) <= ttl);
+        }
+        shard.classes.retain(|_, entries| !entries.is_empty());
+        shard.entries = shard.classes.values().map(HashMap::len).sum();
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn one_shard_never_exceeds_its_allocation_bound() {
+        let store = BeobachtenStore::new();
+        let ttl = Duration::from_secs(60);
+        let target_shard = 0;
+        let mut inserted = 0usize;
+        for suffix in 0..u32::MAX {
+            let ip = IpAddr::V4(std::net::Ipv4Addr::from(suffix));
+            if store.shard_index("scanner", ip) != target_shard {
+                continue;
+            }
+            store.record("scanner", ip, ttl);
+            inserted = inserted.saturating_add(1);
+            if inserted > BEOBACHTEN_ENTRIES_PER_SHARD + 64 {
+                break;
+            }
+        }
+
+        assert_eq!(
+            store.shards[target_shard].lock().entries,
+            BEOBACHTEN_ENTRIES_PER_SHARD
+        );
     }
 }

src/stats/me_counters.rs

@@ -77,11 +77,45 @@ impl Stats {
         if !self.telemetry_me_allows_normal() {
             return;
         }
-        let entry = self
-            .me_handshake_error_codes
-            .entry(code)
-            .or_insert_with(|| AtomicU64::new(0));
-        entry.fetch_add(1, Ordering::Relaxed);
+        if let Some(entry) = self.me_handshake_error_codes.get(&code) {
+            entry.fetch_add(1, Ordering::Relaxed);
+            return;
+        }
+
+        let mut slots = self
+            .me_handshake_error_code_slots
+            .load(Ordering::Acquire);
+        loop {
+            if slots >= ME_HANDSHAKE_ERROR_CODE_MAX {
+                if let Some(entry) = self.me_handshake_error_codes.get(&code) {
+                    entry.fetch_add(1, Ordering::Relaxed);
+                } else {
+                    self.me_handshake_error_code_overflow_total
+                        .fetch_add(1, Ordering::Relaxed);
+                }
+                return;
+            }
+            match self.me_handshake_error_code_slots.compare_exchange_weak(
+                slots,
+                slots + 1,
+                Ordering::AcqRel,
+                Ordering::Acquire,
+            ) {
+                Ok(_) => break,
+                Err(observed) => slots = observed,
+            }
+        }
+
+        match self.me_handshake_error_codes.entry(code) {
+            dashmap::mapref::entry::Entry::Occupied(entry) => {
+                self.me_handshake_error_code_slots
+                    .fetch_sub(1, Ordering::AcqRel);
+                entry.get().fetch_add(1, Ordering::Relaxed);
+            }
+            dashmap::mapref::entry::Entry::Vacant(entry) => {
+                entry.insert(AtomicU64::new(1));
+            }
+        }
     }
     pub fn increment_me_reader_eof_total(&self) {
         if self.telemetry_me_allows_normal() {
@@ -440,3 +474,44 @@ impl Stats {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use std::sync::Arc;
+
+    use super::*;
+
+    #[test]
+    fn handshake_error_code_cardinality_is_bounded_under_concurrency() {
+        const WORKERS: usize = 8;
+        const CODES_PER_WORKER: usize = 128;
+
+        let stats = Arc::new(Stats::new());
+        std::thread::scope(|scope| {
+            for worker in 0..WORKERS {
+                let stats = stats.clone();
+                scope.spawn(move || {
+                    for code in 0..CODES_PER_WORKER {
+                        stats.increment_me_handshake_error_code(
+                            (worker * CODES_PER_WORKER + code) as i32,
+                        );
+                    }
+                });
+            }
+        });
+
+        let counts = stats.get_me_handshake_error_code_counts();
+        assert!(counts.len() <= ME_HANDSHAKE_ERROR_CODE_MAX);
+        assert_eq!(
+            counts.iter().map(|(_, total)| *total).sum::<u64>()
+                + stats.get_me_handshake_error_code_overflow_total(),
+            (WORKERS * CODES_PER_WORKER) as u64
+        );
+        assert_eq!(
+            stats
+                .me_handshake_error_code_slots
+                .load(Ordering::Acquire),
+            counts.len()
+        );
+    }
+}

src/stats/me_getters.rs

@@ -10,6 +10,10 @@ impl Stats {
         out.sort_by_key(|(code, _)| *code);
         out
     }
+    pub fn get_me_handshake_error_code_overflow_total(&self) -> u64 {
+        self.me_handshake_error_code_overflow_total
+            .load(Ordering::Relaxed)
+    }
     pub fn get_me_route_drop_no_conn(&self) -> u64 {
         self.me_route_drop_no_conn.load(Ordering::Relaxed)
     }

src/stats/mod.rs

@@ -18,7 +18,7 @@ mod writer_counters;
 use dashmap::DashMap;
 use std::collections::HashMap;
 use std::sync::Arc;
-use std::sync::atomic::{AtomicBool, AtomicU8, AtomicU64, Ordering};
+use std::sync::atomic::{AtomicBool, AtomicU8, AtomicU64, AtomicUsize, Ordering};
 use std::time::Instant;
 
 pub(crate) use self::quota_store::QuotaStore;
@@ -28,6 +28,8 @@ use self::telemetry::TelemetryPolicy;
 pub use self::tls_fingerprints::TlsFingerprintSnapshotRow;
 use crate::config::MeWriterPickMode;
 
+const ME_HANDSHAKE_ERROR_CODE_MAX: usize = 64;
+
 #[derive(Clone, Copy)]
 enum RouteConnectionGauge {
     Direct,
@@ -219,6 +221,8 @@ pub struct Stats {
     me_floor_swap_idle_total: AtomicU64,
     me_floor_swap_idle_failed_total: AtomicU64,
     me_handshake_error_codes: DashMap<i32, AtomicU64>,
+    me_handshake_error_code_slots: AtomicUsize,
+    me_handshake_error_code_overflow_total: AtomicU64,
     me_route_drop_no_conn: AtomicU64,
     me_route_drop_channel_closed: AtomicU64,
     me_route_drop_queue_full: AtomicU64,

src/stats/quota_store.rs

@@ -56,6 +56,10 @@ impl QuotaStore {
         }
     }
 
+    pub(crate) fn remove(&self, user: &str) {
+        self.users.remove(user);
+    }
+
     pub(crate) fn snapshot(&self) -> HashMap<String, UserQuotaSnapshot> {
         let mut out = HashMap::new();
         for entry in self.users.iter() {