umsh_mac/
peers.rs

1use heapless::{LinearMap, Vec};
2use umsh_core::{ChannelId, ChannelKey, NodeHint, PublicKey, RouterHint};
3use umsh_crypto::{DerivedChannelKeys, PairwiseKeys};
4
5use crate::{CapacityError, cache::ReplayWindow};
6
7/// Opaque identifier for one remote peer.
8#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
9pub struct PeerId(pub u8);
10
11/// Learned routing information for a remote peer.
12#[derive(Clone, Debug, PartialEq, Eq)]
13pub enum CachedRoute {
14    /// Peer is directly reachable without any intermediate routers.
15    ///
16    /// Inferred when a packet arrives with no source-route or traceroute option
17    /// (or an empty traceroute) and `FHOPS_ACC == 0`.
18    Direct,
19    /// Explicit source route derived by reversing the inbound traceroute.
20    Source(Vec<RouterHint, 15>),
21    /// Flood-delivery parameters learned from an inbound packet.
22    Flood { hops: u8, regions: Vec<[u8; 2], 8> },
23}
24
25/// Shared metadata tracked for a remote peer.
26#[derive(Clone, Debug, PartialEq, Eq)]
27pub struct PeerInfo {
28    /// Full public key.
29    pub public_key: PublicKey,
30    /// Whether this peer was explicitly configured by the local application.
31    pub pinned: bool,
32    /// Most recent learned route, if any.
33    pub route: Option<CachedRoute>,
34    /// Most recent observation timestamp.
35    pub last_seen_ms: u64,
36    /// Highest RX frame counter loaded from persistent storage at boot.
37    ///
38    /// Non-zero means a stored boundary was found. When pairwise keys are
39    /// first installed for this peer, the replay window is initialised to
40    /// this value so that frames from before the reboot are rejected.
41    pub initial_rx_counter: u32,
42}
43
44/// Outcome of inserting or updating an auto-learned peer.
45#[derive(Clone, Copy, Debug, PartialEq, Eq)]
46pub struct AutoPeerUpdate {
47    /// Slot assigned to the peer.
48    pub peer_id: PeerId,
49    /// Previous peer key displaced from this slot, if any.
50    pub evicted_key: Option<PublicKey>,
51}
52
53/// Outcome of removing a peer from the registry.
54#[derive(Clone, Copy, Debug, PartialEq, Eq)]
55pub struct PeerRemoval {
56    /// Public key of the removed peer.
57    pub removed_key: PublicKey,
58    /// When removal swap-moved the last entry into the freed slot: that
59    /// entry's `(old, new)` identifiers. Every `PeerId`-keyed structure must
60    /// be re-keyed accordingly.
61    pub moved: Option<(PeerId, PeerId)>,
62}
63
64/// Fixed-capacity registry of remote peers.
65#[derive(Clone, Debug)]
66pub struct PeerRegistry<const N: usize> {
67    peers: Vec<PeerInfo, N>,
68}
69
70impl<const N: usize> Default for PeerRegistry<N> {
71    fn default() -> Self {
72        Self::new()
73    }
74}
75
76impl<const N: usize> PeerRegistry<N> {
77    /// Create an empty peer registry.
78    pub fn new() -> Self {
79        Self { peers: Vec::new() }
80    }
81
82    /// Iterate over peers whose derived hint matches `hint`.
83    pub fn lookup_by_hint(&self, hint: &NodeHint) -> impl Iterator<Item = (PeerId, &PeerInfo)> {
84        self.peers
85            .iter()
86            .enumerate()
87            .filter(move |(_, peer)| peer.public_key.hint() == *hint)
88            .map(|(index, peer)| (PeerId(index as u8), peer))
89    }
90
91    /// Look up a peer by full public key.
92    pub fn lookup_by_key(&self, key: &PublicKey) -> Option<(PeerId, &PeerInfo)> {
93        self.peers
94            .iter()
95            .enumerate()
96            .find(|(_, peer)| peer.public_key == *key)
97            .map(|(index, peer)| (PeerId(index as u8), peer))
98    }
99
100    /// Iterate over all registered peers.
101    pub fn iter(&self) -> impl Iterator<Item = (PeerId, &PeerInfo)> {
102        self.peers
103            .iter()
104            .enumerate()
105            .map(|(index, peer)| (PeerId(index as u8), peer))
106    }
107
108    /// Borrow peer metadata by identifier.
109    pub fn get(&self, id: PeerId) -> Option<&PeerInfo> {
110        self.peers.get(id.0 as usize)
111    }
112
113    /// Mutably borrow peer metadata by identifier.
114    pub fn get_mut(&mut self, id: PeerId) -> Option<&mut PeerInfo> {
115        self.peers.get_mut(id.0 as usize)
116    }
117
118    /// Insert or refresh an explicitly configured peer entry.
119    pub fn try_insert_or_update(&mut self, key: PublicKey) -> Result<PeerId, CapacityError> {
120        if let Some((id, peer)) = self
121            .peers
122            .iter_mut()
123            .enumerate()
124            .find(|(_, peer)| peer.public_key == key)
125        {
126            peer.public_key = key;
127            peer.pinned = true;
128            return Ok(PeerId(id as u8));
129        }
130
131        self.peers
132            .push(PeerInfo {
133                public_key: key,
134                pinned: true,
135                route: None,
136                last_seen_ms: 0,
137                initial_rx_counter: 0,
138            })
139            .map_err(|_| CapacityError)?;
140        Ok(PeerId((self.peers.len() - 1) as u8))
141    }
142
143    /// Insert or refresh an opportunistically learned peer entry.
144    ///
145    /// When the registry is full, this may recycle the oldest non-pinned entry in place
146    /// rather than failing. Explicitly configured (`pinned`) peers are never displaced.
147    pub fn try_insert_or_update_auto(
148        &mut self,
149        key: PublicKey,
150        now_ms: u64,
151    ) -> Result<AutoPeerUpdate, CapacityError> {
152        if let Some((id, peer)) = self
153            .peers
154            .iter_mut()
155            .enumerate()
156            .find(|(_, peer)| peer.public_key == key)
157        {
158            peer.last_seen_ms = now_ms;
159            return Ok(AutoPeerUpdate {
160                peer_id: PeerId(id as u8),
161                evicted_key: None,
162            });
163        }
164
165        if self.peers.len() < N {
166            self.peers
167                .push(PeerInfo {
168                    public_key: key,
169                    pinned: false,
170                    route: None,
171                    last_seen_ms: now_ms,
172                    initial_rx_counter: 0,
173                })
174                .map_err(|_| CapacityError)?;
175            return Ok(AutoPeerUpdate {
176                peer_id: PeerId((self.peers.len() - 1) as u8),
177                evicted_key: None,
178            });
179        }
180
181        let Some((index, oldest)) = self
182            .peers
183            .iter()
184            .enumerate()
185            .filter(|(_, peer)| !peer.pinned)
186            .min_by_key(|(_, peer)| peer.last_seen_ms)
187        else {
188            return Err(CapacityError);
189        };
190
191        let evicted_key = oldest.public_key;
192        self.peers[index] = PeerInfo {
193            public_key: key,
194            pinned: false,
195            route: None,
196            last_seen_ms: now_ms,
197            initial_rx_counter: 0,
198        };
199        Ok(AutoPeerUpdate {
200            peer_id: PeerId(index as u8),
201            evicted_key: Some(evicted_key),
202        })
203    }
204
205    /// Remove a peer, freeing its slot for reuse.
206    ///
207    /// The registry is dense — a `PeerId` is an index — so removal swap-moves
208    /// the last entry into the freed slot. The returned record names that
209    /// move so the caller can re-key any state held under the moved peer's
210    /// old identifier.
211    pub fn remove(&mut self, id: PeerId) -> Option<PeerRemoval> {
212        let index = id.0 as usize;
213        if index >= self.peers.len() {
214            return None;
215        }
216        let last = self.peers.len() - 1;
217        let removed = self.peers.swap_remove(index);
218        let moved = (index != last).then_some((PeerId(last as u8), id));
219        Some(PeerRemoval {
220            removed_key: removed.public_key,
221            moved,
222        })
223    }
224
225    /// Update the cached route for `id`.
226    pub fn update_route(&mut self, id: PeerId, route: CachedRoute) {
227        if let Some(peer) = self.get_mut(id) {
228            peer.route = Some(route);
229        }
230    }
231
232    /// Forget the cached route for `id`, returning whether one was held.
233    ///
234    /// The peer itself stays registered; subsequent sends fall back to the
235    /// default delivery mode until a fresh inbound packet teaches a route.
236    pub fn clear_route(&mut self, id: PeerId) -> bool {
237        self.get_mut(id)
238            .map(|peer| peer.route.take().is_some())
239            .unwrap_or(false)
240    }
241
242    /// Refresh the last-seen timestamp for `id`.
243    pub fn touch(&mut self, id: PeerId, now_ms: u64) {
244        if let Some(peer) = self.get_mut(id) {
245            peer.last_seen_ms = now_ms;
246        }
247    }
248}
249
250/// Per-peer secure transport state.
251#[derive(Clone)]
252pub struct PeerCryptoState {
253    /// Pairwise encryption and MIC keys.
254    pub pairwise_keys: PairwiseKeys,
255    /// Replay state for traffic from this peer.
256    pub replay_window: ReplayWindow,
257    /// Highest `last_accepted` value written to persistent storage.
258    /// Updated by `service_rx_counter_persistence` after each flush.
259    pub persisted_rx_counter: u32,
260    /// Set when `last_accepted` has advanced `COUNTER_PERSIST_BLOCK_SIZE`
261    /// beyond `persisted_rx_counter`. Cleared after the next flush.
262    pub needs_rx_persist: bool,
263}
264
265/// Fixed-capacity map of per-peer secure transport state.
266#[derive(Clone)]
267pub struct PeerCryptoMap<const N: usize> {
268    entries: LinearMap<PeerId, PeerCryptoState, N>,
269}
270
271impl<const N: usize> Default for PeerCryptoMap<N> {
272    fn default() -> Self {
273        Self::new()
274    }
275}
276
277impl<const N: usize> PeerCryptoMap<N> {
278    /// Create an empty peer-crypto map.
279    pub fn new() -> Self {
280        Self {
281            entries: LinearMap::new(),
282        }
283    }
284
285    /// Borrow one peer state.
286    pub fn get(&self, id: &PeerId) -> Option<&PeerCryptoState> {
287        self.entries.get(id)
288    }
289
290    /// Mutably borrow one peer state.
291    pub fn get_mut(&mut self, id: &PeerId) -> Option<&mut PeerCryptoState> {
292        self.entries.get_mut(id)
293    }
294
295    /// Insert or replace state for a peer.
296    pub fn insert(
297        &mut self,
298        id: PeerId,
299        state: PeerCryptoState,
300    ) -> Result<Option<PeerCryptoState>, CapacityError> {
301        self.entries.insert(id, state).map_err(|_| CapacityError)
302    }
303
304    /// Remove state for a peer.
305    pub fn remove(&mut self, id: &PeerId) -> Option<PeerCryptoState> {
306        self.entries.remove(id)
307    }
308
309    /// Iterate over all peer crypto entries.
310    pub fn iter(&self) -> impl Iterator<Item = (&PeerId, &PeerCryptoState)> {
311        self.entries.iter()
312    }
313}
314
315/// Replay state for a sender known only by hint.
316#[derive(Clone)]
317pub struct HintReplayState {
318    /// Replay window for the hint-only sender.
319    pub window: ReplayWindow,
320    /// Most recent observation timestamp.
321    pub last_seen_ms: u64,
322}
323
324/// Shared state for one multicast channel.
325#[derive(Clone)]
326pub struct ChannelState<const RN: usize = 8, const HN: usize = 8> {
327    /// Raw channel key.
328    pub channel_key: ChannelKey,
329    /// Derived transport keys and identifier.
330    pub derived: DerivedChannelKeys,
331    /// Replay windows for peers resolved to full identities.
332    pub replay: LinearMap<PeerId, ReplayWindow, RN>,
333    /// Replay windows for senders known only by hint.
334    pub hint_replay: LinearMap<NodeHint, HintReplayState, HN>,
335}
336
337impl<const RN: usize, const HN: usize> ChannelState<RN, HN> {
338    /// Create a new channel-state record.
339    pub fn new(channel_key: ChannelKey, derived: DerivedChannelKeys) -> Self {
340        Self {
341            channel_key,
342            derived,
343            replay: LinearMap::new(),
344            hint_replay: LinearMap::new(),
345        }
346    }
347}
348
349/// Fixed-capacity channel table shared by the MAC coordinator.
350#[derive(Clone)]
351pub struct ChannelTable<const N: usize, const RN: usize = 8, const HN: usize = 8> {
352    channels: Vec<ChannelState<RN, HN>, N>,
353}
354
355impl<const N: usize, const RN: usize, const HN: usize> Default for ChannelTable<N, RN, HN> {
356    fn default() -> Self {
357        Self::new()
358    }
359}
360
361impl<const N: usize, const RN: usize, const HN: usize> ChannelTable<N, RN, HN> {
362    /// Create an empty channel table.
363    pub fn new() -> Self {
364        Self {
365            channels: Vec::new(),
366        }
367    }
368
369    /// Return the number of configured channels.
370    pub fn len(&self) -> usize {
371        self.channels.len()
372    }
373
374    /// Return whether no channels are configured.
375    pub fn is_empty(&self) -> bool {
376        self.channels.is_empty()
377    }
378
379    /// Iterate over channels whose derived identifier matches `id`.
380    pub fn lookup_by_id(&self, id: &ChannelId) -> impl Iterator<Item = &ChannelState<RN, HN>> {
381        self.channels
382            .iter()
383            .filter(move |channel| channel.derived.channel_id == *id)
384    }
385
386    /// Mutably borrow the first channel whose derived identifier matches `id`.
387    pub fn get_mut_by_id(&mut self, id: &ChannelId) -> Option<&mut ChannelState<RN, HN>> {
388        self.channels
389            .iter_mut()
390            .find(|channel| channel.derived.channel_id == *id)
391    }
392
393    /// Mutably iterate over all channel states.
394    pub fn iter_mut(&mut self) -> impl Iterator<Item = &mut ChannelState<RN, HN>> {
395        self.channels.iter_mut()
396    }
397
398    /// Remove the channel holding this exact key, discarding its replay
399    /// state with it (re-adding the key later starts at first contact).
400    /// Returns whether a channel was removed.
401    pub fn remove_by_key(&mut self, key: &ChannelKey) -> bool {
402        let Some(index) = self
403            .channels
404            .iter()
405            .position(|channel| channel.channel_key.0 == key.0)
406        else {
407            return false;
408        };
409        self.channels.swap_remove(index);
410        true
411    }
412
413    /// Add or replace a channel entry.
414    pub fn try_add(
415        &mut self,
416        key: ChannelKey,
417        derived: DerivedChannelKeys,
418    ) -> Result<(), CapacityError> {
419        if let Some(channel) = self.get_mut_by_id(&derived.channel_id) {
420            channel.channel_key = key;
421            channel.derived = derived;
422            return Ok(());
423        }
424
425        self.channels
426            .push(ChannelState::new(key, derived))
427            .map_err(|_| CapacityError)
428    }
429}
umsh_mac/peers.rs

umsh_mac/
peers.rs
