1use std::collections::HashMap;
4use std::net::IpAddr;
5
6use super::{E2ECheckStatus, E2EKey, E2EProfile, E2EState, Error, e2e_check, e2e_protect};
7
8#[derive(Debug)]
22pub struct E2ERegistry {
23 configs: HashMap<E2EKey, E2EProfile>,
25 rx_states: HashMap<(IpAddr, E2EKey), E2EState>,
27 tx_states: HashMap<E2EKey, E2EState>,
29}
30
31impl E2ERegistry {
32 #[must_use]
34 pub fn new() -> Self {
35 Self {
36 configs: HashMap::new(),
37 rx_states: HashMap::new(),
38 tx_states: HashMap::new(),
39 }
40 }
41
42 pub fn register(&mut self, key: E2EKey, profile: E2EProfile) {
45 self.tx_states.insert(key, E2EState::from_profile(&profile));
46 self.rx_states.retain(|(_, k), _| *k != key);
47 self.configs.insert(key, profile);
48 }
49
50 pub fn unregister(&mut self, key: &E2EKey) {
52 self.configs.remove(key);
53 self.tx_states.remove(key);
54 self.rx_states.retain(|(_, k), _| k != key);
55 }
56
57 #[must_use]
59 pub fn contains_key(&self, key: &E2EKey) -> bool {
60 self.configs.contains_key(key)
61 }
62
63 pub fn check<'a>(
70 &mut self,
71 source: IpAddr,
72 key: E2EKey,
73 payload: &'a [u8],
74 upper_header: [u8; 8],
75 ) -> Option<(E2ECheckStatus, &'a [u8])> {
76 let profile = self.configs.get(&key)?;
77 let state = self
78 .rx_states
79 .entry((source, key))
80 .or_insert_with(|| E2EState::from_profile(profile));
81 Some(e2e_check(profile, state, payload, upper_header))
82 }
83
84 pub fn protect(
88 &mut self,
89 key: E2EKey,
90 payload: &[u8],
91 upper_header: [u8; 8],
92 output: &mut [u8],
93 ) -> Option<Result<usize, Error>> {
94 let profile = self.configs.get(&key)?;
95 let state = self.tx_states.get_mut(&key)?;
96 Some(e2e_protect(profile, state, payload, upper_header, output))
97 }
98
99 pub fn reset_source(&mut self, source: IpAddr) {
103 self.rx_states.retain(|(s, _), _| *s != source);
104 }
105}
106
107impl Default for E2ERegistry {
108 fn default() -> Self {
109 Self::new()
110 }
111}
112
113#[cfg(test)]
114mod tests {
115 use super::*;
116 use crate::e2e::{Profile4Config, Profile5Config};
117 use std::net::Ipv4Addr;
118
119 fn make_key() -> E2EKey {
120 E2EKey::new(0x1234, 0x5678)
121 }
122
123 fn src() -> IpAddr {
124 IpAddr::V4(Ipv4Addr::LOCALHOST)
125 }
126
127 fn make_profile5() -> E2EProfile {
128 E2EProfile::Profile5(Profile5Config::new(0x1234, 20, 15))
129 }
130
131 fn protect_next(sender: &mut E2ERegistry, key: E2EKey, out: &mut [u8; 64]) -> usize {
135 let mut payload = [0u8; 20];
136 payload[..5].copy_from_slice(b"Hello");
137 sender.protect(key, &payload, [0; 8], out).unwrap().unwrap()
138 }
139
140 #[test]
141 fn register_and_check_profile4() {
142 let mut reg = E2ERegistry::new();
143 let key = make_key();
144 let config = Profile4Config::new(0x12345678, 15);
145 reg.register(key, E2EProfile::Profile4(config.clone()));
146 assert!(reg.contains_key(&key));
147
148 let payload = b"Hello";
150 let mut out = [0u8; 64];
151 let len = reg
152 .protect(key, payload, [0; 8], &mut out)
153 .unwrap()
154 .unwrap();
155
156 let (status, stripped) = reg.check(src(), key, &out[..len], [0; 8]).unwrap();
158 assert_eq!(status, E2ECheckStatus::Ok);
159 assert_eq!(stripped, payload);
160 }
161
162 #[test]
163 fn register_and_check_profile5() {
164 let mut reg = E2ERegistry::new();
165 let key = make_key();
166 reg.register(key, make_profile5());
167
168 let mut payload = [0u8; 20];
169 payload[..5].copy_from_slice(b"Hello");
170 let mut out = [0u8; 64];
171 let len = reg
172 .protect(key, &payload, [0; 8], &mut out)
173 .unwrap()
174 .unwrap();
175
176 let (status, stripped) = reg.check(src(), key, &out[..len], [0; 8]).unwrap();
177 assert_eq!(status, E2ECheckStatus::Ok);
178 assert_eq!(stripped, &payload);
179 }
180
181 #[test]
182 fn distinct_sources_have_independent_e2e_state() {
183 let a = IpAddr::V4(Ipv4Addr::new(192, 168, 11, 101));
184 let b = IpAddr::V4(Ipv4Addr::new(192, 168, 11, 102));
185 let key = make_key();
186
187 let mut sender = E2ERegistry::new();
189 sender.register(key, make_profile5());
190 let mut b0 = [0u8; 64];
191 let l0 = protect_next(&mut sender, key, &mut b0);
192 let mut b1 = [0u8; 64];
193 let l1 = protect_next(&mut sender, key, &mut b1);
194
195 let mut recv = E2ERegistry::new();
196 recv.register(key, make_profile5());
197
198 assert_eq!(
200 recv.check(a, key, &b0[..l0], [0; 8]).unwrap().0,
201 E2ECheckStatus::Ok
202 );
203 assert_eq!(
204 recv.check(a, key, &b1[..l1], [0; 8]).unwrap().0,
205 E2ECheckStatus::Ok
206 );
207 assert_eq!(
211 recv.check(b, key, &b0[..l0], [0; 8]).unwrap().0,
212 E2ECheckStatus::Ok,
213 "source B's receive counter must be independent of source A's"
214 );
215 assert_eq!(
216 recv.check(b, key, &b1[..l1], [0; 8]).unwrap().0,
217 E2ECheckStatus::Ok
218 );
219 }
220
221 #[test]
222 fn reset_source_clears_only_that_source() {
223 let a = IpAddr::V4(Ipv4Addr::new(192, 168, 11, 101));
224 let key = make_key();
225
226 let mut sender = E2ERegistry::new();
227 sender.register(key, make_profile5());
228 let mut b0 = [0u8; 64];
229 let l0 = protect_next(&mut sender, key, &mut b0);
230 let mut b1 = [0u8; 64];
231 let l1 = protect_next(&mut sender, key, &mut b1);
232
233 let mut recv = E2ERegistry::new();
234 recv.register(key, make_profile5());
235 recv.check(a, key, &b0[..l0], [0; 8]);
236 recv.check(a, key, &b1[..l1], [0; 8]);
237
238 recv.reset_source(a);
241 assert_eq!(
242 recv.check(a, key, &b0[..l0], [0; 8]).unwrap().0,
243 E2ECheckStatus::Ok,
244 "reset_source(a) restarts A's receive counter sequence"
245 );
246 }
247
248 #[test]
249 fn unregistered_key_returns_none() {
250 let mut reg = E2ERegistry::new();
251 let key = make_key();
252 assert!(!reg.contains_key(&key));
253 assert!(reg.check(src(), key, b"test", [0; 8]).is_none());
254 assert!(reg.protect(key, b"test", [0; 8], &mut [0; 64]).is_none());
255 }
256
257 #[test]
258 fn unregister_removes_key() {
259 let mut reg = E2ERegistry::new();
260 let key = make_key();
261 reg.register(key, E2EProfile::Profile4(Profile4Config::new(0, 15)));
262 assert!(reg.contains_key(&key));
263 reg.unregister(&key);
264 assert!(!reg.contains_key(&key));
265 }
266
267 #[test]
268 fn default_is_empty() {
269 let reg = E2ERegistry::default();
270 assert!(!reg.contains_key(&make_key()));
271 }
272}