fix(create): export JPEG images behind stream filters (#7168)

crates/create/src/extract.rs

@@ -145,13 +145,23 @@ fn decoded(doc: &Document, s: &Stream, budget: &mut ImageBudget) -> Result<Vec<u
 /// alpha, JPEG images included, so cut-out images keep their transparent background.
 fn image(doc: &Document, s: &Stream, budget: &mut ImageBudget) -> Result<(&'static str, Vec<u8>), String> {
     let f = filters(s);
-    let jpeg = matches!(f.last().map(Vec::as_slice), Some(b"DCTDecode" | b"DCT")) && f.len() == 1;
-    match f.last().map(Vec::as_slice) {
-        Some(b"JPXDecode") => return Err("JPEG 2000 images can't be exported yet".into()),
-        Some(b"DCTDecode" | b"DCT") if !jpeg => return Err("JPEG images inside other filters can't be exported yet".into()),
-        _ => {}
+    let jpeg = matches!(f.last().map(Vec::as_slice), Some(b"DCTDecode" | b"DCT"));
+    if matches!(f.last().map(Vec::as_slice), Some(b"JPXDecode")) {
+        return Err("JPEG 2000 images can't be exported yet".into());
     }
-    let as_is = || Ok(("jpg", s.raw.to_vec()));
+    // Decode only the stages before DCT, within the same work budget as other image filters.
+    // The JPEG decoder and as-is export below both need the unwrapped JPEG bytes.
+    let jpeg_data = if jpeg && f.len() > 1 {
+        let direct = doc.with_direct_parms(s);
+        let chain = direct.filters();
+        let Some(((pdfcraft_filters::Filter::Dct, _), prefix)) = chain.split_last() else {
+            return Err("invalid JPEG filter chain".into());
+        };
+        std::borrow::Cow::Owned(budget.decode_chain(prefix, &s.raw, MAX_DECODED).map_err(|e| e.to_string())?)
+    } else {
+        std::borrow::Cow::Borrowed(s.raw.as_ref())
+    };
+    let as_is = || Ok(("jpg", jpeg_data.to_vec()));
     let (w, h) = (s.dict.int(b"Width").unwrap_or(0), s.dict.int(b"Height").unwrap_or(0));
     let Some((w, h)) = size(w, h) else {
         return match () {
@@ -195,12 +205,12 @@ fn image(doc: &Document, s: &Stream, budget: &mut ImageBudget) -> Result<(&'stat
         None => space(doc, s.dict.get(b"ColorSpace"), budget)?,
     };
     let (samples, bpc) = if jpeg {
-        if !budget.pay_jpeg((w as u64).saturating_mul(h as u64).saturating_mul(space.components() as u64), &s.raw) {
+        if !budget.pay_jpeg((w as u64).saturating_mul(h as u64).saturating_mul(space.components() as u64), &jpeg_data) {
             return Err(FilterError::Spent.to_string());
         }
         // Grey, RGB or CMYK samples (YCCK is converted); a JPEG that doesn't decode, or an
         // indexed one (not valid), is written as is, without its transparency.
-        let Some((samples, decoded)) = jpeg_samples(&s.raw, &space, w, h) else { return as_is() };
+        let Some((samples, decoded)) = jpeg_samples(&jpeg_data, &space, w, h) else { return as_is() };
         space = decoded;
         (samples, 8)
     } else {

crates/create/src/tests.rs

@@ -3073,3 +3073,49 @@ fn office_documents_are_their_own_kind() {
     assert_eq!(source_kind("Report.docx", b"%PDF-1.7\n"), Some(SourceKind::Pdf));
     assert_eq!(source_kind("photos.zip", b"PK\x03\x04\x14\x00"), None);
 }
+
+#[test]
+fn flate_wrapped_jpeg_exports_with_and_without_a_soft_mask() {
+    for masked in [false, true] {
+        let jpeg = real_jpeg();
+        let mut doc = from_images(&[("flat.png".into(), png_bytes(false))]).unwrap();
+        with_image(&mut doc, |doc| {
+            let mut d = image_dict(4, 2, "DeviceGray", 8);
+            d.set(b"Filter".to_vec(), Object::Array(vec![Object::name("FlateDecode"), Object::name("DCTDecode")]));
+            if masked {
+                let mask = Stream::from_raw(image_dict(4, 2, "DeviceGray", 8), vec![0, 255, 0, 255, 255, 0, 255, 0]);
+                d.set(b"SMask".to_vec(), Object::Ref(doc.add(Object::Stream(mask))));
+            }
+            Stream::from_raw(d, pdfcraft_filters::encode_flate(&jpeg))
+        });
+        let doc = reopen(&doc);
+        let out = crate::extract::extract_images(&doc, &[0], 0);
+        assert!(out.skipped.is_empty(), "{:?}", out.skipped);
+        assert_eq!(out.images.len(), 1);
+        if masked {
+            assert_eq!(out.images[0].extension, "png");
+            assert_eq!(alphas(&exported_png_pixels(&out.images[0].data)), vec![0, 255, 0, 255, 255, 0, 255, 0]);
+        } else {
+            assert_eq!(out.images[0].extension, "jpg");
+            assert_eq!(out.images[0].data, jpeg);
+        }
+    }
+}
+
+#[test]
+fn flate_wrapped_jpeg_rejects_corruption_and_spent_budget() {
+    for corrupt in [false, true] {
+        let mut doc = from_images(&[("flat.png".into(), png_bytes(false))]).unwrap();
+        with_image(&mut doc, |_| {
+            let mut d = image_dict(4, 2, "DeviceGray", 8);
+            d.set(b"Filter".to_vec(), Object::Array(vec![Object::name("FlateDecode"), Object::name("DCTDecode")]));
+            let bytes = if corrupt { vec![1, 2, 3] } else { pdfcraft_filters::encode_flate(&real_jpeg()) };
+            Stream::from_raw(d, bytes)
+        });
+        let doc = reopen(&doc);
+        let mut budget = pdfcraft_cos::ImageBudget::with_work(if corrupt { u64::MAX } else { 0 });
+        let out = crate::extract::extract_images_within(&doc, &[0], 0, &mut budget, usize::MAX);
+        assert!(out.images.is_empty());
+        assert_eq!(out.skipped.len(), 1, "{:?}", out.skipped);
+    }
+}