Bound JPEG XL input feeds through the final frame (#1517)

Retain the pre-frame size check and allocation budget while avoiding untracked input-sized buffers. Add the reported heap regression and container/progressive compatibility checks.

Co-authored-by: Codex <noreply@openai.com>

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

@@ -25,8 +25,8 @@ const TRACKED_BYTES_PER_PIXEL: u64 = 4 * 4 * 8;
 /// overhead (group buffers, entropy tables) whatever `max_pixels` is.
 const MIN_TRACKED_BYTES: u64 = 64 << 20;
 
-/// Bytes handed to the header stage at a time (the size `JxlImageBuilder::read` reads in).
-const HEADER_CHUNK: usize = 4096;
+/// Bytes handed to the decoder at a time (the size `JxlImageBuilder::read` reads in).
+const INPUT_CHUNK: usize = 4096;
 
 /// Parses the header, checks the image size against `opts`, and only then reads the frames, with
 /// an allocation budget, so a huge image is refused before any of its pixels are decoded
@@ -38,7 +38,7 @@ fn open(bytes: &[u8], opts: &DecodeOptions) -> Result<JxlImage> {
     // Small slices, as `read` does: whatever the header stage has buffered when it initializes is
     // parsed as frame data right away, so feeding the whole file here would skip the size check.
     let mut rest = bytes;
-    let mut window = HEADER_CHUNK;
+    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();
@@ -55,8 +55,11 @@ fn open(bytes: &[u8], opts: &DecodeOptions) -> Result<JxlImage> {
         }
     };
     check_size(F, img.width() as u64, img.height() as u64, opts)?;
-    while !rest.is_empty() {
-        let consumed = img.feed_bytes(rest).map_err(err)?;
+    // 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;
         }

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