openipc_core/
wfb.rs

1use std::collections::BTreeMap;
2
3use crypto_box::aead::Aead;
4use crypto_box::{Nonce as BoxNonce, PublicKey, SalsaBox, SecretKey};
5
6use crate::channel::ChannelId;
7use crate::crypto::decrypt_chacha20poly1305_legacy;
8use crate::fec::FecCode;
9
10/// WFB WiFi MTU used by OpenIPC forwarder packets.
11pub const WIFI_MTU: usize = 4045;
12/// 802.11 header length subtracted from WFB packet capacity.
13pub const IEEE80211_HEADER_LEN: usize = 24;
14/// crypto_box secret key length.
15pub const CRYPTO_BOX_SECRETKEY_LEN: usize = 32;
16/// crypto_box public key length.
17pub const CRYPTO_BOX_PUBLICKEY_LEN: usize = 32;
18/// crypto_box nonce length.
19pub const CRYPTO_BOX_NONCE_LEN: usize = 24;
20/// crypto_box authentication tag length.
21pub const CRYPTO_BOX_TAG_LEN: usize = 16;
22/// WFB session packet header length.
23pub const WSESSION_HDR_LEN: usize = 1 + CRYPTO_BOX_NONCE_LEN;
24/// Plain WFB session body length before crypto_box encryption.
25pub const WSESSION_DATA_LEN: usize = 8 + 4 + 1 + 1 + 1 + CHACHA20_POLY1305_KEY_LEN;
26/// WFB data-block header length.
27pub const WBLOCK_HDR_LEN: usize = 9;
28/// Plain WFB payload-fragment header length.
29pub const WPACKET_HDR_LEN: usize = 3;
30/// WFB session ChaCha20-Poly1305 key length.
31pub const CHACHA20_POLY1305_KEY_LEN: usize = 32;
32/// WFB session ChaCha20-Poly1305 authentication tag length.
33pub const CHACHA20_POLY1305_TAG_LEN: usize = 16;
34/// Maximum encrypted FEC fragment payload carried by one WFB data packet.
35pub const MAX_FEC_PAYLOAD: usize =
36    WIFI_MTU - IEEE80211_HEADER_LEN - WBLOCK_HDR_LEN - CHACHA20_POLY1305_TAG_LEN;
37/// Maximum application payload before WFB fragment headers are added.
38pub const MAX_PAYLOAD_SIZE: usize = MAX_FEC_PAYLOAD - WPACKET_HDR_LEN;
39/// Maximum WFB forwarder packet payload after the 802.11 header.
40pub const MAX_FORWARDER_PACKET_SIZE: usize = WIFI_MTU - IEEE80211_HEADER_LEN;
41/// Largest WFB block index before a transmitter must rotate session keys.
42pub const MAX_BLOCK_IDX: u64 = (1u64 << 55) - 1;
43
44/// WFB packet type for encrypted data fragments.
45pub const WFB_PACKET_DATA: u8 = 0x01;
46/// WFB packet type for encrypted session-key packets.
47pub const WFB_PACKET_KEY: u8 = 0x02;
48/// FEC type used by WFB's Vandermonde Reed-Solomon blocks.
49pub const WFB_FEC_VDM_RS: u8 = 0x01;
50/// Flag marking a WFB packet as parity-only FEC data.
51pub const WFB_PACKET_FEC_ONLY: u8 = 0x01;
52
53/// Error returned while parsing, decrypting, or assembling WFB packets.
54#[derive(Debug, Clone, PartialEq, Eq)]
55pub enum WfbError {
56    /// Packet buffer is empty.
57    Empty,
58    /// Packet exceeds WFB forwarder size.
59    TooLong,
60    /// Data packet is too short.
61    ShortDataPacket,
62    /// Session packet is too short.
63    ShortSessionPacket,
64    /// WFB keypair is not the expected 64-byte file shape.
65    InvalidKeypair,
66    /// Session-key encryption failed.
67    SessionEncryptFailed,
68    /// Session-key decryption failed.
69    SessionDecryptFailed,
70    /// Data encryption failed.
71    DataEncryptFailed,
72    /// Data decryption failed.
73    DataDecryptFailed,
74    /// Session epoch was older than the configured minimum.
75    SessionEpochTooOld {
76        /// Epoch from the received session packet.
77        session_epoch: u64,
78        /// Minimum epoch accepted by the receiver.
79        minimum_epoch: u64,
80    },
81    /// Session packet was for a different WFB channel.
82    SessionChannelMismatch {
83        /// Expected channel id.
84        expected: u32,
85        /// Actual channel id in the session packet.
86        actual: u32,
87    },
88    /// FEC type is not the supported VDM Reed-Solomon mode.
89    UnsupportedFecType(u8),
90    /// Forwarder packet type is unknown.
91    UnknownPacketType(u8),
92    /// FEC parameters are invalid.
93    InvalidFecParameters,
94    /// Fragment index is outside the current FEC block.
95    InvalidFragmentIndex,
96    /// Data nonce encoded a block index beyond the supported range.
97    BlockIndexOverflow,
98    /// Decrypted plain packet is malformed.
99    InvalidPlainPacket,
100    /// Plain payload exceeds the WFB maximum.
101    PayloadTooLarge,
102    /// Encrypted data packet arrived before a session key.
103    MissingSession,
104    /// FEC recovery failed.
105    FecRecoveryFailed,
106}
107
108impl std::fmt::Display for WfbError {
109    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
110        match self {
111            Self::Empty => write!(f, "empty WFB packet"),
112            Self::TooLong => write!(f, "WFB packet exceeds maximum forwarder size"),
113            Self::ShortDataPacket => write!(f, "short WFB data packet"),
114            Self::ShortSessionPacket => write!(f, "invalid WFB session packet size"),
115            Self::InvalidKeypair => write!(f, "WFB keypair must be 64 bytes"),
116            Self::SessionEncryptFailed => write!(f, "unable to encrypt WFB session key"),
117            Self::SessionDecryptFailed => write!(f, "unable to decrypt WFB session key"),
118            Self::DataEncryptFailed => write!(f, "unable to encrypt WFB data packet"),
119            Self::DataDecryptFailed => write!(f, "unable to decrypt WFB data packet"),
120            Self::SessionEpochTooOld {
121                session_epoch,
122                minimum_epoch,
123            } => write!(
124                f,
125                "WFB session epoch {session_epoch} is older than minimum {minimum_epoch}"
126            ),
127            Self::SessionChannelMismatch { expected, actual } => write!(
128                f,
129                "WFB session channel mismatch: expected 0x{expected:08x}, got 0x{actual:08x}"
130            ),
131            Self::UnsupportedFecType(fec_type) => {
132                write!(f, "unsupported WFB FEC type {fec_type}")
133            }
134            Self::UnknownPacketType(packet_type) => {
135                write!(f, "unknown WFB packet type 0x{packet_type:02x}")
136            }
137            Self::InvalidFecParameters => write!(f, "invalid WFB FEC parameters"),
138            Self::InvalidFragmentIndex => write!(f, "invalid WFB fragment index"),
139            Self::BlockIndexOverflow => write!(f, "WFB block index overflow"),
140            Self::InvalidPlainPacket => write!(f, "invalid decrypted WFB packet"),
141            Self::PayloadTooLarge => write!(f, "decrypted WFB payload is too large"),
142            Self::MissingSession => write!(f, "WFB data packet arrived before session key"),
143            Self::FecRecoveryFailed => write!(f, "WFB FEC recovery failed"),
144        }
145    }
146}
147
148impl std::error::Error for WfbError {}
149
150/// Borrowed WFB forwarder packet.
151#[derive(Debug, Clone, Copy, PartialEq, Eq)]
152pub enum WfbPacket<'a> {
153    /// Encrypted WFB data/FEC fragment packet.
154    Data {
155        /// Data nonce; high bits are block index and low byte is fragment index.
156        data_nonce: u64,
157        /// Encrypted fragment payload plus authentication tag.
158        encrypted_payload: &'a [u8],
159        /// Associated data used for WFB data authentication.
160        associated_data: &'a [u8],
161    },
162    /// Encrypted WFB session-key packet.
163    SessionKey {
164        /// crypto_box session nonce.
165        session_nonce: &'a [u8],
166        /// Encrypted session data.
167        encrypted_session: &'a [u8],
168    },
169}
170
171/// Ground-station WFB keypair file contents.
172#[derive(Debug, Clone, Copy, PartialEq, Eq)]
173pub struct WfbKeypair {
174    /// Ground-station receive secret key.
175    pub rx_secretkey: [u8; CRYPTO_BOX_SECRETKEY_LEN],
176    /// Air-unit transmit public key.
177    pub tx_publickey: [u8; CRYPTO_BOX_PUBLICKEY_LEN],
178}
179
180impl WfbKeypair {
181    /// Parse the 64-byte `gs.key` style keypair.
182    pub fn from_bytes(bytes: &[u8]) -> Result<Self, WfbError> {
183        if bytes.len() != CRYPTO_BOX_SECRETKEY_LEN + CRYPTO_BOX_PUBLICKEY_LEN {
184            return Err(WfbError::InvalidKeypair);
185        }
186        let mut rx_secretkey = [0; CRYPTO_BOX_SECRETKEY_LEN];
187        let mut tx_publickey = [0; CRYPTO_BOX_PUBLICKEY_LEN];
188        rx_secretkey.copy_from_slice(&bytes[..CRYPTO_BOX_SECRETKEY_LEN]);
189        tx_publickey.copy_from_slice(&bytes[CRYPTO_BOX_SECRETKEY_LEN..]);
190        Ok(Self {
191            rx_secretkey,
192            tx_publickey,
193        })
194    }
195}
196
197/// Cumulative FEC counters for a WFB receiver or assembler.
198#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
199pub struct FecCounters {
200    /// Total data fragments observed.
201    pub total_packets: u64,
202    /// Primary fragments recovered by FEC.
203    pub recovered_packets: u64,
204    /// Primary fragments considered lost.
205    pub lost_packets: u64,
206    /// Malformed or unrecoverable fragments.
207    pub bad_packets: u64,
208}
209
210/// Decrypted WFB session parameters.
211#[derive(Debug, Clone, PartialEq, Eq)]
212pub struct WfbSession {
213    /// Session epoch.
214    pub epoch: u64,
215    /// Channel id this session applies to.
216    pub channel_id: ChannelId,
217    /// WFB FEC type.
218    pub fec_type: u8,
219    /// Primary fragment count.
220    pub fec_k: usize,
221    /// Total primary plus parity fragment count.
222    pub fec_n: usize,
223    /// Symmetric key used for WFB data packets.
224    pub session_key: [u8; CHACHA20_POLY1305_KEY_LEN],
225}
226
227impl WfbSession {
228    fn parse(
229        plaintext: &[u8],
230        expected_channel_id: ChannelId,
231        minimum_epoch: u64,
232    ) -> Result<Self, WfbError> {
233        if plaintext.len() < WSESSION_DATA_LEN {
234            return Err(WfbError::SessionDecryptFailed);
235        }
236        let epoch = u64::from_be_bytes(plaintext[0..8].try_into().expect("checked length"));
237        if epoch < minimum_epoch {
238            return Err(WfbError::SessionEpochTooOld {
239                session_epoch: epoch,
240                minimum_epoch,
241            });
242        }
243
244        let raw_channel = u32::from_be_bytes(plaintext[8..12].try_into().expect("checked length"));
245        let channel_id = ChannelId::new(raw_channel);
246        if channel_id != expected_channel_id {
247            return Err(WfbError::SessionChannelMismatch {
248                expected: expected_channel_id.raw(),
249                actual: raw_channel,
250            });
251        }
252
253        let fec_type = plaintext[12];
254        if fec_type != WFB_FEC_VDM_RS {
255            return Err(WfbError::UnsupportedFecType(fec_type));
256        }
257        let fec_k = plaintext[13] as usize;
258        let fec_n = plaintext[14] as usize;
259        if fec_k == 0 || fec_n == 0 || fec_k > fec_n || fec_n >= 256 {
260            return Err(WfbError::InvalidFecParameters);
261        }
262
263        let mut session_key = [0; CHACHA20_POLY1305_KEY_LEN];
264        session_key.copy_from_slice(&plaintext[15..47]);
265        Ok(Self {
266            epoch,
267            channel_id,
268            fec_type,
269            fec_k,
270            fec_n,
271            session_key,
272        })
273    }
274}
275
276/// Event emitted by an encrypted WFB receiver.
277#[derive(Debug, Clone, PartialEq, Eq)]
278pub enum WfbEvent {
279    /// A session key was accepted.
280    Session(WfbSession),
281    /// One recovered payload was emitted.
282    Payload(WfbOutput),
283}
284
285/// Encrypted WFB receiver for one channel.
286#[derive(Debug, Clone)]
287pub struct WfbReceiver {
288    channel_id: ChannelId,
289    minimum_epoch: u64,
290    keypair: WfbKeypair,
291    session: Option<WfbSession>,
292    assembler: Option<PlainAssembler>,
293    incoming_packets: u64,
294    session_packets: u64,
295    data_packets: u64,
296}
297
298impl WfbReceiver {
299    /// Create a receiver for one channel and keypair.
300    pub fn new(channel_id: ChannelId, keypair: WfbKeypair, minimum_epoch: u64) -> Self {
301        Self {
302            channel_id,
303            minimum_epoch,
304            keypair,
305            session: None,
306            assembler: None,
307            incoming_packets: 0,
308            session_packets: 0,
309            data_packets: 0,
310        }
311    }
312
313    /// Return the currently accepted WFB session, if any.
314    pub fn session(&self) -> Option<&WfbSession> {
315        self.session.as_ref()
316    }
317
318    /// Return cumulative receive/FEC counters.
319    pub fn counters(&self) -> FecCounters {
320        let assembler = self
321            .assembler
322            .as_ref()
323            .map(PlainAssembler::counters)
324            .unwrap_or_default();
325        FecCounters {
326            total_packets: self.incoming_packets,
327            recovered_packets: assembler.recovered_packets,
328            lost_packets: assembler.lost_packets,
329            bad_packets: assembler.bad_packets,
330        }
331    }
332
333    /// Push one WFB forwarder packet payload.
334    pub fn push_forwarder_packet(&mut self, buf: &[u8]) -> Result<Vec<WfbEvent>, WfbError> {
335        match parse_forwarder_packet(buf)? {
336            WfbPacket::SessionKey {
337                session_nonce,
338                encrypted_session,
339            } => {
340                self.incoming_packets += 1;
341                self.session_packets += 1;
342                let session = self.decrypt_session(session_nonce, encrypted_session)?;
343                let changed = self
344                    .session
345                    .as_ref()
346                    .map(|current| current.session_key != session.session_key)
347                    .unwrap_or(true);
348                if changed {
349                    self.assembler = Some(PlainAssembler::new(session.fec_k, session.fec_n)?);
350                    self.session = Some(session.clone());
351                    Ok(vec![WfbEvent::Session(session)])
352                } else {
353                    Ok(Vec::new())
354                }
355            }
356            WfbPacket::Data {
357                data_nonce,
358                encrypted_payload,
359                associated_data,
360            } => {
361                self.incoming_packets += 1;
362                self.data_packets += 1;
363                let session = self.session.as_ref().ok_or(WfbError::MissingSession)?;
364                let nonce = &associated_data[1..WBLOCK_HDR_LEN];
365                let decrypted = decrypt_chacha20poly1305_legacy(
366                    &session.session_key,
367                    nonce,
368                    associated_data,
369                    encrypted_payload,
370                )
371                .map_err(|_| WfbError::DataDecryptFailed)?;
372                let assembler = self.assembler.as_mut().ok_or(WfbError::MissingSession)?;
373                Ok(assembler
374                    .push_decrypted_fragment(data_nonce, &decrypted)?
375                    .into_iter()
376                    .map(WfbEvent::Payload)
377                    .collect())
378            }
379        }
380    }
381
382    fn decrypt_session(
383        &self,
384        session_nonce: &[u8],
385        encrypted_session: &[u8],
386    ) -> Result<WfbSession, WfbError> {
387        let nonce: [u8; CRYPTO_BOX_NONCE_LEN] = session_nonce
388            .try_into()
389            .map_err(|_| WfbError::ShortSessionPacket)?;
390        let rx_secret = SecretKey::from(self.keypair.rx_secretkey);
391        let tx_public = PublicKey::from(self.keypair.tx_publickey);
392        let cipher = SalsaBox::new(&tx_public, &rx_secret);
393        let plaintext = cipher
394            .decrypt(BoxNonce::from_slice(&nonce), encrypted_session)
395            .map_err(|_| WfbError::SessionDecryptFailed)?;
396        let minimum_epoch = self
397            .session
398            .as_ref()
399            .map(|session| session.epoch.max(self.minimum_epoch))
400            .unwrap_or(self.minimum_epoch);
401        WfbSession::parse(&plaintext, self.channel_id, minimum_epoch)
402    }
403}
404
405/// Parse a WFB forwarder packet as data or session-key payload.
406pub fn parse_forwarder_packet(buf: &[u8]) -> Result<WfbPacket<'_>, WfbError> {
407    if buf.is_empty() {
408        return Err(WfbError::Empty);
409    }
410    if buf.len() > MAX_FORWARDER_PACKET_SIZE {
411        return Err(WfbError::TooLong);
412    }
413
414    match buf[0] {
415        WFB_PACKET_DATA => {
416            if buf.len() < WBLOCK_HDR_LEN + WPACKET_HDR_LEN + CHACHA20_POLY1305_TAG_LEN {
417                return Err(WfbError::ShortDataPacket);
418            }
419            let mut nonce = [0; 8];
420            nonce.copy_from_slice(&buf[1..9]);
421            Ok(WfbPacket::Data {
422                data_nonce: u64::from_be_bytes(nonce),
423                encrypted_payload: &buf[WBLOCK_HDR_LEN..],
424                associated_data: &buf[..WBLOCK_HDR_LEN],
425            })
426        }
427        WFB_PACKET_KEY => {
428            if buf.len() < WSESSION_HDR_LEN + WSESSION_DATA_LEN + CRYPTO_BOX_TAG_LEN {
429                return Err(WfbError::ShortSessionPacket);
430            }
431            Ok(WfbPacket::SessionKey {
432                session_nonce: &buf[1..WSESSION_HDR_LEN],
433                encrypted_session: &buf[WSESSION_HDR_LEN..],
434            })
435        }
436        other => Err(WfbError::UnknownPacketType(other)),
437    }
438}
439
440/// Recovered WFB application payload.
441#[derive(Debug, Clone, PartialEq, Eq)]
442pub struct WfbOutput {
443    /// Recovered packet sequence number.
444    pub packet_seq: u64,
445    /// Raw application payload bytes.
446    pub payload: Vec<u8>,
447}
448
449#[derive(Debug, Clone)]
450struct Block {
451    fragments: Vec<Option<Vec<u8>>>,
452    received: usize,
453    next_fragment: usize,
454}
455
456impl Block {
457    fn new(n: usize) -> Self {
458        Self {
459            fragments: vec![None; n],
460            received: 0,
461            next_fragment: 0,
462        }
463    }
464}
465
466/// Plain WFB FEC assembler.
467///
468/// This is used after data decryption, or directly for tests/pre-decrypted
469/// captures. It accepts primary and parity fragments and emits recovered
470/// application payloads in order.
471#[derive(Debug, Clone)]
472pub struct PlainAssembler {
473    fec_k: usize,
474    fec_n: usize,
475    fec: FecCode,
476    blocks: BTreeMap<u64, Block>,
477    next_block: Option<u64>,
478    /// Total fragments observed.
479    pub total_packets: u64,
480    /// Primary fragments considered lost.
481    pub lost_packets: u64,
482    /// Primary fragments recovered by FEC.
483    pub recovered_packets: u64,
484    /// Malformed or unrecoverable fragments.
485    pub bad_packets: u64,
486}
487
488impl PlainAssembler {
489    /// Create a plain assembler for `fec_k` primary and `fec_n` total fragments.
490    pub fn new(fec_k: usize, fec_n: usize) -> Result<Self, WfbError> {
491        if fec_k == 0 || fec_n == 0 || fec_k > fec_n || fec_n > 255 {
492            return Err(WfbError::InvalidFecParameters);
493        }
494        let fec = FecCode::new(fec_k, fec_n).map_err(|_| WfbError::InvalidFecParameters)?;
495        Ok(Self {
496            fec_k,
497            fec_n,
498            fec,
499            blocks: BTreeMap::new(),
500            next_block: None,
501            total_packets: 0,
502            lost_packets: 0,
503            recovered_packets: 0,
504            bad_packets: 0,
505        })
506    }
507
508    /// Return the primary fragment count.
509    pub const fn fec_k(&self) -> usize {
510        self.fec_k
511    }
512
513    /// Return the total primary plus parity fragment count.
514    pub const fn fec_n(&self) -> usize {
515        self.fec_n
516    }
517
518    /// Reset assembler state and FEC parameters.
519    pub fn reset_fec(&mut self, fec_k: usize, fec_n: usize) -> Result<(), WfbError> {
520        *self = Self::new(fec_k, fec_n)?;
521        Ok(())
522    }
523
524    /// Return cumulative FEC counters.
525    pub fn counters(&self) -> FecCounters {
526        FecCounters {
527            total_packets: self.total_packets,
528            recovered_packets: self.recovered_packets,
529            lost_packets: self.lost_packets,
530            bad_packets: self.bad_packets,
531        }
532    }
533
534    /// Push a decrypted WFB FEC fragment.
535    pub fn push_decrypted_fragment(
536        &mut self,
537        data_nonce: u64,
538        fragment: &[u8],
539    ) -> Result<Vec<WfbOutput>, WfbError> {
540        let block_idx = data_nonce >> 8;
541        let fragment_idx = (data_nonce & 0xff) as usize;
542
543        if block_idx > MAX_BLOCK_IDX {
544            return Err(WfbError::BlockIndexOverflow);
545        }
546        if fragment_idx >= self.fec_n {
547            return Err(WfbError::InvalidFragmentIndex);
548        }
549        self.total_packets += 1;
550
551        if self.next_block.is_none() {
552            self.next_block = Some(block_idx);
553        }
554        if self
555            .next_block
556            .map(|next_block| block_idx < next_block)
557            .unwrap_or(false)
558        {
559            return Ok(Vec::new());
560        }
561
562        let block = self
563            .blocks
564            .entry(block_idx)
565            .or_insert_with(|| Block::new(self.fec_n));
566        if block.fragments[fragment_idx].is_none() {
567            let mut padded = vec![0; MAX_FEC_PAYLOAD];
568            let len = fragment.len().min(MAX_FEC_PAYLOAD);
569            padded[..len].copy_from_slice(&fragment[..len]);
570            block.fragments[fragment_idx] = Some(padded);
571            block.received += 1;
572        }
573
574        Ok(self.drain_ready_blocks())
575    }
576
577    fn drain_ready_blocks(&mut self) -> Vec<WfbOutput> {
578        let mut out = Vec::new();
579        while let Some(block_idx) = self.next_block {
580            if !self.blocks.contains_key(&block_idx) {
581                if self.should_skip_missing_block(block_idx) {
582                    self.lost_packets += self.fec_k as u64;
583                    self.next_block = Some(block_idx + 1);
584                    continue;
585                }
586                break;
587            }
588
589            self.emit_contiguous_primary(block_idx, &mut out);
590            let complete = self
591                .blocks
592                .get(&block_idx)
593                .map(|block| block.next_fragment == self.fec_k)
594                .unwrap_or(false);
595            if complete {
596                self.blocks.remove(&block_idx);
597                self.next_block = Some(block_idx + 1);
598                continue;
599            }
600
601            let can_recover = self
602                .blocks
603                .get(&block_idx)
604                .map(|block| block.received >= self.fec_k)
605                .unwrap_or(false);
606            if can_recover {
607                if let Some(block) = self.blocks.get_mut(&block_idx) {
608                    match self
609                        .fec
610                        .recover_primary(&mut block.fragments, MAX_FEC_PAYLOAD)
611                    {
612                        Ok(recovered) => {
613                            self.recovered_packets += recovered as u64;
614                        }
615                        Err(_) => {
616                            self.bad_packets += 1;
617                            self.force_flush_block(block_idx, &mut out);
618                            continue;
619                        }
620                    }
621                }
622                self.emit_contiguous_primary(block_idx, &mut out);
623                self.blocks.remove(&block_idx);
624                self.next_block = Some(block_idx + 1);
625                continue;
626            }
627
628            if self.should_force_flush(block_idx) {
629                self.force_flush_block(block_idx, &mut out);
630                continue;
631            }
632
633            break;
634        }
635        out
636    }
637
638    fn should_skip_missing_block(&self, block_idx: u64) -> bool {
639        let Some((&next_present_block, block)) = self.blocks.range((block_idx + 1)..).next() else {
640            return false;
641        };
642
643        block.received >= self.fec_k
644            || self.blocks.len() > 40
645            || next_present_block.saturating_sub(block_idx) >= 40
646    }
647
648    fn emit_contiguous_primary(&mut self, block_idx: u64, out: &mut Vec<WfbOutput>) {
649        let Some(block) = self.blocks.get_mut(&block_idx) else {
650            return;
651        };
652        while block.next_fragment < self.fec_k {
653            let fragment_idx = block.next_fragment;
654            let Some(fragment) = block.fragments[fragment_idx].as_deref() else {
655                break;
656            };
657            let packet_seq = block_idx * self.fec_k as u64 + fragment_idx as u64;
658            match parse_plain_packet(fragment) {
659                Ok(Some(payload)) => out.push(WfbOutput {
660                    packet_seq,
661                    payload: payload.to_vec(),
662                }),
663                Ok(None) => {}
664                Err(_) => {
665                    self.bad_packets += 1;
666                }
667            }
668            block.next_fragment += 1;
669        }
670    }
671
672    fn should_force_flush(&self, block_idx: u64) -> bool {
673        if self.blocks.len() > 40 {
674            return true;
675        }
676        self.blocks
677            .range((block_idx + 1)..)
678            .any(|(_, block)| block.received >= self.fec_k)
679    }
680
681    fn force_flush_block(&mut self, block_idx: u64, out: &mut Vec<WfbOutput>) {
682        if let Some(block) = self.blocks.remove(&block_idx) {
683            for fragment_idx in block.next_fragment..self.fec_k {
684                let packet_seq = block_idx * self.fec_k as u64 + fragment_idx as u64;
685                match block.fragments[fragment_idx].as_deref() {
686                    Some(fragment) => match parse_plain_packet(fragment) {
687                        Ok(Some(payload)) => out.push(WfbOutput {
688                            packet_seq,
689                            payload: payload.to_vec(),
690                        }),
691                        Ok(None) => {}
692                        Err(_) => {
693                            self.bad_packets += 1;
694                        }
695                    },
696                    None => {
697                        self.lost_packets += 1;
698                    }
699                }
700            }
701            self.next_block = Some(block_idx + 1);
702        }
703    }
704}
705
706/// Parse a decrypted WFB plain packet and return payload bytes when present.
707pub fn parse_plain_packet(fragment: &[u8]) -> Result<Option<&[u8]>, WfbError> {
708    if fragment.len() < WPACKET_HDR_LEN {
709        return Err(WfbError::InvalidPlainPacket);
710    }
711    let flags = fragment[0];
712    let packet_size = u16::from_be_bytes([fragment[1], fragment[2]]) as usize;
713    if packet_size > MAX_PAYLOAD_SIZE || WPACKET_HDR_LEN + packet_size > fragment.len() {
714        return Err(WfbError::PayloadTooLarge);
715    }
716    if flags & WFB_PACKET_FEC_ONLY != 0 {
717        return Ok(None);
718    }
719    Ok(Some(
720        &fragment[WPACKET_HDR_LEN..WPACKET_HDR_LEN + packet_size],
721    ))
722}
723
724#[cfg(test)]
725mod tests {
726    use super::*;
727    use crate::crypto::encrypt_chacha20poly1305_legacy;
728    use crypto_box::aead::Aead;
729
730    fn plain(payload: &[u8]) -> Vec<u8> {
731        let mut out = Vec::new();
732        out.push(0);
733        out.extend_from_slice(&(payload.len() as u16).to_be_bytes());
734        out.extend_from_slice(payload);
735        out
736    }
737
738    fn padded(fragment: &[u8]) -> Vec<u8> {
739        let mut out = vec![0; MAX_FEC_PAYLOAD];
740        out[..fragment.len()].copy_from_slice(fragment);
741        out
742    }
743
744    #[test]
745    fn parses_forwarder_data_packet() {
746        let mut packet = vec![WFB_PACKET_DATA];
747        packet.extend_from_slice(&0x0102_0304_0506_0708u64.to_be_bytes());
748        let encrypted = [9; WPACKET_HDR_LEN + CHACHA20_POLY1305_TAG_LEN];
749        packet.extend_from_slice(&encrypted);
750
751        let parsed = parse_forwarder_packet(&packet).unwrap();
752        match parsed {
753            WfbPacket::Data {
754                data_nonce,
755                encrypted_payload,
756                associated_data,
757            } => {
758                assert_eq!(data_nonce, 0x0102_0304_0506_0708);
759                assert_eq!(encrypted_payload, encrypted);
760                assert_eq!(associated_data.len(), WBLOCK_HDR_LEN);
761            }
762            WfbPacket::SessionKey { .. } => panic!("expected data"),
763        }
764    }
765
766    #[test]
767    fn rejects_data_packets_without_encrypted_plain_header_and_tag() {
768        let mut packet = vec![WFB_PACKET_DATA];
769        packet.extend_from_slice(&0x0102_0304_0506_0708u64.to_be_bytes());
770        packet.extend_from_slice(&[0; WPACKET_HDR_LEN + CHACHA20_POLY1305_TAG_LEN - 1]);
771
772        assert_eq!(
773            parse_forwarder_packet(&packet),
774            Err(WfbError::ShortDataPacket)
775        );
776    }
777
778    #[test]
779    fn emits_primary_fragments_in_order() {
780        let mut assembler = PlainAssembler::new(2, 4).unwrap();
781        let first = assembler
782            .push_decrypted_fragment(0, &plain(b"first"))
783            .unwrap();
784        assert_eq!(first.len(), 1);
785        assert_eq!(first[0].payload, b"first");
786        let out = assembler
787            .push_decrypted_fragment(1, &plain(b"second"))
788            .unwrap();
789        assert_eq!(out.len(), 1);
790        assert_eq!(out[0].payload, b"second");
791    }
792
793    #[test]
794    fn recovers_missing_primary_fragment_from_fec() {
795        let fec = FecCode::new(3, 5).unwrap();
796        let primary = vec![
797            padded(&plain(b"first")),
798            padded(&plain(b"second")),
799            padded(&plain(b"third")),
800        ];
801        let parity = fec.encode(&primary, MAX_FEC_PAYLOAD).unwrap();
802
803        let mut assembler = PlainAssembler::new(3, 5).unwrap();
804        let first = assembler.push_decrypted_fragment(0, &primary[0]).unwrap();
805        assert_eq!(first[0].payload, b"first");
806        assert!(assembler
807            .push_decrypted_fragment(2, &primary[2])
808            .unwrap()
809            .is_empty());
810        let recovered = assembler.push_decrypted_fragment(3, &parity[0]).unwrap();
811        assert_eq!(recovered.len(), 2);
812        assert_eq!(recovered[0].payload, b"second");
813        assert_eq!(recovered[1].payload, b"third");
814        assert_eq!(assembler.recovered_packets, 1);
815    }
816
817    #[test]
818    fn skips_fully_missing_blocks_when_later_block_is_ready() {
819        let mut assembler = PlainAssembler::new(2, 2).unwrap();
820
821        let first = assembler
822            .push_decrypted_fragment(0, &plain(b"b0-f0"))
823            .unwrap();
824        assert_eq!(first[0].payload, b"b0-f0");
825
826        assert!(assembler
827            .push_decrypted_fragment(2 << 8, &plain(b"b2-f0"))
828            .unwrap()
829            .is_empty());
830        let out = assembler
831            .push_decrypted_fragment((2 << 8) | 1, &plain(b"b2-f1"))
832            .unwrap();
833
834        assert_eq!(out.len(), 2);
835        assert_eq!(out[0].payload, b"b2-f0");
836        assert_eq!(out[1].payload, b"b2-f1");
837        assert_eq!(assembler.lost_packets, 3);
838    }
839
840    #[test]
841    fn ignores_late_fragments_from_already_flushed_blocks() {
842        let mut assembler = PlainAssembler::new(2, 2).unwrap();
843
844        assembler
845            .push_decrypted_fragment(0, &plain(b"b0-f0"))
846            .unwrap();
847        assembler
848            .push_decrypted_fragment(2 << 8, &plain(b"b2-f0"))
849            .unwrap();
850        assembler
851            .push_decrypted_fragment((2 << 8) | 1, &plain(b"b2-f1"))
852            .unwrap();
853
854        let late = assembler
855            .push_decrypted_fragment(1 << 8, &plain(b"late-b1-f0"))
856            .unwrap();
857        assert!(late.is_empty());
858    }
859
860    #[test]
861    fn skips_fec_only_plain_packets() {
862        let mut fragment = vec![WFB_PACKET_FEC_ONLY];
863        fragment.extend_from_slice(&4u16.to_be_bytes());
864        fragment.extend_from_slice(b"skip");
865        assert!(parse_plain_packet(&fragment).unwrap().is_none());
866    }
867
868    #[test]
869    fn receiver_decrypts_session_and_data_packet() {
870        let rx_secret = SecretKey::from([1; CRYPTO_BOX_SECRETKEY_LEN]);
871        let tx_secret = SecretKey::from([2; CRYPTO_BOX_SECRETKEY_LEN]);
872        let keypair = WfbKeypair {
873            rx_secretkey: rx_secret.to_bytes(),
874            tx_publickey: *tx_secret.public_key().as_bytes(),
875        };
876        let channel_id = ChannelId::default_video();
877        let session_key = [7; CHACHA20_POLY1305_KEY_LEN];
878
879        let mut session_plain = Vec::new();
880        session_plain.extend_from_slice(&1u64.to_be_bytes());
881        session_plain.extend_from_slice(&channel_id.raw().to_be_bytes());
882        session_plain.push(WFB_FEC_VDM_RS);
883        session_plain.push(1);
884        session_plain.push(1);
885        session_plain.extend_from_slice(&session_key);
886        assert_eq!(session_plain.len(), WSESSION_DATA_LEN);
887        // wfb-ng allows encrypted optional session TLVs after the fixed fields.
888        session_plain.extend_from_slice(&[0x42, 0x00, 0x01, 0x99]);
889
890        let session_nonce = [3; CRYPTO_BOX_NONCE_LEN];
891        let tx_box = SalsaBox::new(&rx_secret.public_key(), &tx_secret);
892        let encrypted_session = tx_box
893            .encrypt(
894                BoxNonce::from_slice(&session_nonce),
895                session_plain.as_slice(),
896            )
897            .unwrap();
898        let mut session_packet = vec![WFB_PACKET_KEY];
899        session_packet.extend_from_slice(&session_nonce);
900        session_packet.extend_from_slice(&encrypted_session);
901
902        let mut receiver = WfbReceiver::new(channel_id, keypair, 0);
903        let session_events = receiver.push_forwarder_packet(&session_packet).unwrap();
904        assert!(matches!(session_events.as_slice(), [WfbEvent::Session(_)]));
905
906        let data_nonce = 0u64;
907        let mut block_header = vec![WFB_PACKET_DATA];
908        block_header.extend_from_slice(&data_nonce.to_be_bytes());
909        let encrypted_data = encrypt_chacha20poly1305_legacy(
910            &session_key,
911            &block_header[1..WBLOCK_HDR_LEN],
912            &block_header,
913            &plain(b"rtp payload"),
914        )
915        .unwrap();
916        let mut data_packet = block_header;
917        data_packet.extend_from_slice(&encrypted_data);
918
919        let payload_events = receiver.push_forwarder_packet(&data_packet).unwrap();
920        match payload_events.as_slice() {
921            [WfbEvent::Payload(payload)] => assert_eq!(payload.payload, b"rtp payload"),
922            other => panic!("unexpected events: {other:?}"),
923        }
924
925        let mut older_session_plain = Vec::new();
926        older_session_plain.extend_from_slice(&0u64.to_be_bytes());
927        older_session_plain.extend_from_slice(&channel_id.raw().to_be_bytes());
928        older_session_plain.push(WFB_FEC_VDM_RS);
929        older_session_plain.push(1);
930        older_session_plain.push(1);
931        older_session_plain.extend_from_slice(&[8; CHACHA20_POLY1305_KEY_LEN]);
932        let older_session_nonce = [4; CRYPTO_BOX_NONCE_LEN];
933        let encrypted_older_session = tx_box
934            .encrypt(
935                BoxNonce::from_slice(&older_session_nonce),
936                older_session_plain.as_slice(),
937            )
938            .unwrap();
939        let mut older_session_packet = vec![WFB_PACKET_KEY];
940        older_session_packet.extend_from_slice(&older_session_nonce);
941        older_session_packet.extend_from_slice(&encrypted_older_session);
942
943        assert_eq!(
944            receiver.push_forwarder_packet(&older_session_packet),
945            Err(WfbError::SessionEpochTooOld {
946                session_epoch: 0,
947                minimum_epoch: 1,
948            })
949        );
950    }
951}
openipc_core/wfb.rs

openipc_core/
wfb.rs
