fix(edit): handle blank pages and bound image cleanup rewrites

crates/edit/src/images.rs

@@ -114,6 +114,38 @@ const IMAGE_CENSUS_MAX_OPERATORS: usize = 5_000_000;
 const IMAGE_CENSUS_MAX_VISITS: usize = 50_000;
 const IMAGE_CENSUS_MAX_DEPTH: usize = 64;
 const IMAGE_CENSUS_MAX_ALIASES: usize = 255;
+const IMAGE_REWRITE_MAX_ENTRIES: usize = 500_000;
+const IMAGE_REWRITE_MAX_CONTAINER_ENTRIES: usize = 100_000;
+const IMAGE_REWRITE_MAX_KEY_BYTES: usize = 16 * 1024 * 1024;
+const IMAGE_REWRITE_MAX_STAGED_OBJECTS: usize = IMAGE_CENSUS_MAX_ALIASES;
+
+/// Separate budget for the conservative cleanup rewrite walk. It counts every array item and
+/// dictionary entry inspected, including primitive values, without recursively visiting scalar
+/// leaves. This work is distinct from the executable-content census budget above.
+#[derive(Default)]
+struct ImageRewriteBudget {
+    entries: usize,
+    key_bytes: usize,
+}
+
+impl ImageRewriteBudget {
+    fn charge_entries(&mut self, count: usize) -> Result<(), ()> {
+        if count > IMAGE_REWRITE_MAX_CONTAINER_ENTRIES {
+            return Err(());
+        }
+        self.entries = self.entries.saturating_add(count);
+        if self.entries > IMAGE_REWRITE_MAX_ENTRIES { Err(()) } else { Ok(()) }
+    }
+
+    fn charge_dict(&mut self, dict: &Dict) -> Result<(), ()> {
+        self.charge_entries(dict.len())
+    }
+
+    fn charge_key_bytes(&mut self, count: usize) -> Result<(), ()> {
+        self.key_bytes = self.key_bytes.saturating_add(count);
+        if self.key_bytes > IMAGE_REWRITE_MAX_KEY_BYTES { Err(()) } else { Ok(()) }
+    }
+}
 
 /// A full-document, bounded census of executable resource consumers. `incomplete` means a
 /// conservative cleanup cannot prove that an image is unused, so callers must leave its bytes.
