Video: check the whole picture store against a memory budget (#398)

Per-axis size caps let AV1 accept 65536x65536 frames (over 8 GiB of luma)
and did not count references or pictures in flight. A shared budget module
in deckcraft-bitstream sizes a picture with checked arithmetic, and the
AV1, VP9, HEVC and H.264 decoders refuse a picture size whose current
picture, reference slots and in-flight pictures would exceed 4 GiB (1 GiB
on 32-bit), before allocating. set_memory_budget lowers it for hosts and tests.

Co-authored-by: Cursor <cursoragent@cursor.com>

crates/av1/src/decoder.rs

@@ -206,12 +206,35 @@ pub struct Decoder {
     /// Frames submitted to the workers and not known to be finished (oldest first).
     in_flight: VecDeque<Arc<FrameSlot>>,
     max_in_flight: usize,
+    /// Bytes the picture store may take ([`deckcraft_bitstream::budget`]).
+    memory_budget: u64,
     /// Shown pictures not yet returned, in output order.
     out: VecDeque<Output>,
     /// pts of the temporal unit being decoded.
     pts: i64,
 }
 
+/// Whether `pictures` frames of `width` x `height`, padded as `FrameBuf::new` pads them, fit in
+/// `budget` bytes.
+fn store_fits(width: u32, height: u32, color: &crate::header::ColorConfig, pictures: usize, budget: u64) -> bool {
+    let shape = deckcraft_bitstream::budget::PictureShape {
+        width: u64::from(width).div_ceil(128) * 128 + 160,
+        height: u64::from(height).div_ceil(128) * 128 + 160,
+        shift_x: u32::from(color.subsampling_x),
+        shift_y: u32::from(color.subsampling_y),
+        chroma_planes: if color.mono_chrome { 0 } else { 2 },
+        bytes_per_sample: 2,
+    };
+    shape.fits(pictures as u64, budget)
+}
+
+impl Decoder {
+    /// Cap the bytes of decoded pictures (references, current, in flight) this decoder may hold.
+    pub fn set_memory_budget(&mut self, bytes: u64) {
+        self.memory_budget = bytes;
+    }
+}
+
 fn default_threads() -> usize {
     #[cfg(all(feature = "threads", not(target_arch = "wasm32")))]
     {
@@ -268,6 +291,7 @@ impl Decoder {
             workers,
             in_flight: VecDeque::new(),
             max_in_flight: fthreads + 2,
+            memory_budget: deckcraft_bitstream::budget::DECODE_MEMORY_BUDGET,
             out: VecDeque::new(),
             pts: 0,
         }
@@ -508,6 +532,12 @@ impl Decoder {
         if fh.upscaled_width > 65536 || fh.frame_height > 65536 {
             return Err(Error::Invalid("frame size"));
         }
+        // The axis cap alone allows 8 GiB pictures: check the whole store (every reference slot,
+        // the new frame and those in flight, as `FrameBuf::new` pads them) before allocating.
+        let width = fh.upscaled_width.max(fh.frame_width);
+        if !store_fits(width, fh.frame_height, &seq.color, NUM_REF_FRAMES + 1 + self.max_in_flight, self.memory_budget) {
+            return Err(Error::Invalid("frame size exceeds the decoder memory budget"));
+        }
         if !fh.frame_is_intra {
             for i in 0..REFS_PER_FRAME {
                 if self.ref_frames[fh.ref_frame_idx[i]].is_none() {
@@ -947,3 +977,32 @@ fn project(fs: &mut FrameState, src: usize, dst_sign: i32) -> bool {
     }
     true
 }
+
+#[cfg(test)]
+mod budget_tests {
+    use super::*;
+    use deckcraft_bitstream::budget::DECODE_MEMORY_BUDGET;
+
+    fn color(mono: bool) -> crate::header::ColorConfig {
+        crate::header::ColorConfig { mono_chrome: mono, subsampling_x: 1, subsampling_y: 1, ..Default::default() }
+    }
+
+    #[test]
+    fn the_axis_maximum_is_over_budget_and_4k_is_not() {
+        let refs = NUM_REF_FRAMES + 1 + 4;
+        assert!(!store_fits(65536, 65536, &color(false), refs, DECODE_MEMORY_BUDGET));
+        assert!(!store_fits(65536, 65536, &color(true), 1, DECODE_MEMORY_BUDGET));
+        assert!(store_fits(3840, 2160, &color(false), refs, DECODE_MEMORY_BUDGET));
+        assert!(store_fits(1920, 1080, &color(false), refs, DECODE_MEMORY_BUDGET));
+    }
+
+    #[test]
+    fn exact_budget_boundary_and_reference_count() {
+        // 8x8 pads to 288x288: luma 82944 + chroma 2 * 20736 samples, 2 bytes each.
+        let one = (288 * 288 + 2 * 144 * 144) * 2;
+        assert!(store_fits(8, 8, &color(false), 9, 9 * one));
+        assert!(!store_fits(8, 8, &color(false), 9, 9 * one - 1));
+        assert!(!store_fits(8, 8, &color(false), 10, 9 * one));
+        assert!(store_fits(8, 8, &color(true), 9, 9 * 288 * 288 * 2));
+    }
+}

crates/bitstream/src/budget.rs

@@ -0,0 +1,82 @@
+//! Memory budget for decoded pictures, shared by the video decoders.
+//!
+//! A per-axis size cap does not bound memory: a decoder holds the current picture, every
+//! reference picture and the pictures still being decoded on worker threads. Each decoder checks
+//! the whole picture store against [`DECODE_MEMORY_BUDGET`] with checked arithmetic before it
+//! allocates for a new picture size, so a hostile header is an error, not an out-of-memory abort.
+
+/// Bytes the decoded pictures of one decoder may take: 4 GiB on 64-bit targets, 1 GiB on 32-bit
+/// ones (wasm32 has a 4 GiB address space in all).
+pub const DECODE_MEMORY_BUDGET: u64 = if cfg!(target_pointer_width = "64") { 4 << 30 } else { 1 << 30 };
+
+/// The shape of one decoded picture as a decoder allocates it.
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub struct PictureShape {
+    /// Allocated luma samples per row and rows (padding included).
+    pub width: u64,
+    pub height: u64,
+    /// Chroma subsampling shifts (0 or 1); ignored when `chroma_planes` is 0.
+    pub shift_x: u32,
+    pub shift_y: u32,
+    /// Chroma planes (0 for monochrome, 2 otherwise).
+    pub chroma_planes: u64,
+    /// Bytes per sample (1 for 8-bit buffers, 2 for 16-bit ones).
+    pub bytes_per_sample: u64,
+}
+
+impl PictureShape {
+    /// Bytes of one picture, `None` if it does not fit in a `u64`.
+    pub fn bytes(&self) -> Option<u64> {
+        let luma = self.width.checked_mul(self.height)?;
+        let cw = self.width.div_ceil(1 << self.shift_x.min(1));
+        let ch = self.height.div_ceil(1 << self.shift_y.min(1));
+        let chroma = cw.checked_mul(ch)?.checked_mul(self.chroma_planes)?;
+        luma.checked_add(chroma)?.checked_mul(self.bytes_per_sample)
+    }
+
+    /// Bytes of `pictures` such pictures, `None` on overflow.
+    pub fn store_bytes(&self, pictures: u64) -> Option<u64> {
+        self.bytes()?.checked_mul(pictures)
+    }
+
+    /// Whether `pictures` such pictures fit in `budget` bytes.
+    pub fn fits(&self, pictures: u64, budget: u64) -> bool {
+        self.store_bytes(pictures).is_some_and(|b| b <= budget)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn shape(w: u64, h: u64) -> PictureShape {
+        PictureShape { width: w, height: h, shift_x: 1, shift_y: 1, chroma_planes: 2, bytes_per_sample: 2 }
+    }
+
+    #[test]
+    fn bytes_count_luma_and_subsampled_chroma() {
+        assert_eq!(shape(4, 4).bytes(), Some((16 + 2 * 4) * 2));
+        // Odd sizes round chroma up.
+        assert_eq!(shape(3, 3).bytes(), Some((9 + 2 * 4) * 2));
+        let mono = PictureShape { chroma_planes: 0, bytes_per_sample: 1, ..shape(10, 10) };
+        assert_eq!(mono.bytes(), Some(100));
+        let full = PictureShape { shift_x: 0, shift_y: 0, ..shape(10, 10) };
+        assert_eq!(full.bytes(), Some(600));
+    }
+
+    #[test]
+    fn exact_budget_boundary() {
+        let s = shape(4, 4); // 48 bytes
+        assert!(s.fits(10, 480));
+        assert!(!s.fits(10, 479));
+        assert!(!s.fits(11, 480));
+    }
+
+    #[test]
+    fn overflow_is_never_within_budget() {
+        let huge = shape(u64::MAX / 2, 3);
+        assert_eq!(huge.bytes(), None);
+        assert!(!huge.fits(1, u64::MAX));
+        assert!(!shape(1 << 31, 1 << 31).fits(u64::MAX, u64::MAX));
+    }
+}

crates/bitstream/src/lib.rs

@@ -3,12 +3,15 @@
 //! - [`BitReader`]: MSB-first reader with Exp-Golomb (`ue(v)`, `se(v)`), as used by H.264/HEVC/AV1 headers.
 //! - [`BitWriter`]: the matching writer (for encoders and header rewriting).
 //! - [`unescape_rbsp`] / [`escape_rbsp`]: remove/insert H.264/HEVC emulation-prevention bytes.
+//! - [`budget`]: the memory budget video decoders check their picture store against.
 //! - [`annexb_nals`] / [`length_prefixed_nals`]: split Annex-B byte streams and AVCC/HVCC samples into NAL units.
 
 #![cfg_attr(not(test), deny(clippy::unwrap_used, clippy::expect_used, clippy::panic, clippy::unimplemented, clippy::todo, clippy::unreachable))]
 
 use std::fmt;
 
+pub mod budget;
+
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 pub enum BitError {
     /// Tried to read past the end of the buffer.

crates/h264/src/decoder.rs

@@ -95,6 +95,8 @@ pub struct Decoder {
     pool: Option<rayon::ThreadPool>,
     in_flight: VecDeque<FrameRef>,
     max_in_flight: usize,
+    /// Bytes the picture store may take ([`deckcraft_bitstream::budget`]).
+    memory_budget: u64,
     threads: usize,
     /// Draft mode: skip deblocking of non-reference pictures ([`Decoder::set_draft`]).
     draft: bool,
@@ -178,6 +180,7 @@ impl Decoder {
             pool,
             in_flight: VecDeque::new(),
             max_in_flight: threads.max(1) + 2,
+            memory_budget: deckcraft_bitstream::budget::DECODE_MEMORY_BUDGET,
             threads: workers,
             draft: false,
         }
@@ -421,7 +424,34 @@ impl Decoder {
         Ok(())
     }
 
+    /// Cap the bytes of decoded pictures (current, references, in flight) this decoder may hold.
+    pub fn set_memory_budget(&mut self, bytes: u64) {
+        self.memory_budget = bytes;
+    }
+
+    /// Refuse an SPS whose pictures would not fit the memory budget, before anything is allocated.
+    fn check_budget(&self, sps: &Sps) -> Result<()> {
+        let chroma = sps.chroma_array_type();
+        let shape = deckcraft_bitstream::budget::PictureShape {
+            width: u64::from(sps.width()),
+            height: u64::from(sps.height()),
+            shift_x: u32::from(chroma != 3),
+            shift_y: u32::from(chroma == 1),
+            chroma_planes: if chroma == 0 { 0 } else { 2 },
+            bytes_per_sample: 1,
+        };
+        let pictures = (sps.max_dpb_frames() + 1 + self.max_in_flight) as u64;
+        ensure!(
+            shape.fits(pictures, self.memory_budget),
+            "{}x{} pictures with {pictures} in the decoder exceed the memory budget",
+            sps.width(),
+            sps.height()
+        );
+        Ok(())
+    }
+
     fn start_picture(&mut self, sh: &SliceHeader, sps: &Arc<Sps>, pts: i64) -> Result<()> {
+        self.check_budget(sps)?;
         // Activate SPS; a resolution change flushes the DPB.
         let changed = match &self.active_sps {
             None => true,

crates/h264/src/synth_tests.rs

@@ -539,3 +539,43 @@ fn poc_type1_reordering() {
         }
     }
 }
+
+/// One IDR picture of PCM macroblocks.
+fn idr_stream(max_refs: u32) -> Vec<u8> {
+    let mut seed = 5u32;
+    let mbs: Vec<Option<Vec<u8>>> = (0..NMB).map(|_| Some(random_mb(&mut seed))).collect();
+    let mut stream = sps(66, max_refs, false, PocType::T2);
+    stream.extend(pps(false, false));
+    stream.extend(slice_nal(&Pic { idr: true, frame_num: 0, ref_idc: 3, mbs, ..Default::default() }, false));
+    stream
+}
+
+fn decodes_within(stream: &[u8], budget: u64) -> bool {
+    let mut dec = Decoder::with_threads(1);
+    dec.set_memory_budget(budget);
+    let Ok(mut out) = dec.decode(stream, 0) else { return false };
+    out.extend(dec.flush());
+    dec.take_error().is_none() && out.len() == 1
+}
+
+#[test]
+fn picture_store_is_checked_against_the_memory_budget() {
+    let stream = idr_stream(4);
+    let sps_nal = sps(66, 4, false, PocType::T2);
+    let parsed = crate::params::Sps::parse(&deckcraft_bitstream::unescape_rbsp(&sps_nal[5..])).unwrap();
+    // 4:2:0, 8-bit: luma plus two quarter-size chroma planes per picture; the DPB, the current
+    // picture and threads + 2 in flight.
+    let one = (NMB * 256 + 2 * NMB * 64) as u64;
+    let pictures = (parsed.max_dpb_frames() + 1 + 3) as u64;
+    assert!(decodes_within(&stream, pictures * one));
+    assert!(!decodes_within(&stream, pictures * one - 1));
+    assert!(decodes_within(&stream, deckcraft_bitstream::budget::DECODE_MEMORY_BUDGET));
+}
+
+#[test]
+fn largest_accepted_h264_size_is_over_the_default_budget() {
+    // 1024 x 1024 macroblocks (16384 x 16384) with 16 references and 3 in flight.
+    let shape =
+        deckcraft_bitstream::budget::PictureShape { width: 16384, height: 16384, shift_x: 1, shift_y: 1, chroma_planes: 2, bytes_per_sample: 1 };
+    assert!(!shape.fits(16 + 1 + 3, deckcraft_bitstream::budget::DECODE_MEMORY_BUDGET));
+}

crates/hevc/src/decoder.rs

@@ -88,6 +88,8 @@ pub struct Decoder {
     pool: Option<rayon::ThreadPool>,
     in_flight: VecDeque<FrameRef>,
     max_in_flight: usize,
+    /// Bytes the picture store may take ([`deckcraft_bitstream::budget`]).
+    memory_budget: u64,
     draft: bool,
 }
 
@@ -142,6 +144,7 @@ impl Decoder {
             pool,
             in_flight: VecDeque::new(),
             max_in_flight: threads.max(1) + 2,
+            memory_budget: deckcraft_bitstream::budget::DECODE_MEMORY_BUDGET,
             draft: false,
         }
     }
@@ -372,7 +375,34 @@ impl Decoder {
         Ok(())
     }
 
+    /// Cap the bytes of decoded pictures (current, references, in flight) this decoder may hold.
+    pub fn set_memory_budget(&mut self, bytes: u64) {
+        self.memory_budget = bytes;
+    }
+
+    /// Refuse an SPS whose pictures would not fit the memory budget, before anything is allocated.
+    fn check_budget(&self, sps: &Sps) -> Result<()> {
+        let chroma = sps.chroma_array_type();
+        let shape = deckcraft_bitstream::budget::PictureShape {
+            width: u64::from(sps.width),
+            height: u64::from(sps.height),
+            shift_x: u32::from(sps.sub_width_c() == 2),
+            shift_y: u32::from(sps.sub_height_c() == 2),
+            chroma_planes: if chroma == 0 { 0 } else { 2 },
+            bytes_per_sample: 2,
+        };
+        let pictures = u64::from(sps.max_dec_pic_buffering) + 1 + self.max_in_flight as u64;
+        ensure!(
+            shape.fits(pictures, self.memory_budget),
+            "{}x{} pictures with {pictures} in the decoder exceed the memory budget",
+            sps.width,
+            sps.height
+        );
+        Ok(())
+    }
+
     fn start_picture(&mut self, sh: &SliceHeader, sps: &Arc<Sps>, pps: &Arc<Pps>, pts: i64) -> Result<()> {
+        self.check_budget(sps)?;
         ensure!(sh.first_slice_segment_in_pic && !sh.dependent, "invalid first slice segment");
         let hdr = sh.nal;
         let changed = match &self.active_sps {

crates/hevc/src/synth_tests.rs

@@ -618,3 +618,29 @@ fn slices_without_tiles() {
     ];
     run("slices", None, false, segs);
 }
+
+/// Decode with a memory budget of `budget` bytes; true if every picture came out without error.
+fn decodes_within(stream: &[u8], budget: u64) -> bool {
+    let mut dec = Decoder::with_threads(1);
+    dec.set_memory_budget(budget);
+    let mut pics = match dec.decode(stream, 0) {
+        Ok(p) => p,
+        Err(_) => return false,
+    };
+    pics.extend(dec.flush());
+    dec.take_error().is_none() && pics.len() == 5
+}
+
+#[test]
+fn picture_store_is_checked_against_the_memory_budget() {
+    let (stream, _) = sequence(&None, true, whole(WC * HC));
+    let sps_nal = sps(true);
+    let parsed = crate::params::Sps::parse(&deckcraft_bitstream::unescape_rbsp(&sps_nal[6..])).unwrap();
+    // 96x64 4:2:0 in 16-bit samples; the decoder holds the DPB, the current picture and up to
+    // threads + 2 in flight.
+    let one = (W * H + 2 * (W / 2) * (H / 2)) as u64 * 2;
+    let pictures = u64::from(parsed.max_dec_pic_buffering) + 1 + 3;
+    assert!(decodes_within(&stream, pictures * one));
+    assert!(!decodes_within(&stream, pictures * one - 1));
+    assert!(decodes_within(&stream, deckcraft_bitstream::budget::DECODE_MEMORY_BUDGET));
+}

crates/vp9/src/decoder.rs

@@ -73,6 +73,8 @@ pub struct Decoder {
     jobs: VecDeque<std::thread::JoinHandle<()>>,
     #[cfg(feature = "threads")]
     pool: Option<rayon::ThreadPool>,
+    /// Bytes the picture store may take ([`deckcraft_bitstream::budget`]).
+    memory_budget: u64,
 }
 
 impl Default for Decoder {
@@ -104,6 +106,11 @@ impl Decoder {
         Self::with_threads(default_threads())
     }
 
+    /// Cap the bytes of decoded pictures (references, current, in flight) this decoder may hold.
+    pub fn set_memory_budget(&mut self, bytes: u64) {
+        self.memory_budget = bytes;
+    }
+
     /// A decoder using up to `threads` worker threads (1 = decode on the calling thread, and
     /// every picture is returned by the `decode` call that completes it).
     pub fn with_threads(threads: usize) -> Self {
@@ -122,6 +129,7 @@ impl Decoder {
             last_show_frame: false,
             stats: DecodeStats::default(),
             threads: threads.max(1),
+            memory_budget: deckcraft_bitstream::budget::DECODE_MEMORY_BUDGET,
             bufs: Arc::default(),
             draft: false,
             pending: VecDeque::new(),
@@ -256,6 +264,13 @@ impl Decoder {
             return Ok(());
         }
         ensure!(h.width <= 16384 && h.height <= 16384, "frame size {}x{} too large", h.width, h.height);
+        // Eight reference slots, the new frame and those still decoding, on the 64-sample grid.
+        ensure!(
+            store_fits(h.width, h.height, &h.color, 8 + 1 + self.threads + 2, self.memory_budget),
+            "frame size {}x{} exceeds the decoder memory budget",
+            h.width,
+            h.height
+        );
         let (mi_rows, mi_cols) = (h.mi_rows as usize, h.mi_cols as usize);
         // compute_image_size semantics (7.2.6).
         let size = (h.width, h.height);
@@ -760,3 +775,41 @@ fn picture_of(f: &Frame, pts: i64, draft: bool) -> Picture {
     }
     b.finish()
 }
+
+/// Whether `pictures` frames of `width` x `height` (on the 64-sample superblock grid) fit in
+/// `budget` bytes.
+fn store_fits(width: u32, height: u32, color: &crate::header::ColorConfig, pictures: usize, budget: u64) -> bool {
+    let shape = deckcraft_bitstream::budget::PictureShape {
+        width: u64::from(width).div_ceil(64) * 64,
+        height: u64::from(height).div_ceil(64) * 64,
+        shift_x: u32::from(color.subsampling_x),
+        shift_y: u32::from(color.subsampling_y),
+        chroma_planes: 2,
+        bytes_per_sample: 2,
+    };
+    shape.fits(pictures as u64, budget)
+}
+
+#[cfg(test)]
+mod budget_tests {
+    use super::*;
+    use deckcraft_bitstream::budget::DECODE_MEMORY_BUDGET;
+
+    fn c420() -> crate::header::ColorConfig {
+        crate::header::ColorConfig { subsampling_x: true, subsampling_y: true, ..Default::default() }
+    }
+
+    #[test]
+    fn the_axis_maximum_is_over_budget_and_4k_is_not() {
+        assert!(!store_fits(16384, 16384, &c420(), 8 + 1 + 1 + 2, DECODE_MEMORY_BUDGET));
+        assert!(store_fits(3840, 2160, &c420(), 8 + 1 + 8 + 2, DECODE_MEMORY_BUDGET));
+    }
+
+    #[test]
+    fn exact_budget_boundary_and_reference_count() {
+        let one = (64 * 64 + 2 * 32 * 32) * 2;
+        assert!(store_fits(10, 10, &c420(), 11, 11 * one));
+        assert!(!store_fits(10, 10, &c420(), 11, 11 * one - 1));
+        assert!(!store_fits(10, 10, &c420(), 12, 11 * one));
+    }
+}