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)]
18pub struct WfbTxKeypair {
19 pub tx_secretkey: [u8; CRYPTO_BOX_SECRETKEY_LEN],
20 pub rx_publickey: [u8; CRYPTO_BOX_PUBLICKEY_LEN],
21}
22
23impl WfbTxKeypair {
24 pub fn from_bytes(bytes: &[u8]) -> Result<Self, WfbError> {
25 if bytes.len() != CRYPTO_BOX_SECRETKEY_LEN + CRYPTO_BOX_PUBLICKEY_LEN {
26 return Err(WfbError::InvalidKeypair);
27 }
28 let mut tx_secretkey = [0; CRYPTO_BOX_SECRETKEY_LEN];
29 let mut rx_publickey = [0; CRYPTO_BOX_PUBLICKEY_LEN];
30 tx_secretkey.copy_from_slice(&bytes[..CRYPTO_BOX_SECRETKEY_LEN]);
31 rx_publickey.copy_from_slice(&bytes[CRYPTO_BOX_SECRETKEY_LEN..]);
32 Ok(Self {
33 tx_secretkey,
34 rx_publickey,
35 })
36 }
37}
38
39#[derive(Debug, Clone)]
40pub struct WfbTransmitter {
41 channel_id: ChannelId,
42 keypair: WfbTxKeypair,
43 epoch: u64,
44 fec_k: usize,
45 fec_n: usize,
46 fec: FecCode,
47 block: Vec<Vec<u8>>,
48 block_index: u64,
49 fragment_index: usize,
50 max_packet_size: usize,
51 session_key: [u8; CHACHA20_POLY1305_KEY_LEN],
52 session_packet: Vec<u8>,
53 sequence_control: u16,
54}
55
56impl WfbTransmitter {
57 pub fn new(
58 channel_id: ChannelId,
59 keypair: WfbTxKeypair,
60 epoch: u64,
61 fec_k: usize,
62 fec_n: usize,
63 ) -> Result<Self, WfbError> {
64 if fec_k == 0 || fec_n == 0 || fec_k > fec_n || fec_n > 255 {
65 return Err(WfbError::InvalidFecParameters);
66 }
67 let fec = FecCode::new(fec_k, fec_n).map_err(|_| WfbError::InvalidFecParameters)?;
68 let mut tx = Self {
69 channel_id,
70 keypair,
71 epoch,
72 fec_k,
73 fec_n,
74 fec,
75 block: vec![vec![0; MAX_FEC_PAYLOAD]; fec_n],
76 block_index: 0,
77 fragment_index: 0,
78 max_packet_size: 0,
79 session_key: [0; CHACHA20_POLY1305_KEY_LEN],
80 session_packet: Vec::new(),
81 sequence_control: 0,
82 };
83 tx.rotate_session_key()?;
84 Ok(tx)
85 }
86
87 pub const fn channel_id(&self) -> ChannelId {
88 self.channel_id
89 }
90
91 pub const fn fec_k(&self) -> usize {
92 self.fec_k
93 }
94
95 pub const fn fec_n(&self) -> usize {
96 self.fec_n
97 }
98
99 pub fn session_forwarder_packet(&self) -> &[u8] {
100 &self.session_packet
101 }
102
103 pub fn session_radio_packet(&mut self, params: TxRadioParams) -> Vec<u8> {
104 let packet = self.session_packet.clone();
105 self.wrap_forwarder_packet(&packet, params)
106 }
107
108 pub fn radio_packets_for_payload(
109 &mut self,
110 payload: &[u8],
111 params: TxRadioParams,
112 ) -> Result<Vec<Vec<u8>>, WfbError> {
113 let packets = self.forwarder_packets_for_payload(payload, 0)?;
114 Ok(packets
115 .into_iter()
116 .map(|packet| self.wrap_forwarder_packet(&packet, params))
117 .collect())
118 }
119
120 pub fn forwarder_packets_for_payload(
121 &mut self,
122 payload: &[u8],
123 flags: u8,
124 ) -> Result<Vec<Vec<u8>>, WfbError> {
125 if payload.len() > MAX_PAYLOAD_SIZE {
126 return Err(WfbError::PayloadTooLarge);
127 }
128
129 let fragment_index = self.fragment_index;
130 let fragment = &mut self.block[fragment_index];
131 fragment.fill(0);
132 fragment[0] = flags;
133 fragment[1..3].copy_from_slice(&(payload.len() as u16).to_be_bytes());
134 fragment[WPACKET_HDR_LEN..WPACKET_HDR_LEN + payload.len()].copy_from_slice(payload);
135 let packet_size = WPACKET_HDR_LEN + payload.len();
136
137 let mut out = vec![self.encrypt_block_fragment(fragment_index, packet_size)?];
138 self.max_packet_size = self.max_packet_size.max(packet_size);
139 self.fragment_index += 1;
140
141 if self.fragment_index == self.fec_k {
142 if self.fec_n > self.fec_k {
143 let parity = self
144 .fec
145 .encode(&self.block[..self.fec_k], self.max_packet_size)
146 .map_err(|_| WfbError::FecRecoveryFailed)?;
147 for (offset, parity_fragment) in parity.into_iter().enumerate() {
148 let idx = self.fec_k + offset;
149 self.block[idx].fill(0);
150 self.block[idx][..parity_fragment.len()].copy_from_slice(&parity_fragment);
151 out.push(self.encrypt_block_fragment(idx, self.max_packet_size)?);
152 }
153 }
154 self.finish_block()?;
155 }
156
157 Ok(out)
158 }
159
160 fn finish_block(&mut self) -> Result<(), WfbError> {
161 self.block_index += 1;
162 self.fragment_index = 0;
163 self.max_packet_size = 0;
164 if self.block_index > MAX_BLOCK_IDX {
165 self.block_index = 0;
166 self.rotate_session_key()?;
167 }
168 Ok(())
169 }
170
171 fn encrypt_block_fragment(
172 &self,
173 fragment_index: usize,
174 packet_size: usize,
175 ) -> Result<Vec<u8>, WfbError> {
176 let data_nonce = ((self.block_index & MAX_BLOCK_IDX) << 8) | fragment_index as u64;
177 let mut block_header = [0u8; WBLOCK_HDR_LEN];
178 block_header[0] = WFB_PACKET_DATA;
179 block_header[1..].copy_from_slice(&data_nonce.to_be_bytes());
180 let nonce = &block_header[1..WBLOCK_HDR_LEN];
181 let encrypted = encrypt_chacha20poly1305_legacy(
182 &self.session_key,
183 nonce,
184 &block_header,
185 &self.block[fragment_index][..packet_size],
186 )
187 .map_err(|_| WfbError::DataEncryptFailed)?;
188
189 let mut out = Vec::with_capacity(WBLOCK_HDR_LEN + encrypted.len());
190 out.extend_from_slice(&block_header);
191 out.extend_from_slice(&encrypted);
192 Ok(out)
193 }
194
195 fn rotate_session_key(&mut self) -> Result<(), WfbError> {
196 OsRng.fill_bytes(&mut self.session_key);
197 self.session_packet = self.build_session_packet()?;
198 Ok(())
199 }
200
201 fn build_session_packet(&self) -> Result<Vec<u8>, WfbError> {
202 let mut nonce = [0u8; CRYPTO_BOX_NONCE_LEN];
203 OsRng.fill_bytes(&mut nonce);
204
205 let mut session_data = [0u8; WSESSION_DATA_LEN];
206 session_data[0..8].copy_from_slice(&self.epoch.to_be_bytes());
207 session_data[8..12].copy_from_slice(&self.channel_id.raw().to_be_bytes());
208 session_data[12] = WFB_FEC_VDM_RS;
209 session_data[13] = self.fec_k as u8;
210 session_data[14] = self.fec_n as u8;
211 session_data[15..47].copy_from_slice(&self.session_key);
212
213 let tx_secret = SecretKey::from(self.keypair.tx_secretkey);
214 let rx_public = PublicKey::from(self.keypair.rx_publickey);
215 let cipher = SalsaBox::new(&rx_public, &tx_secret);
216 let encrypted = cipher
217 .encrypt(BoxNonce::from_slice(&nonce), session_data.as_slice())
218 .map_err(|_| WfbError::SessionEncryptFailed)?;
219
220 let mut out = Vec::with_capacity(WSESSION_HDR_LEN + encrypted.len());
221 out.push(WFB_PACKET_KEY);
222 out.extend_from_slice(&nonce);
223 out.extend_from_slice(&encrypted);
224 Ok(out)
225 }
226
227 fn wrap_forwarder_packet(&mut self, forwarder_packet: &[u8], params: TxRadioParams) -> Vec<u8> {
228 let mut out = build_radiotap_header(params);
229 let seq = self.sequence_control.to_le_bytes();
230 out.extend_from_slice(&build_wfb_header_with_frame_type(
231 self.channel_id,
232 seq,
233 params.frame_type,
234 ));
235 out.extend_from_slice(forwarder_packet);
236 self.sequence_control = self.sequence_control.wrapping_add(16);
237 out
238 }
239}
240
241#[cfg(test)]
242mod tests {
243 use super::*;
244 use crate::wfb::{WfbKeypair, WfbReceiver};
245 use crypto_box::SecretKey;
246
247 fn linked_keypairs() -> (WfbTxKeypair, WfbKeypair) {
248 let ground_secret = SecretKey::from([3u8; 32]);
249 let air_secret = SecretKey::from([9u8; 32]);
250 let ground_public = ground_secret.public_key();
251 let air_public = air_secret.public_key();
252 (
253 WfbTxKeypair {
254 tx_secretkey: ground_secret.to_bytes(),
255 rx_publickey: air_public.to_bytes(),
256 },
257 WfbKeypair {
258 rx_secretkey: air_secret.to_bytes(),
259 tx_publickey: ground_public.to_bytes(),
260 },
261 )
262 }
263
264 #[test]
265 fn transmitted_session_and_payload_roundtrip() {
266 let channel = ChannelId::from_link_port(0x112233, crate::RadioPort::MavlinkTx);
267 let (tx_keys, rx_keys) = linked_keypairs();
268 let mut tx = WfbTransmitter::new(channel, tx_keys, 42, 1, 1).unwrap();
269 let mut rx = WfbReceiver::new(channel, rx_keys, 0);
270
271 let session_events = rx
272 .push_forwarder_packet(tx.session_forwarder_packet())
273 .unwrap();
274 assert_eq!(session_events.len(), 1);
275
276 let data_packets = tx.forwarder_packets_for_payload(b"hello", 0).unwrap();
277 assert_eq!(data_packets.len(), 1);
278 let events = rx.push_forwarder_packet(&data_packets[0]).unwrap();
279 assert_eq!(events.len(), 1);
280 match &events[0] {
281 crate::wfb::WfbEvent::Payload(payload) => assert_eq!(payload.payload, b"hello"),
282 other => panic!("unexpected event: {other:?}"),
283 }
284 }
285}