diff --git a/crates/av1/src/decoder.rs b/crates/av1/src/decoder.rs index 66c96698..88a29d99 100644 --- a/crates/av1/src/decoder.rs +++ b/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>, 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, /// 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)); + } +} diff --git a/crates/bitstream/src/budget.rs b/crates/bitstream/src/budget.rs new file mode 100644 index 00000000..cf98df2c --- /dev/null +++ b/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 { + 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 { + 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)); + } +} diff --git a/crates/bitstream/src/lib.rs b/crates/bitstream/src/lib.rs index 48884673..f25a3d33 100644 --- a/crates/bitstream/src/lib.rs +++ b/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. diff --git a/crates/h264/src/decoder.rs b/crates/h264/src/decoder.rs index 5f8941ea..ba105539 100644 --- a/crates/h264/src/decoder.rs +++ b/crates/h264/src/decoder.rs @@ -95,6 +95,8 @@ pub struct Decoder { pool: Option, in_flight: VecDeque, 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, pts: i64) -> Result<()> { + self.check_budget(sps)?; // Activate SPS; a resolution change flushes the DPB. let changed = match &self.active_sps { None => true, diff --git a/crates/h264/src/synth_tests.rs b/crates/h264/src/synth_tests.rs index b0e0d5d0..99d24303 100644 --- a/crates/h264/src/synth_tests.rs +++ b/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 { + let mut seed = 5u32; + let mbs: Vec>> = (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)); +} diff --git a/crates/hevc/src/decoder.rs b/crates/hevc/src/decoder.rs index cd5450fd..9038c3ad 100644 --- a/crates/hevc/src/decoder.rs +++ b/crates/hevc/src/decoder.rs @@ -88,6 +88,8 @@ pub struct Decoder { pool: Option, in_flight: VecDeque, 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, pps: &Arc, 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 { diff --git a/crates/hevc/src/synth_tests.rs b/crates/hevc/src/synth_tests.rs index 2d5b9fcf..dc225571 100644 --- a/crates/hevc/src/synth_tests.rs +++ b/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)); +} diff --git a/crates/vp9/src/decoder.rs b/crates/vp9/src/decoder.rs index bd9cd1a6..28e9d6d2 100644 --- a/crates/vp9/src/decoder.rs +++ b/crates/vp9/src/decoder.rs @@ -73,6 +73,8 @@ pub struct Decoder { jobs: VecDeque>, #[cfg(feature = "threads")] pool: Option, + /// 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)); + } +}