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    /// Explicit source route.
15    Source(Vec<RouterHint, 15>),
16    /// Flood-distance estimate.
17    Flood { hops: u8 },
18}
19
20/// Shared metadata tracked for a remote peer.
21#[derive(Clone, Debug, PartialEq, Eq)]
22pub struct PeerInfo {
23    /// Full public key.
24    pub public_key: PublicKey,
25    /// Whether this peer was explicitly configured by the local application.
26    pub pinned: bool,
27    /// Most recent learned route, if any.
28    pub route: Option<CachedRoute>,
29    /// Most recent observation timestamp.
30    pub last_seen_ms: u64,
31}
32
33/// Outcome of inserting or updating an auto-learned peer.
34#[derive(Clone, Copy, Debug, PartialEq, Eq)]
35pub struct AutoPeerUpdate {
36    /// Slot assigned to the peer.
37    pub peer_id: PeerId,
38    /// Previous peer key displaced from this slot, if any.
39    pub evicted_key: Option<PublicKey>,
40}
41
42/// Fixed-capacity registry of remote peers.
43#[derive(Clone, Debug)]
44pub struct PeerRegistry<const N: usize> {
45    peers: Vec<PeerInfo, N>,
46}
47
48impl<const N: usize> Default for PeerRegistry<N> {
49    fn default() -> Self {
50        Self::new()
51    }
52}
53
54impl<const N: usize> PeerRegistry<N> {
55    /// Create an empty peer registry.
56    pub fn new() -> Self {
57        Self { peers: Vec::new() }
58    }
59
60    /// Iterate over peers whose derived hint matches `hint`.
61    pub fn lookup_by_hint(&self, hint: &NodeHint) -> impl Iterator<Item = (PeerId, &PeerInfo)> {
62        self.peers
63            .iter()
64            .enumerate()
65            .filter(move |(_, peer)| peer.public_key.hint() == *hint)
66            .map(|(index, peer)| (PeerId(index as u8), peer))
67    }
68
69    /// Look up a peer by full public key.
70    pub fn lookup_by_key(&self, key: &PublicKey) -> Option<(PeerId, &PeerInfo)> {
71        self.peers
72            .iter()
73            .enumerate()
74            .find(|(_, peer)| peer.public_key == *key)
75            .map(|(index, peer)| (PeerId(index as u8), peer))
76    }
77
78    /// Borrow peer metadata by identifier.
79    pub fn get(&self, id: PeerId) -> Option<&PeerInfo> {
80        self.peers.get(id.0 as usize)
81    }
82
83    /// Mutably borrow peer metadata by identifier.
84    pub fn get_mut(&mut self, id: PeerId) -> Option<&mut PeerInfo> {
85        self.peers.get_mut(id.0 as usize)
86    }
87
88    /// Insert or refresh an explicitly configured peer entry.
89    pub fn try_insert_or_update(&mut self, key: PublicKey) -> Result<PeerId, CapacityError> {
90        if let Some((id, peer)) = self
91            .peers
92            .iter_mut()
93            .enumerate()
94            .find(|(_, peer)| peer.public_key == key)
95        {
96            peer.public_key = key;
97            peer.pinned = true;
98            return Ok(PeerId(id as u8));
99        }
100
101        self.peers
102            .push(PeerInfo {
103                public_key: key,
104                pinned: true,
105                route: None,
106                last_seen_ms: 0,
107            })
108            .map_err(|_| CapacityError)?;
109        Ok(PeerId((self.peers.len() - 1) as u8))
110    }
111
112    /// Insert or refresh an opportunistically learned peer entry.
113    ///
114    /// When the registry is full, this may recycle the oldest non-pinned entry in place
115    /// rather than failing. Explicitly configured (`pinned`) peers are never displaced.
116    pub fn try_insert_or_update_auto(
117        &mut self,
118        key: PublicKey,
119        now_ms: u64,
120    ) -> Result<AutoPeerUpdate, CapacityError> {
121        if let Some((id, peer)) = self
122            .peers
123            .iter_mut()
124            .enumerate()
125            .find(|(_, peer)| peer.public_key == key)
126        {
127            peer.last_seen_ms = now_ms;
128            return Ok(AutoPeerUpdate {
129                peer_id: PeerId(id as u8),
130                evicted_key: None,
131            });
132        }
133
134        if self.peers.len() < N {
135            self.peers
136                .push(PeerInfo {
137                    public_key: key,
138                    pinned: false,
139                    route: None,
140                    last_seen_ms: now_ms,
141                })
142                .map_err(|_| CapacityError)?;
143            return Ok(AutoPeerUpdate {
144                peer_id: PeerId((self.peers.len() - 1) as u8),
145                evicted_key: None,
146            });
147        }
148
149        let Some((index, oldest)) = self
150            .peers
151            .iter()
152            .enumerate()
153            .filter(|(_, peer)| !peer.pinned)
154            .min_by_key(|(_, peer)| peer.last_seen_ms)
155        else {
156            return Err(CapacityError);
157        };
158
159        let evicted_key = oldest.public_key;
160        self.peers[index] = PeerInfo {
161            public_key: key,
162            pinned: false,
163            route: None,
164            last_seen_ms: now_ms,
165        };
166        Ok(AutoPeerUpdate {
167            peer_id: PeerId(index as u8),
168            evicted_key: Some(evicted_key),
169        })
170    }
171
172    /// Update the cached route for `id`.
173    pub fn update_route(&mut self, id: PeerId, route: CachedRoute) {
174        if let Some(peer) = self.get_mut(id) {
175            peer.route = Some(route);
176        }
177    }
178
179    /// Refresh the last-seen timestamp for `id`.
180    pub fn touch(&mut self, id: PeerId, now_ms: u64) {
181        if let Some(peer) = self.get_mut(id) {
182            peer.last_seen_ms = now_ms;
183        }
184    }
185}
186
187/// Per-peer secure transport state.
188#[derive(Clone)]
189pub struct PeerCryptoState {
190    /// Pairwise encryption and MIC keys.
191    pub pairwise_keys: PairwiseKeys,
192    /// Replay state for traffic from this peer.
193    pub replay_window: ReplayWindow,
194}
195
196/// Fixed-capacity map of per-peer secure transport state.
197#[derive(Clone)]
198pub struct PeerCryptoMap<const N: usize> {
199    entries: LinearMap<PeerId, PeerCryptoState, N>,
200}
201
202impl<const N: usize> Default for PeerCryptoMap<N> {
203    fn default() -> Self {
204        Self::new()
205    }
206}
207
208impl<const N: usize> PeerCryptoMap<N> {
209    /// Create an empty peer-crypto map.
210    pub fn new() -> Self {
211        Self {
212            entries: LinearMap::new(),
213        }
214    }
215
216    /// Borrow one peer state.
217    pub fn get(&self, id: &PeerId) -> Option<&PeerCryptoState> {
218        self.entries.get(id)
219    }
220
221    /// Mutably borrow one peer state.
222    pub fn get_mut(&mut self, id: &PeerId) -> Option<&mut PeerCryptoState> {
223        self.entries.get_mut(id)
224    }
225
226    /// Insert or replace state for a peer.
227    pub fn insert(
228        &mut self,
229        id: PeerId,
230        state: PeerCryptoState,
231    ) -> Result<Option<PeerCryptoState>, CapacityError> {
232        self.entries.insert(id, state).map_err(|_| CapacityError)
233    }
234
235    /// Remove state for a peer.
236    pub fn remove(&mut self, id: &PeerId) -> Option<PeerCryptoState> {
237        self.entries.remove(id)
238    }
239}
240
241/// Replay state for a sender known only by hint.
242#[derive(Clone)]
243pub struct HintReplayState {
244    /// Replay window for the hint-only sender.
245    pub window: ReplayWindow,
246    /// Most recent observation timestamp.
247    pub last_seen_ms: u64,
248}
249
250/// Shared state for one multicast channel.
251#[derive(Clone)]
252pub struct ChannelState<const RN: usize = 8, const HN: usize = 8> {
253    /// Raw channel key.
254    pub channel_key: ChannelKey,
255    /// Derived transport keys and identifier.
256    pub derived: DerivedChannelKeys,
257    /// Replay windows for peers resolved to full identities.
258    pub replay: LinearMap<PeerId, ReplayWindow, RN>,
259    /// Replay windows for senders known only by hint.
260    pub hint_replay: LinearMap<NodeHint, HintReplayState, HN>,
261}
262
263impl<const RN: usize, const HN: usize> ChannelState<RN, HN> {
264    /// Create a new channel-state record.
265    pub fn new(channel_key: ChannelKey, derived: DerivedChannelKeys) -> Self {
266        Self {
267            channel_key,
268            derived,
269            replay: LinearMap::new(),
270            hint_replay: LinearMap::new(),
271        }
272    }
273}
274
275/// Fixed-capacity channel table shared by the MAC coordinator.
276#[derive(Clone)]
277pub struct ChannelTable<const N: usize, const RN: usize = 8, const HN: usize = 8> {
278    channels: Vec<ChannelState<RN, HN>, N>,
279}
280
281impl<const N: usize, const RN: usize, const HN: usize> Default for ChannelTable<N, RN, HN> {
282    fn default() -> Self {
283        Self::new()
284    }
285}
286
287impl<const N: usize, const RN: usize, const HN: usize> ChannelTable<N, RN, HN> {
288    /// Create an empty channel table.
289    pub fn new() -> Self {
290        Self {
291            channels: Vec::new(),
292        }
293    }
294
295    /// Return the number of configured channels.
296    pub fn len(&self) -> usize {
297        self.channels.len()
298    }
299
300    /// Return whether no channels are configured.
301    pub fn is_empty(&self) -> bool {
302        self.channels.is_empty()
303    }
304
305    /// Iterate over channels whose derived identifier matches `id`.
306    pub fn lookup_by_id(&self, id: &ChannelId) -> impl Iterator<Item = &ChannelState<RN, HN>> {
307        self.channels
308            .iter()
309            .filter(move |channel| channel.derived.channel_id == *id)
310    }
311
312    /// Mutably borrow the first channel whose derived identifier matches `id`.
313    pub fn get_mut_by_id(&mut self, id: &ChannelId) -> Option<&mut ChannelState<RN, HN>> {
314        self.channels
315            .iter_mut()
316            .find(|channel| channel.derived.channel_id == *id)
317    }
318
319    /// Mutably iterate over all channel states.
320    pub fn iter_mut(&mut self) -> impl Iterator<Item = &mut ChannelState<RN, HN>> {
321        self.channels.iter_mut()
322    }
323
324    /// Add or replace a channel entry.
325    pub fn try_add(
326        &mut self,
327        key: ChannelKey,
328        derived: DerivedChannelKeys,
329    ) -> Result<(), CapacityError> {
330        if let Some(channel) = self.get_mut_by_id(&derived.channel_id) {
331            channel.channel_key = key;
332            channel.derived = derived;
333            return Ok(());
334        }
335
336        self.channels
337            .push(ChannelState::new(key, derived))
338            .map_err(|_| CapacityError)
339    }
340}
umsh_mac/peers.rs

umsh_mac/
peers.rs
