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

@@ -401,11 +401,13 @@ fn parse_frame(marker: u8, seg: &[u8], limits: &Limits) -> Result<Frame, Error>
     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::<UpsampleStep>() 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<UpsampleStep> {
     let last = samples.saturating_sub(1);
     (0..out)
         .map(|x| {

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