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#[derive(Debug, Clone, Copy, PartialEq, Eq)]
22pub struct WfbTxKeypair {
23 pub tx_secretkey: [u8; CRYPTO_BOX_SECRETKEY_LEN],
25 pub rx_publickey: [u8; CRYPTO_BOX_PUBLICKEY_LEN],
27}
28
29impl WfbTxKeypair {
30 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#[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 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 pub const fn channel_id(&self) -> ChannelId {
106 self.channel_id
107 }
108
109 pub const fn fec_k(&self) -> usize {
111 self.fec_k
112 }
113
114 pub const fn fec_n(&self) -> usize {
116 self.fec_n
117 }
118
119 pub fn session_forwarder_packet(&self) -> &[u8] {
124 &self.session_packet
125 }
126
127 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 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 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 pub const fn has_open_fec_block(&self) -> bool {
161 self.fragment_index != 0
162 }
163
164 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 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 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}