From 3eaf216405b4e2d19b9e38bc4ef0f5d9ea5ddde6 Mon Sep 17 00:00:00 2001
From: "Dylan M. Taylor" <dylan@dylanmtaylor.com>
Date: Wed, 4 Nov 2026 01:48:07 +0000
Subject: [PATCH] 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 | 19 ++++++++++++++-----
 crates/video/src/tests.rs | 16 ++++++++++++----
 2 files changed, 26 insertions(+), 9 deletions(-)

diff --git a/crates/video/src/movie.rs b/crates/video/src/movie.rs
index eefceb50..1a3d3de9 100644
--- a/crates/video/src/movie.rs
+++ b/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);
@@ -502,17 +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());
-                if fed > Self::MAX_REQUEST_BYTES {
-                    return Err(VideoError::Decode(format!("frame {index} needs more than {} MiB of samples", Self::MAX_REQUEST_BYTES >> 20)));
-                }
                 let Some(st) = self.h264.as_mut() else { break };
                 match st.dec.decode(&data, pts) {
                     Ok(p) => p,
diff --git a/crates/video/src/tests.rs b/crates/video/src/tests.rs
index d119d1ec..141d6c8e 100644
--- a/crates/video/src/tests.rs
+++ b/crates/video/src/tests.rs
@@ -167,15 +167,23 @@ fn one_frame_request_reads_a_bounded_amount_of_samples() {
     // 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();
-    assert_eq!(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());
+    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]