M2.6: a JBIG2 image stored again without its globals decodes within the run's allowance

crates/redact/src/image.rs

@@ -211,12 +211,12 @@ fn stored_len(s: &Stream) -> usize {
 
 /// JBIG2 image `s` stored again without its `/JBIG2Globals`: decoded through them, as it is drawn,
 /// and written as the same samples, Flate, Group 4 or JBIG2 coded alone, whichever is smallest
-/// (`smallest`), its other entries kept. `None` when it can't
-/// be decoded for certain (too large, data that doesn't decode, a JBIG2 page of another size than
-/// the image, or not a one-bit image of one component). Redaction does this to every image that
+/// (`smallest`), its other entries kept. `None` when it can't be decoded for certain (too large,
+/// past what is left of `budget`, data that doesn't decode, a JBIG2 page of another size than the
+/// image, or not a one-bit image of one component). Redaction does this to every image that
 /// shares a redacted image's globals, since their symbol dictionary holds the shapes of marks drawn
 /// under the area too.
-pub(crate) fn without_globals(doc: &Document, s: &Stream) -> Option<Stream> {
+pub(crate) fn without_globals(doc: &Document, s: &Stream, budget: &mut ImageBudget) -> Option<Stream> {
     let (w, h) = (s.dict.int(b"Width")?, s.dict.int(b"Height")?);
     if w <= 0 || h <= 0 || (w as u64) * (h as u64) > MAX_PIXELS {
         return None;
@@ -231,7 +231,7 @@ pub(crate) fn without_globals(doc: &Document, s: &Stream) -> Option<Stream> {
     // As for `clear`: the filters before the JBIG2 decoder may give up to three bytes for each one
     // of the raster and a byte a row.
     let allowance = bytes.checked_add(h)?.checked_mul(3)?.checked_add(4096)?;
-    let mut data = doc.decoded_image_within(s, allowance).ok()?;
+    let mut data = budget.decode(doc, s, allowance).ok()?;
     if data.len() < bytes {
         return None;
     }

crates/redact/src/lib.rs

@@ -449,7 +449,7 @@ fn apply_within(doc: &mut Document, pages: Option<&[usize]>, images: pdfcraft_co
     if applying.images.refused() > 0 {
         return Err(RedactError::ImagesPastAllowance);
     }
-    let Applying { mut report, doomed_fields, doomed_xobjects, retired_properties, mut removed_everywhere, .. } = applying;
+    let Applying { mut report, doomed_fields, doomed_xobjects, retired_properties, mut removed_everywhere, mut images, .. } = applying;
 
     // 2b. Replies and pop-ups of a removed annotation can be listed on another page than their
     // parent: their `/IRT` or `/Parent` would keep the removed comment reachable.
@@ -532,7 +532,7 @@ fn apply_within(doc: &mut Document, pages: Option<&[usize]>, images: pdfcraft_co
     // A removed image or form must not stay reachable through another resource dictionary (an
     // inherited or shared /Resources): a full save keeps whatever is still referenced.
     interp::unlink_doomed(doc, &doomed_xobjects, &retired_properties)?;
-    report.images_restored = restore_shared_jbig2(doc, &doomed_xobjects)?;
+    report.images_restored = restore_shared_jbig2(doc, &doomed_xobjects, &mut images)?;
 
     // 5. Verify: nothing readable may remain under a region.
     let mut residue = 0;
@@ -570,8 +570,10 @@ const MAX_JBIG2_WALK: usize = 100_000;
 /// and stored again in place without them, as a cleared copy is, and a full save then leaves them
 /// out. An image that can't be stored again refuses the redaction (nothing changes, as `apply`
 /// works on a copy), but for a removed or replaced one, which is either unlinked or still drawn
-/// whole on another page. Returns how many images that weren't removed or replaced were stored again.
-fn restore_shared_jbig2(doc: &mut Document, doomed: &[ObjRef]) -> Result<usize, RedactError> {
+/// whole on another page. They are decoded within what is left of the run's allowance (`images`);
+/// past it the redaction is refused as one clearing too many images is. Returns how many images
+/// that weren't removed or replaced were stored again.
+fn restore_shared_jbig2(doc: &mut Document, doomed: &[ObjRef], images: &mut pdfcraft_cos::ImageBudget) -> Result<usize, RedactError> {
     let mut globals = HashSet::new();
     let mut seen = HashSet::new();
     let mut stack: Vec<ObjRef> = doomed.to_vec();
@@ -603,19 +605,24 @@ fn restore_shared_jbig2(doc: &mut Document, doomed: &[ObjRef]) -> Result<usize,
         .map(|(r, _)| r)
         .collect();
     let doomed: HashSet<ObjRef> = doomed.iter().copied().collect();
-    let (mut restored, mut stuck) = (0, 0);
+    let (mut restored, mut stuck, mut past_allowance) = (0, 0, false);
     for r in sharing {
         let object = doc.get(r);
         let Object::Stream(s) = &*object else { continue };
-        match image::without_globals(doc, s) {
+        let refused = images.refused();
+        match image::without_globals(doc, s, images) {
             Some(new) => {
                 doc.set(r, Object::Stream(new));
                 restored += usize::from(!doomed.contains(&r));
             }
             None if doomed.contains(&r) => {}
+            None if images.refused() > refused => past_allowance = true,
             None => stuck += 1,
         }
     }
+    if past_allowance {
+        return Err(RedactError::ImagesPastAllowance);
+    }
     if stuck > 0 {
         return Err(RedactError::SharedJbig2Globals(stuck));
     }

crates/redact/src/tests.rs

@@ -1303,6 +1303,15 @@ fn globals_a_redacted_jbig2_image_shares_with_another_page_leave_the_file() {
     let mut doc = build(&image.replace(&format!("/Height {h}"), &format!("/Height {}", h + 1)));
     mark(&mut doc, 0, &[[0.25, 113.75, 49.75, 120.0]], "");
     assert!(matches!(apply(&mut doc, None), Err(RedactError::SharedJbig2Globals(1))));
+    // Storing the second page's image again decodes it, within what is left of the run's
+    // allowance: with nothing left the redaction is refused as one past the allowance, not
+    // reported as an image that can't be read.
+    let mut doc = build(&image);
+    let first_image = [ObjRef::new(7, 0)];
+    let refused = crate::restore_shared_jbig2(&mut doc, &first_image, &mut pdfcraft_cos::ImageBudget::with_work(0));
+    assert!(matches!(refused, Err(RedactError::ImagesPastAllowance)), "{refused:?}");
+    let restored = crate::restore_shared_jbig2(&mut doc, &first_image, &mut pdfcraft_cos::ImageBudget::with_work(u64::MAX));
+    assert_eq!(restored.ok(), Some(1));
 }
 
 /// Apply Redactions gives the page its own pruned /Resources, but the dictionary it drew through