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