diff --git a/Cargo.lock b/Cargo.lock index 30ce9e1f7..5d7f5111a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3140,6 +3140,7 @@ version = "0.9.1" dependencies = [ "avif-serialize", "bytemuck", + "dhat", "image", "image-webp", "jixel", diff --git a/crates/codecs/Cargo.toml b/crates/codecs/Cargo.toml index 1ae016bc8..589d285fa 100644 --- a/crates/codecs/Cargo.toml +++ b/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 } diff --git a/crates/codecs/src/jxl.rs b/crates/codecs/src/jxl.rs index 26e144aed..adf59c0aa 100644 --- a/crates/codecs/src/jxl.rs +++ b/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 { + 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 { - 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 { .. }); diff --git a/crates/codecs/tests/jxl_streaming.rs b/crates/codecs/tests/jxl_streaming.rs new file mode 100644 index 000000000..298034656 --- /dev/null +++ b/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 { + 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, 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""); + 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()); + } + } +} diff --git a/crates/codecs/tests/jxl_tail_heap.rs b/crates/codecs/tests/jxl_tail_heap.rs new file mode 100644 index 000000000..068ebc952 --- /dev/null +++ b/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"); +} diff --git a/crates/codecs/tests/roundtrip.rs b/crates/codecs/tests/roundtrip.rs index af689ae5d..3391e3a91 100644 --- a/crates/codecs/tests/roundtrip.rs +++ b/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() {