1use crate::channel::ChannelId;
2use crate::ieee80211::{FrameLayout, WifiFrame};
3use crate::wfb::{
4 parse_forwarder_packet, FecCounters, PlainAssembler, WfbError, WfbEvent, WfbKeypair, WfbPacket,
5 WfbReceiver,
6};
7
8#[derive(Debug, Clone, PartialEq, Eq)]
10pub struct RecoveredPayload {
11 pub channel_id: ChannelId,
13 pub packet_seq: u64,
15 pub data: Vec<u8>,
17}
18
19#[derive(Debug, Clone, PartialEq, Eq)]
21pub enum PayloadPipelineEvent {
22 IgnoredFrame,
24 SessionEstablished {
26 epoch: u64,
28 fec_k: usize,
30 fec_n: usize,
32 },
33 Payload(RecoveredPayload),
35}
36
37#[derive(Debug, Clone)]
42pub struct PayloadPipeline {
43 channel_id: ChannelId,
44 frame_layout: FrameLayout,
45 assembler: PlainAssembler,
46 wfb_receiver: Option<WfbReceiver>,
47}
48
49impl PayloadPipeline {
50 pub fn new(
52 channel_id: ChannelId,
53 frame_layout: FrameLayout,
54 fec_k: usize,
55 fec_n: usize,
56 ) -> Result<Self, WfbError> {
57 Ok(Self {
58 channel_id,
59 frame_layout,
60 assembler: PlainAssembler::new(fec_k, fec_n)?,
61 wfb_receiver: None,
62 })
63 }
64
65 pub fn with_keypair(
67 channel_id: ChannelId,
68 frame_layout: FrameLayout,
69 keypair: WfbKeypair,
70 minimum_epoch: u64,
71 ) -> Result<Self, WfbError> {
72 Ok(Self {
73 channel_id,
74 frame_layout,
75 assembler: PlainAssembler::new(1, 1)?,
76 wfb_receiver: Some(WfbReceiver::new(channel_id, keypair, minimum_epoch)),
77 })
78 }
79
80 pub const fn channel_id(&self) -> ChannelId {
82 self.channel_id
83 }
84
85 pub fn fec_counters(&self) -> FecCounters {
87 self.wfb_receiver
88 .as_ref()
89 .map(WfbReceiver::counters)
90 .unwrap_or_else(|| self.assembler.counters())
91 }
92
93 pub fn accepts_80211_frame(&self, frame: &[u8]) -> bool {
95 WifiFrame::parse(frame, self.frame_layout)
96 .map(|frame| frame.matches_channel_id(self.channel_id))
97 .unwrap_or(false)
98 }
99
100 pub fn push_80211_frame(
107 &mut self,
108 frame: &[u8],
109 ) -> Result<Vec<PayloadPipelineEvent>, WfbError> {
110 let Ok(frame) = WifiFrame::parse(frame, self.frame_layout) else {
111 return Ok(vec![PayloadPipelineEvent::IgnoredFrame]);
112 };
113 if !frame.matches_channel_id(self.channel_id) {
114 return Ok(vec![PayloadPipelineEvent::IgnoredFrame]);
115 }
116
117 self.push_matched_payload(frame.payload())
118 }
119
120 pub(crate) fn push_matched_payload(
122 &mut self,
123 payload: &[u8],
124 ) -> Result<Vec<PayloadPipelineEvent>, WfbError> {
125 let mut events = Vec::new();
126 self.push_matched_payload_into(payload, &mut events)?;
127 Ok(events)
128 }
129
130 pub(crate) fn push_matched_payload_into(
131 &mut self,
132 payload: &[u8],
133 events: &mut Vec<PayloadPipelineEvent>,
134 ) -> Result<(), WfbError> {
135 events.clear();
136 if let Some(receiver) = self.wfb_receiver.as_mut() {
137 let channel_id = self.channel_id;
138 receiver.push_forwarder_packet_with(payload, &mut |event| {
139 events.push(map_wfb_event(channel_id, event));
140 })?;
141 return Ok(());
142 }
143
144 match parse_forwarder_packet(payload)? {
145 WfbPacket::Data {
146 data_nonce,
147 encrypted_payload,
148 ..
149 } => {
150 let channel_id = self.channel_id;
151 self.assembler.push_decrypted_fragment_with(
152 data_nonce,
153 encrypted_payload,
154 &mut |payload| {
155 events.push(PayloadPipelineEvent::Payload(RecoveredPayload {
156 channel_id,
157 packet_seq: payload.packet_seq,
158 data: payload.payload,
159 }));
160 },
161 )?;
162 Ok(())
163 }
164 WfbPacket::SessionKey { .. } => Ok(()),
165 }
166 }
167
168 pub fn push_decrypted_80211_frame(
170 &mut self,
171 frame: &[u8],
172 decrypted_fragment: &[u8],
173 ) -> Result<Vec<PayloadPipelineEvent>, WfbError> {
174 let Ok(frame) = WifiFrame::parse(frame, self.frame_layout) else {
175 return Ok(vec![PayloadPipelineEvent::IgnoredFrame]);
176 };
177 if !frame.matches_channel_id(self.channel_id) {
178 return Ok(vec![PayloadPipelineEvent::IgnoredFrame]);
179 }
180
181 let packet = match parse_forwarder_packet(frame.payload())? {
182 WfbPacket::Data { data_nonce, .. } => data_nonce,
183 WfbPacket::SessionKey { .. } => return Ok(Vec::new()),
184 };
185 self.push_decrypted_fragment(packet, decrypted_fragment)
186 }
187
188 pub fn push_decrypted_fragment(
190 &mut self,
191 data_nonce: u64,
192 decrypted_fragment: &[u8],
193 ) -> Result<Vec<PayloadPipelineEvent>, WfbError> {
194 let payloads = self
195 .assembler
196 .push_decrypted_fragment(data_nonce, decrypted_fragment)?;
197 Ok(payloads
198 .into_iter()
199 .map(|payload| {
200 PayloadPipelineEvent::Payload(RecoveredPayload {
201 channel_id: self.channel_id,
202 packet_seq: payload.packet_seq,
203 data: payload.payload,
204 })
205 })
206 .collect())
207 }
208}
209
210fn map_wfb_event(channel_id: ChannelId, event: WfbEvent) -> PayloadPipelineEvent {
211 match event {
212 WfbEvent::Session(session) => PayloadPipelineEvent::SessionEstablished {
213 epoch: session.epoch,
214 fec_k: session.fec_k,
215 fec_n: session.fec_n,
216 },
217 WfbEvent::Payload(payload) => PayloadPipelineEvent::Payload(RecoveredPayload {
218 channel_id,
219 packet_seq: payload.packet_seq,
220 data: payload.payload,
221 }),
222 }
223}
224
225#[derive(Debug, Clone)]
233pub struct MockPayloadPipeline {
234 channel_id: ChannelId,
235}
236
237impl MockPayloadPipeline {
238 pub const fn new(channel_id: ChannelId) -> Self {
240 Self { channel_id }
241 }
242
243 pub const fn channel_id(&self) -> ChannelId {
245 self.channel_id
246 }
247
248 pub const fn fec_counters(&self) -> FecCounters {
250 FecCounters {
251 total_packets: 0,
252 recovered_packets: 0,
253 lost_packets: 0,
254 bad_packets: 0,
255 }
256 }
257
258 pub fn push_payload(&mut self, packet_seq: u64, data: &[u8]) -> Vec<PayloadPipelineEvent> {
260 vec![PayloadPipelineEvent::Payload(RecoveredPayload {
261 channel_id: self.channel_id,
262 packet_seq,
263 data: data.to_vec(),
264 })]
265 }
266}
267
268#[cfg(test)]
269mod tests {
270 use super::*;
271
272 fn plain(payload: &[u8]) -> Vec<u8> {
273 let mut out = Vec::new();
274 out.push(0);
275 out.extend_from_slice(&(payload.len() as u16).to_be_bytes());
276 out.extend_from_slice(payload);
277 out
278 }
279
280 #[test]
281 fn payload_pipeline_emits_raw_recovered_payloads() {
282 let mut pipeline =
283 PayloadPipeline::new(ChannelId::default_video(), FrameLayout::WithFcs, 1, 1).unwrap();
284 let events = pipeline
285 .push_decrypted_fragment(0, &plain(b"telemetry bytes"))
286 .unwrap();
287 assert_eq!(
288 events,
289 vec![PayloadPipelineEvent::Payload(RecoveredPayload {
290 channel_id: ChannelId::default_video(),
291 packet_seq: 0,
292 data: b"telemetry bytes".to_vec(),
293 })]
294 );
295 }
296
297 #[test]
298 fn recovered_payloads_carry_the_pipeline_channel_id() {
299 let channel_id = ChannelId::from_link_port(0x112233, crate::RadioPort::TunnelRx);
300 let mut pipeline = PayloadPipeline::new(channel_id, FrameLayout::WithFcs, 1, 1).unwrap();
301 let events = pipeline
302 .push_decrypted_fragment(0, &plain(b"data bytes"))
303 .unwrap();
304 let PayloadPipelineEvent::Payload(payload) = &events[0] else {
305 panic!("expected payload event");
306 };
307 assert_eq!(payload.channel_id, channel_id);
308 assert_eq!(payload.data, b"data bytes");
309 }
310
311 #[test]
312 fn mock_payload_pipeline_emits_recovered_payloads_without_wfb() {
313 let channel_id = ChannelId::default_video();
314 let mut pipeline = MockPayloadPipeline::new(channel_id);
315 let events = pipeline.push_payload(42, b"mock rtp bytes");
316
317 assert_eq!(
318 events,
319 vec![PayloadPipelineEvent::Payload(RecoveredPayload {
320 channel_id,
321 packet_seq: 42,
322 data: b"mock rtp bytes".to_vec(),
323 })]
324 );
325 assert_eq!(pipeline.fec_counters(), FecCounters::default());
326 }
327}