1use core::net::IpAddr;
11
12use heapless::index_map::{Entry, FnvIndexMap};
13
14use super::{E2ECheckStatus, E2EKey, E2EProfile, E2EState, Error, e2e_check, e2e_protect};
15
16pub const E2E_REGISTRY_CAP: usize = 32;
23
24const _: () = assert!(
25 E2E_REGISTRY_CAP.is_power_of_two(),
26 "E2E_REGISTRY_CAP must be a power of two for heapless::FnvIndexMap"
27);
28
29pub const E2E_RX_STATE_CAP: usize = 64;
44
45const _: () = assert!(
46 E2E_RX_STATE_CAP.is_power_of_two(),
47 "E2E_RX_STATE_CAP must be a power of two for heapless::FnvIndexMap"
48);
49
50#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
57#[error("e2e registry at capacity ({0})")]
58pub struct E2ERegistryFull(pub usize);
59
60#[derive(Debug)]
80pub struct E2ERegistry {
81 configs: FnvIndexMap<E2EKey, E2EProfile, E2E_REGISTRY_CAP>,
83 rx_states: FnvIndexMap<(IpAddr, E2EKey), E2EState, E2E_RX_STATE_CAP>,
85 tx_states: FnvIndexMap<E2EKey, E2EState, E2E_REGISTRY_CAP>,
87 rx_saturation_warned: bool,
90}
91
92impl E2ERegistry {
93 #[must_use]
96 pub const fn new() -> Self {
97 Self {
98 configs: FnvIndexMap::new(),
99 rx_states: FnvIndexMap::new(),
100 tx_states: FnvIndexMap::new(),
101 rx_saturation_warned: false,
102 }
103 }
104
105 pub fn register(&mut self, key: E2EKey, profile: E2EProfile) -> Result<(), E2ERegistryFull> {
120 let state = E2EState::from_profile(&profile);
121 if self.configs.insert(key, profile).is_err() {
127 return Err(E2ERegistryFull(E2E_REGISTRY_CAP));
128 }
129 let _ = self.tx_states.insert(key, state);
130 self.rx_states.retain(|(_, k), _| *k != key);
133 Ok(())
134 }
135
136 pub fn unregister(&mut self, key: &E2EKey) {
138 self.configs.remove(key);
139 self.tx_states.remove(key);
140 self.rx_states.retain(|(_, k), _| k != key);
141 }
142
143 #[must_use]
145 pub fn contains_key(&self, key: &E2EKey) -> bool {
146 self.configs.contains_key(key)
147 }
148
149 pub fn check<'a>(
156 &mut self,
157 source: IpAddr,
158 key: E2EKey,
159 payload: &'a [u8],
160 upper_header: [u8; 8],
161 ) -> Option<(E2ECheckStatus, &'a [u8])> {
162 let profile = self.configs.get(&key)?;
163 match self.rx_states.entry((source, key)) {
169 Entry::Occupied(occupied) => {
170 let state = occupied.into_mut();
171 Some(e2e_check(profile, state, payload, upper_header))
172 }
173 Entry::Vacant(vacant) => match vacant.insert(E2EState::from_profile(profile)) {
174 Ok(state) => Some(e2e_check(profile, state, payload, upper_header)),
175 Err(_full) => {
176 if !self.rx_saturation_warned {
177 self.rx_saturation_warned = true;
178 crate::log::warn!(
179 "E2E rx_states at capacity ({}); source {} falls back to a \
180 transient counter — sequence continuity untracked until a slot frees",
181 E2E_RX_STATE_CAP,
182 source
183 );
184 }
185 let mut transient = E2EState::from_profile(profile);
186 Some(e2e_check(profile, &mut transient, payload, upper_header))
187 }
188 },
189 }
190 }
191
192 pub fn protect(
196 &mut self,
197 key: E2EKey,
198 payload: &[u8],
199 upper_header: [u8; 8],
200 output: &mut [u8],
201 ) -> Option<Result<usize, Error>> {
202 let profile = self.configs.get(&key)?;
203 let state = self.tx_states.get_mut(&key)?;
204 Some(e2e_protect(profile, state, payload, upper_header, output))
205 }
206
207 pub fn reset_source(&mut self, source: IpAddr) {
211 self.rx_states.retain(|(s, _), _| *s != source);
212 }
213}
214
215impl Default for E2ERegistry {
216 fn default() -> Self {
217 Self::new()
218 }
219}
220
221#[cfg(test)]
222mod tests {
223 use super::*;
224 use crate::e2e::{Profile4Config, Profile5Config};
225 use core::net::Ipv4Addr;
226
227 fn make_key() -> E2EKey {
228 E2EKey::new(0x1234, 0x5678)
229 }
230
231 fn src() -> IpAddr {
232 IpAddr::V4(Ipv4Addr::LOCALHOST)
233 }
234
235 fn make_profile5() -> E2EProfile {
236 E2EProfile::Profile5(Profile5Config::new(0x1234, 20, 15))
237 }
238
239 fn protect_next(sender: &mut E2ERegistry, key: E2EKey, out: &mut [u8; 64]) -> usize {
243 let mut payload = [0u8; 20];
244 payload[..5].copy_from_slice(b"Hello");
245 sender.protect(key, &payload, [0; 8], out).unwrap().unwrap()
246 }
247
248 #[test]
249 fn register_and_check_profile4() {
250 let mut reg = E2ERegistry::new();
251 let key = make_key();
252 let config = Profile4Config::new(0x12345678, 15);
253 reg.register(key, E2EProfile::Profile4(config.clone()))
254 .expect("register fits within E2E_REGISTRY_CAP");
255 assert!(reg.contains_key(&key));
256
257 let payload = b"Hello";
259 let mut out = [0u8; 64];
260 let len = reg
261 .protect(key, payload, [0; 8], &mut out)
262 .unwrap()
263 .unwrap();
264
265 let (status, stripped) = reg.check(src(), key, &out[..len], [0; 8]).unwrap();
267 assert_eq!(status, E2ECheckStatus::Ok);
268 assert_eq!(stripped, payload);
269 }
270
271 #[test]
272 fn register_and_check_profile5() {
273 let mut reg = E2ERegistry::new();
274 let key = make_key();
275 reg.register(key, make_profile5())
276 .expect("register fits within E2E_REGISTRY_CAP");
277
278 let mut payload = [0u8; 20];
279 payload[..5].copy_from_slice(b"Hello");
280 let mut out = [0u8; 64];
281 let len = reg
282 .protect(key, &payload, [0; 8], &mut out)
283 .unwrap()
284 .unwrap();
285
286 let (status, stripped) = reg.check(src(), key, &out[..len], [0; 8]).unwrap();
287 assert_eq!(status, E2ECheckStatus::Ok);
288 assert_eq!(stripped, &payload);
289 }
290
291 #[test]
292 fn distinct_sources_have_independent_e2e_state() {
293 let a = IpAddr::V4(Ipv4Addr::new(192, 168, 11, 101));
294 let b = IpAddr::V4(Ipv4Addr::new(192, 168, 11, 102));
295 let key = make_key();
296
297 let mut sender = E2ERegistry::new();
299 sender
300 .register(key, make_profile5())
301 .expect("register fits within E2E_REGISTRY_CAP");
302 let mut b0 = [0u8; 64];
303 let l0 = protect_next(&mut sender, key, &mut b0);
304 let mut b1 = [0u8; 64];
305 let l1 = protect_next(&mut sender, key, &mut b1);
306
307 let mut recv = E2ERegistry::new();
308 recv.register(key, make_profile5())
309 .expect("register fits within E2E_REGISTRY_CAP");
310
311 assert_eq!(
313 recv.check(a, key, &b0[..l0], [0; 8]).unwrap().0,
314 E2ECheckStatus::Ok
315 );
316 assert_eq!(
317 recv.check(a, key, &b1[..l1], [0; 8]).unwrap().0,
318 E2ECheckStatus::Ok
319 );
320 assert_eq!(
324 recv.check(b, key, &b0[..l0], [0; 8]).unwrap().0,
325 E2ECheckStatus::Ok,
326 "source B's receive counter must be independent of source A's"
327 );
328 assert_eq!(
329 recv.check(b, key, &b1[..l1], [0; 8]).unwrap().0,
330 E2ECheckStatus::Ok
331 );
332 }
333
334 #[test]
335 fn reset_source_clears_only_that_source() {
336 let a = IpAddr::V4(Ipv4Addr::new(192, 168, 11, 101));
337 let key = make_key();
338
339 let mut sender = E2ERegistry::new();
340 sender
341 .register(key, make_profile5())
342 .expect("register fits within E2E_REGISTRY_CAP");
343 let mut b0 = [0u8; 64];
344 let l0 = protect_next(&mut sender, key, &mut b0);
345 let mut b1 = [0u8; 64];
346 let l1 = protect_next(&mut sender, key, &mut b1);
347
348 let mut recv = E2ERegistry::new();
349 recv.register(key, make_profile5())
350 .expect("register fits within E2E_REGISTRY_CAP");
351 recv.check(a, key, &b0[..l0], [0; 8]);
352 recv.check(a, key, &b1[..l1], [0; 8]);
353
354 recv.reset_source(a);
357 assert_eq!(
358 recv.check(a, key, &b0[..l0], [0; 8]).unwrap().0,
359 E2ECheckStatus::Ok,
360 "reset_source(a) restarts A's receive counter sequence"
361 );
362 }
363
364 #[test]
365 fn unregistered_key_returns_none() {
366 let mut reg = E2ERegistry::new();
367 let key = make_key();
368 assert!(!reg.contains_key(&key));
369 assert!(reg.check(src(), key, b"test", [0; 8]).is_none());
370 assert!(reg.protect(key, b"test", [0; 8], &mut [0; 64]).is_none());
371 }
372
373 #[test]
374 fn unregister_removes_key() {
375 let mut reg = E2ERegistry::new();
376 let key = make_key();
377 reg.register(key, E2EProfile::Profile4(Profile4Config::new(0, 15)))
378 .expect("register fits within E2E_REGISTRY_CAP");
379 assert!(reg.contains_key(&key));
380 reg.unregister(&key);
381 assert!(!reg.contains_key(&key));
382 }
383
384 #[test]
385 fn default_is_empty() {
386 let reg = E2ERegistry::default();
387 assert!(!reg.contains_key(&make_key()));
388 }
389
390 #[test]
395 fn register_replacement_succeeds_when_full() {
396 let mut reg = E2ERegistry::new();
397 for i in 0..E2E_REGISTRY_CAP {
398 let key = E2EKey::new(0x1000 + u16::try_from(i).unwrap(), 0);
399 reg.register(key, E2EProfile::Profile4(Profile4Config::new(0, 15)))
400 .expect("filling to cap");
401 }
402 let key0 = E2EKey::new(0x1000, 0);
404 let result = reg.register(key0, E2EProfile::Profile4(Profile4Config::new(42, 15)));
405 assert!(
406 result.is_ok(),
407 "replacing an existing entry must succeed even at capacity"
408 );
409 }
410
411 #[test]
416 fn register_overflow_returns_err_and_does_not_mutate() {
417 let mut reg = E2ERegistry::new();
418 for i in 0..E2E_REGISTRY_CAP {
419 reg.register(
420 E2EKey::new(0x2000 + u16::try_from(i).unwrap(), 0),
421 E2EProfile::Profile4(Profile4Config::new(0, 15)),
422 )
423 .expect("filling to cap");
424 }
425 let overflow_key = E2EKey::new(0xFFFE, 0);
427 let err = reg
428 .register(
429 overflow_key,
430 E2EProfile::Profile4(Profile4Config::new(0, 15)),
431 )
432 .expect_err("registering the (cap+1)-th key must overflow");
433 assert_eq!(err, E2ERegistryFull(E2E_REGISTRY_CAP));
434 assert!(!reg.contains_key(&overflow_key));
436 }
437}