Skip to main content

iso_bmff/bitstream/
avc.rs

1//! H.264 Annex-B ↔ AVCC.
2
3#![forbid(unsafe_code)]
4
5use bytes::Bytes;
6use memchr::memmem;
7
8/// AVCC conversion result.
9#[derive(Debug, Clone, PartialEq, Eq)]
10pub struct AvccOut {
11    /// Length-prefixed access unit.
12    pub payload: Bytes,
13    /// Fresh `avcC` when SPS+PPS were present.
14    pub avcc: Option<Bytes>,
15}
16
17/// Convert Annex-B to 4-byte length-prefixed AVCC, or pass through.
18#[must_use]
19pub fn to_avcc(data: &[u8]) -> AvccOut {
20    if !is_annex_b(data) {
21        return AvccOut {
22            payload: Bytes::copy_from_slice(data),
23            avcc: None,
24        };
25    }
26
27    let mut out = Vec::with_capacity(data.len());
28    let mut sps: Option<&[u8]> = None;
29    let mut pps: Option<&[u8]> = None;
30
31    for nal in NalIter::new(data) {
32        if nal.is_empty() {
33            continue;
34        }
35        match nal[0] & 0x1f {
36            7 => sps = Some(nal),
37            8 => pps = Some(nal),
38            _ => {}
39        }
40        let len = u32::try_from(nal.len()).unwrap_or(u32::MAX);
41        out.extend_from_slice(&len.to_be_bytes());
42        out.extend_from_slice(nal);
43    }
44
45    let avcc = match (sps, pps) {
46        (Some(s), Some(p)) => Some(build_avcc(s, p)),
47        _ => None,
48    };
49
50    AvccOut {
51        payload: Bytes::from(out),
52        avcc,
53    }
54}
55
56fn is_annex_b(data: &[u8]) -> bool {
57    matches!(find_start_code(data), Some((0, _)))
58}
59
60/// Parsed `lengthSizeMinusOne` + SPS/PPS from an `AVCDecoderConfigurationRecord`
61/// (the raw `avcC` box payload, ISO/IEC 14496-15).
62#[derive(Debug, Clone, PartialEq, Eq)]
63pub struct AvcDecoderConfig {
64    /// NAL length prefix size in bytes (1, 2, or 4) used by AVCC-framed samples.
65    pub nal_length_size: u8,
66    /// One or more SPS NAL units (without start code or length prefix).
67    pub sps: Vec<Bytes>,
68    /// One or more PPS NAL units (without start code or length prefix).
69    pub pps: Vec<Bytes>,
70}
71
72/// Parse an `AVCDecoderConfigurationRecord` (`avcC` box payload). Returns `None` on
73/// malformed/truncated input rather than panicking — this reads demuxer-sourced,
74/// otherwise-untrusted bytes.
75#[must_use]
76pub fn parse_avc_decoder_config(record: &[u8]) -> Option<AvcDecoderConfig> {
77    if record.len() < 7 || record[0] != 1 {
78        return None;
79    }
80    let nal_length_size = (record[4] & 0x03) + 1;
81    let num_sps = record[5] & 0x1f;
82    let mut pos = 6usize;
83    let mut sps = Vec::new();
84    for _ in 0..num_sps {
85        let (nal, next) = read_length_prefixed(record, pos)?;
86        sps.push(nal);
87        pos = next;
88    }
89    let pps_count = *record.get(pos)?;
90    pos += 1;
91    let mut pps = Vec::new();
92    for _ in 0..pps_count {
93        let (nal, next) = read_length_prefixed(record, pos)?;
94        pps.push(nal);
95        pos = next;
96    }
97    Some(AvcDecoderConfig {
98        nal_length_size,
99        sps,
100        pps,
101    })
102}
103
104fn read_length_prefixed(data: &[u8], pos: usize) -> Option<(Bytes, usize)> {
105    let len_bytes = data.get(pos..pos + 2)?;
106    let len = usize::from(u16::from_be_bytes([len_bytes[0], len_bytes[1]]));
107    let start = pos + 2;
108    let nal = data.get(start..start + len)?;
109    Some((Bytes::copy_from_slice(nal), start + len))
110}
111
112/// Concatenated Annex-B (4-byte start code) SPS + PPS from a parsed decoder config —
113/// the sequence-header form Windows Media Foundation's `MF_MT_MPEG_SEQUENCE_HEADER`
114/// attribute expects.
115#[must_use]
116pub fn annex_b_sequence_header(config: &AvcDecoderConfig) -> Bytes {
117    let mut out = Vec::new();
118    for nal in config.sps.iter().chain(config.pps.iter()) {
119        out.extend_from_slice(&[0, 0, 0, 1]);
120        out.extend_from_slice(nal);
121    }
122    Bytes::from(out)
123}
124
125/// Convert one AVCC length-prefixed access unit to Annex-B (4-byte start codes).
126///
127/// Stops at the first malformed/truncated NAL length rather than panicking; whatever
128/// converted so far is returned (matches `to_avcc`'s best-effort, non-`Result` style).
129#[must_use]
130pub fn avcc_payload_to_annex_b(data: &[u8], nal_length_size: u8) -> Bytes {
131    let nls = usize::from(nal_length_size).clamp(1, 4);
132    let mut out = Vec::with_capacity(data.len() + 16);
133    let mut pos = 0usize;
134    while pos + nls <= data.len() {
135        let Some(len) = read_nal_length(&data[pos..pos + nls]) else {
136            break;
137        };
138        pos += nls;
139        let Some(nal) = data.get(pos..pos + len) else {
140            break;
141        };
142        out.extend_from_slice(&[0, 0, 0, 1]);
143        out.extend_from_slice(nal);
144        pos += len;
145    }
146    Bytes::from(out)
147}
148
149fn read_nal_length(len_bytes: &[u8]) -> Option<usize> {
150    match len_bytes.len() {
151        1 => Some(usize::from(len_bytes[0])),
152        2 => Some(usize::from(u16::from_be_bytes([
153            len_bytes[0],
154            len_bytes[1],
155        ]))),
156        4 => Some(
157            usize::try_from(u32::from_be_bytes([
158                len_bytes[0],
159                len_bytes[1],
160                len_bytes[2],
161                len_bytes[3],
162            ]))
163            .unwrap_or(usize::MAX),
164        ),
165        _ => None,
166    }
167}
168
169fn build_avcc(sps: &[u8], pps: &[u8]) -> Bytes {
170    let mut v = Vec::with_capacity(11 + sps.len() + pps.len());
171    v.push(1);
172    if sps.len() >= 4 {
173        v.extend_from_slice(&sps[1..4]);
174    } else {
175        v.extend_from_slice(&[0x42, 0x00, 0x1e]);
176    }
177    v.push(0xff);
178    v.push(0xe1);
179    v.extend_from_slice(&(u16::try_from(sps.len()).unwrap_or(u16::MAX)).to_be_bytes());
180    v.extend_from_slice(sps);
181    v.push(1);
182    v.extend_from_slice(&(u16::try_from(pps.len()).unwrap_or(u16::MAX)).to_be_bytes());
183    v.extend_from_slice(pps);
184    Bytes::from(v)
185}
186
187/// Next Annex-B start code: `(offset, code_len)` with `code_len` in `{3, 4}`.
188/// Prefers the 4-byte code when both would match at the same NAL boundary.
189fn find_start_code(hay: &[u8]) -> Option<(usize, usize)> {
190    let i4 = memmem::find(hay, &[0, 0, 0, 1]);
191    let i3 = memmem::find(hay, &[0, 0, 1]);
192    match (i4, i3) {
193        (Some(a), Some(b)) if a <= b => Some((a, 4)),
194        (Some(a), None) => Some((a, 4)),
195        (_, Some(b)) => Some((b, 3)),
196        (None, None) => None,
197    }
198}
199
200/// Zero-alloc Annex-B NAL iterator (yields slices into the input).
201struct NalIter<'a> {
202    data: &'a [u8],
203    pos: usize,
204}
205
206impl<'a> NalIter<'a> {
207    const fn new(data: &'a [u8]) -> Self {
208        Self { data, pos: 0 }
209    }
210}
211
212impl<'a> Iterator for NalIter<'a> {
213    type Item = &'a [u8];
214
215    fn next(&mut self) -> Option<Self::Item> {
216        let data = self.data;
217        loop {
218            if self.pos >= data.len() {
219                return None;
220            }
221            let (sc_at, sc_len) = find_start_code(&data[self.pos..])?;
222            let start = self.pos + sc_at + sc_len;
223            let end = match find_start_code(&data[start..]) {
224                Some((rel, _)) => start + rel,
225                None => data.len(),
226            };
227            self.pos = end;
228            if start < end {
229                return Some(&data[start..end]);
230            }
231        }
232    }
233}
234
235#[cfg(test)]
236#[path = "avc_tests.rs"]
237mod tests;