saiz/saio/hvcc: bound Vec::with_capacity on counts

src/moov/trak/mdia/minf/stbl/saiz.rs

@@ -1,5 +1,7 @@
 use crate::*;
 
+/// Auxiliary information type and parameter shared by `saiz` and `saio`,
+/// ISO/IEC 14496-12:2022 Sect 8.7.8 / 8.7.9.
 #[derive(Debug, Clone, PartialEq, Eq)]
 #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
 pub struct AuxInfo {
@@ -51,7 +53,12 @@ impl AtomExt for Saiz {
         let default_sample_info_size = u8::decode(buf)?;
         let sample_count = u32::decode(buf)?;
         if default_sample_info_size == 0 {
-            let mut sample_info_size = Vec::with_capacity(sample_count as usize);
+            // Each entry is a single byte; reject counts that cannot
+            // possibly fit in the remaining buffer before allocating.
+            if sample_count as usize > buf.remaining() {
+                return Err(Error::OutOfBounds);
+            }
+            let mut sample_info_size = Vec::with_capacity((sample_count as usize).min(4096));
             for _ in 0..sample_count {
                 sample_info_size.push(u8::decode(buf)?);
             }
@@ -115,7 +122,13 @@ impl AtomExt for Saio {
             });
         }
         let entry_count = u32::decode(buf)?;
-        let mut offsets = Vec::with_capacity(entry_count as usize);
+        // Entries are 4 bytes (v0) or 8 bytes (v1); reject counts that
+        // cannot possibly fit in the remaining buffer before allocating.
+        let per_entry = if ext.version == SaioVersion::V1 { 8 } else { 4 };
+        if entry_count as usize > buf.remaining() / per_entry {
+            return Err(Error::OutOfBounds);
+        }
+        let mut offsets = Vec::with_capacity((entry_count as usize).min(4096));
         for _ in 0..entry_count {
             if ext.version == SaioVersion::V0 {
                 let offset = u32::decode(buf)? as u64;
@@ -427,6 +440,32 @@ mod tests {
         );
     }
 
+    // Regression for issue #156: a u32::MAX sample_count must fail
+    // cleanly without attempting a ~4 GiB upfront allocation.
+    const ENCODED_SAIZ_HUGE_COUNT: &[u8] = &[
+        0x00, 0x00, 0x00, 0x11, 0x73, 0x61, 0x69, 0x7a, 0x00, 0x00, 0x00, 0x00, 0x00, 0xFF, 0xFF,
+        0xFF, 0xFF,
+    ];
+
+    #[test]
+    fn test_saiz_huge_count() {
+        let buf: &mut std::io::Cursor<&&[u8]> = &mut std::io::Cursor::new(&ENCODED_SAIZ_HUGE_COUNT);
+        assert!(matches!(Saiz::decode(buf), Err(Error::OverDecode(_))));
+    }
+
+    // Regression for issue #156: a u32::MAX entry_count must fail
+    // cleanly without attempting a multi-GiB upfront allocation.
+    const ENCODED_SAIO_HUGE_COUNT: &[u8] = &[
+        0x00, 0x00, 0x00, 0x10, 0x73, 0x61, 0x69, 0x6f, 0x01, 0x00, 0x00, 0x00, 0xFF, 0xFF, 0xFF,
+        0xFF,
+    ];
+
+    #[test]
+    fn test_saio_huge_count() {
+        let buf: &mut std::io::Cursor<&&[u8]> = &mut std::io::Cursor::new(&ENCODED_SAIO_HUGE_COUNT);
+        assert!(matches!(Saio::decode(buf), Err(Error::OverDecode(_))));
+    }
+
     #[test]
     fn test_saio_encode_cenc() {
         let saio = Saio {

src/moov/trak/mdia/minf/stbl/stsd/hevc/hvcc.rs

@@ -1,5 +1,7 @@
 use crate::*;
 
+/// HEVCConfigurationBox / `HEVCDecoderConfigurationRecord`,
+/// ISO/IEC 14496-15:2022 Sect 8.3.3.
 #[derive(Default, Debug, Clone, PartialEq, Eq)]
 #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
 pub struct Hvcc {
@@ -24,6 +26,9 @@ pub struct Hvcc {
 }
 
 impl Hvcc {
+    /// Returns an `Hvcc` with `configuration_version` set to 1 and all
+    /// other fields zeroed, matching `HEVCDecoderConfigurationRecord`'s
+    /// only currently-defined version.
     pub fn new() -> Self {
         Self {
             configuration_version: 1,
@@ -32,6 +37,8 @@ impl Hvcc {
     }
 }
 
+/// One entry of `HEVCDecoderConfigurationRecord.arrays`: a group of NAL
+/// units of a single `nal_unit_type` (typically VPS, SPS, PPS, or SEI).
 #[derive(Debug, Clone, PartialEq, Eq, Default)]
 #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
 pub struct HvcCArray {
@@ -72,7 +79,14 @@ impl Atom for Hvcc {
         for _ in 0..num_of_arrays {
             let params = u8::decode(buf)?;
             let num_nalus = u16::decode(buf)?;
-            let mut nalus = Vec::with_capacity(num_nalus as usize);
+            // Each NALU has at least a u16 length prefix (2 bytes); reject
+            // counts that cannot possibly fit in the remaining buffer before
+            // allocating. Real HEVC arrays hold a handful of VPS/SPS/PPS/SEI
+            // entries, so cap the upfront reservation at 64.
+            if num_nalus as usize > buf.remaining() / 2 {
+                return Err(Error::OutOfBounds);
+            }
+            let mut nalus = Vec::with_capacity((num_nalus as usize).min(64));
 
             for _ in 0..num_nalus {
                 let size = u16::decode(buf)? as usize;
@@ -420,6 +434,36 @@ mod tests {
         assert_eq!(decoded, hvcc);
     }
 
+    // Regression for issue #156: a u16::MAX num_nalus must fail cleanly
+    // without attempting a ~1.5 MiB upfront allocation for empty input.
+    const ENCODED_HVCC_HUGE_NALU_COUNT: &[u8] = &[
+        0x00, 0x00, 0x00, 0x22, // size = 34
+        0x68, 0x76, 0x63, 0x43, // "hvcC"
+        // 22-byte configuration record up through length_size_minus_one:
+        0x01, // configuration_version
+        0x00, // profile_space|tier|profile_idc
+        0x00, 0x00, 0x00, 0x00, // general_profile_compatibility_flags
+        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // general_constraint_indicator_flags
+        0x00, // general_level_idc
+        0x00, 0x00, // min_spatial_segmentation_idc
+        0x00, // parallelism_type
+        0x00, // chroma_format_idc
+        0x00, // bit_depth_luma_minus8
+        0x00, // bit_depth_chroma_minus8
+        0x00, 0x00, // avg_frame_rate
+        0x00, // constant_frame_rate|num_temporal_layers|nested|length_size
+        0x01, // num_of_arrays = 1
+        0x00, // first array params byte
+        0xFF, 0xFF, // num_nalus = u16::MAX
+    ];
+
+    #[test]
+    fn test_hvcc_huge_nalu_count() {
+        let buf: &mut std::io::Cursor<&&[u8]> =
+            &mut std::io::Cursor::new(&ENCODED_HVCC_HUGE_NALU_COUNT);
+        assert!(matches!(Hvcc::decode(buf), Err(Error::OverDecode(_))));
+    }
+
     #[test]
     fn test_hvcc_mpeg_encode() {
         let hvcc = Hvcc {