mod export;
mod import;
mod split;
pub use export::*;
pub use import::*;
pub use split::*;
use bytes::{Buf, BufMut, Bytes, BytesMut};
const NAL_TYPE_SPS: u8 = 7;
const NAL_TYPE_PPS: u8 = 8;
pub(crate) fn avc1_frame(
data: &[u8],
length_size: usize,
pts: crate::container::Timestamp,
) -> crate::Result<crate::container::Frame> {
Ok(crate::container::Frame {
timestamp: pts,
payload: data.to_vec().into(),
keyframe: avc1_is_keyframe(data, length_size),
duration: None,
})
}
fn avc1_is_keyframe(data: &[u8], length_size: usize) -> bool {
let mut offset = 0;
while offset + length_size <= data.len() {
let nal_len = match length_size {
1 => data[offset] as usize,
2 => u16::from_be_bytes([data[offset], data[offset + 1]]) as usize,
3 => u32::from_be_bytes([0, data[offset], data[offset + 1], data[offset + 2]]) as usize,
4 => u32::from_be_bytes([data[offset], data[offset + 1], data[offset + 2], data[offset + 3]]) as usize,
_ => return false,
};
offset += length_size;
if offset + nal_len > data.len() {
break;
}
if nal_len > 0 && data[offset] & 0x1f == 5 {
return true; }
offset += nal_len;
}
false
}
#[derive(Debug, Clone, thiserror::Error)]
#[non_exhaustive]
pub enum Error {
#[error("SPS NAL too short")]
SpsTooShort,
#[error("failed to parse SPS")]
SpsParse,
#[error("AVCDecoderConfigurationRecord too short")]
AvccTooShort,
#[error("AVCDecoderConfigurationRecord truncated")]
AvccTruncated,
#[error("avc1 description for rendition {name:?} is missing SPS or PPS (sps={sps}, pps={pps})")]
MissingParamSets { name: String, sps: usize, pps: usize },
#[error("SPS too large for avcC length field ({0} > {max})", max = u16::MAX)]
SpsTooLarge(usize),
#[error("PPS too large for avcC length field ({0} > {max})", max = u16::MAX)]
PpsTooLarge(usize),
#[error("avcC requires at least one SPS")]
MissingSps,
#[error("too many SPS for avcC ({0} > 31)")]
TooManySps(usize),
#[error("too many PPS for avcC ({0} > 255)")]
TooManyPps(usize),
#[error("NAL too large for 4-byte length prefix")]
NalTooLarge,
#[error("NAL unit is too short")]
NalTooShort,
#[error("forbidden zero bit is not zero")]
ForbiddenZeroBit,
#[error("not initialized")]
NotInitialized,
#[error("avc3 track not created")]
Avc3TrackNotCreated,
#[error("missing timestamp")]
MissingTimestamp,
#[error("annexb: {0}")]
Annexb(#[from] crate::codec::annexb::Error),
}
pub type Result<T> = std::result::Result<T, Error>;
#[derive(Debug, Clone)]
pub struct Sps {
pub profile: u8,
pub constraints: u8,
pub level: u8,
pub coded_width: u32,
pub coded_height: u32,
}
impl Sps {
pub fn parse(nal: &[u8]) -> Result<Self> {
if nal.len() < 4 {
return Err(Error::SpsTooShort);
}
let rbsp = h264_parser::nal::ebsp_to_rbsp(&nal[1..]);
let sps = h264_parser::Sps::parse(&rbsp).map_err(|_| Error::SpsParse)?;
Ok(Self {
profile: sps.profile_idc,
constraints: pack_constraint_flags(&sps),
level: sps.level_idc,
coded_width: sps.width,
coded_height: sps.height,
})
}
}
#[derive(Debug, Clone)]
#[non_exhaustive]
pub struct Avcc {
pub profile: u8,
pub constraints: u8,
pub level: u8,
pub length_size: usize,
pub sps: Vec<Bytes>,
pub pps: Vec<Bytes>,
pub coded_width: Option<u32>,
pub coded_height: Option<u32>,
}
impl Avcc {
pub fn parse(avcc: &[u8]) -> Result<Self> {
if avcc.len() < 7 {
return Err(Error::AvccTooShort);
}
let profile = avcc[1];
let constraints = avcc[2];
let level = avcc[3];
let length_size = (avcc[4] & 0x03) as usize + 1;
let num_sps = (avcc[5] & 0x1f) as usize;
let mut pos = 6;
let sps = read_param_sets(avcc, &mut pos, num_sps)?;
if avcc.len() <= pos {
return Err(Error::AvccTruncated);
}
let num_pps = avcc[pos] as usize;
pos += 1;
let pps = read_param_sets(avcc, &mut pos, num_pps)?;
let (mut coded_width, mut coded_height) = (None, None);
if let Some(first) = sps.first()
&& first.len() > 1
&& let Ok(parsed) = Sps::parse(first)
{
coded_width = Some(parsed.coded_width);
coded_height = Some(parsed.coded_height);
}
Ok(Self {
profile,
constraints,
level,
length_size,
sps,
pps,
coded_width,
coded_height,
})
}
}
fn pack_constraint_flags(sps: &h264_parser::Sps) -> u8 {
((sps.constraint_set0_flag as u8) << 7)
| ((sps.constraint_set1_flag as u8) << 6)
| ((sps.constraint_set2_flag as u8) << 5)
| ((sps.constraint_set3_flag as u8) << 4)
| ((sps.constraint_set4_flag as u8) << 3)
| ((sps.constraint_set5_flag as u8) << 2)
}
pub(crate) fn build_avcc(sps_nals: &[Bytes], pps_nals: &[Bytes]) -> Result<Bytes> {
let first_sps = sps_nals.first().ok_or(Error::MissingSps)?;
if first_sps.len() < 4 {
return Err(Error::SpsTooShort);
}
if sps_nals.len() > 0x1f {
return Err(Error::TooManySps(sps_nals.len()));
}
if pps_nals.len() > u8::MAX as usize {
return Err(Error::TooManyPps(pps_nals.len()));
}
for sps in sps_nals {
if sps.len() > u16::MAX as usize {
return Err(Error::SpsTooLarge(sps.len()));
}
}
for pps in pps_nals {
if pps.len() > u16::MAX as usize {
return Err(Error::PpsTooLarge(pps.len()));
}
}
let profile_idc = first_sps[1];
let constraints = first_sps[2];
let level_idc = first_sps[3];
let payload: usize = sps_nals.iter().chain(pps_nals).map(|n| 2 + n.len()).sum();
let mut out = BytesMut::with_capacity(7 + payload);
out.put_u8(1); out.put_u8(profile_idc);
out.put_u8(constraints);
out.put_u8(level_idc);
out.put_u8(0xff); out.put_u8(0xe0 | sps_nals.len() as u8); for sps in sps_nals {
out.put_u16(sps.len() as u16);
out.put_slice(sps);
}
out.put_u8(pps_nals.len() as u8); for pps in pps_nals {
out.put_u16(pps.len() as u16);
out.put_slice(pps);
}
Ok(out.freeze())
}
fn read_param_sets(buf: &[u8], pos: &mut usize, count: usize) -> Result<Vec<Bytes>> {
let mut out = Vec::with_capacity(count);
for _ in 0..count {
let after_len = pos.checked_add(2).ok_or(Error::AvccTruncated)?;
if buf.len() < after_len {
return Err(Error::AvccTruncated);
}
let len = u16::from_be_bytes([buf[*pos], buf[*pos + 1]]) as usize;
let after_nal = after_len.checked_add(len).ok_or(Error::AvccTruncated)?;
if buf.len() < after_nal {
return Err(Error::AvccTruncated);
}
out.push(Bytes::copy_from_slice(&buf[after_len..after_nal]));
*pos = after_nal;
}
Ok(out)
}
pub(crate) fn avcc_params(avcc: &[u8]) -> anyhow::Result<(usize, Vec<Bytes>)> {
anyhow::ensure!(avcc.len() >= 6, "AVCDecoderConfigurationRecord too short");
let length_size = (avcc[4] & 0x03) as usize + 1;
let mut params = Vec::new();
let num_sps = avcc[5] & 0x1f;
let mut pos = read_param_set_array(avcc, 6, num_sps as usize, &mut params)?;
anyhow::ensure!(avcc.len() > pos, "avcC missing PPS count");
let num_pps = avcc[pos];
pos += 1;
read_param_set_array(avcc, pos, num_pps as usize, &mut params)?;
Ok((length_size, params))
}
fn read_param_set_array(buf: &[u8], mut pos: usize, count: usize, params: &mut Vec<Bytes>) -> anyhow::Result<usize> {
for _ in 0..count {
anyhow::ensure!(buf.len() >= pos + 2, "truncated parameter-set length");
let len = u16::from_be_bytes([buf[pos], buf[pos + 1]]) as usize;
pos += 2;
anyhow::ensure!(buf.len() >= pos + len, "parameter-set NAL exceeds buffer");
params.push(Bytes::copy_from_slice(&buf[pos..pos + len]));
pos += len;
}
Ok(pos)
}
pub struct Avc1 {
avcc: Option<Bytes>,
sps: Vec<Bytes>,
pps: Vec<Bytes>,
}
impl Default for Avc1 {
fn default() -> Self {
Self::new()
}
}
impl Avc1 {
pub fn new() -> Self {
Self {
avcc: None,
sps: Vec::new(),
pps: Vec::new(),
}
}
pub fn avcc(&self) -> Option<&Bytes> {
self.avcc.as_ref()
}
pub fn transform(&mut self, payload: Bytes) -> Result<Option<Bytes>> {
let mut buf = payload.clone();
let mut nal_iter = crate::codec::annexb::NalIterator::new(&mut buf);
let mut out = BytesMut::with_capacity(payload.remaining());
let mut frame_sps: Vec<Bytes> = Vec::new();
let mut frame_pps: Vec<Bytes> = Vec::new();
let mut emitted_any_slice = false;
loop {
let nal = match nal_iter.next() {
Some(Ok(n)) => n,
Some(Err(e)) => return Err(e.into()),
None => break,
};
if process_nal(&nal, &mut out, &mut frame_sps, &mut frame_pps)? {
emitted_any_slice = true;
}
}
if let Some(nal) = nal_iter.flush()? {
if process_nal(&nal, &mut out, &mut frame_sps, &mut frame_pps)? {
emitted_any_slice = true;
}
}
let mut changed = false;
if !frame_sps.is_empty() && frame_sps != self.sps {
self.sps = frame_sps;
changed = true;
}
if !frame_pps.is_empty() && frame_pps != self.pps {
self.pps = frame_pps;
changed = true;
}
if changed {
self.rebuild_avcc()?;
}
if !emitted_any_slice {
return Ok(None);
}
Ok(Some(out.freeze()))
}
fn rebuild_avcc(&mut self) -> Result<()> {
if self.sps.is_empty() || self.pps.is_empty() {
return Ok(());
}
self.avcc = Some(build_avcc(&self.sps, &self.pps)?);
Ok(())
}
}
fn process_nal(
nal: &Bytes,
out: &mut BytesMut,
frame_sps: &mut Vec<Bytes>,
frame_pps: &mut Vec<Bytes>,
) -> Result<bool> {
if nal.is_empty() {
return Ok(false);
}
match nal[0] & 0x1f {
NAL_TYPE_SPS => {
crate::codec::annexb::push_distinct(frame_sps, nal);
Ok(false)
}
NAL_TYPE_PPS => {
crate::codec::annexb::push_distinct(frame_pps, nal);
Ok(false)
}
_ => {
let len = u32::try_from(nal.len()).map_err(|_| Error::NalTooLarge)?;
out.extend_from_slice(&len.to_be_bytes());
out.extend_from_slice(nal);
Ok(true)
}
}
}
#[cfg(test)]
mod tests {
use super::*;
const SC4: &[u8] = &[0, 0, 0, 1];
fn annexb_frame(nals: &[&[u8]]) -> Bytes {
let mut buf = BytesMut::new();
for nal in nals {
buf.extend_from_slice(SC4);
buf.extend_from_slice(nal);
}
buf.freeze()
}
#[test]
fn avc1_frame_keyframe() {
let idr: &[u8] = &[0x65, 0x88, 0x84, 0x21];
let mut au = BytesMut::new();
au.extend_from_slice(&(idr.len() as u32).to_be_bytes());
au.extend_from_slice(idr);
let frame = avc1_frame(&au, 4, crate::container::Timestamp::from_micros(0).unwrap()).unwrap();
assert!(frame.keyframe);
assert_eq!(frame.payload[4..], *idr);
}
#[test]
fn avc1_frame_delta() {
let pslice: &[u8] = &[0x61, 0xe0, 0x12, 0x34];
let mut au = BytesMut::new();
au.extend_from_slice(&(pslice.len() as u32).to_be_bytes());
au.extend_from_slice(pslice);
let frame = avc1_frame(&au, 4, crate::container::Timestamp::from_micros(0).unwrap()).unwrap();
assert!(!frame.keyframe);
}
#[test]
fn avc3_strips_sps_pps_and_builds_avcc() {
let sps = &[0x67, 0x42, 0xc0, 0x1f, 0xde][..];
let pps = &[0x68, 0xce, 0x3c, 0x80][..];
let idr = &[0x65, 0x88, 0x84, 0x21][..];
let mut tx = Avc1::new();
assert!(tx.avcc().is_none());
let frame = annexb_frame(&[sps, pps, idr]);
let out = tx.transform(frame).expect("transform").expect("expected output");
let avcc = tx.avcc().expect("avcC available").clone();
assert_eq!(avcc[0], 1);
assert_eq!(avcc[1], sps[1]);
assert_eq!(avcc[3], sps[3]);
let mut expected = BytesMut::new();
expected.extend_from_slice(&(idr.len() as u32).to_be_bytes());
expected.extend_from_slice(idr);
assert_eq!(out.as_ref(), expected.as_ref());
}
#[test]
fn avcc_params_roundtrips_build_avcc() {
let sps = Bytes::from_static(&[0x67, 0x42, 0xc0, 0x1f, 0xde]);
let pps = Bytes::from_static(&[0x68, 0xce, 0x3c, 0x80]);
let avcc = build_avcc(std::slice::from_ref(&sps), std::slice::from_ref(&pps)).unwrap();
let (length_size, params) = avcc_params(&avcc).unwrap();
assert_eq!(length_size, 4);
assert_eq!(params.len(), 2);
assert_eq!(params[0], sps);
assert_eq!(params[1], pps);
}
#[test]
fn build_avcc_carries_multiple_pps() {
let sps = Bytes::from_static(&[0x67, 0x42, 0xc0, 0x1f, 0xde]);
let pps0 = Bytes::from_static(&[0x68, 0xce, 0x3c, 0x80]);
let pps1 = Bytes::from_static(&[0x68, 0xce, 0x3c, 0x81]);
let avcc = build_avcc(std::slice::from_ref(&sps), &[pps0.clone(), pps1.clone()]).unwrap();
assert_eq!(avcc[5] & 0x1f, 1);
let (_, params) = avcc_params(&avcc).unwrap();
assert_eq!(params, vec![sps, pps0, pps1]);
}
#[test]
fn avc3_keyframe_with_two_pps_keeps_both() {
let sps = &[0x67, 0x42, 0xc0, 0x1f, 0xde][..];
let pps0 = &[0x68, 0xce, 0x3c, 0x80][..];
let pps1 = &[0x68, 0xce, 0x3c, 0x81][..];
let idr = &[0x65, 0x88][..];
let mut tx = Avc1::new();
tx.transform(annexb_frame(&[sps, pps0, pps1, idr])).unwrap();
let avcc = tx.avcc().expect("avcC available");
let (_, params) = avcc_params(avcc).unwrap();
assert_eq!(
params.iter().map(|p| p.as_ref()).collect::<Vec<_>>(),
vec![sps, pps0, pps1]
);
}
#[test]
fn avc3_reinit_drops_superseded_pps() {
let sps = &[0x67, 0x42, 0xc0, 0x1f, 0xde][..];
let pps0 = &[0x68, 0xce, 0x3c, 0x80][..];
let pps1 = &[0x68, 0xce, 0x3c, 0x81][..];
let idr = &[0x65, 0x88][..];
let mut tx = Avc1::new();
tx.transform(annexb_frame(&[sps, pps0, idr])).unwrap();
tx.transform(annexb_frame(&[sps, pps1, idr])).unwrap();
let avcc = tx.avcc().expect("avcC available");
let (_, params) = avcc_params(avcc).unwrap();
assert_eq!(
params.iter().map(|p| p.as_ref()).collect::<Vec<_>>(),
vec![sps, pps1],
"reinit must drop the superseded PPS"
);
}
#[test]
fn avc3_parameter_only_frame_returns_none() {
let sps = &[0x67, 0x42, 0xc0, 0x1f, 0xde][..];
let pps = &[0x68, 0xce, 0x3c, 0x80][..];
let mut tx = Avc1::new();
let frame = annexb_frame(&[sps, pps]);
assert!(tx.transform(frame).unwrap().is_none());
assert!(tx.avcc().is_some());
}
#[test]
fn avc3_subsequent_frame_uses_cached_avcc() {
let sps = &[0x67, 0x42, 0xc0, 0x1f, 0xde][..];
let pps = &[0x68, 0xce, 0x3c, 0x80][..];
let idr = &[0x65, 0x88][..];
let p = &[0x61, 0xe0, 0x12][..];
let mut tx = Avc1::new();
tx.transform(annexb_frame(&[sps, pps, idr])).unwrap();
let avcc_v1 = tx.avcc().unwrap().clone();
let out = tx.transform(annexb_frame(&[p])).unwrap().unwrap();
assert_eq!(tx.avcc().unwrap(), &avcc_v1);
let mut expected = BytesMut::new();
expected.extend_from_slice(&(p.len() as u32).to_be_bytes());
expected.extend_from_slice(p);
assert_eq!(out.as_ref(), expected.as_ref());
}
#[test]
fn avc3_export_e2e_payload_shape() {
let sps = &[0x67u8, 0x42, 0xc0, 0x1f, 0xde, 0xad, 0xbe, 0xef][..];
let pps = &[0x68u8, 0xce, 0x3c, 0x80][..];
let idr = &[0x65u8, 0x88, 0x84, 0x21, 0x00, 0x11, 0x22, 0x33][..];
let pslice = &[0x61u8, 0xe0, 0x12, 0x34][..];
let mut tx = Avc1::new();
let key = annexb_frame(&[sps, pps, idr]);
let key_out = tx.transform(key).expect("transform key").expect("output");
assert!(tx.avcc().is_some());
assert_eq!(key_out.len(), 4 + idr.len());
assert_eq!(&key_out[4..], idr);
let p = annexb_frame(&[pslice]);
let p_out = tx.transform(p).expect("transform p").expect("output");
assert_eq!(p_out.len(), 4 + pslice.len());
assert_eq!(&p_out[4..], pslice);
}
}