Merge pull request #1514

Reject picture slices in codec configuration records

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.

crates/h264/src/synth_tests.rs

@@ -82,6 +82,56 @@ fn pps(cabac: bool, weighted: bool) -> Vec<u8> {
     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());
+}
+
+#[test]
+fn avcc_rejects_a_valid_picture_slice_without_starting_a_picture() {
+    let sequence = sps(66, 4, false, PocType::T0);
+    let picture = pps(false, false);
+    let slice = slice_nal(
+        &Pic {
+            idr: true,
+            frame_num: 0,
+            ref_idc: 3,
+            long_term_ref: false,
+            mods: vec![],
+            mmcos: vec![],
+            mbs: vec![Some(vec![128; 384]); NMB],
+            b: false,
+            poc_lsb: Some(0),
+            poc_delta: None,
+            wt: None,
+        },
+        false,
+    );
+    let mut avcc = vec![1, 66, 0, 30, 0xff, 0xe1];
+    let nal = &sequence[4..];
+    avcc.extend_from_slice(&(nal.len() as u16).to_be_bytes());
+    avcc.extend_from_slice(nal);
+    avcc.push(2);
+    avcc.extend_from_slice(&((picture.len() - 4) as u16).to_be_bytes());
+    avcc.extend_from_slice(&picture[4..]);
+    avcc.extend_from_slice(&((slice.len() - 4) as u16).to_be_bytes());
+    avcc.extend_from_slice(&slice[4..]);
+
+    let mut decoder = Decoder::with_threads(1);
+    assert!(decoder.configure_avcc(&avcc).is_err());
+    assert!(decoder.flush().is_empty());
+    assert_eq!(decoder.stats().pictures, 0);
+}
+
 #[derive(Default, Clone)]
 struct Pic {
     idr: bool,

crates/hevc/src/decoder.rs

@@ -160,22 +160,35 @@ 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;
-        for _ in 0..num_arrays {
-            ensure!(pos + 3 <= hvcc.len(), "hvcC truncated");
-            let n = u16::from_be_bytes([hvcc[pos + 1], hvcc[pos + 2]]) as usize;
-            pos += 3;
-            for _ in 0..n {
-                ensure!(pos + 2 <= hvcc.len(), "hvcC truncated");
-                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)?;
-                pos += len;
+        for pass in 0..2 {
+            let mut pos = 23;
+            for _ in 0..num_arrays {
+                ensure!(pos <= hvcc.len() && hvcc.len() - pos >= 3, "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 {
+                    ensure!(pos <= hvcc.len() && hvcc.len() - pos >= 2, "hvcC truncated");
+                    let len = u16::from_be_bytes([hvcc[pos], hvcc[pos + 1]]) as usize;
+                    pos += 2;
+                    ensure!(len <= hvcc.len() - pos, "hvcC truncated");
+                    let nal = &hvcc[pos..pos + len];
+                    if pass == 0 {
+                        let header = NalHeader::parse(nal)?;
+                        ensure!(
+                            [nal_type::VPS, nal_type::SPS, nal_type::PPS, nal_type::PREFIX_SEI, nal_type::SUFFIX_SEI].contains(&array_type),
+                            "hvcC array has unsupported NAL type {array_type}"
+                        );
+                        ensure!(header.nal_type == array_type, "hvcC array type {array_type} disagrees with NAL header type {}", header.nal_type);
+                    } else {
+                        self.handle_nal(nal, 0)?;
+                    }
+                    pos += len;
+                }
             }
         }
+        self.nal_length_size = Some((hvcc[21] & 3) as usize + 1);
         Ok(())
     }
 
@@ -653,3 +666,20 @@ fn make_picture(o: &Output) -> Picture {
         draft: o.meta.draft,
     }
 }
+
+#[cfg(test)]
+mod config_tests {
+    use super::Decoder;
+
+    #[test]
+    fn hvcc_skips_prefix_and_suffix_sei_arrays() {
+        let mut config = vec![0; 23];
+        config[0] = 1;
+        config[22] = 2;
+        config.extend_from_slice(&[39, 0, 1, 0, 2, 0x4e, 0x01]);
+        config.extend_from_slice(&[40, 0, 1, 0, 2, 0x50, 0x01]);
+
+        let decoder = Decoder::from_hvcc(&config);
+        assert!(decoder.is_ok());
+    }
+}

crates/hevc/src/synth_tests.rs

@@ -644,3 +644,62 @@ 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());
+}
+
+#[test]
+fn hvcc_rejects_a_valid_picture_slice_without_starting_a_picture() {
+    let sets = [vps(), sps(true), pps(0, &None, false)];
+    let mut stream = Vec::new();
+    let mut saved = None;
+    encode_picture(
+        &PicDesc {
+            nal_type: 19,
+            poc: 0,
+            intra: true,
+            cu: Cu::Pcm,
+            pps_id: 0,
+            st: vec![],
+            lt: vec![],
+            num_ref: 0,
+            list_entry: None,
+            segments: whole(WC * HC),
+        },
+        &None,
+        &content(0),
+        &mut stream,
+        &mut saved,
+    );
+
+    let mut hvcc = vec![0; 23];
+    hvcc[0] = 1;
+    hvcc[22] = 4;
+    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);
+    }
+    let slice = &stream[4..];
+    hvcc.extend_from_slice(&[19, 0, 1]);
+    hvcc.extend_from_slice(&(slice.len() as u16).to_be_bytes());
+    hvcc.extend_from_slice(slice);
+
+    let mut decoder = Decoder::with_threads(1);
+    assert!(decoder.configure_hvcc(&hvcc).is_err());
+    assert!(decoder.flush().is_empty());
+    assert_eq!(decoder.stats().pictures, 0);
+}