umsh_journal_store/
counter.rs1pub const MAX_KEY_LEN: usize = 40;
21
22pub const MAX_ENTRIES: usize = 12;
27
28pub 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#[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 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 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 pub fn clear(&mut self) {
80 self.entries.clear();
81 }
82
83 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 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 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 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 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 assert_eq!(map.set(&[0xFF; 32], 1), Err(MapFull));
181 assert_eq!(map.set(&[0x00; 32], 7), Ok(true));
182
183 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; assert_eq!(CounterMap::decode(&valid[..len + 1]), None);
196
197 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 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}