use core::net::IpAddr;
use heapless::index_map::{Entry, FnvIndexMap};
use super::{E2ECheckStatus, E2EKey, E2EProfile, E2EState, Error, e2e_check, e2e_protect};
pub const E2E_REGISTRY_CAP: usize = 32;
const _: () = assert!(
E2E_REGISTRY_CAP.is_power_of_two(),
"E2E_REGISTRY_CAP must be a power of two for heapless::FnvIndexMap"
);
pub const E2E_RX_STATE_CAP: usize = 64;
const _: () = assert!(
E2E_RX_STATE_CAP.is_power_of_two(),
"E2E_RX_STATE_CAP must be a power of two for heapless::FnvIndexMap"
);
#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
#[error("e2e registry at capacity ({0})")]
pub struct E2ERegistryFull(pub usize);
#[derive(Debug)]
pub struct E2ERegistry {
configs: FnvIndexMap<E2EKey, E2EProfile, E2E_REGISTRY_CAP>,
rx_states: FnvIndexMap<(IpAddr, E2EKey), E2EState, E2E_RX_STATE_CAP>,
tx_states: FnvIndexMap<E2EKey, E2EState, E2E_REGISTRY_CAP>,
rx_saturation_warned: bool,
}
impl E2ERegistry {
#[must_use]
pub const fn new() -> Self {
Self {
configs: FnvIndexMap::new(),
rx_states: FnvIndexMap::new(),
tx_states: FnvIndexMap::new(),
rx_saturation_warned: false,
}
}
pub fn register(&mut self, key: E2EKey, profile: E2EProfile) -> Result<(), E2ERegistryFull> {
let state = E2EState::from_profile(&profile);
if self.configs.insert(key, profile).is_err() {
return Err(E2ERegistryFull(E2E_REGISTRY_CAP));
}
let _ = self.tx_states.insert(key, state);
self.rx_states.retain(|(_, k), _| *k != key);
Ok(())
}
pub fn unregister(&mut self, key: &E2EKey) {
self.configs.remove(key);
self.tx_states.remove(key);
self.rx_states.retain(|(_, k), _| k != key);
}
#[must_use]
pub fn contains_key(&self, key: &E2EKey) -> bool {
self.configs.contains_key(key)
}
pub fn check<'a>(
&mut self,
source: IpAddr,
key: E2EKey,
payload: &'a [u8],
upper_header: [u8; 8],
) -> Option<(E2ECheckStatus, &'a [u8])> {
let profile = self.configs.get(&key)?;
match self.rx_states.entry((source, key)) {
Entry::Occupied(occupied) => {
let state = occupied.into_mut();
Some(e2e_check(profile, state, payload, upper_header))
}
Entry::Vacant(vacant) => match vacant.insert(E2EState::from_profile(profile)) {
Ok(state) => Some(e2e_check(profile, state, payload, upper_header)),
Err(_full) => {
if !self.rx_saturation_warned {
self.rx_saturation_warned = true;
crate::log::warn!(
"E2E rx_states at capacity ({}); source {} falls back to a \
transient counter — sequence continuity untracked until a slot frees",
E2E_RX_STATE_CAP,
source
);
}
let mut transient = E2EState::from_profile(profile);
Some(e2e_check(profile, &mut transient, payload, upper_header))
}
},
}
}
pub fn protect(
&mut self,
key: E2EKey,
payload: &[u8],
upper_header: [u8; 8],
output: &mut [u8],
) -> Option<Result<usize, Error>> {
let profile = self.configs.get(&key)?;
let state = self.tx_states.get_mut(&key)?;
Some(e2e_protect(profile, state, payload, upper_header, output))
}
pub fn reset_source(&mut self, source: IpAddr) {
self.rx_states.retain(|(s, _), _| *s != source);
}
}
impl Default for E2ERegistry {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::e2e::{Profile4Config, Profile5Config};
use core::net::Ipv4Addr;
fn make_key() -> E2EKey {
E2EKey::new(0x1234, 0x5678)
}
fn src() -> IpAddr {
IpAddr::V4(Ipv4Addr::LOCALHOST)
}
fn make_profile5() -> E2EProfile {
E2EProfile::Profile5(Profile5Config::new(0x1234, 20, 15))
}
fn protect_next(sender: &mut E2ERegistry, key: E2EKey, out: &mut [u8; 64]) -> usize {
let mut payload = [0u8; 20];
payload[..5].copy_from_slice(b"Hello");
sender.protect(key, &payload, [0; 8], out).unwrap().unwrap()
}
#[test]
fn register_and_check_profile4() {
let mut reg = E2ERegistry::new();
let key = make_key();
let config = Profile4Config::new(0x12345678, 15);
reg.register(key, E2EProfile::Profile4(config.clone()))
.expect("register fits within E2E_REGISTRY_CAP");
assert!(reg.contains_key(&key));
let payload = b"Hello";
let mut out = [0u8; 64];
let len = reg
.protect(key, payload, [0; 8], &mut out)
.unwrap()
.unwrap();
let (status, stripped) = reg.check(src(), key, &out[..len], [0; 8]).unwrap();
assert_eq!(status, E2ECheckStatus::Ok);
assert_eq!(stripped, payload);
}
#[test]
fn register_and_check_profile5() {
let mut reg = E2ERegistry::new();
let key = make_key();
reg.register(key, make_profile5())
.expect("register fits within E2E_REGISTRY_CAP");
let mut payload = [0u8; 20];
payload[..5].copy_from_slice(b"Hello");
let mut out = [0u8; 64];
let len = reg
.protect(key, &payload, [0; 8], &mut out)
.unwrap()
.unwrap();
let (status, stripped) = reg.check(src(), key, &out[..len], [0; 8]).unwrap();
assert_eq!(status, E2ECheckStatus::Ok);
assert_eq!(stripped, &payload);
}
#[test]
fn distinct_sources_have_independent_e2e_state() {
let a = IpAddr::V4(Ipv4Addr::new(192, 168, 11, 101));
let b = IpAddr::V4(Ipv4Addr::new(192, 168, 11, 102));
let key = make_key();
let mut sender = E2ERegistry::new();
sender
.register(key, make_profile5())
.expect("register fits within E2E_REGISTRY_CAP");
let mut b0 = [0u8; 64];
let l0 = protect_next(&mut sender, key, &mut b0);
let mut b1 = [0u8; 64];
let l1 = protect_next(&mut sender, key, &mut b1);
let mut recv = E2ERegistry::new();
recv.register(key, make_profile5())
.expect("register fits within E2E_REGISTRY_CAP");
assert_eq!(
recv.check(a, key, &b0[..l0], [0; 8]).unwrap().0,
E2ECheckStatus::Ok
);
assert_eq!(
recv.check(a, key, &b1[..l1], [0; 8]).unwrap().0,
E2ECheckStatus::Ok
);
assert_eq!(
recv.check(b, key, &b0[..l0], [0; 8]).unwrap().0,
E2ECheckStatus::Ok,
"source B's receive counter must be independent of source A's"
);
assert_eq!(
recv.check(b, key, &b1[..l1], [0; 8]).unwrap().0,
E2ECheckStatus::Ok
);
}
#[test]
fn reset_source_clears_only_that_source() {
let a = IpAddr::V4(Ipv4Addr::new(192, 168, 11, 101));
let key = make_key();
let mut sender = E2ERegistry::new();
sender
.register(key, make_profile5())
.expect("register fits within E2E_REGISTRY_CAP");
let mut b0 = [0u8; 64];
let l0 = protect_next(&mut sender, key, &mut b0);
let mut b1 = [0u8; 64];
let l1 = protect_next(&mut sender, key, &mut b1);
let mut recv = E2ERegistry::new();
recv.register(key, make_profile5())
.expect("register fits within E2E_REGISTRY_CAP");
recv.check(a, key, &b0[..l0], [0; 8]);
recv.check(a, key, &b1[..l1], [0; 8]);
recv.reset_source(a);
assert_eq!(
recv.check(a, key, &b0[..l0], [0; 8]).unwrap().0,
E2ECheckStatus::Ok,
"reset_source(a) restarts A's receive counter sequence"
);
}
#[test]
fn unregistered_key_returns_none() {
let mut reg = E2ERegistry::new();
let key = make_key();
assert!(!reg.contains_key(&key));
assert!(reg.check(src(), key, b"test", [0; 8]).is_none());
assert!(reg.protect(key, b"test", [0; 8], &mut [0; 64]).is_none());
}
#[test]
fn unregister_removes_key() {
let mut reg = E2ERegistry::new();
let key = make_key();
reg.register(key, E2EProfile::Profile4(Profile4Config::new(0, 15)))
.expect("register fits within E2E_REGISTRY_CAP");
assert!(reg.contains_key(&key));
reg.unregister(&key);
assert!(!reg.contains_key(&key));
}
#[test]
fn default_is_empty() {
let reg = E2ERegistry::default();
assert!(!reg.contains_key(&make_key()));
}
#[test]
fn register_replacement_succeeds_when_full() {
let mut reg = E2ERegistry::new();
for i in 0..E2E_REGISTRY_CAP {
let key = E2EKey::new(0x1000 + u16::try_from(i).unwrap(), 0);
reg.register(key, E2EProfile::Profile4(Profile4Config::new(0, 15)))
.expect("filling to cap");
}
let key0 = E2EKey::new(0x1000, 0);
let result = reg.register(key0, E2EProfile::Profile4(Profile4Config::new(42, 15)));
assert!(
result.is_ok(),
"replacing an existing entry must succeed even at capacity"
);
}
#[test]
fn register_overflow_returns_err_and_does_not_mutate() {
let mut reg = E2ERegistry::new();
for i in 0..E2E_REGISTRY_CAP {
reg.register(
E2EKey::new(0x2000 + u16::try_from(i).unwrap(), 0),
E2EProfile::Profile4(Profile4Config::new(0, 15)),
)
.expect("filling to cap");
}
let overflow_key = E2EKey::new(0xFFFE, 0);
let err = reg
.register(
overflow_key,
E2EProfile::Profile4(Profile4Config::new(0, 15)),
)
.expect_err("registering the (cap+1)-th key must overflow");
assert_eq!(err, E2ERegistryFull(E2E_REGISTRY_CAP));
assert!(!reg.contains_key(&overflow_key));
}
}