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}"); }