1use 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
47const 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_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 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; }
102 }
103 self.last_seq = Some(seq);
104
105 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 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 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 Codec::H264 => nal[0] & 0x1f == 5,
145 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 fn push_payload_h264(&mut self, payload: &[u8]) {
161 match payload[0] & 0x1f {
162 1..=23 => self.push_nal(payload),
163 24 => {
164 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 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 _ => {} }
199 }
200
201 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 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 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 _ => {} }
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 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 assert!(d.push(&rtp(10, 2000, false, &[0x5c, 0x81, 9, 9])).is_none()); 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(); assert!(!au.keyframe);
300 assert_eq!(&au.data[..], &[0, 0, 0, 1, 0x41, 9, 9, 8, 7]);
301 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#[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 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 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 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, 0, 0, 0, 1, 66, 1, 0xaa, 0, 0, 0, 1, 68, 1, 0xaa, 0, 0, 0, 1, 38, 1, 1, 2, 3, ]
362 );
363 }
364
365 #[test]
366 fn cra_counts_as_a_keyframe_but_a_trailing_picture_does_not() {
367 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.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 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 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 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}