rwhod: put max cap on status update storage

Without a max cap, someone could potentially flood our memory by posting
tons and tons of spoofed status updates.

nix/module.nix

@@ -69,6 +69,16 @@ in {
               '';
             };
 
+            max_status_entries = lib.mkOption {
+              type = lib.types.ints.positive;
+              default = 4096;
+              description = ''
+                Maximum number of distinct hosts to keep rwhod status records for
+                at the same time. Once at capacity, the least recently updated
+                record is evicted to make room for a newly-seen host.
+              '';
+            };
+
             ignoreUsers = lib.mkOption {
               type = with lib.types; listOf (either str ints.unsigned);
               default = [ ];

src/bin/roowhod.rs

@@ -15,10 +15,15 @@ use tracing_subscriber::layer::SubscriberExt;
 use roowho2_lib::server::{
     config::{DEFAULT_CONFIG_PATH, LogLevel},
     ignore_list::IgnoreList,
-    rwhod::{RwhodStatusStore, rwhod_packet_receiver_task, rwhod_packet_sender_task},
+    rwhod::{
+        RwhodStatusRegistry, RwhodStatusStore, rwhod_packet_receiver_task, rwhod_packet_sender_task,
+    },
     varlink_api::varlink_client_server_task,
 };
 
+/// Default maximum number of distinct hosts to keep rwhod status records for at the same time.
+const DEFAULT_MAX_STATUS_ENTRIES: usize = 4096;
+
 #[derive(Parser)]
 #[command(
     author = "Programvareverkstedet <projects@pvv.ntnu.no>",
@@ -75,7 +80,19 @@ async fn main() -> anyhow::Result<()> {
 
     let mut join_set = tokio::task::JoinSet::new();
 
-    let whod_status_store = Arc::new(RwLock::new(HashMap::new()));
+    let max_status_entries = match config.rwhod.max_status_entries {
+        Some(0) => {
+            tracing::warn!(
+                "rwhod.max_status_entries is set to 0, remapping to default value of {}",
+                DEFAULT_MAX_STATUS_ENTRIES
+            );
+            DEFAULT_MAX_STATUS_ENTRIES
+        }
+        Some(max_status_entries) => max_status_entries,
+        None => DEFAULT_MAX_STATUS_ENTRIES,
+    };
+    let whod_status_store: RwhodStatusStore =
+        Arc::new(RwLock::new(RwhodStatusRegistry::new(max_status_entries)));
 
     let client_server_token = CancellationToken::new();
     let client_server_token_ = client_server_token.clone();

src/server/config.rs

@@ -48,6 +48,13 @@ pub struct RwhodConfig {
     ///
     /// If left as `None`, defaults to 60 seconds.
     pub send_interval_seconds: Option<u64>,
+
+    /// Maximum number of distinct hosts to keep rwhod status records for at
+    /// the same time. Once at capacity, the least recently updated record
+    /// is evicted to make room for a newly-seen host.
+    ///
+    /// If left as `None`, defaults to 4096.
+    pub max_status_entries: Option<usize>,
 }
 
 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]

src/server/rwhod.rs

@@ -1,14 +1,14 @@
 mod packet_receiver;
 mod packet_sender;
 mod rwhod_status;
+mod status_registry;
 
 pub use packet_receiver::rwhod_packet_receiver_task;
 pub use packet_sender::{determine_relevant_interfaces, rwhod_packet_sender_task};
 pub use rwhod_status::{generate_rwhod_status_update, generate_rwhod_user_entries};
+pub use status_registry::RwhodStatusRegistry;
 
-use std::{collections::HashMap, sync::Arc};
+use std::sync::Arc;
 use tokio::sync::RwLock;
 
-use crate::proto::WhodStatusUpdate;
-
-pub type RwhodStatusStore = Arc<RwLock<HashMap<String, WhodStatusUpdate>>>;
+pub type RwhodStatusStore = Arc<RwLock<RwhodStatusRegistry>>;

src/server/rwhod/packet_receiver.rs

@@ -61,7 +61,7 @@ pub async fn rwhod_packet_receiver_task(
                 }
 
                 let mut store = whod_status_store.write().await;
-                store.insert(status_update.hostname.clone(), status_update);
+                store.upsert(status_update);
             }
             Err(err) => {
                 tracing::error!("Error processing whod packet from {src}: {err}");

src/server/rwhod/status_registry.rs

@@ -0,0 +1,155 @@
+use std::collections::{BTreeSet, HashMap};
+
+use chrono::{DateTime, Utc};
+
+use crate::proto::WhodStatusUpdate;
+
+/// A bounded collection of rwhod status updates, keyed by hostname.
+///
+/// This exists to put a hard cap on memory usage regardless of how many
+/// distinct hostnames a (potentially malicious) sender on the LAN tries to
+/// report. When the registry is at capacity and an update for a
+/// previously-unseen hostname comes in, the entry with the oldest
+/// `recvtime` is evicted to make room for it. Hosts that keep broadcasting
+/// periodically naturally stay "recent" and are never evicted by this.
+#[derive(Debug)]
+pub struct RwhodStatusRegistry {
+    max_entries: usize,
+    by_hostname: HashMap<String, WhodStatusUpdate>,
+    by_recvtime: BTreeSet<(DateTime<Utc>, String)>,
+}
+
+impl RwhodStatusRegistry {
+    pub fn new(max_entries: usize) -> Self {
+        Self {
+            max_entries,
+            by_hostname: HashMap::new(),
+            by_recvtime: BTreeSet::new(),
+        }
+    }
+
+    /// Insert or refresh a status update, keyed by its hostname.
+    ///
+    /// If the registry is at capacity and `status_update` is for a
+    /// previously-unseen hostname, the entry with the oldest `recvtime`
+    /// will be evicted to make room for it.
+    pub fn upsert(&mut self, status_update: WhodStatusUpdate) {
+        let Some(recvtime) = status_update.recvtime else {
+            tracing::warn!(
+                "Refusing to store whod status update from '{}' with no recvtime set",
+                status_update.hostname
+            );
+            return;
+        };
+
+        if self.max_entries == 0 {
+            tracing::warn!(
+                "rwhod status registry capacity is 0; dropping update from '{}'",
+                status_update.hostname
+            );
+            return;
+        }
+
+        let is_known_hostname =
+            if let Some(previous) = self.by_hostname.get(&status_update.hostname) {
+                if let Some(previous_recvtime) = previous.recvtime {
+                    self.by_recvtime
+                        .remove(&(previous_recvtime, status_update.hostname.clone()));
+                }
+                true
+            } else {
+                false
+            };
+
+        if !is_known_hostname
+            && self.by_hostname.len() >= self.max_entries
+            && let Some((stale_recvtime, stale_hostname)) = self.by_recvtime.pop_first()
+        {
+            tracing::warn!(
+                "rwhod status registry at capacity ({} entries); evicting stalest entry '{}' (last seen {}) to make room for '{}'",
+                self.max_entries,
+                stale_hostname,
+                stale_recvtime,
+                status_update.hostname
+            );
+            self.by_hostname.remove(&stale_hostname);
+        }
+
+        self.by_recvtime
+            .insert((recvtime, status_update.hostname.clone()));
+        self.by_hostname
+            .insert(status_update.hostname.clone(), status_update);
+    }
+
+    pub fn len(&self) -> usize {
+        self.by_hostname.len()
+    }
+
+    pub fn is_empty(&self) -> bool {
+        self.by_hostname.is_empty()
+    }
+
+    pub fn values(&self) -> impl Iterator<Item = &WhodStatusUpdate> {
+        self.by_hostname.values()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn status_update_at(hostname: &str, recvtime_secs: i64) -> WhodStatusUpdate {
+        let now = DateTime::from_timestamp(recvtime_secs, 0).unwrap();
+        WhodStatusUpdate::new(now, Some(now), hostname.to_string(), (0, 0, 0), now, vec![])
+    }
+
+    #[test]
+    fn test_upsert_evicts_stalest_entry_when_at_capacity() {
+        let mut registry = RwhodStatusRegistry::new(2);
+
+        registry.upsert(status_update_at("a", 1));
+        registry.upsert(status_update_at("b", 2));
+        assert_eq!(registry.len(), 2);
+
+        registry.upsert(status_update_at("c", 3));
+
+        assert_eq!(registry.len(), 2);
+        let hostnames: std::collections::HashSet<_> =
+            registry.values().map(|u| u.hostname.as_str()).collect();
+        assert_eq!(hostnames, std::collections::HashSet::from(["b", "c"]));
+    }
+
+    #[test]
+    fn test_upsert_refresh_does_not_evict_itself() {
+        let mut registry = RwhodStatusRegistry::new(2);
+
+        registry.upsert(status_update_at("a", 1));
+        registry.upsert(status_update_at("a", 2));
+
+        assert_eq!(registry.len(), 1);
+        assert_eq!(
+            registry.values().next().unwrap().recvtime,
+            Some(DateTime::from_timestamp(2, 0).unwrap())
+        );
+    }
+
+    #[test]
+    fn test_upsert_with_zero_capacity_drops_everything() {
+        let mut registry = RwhodStatusRegistry::new(0);
+
+        registry.upsert(status_update_at("a", 1));
+
+        assert!(registry.is_empty());
+    }
+
+    #[test]
+    fn test_refreshing_a_hostname_does_not_count_against_capacity() {
+        let mut registry = RwhodStatusRegistry::new(3);
+
+        registry.upsert(status_update_at("a", 1));
+        registry.upsert(status_update_at("a", 2));
+        registry.upsert(status_update_at("a", 3));
+
+        assert_eq!(registry.len(), 1);
+    }
+}