roowho2_lib/server/rwhod/
status_registry.rs

1use std::collections::{BTreeSet, HashMap};
2
3use chrono::{DateTime, Utc};
4
5use crate::proto::WhodStatusUpdate;
6
7/// A bounded collection of rwhod status updates, keyed by hostname.
8///
9/// This exists to put a hard cap on memory usage regardless of how many
10/// distinct hostnames a (potentially malicious) sender on the LAN tries to
11/// report. When the registry is at capacity and an update for a
12/// previously-unseen hostname comes in, the entry with the oldest
13/// `recvtime` is evicted to make room for it. Hosts that keep broadcasting
14/// periodically naturally stay "recent" and are never evicted by this.
15#[derive(Debug)]
16pub struct RwhodStatusRegistry {
17    max_entries: usize,
18    by_hostname: HashMap<String, WhodStatusUpdate>,
19    by_recvtime: BTreeSet<(DateTime<Utc>, String)>,
20}
21
22impl RwhodStatusRegistry {
23    pub fn new(max_entries: usize) -> Self {
24        Self {
25            max_entries,
26            by_hostname: HashMap::new(),
27            by_recvtime: BTreeSet::new(),
28        }
29    }
30
31    /// Insert or refresh a status update, keyed by its hostname.
32    ///
33    /// If the registry is at capacity and `status_update` is for a
34    /// previously-unseen hostname, the entry with the oldest `recvtime`
35    /// will be evicted to make room for it.
36    pub fn upsert(&mut self, status_update: WhodStatusUpdate) {
37        let Some(recvtime) = status_update.recvtime else {
38            tracing::warn!(
39                "Refusing to store whod status update from '{}' with no recvtime set",
40                status_update.hostname
41            );
42            return;
43        };
44
45        if self.max_entries == 0 {
46            tracing::warn!(
47                "rwhod status registry capacity is 0; dropping update from '{}'",
48                status_update.hostname
49            );
50            return;
51        }
52
53        let is_known_hostname =
54            if let Some(previous) = self.by_hostname.get(&status_update.hostname) {
55                if let Some(previous_recvtime) = previous.recvtime {
56                    self.by_recvtime
57                        .remove(&(previous_recvtime, status_update.hostname.clone()));
58                }
59                true
60            } else {
61                false
62            };
63
64        if !is_known_hostname
65            && self.by_hostname.len() >= self.max_entries
66            && let Some((stale_recvtime, stale_hostname)) = self.by_recvtime.pop_first()
67        {
68            tracing::warn!(
69                "rwhod status registry at capacity ({} entries); evicting stalest entry '{}' (last seen {}) to make room for '{}'",
70                self.max_entries,
71                stale_hostname,
72                stale_recvtime,
73                status_update.hostname
74            );
75            self.by_hostname.remove(&stale_hostname);
76        }
77
78        self.by_recvtime
79            .insert((recvtime, status_update.hostname.clone()));
80        self.by_hostname
81            .insert(status_update.hostname.clone(), status_update);
82    }
83
84    pub fn len(&self) -> usize {
85        self.by_hostname.len()
86    }
87
88    pub fn is_empty(&self) -> bool {
89        self.by_hostname.is_empty()
90    }
91
92    pub fn values(&self) -> impl Iterator<Item = &WhodStatusUpdate> {
93        self.by_hostname.values()
94    }
95}
96
97#[cfg(test)]
98mod tests {
99    use super::*;
100
101    fn status_update_at(hostname: &str, recvtime_secs: i64) -> WhodStatusUpdate {
102        let now = DateTime::from_timestamp(recvtime_secs, 0).unwrap();
103        WhodStatusUpdate::new(now, Some(now), hostname.to_string(), (0, 0, 0), now, vec![])
104    }
105
106    #[test]
107    fn test_upsert_evicts_stalest_entry_when_at_capacity() {
108        let mut registry = RwhodStatusRegistry::new(2);
109
110        registry.upsert(status_update_at("a", 1));
111        registry.upsert(status_update_at("b", 2));
112        assert_eq!(registry.len(), 2);
113
114        registry.upsert(status_update_at("c", 3));
115
116        assert_eq!(registry.len(), 2);
117        let hostnames: std::collections::HashSet<_> =
118            registry.values().map(|u| u.hostname.as_str()).collect();
119        assert_eq!(hostnames, std::collections::HashSet::from(["b", "c"]));
120    }
121
122    #[test]
123    fn test_upsert_refresh_does_not_evict_itself() {
124        let mut registry = RwhodStatusRegistry::new(2);
125
126        registry.upsert(status_update_at("a", 1));
127        registry.upsert(status_update_at("a", 2));
128
129        assert_eq!(registry.len(), 1);
130        assert_eq!(
131            registry.values().next().unwrap().recvtime,
132            Some(DateTime::from_timestamp(2, 0).unwrap())
133        );
134    }
135
136    #[test]
137    fn test_upsert_with_zero_capacity_drops_everything() {
138        let mut registry = RwhodStatusRegistry::new(0);
139
140        registry.upsert(status_update_at("a", 1));
141
142        assert!(registry.is_empty());
143    }
144
145    #[test]
146    fn test_refreshing_a_hostname_does_not_count_against_capacity() {
147        let mut registry = RwhodStatusRegistry::new(3);
148
149        registry.upsert(status_update_at("a", 1));
150        registry.upsert(status_update_at("a", 2));
151        registry.upsert(status_update_at("a", 3));
152
153        assert_eq!(registry.len(), 1);
154    }
155}
roowho2_lib/server/rwhod/status_registry.rs

roowho2_lib/server/rwhod/
status_registry.rs
