From f45f71cb483ba53cfbb2080f10882900242a4128 Mon Sep 17 00:00:00 2001 From: "Dylan M. Taylor" Date: Tue, 10 Nov 2026 18:35:50 -0500 Subject: [PATCH] jpeg: count the upsampling maps against max_alloc (#540) render keeps a column map and a row map for each component while it writes the output, but the check before allocating counted only the coefficients, the reconstructed planes and the output. The maps are now counted too, at their real element size. --- crates/jpeg/src/lib.rs | 11 ++++++++--- crates/jpeg/src/tests.rs | 7 ++++--- 2 files changed, 12 insertions(+), 6 deletions(-) diff --git a/crates/jpeg/src/lib.rs b/crates/jpeg/src/lib.rs index d9ef3bd3..1de2b17a 100644 --- a/crates/jpeg/src/lib.rs +++ b/crates/jpeg/src/lib.rs @@ -401,11 +401,13 @@ fn parse_frame(marker: u8, seg: &[u8], limits: &Limits) -> Result let mcux = width.div_ceil(8 * hmax); let mcuy = height.div_ceil(8 * vmax); // `render` holds the coefficients (4 bytes) and the reconstructed plane (4 bytes) of every - // padded sample while it writes the output, so all three count against `max_alloc` together. + // padded sample, and a column and a row map for upsampling each component, while it writes + // the output, so all of them count against `max_alloc` together. let mut total = (width as u64 * height as u64).saturating_mul(bytes_per_pixel); + let maps = (width as u64 + height as u64).saturating_mul(size_of::() as u64); for &(_, h, v, _) in &raw { let samples = (mcux * h * 8) as u64 * (mcuy * v * 8) as u64; - total = total.saturating_add(samples.saturating_mul(8)); + total = total.saturating_add(samples.saturating_mul(8)).saturating_add(maps); } // No allocation may pass `isize::MAX` bytes, which matters on 32-bit targets when a caller // raises `max_alloc`. @@ -1242,9 +1244,12 @@ fn reconstruct(comp: &Component, q: &[u16; 64], basis: &[[f32; 8]; 8], max: f32) plane } +/// An output coordinate's two neighboring component samples and the weight of the second. +type UpsampleStep = (usize, usize, f32); + /// For each output coordinate: the two neighboring component samples and the weight of the /// second (centered linear interpolation of a subsampled component). -fn upsample_map(out: usize, samples: usize, factor: usize, fmax: usize) -> Vec<(usize, usize, f32)> { +fn upsample_map(out: usize, samples: usize, factor: usize, fmax: usize) -> Vec { let last = samples.saturating_sub(1); (0..out) .map(|x| { diff --git a/crates/jpeg/src/tests.rs b/crates/jpeg/src/tests.rs index 3398b0e7..fcfb2e66 100644 --- a/crates/jpeg/src/tests.rs +++ b/crates/jpeg/src/tests.rs @@ -814,10 +814,11 @@ fn limits_are_checked_before_allocating() { segment(&mut s, 0xC9, &[8, 0xFF, 0xFF, 0xFF, 0xFF, 3, 1, 0x44, 0, 2, 0x11, 0, 3, 0x11, 0]); let small = Limits { max_alloc: 1 << 20, ..Limits::default() }; assert!(matches!(decode(&s, &small, None), Err(Error::LimitExceeded(_)))); - // 64×64 gray: 32 KiB of coefficients and plane plus 4 KiB of output, each under the limit - // alone but not together. + // 64×64 gray: 32 KiB of coefficients and plane, 4 KiB of output and 128 upsampling steps, + // each under the limit alone but not together. let stream = arith_stream(8, 8, 8, &coefficients(64, 1, 100), &[(0, 63, 0, 0)], 0); - for (max_alloc, ok) in [(36 << 10, true), ((36 << 10) - 1, false), (32 << 10, false)] { + let all = (36 << 10) + 128 * size_of::() as u64; + for (max_alloc, ok) in [(all, true), (all - 1, false), (36 << 10, false), (32 << 10, false)] { let limits = Limits { max_alloc, ..Limits::default() }; assert_eq!(decode(&stream, &limits, None).is_ok(), ok, "max_alloc {max_alloc}"); }