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

umsh_mac/
peers.rs
