1mod export;
14mod import;
15mod split;
16
17pub use export::*;
18pub use import::*;
19pub use split::*;
20
21use bytes::{Buf, BufMut, Bytes, BytesMut};
22
23const NAL_TYPE_SPS: u8 = 7;
25const NAL_TYPE_PPS: u8 = 8;
26
27pub(crate) fn avc1_frame(
35 data: &[u8],
36 length_size: usize,
37 pts: crate::container::Timestamp,
38) -> crate::Result<crate::container::Frame> {
39 Ok(crate::container::Frame {
40 timestamp: pts,
41 payload: data.to_vec().into(),
42 keyframe: avc1_is_keyframe(data, length_size),
43 duration: None,
44 })
45}
46
47fn avc1_is_keyframe(data: &[u8], length_size: usize) -> bool {
49 let mut offset = 0;
50 while offset + length_size <= data.len() {
51 let nal_len = match length_size {
52 1 => data[offset] as usize,
53 2 => u16::from_be_bytes([data[offset], data[offset + 1]]) as usize,
54 3 => u32::from_be_bytes([0, data[offset], data[offset + 1], data[offset + 2]]) as usize,
55 4 => u32::from_be_bytes([data[offset], data[offset + 1], data[offset + 2], data[offset + 3]]) as usize,
56 _ => return false,
57 };
58 offset += length_size;
59 if offset + nal_len > data.len() {
60 break;
61 }
62 if nal_len > 0 && data[offset] & 0x1f == 5 {
63 return true; }
65 offset += nal_len;
66 }
67 false
68}
69
70#[derive(Debug, Clone, thiserror::Error)]
72#[non_exhaustive]
73pub enum Error {
74 #[error("SPS NAL too short")]
75 SpsTooShort,
76
77 #[error("failed to parse SPS")]
78 SpsParse,
79
80 #[error("AVCDecoderConfigurationRecord too short")]
81 AvccTooShort,
82
83 #[error("AVCDecoderConfigurationRecord truncated")]
84 AvccTruncated,
85
86 #[error("avc1 description for rendition {name:?} is missing SPS or PPS (sps={sps}, pps={pps})")]
87 MissingParamSets { name: String, sps: usize, pps: usize },
88
89 #[error("SPS too large for avcC length field ({0} > {max})", max = u16::MAX)]
90 SpsTooLarge(usize),
91
92 #[error("PPS too large for avcC length field ({0} > {max})", max = u16::MAX)]
93 PpsTooLarge(usize),
94
95 #[error("avcC requires at least one SPS")]
96 MissingSps,
97
98 #[error("too many SPS for avcC ({0} > 31)")]
99 TooManySps(usize),
100
101 #[error("too many PPS for avcC ({0} > 255)")]
102 TooManyPps(usize),
103
104 #[error("NAL too large for 4-byte length prefix")]
105 NalTooLarge,
106
107 #[error("NAL unit is too short")]
108 NalTooShort,
109
110 #[error("forbidden zero bit is not zero")]
111 ForbiddenZeroBit,
112
113 #[error("not initialized")]
114 NotInitialized,
115
116 #[error("avc3 track not created")]
117 Avc3TrackNotCreated,
118
119 #[error("missing timestamp")]
120 MissingTimestamp,
121
122 #[error("annexb: {0}")]
123 Annexb(#[from] crate::codec::annexb::Error),
124}
125
126pub type Result<T> = std::result::Result<T, Error>;
127
128#[derive(Debug, Clone)]
134pub struct Sps {
135 pub profile: u8,
136 pub constraints: u8,
137 pub level: u8,
138 pub coded_width: u32,
139 pub coded_height: u32,
140}
141
142impl Sps {
143 pub fn parse(nal: &[u8]) -> Result<Self> {
145 if nal.len() < 4 {
146 return Err(Error::SpsTooShort);
147 }
148 let rbsp = h264_parser::nal::ebsp_to_rbsp(&nal[1..]);
149 let sps = h264_parser::Sps::parse(&rbsp).map_err(|_| Error::SpsParse)?;
150 Ok(Self {
151 profile: sps.profile_idc,
152 constraints: pack_constraint_flags(&sps),
153 level: sps.level_idc,
154 coded_width: sps.width,
155 coded_height: sps.height,
156 })
157 }
158}
159
160#[derive(Debug, Clone)]
166#[non_exhaustive]
167pub struct Avcc {
168 pub profile: u8,
170 pub constraints: u8,
172 pub level: u8,
174 pub length_size: usize,
176 pub sps: Vec<Bytes>,
178 pub pps: Vec<Bytes>,
180 pub coded_width: Option<u32>,
182 pub coded_height: Option<u32>,
183}
184
185impl Avcc {
186 pub fn parse(avcc: &[u8]) -> Result<Self> {
188 if avcc.len() < 7 {
189 return Err(Error::AvccTooShort);
190 }
191
192 let profile = avcc[1];
193 let constraints = avcc[2];
194 let level = avcc[3];
195 let length_size = (avcc[4] & 0x03) as usize + 1;
196 let num_sps = (avcc[5] & 0x1f) as usize;
197
198 let mut pos = 6;
199 let sps = read_param_sets(avcc, &mut pos, num_sps)?;
200
201 if avcc.len() <= pos {
202 return Err(Error::AvccTruncated);
203 }
204 let num_pps = avcc[pos] as usize;
205 pos += 1;
206 let pps = read_param_sets(avcc, &mut pos, num_pps)?;
207
208 let (mut coded_width, mut coded_height) = (None, None);
210 if let Some(first) = sps.first()
211 && first.len() > 1
212 && let Ok(parsed) = Sps::parse(first)
213 {
214 coded_width = Some(parsed.coded_width);
215 coded_height = Some(parsed.coded_height);
216 }
217
218 Ok(Self {
219 profile,
220 constraints,
221 level,
222 length_size,
223 sps,
224 pps,
225 coded_width,
226 coded_height,
227 })
228 }
229}
230
231fn pack_constraint_flags(sps: &h264_parser::Sps) -> u8 {
232 ((sps.constraint_set0_flag as u8) << 7)
233 | ((sps.constraint_set1_flag as u8) << 6)
234 | ((sps.constraint_set2_flag as u8) << 5)
235 | ((sps.constraint_set3_flag as u8) << 4)
236 | ((sps.constraint_set4_flag as u8) << 3)
237 | ((sps.constraint_set5_flag as u8) << 2)
238}
239
240pub(crate) fn build_avcc(sps_nals: &[Bytes], pps_nals: &[Bytes]) -> Result<Bytes> {
246 let first_sps = sps_nals.first().ok_or(Error::MissingSps)?;
247 if first_sps.len() < 4 {
248 return Err(Error::SpsTooShort);
249 }
250 if sps_nals.len() > 0x1f {
252 return Err(Error::TooManySps(sps_nals.len()));
253 }
254 if pps_nals.len() > u8::MAX as usize {
255 return Err(Error::TooManyPps(pps_nals.len()));
256 }
257 for sps in sps_nals {
258 if sps.len() > u16::MAX as usize {
259 return Err(Error::SpsTooLarge(sps.len()));
260 }
261 }
262 for pps in pps_nals {
263 if pps.len() > u16::MAX as usize {
264 return Err(Error::PpsTooLarge(pps.len()));
265 }
266 }
267
268 let profile_idc = first_sps[1];
269 let constraints = first_sps[2];
270 let level_idc = first_sps[3];
271
272 let payload: usize = sps_nals.iter().chain(pps_nals).map(|n| 2 + n.len()).sum();
273 let mut out = BytesMut::with_capacity(7 + payload);
274 out.put_u8(1); out.put_u8(profile_idc);
276 out.put_u8(constraints);
277 out.put_u8(level_idc);
278 out.put_u8(0xff); out.put_u8(0xe0 | sps_nals.len() as u8); for sps in sps_nals {
281 out.put_u16(sps.len() as u16);
282 out.put_slice(sps);
283 }
284 out.put_u8(pps_nals.len() as u8); for pps in pps_nals {
286 out.put_u16(pps.len() as u16);
287 out.put_slice(pps);
288 }
289 Ok(out.freeze())
290}
291
292fn read_param_sets(buf: &[u8], pos: &mut usize, count: usize) -> Result<Vec<Bytes>> {
296 let mut out = Vec::with_capacity(count);
297 for _ in 0..count {
298 let after_len = pos.checked_add(2).ok_or(Error::AvccTruncated)?;
299 if buf.len() < after_len {
300 return Err(Error::AvccTruncated);
301 }
302 let len = u16::from_be_bytes([buf[*pos], buf[*pos + 1]]) as usize;
303 let after_nal = after_len.checked_add(len).ok_or(Error::AvccTruncated)?;
304 if buf.len() < after_nal {
305 return Err(Error::AvccTruncated);
306 }
307 out.push(Bytes::copy_from_slice(&buf[after_len..after_nal]));
308 *pos = after_nal;
309 }
310 Ok(out)
311}
312
313pub(crate) fn avcc_params(avcc: &[u8]) -> anyhow::Result<(usize, Vec<Bytes>)> {
317 anyhow::ensure!(avcc.len() >= 6, "AVCDecoderConfigurationRecord too short");
318 let length_size = (avcc[4] & 0x03) as usize + 1;
319
320 let mut params = Vec::new();
321 let num_sps = avcc[5] & 0x1f;
322 let mut pos = read_param_set_array(avcc, 6, num_sps as usize, &mut params)?;
323
324 anyhow::ensure!(avcc.len() > pos, "avcC missing PPS count");
325 let num_pps = avcc[pos];
326 pos += 1;
327 read_param_set_array(avcc, pos, num_pps as usize, &mut params)?;
328
329 Ok((length_size, params))
330}
331
332fn read_param_set_array(buf: &[u8], mut pos: usize, count: usize, params: &mut Vec<Bytes>) -> anyhow::Result<usize> {
335 for _ in 0..count {
336 anyhow::ensure!(buf.len() >= pos + 2, "truncated parameter-set length");
337 let len = u16::from_be_bytes([buf[pos], buf[pos + 1]]) as usize;
338 pos += 2;
339 anyhow::ensure!(buf.len() >= pos + len, "parameter-set NAL exceeds buffer");
340 params.push(Bytes::copy_from_slice(&buf[pos..pos + len]));
341 pos += len;
342 }
343 Ok(pos)
344}
345
346pub struct Avc1 {
358 avcc: Option<Bytes>,
359 sps: Vec<Bytes>,
361 pps: Vec<Bytes>,
363}
364
365impl Default for Avc1 {
366 fn default() -> Self {
367 Self::new()
368 }
369}
370
371impl Avc1 {
372 pub fn new() -> Self {
374 Self {
375 avcc: None,
376 sps: Vec::new(),
377 pps: Vec::new(),
378 }
379 }
380
381 pub fn avcc(&self) -> Option<&Bytes> {
383 self.avcc.as_ref()
384 }
385
386 pub fn transform(&mut self, payload: Bytes) -> Result<Option<Bytes>> {
394 let mut buf = payload.clone();
398 let mut nal_iter = crate::codec::annexb::NalIterator::new(&mut buf);
399
400 let mut out = BytesMut::with_capacity(payload.remaining());
401 let mut frame_sps: Vec<Bytes> = Vec::new();
402 let mut frame_pps: Vec<Bytes> = Vec::new();
403 let mut emitted_any_slice = false;
404
405 loop {
406 let nal = match nal_iter.next() {
407 Some(Ok(n)) => n,
408 Some(Err(e)) => return Err(e.into()),
409 None => break,
410 };
411 if process_nal(&nal, &mut out, &mut frame_sps, &mut frame_pps)? {
412 emitted_any_slice = true;
413 }
414 }
415
416 if let Some(nal) = nal_iter.flush()? {
417 if process_nal(&nal, &mut out, &mut frame_sps, &mut frame_pps)? {
418 emitted_any_slice = true;
419 }
420 }
421
422 let mut changed = false;
427 if !frame_sps.is_empty() && frame_sps != self.sps {
428 self.sps = frame_sps;
429 changed = true;
430 }
431 if !frame_pps.is_empty() && frame_pps != self.pps {
432 self.pps = frame_pps;
433 changed = true;
434 }
435 if changed {
436 self.rebuild_avcc()?;
437 }
438
439 if !emitted_any_slice {
440 return Ok(None);
441 }
442
443 Ok(Some(out.freeze()))
444 }
445
446 fn rebuild_avcc(&mut self) -> Result<()> {
447 if self.sps.is_empty() || self.pps.is_empty() {
448 return Ok(());
449 }
450 self.avcc = Some(build_avcc(&self.sps, &self.pps)?);
451 Ok(())
452 }
453}
454
455fn process_nal(
459 nal: &Bytes,
460 out: &mut BytesMut,
461 frame_sps: &mut Vec<Bytes>,
462 frame_pps: &mut Vec<Bytes>,
463) -> Result<bool> {
464 if nal.is_empty() {
465 return Ok(false);
466 }
467 match nal[0] & 0x1f {
468 NAL_TYPE_SPS => {
469 crate::codec::annexb::push_distinct(frame_sps, nal);
470 Ok(false)
471 }
472 NAL_TYPE_PPS => {
473 crate::codec::annexb::push_distinct(frame_pps, nal);
474 Ok(false)
475 }
476 _ => {
477 let len = u32::try_from(nal.len()).map_err(|_| Error::NalTooLarge)?;
478 out.extend_from_slice(&len.to_be_bytes());
479 out.extend_from_slice(nal);
480 Ok(true)
481 }
482 }
483}
484
485#[cfg(test)]
486mod tests {
487 use super::*;
488
489 const SC4: &[u8] = &[0, 0, 0, 1];
490
491 fn annexb_frame(nals: &[&[u8]]) -> Bytes {
492 let mut buf = BytesMut::new();
493 for nal in nals {
494 buf.extend_from_slice(SC4);
495 buf.extend_from_slice(nal);
496 }
497 buf.freeze()
498 }
499
500 #[test]
503 fn avc1_frame_keyframe() {
504 let idr: &[u8] = &[0x65, 0x88, 0x84, 0x21];
505 let mut au = BytesMut::new();
506 au.extend_from_slice(&(idr.len() as u32).to_be_bytes());
507 au.extend_from_slice(idr);
508
509 let frame = avc1_frame(&au, 4, crate::container::Timestamp::from_micros(0).unwrap()).unwrap();
510 assert!(frame.keyframe);
511 assert_eq!(frame.payload[4..], *idr);
512 }
513
514 #[test]
516 fn avc1_frame_delta() {
517 let pslice: &[u8] = &[0x61, 0xe0, 0x12, 0x34];
518 let mut au = BytesMut::new();
519 au.extend_from_slice(&(pslice.len() as u32).to_be_bytes());
520 au.extend_from_slice(pslice);
521
522 let frame = avc1_frame(&au, 4, crate::container::Timestamp::from_micros(0).unwrap()).unwrap();
523 assert!(!frame.keyframe);
524 }
525
526 #[test]
527 fn avc3_strips_sps_pps_and_builds_avcc() {
528 let sps = &[0x67, 0x42, 0xc0, 0x1f, 0xde][..];
529 let pps = &[0x68, 0xce, 0x3c, 0x80][..];
530 let idr = &[0x65, 0x88, 0x84, 0x21][..];
531
532 let mut tx = Avc1::new();
533 assert!(tx.avcc().is_none());
534
535 let frame = annexb_frame(&[sps, pps, idr]);
536 let out = tx.transform(frame).expect("transform").expect("expected output");
537
538 let avcc = tx.avcc().expect("avcC available").clone();
539 assert_eq!(avcc[0], 1);
540 assert_eq!(avcc[1], sps[1]);
541 assert_eq!(avcc[3], sps[3]);
542
543 let mut expected = BytesMut::new();
544 expected.extend_from_slice(&(idr.len() as u32).to_be_bytes());
545 expected.extend_from_slice(idr);
546 assert_eq!(out.as_ref(), expected.as_ref());
547 }
548
549 #[test]
550 fn avcc_params_roundtrips_build_avcc() {
551 let sps = Bytes::from_static(&[0x67, 0x42, 0xc0, 0x1f, 0xde]);
552 let pps = Bytes::from_static(&[0x68, 0xce, 0x3c, 0x80]);
553
554 let avcc = build_avcc(std::slice::from_ref(&sps), std::slice::from_ref(&pps)).unwrap();
555 let (length_size, params) = avcc_params(&avcc).unwrap();
556
557 assert_eq!(length_size, 4);
558 assert_eq!(params.len(), 2);
559 assert_eq!(params[0], sps);
560 assert_eq!(params[1], pps);
561 }
562
563 #[test]
564 fn build_avcc_carries_multiple_pps() {
565 let sps = Bytes::from_static(&[0x67, 0x42, 0xc0, 0x1f, 0xde]);
568 let pps0 = Bytes::from_static(&[0x68, 0xce, 0x3c, 0x80]);
569 let pps1 = Bytes::from_static(&[0x68, 0xce, 0x3c, 0x81]);
570
571 let avcc = build_avcc(std::slice::from_ref(&sps), &[pps0.clone(), pps1.clone()]).unwrap();
572 assert_eq!(avcc[5] & 0x1f, 1);
574
575 let (_, params) = avcc_params(&avcc).unwrap();
576 assert_eq!(params, vec![sps, pps0, pps1]);
577 }
578
579 #[test]
580 fn avc3_keyframe_with_two_pps_keeps_both() {
581 let sps = &[0x67, 0x42, 0xc0, 0x1f, 0xde][..];
583 let pps0 = &[0x68, 0xce, 0x3c, 0x80][..];
584 let pps1 = &[0x68, 0xce, 0x3c, 0x81][..];
585 let idr = &[0x65, 0x88][..];
586
587 let mut tx = Avc1::new();
588 tx.transform(annexb_frame(&[sps, pps0, pps1, idr])).unwrap();
589
590 let avcc = tx.avcc().expect("avcC available");
591 let (_, params) = avcc_params(avcc).unwrap();
592 assert_eq!(
593 params.iter().map(|p| p.as_ref()).collect::<Vec<_>>(),
594 vec![sps, pps0, pps1]
595 );
596 }
597
598 #[test]
599 fn avc3_reinit_drops_superseded_pps() {
600 let sps = &[0x67, 0x42, 0xc0, 0x1f, 0xde][..];
603 let pps0 = &[0x68, 0xce, 0x3c, 0x80][..];
604 let pps1 = &[0x68, 0xce, 0x3c, 0x81][..];
605 let idr = &[0x65, 0x88][..];
606
607 let mut tx = Avc1::new();
608 tx.transform(annexb_frame(&[sps, pps0, idr])).unwrap();
609 tx.transform(annexb_frame(&[sps, pps1, idr])).unwrap();
610
611 let avcc = tx.avcc().expect("avcC available");
612 let (_, params) = avcc_params(avcc).unwrap();
613 assert_eq!(
614 params.iter().map(|p| p.as_ref()).collect::<Vec<_>>(),
615 vec![sps, pps1],
616 "reinit must drop the superseded PPS"
617 );
618 }
619
620 #[test]
621 fn avc3_parameter_only_frame_returns_none() {
622 let sps = &[0x67, 0x42, 0xc0, 0x1f, 0xde][..];
623 let pps = &[0x68, 0xce, 0x3c, 0x80][..];
624
625 let mut tx = Avc1::new();
626 let frame = annexb_frame(&[sps, pps]);
627 assert!(tx.transform(frame).unwrap().is_none());
628 assert!(tx.avcc().is_some());
629 }
630
631 #[test]
632 fn avc3_subsequent_frame_uses_cached_avcc() {
633 let sps = &[0x67, 0x42, 0xc0, 0x1f, 0xde][..];
634 let pps = &[0x68, 0xce, 0x3c, 0x80][..];
635 let idr = &[0x65, 0x88][..];
636 let p = &[0x61, 0xe0, 0x12][..];
637
638 let mut tx = Avc1::new();
639 tx.transform(annexb_frame(&[sps, pps, idr])).unwrap();
640 let avcc_v1 = tx.avcc().unwrap().clone();
641
642 let out = tx.transform(annexb_frame(&[p])).unwrap().unwrap();
643 assert_eq!(tx.avcc().unwrap(), &avcc_v1);
644 let mut expected = BytesMut::new();
645 expected.extend_from_slice(&(p.len() as u32).to_be_bytes());
646 expected.extend_from_slice(p);
647 assert_eq!(out.as_ref(), expected.as_ref());
648 }
649
650 #[test]
651 fn avc3_export_e2e_payload_shape() {
652 let sps = &[0x67u8, 0x42, 0xc0, 0x1f, 0xde, 0xad, 0xbe, 0xef][..];
655 let pps = &[0x68u8, 0xce, 0x3c, 0x80][..];
656 let idr = &[0x65u8, 0x88, 0x84, 0x21, 0x00, 0x11, 0x22, 0x33][..];
657 let pslice = &[0x61u8, 0xe0, 0x12, 0x34][..];
658
659 let mut tx = Avc1::new();
660 let key = annexb_frame(&[sps, pps, idr]);
661 let key_out = tx.transform(key).expect("transform key").expect("output");
662 assert!(tx.avcc().is_some());
663
664 assert_eq!(key_out.len(), 4 + idr.len());
665 assert_eq!(&key_out[4..], idr);
666
667 let p = annexb_frame(&[pslice]);
668 let p_out = tx.transform(p).expect("transform p").expect("output");
669 assert_eq!(p_out.len(), 4 + pslice.len());
670 assert_eq!(&p_out[4..], pslice);
671 }
672}