@@ -502,9 +534,16 @@ fn prune_unused_image(doc: &mut Document, target: ObjRef) {
         return;
     }
     for page in &pages {
-        let Some(resources) = page.dict.get(b"Resources").map(|r| doc.resolve(r)).and_then(|r| r.as_dict().cloned()) else {
-            census.incomplete = true;
-            break;
+        let resources = match page.dict.get(b"Resources") {
+            None => Dict::new(),
+            Some(value) => {
+                let resolved = doc.resolve(value);
+                let Some(resources) = resolved.as_dict().cloned() else {
+                    census.incomplete = true;
+                    break;
+                };
+                resources
+            }
         };
         let contents_object = page.dict.get(b"Contents").map(|value| doc.resolve(value));
         let contents: &[Object] = match contents_object.as_deref() {
@@ -578,22 +617,21 @@ fn prune_unused_image(doc: &mut Document, target: ObjRef) {
     let mut indirect_xobject_dicts = std::collections::HashSet::new();
     let mut indirect_alternate_arrays = std::collections::HashSet::new();
     let mut aliases = 0usize;
-    let mut mapping_visits = 0usize;
-    let mut dependency_visits = 0usize;
+    let mut rewrite_budget = ImageRewriteBudget::default();
     for (reference, item) in checked {
         let Ok(object) = item else { return };
         let mut object = (*object).clone();
-        let mappings_changed = match prune_xobject_mappings(&mut object, target, &mut indirect_xobject_dicts, &mut aliases, 0, &mut mapping_visits) {
+        let mappings_changed = match prune_xobject_mappings(&mut object, target, &mut indirect_xobject_dicts, &mut aliases, 0, &mut rewrite_budget) {
             Ok(changed) => changed,
             Err(()) => return,
         };
         let dependencies_changed =
-            match prune_image_dependencies(&mut object, doc, target, &mut indirect_alternate_arrays, &mut aliases, 0, &mut dependency_visits) {
+            match prune_image_dependencies(&mut object, doc, target, &mut indirect_alternate_arrays, &mut aliases, 0, &mut rewrite_budget) {
                 Ok(changed) => changed,
                 Err(()) => return,
             };
         if mappings_changed || dependencies_changed {
-            if aliases > IMAGE_CENSUS_MAX_ALIASES {
+            if aliases > IMAGE_CENSUS_MAX_ALIASES || staged.len() >= IMAGE_REWRITE_MAX_STAGED_OBJECTS {
                 return;
             }
             staged.insert(reference, object);
@@ -602,12 +640,12 @@ fn prune_unused_image(doc: &mut Document, target: ObjRef) {
     for reference in indirect_xobject_dicts {
         let was_staged = staged.contains_key(&reference);
         let mut object = staged.remove(&reference).unwrap_or_else(|| (*doc.get(reference)).clone());
-        let changed = match remove_target_aliases(object.as_dict_mut(), target, &mut aliases) {
+        let changed = match remove_target_aliases(object.as_dict_mut(), target, &mut aliases, &mut rewrite_budget) {
             Ok(changed) => changed,
             Err(()) => return,
         };
         if changed || was_staged {
-            if aliases > IMAGE_CENSUS_MAX_ALIASES {
+            if aliases > IMAGE_CENSUS_MAX_ALIASES || (!was_staged && staged.len() >= IMAGE_REWRITE_MAX_STAGED_OBJECTS) {
                 return;
             }
             staged.insert(reference, object);
@@ -616,12 +654,12 @@ fn prune_unused_image(doc: &mut Document, target: ObjRef) {
     for reference in indirect_alternate_arrays {
         let was_staged = staged.contains_key(&reference);
         let mut object = staged.remove(&reference).unwrap_or_else(|| (*doc.get(reference)).clone());
-        let changed = match remove_target_alternate_entries(&mut object, doc, target, &mut aliases, &mut dependency_visits) {
+        let changed = match remove_target_alternate_entries(&mut object, doc, target, &mut aliases, &mut rewrite_budget) {
             Ok(changed) => changed,
             Err(()) => return,
         };
         if changed {
-            if aliases > IMAGE_CENSUS_MAX_ALIASES {
+            if aliases > IMAGE_CENSUS_MAX_ALIASES || (!was_staged && staged.len() >= IMAGE_REWRITE_MAX_STAGED_OBJECTS) {
                 return;
             }
             staged.insert(reference, object);
@@ -638,13 +676,12 @@ fn prune_unused_image(doc: &mut Document, target: ObjRef) {
     }
 }
 
-fn remove_target_aliases(dict: Option<&mut Dict>, target: ObjRef, aliases: &mut usize) -> Result<bool, ()> {
+fn remove_target_aliases(dict: Option<&mut Dict>, target: ObjRef, aliases: &mut usize, budget: &mut ImageRewriteBudget) -> Result<bool, ()> {
     let Some(dict) = dict else { return Ok(false) };
-    if dict.iter().take(IMAGE_CENSUS_MAX_VISITS + 1).count() > IMAGE_CENSUS_MAX_VISITS {
-        return Err(());
-    }
+    budget.charge_dict(dict)?;
     let mut names = Vec::new();
     for (name, value) in dict.iter() {
+        budget.charge_key_bytes(name.len())?;
         if value.as_ref() == Some(target) {
             if aliases.saturating_add(names.len()) >= IMAGE_CENSUS_MAX_ALIASES {
                 return Err(());
@@ -665,14 +702,17 @@ fn prune_xobject_mappings(
     indirect: &mut std::collections::HashSet<ObjRef>,
     aliases: &mut usize,
     depth: usize,
-    visits: &mut usize,
+    budget: &mut ImageRewriteBudget,
 ) -> Result<bool, ()> {
-    charge_object_walk(depth, visits)?;
+    if depth > IMAGE_CENSUS_MAX_DEPTH {
+        return Err(());
+    }
     let mut changed = false;
     match object {
         Object::Array(values) => {
-            for value in values {
-                changed |= prune_xobject_mappings(value, target, indirect, aliases, depth + 1, visits)?;
+            budget.charge_entries(values.len())?;
+            for value in values.iter_mut().filter(|value| is_container(value)) {
+                changed |= prune_xobject_mappings(value, target, indirect, aliases, depth + 1, budget)?;
             }
         }
         Object::Dict(dict) | Object::Stream(pdfcraft_cos::Stream { dict, .. }) => {
@@ -681,20 +721,21 @@ fn prune_xobject_mappings(
                     Object::Ref(reference) => {
                         indirect.insert(*reference);
                     }
-                    Object::Dict(mappings) => changed |= remove_target_aliases(Some(mappings), target, aliases)?,
+                    Object::Dict(mappings) => changed |= remove_target_aliases(Some(mappings), target, aliases, budget)?,
                     _ => {}
                 }
             }
             let mut keys = Vec::new();
-            for (key, _) in dict.iter() {
-                if keys.len() >= IMAGE_CENSUS_MAX_VISITS {
-                    return Err(());
+            budget.charge_dict(dict)?;
+            for (key, value) in dict.iter() {
+                budget.charge_key_bytes(key.len())?;
+                if is_container(value) {
+                    keys.push(key.clone());
                 }
-                keys.push(key.clone());
             }
             for key in keys {
                 if let Some(value) = dict.get_mut(&key) {
-                    changed |= prune_xobject_mappings(value, target, indirect, aliases, depth + 1, visits)?;
+                    changed |= prune_xobject_mappings(value, target, indirect, aliases, depth + 1, budget)?;
                 }
             }
         }
@@ -703,13 +744,16 @@ fn prune_xobject_mappings(
     Ok(changed)
 }
 
-fn contains_target_ref(object: &Object, target: ObjRef, depth: usize, visits: &mut usize) -> Result<bool, ()> {
-    charge_object_walk(depth, visits)?;
+fn contains_target_ref(object: &Object, target: ObjRef, depth: usize, budget: &mut ImageRewriteBudget) -> Result<bool, ()> {
+    if depth > IMAGE_CENSUS_MAX_DEPTH {
+        return Err(());
+    }
     match object {
         Object::Ref(reference) => Ok(*reference == target),
         Object::Array(values) => {
-            for value in values {
-                if contains_target_ref(value, target, depth + 1, visits)? {
+            budget.charge_entries(values.len())?;
+            for value in values.iter().filter(|value| matches!(value, Object::Array(_))) {
+                if contains_target_ref(value, target, depth + 1, budget)? {
                     return Ok(true);
                 }
             }
@@ -723,12 +767,18 @@ fn alternate_uses_target(object: &Object, doc: &Document, target: ObjRef) -> boo
     doc.resolve(object).as_dict().and_then(|dict| dict.get(b"Image")).is_some_and(|image| image.as_ref() == Some(target))
 }
 
-fn remove_target_alternate_entries(object: &mut Object, doc: &Document, target: ObjRef, aliases: &mut usize, visits: &mut usize) -> Result<bool, ()> {
+fn remove_target_alternate_entries(
+    object: &mut Object,
+    doc: &Document,
+    target: ObjRef,
+    aliases: &mut usize,
+    budget: &mut ImageRewriteBudget,
+) -> Result<bool, ()> {
     let Object::Array(entries) = object else { return Ok(false) };
+    budget.charge_entries(entries.len())?;
     let mut retained = Vec::new();
     let mut removed = 0;
     for entry in entries.iter() {
-        charge_object_walk(1, visits)?;
         if alternate_uses_target(entry, doc, target) {
             if aliases.saturating_add(removed) >= IMAGE_CENSUS_MAX_ALIASES {
                 return Err(());
@@ -755,21 +805,24 @@ fn prune_image_dependencies(
     indirect_alternate_arrays: &mut std::collections::HashSet<ObjRef>,
     aliases: &mut usize,
     depth: usize,
-    visits: &mut usize,
+    budget: &mut ImageRewriteBudget,
 ) -> Result<bool, ()> {
-    charge_object_walk(depth, visits)?;
+    if depth > IMAGE_CENSUS_MAX_DEPTH {
+        return Err(());
+    }
     let mut changed = false;
     match object {
         Object::Array(values) => {
-            for value in values {
-                changed |= prune_image_dependencies(value, doc, target, indirect_alternate_arrays, aliases, depth + 1, visits)?;
+            budget.charge_entries(values.len())?;
+            for value in values.iter_mut().filter(|value| is_container(value)) {
+                changed |= prune_image_dependencies(value, doc, target, indirect_alternate_arrays, aliases, depth + 1, budget)?;
             }
         }
         Object::Dict(dict) | Object::Stream(pdfcraft_cos::Stream { dict, .. }) => {
             if dict.name(b"Subtype") == Some(b"Image") {
                 for key in [b"SMask".as_slice(), b"Mask".as_slice()] {
                     let contains_target =
-                        if let Some(value) = dict.get(key) { contains_target_ref(value, target, depth + 1, visits)? } else { false };
+                        if let Some(value) = dict.get(key) { contains_target_ref(value, target, depth + 1, budget)? } else { false };
                     if contains_target {
                         if aliases.saturating_add(1) > IMAGE_CENSUS_MAX_ALIASES {
                             return Err(());
@@ -782,10 +835,10 @@ fn prune_image_dependencies(
                 match dict.get(b"Alternates") {
                     Some(Object::Array(_)) => {
                         let Some(Object::Array(entries)) = dict.get_mut(b"Alternates") else { return Err(()) };
+                        budget.charge_entries(entries.len())?;
                         let mut retained = Vec::new();
                         let mut removed = 0;
                         for entry in entries.iter() {
-                            charge_object_walk(depth + 1, visits)?;
                             if alternate_uses_target(entry, doc, target) {
                                 if aliases.saturating_add(removed) >= IMAGE_CENSUS_MAX_ALIASES {
                                     return Err(());
@@ -808,15 +861,16 @@ fn prune_image_dependencies(
                 }
             }
             let mut keys = Vec::new();
-            for (key, _) in dict.iter() {
-                if keys.len() >= IMAGE_CENSUS_MAX_VISITS {
-                    return Err(());
+            budget.charge_dict(dict)?;
+            for (key, value) in dict.iter() {
+                budget.charge_key_bytes(key.len())?;
+                if is_container(value) {
+                    keys.push(key.clone());
                 }
-                keys.push(key.clone());
             }
             for key in keys {
                 if let Some(value) = dict.get_mut(&key) {
-                    changed |= prune_image_dependencies(value, doc, target, indirect_alternate_arrays, aliases, depth + 1, visits)?;
+                    changed |= prune_image_dependencies(value, doc, target, indirect_alternate_arrays, aliases, depth + 1, budget)?;
                 }
             }
         }
@@ -825,12 +879,8 @@ fn prune_image_dependencies(
     Ok(changed)
 }
 
-fn charge_object_walk(depth: usize, visits: &mut usize) -> Result<(), ()> {
-    if depth > IMAGE_CENSUS_MAX_DEPTH {
-        return Err(());
-    }
-    *visits = visits.saturating_add(1);
-    if *visits > IMAGE_CENSUS_MAX_VISITS { Err(()) } else { Ok(()) }
+fn is_container(object: &Object) -> bool {
+    matches!(object, Object::Array(_) | Object::Dict(_) | Object::Stream(_))
 }
 
 fn intersection(a: [f64; 4], b: [f64; 4]) -> Option<[f64; 4]> {

crates/edit/src/tests.rs

@@ -2252,6 +2252,62 @@ fn set_image_content(doc: &mut Document, page: usize, pieces: &[&[u8]]) {
     doc.update_dict(page, |d| d.set(b"Contents".to_vec(), Object::Array(contents))).unwrap();
 }
 
+fn add_blank_page(doc: &mut Document, resources: Option<Object>) {
+    let mut page = Dict::from_iter([
+        (b"Type".to_vec(), Object::name("Page")),
+        (b"Parent".to_vec(), Object::Ref(pdfcraft_cos::ObjRef::new(2, 0))),
+        (b"MediaBox".to_vec(), Object::Array([0, 0, 600, 800].map(Object::Int).to_vec())),
+    ]);
+    if let Some(resources) = resources {
+        page.set(b"Resources".to_vec(), resources);
+    }
+    let added = doc.add(Object::Dict(page));
+    doc.update_dict(pdfcraft_cos::ObjRef::new(2, 0), |d| {
+        d.set(b"Kids".to_vec(), Object::Array(vec![Object::Ref(pdfcraft_cos::ObjRef::new(3, 0)), Object::Ref(added)]));
+        d.set(b"Count".to_vec(), Object::Int(2));
+    })
+    .unwrap();
+}
+
+fn image_pruning_fixture(blank_resources: Option<Option<Object>>, font_count: usize, private: &[u8]) -> Document {
+    let mut doc = image_page();
+    let target = images::page_images(&doc, 0).unwrap()[0].object.unwrap();
+    doc.set(target, Object::Stream(identifiable_test_image(private)));
+    if let Some(resources) = blank_resources {
+        add_blank_page(&mut doc, resources);
+    }
+    if font_count > 0 {
+        let mut fonts = doc
+            .get(pdfcraft_model::pages(&doc)[0].obj)
+            .as_dict()
+            .and_then(|d| d.get(b"Resources"))
+            .and_then(Object::as_dict)
+            .and_then(|d| d.get(b"Font"))
+            .and_then(Object::as_dict)
+            .cloned()
+            .unwrap_or_default();
+        for index in 0..font_count {
+            let font = doc.add(Object::Dict(Dict::from_iter([
+                (b"Type".to_vec(), Object::name("Font")),
+                (b"Subtype".to_vec(), Object::name("Type1")),
+                (b"BaseFont".to_vec(), Object::name("Helvetica")),
+                (b"FirstChar".to_vec(), Object::Int(32)),
+                (b"LastChar".to_vec(), Object::Int(255)),
+                (b"Widths".to_vec(), Object::Array(vec![Object::Int(500); 224])),
+            ])));
+            fonts.set(format!("ReviewFont{index}").into_bytes(), Object::Ref(font));
+        }
+        let page = pdfcraft_model::pages(&doc)[0].obj;
+        doc.update_dict(page, |d| {
+            let mut resources = d.get(b"Resources").and_then(Object::as_dict).cloned().unwrap_or_default();
+            resources.set(b"Font".to_vec(), Object::Dict(fonts));
+            d.set(b"Resources".to_vec(), Object::Dict(resources));
+        })
+        .unwrap();
+    }
+    doc
+}
+
 #[test]
 fn deleted_image_is_kept_when_reachable_through_each_non_page_consumer() {
     const PRIVATE: &[u8] = b"PRIVATE_CONTEXT_IMAGE";
@@ -2582,6 +2638,79 @@ fn deleting_image_prunes_unused_shared_resources_across_all_pages() {
     assert!(xobject(&saved, 1, b"Im0").is_none());
 }
 
+#[test]
+fn image_payload_pruning_handles_missing_resources_and_wide_scalar_arrays() {
+    const PRIVATE: &[u8] = b"PRIVATE_REVIEW_DELETED_IMAGE_PAYLOAD";
+    for (resources, font_count) in [(Some(None), 0), (Some(Some(Object::Dict(Dict::new()))), 0), (None, 225)] {
+        for replace in [false, true] {
+            let mut source = image_pruning_fixture(resources.clone(), font_count, PRIVATE);
+            let input = write_full(&source, &SaveOptions::default()).unwrap();
+            assert!(input.windows(PRIVATE.len()).any(|window| window == PRIVATE), "fixture lacked payload before editing");
+            source = Document::open(Arc::new(input)).unwrap();
+            let change = if replace {
+                images::ImageChange::Replace(source.add(Object::Stream(identifiable_test_image(b"PUBLIC_REPLACEMENT"))))
+            } else {
+                images::ImageChange::Delete
+            };
+            images::change_image(&mut source, 0, 0, &change).unwrap();
+
+            let output = write_full(&source, &SaveOptions::default()).unwrap();
+            let reopened = Document::open(Arc::new(output.clone())).unwrap();
+            assert!(!output.windows(PRIVATE.len()).any(|window| window == PRIVATE), "replace={replace}, resources={resources:?}, fonts={font_count}");
+            assert_eq!(images::page_images(&reopened, 0).unwrap().len(), usize::from(replace));
+        }
+    }
+}
+
+#[test]
+fn malformed_present_page_resources_keep_image_payload() {
+    const PRIVATE: &[u8] = b"PRIVATE_MALFORMED_RESOURCES_IMAGE";
+    let mut source = image_pruning_fixture(None, 0, PRIVATE);
+    add_blank_page(&mut source, Some(Object::Int(7)));
+    let input = write_full(&source, &SaveOptions::default()).unwrap();
+    assert!(input.windows(PRIVATE.len()).any(|window| window == PRIVATE), "fixture lacked payload before editing");
+    let mut source = Document::open(Arc::new(input)).unwrap();
+
+    images::change_image(&mut source, 0, 0, &images::ImageChange::Delete).unwrap();
+    let output = write_full(&source, &SaveOptions::default()).unwrap();
+    assert!(output.windows(PRIVATE.len()).any(|window| window == PRIVATE), "malformed present Resources unexpectedly allowed pruning");
+}
+
+#[test]
+fn rewrite_entry_budget_failure_keeps_earlier_staged_image_mappings() {
+    const PRIVATE: &[u8] = b"PRIVATE_REWRITE_BUDGET_IMAGE";
+    const OVERSIZED_CONTAINER: usize = 100_001;
+    let mut source = image_pruning_fixture(None, 0, PRIVATE);
+    let form = source.add(Object::Stream(Stream::flate(
+        Dict::from_iter([
+            (b"Subtype".to_vec(), Object::name("Form")),
+            (b"BBox".to_vec(), Object::Array([0, 0, 10, 10].map(Object::Int).to_vec())),
+            (b"Probe".to_vec(), Object::Array(vec![Object::Int(0); OVERSIZED_CONTAINER])),
+        ]),
+        b"",
+    )));
+    let page = pdfcraft_model::pages(&source)[0].obj;
+    source
+        .update_dict(page, |d| {
+            let mut resources = d.get(b"Resources").and_then(Object::as_dict).cloned().unwrap();
+            let mut xobjects = resources.get(b"XObject").and_then(Object::as_dict).cloned().unwrap();
+            xobjects.set(b"F".to_vec(), Object::Ref(form));
+            resources.set(b"XObject".to_vec(), Object::Dict(xobjects));
+            d.set(b"Resources".to_vec(), Object::Dict(resources));
+        })
+        .unwrap();
+    set_image_content(&mut source, 0, &[b"q 200 0 0 100 100 100 cm /Im0 Do Q /F Do"]);
+    let input = write_full(&source, &SaveOptions::default()).unwrap();
+    assert!(input.windows(PRIVATE.len()).any(|window| window == PRIVATE), "fixture lacked payload before editing");
+    let mut source = Document::open(Arc::new(input)).unwrap();
+
+    images::change_image(&mut source, 0, 0, &images::ImageChange::Delete).unwrap();
+    let output = write_full(&source, &SaveOptions::default()).unwrap();
+    let reopened = Document::open(Arc::new(output.clone())).unwrap();
+    assert!(output.windows(PRIVATE.len()).any(|window| window == PRIVATE), "rewrite budget failure pruned image payload");
+    assert!(xobject(&reopened, 0, b"Im0").is_some(), "partial staging removed the page's image alias");
+}
+
 #[test]
 fn image_mask_dependency_keeps_a_deleted_image_live() {
     const PRIVATE: &[u8] = b"PRIVATE_IMAGE_MASK_DEPENDENCY";