raw: charge the lossless-JPEG intermediate against max_alloc (#2925) (#3194)

Co-authored-by: switch87 <8284204+switch87@users.noreply.github.com>

crates/raw/src/sensor.rs

@@ -225,6 +225,11 @@ pub(crate) fn read_plane(t: &Tiff, ifd: &Ifd, limits: &Limits, layout: JpegLayou
         let s = &segments[i];
         let rows = s.h.min(height - s.y);
         let src = t.bytes(s.offset, s.len).ok_or_else(|| RawError::malformed("raw data lies outside the file"))?;
+        // Every branch below sizes at least one buffer from the segment's own geometry, which
+        // the file declares independently of the image: a tile can be far larger than the image
+        // it belongs to, so the image check above does not bound those buffers. Charge the
+        // segment itself before decoding it (#2925).
+        limits.check(s.w as u64, s.h as u64, samples as u64 * 2)?;
         match compression {
             1 => unpack(src, s.w, rows, samples, bits, le),
             8 => {
@@ -280,10 +285,15 @@ pub(crate) fn read_plane(t: &Tiff, ifd: &Ifd, limits: &Limits, layout: JpegLayou
                 // The 8-bit samples widen to the document's 16-bit range (0xFF -> 0xFFFF).
                 Ok(decoded.into_iter().map(|b| u16::from(b) * 257).collect())
             }
+            // Compression 7 (lossless JPEG): the frame the entropy stream declares is capped
+            // by the budget, not only by the segment, and the `Quads` scatter buffer below
+            // comes out of the same budget. `max_alloc` counts bytes, a sample is two.
             _ => {
                 let need = s.w * rows * samples;
-                let max = s.w.saturating_mul(s.h).saturating_mul(samples).saturating_mul(4);
-                let (f, v) = ljpeg::decode(src, max)?;
+                let scatter = if layout == JpegLayout::Quads && samples == 1 { s.w.saturating_mul(s.h).saturating_mul(2) as u64 } else { 0 };
+                let budget = usize::try_from(limits.max_alloc.saturating_sub(scatter) / 2).unwrap_or(usize::MAX);
+                let geom = s.w.saturating_mul(s.h).saturating_mul(samples).saturating_mul(4);
+                let (f, v) = ljpeg::decode(src, geom.min(budget))?;
                 if v.len() < need {
                     return Err(RawError::malformed("lossless JPEG tile is smaller than its tile"));
                 }

crates/raw/tests/hostile.rs

@@ -150,3 +150,66 @@ fn absurd_dimensions_are_rejected_without_allocating() {
     assert!(patched > 0);
     assert!(matches!(decode(&b, &Limits::default()), Err(RawError::LimitExceeded(_))));
 }
+
+/// An 8×8 CFA image stored as one lossless-JPEG tile that the file declares 512×512: the
+/// reader takes the segment's geometry from the tile tags, not from the image (#2925).
+fn oversized_lj92_tile() -> Vec<u8> {
+    let cfa = mosaic(&scene(8, 8), 8, [0, 1, 1, 2], 64, 4000);
+    let mut d = DngSpec::cfa(8, 8, cfa);
+    d.bits = 12;
+    d.white = 4000;
+    d.black = vec![64];
+    d.storage = DngStorage::Lj92Tiles { width: 512, height: 512 };
+    d.build()
+}
+
+/// Rewrites one inline LONG IFD value; `false` when the tag is not found exactly once.
+fn patch_long(b: &mut [u8], tag: u16, to: u32) -> bool {
+    let mut patched = 0;
+    for i in 0..b.len().saturating_sub(12) {
+        if b[i..i + 2] == tag.to_le_bytes() && b[i + 2..i + 4] == 4u16.to_le_bytes() && b[i + 4..i + 8] == 1u32.to_le_bytes() {
+            b[i + 8..i + 12].copy_from_slice(&to.to_le_bytes());
+            patched += 1;
+        }
+    }
+    patched == 1
+}
+
+/// The lossless-JPEG intermediate is sized from the segment, so the plane check alone
+/// does not bound it: it has to be charged against `max_alloc` too (#2925).
+#[test]
+fn lossless_jpeg_intermediate_is_charged_against_max_alloc() {
+    let b = oversized_lj92_tile();
+    assert!(decode(&b, &Limits::default()).is_ok(), "the file itself is valid");
+    // The image is 8 * 8 * 2 = 128 bytes, so the plane check passes, while the declared
+    // tile's intermediate is 512 * 512 * 2 bytes — far past a 64 KiB budget.
+    let budget = Limits { max_alloc: 64 << 10, ..Limits::default() };
+    assert!(matches!(decode(&b, &budget), Err(RawError::LimitExceeded(_))));
+    // A budget that does hold the intermediate still decodes it unchanged.
+    let roomy = Limits { max_alloc: 4 << 20, ..Limits::default() };
+    assert!(decode(&b, &roomy).is_ok(), "a budget that holds the intermediate still decodes it");
+    // Ordinary geometry stays inside the same budget: the segment check is not stricter
+    // than the image it belongs to.
+    let mut normal = DngSpec::cfa(8, 8, mosaic(&scene(8, 8), 8, [0, 1, 1, 2], 64, 4000));
+    normal.bits = 12;
+    normal.white = 4000;
+    normal.black = vec![64];
+    normal.storage = DngStorage::Lj92Tiles { width: 4, height: 4 };
+    let tight = Limits { max_alloc: 4096, ..Limits::default() };
+    assert!(decode(&normal.build(), &tight).is_ok(), "ordinary geometry stays inside the budget");
+}
+
+/// A frame may declare up to 4× the segment's samples (the headroom the reader passes);
+/// when that does not fit the budget either, the decode is refused rather than allocated.
+#[test]
+fn lossless_jpeg_frame_larger_than_the_budget_is_refused() {
+    let mut b = oversized_lj92_tile();
+    // Shrink the declared tile to 256×256: its geometry then allows exactly 4× its own
+    // samples — the 262144-sample frame this file actually holds (322 = TileWidth,
+    // 323 = TileLength).
+    assert!(patch_long(&mut b, 322, 256), "TileWidth patched once");
+    assert!(patch_long(&mut b, 323, 256), "TileLength patched once");
+    // Holds the image (128 B) and the tile (128 KiB), but not the frame's 512 KiB.
+    let budget = Limits { max_alloc: 256 << 10, ..Limits::default() };
+    assert!(matches!(decode(&b, &budget), Err(RawError::LimitExceeded(_))));
+}