diff --git a/crates/h264/src/decoder.rs b/crates/h264/src/decoder.rs index 37ef1f29..64fcb48b 100644 --- a/crates/h264/src/decoder.rs +++ b/crates/h264/src/decoder.rs @@ -216,26 +216,37 @@ impl Decoder { pub fn configure_avcc(&mut self, avcc: &[u8]) -> Result<()> { ensure!(avcc.len() >= 7, "avcC too short"); ensure!(avcc[0] == 1, "unsupported avcC version {}", avcc[0]); - self.nal_length_size = Some((avcc[4] & 3) as usize + 1); let mut pos = 5; - let read_sets = |pos: &mut usize, count: usize, d: &mut Decoder| -> Result<()> { + let mut sets = Vec::new(); + let read_sets = |pos: &mut usize, count: usize, sets: &mut Vec<(usize, usize)>| -> Result<()> { for _ in 0..count { ensure!(*pos + 2 <= avcc.len(), "avcC truncated"); let len = u16::from_be_bytes([avcc[*pos], avcc[*pos + 1]]) as usize; *pos += 2; ensure!(*pos + len <= avcc.len(), "avcC truncated"); - d.handle_nal(&avcc[*pos..*pos + len], 0)?; + sets.push((*pos, len)); *pos += len; } Ok(()) }; let nsps = (avcc[pos] & 0x1f) as usize; pos += 1; - read_sets(&mut pos, nsps, self)?; + read_sets(&mut pos, nsps, &mut sets)?; ensure!(pos < avcc.len(), "avcC truncated"); let npps = avcc[pos] as usize; pos += 1; - read_sets(&mut pos, npps, self) + read_sets(&mut pos, npps, &mut sets)?; + for (i, (start, len)) in sets.iter().enumerate() { + let nal = &avcc[*start..*start + *len]; + let header = nal.first().copied().ok_or_else(|| Error::Invalid("empty avcC parameter set".into()))?; + let expected = if i < nsps { nal_type::SPS } else { nal_type::PPS }; + ensure!(header & 0x1f == expected, "avcC parameter set has NAL type {}, expected {expected}", header & 0x1f); + } + for (start, len) in sets { + self.handle_nal(&avcc[start..start + len], 0)?; + } + self.nal_length_size = Some((avcc[4] & 3) as usize + 1); + Ok(()) } /// Length of NAL length prefixes (from avcC), or None for Annex-B input. @@ -675,3 +686,17 @@ fn make_picture(o: &Output) -> Picture { draft: o.meta.draft, } } + +#[cfg(test)] +mod config_tests { + use super::Decoder; + + #[test] + fn avcc_rejects_picture_slices_without_starting_a_picture() { + let config = [1, 66, 0, 30, 0xff, 0xe1, 0, 2, 0x41, 0x80, 0]; + let mut decoder = Decoder::with_threads(1); + assert!(decoder.configure_avcc(&config).is_err()); + assert!(decoder.pending.is_none()); + assert_eq!(decoder.stats().pictures, 0); + } +} diff --git a/crates/h264/src/synth_tests.rs b/crates/h264/src/synth_tests.rs index a0c9da2f..8588c540 100644 --- a/crates/h264/src/synth_tests.rs +++ b/crates/h264/src/synth_tests.rs @@ -82,6 +82,20 @@ fn pps(cabac: bool, weighted: bool) -> Vec { nal(3, 8, &w.finish()) } +#[test] +fn avcc_accepts_generated_sps_and_pps() { + let sps = sps(66, 4, false, PocType::T2); + let pps = pps(false, false); + let (sps, pps) = (&sps[4..], &pps[4..]); + let mut avcc = vec![1, 66, 0, 30, 0xff, 0xe1]; + avcc.extend_from_slice(&(sps.len() as u16).to_be_bytes()); + avcc.extend_from_slice(sps); + avcc.push(1); + avcc.extend_from_slice(&(pps.len() as u16).to_be_bytes()); + avcc.extend_from_slice(pps); + assert!(Decoder::from_avcc(&avcc).is_ok()); +} + #[derive(Default, Clone)] struct Pic { idr: bool, diff --git a/crates/hevc/src/decoder.rs b/crates/hevc/src/decoder.rs index 4b46044d..a78600f2 100644 --- a/crates/hevc/src/decoder.rs +++ b/crates/hevc/src/decoder.rs @@ -160,11 +160,12 @@ impl Decoder { pub fn configure_hvcc(&mut self, hvcc: &[u8]) -> Result<()> { ensure!(hvcc.len() >= 23, "hvcC too short"); ensure!(hvcc[0] == 1, "unsupported hvcC version {}", hvcc[0]); - self.nal_length_size = Some((hvcc[21] & 3) as usize + 1); let num_arrays = hvcc[22] as usize; let mut pos = 23; + let mut sets = Vec::new(); for _ in 0..num_arrays { ensure!(pos + 3 <= hvcc.len(), "hvcC truncated"); + let array_type = hvcc[pos] & 0x3f; let n = u16::from_be_bytes([hvcc[pos + 1], hvcc[pos + 2]]) as usize; pos += 3; for _ in 0..n { @@ -172,10 +173,19 @@ impl Decoder { let len = u16::from_be_bytes([hvcc[pos], hvcc[pos + 1]]) as usize; pos += 2; ensure!(pos + len <= hvcc.len(), "hvcC truncated"); - self.handle_nal(&hvcc[pos..pos + len], 0)?; + sets.push((array_type, &hvcc[pos..pos + len])); pos += len; } } + for (array_type, nal) in &sets { + let header = NalHeader::parse(nal)?; + ensure!([nal_type::VPS, nal_type::SPS, nal_type::PPS].contains(array_type), "hvcC array has non-parameter-set NAL type {array_type}"); + ensure!(header.nal_type == *array_type, "hvcC array type {array_type} disagrees with NAL header type {}", header.nal_type); + } + for (_, nal) in sets { + self.handle_nal(nal, 0)?; + } + self.nal_length_size = Some((hvcc[21] & 3) as usize + 1); Ok(()) } @@ -651,3 +661,20 @@ fn make_picture(o: &Output) -> Picture { draft: o.meta.draft, } } + +#[cfg(test)] +mod config_tests { + use super::Decoder; + + #[test] + fn hvcc_rejects_picture_slices_without_starting_a_picture() { + let mut config = vec![0; 23]; + config[0] = 1; + config[22] = 1; + config.extend_from_slice(&[32, 0, 1, 0, 2, 0x02, 0x01]); + let mut decoder = Decoder::with_threads(1); + assert!(decoder.configure_hvcc(&config).is_err()); + assert!(decoder.pending.is_none()); + assert_eq!(decoder.stats().pictures, 0); + } +} diff --git a/crates/hevc/src/synth_tests.rs b/crates/hevc/src/synth_tests.rs index 4c2f1c96..52a6e177 100644 --- a/crates/hevc/src/synth_tests.rs +++ b/crates/hevc/src/synth_tests.rs @@ -644,3 +644,18 @@ fn picture_store_is_checked_against_the_memory_budget() { assert!(!decodes_within(&stream, pictures * one - 1)); assert!(decodes_within(&stream, deckcraft_bitstream::budget::DECODE_MEMORY_BUDGET)); } + +#[test] +fn hvcc_accepts_generated_vps_sps_and_pps() { + let sets = [vps(), sps(true), pps(0, &None, false)]; + let mut hvcc = vec![0; 23]; + hvcc[0] = 1; + hvcc[22] = sets.len() as u8; + for (kind, set) in [32, 33, 34].into_iter().zip(sets) { + let nal = &set[4..]; + hvcc.extend_from_slice(&[kind, 0, 1]); + hvcc.extend_from_slice(&(nal.len() as u16).to_be_bytes()); + hvcc.extend_from_slice(nal); + } + assert!(Decoder::from_hvcc(&hvcc).is_ok()); +}