Skip to main content

cortexkit_bus_trait/
census.rs

1use std::collections::BTreeSet;
2
3use crate::{BusError, BusResult, RegisterUpdate, WatchEvent};
4
5/// Tracks presence and deletion observations made on this connection as a local safeguard;
6/// the server-side credential revocation remains authoritative.
7#[derive(Debug, Default, Clone)]
8pub struct CensusWatchGuard {
9    present: BTreeSet<String>,
10    observed_deleted: BTreeSet<String>,
11    has_snapshot: bool,
12    connection_dropped: bool,
13}
14
15impl CensusWatchGuard {
16    pub fn new() -> Self {
17        Self::default()
18    }
19
20    pub fn observe(&mut self, event: &WatchEvent) {
21        if self.connection_dropped && !matches!(event, WatchEvent::Snapshot { .. }) {
22            return;
23        }
24        match event {
25            WatchEvent::Snapshot { entries, .. } => {
26                self.connection_dropped = false;
27                self.present = entries.iter().map(|entry| entry.key.clone()).collect();
28                self.observed_deleted.clear();
29                self.has_snapshot = true;
30            }
31            WatchEvent::Update(RegisterUpdate::Put(entry)) => {
32                self.present.insert(entry.key.clone());
33                self.observed_deleted.remove(&entry.key);
34            }
35            WatchEvent::Update(RegisterUpdate::Delete(tombstone)) => {
36                self.present.remove(&tombstone.key);
37                self.observed_deleted.insert(tombstone.key.clone());
38            }
39        }
40    }
41
42    /// Clears observations after disconnect so missing events are not treated as proof of absence.
43    pub fn connection_dropped(&mut self) {
44        self.present.clear();
45        self.observed_deleted.clear();
46        self.has_snapshot = false;
47        self.connection_dropped = true;
48    }
49
50    pub fn has_snapshot(&self) -> bool {
51        self.has_snapshot
52    }
53
54    pub fn check_publish(&self, census_key: &str, subject: &str) -> BusResult<()> {
55        if self.connection_dropped {
56            return Ok(());
57        }
58
59        let confirmed_absent = self.observed_deleted.contains(census_key)
60            || (self.has_snapshot && !self.present.contains(census_key));
61        if confirmed_absent {
62            return Err(BusError::denied(
63                subject,
64                format!("confirmed absent census key {census_key}"),
65            ));
66        }
67        Ok(())
68    }
69}
70
71#[cfg(test)]
72mod tests {
73    use super::*;
74    use crate::{RegisterEntry, Revision, Tombstone};
75
76    #[test]
77    fn no_snapshot_and_dropped_connection_are_both_neutral() {
78        let mut guard = CensusWatchGuard::new();
79        assert_eq!(
80            guard.check_publish("module-a", "ck.box.room.r.post"),
81            Ok(())
82        );
83
84        guard.observe(&WatchEvent::Snapshot {
85            entries: Vec::new(),
86            revision: Revision(0),
87        });
88        assert!(matches!(
89            guard.check_publish("module-a", "ck.box.room.r.post"),
90            Err(BusError::Denied { .. })
91        ));
92
93        guard.connection_dropped();
94        guard.observe(&WatchEvent::Update(RegisterUpdate::Delete(Tombstone {
95            key: "module-a".into(),
96            revision: Revision(1),
97        })));
98        assert_eq!(
99            guard.check_publish("module-a", "ck.box.room.r.post"),
100            Ok(())
101        );
102    }
103
104    #[test]
105    fn observed_delete_confirms_absence_even_without_a_snapshot() {
106        let mut guard = CensusWatchGuard::new();
107        guard.observe(&WatchEvent::Update(RegisterUpdate::Delete(Tombstone {
108            key: "module-a".into(),
109            revision: Revision(2),
110        })));
111
112        assert!(matches!(
113            guard.check_publish("module-a", "ck.box.wake.agent.fire"),
114            Err(BusError::Denied { subject, .. }) if subject == "ck.box.wake.agent.fire"
115        ));
116
117        guard.observe(&WatchEvent::Update(RegisterUpdate::Put(RegisterEntry {
118            key: "module-a".into(),
119            value: Vec::new(),
120            revision: Revision(3),
121        })));
122        assert_eq!(
123            guard.check_publish("module-a", "ck.box.wake.agent.fire"),
124            Ok(())
125        );
126    }
127}