diff --git a/nix/module.nix b/nix/module.nix index fd11f79..4a12210 100644 --- a/nix/module.nix +++ b/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 = [ ]; diff --git a/src/bin/roowhod.rs b/src/bin/roowhod.rs index c090573..bbe91fb 100644 --- a/src/bin/roowhod.rs +++ b/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 ", @@ -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(); diff --git a/src/server/config.rs b/src/server/config.rs index 02ac079..e496360 100644 --- a/src/server/config.rs +++ b/src/server/config.rs @@ -48,6 +48,13 @@ pub struct RwhodConfig { /// /// If left as `None`, defaults to 60 seconds. pub send_interval_seconds: Option, + + /// 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, } #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] diff --git a/src/server/rwhod.rs b/src/server/rwhod.rs index 9f3948a..43a5603 100644 --- a/src/server/rwhod.rs +++ b/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>>; +pub type RwhodStatusStore = Arc>; diff --git a/src/server/rwhod/packet_receiver.rs b/src/server/rwhod/packet_receiver.rs index 78209db..10da2d7 100644 --- a/src/server/rwhod/packet_receiver.rs +++ b/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}"); diff --git a/src/server/rwhod/status_registry.rs b/src/server/rwhod/status_registry.rs new file mode 100644 index 0000000..42b541b --- /dev/null +++ b/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, + by_recvtime: BTreeSet<(DateTime, 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 { + 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); + } +}