openipc_core/
diversity.rs

1//! Packet-selection diversity for multiple independent receive radios.
2//!
3//! Radios tuned to the same channel often receive the same WFB fragment. The
4//! combiner forwards the first valid copy immediately and rejects later copies
5//! before session decryption and FEC assembly. It intentionally does not wait
6//! for a stronger copy, so enabling diversity adds no comparison window to the
7//! media path.
8
9use std::collections::{BTreeMap, HashMap, HashSet, VecDeque};
10
11use crate::{
12    ieee80211::{FrameLayout, WifiFrame},
13    wfb::{parse_forwarder_packet, WfbPacket, CRYPTO_BOX_NONCE_LEN},
14    ChannelId,
15};
16
17/// Stable index assigned to one receive radio for the lifetime of a receiver.
18#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
19pub struct DiversitySourceId(u16);
20
21impl DiversitySourceId {
22    /// Construct a source id from an application-assigned index.
23    pub const fn new(value: u16) -> Self {
24        Self(value)
25    }
26
27    /// Return the application-assigned source index.
28    pub const fn get(self) -> u16 {
29        self.0
30    }
31}
32
33/// Result of examining one valid 802.11 receive frame.
34#[derive(Debug, Clone, Copy, PartialEq, Eq)]
35pub enum DiversityDecision {
36    /// First observed copy of a WFB session or data packet.
37    Accept,
38    /// A previously accepted radio already delivered this WFB packet.
39    Duplicate,
40    /// The frame is not identifiable as WFB and should continue normally.
41    Passthrough,
42}
43
44impl DiversityDecision {
45    /// Return true when the frame should enter the shared receiver pipeline.
46    pub const fn should_forward(self) -> bool {
47        !matches!(self, Self::Duplicate)
48    }
49}
50
51/// Cumulative packet-selection counters for one receive radio.
52#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
53pub struct DiversitySourceStats {
54    /// Valid frames examined from this source.
55    pub observed: u64,
56    /// First WFB copies contributed by this source.
57    pub accepted: u64,
58    /// Copies discarded because another source arrived first.
59    pub duplicates: u64,
60    /// Frames that could not be classified as a WFB packet.
61    pub passthrough: u64,
62}
63
64/// Cumulative state of the packet diversity combiner.
65#[derive(Debug, Clone, Default, PartialEq, Eq)]
66pub struct DiversityStats {
67    /// First WFB copies forwarded to the protocol pipeline.
68    pub accepted: u64,
69    /// Duplicate WFB copies discarded before decryption.
70    pub duplicates: u64,
71    /// Unclassified frames forwarded without deduplication.
72    pub passthrough: u64,
73    /// Current number of packet identities in the bounded deduplication cache.
74    pub cached_packets: usize,
75    /// Per-radio contribution counters.
76    pub sources: BTreeMap<DiversitySourceId, DiversitySourceStats>,
77}
78
79#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
80enum PacketIdentity {
81    Data {
82        channel_id: ChannelId,
83        session_generation: u64,
84        data_nonce: u64,
85    },
86    Session {
87        channel_id: ChannelId,
88        nonce: [u8; CRYPTO_BOX_NONCE_LEN],
89    },
90}
91
92/// Bounded first-valid-copy selector shared by all receive adapters.
93#[derive(Debug, Clone)]
94pub struct DiversityCombiner {
95    capacity: usize,
96    seen: HashSet<PacketIdentity>,
97    insertion_order: VecDeque<PacketIdentity>,
98    session_generations: HashMap<ChannelId, u64>,
99    current_sessions: HashMap<ChannelId, [u8; CRYPTO_BOX_NONCE_LEN]>,
100    session_order: VecDeque<ChannelId>,
101    stats: DiversityStats,
102}
103
104impl Default for DiversityCombiner {
105    fn default() -> Self {
106        Self::new(8_192)
107    }
108}
109
110impl DiversityCombiner {
111    /// Create a combiner retaining at most `capacity` recent WFB identities.
112    pub fn new(capacity: usize) -> Self {
113        Self {
114            capacity: capacity.max(1),
115            seen: HashSet::with_capacity(capacity.min(8_192)),
116            insertion_order: VecDeque::with_capacity(capacity.min(8_192)),
117            session_generations: HashMap::new(),
118            current_sessions: HashMap::new(),
119            session_order: VecDeque::new(),
120            stats: DiversityStats::default(),
121        }
122    }
123
124    /// Examine one CRC/ICV-valid frame and select the first WFB copy.
125    ///
126    /// Corrupt descriptor packets must be rejected by the caller before this
127    /// method. Unrecognized frames are passed through so diversity does not
128    /// change route-manager behavior for future payload formats.
129    pub fn observe_frame(
130        &mut self,
131        source: DiversitySourceId,
132        frame: &[u8],
133        layout: FrameLayout,
134    ) -> DiversityDecision {
135        let Ok(frame) = WifiFrame::parse(frame, layout) else {
136            self.source_mut(source).observed += 1;
137            self.stats.passthrough += 1;
138            self.source_mut(source).passthrough += 1;
139            return DiversityDecision::Passthrough;
140        };
141        self.observe_wifi_frame(source, frame)
142    }
143
144    /// Examine an already-validated WiFi frame without parsing it again.
145    pub fn observe_wifi_frame(
146        &mut self,
147        source: DiversitySourceId,
148        frame: WifiFrame<'_>,
149    ) -> DiversityDecision {
150        self.source_mut(source).observed += 1;
151        let Some(identity) = self.packet_identity(frame) else {
152            self.stats.passthrough += 1;
153            self.source_mut(source).passthrough += 1;
154            return DiversityDecision::Passthrough;
155        };
156
157        if self.seen.contains(&identity) {
158            self.stats.duplicates += 1;
159            self.source_mut(source).duplicates += 1;
160            return DiversityDecision::Duplicate;
161        }
162
163        if let PacketIdentity::Session { channel_id, nonce } = identity {
164            let changed = self
165                .current_sessions
166                .get(&channel_id)
167                .map(|current| current != &nonce)
168                .unwrap_or(true);
169            if changed {
170                if !self.current_sessions.contains_key(&channel_id) {
171                    self.remember_session_channel(channel_id);
172                }
173                let generation = self.session_generations.entry(channel_id).or_default();
174                *generation = generation.wrapping_add(1);
175                self.current_sessions.insert(channel_id, nonce);
176            }
177        }
178
179        self.remember(identity);
180        self.stats.accepted += 1;
181        self.source_mut(source).accepted += 1;
182        DiversityDecision::Accept
183    }
184
185    /// Return a snapshot of cumulative diversity counters.
186    pub fn stats(&self) -> DiversityStats {
187        let mut stats = self.stats.clone();
188        stats.cached_packets = self.seen.len();
189        stats
190    }
191
192    /// Clear packet identities and counters while keeping the configured size.
193    pub fn reset(&mut self) {
194        self.seen.clear();
195        self.insertion_order.clear();
196        self.session_generations.clear();
197        self.current_sessions.clear();
198        self.session_order.clear();
199        self.stats = DiversityStats::default();
200    }
201
202    fn packet_identity(&self, frame: WifiFrame<'_>) -> Option<PacketIdentity> {
203        let channel_id = frame.channel_id()?;
204        match parse_forwarder_packet(frame.payload()).ok()? {
205            WfbPacket::Data { data_nonce, .. } => Some(PacketIdentity::Data {
206                channel_id,
207                session_generation: self
208                    .session_generations
209                    .get(&channel_id)
210                    .copied()
211                    .unwrap_or(0),
212                data_nonce,
213            }),
214            WfbPacket::SessionKey { session_nonce, .. } => {
215                let nonce = session_nonce.try_into().ok()?;
216                Some(PacketIdentity::Session { channel_id, nonce })
217            }
218        }
219    }
220
221    fn remember(&mut self, identity: PacketIdentity) {
222        if self.seen.insert(identity) {
223            self.insertion_order.push_back(identity);
224        }
225        while self.insertion_order.len() > self.capacity {
226            if let Some(expired) = self.insertion_order.pop_front() {
227                self.seen.remove(&expired);
228            }
229        }
230    }
231
232    fn source_mut(&mut self, source: DiversitySourceId) -> &mut DiversitySourceStats {
233        self.stats.sources.entry(source).or_default()
234    }
235
236    fn remember_session_channel(&mut self, channel_id: ChannelId) {
237        while self.session_order.len() >= self.capacity {
238            if let Some(expired) = self.session_order.pop_front() {
239                self.current_sessions.remove(&expired);
240                self.session_generations.remove(&expired);
241            }
242        }
243        self.session_order.push_back(channel_id);
244    }
245}
246
247#[cfg(test)]
248mod tests {
249    use super::*;
250    use crate::{
251        fec::FecCode,
252        ieee80211::build_wfb_header,
253        wfb::{PlainAssembler, CHACHA20_POLY1305_TAG_LEN, MAX_FEC_PAYLOAD, WPACKET_HDR_LEN},
254        PayloadPipeline, PayloadPipelineEvent, WfbKeypair, WfbTransmitter, WfbTxKeypair,
255    };
256    use crypto_box::SecretKey;
257
258    fn data_frame(channel: ChannelId, data_nonce: u64) -> Vec<u8> {
259        let mut frame = Vec::from(build_wfb_header(channel, [0, 0]));
260        frame.push(1);
261        frame.extend_from_slice(&data_nonce.to_be_bytes());
262        frame.resize(frame.len() + WPACKET_HDR_LEN + CHACHA20_POLY1305_TAG_LEN, 0);
263        frame.extend_from_slice(&[0; 4]);
264        frame
265    }
266
267    fn session_frame(channel: ChannelId, marker: u8) -> Vec<u8> {
268        let mut frame = Vec::from(build_wfb_header(channel, [0, 0]));
269        frame.push(2);
270        frame.extend_from_slice(&[marker; CRYPTO_BOX_NONCE_LEN]);
271        frame.resize(frame.len() + crate::wfb::WSESSION_DATA_LEN + 16, 0);
272        frame.extend_from_slice(&[0; 4]);
273        frame
274    }
275
276    fn plain(payload: &[u8]) -> Vec<u8> {
277        let mut out = vec![0];
278        out.extend_from_slice(&(payload.len() as u16).to_be_bytes());
279        out.extend_from_slice(payload);
280        out.resize(MAX_FEC_PAYLOAD, 0);
281        out
282    }
283
284    fn linked_keypairs() -> (WfbTxKeypair, WfbKeypair) {
285        let transmitter = SecretKey::from([3u8; 32]);
286        let receiver = SecretKey::from([9u8; 32]);
287        (
288            WfbTxKeypair {
289                tx_secretkey: transmitter.to_bytes(),
290                rx_publickey: receiver.public_key().to_bytes(),
291            },
292            WfbKeypair {
293                rx_secretkey: receiver.to_bytes(),
294                tx_publickey: transmitter.public_key().to_bytes(),
295            },
296        )
297    }
298
299    fn wrap_forwarder_packet(channel: ChannelId, packet: &[u8]) -> Vec<u8> {
300        let mut frame = Vec::from(build_wfb_header(channel, [0, 0]));
301        frame.extend_from_slice(packet);
302        frame.extend_from_slice(&[0; 4]);
303        frame
304    }
305
306    #[test]
307    fn first_valid_radio_wins_without_delaying_the_packet() {
308        let mut combiner = DiversityCombiner::default();
309        let frame = data_frame(ChannelId::default_video(), 42);
310
311        assert_eq!(
312            combiner.observe_frame(DiversitySourceId::new(1), &frame, FrameLayout::WithFcs),
313            DiversityDecision::Accept
314        );
315        assert_eq!(
316            combiner.observe_frame(DiversitySourceId::new(2), &frame, FrameLayout::WithFcs),
317            DiversityDecision::Duplicate
318        );
319        let stats = combiner.stats();
320        assert_eq!(stats.accepted, 1);
321        assert_eq!(stats.duplicates, 1);
322        assert_eq!(stats.sources[&DiversitySourceId::new(1)].accepted, 1);
323        assert_eq!(stats.sources[&DiversitySourceId::new(2)].duplicates, 1);
324    }
325
326    #[test]
327    fn a_new_session_can_reuse_data_nonces() {
328        let mut combiner = DiversityCombiner::default();
329        let channel = ChannelId::default_video();
330        let data = data_frame(channel, 7);
331
332        assert_eq!(
333            combiner.observe_frame(
334                DiversitySourceId::new(0),
335                &session_frame(channel, 1),
336                FrameLayout::WithFcs,
337            ),
338            DiversityDecision::Accept
339        );
340        assert_eq!(
341            combiner.observe_frame(DiversitySourceId::new(0), &data, FrameLayout::WithFcs),
342            DiversityDecision::Accept
343        );
344        assert_eq!(
345            combiner.observe_frame(DiversitySourceId::new(1), &data, FrameLayout::WithFcs),
346            DiversityDecision::Duplicate
347        );
348        assert_eq!(
349            combiner.observe_frame(
350                DiversitySourceId::new(1),
351                &session_frame(channel, 2),
352                FrameLayout::WithFcs,
353            ),
354            DiversityDecision::Accept
355        );
356        assert_eq!(
357            combiner.observe_frame(DiversitySourceId::new(1), &data, FrameLayout::WithFcs),
358            DiversityDecision::Accept
359        );
360    }
361
362    #[test]
363    fn session_tracking_is_bounded_with_the_packet_cache() {
364        let mut combiner = DiversityCombiner::new(1);
365        let first = ChannelId::new(1);
366        let second = ChannelId::new(2);
367
368        combiner.observe_frame(
369            DiversitySourceId::new(0),
370            &session_frame(first, 1),
371            FrameLayout::WithFcs,
372        );
373        combiner.observe_frame(
374            DiversitySourceId::new(0),
375            &session_frame(second, 2),
376            FrameLayout::WithFcs,
377        );
378
379        assert_eq!(combiner.current_sessions.len(), 1);
380        assert!(!combiner.current_sessions.contains_key(&first));
381        assert!(combiner.current_sessions.contains_key(&second));
382    }
383
384    #[test]
385    fn fragments_from_two_radios_recover_one_shared_fec_block() {
386        let channel = ChannelId::default_video();
387        let primary = [plain(b"first"), plain(b"second"), plain(b"third")];
388        let parity = FecCode::new(3, 5)
389            .unwrap()
390            .encode(&primary, MAX_FEC_PAYLOAD)
391            .unwrap();
392        let arrivals = [
393            (DiversitySourceId::new(0), 0, primary[0].as_slice()),
394            (DiversitySourceId::new(1), 0, primary[0].as_slice()),
395            (DiversitySourceId::new(0), 2, primary[2].as_slice()),
396            (DiversitySourceId::new(1), 3, parity[0].as_slice()),
397        ];
398        let mut combiner = DiversityCombiner::default();
399        let mut assembler = PlainAssembler::new(3, 5).unwrap();
400        let mut output = Vec::new();
401
402        for (source, nonce, fragment) in arrivals {
403            let frame = data_frame(channel, nonce);
404            if combiner
405                .observe_frame(source, &frame, FrameLayout::WithFcs)
406                .should_forward()
407            {
408                output.extend(assembler.push_decrypted_fragment(nonce, fragment).unwrap());
409            }
410        }
411
412        assert_eq!(
413            output
414                .into_iter()
415                .map(|packet| packet.payload)
416                .collect::<Vec<_>>(),
417            vec![b"first".to_vec(), b"second".to_vec(), b"third".to_vec()]
418        );
419        assert_eq!(combiner.stats().duplicates, 1);
420        assert_eq!(assembler.recovered_packets, 1);
421    }
422
423    #[test]
424    fn encrypted_fragments_from_two_radios_share_one_pipeline() {
425        let channel = ChannelId::default_video();
426        let (tx_keys, rx_keys) = linked_keypairs();
427        let mut transmitter = WfbTransmitter::new(channel, tx_keys, 42, 2, 3).unwrap();
428        let mut pipeline =
429            PayloadPipeline::with_keypair(channel, FrameLayout::WithFcs, rx_keys, 0).unwrap();
430        let mut combiner = DiversityCombiner::default();
431
432        let session = wrap_forwarder_packet(channel, transmitter.session_forwarder_packet());
433        assert!(combiner
434            .observe_frame(DiversitySourceId::new(0), &session, FrameLayout::WithFcs)
435            .should_forward());
436        let events = pipeline.push_80211_frame(&session).unwrap();
437        assert!(matches!(
438            events.as_slice(),
439            [PayloadPipelineEvent::SessionEstablished {
440                epoch: 42,
441                fec_k: 2,
442                fec_n: 3
443            }]
444        ));
445
446        let missing_primary = transmitter
447            .forwarder_packets_for_payload(b"first", 0)
448            .unwrap();
449        assert_eq!(missing_primary.len(), 1);
450        let second_and_parity = transmitter
451            .forwarder_packets_for_payload(b"second", 0)
452            .unwrap();
453        assert_eq!(second_and_parity.len(), 2);
454        let second = wrap_forwarder_packet(channel, &second_and_parity[0]);
455        let parity = wrap_forwarder_packet(channel, &second_and_parity[1]);
456
457        let arrivals = [
458            (DiversitySourceId::new(0), second.as_slice()),
459            (DiversitySourceId::new(1), second.as_slice()),
460            (DiversitySourceId::new(1), parity.as_slice()),
461        ];
462        let mut payloads = Vec::new();
463        for (source, frame) in arrivals {
464            if combiner
465                .observe_frame(source, frame, FrameLayout::WithFcs)
466                .should_forward()
467            {
468                for event in pipeline.push_80211_frame(frame).unwrap() {
469                    if let PayloadPipelineEvent::Payload(payload) = event {
470                        payloads.push(payload.data);
471                    }
472                }
473            }
474        }
475
476        assert_eq!(payloads, [b"first".to_vec(), b"second".to_vec()]);
477        let stats = combiner.stats();
478        assert_eq!(stats.duplicates, 1);
479        assert_eq!(stats.sources[&DiversitySourceId::new(1)].accepted, 1);
480    }
481}
openipc_core/diversity.rs

openipc_core/
diversity.rs
