address review: check the sample budget before reading, keep the decoder from skipping a refused sample, stop on a failed read

crates/video/src/movie.rs

@@ -380,6 +380,11 @@ impl Movie {
             .collect()
     }
 
+    /// Size the sample table declares for `sample` (0 when there is none).
+    fn declared_size(&self, sample: usize) -> usize {
+        self.mp4.tracks.get(self.track).and_then(|t| t.samples.get(sample)).map_or(0, |s| s.size as usize)
+    }
+
     fn read(&self, sample: usize) -> Result<Vec<u8>> {
         let t = self.mp4.tracks.get(self.track).ok_or(VideoError::NoVideoTrack)?;
         let size = t.samples.get(sample).map_or(0, |s| s.size);
@@ -415,6 +420,14 @@ impl Movie {
                 Ok(Frame { width: f.width, height: f.height, rgba, index, time })
             }
             Codec::Jpeg => {
+                // The frame's own size, from its header, before a pixel is decoded: the track's
+                // size was checked against `MAX_PIXELS` when the movie opened, and a frame that
+                // differs from it is refused, as a ProRes frame is.
+                let reader = image::ImageReader::with_format(std::io::Cursor::new(&data), image::ImageFormat::Jpeg);
+                let (w, h) = reader.into_dimensions().map_err(|e| VideoError::Decode(e.to_string()))?;
+                if (w, h) != (self.info.width, self.info.height) {
+                    return Err(VideoError::Decode(format!("{w}x{h} frame in a {}x{} track", self.info.width, self.info.height)));
+                }
                 let img = image::load_from_memory_with_format(&data, image::ImageFormat::Jpeg).map_err(|e| VideoError::Decode(e.to_string()))?;
                 let rgba = img.to_rgba8();
                 let (width, height) = rgba.dimensions();
@@ -457,6 +470,14 @@ impl Movie {
         Ok(self.h264_to_frame(p, index))
     }
 
+    /// Most sample bytes one frame request feeds the decoder (two samples at the 256 MiB cap):
+    /// a request decodes from the previous keyframe, and a sample table whose entries all name
+    /// one large payload would otherwise read and parse it once per sample.
+    const MAX_REQUEST_BYTES: usize = 512 << 20;
+    /// Most samples one frame request feeds without a picture coming out: the reorder depth is
+    /// at most 16, so a stream that produces nothing for this long will produce nothing.
+    const MAX_SAMPLES_WITHOUT_PICTURE: usize = 256;
+
     fn decode_h264(&mut self, index: usize) -> Result<Arc<Frame>> {
         let Codec::H264 { avcc } = &self.codec else { return Err(VideoError::Unsupported("not H.264".into())) };
         let target = self.order.get(index).copied().ok_or(VideoError::OutOfRange)?;
@@ -471,6 +492,7 @@ impl Movie {
         }
         let total = self.order.len();
         let mut last_err: Option<String> = None;
+        let (mut fed, mut without_picture) = (0usize, 0usize);
         // Bounded: each iteration feeds one sample, or flushes once at the end.
         for _ in 0..=total.saturating_add(1) {
             let Some(st) = self.h264.as_mut() else { break };
@@ -485,13 +507,21 @@ impl Movie {
                 let s = st.next;
                 st.next += 1;
                 let pts = self.rank.get(s).copied().unwrap_or(0) as i64;
+                // The declared size is checked before the read allocates it. An error leaves no decoder
+                // behind that has stepped past this sample: the next request starts from its keyframe.
+                if fed.saturating_add(self.declared_size(s)) > Self::MAX_REQUEST_BYTES {
+                    self.h264 = None;
+                    return Err(VideoError::Decode(format!("frame {index} needs more than {} MiB of samples", Self::MAX_REQUEST_BYTES >> 20)));
+                }
                 let data = match self.read(s) {
                     Ok(d) => d,
                     Err(e) => {
-                        last_err = Some(e.to_string());
-                        continue;
+                        // a failing source is not probed sample after sample (each probe may allocate a whole sample)
+                        self.h264 = None;
+                        return Err(e);
                     }
                 };
+                fed = fed.saturating_add(data.len());
                 let Some(st) = self.h264.as_mut() else { break };
                 match st.dec.decode(&data, pts) {
                     Ok(p) => p,
@@ -506,6 +536,10 @@ impl Movie {
                 st.dec.flush()
             };
             let Some(st) = self.h264.as_mut() else { break };
+            without_picture = if pics.is_empty() { without_picture + 1 } else { 0 };
+            if without_picture > Self::MAX_SAMPLES_WITHOUT_PICTURE {
+                return Err(VideoError::Decode(format!("no picture in {} samples before frame {index}", Self::MAX_SAMPLES_WITHOUT_PICTURE)));
+            }
             for p in pics {
                 let q = usize::try_from(p.pts).unwrap_or(usize::MAX);
                 st.last_out = Some(q);

crates/video/src/tests.rs

@@ -121,7 +121,7 @@ fn prores_frame_size_is_checked_before_decoding() {
 }
 
 /// In-memory movie bytes, zero-padded up to `len` without storing the padding. Records the
-/// largest read asked of it and refuses reads over 1 MiB (the test movies are much smaller).
+/// largest read asked of it and refuses reads over 32 MiB (the test movies are much smaller).
 struct PaddedSource {
     bytes: Vec<u8>,
     len: u64,
@@ -134,7 +134,7 @@ impl ByteSource for PaddedSource {
     }
     fn read_at(&self, offset: u64, buf: &mut [u8]) -> std::io::Result<()> {
         self.largest.fetch_max(buf.len(), Ordering::Relaxed);
-        if buf.len() > 1 << 20 {
+        if buf.len() > 32 << 20 {
             return Err(std::io::Error::other("read too large for the test source"));
         }
         let start = usize::try_from(offset).unwrap_or(usize::MAX);
@@ -161,6 +161,43 @@ fn h264_sample_reads_are_capped() {
     assert!(largest <= 256 << 20, "a {largest}-byte read reached the source");
 }
 
+#[test]
+fn one_frame_request_reads_a_bounded_amount_of_samples() {
+    // 17 frames whose `stsz` entries each claim 32 MiB: under the per-sample cap, but together
+    // past what one request may feed the decoder.
+    let mut bytes = h264_mp4(17, 17).unwrap();
+    let at = bytes.windows(4).position(|w| w == b"stsz").unwrap();
+    if bytes[at + 8..at + 12] == [0; 4] {
+        // one size per sample
+        for i in 0..17 {
+            let field = at + 16 + 4 * i;
+            bytes[field..field + 4].copy_from_slice(&(32u32 << 20).to_be_bytes());
+        }
+    } else {
+        // the writer stored one size for all the (equally large) samples
+        bytes[at + 8..at + 12].copy_from_slice(&(32u32 << 20).to_be_bytes());
+    }
+    let src = Arc::new(PaddedSource { bytes, len: 1 << 30, largest: AtomicUsize::new(0) });
+    let mut m = Movie::open_source(src).unwrap();
+    let r = m.frame(16);
+    assert!(matches!(&r, Err(VideoError::Decode(e)) if e.contains("needs more than 512 MiB of samples")), "{r:?}");
+    // the sample that was refused is not skipped: asking again fails the same way instead of decoding past it
+    let again = m.frame(16);
+    assert!(matches!(&again, Err(VideoError::Decode(e)) if e.contains("needs more than 512 MiB of samples")), "{again:?}");
+}
+
+#[test]
+fn mjpeg_frame_size_is_checked_before_decoding() {
+    // The first frame's SOF0 header claims 16 000 x 16 000: refused from the header, not decoded.
+    let mut bytes = mjpeg_mov(2).unwrap();
+    let sof = bytes.windows(2).position(|w| w == [0xFF, 0xC0]).unwrap();
+    bytes[sof + 5..sof + 9].copy_from_slice(&[0x3E, 0x80, 0x3E, 0x80]);
+    let mut m = Movie::open(bytes).unwrap();
+    let r = m.frame(0);
+    assert!(matches!(&r, Err(VideoError::Decode(e)) if e.contains("16000x16000 frame in a 32x24 track")), "{r:?}");
+    assert!(m.frame(1).is_ok(), "the untouched frame still decodes");
+}
+
 #[test]
 fn unsupported_codec_opens_but_does_not_decode() {
     let mut mw = Mp4Writer::new(Cursor::new(Vec::new()), WriterOptions::new(Brand::Mp4)).unwrap();