M2.6: image export pays for a JPEG mask, and reads a JPEG's size first

A JPEG that was itself a mask paid nothing, and images can share one: 300
images in 66 KB took 135 s with none refused, about 0.45 s for 220 bytes. It
now pays for its samples each time an image reads it.

Export priced a JPEG by its dictionary and decoded the size in its header,
comparing them only afterwards: an image of one pixel over a header of 500
megapixels cost 2 steps and 0.3 s. Both paths now read the header first and
leave a JPEG that disagrees undecoded, as Optimize and the CMYK path do.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01AmsdeUCdTaD72jkDw2pJKY

crates/create/README.md

@@ -32,8 +32,8 @@ other filters too, with the `/JBIG2Globals` of their own `/DecodeParms` entry, a
 whose JBIG2 page isn't its size (PdfCraft's renderer draws the whole page into the image; the page is refused before it is decoded),
 or whose globals can't be read or decode to more than 32 MiB (they are decoded before the page's size is known, so they have a
 limit of their own), is skipped, and such a mask is left out and the image exported opaque, as JPX and CCITT masks are.
-Past what one export may decode, which its images share, filters, decoders and pixels alike (`ImageBudget`; a JPEG only when a mask
-means decoding it), an image is skipped, and so is one whose mask is past it, which would otherwise come out showing what the mask hides. An export also stops at 2 GiB of files held, or twice the document if that
+Past what one export may decode, which its images share, filters, decoders and pixels alike (`ImageBudget`; a JPEG when it has to be
+decoded, as a mask or to apply one, and not when it is written as it is), an image is skipped, and so is one whose mask is past it, which would otherwise come out showing what the mask hides. An export also stops at 2 GiB of files held, or twice the document if that
 is more, since they stay in memory until it ends.
 Indexed Decode values select palette entries after rounding and clipping. Color-key masks compare the original samples.
 Missing, null or malformed Decode arrays keep the default mapping.

crates/create/src/extract.rs

@@ -285,6 +285,32 @@ fn size(w: i64, h: i64) -> Option<(usize, usize)> {
     (w > 0 && h > 0 && pixels <= MAX_PIXELS).then_some((w, h))
 }
 
