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