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#[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.wrap_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.wrap_forwarder_packet(&packet, params))
143 .collect())
144 }
145
146 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}