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