Skip to main content

seydd/input/
rtp.rs

1// Copyright 2026 Anton Gravestam
2// SPDX-License-Identifier: Apache-2.0
3//! Bare RTP over UDP — what `sim/video-source.sh` and most GStreamer/FFmpeg
4//! pipelines emit. In-house depacketizer for H.264 (RFC 6184: single NAL,
5//! STAP-A, FU-A) and H.265 (RFC 7798: single NAL, AP, FU); access units
6//! delimited by the RTP marker bit or a timestamp change. Emits Annex B.
7//! No decoding.
8//!
9//! The RTP framing — header, sequence, timestamp, marker — is identical for
10//! both. Only the payload header differs, so the codec split is confined to
11//! `push_payload` and the keyframe test in `push_nal`.
12
13use super::{Codec, VideoAu};
14use bytes::{BufMut, Bytes, BytesMut};
15use tokio::net::UdpSocket;
16use tokio::sync::mpsc;
17
18const START_CODE: [u8; 4] = [0, 0, 0, 1];
19
20pub async fn run(addr: String, codec: Codec, tx: mpsc::Sender<VideoAu>) {
21    let sock = match UdpSocket::bind(&addr).await {
22        Ok(s) => s,
23        Err(e) => {
24            tracing::error!(%addr, error = %e, "rtp input: bind failed");
25            return;
26        }
27    };
28    tracing::info!(%addr, ?codec, "rtp input listening");
29    let mut buf = vec![0u8; 65536];
30    let mut dp = Depacketizer::new(codec);
31    loop {
32        let n = match sock.recv(&mut buf).await {
33            Ok(n) => n,
34            Err(e) => {
35                tracing::warn!(error = %e, "rtp recv");
36                continue;
37            }
38        };
39        if let Some(au) = dp.push(&buf[..n]) {
40            if tx.send(au).await.is_err() {
41                return;
42            }
43        }
44    }
45}
46
47/// H.264 RTP sampling clock (RFC 6184 §8.2.1).
48const RTP_CLOCK_HZ: u64 = 90_000;
49
50#[derive(Default)]
51pub struct Depacketizer {
52    codec: Codec,
53    au: BytesMut,
54    au_ts: Option<u32>,
55    /// `au_ts` unwrapped past the 32-bit rollover (~13 h at 90 kHz), relative
56    /// to the first packet — the source's sampling timeline, which the pilot
57    /// paces presentation on.
58    au_ts_ext: u64,
59    ts_base: Option<u32>,
60    ts_wraps: u64,
61    keyframe: bool,
62    fu: Option<BytesMut>,
63    last_seq: Option<u16>,
64    lost: u16,
65}
66
67impl Depacketizer {
68    pub fn new(codec: Codec) -> Self {
69        Self {
70            codec,
71            ..Default::default()
72        }
73    }
74
75    /// Feed one RTP packet; returns a complete access unit when one closes.
76    pub fn push(&mut self, pkt: &[u8]) -> Option<VideoAu> {
77        if pkt.len() < 12 || pkt[0] >> 6 != 2 {
78            return None;
79        }
80        let cc = (pkt[0] & 0x0f) as usize;
81        let ext = pkt[0] & 0x10 != 0;
82        let marker = pkt[1] & 0x80 != 0;
83        let seq = u16::from_be_bytes([pkt[2], pkt[3]]);
84        let ts = u32::from_be_bytes([pkt[4], pkt[5], pkt[6], pkt[7]]);
85        let mut off = 12 + cc * 4;
86        if ext {
87            if pkt.len() < off + 4 {
88                return None;
89            }
90            let words = u16::from_be_bytes([pkt[off + 2], pkt[off + 3]]) as usize;
91            off += 4 + words * 4;
92        }
93        if pkt.len() <= off {
94            return None;
95        }
96        if let Some(last) = self.last_seq {
97            let gap = seq.wrapping_sub(last);
98            if gap != 1 && gap < 0x8000 {
99                self.lost = self.lost.saturating_add(gap - 1);
100                self.fu = None; // a fragment with a hole is unusable
101            }
102        }
103        self.last_seq = Some(seq);
104
105        // A timestamp change without a marker still closes the previous AU.
106        let mut finished = None;
107        if let Some(prev) = self.au_ts {
108            if prev != ts && !self.au.is_empty() {
109                finished = self.finish();
110            }
111            // Backwards by more than half the space is a rollover, not reorder.
112            if ts < prev && prev - ts > 0x8000_0000 {
113                self.ts_wraps += 1;
114            }
115        }
116        let base = *self.ts_base.get_or_insert(ts);
117        self.au_ts_ext = (self.ts_wraps << 32)
118            .wrapping_add(ts as u64)
119            .wrapping_sub(base as u64);
120        self.au_ts = Some(ts);
121
122        let payload = &pkt[off..];
123        match self.codec {
124            Codec::H264 => self.push_payload_h264(payload),
125            Codec::H265 => self.push_payload_h265(payload),
126            // Rejected in spawn_video: RFC 2435 needs JFIF reconstruction that
127            // only the RTSP path (retina) performs.
128            Codec::Mjpeg => {}
129        }
130        if marker {
131            if let Some(au) = self.finish() {
132                return Some(au);
133            }
134        }
135        finished
136    }
137
138    fn push_nal(&mut self, nal: &[u8]) {
139        if nal.is_empty() {
140            return;
141        }
142        let is_key = match self.codec {
143            // H.264: nal_unit_type in the low 5 bits; 5 is an IDR slice.
144            Codec::H264 => nal[0] & 0x1f == 5,
145            // H.265: nal_unit_type is bits 6..1 of the first header byte, and
146            // any IRAP picture (BLA_W_LP=16 … RSV_IRAP_VCL23=23) is a random
147            // access point. Narrowing this to IDR alone would miss CRA
148            // pictures, which several camera encoders emit instead.
149            Codec::H265 => nal.len() >= 2 && matches!((nal[0] >> 1) & 0x3f, 16..=23),
150            Codec::Mjpeg => true,
151        };
152        if is_key {
153            self.keyframe = true;
154        }
155        self.au.extend_from_slice(&START_CODE);
156        self.au.extend_from_slice(nal);
157    }
158
159    /// RFC 6184: single NAL unit, STAP-A aggregation, FU-A fragmentation.
160    fn push_payload_h264(&mut self, payload: &[u8]) {
161        match payload[0] & 0x1f {
162            1..=23 => self.push_nal(payload),
163            24 => {
164                // STAP-A: 1-byte header, then repeated 2-byte length + NAL.
165                let mut p = &payload[1..];
166                while p.len() >= 2 {
167                    let len = u16::from_be_bytes([p[0], p[1]]) as usize;
168                    p = &p[2..];
169                    if p.len() < len {
170                        break;
171                    }
172                    self.push_nal(&p[..len]);
173                    p = &p[len..];
174                }
175            }
176            28 => {
177                // FU-A: rebuild the 1-byte NAL header from the indicator's
178                // top 3 bits and the FU header's type.
179                if payload.len() < 2 {
180                    return;
181                }
182                let fu = payload[1];
183                if fu & 0x80 != 0 {
184                    let mut b = BytesMut::with_capacity(payload.len() + 1);
185                    b.put_u8((payload[0] & 0xe0) | (fu & 0x1f));
186                    b.extend_from_slice(&payload[2..]);
187                    self.fu = Some(b);
188                } else if let Some(b) = self.fu.as_mut() {
189                    b.extend_from_slice(&payload[2..]);
190                }
191                if fu & 0x40 != 0 {
192                    if let Some(b) = self.fu.take() {
193                        self.push_nal(&b.freeze());
194                    }
195                }
196            }
197            _ => {} // FU-B, STAP-B, MTAP: not produced by the encoders we target
198        }
199    }
200
201    /// RFC 7798: the header is two bytes, and the packet type sits where
202    /// H.264 keeps its NAL type — hence 48 (AP) and 49 (FU) rather than 24
203    /// and 28. `sprop-max-don-diff` is assumed 0, so there is no DONL field;
204    /// no camera encoder we target sets it.
205    fn push_payload_h265(&mut self, payload: &[u8]) {
206        if payload.len() < 2 {
207            return;
208        }
209        match (payload[0] >> 1) & 0x3f {
210            0..=47 => self.push_nal(payload),
211            48 => {
212                // AP: 2-byte header, then repeated 2-byte length + NAL.
213                let mut p = &payload[2..];
214                while p.len() >= 2 {
215                    let len = u16::from_be_bytes([p[0], p[1]]) as usize;
216                    p = &p[2..];
217                    if p.len() < len {
218                        break;
219                    }
220                    self.push_nal(&p[..len]);
221                    p = &p[len..];
222                }
223            }
224            49 => {
225                // FU: 2-byte header + 1-byte FU header. Rebuild the original
226                // 2-byte NAL header by putting the FU type back in place.
227                if payload.len() < 3 {
228                    return;
229                }
230                let fu = payload[2];
231                if fu & 0x80 != 0 {
232                    let mut b = BytesMut::with_capacity(payload.len());
233                    b.put_u8((payload[0] & 0x81) | ((fu & 0x3f) << 1));
234                    b.put_u8(payload[1]);
235                    b.extend_from_slice(&payload[3..]);
236                    self.fu = Some(b);
237                } else if let Some(b) = self.fu.as_mut() {
238                    b.extend_from_slice(&payload[3..]);
239                }
240                if fu & 0x40 != 0 {
241                    if let Some(b) = self.fu.take() {
242                        self.push_nal(&b.freeze());
243                    }
244                }
245            }
246            _ => {} // PACI (50) and reserved: not produced by the encoders we target
247        }
248    }
249
250    fn finish(&mut self) -> Option<VideoAu> {
251        if self.au.is_empty() {
252            return None;
253        }
254        let data: Bytes = self.au.split().freeze();
255        let au = VideoAu {
256            data,
257            keyframe: std::mem::take(&mut self.keyframe),
258            capture_ts_us: self.au_ts_ext * 1_000_000 / RTP_CLOCK_HZ,
259            input_loss: std::mem::take(&mut self.lost),
260        };
261        Some(au)
262    }
263}
264
265#[cfg(test)]
266mod tests {
267    use super::*;
268
269    fn rtp(seq: u16, ts: u32, marker: bool, payload: &[u8]) -> Vec<u8> {
270        let mut p = vec![0x80, if marker { 0x80 | 96 } else { 96 }];
271        p.extend_from_slice(&seq.to_be_bytes());
272        p.extend_from_slice(&ts.to_be_bytes());
273        p.extend_from_slice(&[0, 0, 0, 1]);
274        p.extend_from_slice(payload);
275        p
276    }
277
278    #[test]
279    fn single_nal_and_stap_a_form_one_au() {
280        let mut d = Depacketizer::new(Codec::H264);
281        // STAP-A carrying SPS (7) and PPS (8), then IDR (5) with marker.
282        let stap = [24u8, 0, 2, 0x67, 0xaa, 0, 2, 0x68, 0xbb];
283        assert!(d.push(&rtp(1, 1000, false, &stap)).is_none());
284        let au = d.push(&rtp(2, 1000, true, &[0x65, 1, 2, 3])).unwrap();
285        assert!(au.keyframe);
286        assert_eq!(
287            &au.data[..],
288            &[0, 0, 0, 1, 0x67, 0xaa, 0, 0, 0, 1, 0x68, 0xbb, 0, 0, 0, 1, 0x65, 1, 2, 3]
289        );
290    }
291
292    #[test]
293    fn fu_a_reassembles_and_loss_counts() {
294        let mut d = Depacketizer::new(Codec::H264);
295        // NAL type 1 (non-IDR), nri 2 → indicator 0x41, header bits
296        assert!(d.push(&rtp(10, 2000, false, &[0x5c, 0x81, 9, 9])).is_none()); // start
297        assert!(d.push(&rtp(11, 2000, false, &[0x5c, 0x01, 8])).is_none());
298        let au = d.push(&rtp(12, 2000, true, &[0x5c, 0x41, 7])).unwrap(); // end
299        assert!(!au.keyframe);
300        assert_eq!(&au.data[..], &[0, 0, 0, 1, 0x41, 9, 9, 8, 7]);
301        // gap: seq 12 → 15
302        let au2 = d.push(&rtp(15, 3000, true, &[0x41, 1])).unwrap();
303        assert_eq!(au2.input_loss, 2);
304    }
305
306    #[test]
307    fn timestamp_change_closes_au_without_marker() {
308        let mut d = Depacketizer::new(Codec::H265);
309        assert!(d.push(&rtp(1, 100, false, &[0x41, 1])).is_none());
310        let au = d.push(&rtp(2, 200, false, &[0x41, 2])).unwrap();
311        assert_eq!(&au.data[..], &[0, 0, 0, 1, 0x41, 1]);
312    }
313}
314
315// ── H.265 (RFC 7798) ────────────────────────────────────────────────────────
316#[cfg(test)]
317mod h265_tests {
318    use super::*;
319
320    fn rtp(seq: u16, ts: u32, marker: bool, payload: &[u8]) -> Vec<u8> {
321        let mut p = vec![0x80, if marker { 0x80 | 96 } else { 96 }];
322        p.extend_from_slice(&seq.to_be_bytes());
323        p.extend_from_slice(&ts.to_be_bytes());
324        p.extend_from_slice(&[0, 0, 0, 1]);
325        p.extend_from_slice(payload);
326        p
327    }
328
329    /// Two-byte H.265 NAL header for a given nal_unit_type, layer 0, tid 1.
330    fn hdr(nal_type: u8) -> [u8; 2] {
331        [nal_type << 1, 0x01]
332    }
333
334    #[test]
335    fn aggregation_packet_then_idr_forms_one_keyframe_au() {
336        let mut d = Depacketizer::new(Codec::H265);
337        // AP (type 48) carrying VPS(32), SPS(33), PPS(34).
338        let mut ap = vec![48u8 << 1, 0x01];
339        for t in [32u8, 33, 34] {
340            ap.extend_from_slice(&[0, 3]);
341            ap.extend_from_slice(&hdr(t));
342            ap.push(0xaa);
343        }
344        assert!(d.push(&rtp(1, 1000, false, &ap)).is_none());
345
346        // IDR_W_RADL is type 19 — an IRAP, so a keyframe.
347        let mut idr = hdr(19).to_vec();
348        idr.extend_from_slice(&[1, 2, 3]);
349        let au = d.push(&rtp(2, 1000, true, &idr)).unwrap();
350        assert!(
351            au.keyframe,
352            "IDR_W_RADL must mark the access unit as a keyframe"
353        );
354        assert_eq!(
355            &au.data[..],
356            &[
357                0, 0, 0, 1, 64, 1, 0xaa, // VPS
358                0, 0, 0, 1, 66, 1, 0xaa, // SPS
359                0, 0, 0, 1, 68, 1, 0xaa, // PPS
360                0, 0, 0, 1, 38, 1, 1, 2, 3, // IDR
361            ]
362        );
363    }
364
365    #[test]
366    fn cra_counts_as_a_keyframe_but_a_trailing_picture_does_not() {
367        // Several camera encoders open a GOP with CRA (21) rather than IDR;
368        // treating only IDR as a random access point would strand them.
369        let mut d = Depacketizer::new(Codec::H265);
370        let mut cra = hdr(21).to_vec();
371        cra.push(9);
372        assert!(d.push(&rtp(1, 100, true, &cra)).unwrap().keyframe);
373
374        let mut trail = hdr(1).to_vec(); // TRAIL_R
375        trail.push(9);
376        assert!(!d.push(&rtp(2, 200, true, &trail)).unwrap().keyframe);
377    }
378
379    #[test]
380    fn fragmentation_unit_rebuilds_the_two_byte_header() {
381        let mut d = Depacketizer::new(Codec::H265);
382        // FU (49) carrying a fragmented IDR (19): payload header, FU header, data.
383        let fu = |s: bool, e: bool, data: &[u8]| {
384            let mut p = vec![49u8 << 1, 0x01];
385            p.push((if s { 0x80 } else { 0 }) | (if e { 0x40 } else { 0 }) | 19);
386            p.extend_from_slice(data);
387            p
388        };
389        assert!(d
390            .push(&rtp(10, 500, false, &fu(true, false, &[7, 7])))
391            .is_none());
392        assert!(d
393            .push(&rtp(11, 500, false, &fu(false, false, &[8])))
394            .is_none());
395        let au = d.push(&rtp(12, 500, true, &fu(false, true, &[9]))).unwrap();
396        assert!(au.keyframe);
397        // Reassembled as one NAL with the original IDR header restored.
398        assert_eq!(&au.data[..], &[0, 0, 0, 1, 38, 1, 7, 7, 8, 9]);
399    }
400
401    #[test]
402    fn a_hole_in_a_fragment_discards_it_rather_than_emitting_a_torn_nal() {
403        let mut d = Depacketizer::new(Codec::H265);
404        let mut start = vec![49u8 << 1, 0x01, 0x80 | 19];
405        start.extend_from_slice(&[1, 2]);
406        assert!(d.push(&rtp(1, 900, false, &start)).is_none());
407        // seq jumps 1 → 3: the fragment is unusable and must be dropped.
408        let mut end = vec![49u8 << 1, 0x01, 0x40 | 19];
409        end.push(3);
410        assert!(d.push(&rtp(3, 900, true, &end)).is_none());
411    }
412
413    #[test]
414    fn codec_strings_map_to_the_right_framing() {
415        assert_eq!(
416            Codec::from_codec_string("avc1.42001f").unwrap(),
417            Codec::H264
418        );
419        assert_eq!(
420            Codec::from_codec_string("hev1.1.6.L93.B0").unwrap(),
421            Codec::H265
422        );
423        assert_eq!(
424            Codec::from_codec_string("hvc1.1.6.L93.B0").unwrap(),
425            Codec::H265
426        );
427        assert_eq!(Codec::from_codec_string("H265").unwrap(), Codec::H265);
428        assert!(Codec::from_codec_string("vp09.00.10.08").is_err());
429    }
430}