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_KEY, WPACKET_HDR_LEN, WSESSION_DATA_LEN,
14    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.wrap_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.wrap_forwarder_packet(&packet, params))
143            .collect())
144    }
145
146    /// Fragment, encrypt, and FEC-encode one payload as WFB forwarder packets.
147    ///
148    /// The returned packets do not include radiotap or 802.11 headers, which
149    /// makes this useful when another layer owns radio framing.
150    pub fn forwarder_packets_for_payload(
151        &mut self,
152        payload: &[u8],
153        flags: u8,
154    ) -> Result<Vec<Vec<u8>>, WfbError> {
155        if payload.len() > MAX_PAYLOAD_SIZE {
156            return Err(WfbError::PayloadTooLarge);
157        }
158
159        let fragment_index = self.fragment_index;
160        let fragment = &mut self.block[fragment_index];
161        fragment.fill(0);
162        fragment[0] = flags;
163        fragment[1..3].copy_from_slice(&(payload.len() as u16).to_be_bytes());
164        fragment[WPACKET_HDR_LEN..WPACKET_HDR_LEN + payload.len()].copy_from_slice(payload);
165        let packet_size = WPACKET_HDR_LEN + payload.len();
166
167        let mut out = vec![self.encrypt_block_fragment(fragment_index, packet_size)?];
168        self.max_packet_size = self.max_packet_size.max(packet_size);
169        self.fragment_index += 1;
170
171        if self.fragment_index == self.fec_k {
172            if self.fec_n > self.fec_k {
173                let parity = self
174                    .fec
175                    .encode(&self.block[..self.fec_k], self.max_packet_size)
176                    .map_err(|_| WfbError::FecRecoveryFailed)?;
177                for (offset, parity_fragment) in parity.into_iter().enumerate() {
178                    let idx = self.fec_k + offset;
179                    self.block[idx].fill(0);
180                    self.block[idx][..parity_fragment.len()].copy_from_slice(&parity_fragment);
181                    out.push(self.encrypt_block_fragment(idx, self.max_packet_size)?);
182                }
183            }
184            self.finish_block()?;
185        }
186
187        Ok(out)
188    }
189
190    fn finish_block(&mut self) -> Result<(), WfbError> {
191        self.block_index += 1;
192        self.fragment_index = 0;
193        self.max_packet_size = 0;
194        if self.block_index > MAX_BLOCK_IDX {
195            self.block_index = 0;
196            self.rotate_session_key()?;
197        }
198        Ok(())
199    }
200
201    fn encrypt_block_fragment(
202        &self,
203        fragment_index: usize,
204        packet_size: usize,
205    ) -> Result<Vec<u8>, WfbError> {
206        let data_nonce = ((self.block_index & MAX_BLOCK_IDX) << 8) | fragment_index as u64;
207        let mut block_header = [0u8; WBLOCK_HDR_LEN];
208        block_header[0] = WFB_PACKET_DATA;
209        block_header[1..].copy_from_slice(&data_nonce.to_be_bytes());
210        let nonce = &block_header[1..WBLOCK_HDR_LEN];
211        let encrypted = encrypt_chacha20poly1305_legacy(
212            &self.session_key,
213            nonce,
214            &block_header,
215            &self.block[fragment_index][..packet_size],
216        )
217        .map_err(|_| WfbError::DataEncryptFailed)?;
218
219        let mut out = Vec::with_capacity(WBLOCK_HDR_LEN + encrypted.len());
220        out.extend_from_slice(&block_header);
221        out.extend_from_slice(&encrypted);
222        Ok(out)
223    }
224
225    fn rotate_session_key(&mut self) -> Result<(), WfbError> {
226        OsRng.fill_bytes(&mut self.session_key);
227        self.session_packet = self.build_session_packet()?;
228        Ok(())
229    }
230
231    fn build_session_packet(&self) -> Result<Vec<u8>, WfbError> {
232        let mut nonce = [0u8; CRYPTO_BOX_NONCE_LEN];
233        OsRng.fill_bytes(&mut nonce);
234
235        let mut session_data = [0u8; WSESSION_DATA_LEN];
236        session_data[0..8].copy_from_slice(&self.epoch.to_be_bytes());
237        session_data[8..12].copy_from_slice(&self.channel_id.raw().to_be_bytes());
238        session_data[12] = WFB_FEC_VDM_RS;
239        session_data[13] = self.fec_k as u8;
240        session_data[14] = self.fec_n as u8;
241        session_data[15..47].copy_from_slice(&self.session_key);
242
243        let tx_secret = SecretKey::from(self.keypair.tx_secretkey);
244        let rx_public = PublicKey::from(self.keypair.rx_publickey);
245        let cipher = SalsaBox::new(&rx_public, &tx_secret);
246        let encrypted = cipher
247            .encrypt(BoxNonce::from_slice(&nonce), session_data.as_slice())
248            .map_err(|_| WfbError::SessionEncryptFailed)?;
249
250        let mut out = Vec::with_capacity(WSESSION_HDR_LEN + encrypted.len());
251        out.push(WFB_PACKET_KEY);
252        out.extend_from_slice(&nonce);
253        out.extend_from_slice(&encrypted);
254        Ok(out)
255    }
256
257    fn wrap_forwarder_packet(&mut self, forwarder_packet: &[u8], params: TxRadioParams) -> Vec<u8> {
258        let mut out = build_radiotap_header(params);
259        let seq = self.sequence_control.to_le_bytes();
260        out.extend_from_slice(&build_wfb_header_with_frame_type(
261            self.channel_id,
262            seq,
263            params.frame_type,
264        ));
265        out.extend_from_slice(forwarder_packet);
266        self.sequence_control = self.sequence_control.wrapping_add(16);
267        out
268    }
269}
270
271#[cfg(test)]
272mod tests {
273    use super::*;
274    use crate::wfb::{WfbKeypair, WfbReceiver};
275    use crypto_box::SecretKey;
276
277    fn linked_keypairs() -> (WfbTxKeypair, WfbKeypair) {
278        let ground_secret = SecretKey::from([3u8; 32]);
279        let air_secret = SecretKey::from([9u8; 32]);
280        let ground_public = ground_secret.public_key();
281        let air_public = air_secret.public_key();
282        (
283            WfbTxKeypair {
284                tx_secretkey: ground_secret.to_bytes(),
285                rx_publickey: air_public.to_bytes(),
286            },
287            WfbKeypair {
288                rx_secretkey: air_secret.to_bytes(),
289                tx_publickey: ground_public.to_bytes(),
290            },
291        )
292    }
293
294    #[test]
295    fn transmitted_session_and_payload_roundtrip() {
296        let channel = ChannelId::from_link_port(0x112233, crate::RadioPort::TunnelTx);
297        let (tx_keys, rx_keys) = linked_keypairs();
298        let mut tx = WfbTransmitter::new(channel, tx_keys, 42, 1, 1).unwrap();
299        let mut rx = WfbReceiver::new(channel, rx_keys, 0);
300
301        let session_events = rx
302            .push_forwarder_packet(tx.session_forwarder_packet())
303            .unwrap();
304        assert_eq!(session_events.len(), 1);
305
306        let data_packets = tx.forwarder_packets_for_payload(b"hello", 0).unwrap();
307        assert_eq!(data_packets.len(), 1);
308        let events = rx.push_forwarder_packet(&data_packets[0]).unwrap();
309        assert_eq!(events.len(), 1);
310        match &events[0] {
311            crate::wfb::WfbEvent::Payload(payload) => assert_eq!(payload.payload, b"hello"),
312            other => panic!("unexpected event: {other:?}"),
313        }
314    }
315}