Video: check MJPEG frame size before decoding; budget each H.264 frame request (#1326)

* fix(video): check MJPEG frame size from its header and budget H.264 frame requests

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

* address review: restore lenient MJPEG sizing and unreadable-sample skipping; per-request sample cap 2 GiB

* address review: an unreadable sample counts toward the no-picture budget

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,13 @@ 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: a frame may differ from
+                // the track's size (as before), but not exceed `MAX_PIXELS`.
+                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 u64::from(w) * u64::from(h) > MAX_PIXELS {
+                    return Err(VideoError::Decode("picture too large".into()));
+                }
                 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 +469,14 @@ impl Movie {
         Ok(self.h264_to_frame(p, index))
     }
 
+    /// Most sample bytes one frame request feeds the decoder (eight 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 = 2 << 30;
+    /// 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 +491,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 +506,29 @@ 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,
+                    // an unreadable sample is skipped, as before, and counts as a sample without a picture: the
+                    // request's byte and no-picture budgets bound the walk
                     Err(e) => {
                         last_err = Some(e.to_string());
+                        without_picture += 1;
+                        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
+                            )));
+                        }
                         continue;
                     }
                 };
+                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 +543,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,80 @@ 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() {
+    // 65 frames whose `stsz` entries each claim 32 MiB: under the per-sample cap, but together
+    // past the 2 GiB one request may feed the decoder.
+    let mut bytes = h264_mp4(65, 65).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..65 {
+            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: 3 << 30, largest: AtomicUsize::new(0) });
+    let mut m = Movie::open_source(src).unwrap();
+    let r = m.frame(64);
+    assert!(matches!(&r, Err(VideoError::Decode(e)) if e.contains("needs more than 2048 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(64);
+    assert!(matches!(&again, Err(VideoError::Decode(e)) if e.contains("needs more than 2048 MiB of samples")), "{again:?}");
+}
+
+/// A source that, once `failing` is set, cannot be read at all.
+struct UnreadableSamples {
+    bytes: Vec<u8>,
+    failing: std::sync::atomic::AtomicBool,
+}
+
+impl ByteSource for UnreadableSamples {
+    fn len(&self) -> u64 {
+        self.bytes.len() as u64
+    }
+    fn read_at(&self, offset: u64, buf: &mut [u8]) -> std::io::Result<()> {
+        let start = usize::try_from(offset).unwrap_or(usize::MAX);
+        if self.failing.load(Ordering::Relaxed) || start.saturating_add(buf.len()) > self.bytes.len() {
+            return Err(std::io::Error::other("unreadable"));
+        }
+        buf.copy_from_slice(&self.bytes[start..start + buf.len()]);
+        Ok(())
+    }
+}
+
+#[test]
+fn a_request_over_unreadable_samples_stops_at_the_no_picture_budget() {
+    // 400 frames of one GOP, none of whose samples can be read once the movie is open: the walk gives up after 256 skipped samples
+    let bytes = h264_mp4(400, 400).unwrap();
+    let src = Arc::new(UnreadableSamples { bytes, failing: false.into() });
+    let mut m = Movie::open_source(src.clone()).unwrap();
+    src.failing.store(true, Ordering::Relaxed);
+    let r = m.frame(399);
+    assert!(matches!(&r, Err(VideoError::Decode(e)) if e.contains("no picture in 256 samples")), "{r:?}");
+}
+
+#[test]
+fn mjpeg_frame_size_is_capped_from_its_header_but_may_differ_from_the_track() {
+    // A frame whose header differs from the track's size still decodes (within the dimension cap, as before)...
+    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(&[0, 8, 0, 8]);
+    let mut m = Movie::open(bytes).unwrap();
+    assert!(m.frame(0).is_ok());
+    // ...but one claiming 16 000 x 16 000 is 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("picture too large")), "{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();