diff --git a/crates/edit/src/images.rs b/crates/edit/src/images.rs index 065a5430c..26b2d8e54 100644 --- a/crates/edit/src/images.rs +++ b/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 { +fn remove_target_aliases(dict: Option<&mut Dict>, target: ObjRef, aliases: &mut usize, budget: &mut ImageRewriteBudget) -> Result { 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, aliases: &mut usize, depth: usize, - visits: &mut usize, + budget: &mut ImageRewriteBudget, ) -> Result { - 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 { - charge_object_walk(depth, visits)?; +fn contains_target_ref(object: &Object, target: ObjRef, depth: usize, budget: &mut ImageRewriteBudget) -> Result { + 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 { +fn remove_target_alternate_entries( + object: &mut Object, + doc: &Document, + target: ObjRef, + aliases: &mut usize, + budget: &mut ImageRewriteBudget, +) -> Result { 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, aliases: &mut usize, depth: usize, - visits: &mut usize, + budget: &mut ImageRewriteBudget, ) -> Result { - 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]> { diff --git a/crates/edit/src/tests.rs b/crates/edit/src/tests.rs index d7a25af3a..1cdc6bcca 100644 --- a/crates/edit/src/tests.rs +++ b/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) { + 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>, 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";