Skip to main content

openipc_core/
wfb_tx.rs

1use crypto_box::aead::Aead;
2use crypto_box::{Nonce as BoxNonce, PublicKey, SalsaBox, SecretKey};
3use rand_core::{OsRng, RngCore};
4
5use crate::channel::ChannelId;
6use crate::crypto::encrypt_chacha20poly1305_legacy;
7use crate::fec::FecCode;
8use crate::ieee80211::build_wfb_header_with_frame_type;
9use crate::radiotap::{build_radiotap_header, TxRadioParams};
10use crate::wfb::{
11    WfbError, CHACHA20_POLY1305_KEY_LEN, CRYPTO_BOX_NONCE_LEN, CRYPTO_BOX_PUBLICKEY_LEN,
12    CRYPTO_BOX_SECRETKEY_LEN, MAX_BLOCK_IDX, MAX_FEC_PAYLOAD, MAX_PAYLOAD_SIZE, WBLOCK_HDR_LEN,
13    WFB_FEC_VDM_RS, WFB_PACKET_DATA, WFB_PACKET_FEC_ONLY, WFB_PACKET_KEY, WPACKET_HDR_LEN,
14    WSESSION_DATA_LEN, WSESSION_HDR_LEN,
15};
16
17/// Key material used by the ground station when transmitting WFB uplink data.
18///
19/// This is the inverse of `WfbKeypair`: it contains the transmitter secret key
20/// and the receiver public key needed to encrypt WFB session packets.
21#[derive(Debug, Clone, Copy, PartialEq, Eq)]
22pub struct WfbTxKeypair {
23    /// Secret key for the local transmitter.
24    pub tx_secretkey: [u8; CRYPTO_BOX_SECRETKEY_LEN],
25    /// Public key for the remote receiver.
26    pub rx_publickey: [u8; CRYPTO_BOX_PUBLICKEY_LEN],
27}
28
29impl WfbTxKeypair {
30    /// Parse a concatenated transmitter-secret + receiver-public keypair.
31    pub fn from_bytes(bytes: &[u8]) -> Result<Self, WfbError> {
32        if bytes.len() != CRYPTO_BOX_SECRETKEY_LEN + CRYPTO_BOX_PUBLICKEY_LEN {
33            return Err(WfbError::InvalidKeypair);
34        }
35        let mut tx_secretkey = [0; CRYPTO_BOX_SECRETKEY_LEN];
36        let mut rx_publickey = [0; CRYPTO_BOX_PUBLICKEY_LEN];
37        tx_secretkey.copy_from_slice(&bytes[..CRYPTO_BOX_SECRETKEY_LEN]);
38        rx_publickey.copy_from_slice(&bytes[CRYPTO_BOX_SECRETKEY_LEN..]);
39        Ok(Self {
40            tx_secretkey,
41            rx_publickey,
42        })
43    }
44}
45
46/// Stateful WFB transmitter for adaptive-link and other uplink payloads.
47///
48/// The transmitter owns the current WFB session key, fragments payloads into
49/// FEC blocks, emits parity fragments when configured, encrypts each block
50/// fragment, and can optionally wrap packets in radiotap + 802.11 headers for
51/// direct radio injection.
52#[derive(Debug, Clone)]
53pub struct WfbTransmitter {
54    channel_id: ChannelId,
55    keypair: WfbTxKeypair,
56    epoch: u64,
57    fec_k: usize,
58    fec_n: usize,
59    fec: FecCode,
60    block: Vec<Vec<u8>>,
61    block_index: u64,
62    fragment_index: usize,
63    max_packet_size: usize,
64    session_key: [u8; CHACHA20_POLY1305_KEY_LEN],
65    session_packet: Vec<u8>,
66    sequence_control: u16,
67}
68
69impl WfbTransmitter {
70    /// Create a transmitter for one WFB channel.
71    ///
72    /// `fec_k` is the number of source fragments per block and `fec_n` is the
73    /// total number of source + parity fragments transmitted for that block.
74    pub fn new(
75        channel_id: ChannelId,
76        keypair: WfbTxKeypair,
77        epoch: u64,
78        fec_k: usize,
79        fec_n: usize,
80    ) -> Result<Self, WfbError> {
81        if fec_k == 0 || fec_n == 0 || fec_k > fec_n || fec_n > 255 {
82            return Err(WfbError::InvalidFecParameters);
83        }
84        let fec = FecCode::new(fec_k, fec_n).map_err(|_| WfbError::InvalidFecParameters)?;
85        let mut tx = Self {
86            channel_id,
87            keypair,
88            epoch,
89            fec_k,
90            fec_n,
91            fec,
92            block: vec![vec![0; MAX_FEC_PAYLOAD]; fec_n],
93            block_index: 0,
94            fragment_index: 0,
95            max_packet_size: 0,
96            session_key: [0; CHACHA20_POLY1305_KEY_LEN],
97            session_packet: Vec::new(),
98            sequence_control: 0,
99        };
100        tx.rotate_session_key()?;
101        Ok(tx)
102    }
103
104    /// Return the WFB channel this transmitter writes to.
105    pub const fn channel_id(&self) -> ChannelId {
106        self.channel_id
107    }
108
109    /// Return the number of source fragments in each FEC block.
110    pub const fn fec_k(&self) -> usize {
111        self.fec_k
112    }
113
114    /// Return the total number of source + parity fragments in each FEC block.
115    pub const fn fec_n(&self) -> usize {
116        self.fec_n
117    }
118
119    /// Return the current encrypted WFB session packet without radio headers.
120    ///
121    /// Send this periodically before data packets so receivers can establish or
122    /// refresh the session key for this channel.
123    pub fn session_forwarder_packet(&self) -> &[u8] {
124        &self.session_packet
125    }
126
127    /// Build the current session packet as a radiotap + 802.11 radio packet.
128    pub fn session_radio_packet(&mut self, params: TxRadioParams) -> Vec<u8> {
129        let packet = self.session_packet.clone();
130        self.radio_packet_for_forwarder_packet(&packet, params)
131    }
132
133    /// Fragment, encrypt, FEC-encode, and wrap one payload for radio injection.
134    pub fn radio_packets_for_payload(
135        &mut self,
136        payload: &[u8],
137        params: TxRadioParams,
138    ) -> Result<Vec<Vec<u8>>, WfbError> {
139        let packets = self.forwarder_packets_for_payload(payload, 0)?;
140        Ok(packets
141            .into_iter()
142            .map(|packet| self.radio_packet_for_forwarder_packet(&packet, params))
143            .collect())
144    }
145
146    /// Wrap a WFB forwarder packet in radiotap + 802.11 headers for injection.
147    ///
148    /// This is useful for applications that need access to the raw forwarder
149    /// packet first, for example to mirror it to a debug UDP path or to delay
150    /// parity fragments before the final radio wrapping step.
151    pub fn radio_packet_for_forwarder_packet(
152        &mut self,
153        forwarder_packet: &[u8],
154        params: TxRadioParams,
155    ) -> Vec<u8> {
156        self.wrap_forwarder_packet(forwarder_packet, params)
157    }
158
159    /// Return true when the current FEC block contains at least one fragment.
160    pub const fn has_open_fec_block(&self) -> bool {
161        self.fragment_index != 0
162    }
163
164    /// Emit one FEC-only empty fragment to advance or close a partial block.
165    ///
166    /// WFB-ng uses this when no input payload arrives for `fec_timeout`: each
167    /// timeout contributes one empty source fragment until the current block is
168    /// complete and parity can be generated. If no block is open, no packet is
169    /// emitted.
170    pub fn forwarder_packets_for_fec_only(&mut self) -> Result<Vec<Vec<u8>>, WfbError> {
171        if !self.has_open_fec_block() {
172            return Ok(Vec::new());
173        }
174        self.forwarder_packets_for_payload(&[], WFB_PACKET_FEC_ONLY)
175    }
176
177    /// Emit enough FEC-only empty fragments to close a partial FEC block.
178    ///
179    /// This mirrors the control-path behavior in WFB-ng when FEC settings are
180    /// changed: the old block is completed before a new session is started.
181    pub fn close_fec_block(&mut self) -> Result<Vec<Vec<u8>>, WfbError> {
182        let mut out = Vec::new();
183        while self.has_open_fec_block() {
184            out.extend(self.forwarder_packets_for_fec_only()?);
185        }
186        Ok(out)
187    }
188
189    /// Fragment, encrypt, and FEC-encode one payload as WFB forwarder packets.
190    ///
191    /// The returned packets do not include radiotap or 802.11 headers, which
192    /// makes this useful when another layer owns radio framing.
193    pub fn forwarder_packets_for_payload(
194        &mut self,
195        payload: &[u8],
196        flags: u8,
197    ) -> Result<Vec<Vec<u8>>, WfbError> {
198        if flags & WFB_PACKET_FEC_ONLY != 0 && !self.has_open_fec_block() {
199            return Ok(Vec::new());
200        }
201        if payload.len() > MAX_PAYLOAD_SIZE {
202            return Err(WfbError::PayloadTooLarge);
203        }
204
205        let fragment_index = self.fragment_index;
206        let fragment = &mut self.block[fragment_index];
207        fragment.fill(0);
208        fragment[0] = flags;
209        fragment[1..3].copy_from_slice(&(payload.len() as u16).to_be_bytes());
210        fragment[WPACKET_HDR_LEN..WPACKET_HDR_LEN + payload.len()].copy_from_slice(payload);
211        let packet_size = WPACKET_HDR_LEN + payload.len();
212
213        let mut out = vec![self.encrypt_block_fragment(fragment_index, packet_size)?];
214        self.max_packet_size = self.max_packet_size.max(packet_size);
215        self.fragment_index += 1;
216
217        if self.fragment_index == self.fec_k {
218            if self.fec_n > self.fec_k {
219                let parity = self
220                    .fec
221                    .encode(&self.block[..self.fec_k], self.max_packet_size)
222                    .map_err(|_| WfbError::FecRecoveryFailed)?;
223                for (offset, parity_fragment) in parity.into_iter().enumerate() {
224                    let idx = self.fec_k + offset;
225                    self.block[idx].fill(0);
226                    self.block[idx][..parity_fragment.len()].copy_from_slice(&parity_fragment);
227                    out.push(self.encrypt_block_fragment(idx, self.max_packet_size)?);
228                }
229            }
230            self.finish_block()?;
231        }
232
233        Ok(out)
234    }
235
236    fn finish_block(&mut self) -> Result<(), WfbError> {
237        self.block_index += 1;
238        self.fragment_index = 0;
239        self.max_packet_size = 0;
240        if self.block_index > MAX_BLOCK_IDX {
241            self.block_index = 0;
242            self.rotate_session_key()?;
243        }
244        Ok(())
245    }
246
247    fn encrypt_block_fragment(
248        &self,
249        fragment_index: usize,
250        packet_size: usize,
251    ) -> Result<Vec<u8>, WfbError> {
252        let data_nonce = ((self.block_index & MAX_BLOCK_IDX) << 8) | fragment_index as u64;
253        let mut block_header = [0u8; WBLOCK_HDR_LEN];
254        block_header[0] = WFB_PACKET_DATA;
255        block_header[1..].copy_from_slice(&data_nonce.to_be_bytes());
256        let nonce = &block_header[1..WBLOCK_HDR_LEN];
257        let encrypted = encrypt_chacha20poly1305_legacy(
258            &self.session_key,
259            nonce,
260            &block_header,
261            &self.block[fragment_index][..packet_size],
262        )
263        .map_err(|_| WfbError::DataEncryptFailed)?;
264
265        let mut out = Vec::with_capacity(WBLOCK_HDR_LEN + encrypted.len());
266        out.extend_from_slice(&block_header);
267        out.extend_from_slice(&encrypted);
268        Ok(out)
269    }
270
271    fn rotate_session_key(&mut self) -> Result<(), WfbError> {
272        OsRng.fill_bytes(&mut self.session_key);
273        self.session_packet = self.build_session_packet()?;
274        Ok(())
275    }
276
277    fn build_session_packet(&self) -> Result<Vec<u8>, WfbError> {
278        let mut nonce = [0u8; CRYPTO_BOX_NONCE_LEN];
279        OsRng.fill_bytes(&mut nonce);
280
281        let mut session_data = [0u8; WSESSION_DATA_LEN];
282        session_data[0..8].copy_from_slice(&self.epoch.to_be_bytes());
283        session_data[8..12].copy_from_slice(&self.channel_id.raw().to_be_bytes());
284        session_data[12] = WFB_FEC_VDM_RS;
285        session_data[13] = self.fec_k as u8;
286        session_data[14] = self.fec_n as u8;
287        session_data[15..47].copy_from_slice(&self.session_key);
288
289        let tx_secret = SecretKey::from(self.keypair.tx_secretkey);
290        let rx_public = PublicKey::from(self.keypair.rx_publickey);
291        let cipher = SalsaBox::new(&rx_public, &tx_secret);
292        let encrypted = cipher
293            .encrypt(BoxNonce::from_slice(&nonce), session_data.as_slice())
294            .map_err(|_| WfbError::SessionEncryptFailed)?;
295
296        let mut out = Vec::with_capacity(WSESSION_HDR_LEN + encrypted.len());
297        out.push(WFB_PACKET_KEY);
298        out.extend_from_slice(&nonce);
299        out.extend_from_slice(&encrypted);
300        Ok(out)
301    }
302
303    fn wrap_forwarder_packet(&mut self, forwarder_packet: &[u8], params: TxRadioParams) -> Vec<u8> {
304        let mut out = build_radiotap_header(params);
305        let seq = self.sequence_control.to_le_bytes();
306        out.extend_from_slice(&build_wfb_header_with_frame_type(
307            self.channel_id,
308            seq,
309            params.frame_type,
310        ));
311        out.extend_from_slice(forwarder_packet);
312        self.sequence_control = self.sequence_control.wrapping_add(16);
313        out
314    }
315}
316
317#[cfg(test)]
318mod tests {
319    use super::*;
320    use crate::wfb::{WfbKeypair, WfbReceiver};
321    use crypto_box::SecretKey;
322
323    fn linked_keypairs() -> (WfbTxKeypair, WfbKeypair) {
324        let ground_secret = SecretKey::from([3u8; 32]);
325        let air_secret = SecretKey::from([9u8; 32]);
326        let ground_public = ground_secret.public_key();
327        let air_public = air_secret.public_key();
328        (
329            WfbTxKeypair {
330                tx_secretkey: ground_secret.to_bytes(),
331                rx_publickey: air_public.to_bytes(),
332            },
333            WfbKeypair {
334                rx_secretkey: air_secret.to_bytes(),
335                tx_publickey: ground_public.to_bytes(),
336            },
337        )
338    }
339
340    #[test]
341    fn transmitted_session_and_payload_roundtrip() {
342        let channel = ChannelId::from_link_port(0x112233, crate::RadioPort::TunnelTx);
343        let (tx_keys, rx_keys) = linked_keypairs();
344        let mut tx = WfbTransmitter::new(channel, tx_keys, 42, 1, 1).unwrap();
345        let mut rx = WfbReceiver::new(channel, rx_keys, 0);
346
347        let session_events = rx
348            .push_forwarder_packet(tx.session_forwarder_packet())
349            .unwrap();
350        assert_eq!(session_events.len(), 1);
351
352        let data_packets = tx.forwarder_packets_for_payload(b"hello", 0).unwrap();
353        assert_eq!(data_packets.len(), 1);
354        let events = rx.push_forwarder_packet(&data_packets[0]).unwrap();
355        assert_eq!(events.len(), 1);
356        match &events[0] {
357            crate::wfb::WfbEvent::Payload(payload) => assert_eq!(payload.payload, b"hello"),
358            other => panic!("unexpected event: {other:?}"),
359        }
360    }
361
362    #[test]
363    fn fec_only_closes_partial_blocks_without_payload_events() {
364        let channel = ChannelId::from_link_port(0x112233, crate::RadioPort::TunnelTx);
365        let (tx_keys, rx_keys) = linked_keypairs();
366        let mut tx = WfbTransmitter::new(channel, tx_keys, 42, 2, 3).unwrap();
367        let mut rx = WfbReceiver::new(channel, rx_keys, 0);
368
369        rx.push_forwarder_packet(tx.session_forwarder_packet())
370            .unwrap();
371        let first = tx.forwarder_packets_for_payload(b"hello", 0).unwrap();
372        assert_eq!(first.len(), 1);
373        assert!(tx.has_open_fec_block());
374
375        let close = tx.forwarder_packets_for_fec_only().unwrap();
376        assert_eq!(close.len(), 2);
377        assert!(!tx.has_open_fec_block());
378
379        let mut payloads = Vec::new();
380        for packet in first.iter().chain(close.iter()) {
381            for event in rx.push_forwarder_packet(packet).unwrap() {
382                if let crate::wfb::WfbEvent::Payload(payload) = event {
383                    payloads.push(payload.payload);
384                }
385            }
386        }
387        assert_eq!(payloads, vec![b"hello".to_vec()]);
388    }
389}