+#[cfg(test)]
+thread_local! {
+    /// JPEGs decoded whole, for the tests that keep one from being decoded for nothing.
+    pub(crate) static JPEGS_DECODED: std::cell::Cell<usize> = const { std::cell::Cell::new(0) };
+}
+
+/// The size a JPEG's header gives, which is what decoding it goes by, whatever its dictionary
+/// says: an image of one pixel can hold a header of 500 megapixels.
+fn jpeg_size(data: &[u8]) -> Option<(usize, usize)> {
+    use zune_jpeg::zune_core::{bytestream::ZCursor, options::DecoderOptions};
+    let options = DecoderOptions::default().set_strict_mode(false).set_max_width(usize::from(u16::MAX)).set_max_height(usize::from(u16::MAX));
+    let mut dec = zune_jpeg::JpegDecoder::new_with_options(ZCursor::new(data), options);
+    dec.decode_headers().ok()?;
+    dec.dimensions()
+}
+
+/// JPEG `data`, decoded, if its header says it is `w` × `h`.
+fn jpeg_of_size(data: &[u8], w: usize, h: usize) -> Option<image::DynamicImage> {
+    if jpeg_size(data)? != (w, h) {
+        return None;
+    }
+    #[cfg(test)]
+    JPEGS_DECODED.with(|n| n.set(n.get() + 1));
+    image::load_from_memory_with_format(data, image::ImageFormat::Jpeg).ok()
+}
+
 /// A JPEG's samples (8 bits): grey for a grey image, CMYK as stored for a CMYK one, else RGB,
 /// with the space they are in.
 fn jpeg_samples(data: &[u8], space: &Space, w: usize, h: usize) -> Option<(Vec<u8>, Space)> {
@@ -293,10 +319,7 @@ fn jpeg_samples(data: &[u8], space: &Space, w: usize, h: usize) -> Option<(Vec<u
         Space::Cmyk => return cmyk_jpeg_samples(data, w, h).map(|px| (px, Space::Cmyk)),
         _ => {}
     }
-    let img = image::load_from_memory_with_format(data, image::ImageFormat::Jpeg).ok()?;
-    if img.width() as usize != w || img.height() as usize != h {
-        return None;
-    }
+    let img = jpeg_of_size(data, w, h)?;
     Some(match space {
         Space::Gray => (img.into_luma8().into_raw(), Space::Gray),
         _ => (img.into_rgb8().into_raw(), Space::Rgb),
@@ -521,8 +544,11 @@ fn gray_plane(doc: &Document, m: &Stream, w: usize, h: usize, budget: &mut Image
     let f = filters(m);
     let mut px = match f.last().map(Vec::as_slice) {
         Some(b"DCTDecode" | b"DCT") if f.len() == 1 => {
-            let img = image::load_from_memory_with_format(&m.raw, image::ImageFormat::Jpeg).ok()?.into_luma8();
-            (img.width() as usize == mw && img.height() as usize == mh).then(|| img.into_raw())?
+            // Past the allowance the image goes too: `image` sees this one refused.
+            if jpeg_size(&m.raw)? != (mw, mh) || !budget.pay_jpeg((mw as u64).saturating_mul(mh as u64)) {
+                return None;
+            }
+            jpeg_of_size(&m.raw, mw, mh)?.into_luma8().into_raw()
         }
         Some(b"DCTDecode" | b"DCT" | b"JPXDecode") => return None,
         _ => {

crates/create/src/tests.rs

@@ -2357,6 +2357,84 @@ fn a_jpeg_that_has_to_be_decoded_shares_the_allowance() {
     assert_eq!(budget.refused(), 1);
 }
 
+/// A JPEG that is itself a mask pays for its samples, each time an image reads it: images can
+/// share one, and only the images are ever seen once.
+#[test]
+fn a_jpeg_mask_pays_for_its_samples() {
+    use pdfcraft_cos::{Dict, Document, ImageBudget, Object, Stream};
+    let mut jpeg = Vec::new();
+    image::GrayImage::from_pixel(4, 4, image::Luma([90])).write_to(&mut std::io::Cursor::new(&mut jpeg), image::ImageFormat::Jpeg).unwrap();
+    let mut mask = Dict::new();
+    mask.set(b"Subtype".to_vec(), Object::name("Image"));
+    mask.set(b"Width".to_vec(), Object::Int(4));
+    mask.set(b"Height".to_vec(), Object::Int(4));
+    mask.set(b"BitsPerComponent".to_vec(), Object::Int(8));
+    mask.set(b"ColorSpace".to_vec(), Object::name("DeviceGray"));
+    mask.set(b"Filter".to_vec(), Object::name("DCTDecode"));
+    let mut doc = Document::new_empty();
+    let mask = doc.add(Object::Stream(Stream::from_raw(mask, jpeg)));
+    let mut image = Dict::new();
+    image.set(b"Subtype".to_vec(), Object::name("Image"));
+    image.set(b"Width".to_vec(), Object::Int(1));
+    image.set(b"Height".to_vec(), Object::Int(1));
+    image.set(b"BitsPerComponent".to_vec(), Object::Int(8));
+    image.set(b"ColorSpace".to_vec(), Object::name("DeviceGray"));
+    image.set(b"SMask".to_vec(), Object::Ref(mask));
+    let image = doc.add(Object::Stream(Stream::from_raw(image, vec![0x80])));
+    let mut budget = ImageBudget::with_work(u64::MAX);
+    assert!(crate::extract::image_file_within(&doc, image, &mut budget).is_ok());
+    assert_eq!(u64::MAX - budget.work_left(), 2 * 4 * 4 + 2, "its samples twice, and the image's byte and pixel");
+    // Just short of what the mask costs: it isn't decoded, and takes the image with it.
+    let mut budget = ImageBudget::with_work(2 * 4 * 4 - 1);
+    let refused = crate::extract::image_file_within(&doc, image, &mut budget).unwrap_err();
+    assert!(refused.contains("more image decoding than"), "{refused}");
+    assert_eq!((budget.work_left(), budget.refused()), (0, 1));
+}
+
+/// What decoding a JPEG takes goes by its header, so that is what has to agree with the
+/// dictionary before it is decoded: one image of a pixel can hold a page of 500 megapixels. The
+/// file comes out as it is either way, so what shows is whether it was decoded.
+#[test]
+fn a_jpeg_larger_than_its_dictionary_says_isnt_decoded() {
+    use pdfcraft_cos::{Dict, Document, ImageBudget, Object, Stream};
+    let mut jpeg = Vec::new();
+    image::RgbImage::from_pixel(4, 4, image::Rgb([9, 9, 9])).write_to(&mut std::io::Cursor::new(&mut jpeg), image::ImageFormat::Jpeg).unwrap();
+    // A colour key means decoding it, to find the pixels it hides.
+    let keyed = |side: i64| {
+        let mut d = Dict::new();
+        d.set(b"Subtype".to_vec(), Object::name("Image"));
+        d.set(b"Width".to_vec(), Object::Int(side));
+        d.set(b"Height".to_vec(), Object::Int(side));
+        d.set(b"BitsPerComponent".to_vec(), Object::Int(8));
+        d.set(b"ColorSpace".to_vec(), Object::name("DeviceRGB"));
+        d.set(b"Filter".to_vec(), Object::name("DCTDecode"));
+        d.set(b"Mask".to_vec(), Object::Array(vec![Object::Int(0); 6]));
+        Stream::from_raw(d, jpeg.clone())
+    };
+    let decodes = |doc: &Document, image: pdfcraft_cos::ObjRef| {
+        crate::extract::JPEGS_DECODED.with(|n| n.set(0));
+        assert!(crate::extract::image_file_within(doc, image, &mut ImageBudget::with_work(u64::MAX)).is_ok());
+        crate::extract::JPEGS_DECODED.with(std::cell::Cell::get)
+    };
+    let mut doc = Document::new_empty();
+    let (honest, lying) = (doc.add(Object::Stream(keyed(4))), doc.add(Object::Stream(keyed(1))));
+    assert_eq!(decodes(&doc, honest), 1);
+    assert_eq!(decodes(&doc, lying), 0);
+    // The same as a mask: read as none, without being decoded.
+    let mut mask = keyed(1);
+    mask.dict.remove(b"Mask");
+    let mask = doc.add(Object::Stream(mask));
+    let mut image = Dict::new();
+    image.set(b"Subtype".to_vec(), Object::name("Image"));
+    image.set(b"Width".to_vec(), Object::Int(1));
+    image.set(b"Height".to_vec(), Object::Int(1));
+    image.set(b"BitsPerComponent".to_vec(), Object::Int(8));
+    image.set(b"ColorSpace".to_vec(), Object::name("DeviceGray"));
+    image.set(b"SMask".to_vec(), Object::Ref(mask));
+    let image = doc.add(Object::Stream(Stream::from_raw(image, vec![0x80])));
+    assert_eq!(decodes(&doc, image), 0);
+}
+
 /// A mask past the allowance can't be left out as an unreadable one is, or the image comes out
 /// opaque, showing what the mask hides: the image goes, with the reason.
 #[test]