1#[derive(Debug, Clone, Copy, PartialEq, Eq)]
2pub enum RtpError {
3 TooShort,
4 InvalidVersion(u8),
5 InvalidExtension,
6 InvalidPadding,
7 EmptyPayload,
8 UnsupportedPayload,
9 FragmentOverflow,
10}
11
12#[derive(Debug, Clone, Copy, PartialEq, Eq)]
13pub enum Codec {
14 H264,
15 H265,
16}
17
18pub const RTP_PAYLOAD_TYPE_H264: u8 = 96;
19pub const RTP_PAYLOAD_TYPE_H265: u8 = 97;
20
21#[derive(Debug, Clone, Copy, PartialEq, Eq)]
22pub struct RtpHeader {
23 pub marker: bool,
24 pub payload_type: u8,
25 pub sequence_number: u16,
26 pub timestamp: u32,
27 pub ssrc: u32,
28 pub csrc_count: u8,
29 pub has_extension: bool,
30 pub header_len: usize,
31 pub payload_len: usize,
32}
33
34impl RtpHeader {
35 pub fn parse(packet: &[u8]) -> Result<Self, RtpError> {
36 if packet.len() < 12 {
37 return Err(RtpError::TooShort);
38 }
39 let version = packet[0] >> 6;
40 if version != 2 {
41 return Err(RtpError::InvalidVersion(version));
42 }
43
44 let padding = packet[0] & 0x20 != 0;
45 let extension = packet[0] & 0x10 != 0;
46 let csrc_count = packet[0] & 0x0f;
47 let mut header_len = 12 + csrc_count as usize * 4;
48 if packet.len() < header_len {
49 return Err(RtpError::TooShort);
50 }
51
52 if extension {
53 if packet.len() < header_len + 4 {
54 return Err(RtpError::InvalidExtension);
55 }
56 let ext_words =
57 u16::from_be_bytes([packet[header_len + 2], packet[header_len + 3]]) as usize;
58 header_len += 4 + ext_words * 4;
59 if packet.len() < header_len {
60 return Err(RtpError::InvalidExtension);
61 }
62 }
63
64 let padding_len = if padding {
65 let len = *packet.last().ok_or(RtpError::InvalidPadding)? as usize;
66 if len == 0 || len > packet.len() - header_len {
67 return Err(RtpError::InvalidPadding);
68 }
69 len
70 } else {
71 0
72 };
73
74 let payload_len = packet.len() - header_len - padding_len;
75 if payload_len == 0 {
76 return Err(RtpError::EmptyPayload);
77 }
78
79 Ok(Self {
80 marker: packet[1] & 0x80 != 0,
81 payload_type: packet[1] & 0x7f,
82 sequence_number: u16::from_be_bytes([packet[2], packet[3]]),
83 timestamp: u32::from_be_bytes([packet[4], packet[5], packet[6], packet[7]]),
84 ssrc: u32::from_be_bytes([packet[8], packet[9], packet[10], packet[11]]),
85 csrc_count,
86 has_extension: extension,
87 header_len,
88 payload_len,
89 })
90 }
91
92 pub fn payload<'a>(&self, packet: &'a [u8]) -> &'a [u8] {
93 &packet[self.header_len..self.header_len + self.payload_len]
94 }
95}
96
97#[derive(Debug, Clone, PartialEq, Eq)]
98pub struct DepacketizedFrame {
99 pub data: Vec<u8>,
100 pub timestamp: u32,
101 pub is_keyframe: bool,
102 pub codec: Codec,
103}
104
105#[derive(Debug, Default, Clone)]
106struct FragmentState {
107 data: Vec<u8>,
108 timestamp: u32,
109}
110
111#[derive(Debug, Clone)]
112pub struct RtpDepacketizer {
113 h264: FragmentState,
114 h265: FragmentState,
115 h264_sps: Option<Vec<u8>>,
116 h264_pps: Option<Vec<u8>>,
117 h265_vps: Option<Vec<u8>>,
118 h265_sps: Option<Vec<u8>>,
119 h265_pps: Option<Vec<u8>>,
120 max_fragment: usize,
121}
122
123impl Default for RtpDepacketizer {
124 fn default() -> Self {
125 Self::new()
126 }
127}
128
129impl RtpDepacketizer {
130 pub fn new() -> Self {
131 Self {
132 h264: FragmentState::default(),
133 h265: FragmentState::default(),
134 h264_sps: None,
135 h264_pps: None,
136 h265_vps: None,
137 h265_sps: None,
138 h265_pps: None,
139 max_fragment: 1024 * 1024,
140 }
141 }
142
143 pub fn push(&mut self, packet: &[u8]) -> Result<Option<DepacketizedFrame>, RtpError> {
144 let header = RtpHeader::parse(packet)?;
145 let payload = header.payload(packet);
146 let codec = codec_from_payload_type(header.payload_type)
147 .or_else(|| detect_codec(payload))
148 .ok_or(RtpError::UnsupportedPayload)?;
149 match codec {
150 Codec::H264 => self.push_h264(payload, header),
151 Codec::H265 => self.push_h265(payload, header),
152 }
153 }
154
155 fn push_h264(
156 &mut self,
157 payload: &[u8],
158 header: RtpHeader,
159 ) -> Result<Option<DepacketizedFrame>, RtpError> {
160 let nal_type = payload[0] & 0x1f;
161 match nal_type {
162 7 => {
163 self.h264_sps = Some(payload.to_vec());
164 Ok(None)
165 }
166 8 => {
167 self.h264_pps = Some(payload.to_vec());
168 Ok(None)
169 }
170 24 => self.h264_stap_a(payload, header.timestamp),
171 28 => self.h264_fu_a(payload, header.timestamp),
172 _ if self.has_decoder_config(Codec::H264) => Ok(Some(self.frame_with_prefix(
173 payload,
174 header.timestamp,
175 nal_type == 5,
176 Codec::H264,
177 ))),
178 _ => Ok(None),
179 }
180 }
181
182 fn push_h265(
183 &mut self,
184 payload: &[u8],
185 header: RtpHeader,
186 ) -> Result<Option<DepacketizedFrame>, RtpError> {
187 if payload.len() < 2 {
188 return Err(RtpError::UnsupportedPayload);
189 }
190 let nal_type = (payload[0] >> 1) & 0x3f;
191 match nal_type {
192 32 => {
193 self.h265_vps = Some(payload.to_vec());
194 Ok(None)
195 }
196 33 => {
197 self.h265_sps = Some(payload.to_vec());
198 Ok(None)
199 }
200 34 => {
201 self.h265_pps = Some(payload.to_vec());
202 Ok(None)
203 }
204 48 => self.h265_ap(payload, header.timestamp),
205 49 => self.h265_fu(payload, header.timestamp),
206 _ if self.has_decoder_config(Codec::H265) => Ok(Some(self.frame_with_prefix(
207 payload,
208 header.timestamp,
209 (16..=23).contains(&nal_type),
210 Codec::H265,
211 ))),
212 _ => Ok(None),
213 }
214 }
215
216 fn h264_fu_a(
217 &mut self,
218 payload: &[u8],
219 timestamp: u32,
220 ) -> Result<Option<DepacketizedFrame>, RtpError> {
221 if payload.len() < 2 {
222 return Err(RtpError::UnsupportedPayload);
223 }
224 let fu_indicator = payload[0];
225 let fu_header = payload[1];
226 let start = fu_header & 0x80 != 0;
227 let end = fu_header & 0x40 != 0;
228 let nal_type = fu_header & 0x1f;
229 if start {
230 self.h264.data.clear();
231 self.h264.timestamp = timestamp;
232 self.h264.data.push((fu_indicator & 0xe0) | nal_type);
233 }
234 self.append_fragment(Codec::H264, &payload[2..])?;
235 if end {
236 if !self.has_decoder_config(Codec::H264) {
237 self.h264.data.clear();
238 return Ok(None);
239 }
240 let data = self.h264.data.clone();
241 Ok(Some(self.frame_with_prefix(
242 &data,
243 self.h264.timestamp,
244 nal_type == 5,
245 Codec::H264,
246 )))
247 } else {
248 Ok(None)
249 }
250 }
251
252 fn h265_fu(
253 &mut self,
254 payload: &[u8],
255 timestamp: u32,
256 ) -> Result<Option<DepacketizedFrame>, RtpError> {
257 if payload.len() < 3 {
258 return Err(RtpError::UnsupportedPayload);
259 }
260 let fu_header = payload[2];
261 let start = fu_header & 0x80 != 0;
262 let end = fu_header & 0x40 != 0;
263 let nal_type = fu_header & 0x3f;
264 if start {
265 self.h265.data.clear();
266 self.h265.timestamp = timestamp;
267 self.h265.data.push((nal_type << 1) | (payload[0] & 0x01));
268 self.h265.data.push(payload[1]);
269 }
270 self.append_fragment(Codec::H265, &payload[3..])?;
271 if end {
272 if !self.has_decoder_config(Codec::H265) {
273 self.h265.data.clear();
274 return Ok(None);
275 }
276 let data = self.h265.data.clone();
277 Ok(Some(self.frame_with_prefix(
278 &data,
279 self.h265.timestamp,
280 (16..=23).contains(&nal_type),
281 Codec::H265,
282 )))
283 } else {
284 Ok(None)
285 }
286 }
287
288 fn h264_stap_a(
289 &mut self,
290 payload: &[u8],
291 timestamp: u32,
292 ) -> Result<Option<DepacketizedFrame>, RtpError> {
293 let mut out = Vec::new();
294 let mut offset = 1;
295 let mut keyframe = false;
296 let mut has_slice = false;
297 while offset + 2 <= payload.len() {
298 let len = u16::from_be_bytes([payload[offset], payload[offset + 1]]) as usize;
299 offset += 2;
300 if offset + len > payload.len() {
301 break;
302 }
303 let nalu = &payload[offset..offset + len];
304 let nal_type = nalu.first().map(|b| b & 0x1f).unwrap_or(0);
305 match nal_type {
306 7 => self.h264_sps = Some(nalu.to_vec()),
307 8 => self.h264_pps = Some(nalu.to_vec()),
308 _ => {}
309 }
310 has_slice |= (1..=5).contains(&nal_type);
311 keyframe |= nal_type == 5;
312 append_annex_b(&mut out, nalu);
313 offset += len;
314 }
315 if !has_slice || !self.has_decoder_config(Codec::H264) {
316 return Ok(None);
317 }
318 Ok(Some(DepacketizedFrame {
319 data: out,
320 timestamp,
321 is_keyframe: keyframe,
322 codec: Codec::H264,
323 }))
324 }
325
326 fn h265_ap(
327 &mut self,
328 payload: &[u8],
329 timestamp: u32,
330 ) -> Result<Option<DepacketizedFrame>, RtpError> {
331 let mut out = Vec::new();
332 let mut offset = 2;
333 let mut keyframe = false;
334 let mut has_slice = false;
335 while offset + 2 <= payload.len() {
336 let len = u16::from_be_bytes([payload[offset], payload[offset + 1]]) as usize;
337 offset += 2;
338 if offset + len > payload.len() {
339 break;
340 }
341 let nalu = &payload[offset..offset + len];
342 let nal_type = nalu.first().map(|b| (b >> 1) & 0x3f).unwrap_or(0);
343 match nal_type {
344 32 => self.h265_vps = Some(nalu.to_vec()),
345 33 => self.h265_sps = Some(nalu.to_vec()),
346 34 => self.h265_pps = Some(nalu.to_vec()),
347 _ => {}
348 }
349 has_slice |= !nalu.is_empty() && nal_type <= 31;
350 keyframe |= (16..=23).contains(&nal_type);
351 append_annex_b(&mut out, nalu);
352 offset += len;
353 }
354 if !has_slice || !self.has_decoder_config(Codec::H265) {
355 return Ok(None);
356 }
357 Ok(Some(DepacketizedFrame {
358 data: out,
359 timestamp,
360 is_keyframe: keyframe,
361 codec: Codec::H265,
362 }))
363 }
364
365 fn append_fragment(&mut self, codec: Codec, bytes: &[u8]) -> Result<(), RtpError> {
366 let state = match codec {
367 Codec::H264 => &mut self.h264,
368 Codec::H265 => &mut self.h265,
369 };
370 if state.data.len() + bytes.len() > self.max_fragment {
371 return Err(RtpError::FragmentOverflow);
372 }
373 state.data.extend_from_slice(bytes);
374 Ok(())
375 }
376
377 fn frame_with_prefix(
378 &self,
379 nalu: &[u8],
380 timestamp: u32,
381 is_keyframe: bool,
382 codec: Codec,
383 ) -> DepacketizedFrame {
384 let mut data = Vec::new();
385 if is_keyframe {
386 match codec {
387 Codec::H264 => {
388 if let Some(sps) = &self.h264_sps {
389 append_annex_b(&mut data, sps);
390 }
391 if let Some(pps) = &self.h264_pps {
392 append_annex_b(&mut data, pps);
393 }
394 }
395 Codec::H265 => {
396 if let Some(vps) = &self.h265_vps {
397 append_annex_b(&mut data, vps);
398 }
399 if let Some(sps) = &self.h265_sps {
400 append_annex_b(&mut data, sps);
401 }
402 if let Some(pps) = &self.h265_pps {
403 append_annex_b(&mut data, pps);
404 }
405 }
406 }
407 }
408 append_annex_b(&mut data, nalu);
409 DepacketizedFrame {
410 data,
411 timestamp,
412 is_keyframe,
413 codec,
414 }
415 }
416
417 fn has_decoder_config(&self, codec: Codec) -> bool {
418 match codec {
419 Codec::H264 => self.h264_sps.is_some() && self.h264_pps.is_some(),
420 Codec::H265 => {
421 self.h265_vps.is_some() && self.h265_sps.is_some() && self.h265_pps.is_some()
422 }
423 }
424 }
425}
426
427fn codec_from_payload_type(payload_type: u8) -> Option<Codec> {
428 match payload_type {
429 RTP_PAYLOAD_TYPE_H264 => Some(Codec::H264),
430 RTP_PAYLOAD_TYPE_H265 => Some(Codec::H265),
431 _ => None,
432 }
433}
434
435fn detect_codec(payload: &[u8]) -> Option<Codec> {
436 if payload.is_empty() {
437 return None;
438 }
439 if payload.len() >= 2 {
440 let h265_nal_type = (payload[0] >> 1) & 0x3f;
441 if h265_nal_type == 48 || h265_nal_type == 49 || (32..=40).contains(&h265_nal_type) {
442 return Some(Codec::H265);
443 }
444 }
445 let h264_nal_type = payload[0] & 0x1f;
446 if h264_nal_type == 24 || h264_nal_type == 28 || (1..=12).contains(&h264_nal_type) {
447 return Some(Codec::H264);
448 }
449 None
450}
451
452fn append_annex_b(out: &mut Vec<u8>, nalu: &[u8]) {
453 out.extend_from_slice(&[0, 0, 0, 1]);
454 out.extend_from_slice(nalu);
455}
456
457#[cfg(test)]
458mod tests {
459 use super::*;
460
461 fn rtp(payload: &[u8], marker: bool, seq: u16, timestamp: u32) -> Vec<u8> {
462 rtp_with_payload_type(payload, RTP_PAYLOAD_TYPE_H264, marker, seq, timestamp)
463 }
464
465 fn rtp_with_payload_type(
466 payload: &[u8],
467 payload_type: u8,
468 marker: bool,
469 seq: u16,
470 timestamp: u32,
471 ) -> Vec<u8> {
472 let mut packet = vec![
473 0x80,
474 (if marker { 0x80 } else { 0x00 }) | (payload_type & 0x7f),
475 ];
476 packet.extend_from_slice(&seq.to_be_bytes());
477 packet.extend_from_slice(×tamp.to_be_bytes());
478 packet.extend_from_slice(&0x1122_3344u32.to_be_bytes());
479 packet.extend_from_slice(payload);
480 packet
481 }
482
483 fn stap_a(units: &[&[u8]]) -> Vec<u8> {
484 let mut payload = vec![24];
485 for unit in units {
486 payload.extend_from_slice(&(unit.len() as u16).to_be_bytes());
487 payload.extend_from_slice(unit);
488 }
489 payload
490 }
491
492 fn prime_h264(depay: &mut RtpDepacketizer) {
493 assert!(depay
494 .push(&rtp(&[0x67, 0x64, 0x00, 0x1f], true, 1, 10))
495 .unwrap()
496 .is_none());
497 assert!(depay
498 .push(&rtp(&[0x68, 0xee], true, 2, 10))
499 .unwrap()
500 .is_none());
501 }
502
503 fn prime_h265(depay: &mut RtpDepacketizer) {
504 for (seq, payload) in [
505 (1, &[0x40, 0x01, 0xaa][..]),
506 (2, &[0x42, 0x01, 0xbb][..]),
507 (3, &[0x44, 0x01, 0xcc][..]),
508 ] {
509 assert!(depay
510 .push(&rtp_with_payload_type(
511 payload,
512 RTP_PAYLOAD_TYPE_H265,
513 true,
514 seq,
515 10,
516 ))
517 .unwrap()
518 .is_none());
519 }
520 }
521
522 #[test]
523 fn parses_rtp_header() {
524 let packet = rtp(&[0x65, 1, 2], true, 7, 1234);
525 let header = RtpHeader::parse(&packet).unwrap();
526 assert!(header.marker);
527 assert_eq!(header.payload_type, 96);
528 assert_eq!(header.sequence_number, 7);
529 assert_eq!(header.timestamp, 1234);
530 assert_eq!(header.payload(&packet), &[0x65, 1, 2]);
531 }
532
533 #[test]
534 fn depacketizes_h264_single_nalu_as_annex_b() {
535 let mut depay = RtpDepacketizer::new();
536 prime_h264(&mut depay);
537 let frame = depay
538 .push(&rtp(&[0x65, 0xaa], true, 1, 42))
539 .unwrap()
540 .unwrap();
541 assert_eq!(frame.codec, Codec::H264);
542 assert!(frame.is_keyframe);
543 assert_eq!(
544 frame.data,
545 [
546 &[0, 0, 0, 1, 0x67, 0x64, 0x00, 0x1f][..],
547 &[0, 0, 0, 1, 0x68, 0xee][..],
548 &[0, 0, 0, 1, 0x65, 0xaa][..],
549 ]
550 .concat()
551 );
552 }
553
554 #[test]
555 fn drops_h264_video_until_sps_and_pps_are_seen() {
556 let mut depay = RtpDepacketizer::new();
557 assert!(depay
558 .push(&rtp(&[0x65, 0xaa], true, 1, 42))
559 .unwrap()
560 .is_none());
561 }
562
563 #[test]
564 fn h264_payload_type_prevents_h265_false_positive() {
565 let mut depay = RtpDepacketizer::new();
566 prime_h264(&mut depay);
567 let frame = depay
568 .push(&rtp(&[0x41, 0xaa], true, 1, 42))
569 .unwrap()
570 .unwrap();
571 assert_eq!(frame.codec, Codec::H264);
572 assert!(!frame.is_keyframe);
573 assert_eq!(frame.data, &[0, 0, 0, 1, 0x41, 0xaa]);
574 }
575
576 #[test]
577 fn depacketizes_h265_single_nalu_by_payload_type() {
578 let mut depay = RtpDepacketizer::new();
579 prime_h265(&mut depay);
580 let frame = depay
581 .push(&rtp_with_payload_type(
582 &[0x02, 0x01, 0xaa],
583 RTP_PAYLOAD_TYPE_H265,
584 true,
585 1,
586 42,
587 ))
588 .unwrap()
589 .unwrap();
590 assert_eq!(frame.codec, Codec::H265);
591 assert!(!frame.is_keyframe);
592 assert_eq!(frame.data, &[0, 0, 0, 1, 0x02, 0x01, 0xaa]);
593 }
594
595 #[test]
596 fn h264_stap_a_caches_parameter_sets_for_later_keyframe() {
597 let mut depay = RtpDepacketizer::new();
598 let sps = &[0x67, 0x64, 0x00, 0x1f][..];
599 let pps = &[0x68, 0xee][..];
600 let aggregate = depay.push(&rtp(&stap_a(&[sps, pps]), true, 1, 10)).unwrap();
601 assert!(aggregate.is_none());
602
603 let frame = depay
604 .push(&rtp(&[0x65, 0xaa], true, 2, 20))
605 .unwrap()
606 .unwrap();
607 assert!(frame.is_keyframe);
608 assert_eq!(
609 frame.data,
610 [
611 &[0, 0, 0, 1][..],
612 sps,
613 &[0, 0, 0, 1][..],
614 pps,
615 &[0, 0, 0, 1, 0x65, 0xaa][..],
616 ]
617 .concat()
618 );
619 }
620
621 #[test]
622 fn depacketizes_h264_fu_a() {
623 let mut depay = RtpDepacketizer::new();
624 prime_h264(&mut depay);
625 assert!(depay
626 .push(&rtp(&[0x7c, 0x85, 1, 2], false, 1, 99))
627 .unwrap()
628 .is_none());
629 let frame = depay
630 .push(&rtp(&[0x7c, 0x45, 3, 4], true, 2, 99))
631 .unwrap()
632 .unwrap();
633 assert_eq!(
634 frame.data,
635 [
636 &[0, 0, 0, 1, 0x67, 0x64, 0x00, 0x1f][..],
637 &[0, 0, 0, 1, 0x68, 0xee][..],
638 &[0, 0, 0, 1, 0x65, 1, 2, 3, 4][..],
639 ]
640 .concat()
641 );
642 }
643}