fix(create): image export checks a JPEG's Extended XMP chunks before decoding it (#5860)

crates/create/src/extract.rs

@@ -302,9 +302,14 @@ thread_local! {
 }
 
 /// 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.
+/// says: an image of one pixel can hold a header of 500 megapixels. `None`, as for a JPEG that
+/// can't be read, when its metadata would take the decoder quadratic time
+/// ([`decodable_metadata`](crate::upright_image::decodable_metadata)).
 fn jpeg_size(data: &[u8]) -> Option<(usize, usize)> {
     use zune_jpeg::zune_core::{bytestream::ZCursor, options::DecoderOptions};
+    if !crate::upright_image::decodable_metadata(data) {
+        return None;
+    }
     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()?;
@@ -329,6 +334,15 @@ 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)),
         _ => {}
     }
+    // Past the decoder's metadata bound, drop the Extended XMP chunks rather than the mask:
+    // decoded as it stands, an over-chunked JPEG is written as is, without its /SMask.
+    let stripped;
+    let data = if crate::upright_image::decodable_metadata(data) {
+        data
+    } else {
+        stripped = crate::upright_image::without_extended_xmp(data);
+        &stripped[..]
+    };
     let img = jpeg_of_size(data, w, h)?;
     Some(match space {
         Space::Gray => (img.into_luma8().into_raw(), Space::Gray),
@@ -340,6 +354,15 @@ fn jpeg_samples(data: &[u8], space: &Space, w: usize, h: usize) -> Option<(Vec<u
 /// a decoder converting to RGB inverts them, with YCCK converted to CMYK.
 fn cmyk_jpeg_samples(data: &[u8], w: usize, h: usize) -> Option<Vec<u8>> {
     use zune_jpeg::zune_core::{bytestream::ZCursor, colorspace::ColorSpace, options::DecoderOptions};
+    // Past the decoder's metadata bound, decode the image without its Extended XMP chunks
+    // rather than not at all: written out as stored, a CMYK JPEG shows inverted colours.
+    let stripped;
+    let data = if crate::upright_image::decodable_metadata(data) {
+        data
+    } else {
+        stripped = crate::upright_image::without_extended_xmp(data);
+        &stripped[..]
+    };
     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()?;

crates/create/src/tests.rs

@@ -941,6 +941,43 @@ fn damaged_masks_and_sizes_never_panic() {
     }
 }
 
+#[test]
+fn a_jpeg_with_thousands_of_extended_xmp_chunks_keeps_its_mask() {
+    // Export decodes a JPEG with a soft mask, or a JPEG soft mask, to make a PNG with alpha. A JPEG
+    // whose Extended XMP chunks would take the decoder quadratic time is decoded without them,
+    // which is linear in what is left, so the mask survives; a JPEG mask past the same bound is
+    // still left out, without reaching the decoder.
+    let flood = crate::upright_image::tests::with_extended_xmp_chunks(&real_jpeg(), 5_000);
+    let mut doc = from_images(&[("flat.png".into(), png_bytes(false))]).unwrap();
+    with_image(&mut doc, |doc| {
+        let mask = doc.add(Object::Stream(Stream::from_raw(image_dict(4, 2, "DeviceGray", 8), vec![0; 8])));
+        let mut d = image_dict(4, 2, "DeviceGray", 8);
+        d.set(b"Filter".to_vec(), Object::name("DCTDecode"));
+        d.set(b"SMask".to_vec(), Object::Ref(mask));
+        Stream::from_raw(d, flood.clone())
+    });
+    crate::extract::JPEGS_DECODED.with(|n| n.set(0));
+    let out = extract_images(&doc, &[0], 0);
+    assert_eq!(out.images[0].extension, "png");
+    assert_eq!(crate::extract::JPEGS_DECODED.with(std::cell::Cell::get), 1);
+    // The chunks themselves are gone, so the decoder never saw them.
+    assert!(!out.images[0].data.windows(5).any(|w| w == b"http:"));
+
+    crate::extract::JPEGS_DECODED.with(|n| n.set(0));
+
+    let mut doc = from_images(&[("flat.png".into(), png_bytes(false))]).unwrap();
+    with_image(&mut doc, |doc| {
+        let mut md = image_dict(4, 2, "DeviceGray", 8);
+        md.set(b"Filter".to_vec(), Object::name("DCTDecode"));
+        let mask = doc.add(Object::Stream(Stream::from_raw(md, flood.clone())));
+        let mut d = image_dict(4, 2, "DeviceRGB", 8);
+        d.set(b"SMask".to_vec(), Object::Ref(mask));
+        Stream::from_raw(d, vec![10; 24])
+    });
+    assert_eq!(alphas(&exported_rgba(&doc).1), [255; 8]);
+    assert_eq!(crate::extract::JPEGS_DECODED.with(std::cell::Cell::get), 0);
+}
+
 #[test]
 fn indexed_image_export_clamps_samples_to_hival() {
     for palette in [vec![255, 0, 0, 0, 0, 255], vec![255, 0, 0, 0, 0, 255, 0, 255, 0, 255, 255, 255]] {
@@ -2788,3 +2825,30 @@ fn proportions_up_to_the_limit_scale_as_they_did() {
     let mask = noise(20, 10);
     assert_eq!(crate::extract::scaled(&mask, 20, 10), mask);
 }
+
+/// A CMYK JPEG past the decoder's Extended XMP bound is still decoded (its chunks left out) and
+/// exported with the colours the page shows, not written as stored, which inverts them.
+#[test]
+fn a_cmyk_jpeg_flooded_with_extended_xmp_is_exported_in_its_colours() {
+    let jpeg = crate::upright_image::tests::with_extended_xmp_chunks(&flat_four_component_jpeg([255, 0, 0, 0], 0), 5_000);
+    let mut doc = from_images(&[("flat.png".into(), png_bytes(false))]).unwrap();
+    with_image(&mut doc, |doc| {
+        let mut d = image_dict(8, 8, "DeviceCMYK", 8);
+        d.set(b"Filter".to_vec(), Object::name("DCTDecode"));
+        let m = doc.add(Object::Stream(Stream::from_raw(image_dict(8, 8, "DeviceGray", 8), vec![255; 64])));
+        d.set(b"SMask".to_vec(), Object::Ref(m));
+        Stream::from_raw(d, jpeg.clone())
+    });
+    let (_, rgba) = exported_rgba(&doc);
+    assert!(rgba.as_chunks::<4>().0.iter().all(|p| p[0] <= 6 && p[1] >= 249 && p[2] >= 249), "{:?}", &rgba[..4]);
+}
+
+/// Stripping leaves everything but the Extended XMP chunks, byte for byte.
+#[test]
+fn stripping_extended_xmp_leaves_the_rest_of_the_jpeg() {
+    let plain = flat_four_component_jpeg([255, 0, 0, 0], 0);
+    let flooded = crate::upright_image::tests::with_extended_xmp_chunks(&plain, 3);
+    assert_ne!(flooded, plain);
+    // The helper puts a stray byte before the chunks, which the decoder skips: it stays.
+    assert_eq!(crate::upright_image::without_extended_xmp(&flooded), [&plain[..2], b"\0".as_slice(), &plain[2..]].concat());
+}

crates/create/src/upright_image.rs

@@ -30,7 +30,7 @@ pub(crate) fn upright(bytes: &[u8]) -> Cow<'_, [u8]> {
 }
 
 fn turned(bytes: &[u8]) -> Option<Vec<u8>> {
-    if !bytes.starts_with(&[0xFF, 0xD8]) || extended_xmp_chunks(bytes) > MAX_EXTENDED_XMP_CHUNKS {
+    if !bytes.starts_with(&[0xFF, 0xD8]) || !decodable_metadata(bytes) {
         return None;
     }
     #[cfg(test)]
@@ -65,10 +65,41 @@ fn turned(bytes: &[u8]) -> Option<Vec<u8>> {
 /// decoder re-sorts and rescans every chunk so far after each marker, quadratic in their number.
 const MAX_EXTENDED_XMP_CHUNKS: usize = 1024;
 
+/// Whether JPEG `bytes` may be handed to a decoder (`image`'s or `zune-jpeg`'s, which `image` uses),
+/// headers included: not when it carries more than [`MAX_EXTENDED_XMP_CHUNKS`]. Every path that
+/// decodes a JPEG from a file checks this first: image import here, and image export.
+pub(crate) fn decodable_metadata(bytes: &[u8]) -> bool {
+    extended_xmp_chunks(bytes) <= MAX_EXTENDED_XMP_CHUNKS
+}
+
 /// How many Extended XMP chunks (APP1 segments) come before the image data.
 fn extended_xmp_chunks(bytes: &[u8]) -> usize {
+    let mut chunks = 0usize;
+    each_extended_xmp(bytes, |_| chunks += 1);
+    chunks
+}
+
+/// JPEG `bytes` without their Extended XMP chunks, which nothing here reads: what is left decodes
+/// in time linear in its size however many chunks the file carried.
+pub(crate) fn without_extended_xmp(bytes: &[u8]) -> Vec<u8> {
+    // Only the bytes that are kept: a CMYK JPEG can be mostly Extended XMP, and reserving the
+    // whole file for the little that remains aborts the export under a tight address-space limit.
+    let mut removed = 0usize;
+    each_extended_xmp(bytes, |r| removed = removed.saturating_add(r.len()));
+    let mut out = Vec::with_capacity(bytes.len().saturating_sub(removed));
+    let mut at = 0;
+    each_extended_xmp(bytes, |r| {
+        out.extend_from_slice(bytes.get(at..r.start).unwrap_or_default());
+        at = r.end;
+    });
+    out.extend_from_slice(bytes.get(at..).unwrap_or_default());
+    out
+}
+
+/// Calls `f` with the byte range of each Extended XMP chunk (APP1 segment) before the image data.
+fn each_extended_xmp(bytes: &[u8], mut f: impl FnMut(std::ops::Range<usize>)) {
     const PREFIX: &[u8] = b"http://ns.adobe.com/xmp/extension/\0";
-    let (mut i, mut chunks) = (2usize, 0usize);
+    let mut i = 2usize;
     while let Some(&byte) = bytes.get(i) {
         // Stray bytes between segments, which the decoder skips too.
         if byte != 0xFF {
@@ -86,14 +117,16 @@ fn extended_xmp_chunks(bytes: &[u8]) -> usize {
             _ => {
                 let Some(&[hi, lo]) = bytes.get(i.saturating_add(2)..i.saturating_add(4)) else { break };
                 let len = usize::from(u16::from_be_bytes([hi, lo]));
-                if marker == 0xE1 && bytes.get(i.saturating_add(4)..).is_some_and(|seg| seg.starts_with(PREFIX)) {
-                    chunks += 1;
+                let end = i.saturating_add(2).saturating_add(len.max(2));
+                // The prefix must start the APP1 segment's own payload: a zero-length header
+                // followed by prefix bytes is not a chunk.
+                if marker == 0xE1 && bytes.get(i.saturating_add(4)..end.min(bytes.len())).is_some_and(|seg| seg.starts_with(PREFIX)) {
+                    f(i..end.min(bytes.len()));
                 }
-                i = i.saturating_add(2).saturating_add(len.max(2));
+                i = end;
             }
         }
     }
-    chunks
 }
 
 /// The resolution in a JFIF header, if the file has one in inches or centimetres.
@@ -120,7 +153,7 @@ fn jfif_density(bytes: &[u8]) -> Option<PixelDensity> {
 }
 
 #[cfg(test)]
-mod tests {
+pub(crate) mod tests {
     use super::*;
     use crate::tests::phone_jpeg;
 
@@ -138,22 +171,24 @@ mod tests {
         assert!(matches!(upright(b"\xff\xd8 not a jpeg"), Cow::Borrowed(_)));
     }
 
+    /// `jpeg` with `n` empty chunks of one never-complete Extended XMP packet after its SOI.
+    pub(crate) fn with_extended_xmp_chunks(jpeg: &[u8], n: usize) -> Vec<u8> {
+        let mut seg = b"http://ns.adobe.com/xmp/extension/\0".to_vec();
+        seg.extend_from_slice(&[b'0'; 32]); // the packet's GUID
+        seg.extend_from_slice(&1000u32.to_be_bytes()); // its full length
+        seg.extend_from_slice(&0u32.to_be_bytes()); // this chunk's offset
+        let mut chunk = vec![0xFF, 0xE1];
+        chunk.extend_from_slice(&u16::try_from(seg.len() + 2).unwrap().to_be_bytes());
+        chunk.extend_from_slice(&seg);
+        [&jpeg[..2], b"\0".as_slice(), &chunk.repeat(n), &jpeg[2..]].concat()
+    }
+
     #[test]
     fn a_jpeg_with_thousands_of_extended_xmp_chunks_is_not_decoded() {
         // Empty chunks of one never-complete extended packet: the decoder re-sorts and rescans every
         // chunk so far after each marker, quadratic in their number.
-        let chunk = |n: usize| {
-            let mut seg = b"http://ns.adobe.com/xmp/extension/\0".to_vec();
-            seg.extend_from_slice(&[b'0'; 32]); // the packet's GUID
-            seg.extend_from_slice(&1000u32.to_be_bytes()); // its full length
-            seg.extend_from_slice(&0u32.to_be_bytes()); // this chunk's offset
-            let mut out = vec![0xFF, 0xE1];
-            out.extend_from_slice(&u16::try_from(seg.len() + 2).unwrap().to_be_bytes());
-            out.extend_from_slice(&seg);
-            out.repeat(n)
-        };
         let photo = phone_jpeg(6, 144);
-        let with = |n: usize| [&photo[..2], b"\0".as_slice(), &chunk(n), &photo[2..]].concat();
+        let with = |n: usize| with_extended_xmp_chunks(&photo, n);
         DECODES.with(|c| c.set(0));
         assert!(matches!(upright(&with(4)), Cow::Owned(_)), "a few chunks are an ordinary photo");
         assert!(matches!(upright(&with(5_000)), Cow::Borrowed(_)));