umsh_journal_store/
counter.rs

1//! The device node's persisted frame-counter map (device-node plan
2//! increment 4).
3//!
4//! The MAC's `CounterStore` contract is per-context `store` calls
5//! followed by one `flush` per persistence cycle, already batched to
6//! one write per `COUNTER_PERSIST_BLOCK_SIZE` (128) secured frames.
7//! This map is the RAM image behind that contract: `store` upserts
8//! here, and `flush` serializes the *whole* map into a single
9//! [`proto`](crate::proto) record in the firmware's counter journal.
10//! Whole-map records keep the journal machinery identical to the other
11//! journals — newest generation wins, one committed record is the
12//! entire persisted state — at a size (≤ ~600 bytes) far under the
13//! record payload bound.
14//!
15//! Contexts are the MAC's own key formats: the raw 32-byte identity
16//! public key for TX boundaries and `mac.rx:` + public key (39 bytes)
17//! for per-peer RX boundaries. The map stores them opaquely.
18
19/// Longest stored context: `mac.rx:` (7) + 32-byte key, with headroom.
20pub const MAX_KEY_LEN: usize = 40;
21
22/// One TX boundary for the device identity plus RX boundaries for
23/// `MAX_DEV_PEERS` (8) peers, with slack for a stale generation of
24/// entries surviving until the next journal clear (identity
25/// provisioning and CMD_CLEAR both clear the journal).
26pub const MAX_ENTRIES: usize = 12;
27
28/// Upper bound of [`CounterMap::encode`]'s output.
29pub const ENCODED_MAX: usize = MAX_ENTRIES * (1 + MAX_KEY_LEN + 4);
30
31#[derive(Clone, Debug, PartialEq, Eq)]
32struct Entry {
33    key: heapless::Vec<u8, MAX_KEY_LEN>,
34    value: u32,
35}
36
37/// The map was full and the context could not be added.
38#[derive(Clone, Copy, Debug, PartialEq, Eq)]
39pub struct MapFull;
40
41#[derive(Clone, Debug, Default, PartialEq, Eq)]
42pub struct CounterMap {
43    entries: heapless::Vec<Entry, MAX_ENTRIES>,
44}
45
46impl CounterMap {
47    pub const fn new() -> Self {
48        Self {
49            entries: heapless::Vec::new(),
50        }
51    }
52
53    /// The stored value for `key`, if any.
54    pub fn get(&self, key: &[u8]) -> Option<u32> {
55        self.entries
56            .iter()
57            .find(|entry| entry.key == key)
58            .map(|entry| entry.value)
59    }
60
61    /// Upsert `key` to `value`. Returns whether the map changed (an
62    /// equal-value overwrite is a no-op, so callers can skip a flush).
63    pub fn set(&mut self, key: &[u8], value: u32) -> Result<bool, MapFull> {
64        if let Some(entry) = self.entries.iter_mut().find(|entry| entry.key == key) {
65            if entry.value == value {
66                return Ok(false);
67            }
68            entry.value = value;
69            return Ok(true);
70        }
71        let key = heapless::Vec::from_slice(key).map_err(|_| MapFull)?;
72        self.entries
73            .push(Entry { key, value })
74            .map_err(|_| MapFull)?;
75        Ok(true)
76    }
77
78    /// Drop every entry (factory clear).
79    pub fn clear(&mut self) {
80        self.entries.clear();
81    }
82
83    /// Drop TX-boundary entries belonging to any identity other than
84    /// `keep` and report whether anything was removed. TX contexts are
85    /// the MAC's raw 32-byte identity public key; every other context
86    /// format (the 39-byte `mac.rx:` form) is left alone.
87    pub fn prune_tx_except(&mut self, keep: &[u8; 32]) -> bool {
88        let before = self.entries.len();
89        self.entries
90            .retain(|entry| entry.key.len() != 32 || entry.key == keep);
91        before != self.entries.len()
92    }
93
94    pub fn len(&self) -> usize {
95        self.entries.len()
96    }
97
98    /// Serialize into `out` (sized [`ENCODED_MAX`] or larger); returns
99    /// the encoded length. Layout per entry: key length (1 byte), key,
100    /// little-endian u32 value.
101    pub fn encode(&self, out: &mut [u8]) -> Option<usize> {
102        let mut at = 0;
103        for entry in self.entries.iter() {
104            let needed = 1 + entry.key.len() + 4;
105            if out.len() - at < needed {
106                return None;
107            }
108            out[at] = entry.key.len() as u8;
109            out[at + 1..at + 1 + entry.key.len()].copy_from_slice(&entry.key);
110            out[at + 1 + entry.key.len()..at + needed].copy_from_slice(&entry.value.to_le_bytes());
111            at += needed;
112        }
113        Some(at)
114    }
115
116    /// Parse a persisted payload. Anything malformed — truncated
117    /// entries, oversized keys, more entries than capacity — yields
118    /// `None`, and the mount treats the journal as empty (counters
119    /// reseed, which is the safe direction for TX boundaries).
120    pub fn decode(payload: &[u8]) -> Option<Self> {
121        let mut map = Self::new();
122        let mut at = 0;
123        while at < payload.len() {
124            let key_len = usize::from(payload[at]);
125            if key_len == 0 || key_len > MAX_KEY_LEN {
126                return None;
127            }
128            let end = at + 1 + key_len + 4;
129            if end > payload.len() {
130                return None;
131            }
132            let key = heapless::Vec::from_slice(&payload[at + 1..at + 1 + key_len]).ok()?;
133            let value = u32::from_le_bytes(payload[at + 1 + key_len..end].try_into().ok()?);
134            map.entries.push(Entry { key, value }).ok()?;
135            at = end;
136        }
137        Some(map)
138    }
139}
140
141#[cfg(test)]
142mod tests {
143    use super::*;
144
145    #[test]
146    fn round_trips_and_upserts() {
147        let mut map = CounterMap::new();
148        let tx_key = [0xAA; 32];
149        let mut rx_key = [0u8; 39];
150        rx_key[..7].copy_from_slice(b"mac.rx:");
151        rx_key[7..].fill(0xBB);
152
153        assert_eq!(map.set(&tx_key, 128), Ok(true));
154        assert_eq!(map.set(&rx_key, 256), Ok(true));
155        // Equal-value overwrite reports "unchanged".
156        assert_eq!(map.set(&tx_key, 128), Ok(false));
157        assert_eq!(map.set(&tx_key, 384), Ok(true));
158        assert_eq!(map.get(&tx_key), Some(384));
159        assert_eq!(map.get(&rx_key), Some(256));
160        assert_eq!(map.get(&[0x01; 32]), None);
161
162        let mut buf = [0u8; ENCODED_MAX];
163        let len = map.encode(&mut buf).unwrap();
164        let decoded = CounterMap::decode(&buf[..len]).unwrap();
165        assert_eq!(decoded, map);
166
167        // Empty map round-trips to an empty payload.
168        assert_eq!(CounterMap::new().encode(&mut buf), Some(0));
169        assert_eq!(CounterMap::decode(&[]), Some(CounterMap::new()));
170    }
171
172    #[test]
173    fn capacity_and_malformed_payloads() {
174        let mut map = CounterMap::new();
175        for index in 0..MAX_ENTRIES {
176            let key = [index as u8; 32];
177            assert_eq!(map.set(&key, index as u32), Ok(true));
178        }
179        // Full: a new context is refused, existing ones still update.
180        assert_eq!(map.set(&[0xFF; 32], 1), Err(MapFull));
181        assert_eq!(map.set(&[0x00; 32], 7), Ok(true));
182
183        // Truncated entry, zero-length key, oversized key, trailing
184        // garbage after a valid entry.
185        assert_eq!(CounterMap::decode(&[5, 1, 2]), None);
186        assert_eq!(CounterMap::decode(&[0, 0, 0, 0, 0]), None);
187        let mut oversized = [0u8; 1 + MAX_KEY_LEN + 1 + 4];
188        oversized[0] = MAX_KEY_LEN as u8 + 1;
189        assert_eq!(CounterMap::decode(&oversized), None);
190        let mut valid = [0u8; ENCODED_MAX];
191        let mut one = CounterMap::new();
192        one.set(&[1; 32], 9).unwrap();
193        let len = one.encode(&mut valid).unwrap();
194        valid[len] = 3; // claims a 3-byte key with no data behind it
195        assert_eq!(CounterMap::decode(&valid[..len + 1]), None);
196
197        // clear() empties the map.
198        map.clear();
199        assert_eq!(map.len(), 0);
200        assert_eq!(map.get(&[0x00; 32]), None);
201    }
202
203    #[test]
204    fn prune_drops_only_foreign_tx_entries() {
205        let mut map = CounterMap::new();
206        let old_pk = [0x0A; 32];
207        let new_pk = [0x0B; 32];
208        let mut rx_key = [0u8; 39];
209        rx_key[..7].copy_from_slice(b"mac.rx:");
210        map.set(&old_pk, 128).unwrap();
211        map.set(&rx_key, 256).unwrap();
212
213        assert!(map.prune_tx_except(&new_pk));
214        assert_eq!(map.get(&old_pk), None);
215        assert_eq!(map.get(&rx_key), Some(256));
216
217        // Idempotent, and the surviving identity's own entry stays.
218        assert!(!map.prune_tx_except(&new_pk));
219        map.set(&new_pk, 384).unwrap();
220        assert!(!map.prune_tx_except(&new_pk));
221        assert_eq!(map.get(&new_pk), Some(384));
222    }
223}