codec: validate configuration NAL types

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);
+    }
+}

crates/h264/src/synth_tests.rs

@@ -82,6 +82,20 @@ 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());
+}
+
 #[derive(Default, Clone)]
 struct Pic {
     idr: bool,

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);
+    }
+}

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());
+}