Merge PR #1893: JPEG XL: check dimensions during loading & bound tracked allocations (#1517)

Cargo.lock

@@ -3140,6 +3140,7 @@ version = "0.9.1"
 dependencies = [
  "avif-serialize",
  "bytemuck",
+ "dhat",
  "image",
  "image-webp",
  "jixel",

crates/codecs/Cargo.toml

@@ -55,3 +55,6 @@ jxl-oxide = { workspace = true }
 proptest = { workspace = true }
 # Lossless synthetic HEIC files for the `heif` feature's tests.
 lightcraft-heif = { workspace = true, features = ["testdata"] }
+
+[target.'cfg(not(target_arch = "wasm32"))'.dev-dependencies]
+dhat = { workspace = true }

crates/codecs/src/jxl.rs

@@ -8,7 +8,7 @@ use crate::convert::{Buf, Meta, Model, Raw, check_size, finish};
 use crate::space::{NamedSpace, SourceSpace, SpaceOrigin, Trc};
 use crate::{DecodeOptions, Decoded, Error, Format, Result};
 use jxl_oxide::color::{ColourSpace, EnumColourEncoding, Primaries, RenderingIntent, TransferFunction, WhitePoint};
-use jxl_oxide::{AuxBoxData, JxlImage};
+use jxl_oxide::{AllocTracker, AuxBoxData, InitializeResult, JxlImage};
 
 const F: Format = Format::Jxl;
 
@@ -16,10 +16,60 @@ fn err(e: impl std::fmt::Display) -> Error {
     Error::Malformed(F, e.to_string())
 }
 
+/// Tracked allocation budget per allowed pixel for frames, extra channels and modular trees.
+/// This bounds allocations covered by `AllocTracker`, not every allocation made by the decoder (#1517).
+const TRACKED_BYTES_PER_PIXEL: u64 = 4 * 4 * 8;
+
+/// The allocation budget never drops below this, so small images keep the decoder's fixed
+/// overhead (group buffers, entropy tables) whatever `max_pixels` is.
+const MIN_TRACKED_BYTES: u64 = 64 << 20;
+
+/// Bytes handed to the decoder at a time (the size `JxlImageBuilder::read` reads in).
+const INPUT_CHUNK: usize = 4096;
+
+/// Initializes from bounded input, checks dimensions before rendering, then loads remaining frames
+/// within a tracked allocation budget.
+fn open(bytes: &[u8], opts: &DecodeOptions) -> Result<JxlImage> {
+    let budget = opts.max_pixels.saturating_mul(TRACKED_BYTES_PER_PIXEL).max(MIN_TRACKED_BYTES);
+    let tracker = AllocTracker::with_limit(usize::try_from(budget).unwrap_or(usize::MAX));
+    let mut uninit = JxlImage::builder().alloc_tracker(tracker).build_uninit();
+    // Initialization can load already-buffered frame data, so small feeds limit work before
+    // dimension checking.
+    let mut rest = bytes;
+    let mut window = INPUT_CHUNK;
+    let mut img = loop {
+        let consumed = uninit.feed_bytes(rest.get(..window).unwrap_or(rest)).map_err(err)?;
+        rest = rest.get(consumed..).unwrap_or_default();
+        match uninit.try_init().map_err(err)? {
+            InitializeResult::Initialized(img) => break img,
+            InitializeResult::NeedMoreData(more) => uninit = more,
+        }
+        if consumed == 0 {
+            // the parser wants more than the window holds at once
+            if window >= rest.len() {
+                return Err(err("truncated JPEG XL header"));
+            }
+            window = window.saturating_mul(2);
+        }
+    };
+    check_size(F, img.width() as u64, img.height() as u64, opts)?;
+    // Match `JxlImageBuilder::read`: bounded feeds until the final frame, not merely the first
+    // keyframe. `feed_bytes` copies codestream data into an untracked Vec, including any tail
+    // after the last frame; feeding the remaining file at once bypasses the allocation budget.
+    while !img.is_loading_done() && !rest.is_empty() {
+        let consumed = img.feed_bytes(rest.get(..INPUT_CHUNK).unwrap_or(rest)).map_err(err)?;
+        if consumed == 0 {
+            break;
+        }
+        rest = rest.get(consumed..).unwrap_or_default();
+    }
+    img.finalize().map_err(err)?;
+    Ok(img)
+}
+
 pub(crate) fn decode(bytes: &[u8], opts: &DecodeOptions) -> Result<Decoded> {
-    let mut img = JxlImage::builder().read(std::io::Cursor::new(bytes)).map_err(err)?;
+    let mut img = open(bytes, opts)?;
     let (w, h) = (img.width(), img.height());
-    check_size(F, w as u64, h as u64, opts)?;
     let header = img.image_header();
     let bits = header.metadata.bit_depth.bits_per_sample() as u8;
     let float = matches!(header.metadata.bit_depth, jxl_oxide::image::BitDepth::FloatSample { .. });

crates/codecs/tests/jxl_streaming.rs

@@ -0,0 +1,98 @@
+#![cfg(all(feature = "jxl", not(target_arch = "wasm32")))]
+
+use lightcraft_codecs::exif::minimal_exif;
+use lightcraft_codecs::*;
+use lightcraft_raster::Rgba8;
+
+fn codestream(bytes: &[u8]) -> Vec<u8> {
+    if bytes.starts_with(b"\xff\x0a") {
+        return bytes.to_vec();
+    }
+    let mut raw = Vec::new();
+    let mut at = 12;
+    while at + 8 <= bytes.len() {
+        let len = u32::from_be_bytes(bytes[at..at + 4].try_into().unwrap()) as usize;
+        let end = if len == 0 { bytes.len() } else { at + len };
+        match &bytes[at + 4..at + 8] {
+            b"jxlc" => raw.extend_from_slice(&bytes[at + 8..end]),
+            b"jxlp" => raw.extend_from_slice(&bytes[at + 12..end]),
+            _ => {}
+        }
+        at = end;
+    }
+    assert!(raw.starts_with(b"\xff\x0a"));
+    raw
+}
+
+fn box_bytes(out: &mut Vec<u8>, kind: &[u8; 4], data: &[u8]) {
+    out.extend_from_slice(&u32::try_from(data.len() + 8).unwrap().to_be_bytes());
+    out.extend_from_slice(kind);
+    out.extend_from_slice(data);
+}
+
+#[test]
+fn metadata_in_final_feed_matches_builder() {
+    let pixels = Rgba8::from_fn(8, 8, |x, y| [(x * 20) as u8, (y * 20) as u8, 80, 255]);
+    let encoded = encode_jxl(&EncodeImage::rgba8(&pixels), &JxlOptions { lossless: true, ..Default::default() }, &EncodeMeta::default()).unwrap();
+    let raw = codestream(&encoded);
+    let exif = minimal_exif(1);
+    let mut exif_box = vec![0; 4]; // TIFF header offset, relative to the end of this field
+    exif_box.extend_from_slice(&exif);
+    // Split the codestream across two feeds, with Exif and XMP following its last frame.
+    let mut bytes = b"\0\0\0\x0cJXL \x0d\x0a\x87\x0a".to_vec();
+    box_bytes(&mut bytes, b"ftyp", b"jxl \0\0\0\0jxl ");
+    let split = raw.len() / 2;
+    let mut part = 0u32.to_be_bytes().to_vec();
+    part.extend_from_slice(&raw[..split]);
+    box_bytes(&mut bytes, b"jxlp", &part);
+    let padding = vec![0; 4096 - bytes.len() - 8];
+    box_bytes(&mut bytes, b"free", &padding);
+    let mut part = 0x8000_0001u32.to_be_bytes().to_vec();
+    part.extend_from_slice(&raw[split..]);
+    box_bytes(&mut bytes, b"jxlp", &part);
+    box_bytes(&mut bytes, b"Exif", &exif_box);
+    box_bytes(&mut bytes, b"xml ", b"<xmp/>");
+    assert!(bytes.len() < 8192, "metadata must fit in the completing input chunk");
+    let old = jxl_oxide::JxlImage::builder().read(bytes.as_slice()).unwrap();
+    let old_exif = match old.aux_boxes().first_exif().unwrap() {
+        jxl_oxide::AuxBoxData::Data(ex) => ex.payload().to_vec(),
+        _ => panic!("builder must preserve Exif in its final feed"),
+    };
+    let old_xml = match old.aux_boxes().first_xml() {
+        jxl_oxide::AuxBoxData::Data(xml) => xml.to_vec(),
+        _ => panic!("builder must preserve XMP in its final feed"),
+    };
+    let decoded = decode(&bytes, DecodeOptions::default()).unwrap();
+    assert_eq!(decoded.exif.as_deref(), Some(old_exif.as_slice()));
+    assert_eq!(decoded.xmp.as_deref(), Some(std::str::from_utf8(&old_xml).unwrap()));
+    assert_eq!(decoded.to_srgb8(), pixels);
+}
+
+#[test]
+fn progressive_truncation_matches_builder() {
+    let pixels = Rgba8::from_fn(96, 64, |x, y| {
+        let mut seed = (y * 96 + x + 1) as u32;
+        let mut rgb = [0; 4];
+        for channel in &mut rgb[..3] {
+            seed ^= seed << 13;
+            seed ^= seed >> 17;
+            seed ^= seed << 5;
+            *channel = seed as u8;
+        }
+        rgb[3] = 255;
+        rgb
+    });
+    for lossless in [false, true] {
+        let encoded =
+            encode_jxl(&EncodeImage::rgba8(&pixels), &JxlOptions { lossless, progressive: true, ..Default::default() }, &EncodeMeta::default())
+                .unwrap();
+        let raw = codestream(&encoded);
+        assert!(raw.len() > 4096, "fixture must exercise post-header feeds");
+        for cut in [0, 1, 2, 12, 4095, 4096, 4097, raw.len() / 2, raw.len() - 1, raw.len()] {
+            let bytes = &raw[..cut];
+            // `finalize` permits partial images: rendering determines whether frame 0 is usable.
+            let old_ok = jxl_oxide::JxlImage::builder().read(bytes).is_ok_and(|image| image.render_frame(0).is_ok());
+            assert_eq!(decode(bytes, DecodeOptions::default()).is_ok(), old_ok, "lossless={lossless}, cut={cut}/{}", raw.len());
+        }
+    }
+}

crates/codecs/tests/jxl_tail_heap.rs

@@ -0,0 +1,43 @@
+#![cfg(all(feature = "jxl", not(target_arch = "wasm32")))]
+
+use lightcraft_codecs::*;
+use lightcraft_raster::Rgba8;
+
+#[global_allocator]
+static ALLOC: dhat::Alloc = dhat::Alloc;
+
+fn codestream(bytes: &[u8]) -> Vec<u8> {
+    if bytes.starts_with(b"\xff\x0a") {
+        return bytes.to_vec();
+    }
+    let mut raw = Vec::new();
+    let mut at = 12;
+    while at + 8 <= bytes.len() {
+        let len = u32::from_be_bytes(bytes[at..at + 4].try_into().unwrap()) as usize;
+        let end = if len == 0 { bytes.len() } else { at + len };
+        match &bytes[at + 4..at + 8] {
+            b"jxlc" => raw.extend_from_slice(&bytes[at + 8..end]),
+            b"jxlp" => raw.extend_from_slice(&bytes[at + 12..end]),
+            _ => {}
+        }
+        at = end;
+    }
+    assert!(raw.starts_with(b"\xff\x0a"));
+    raw
+}
+
+#[test]
+fn tiny_jxl_tail_does_not_allocate_a_tail_sized_buffer() {
+    let image = Rgba8::from_fn(8, 8, |x, y| [(x * 20) as u8, (y * 20) as u8, 80, 255]);
+    let encoded = encode_jxl(&EncodeImage::rgba8(&image), &JxlOptions { lossless: true, ..Default::default() }, &EncodeMeta::default()).unwrap();
+    let mut raw = codestream(&encoded);
+    let opts = DecodeOptions { max_pixels: 64, ..Default::default() };
+    decode(&raw, opts.clone()).unwrap();
+    raw.resize(raw.len() + (1 << 20), 0);
+    let _profiler = dhat::Profiler::builder().testing().build();
+    let decoded = decode(&raw, opts).unwrap();
+    assert_eq!((decoded.width, decoded.height), (8, 8));
+    let peak = dhat::HeapStats::get().max_bytes;
+    eprintln!("peak heap bytes: {peak}");
+    assert!(peak < (128 << 10), "tiny image used {peak} heap bytes");
+}

crates/codecs/tests/roundtrip.rs

@@ -734,6 +734,21 @@ fn jxl_lossless_exact_with_metadata() {
     assert!(max_diff(&d.image, &want.image) < 1e-5, "16-bit lossless differs by {}", max_diff(&d.image, &want.image));
 }
 
+/// #1517: the pixel limit rejects an oversized image; selected truncations never panic.
+#[test]
+#[cfg(all(feature = "jxl", not(target_arch = "wasm32")))]
+fn jxl_size_limit_and_truncation() {
+    let b = encode_jxl(&EncodeImage::rgba8(&gradient(97, 61)), &JxlOptions { lossless: true, ..Default::default() }, &EncodeMeta::default()).unwrap();
+    let r = decode(&b, DecodeOptions { max_pixels: 97 * 61 - 1, ..Default::default() });
+    assert!(matches!(r, Err(Error::TooLarge(97, 61))), "{r:?}");
+    let d = decode(&b, DecodeOptions { max_pixels: 97 * 61, ..Default::default() }).unwrap();
+    assert_eq!((d.width, d.height), (97, 61));
+    for cut in [0, 1, 2, 5, 12, 40, b.len() / 2, b.len() - 1] {
+        let _ = decode(&b[..cut], DecodeOptions::default());
+    }
+    assert!(decode(&b[..12.min(b.len())], DecodeOptions::default()).is_err());
+}
+
 #[test]
 #[cfg(feature = "jxl")]
 fn jxl_lossy_quality_and_effort() {