Search each shared name tree node once when finding the editing file (#1702)

crates/pdf/src/editing.rs

@@ -2,10 +2,11 @@
 //! of the pages it was written with, so a PDF whose pages another app changed is noticed (and
 //! opens as its artwork instead of the stale native document).
 
+use std::collections::HashSet;
 use std::sync::Arc;
 
 use hayro_syntax::Pdf;
-use hayro_syntax::object::{Array, Dict, Object, Stream, String as PdfString};
+use hayro_syntax::object::{Array, Dict, Object, ObjectIdentifier, Stream, String as PdfString};
 use krilla::embed::{AssociationKind, EmbeddedFile, MimeType};
 use krilla::metadata::DateTime;
 
@@ -77,7 +78,7 @@ pub(crate) fn editing_in(pdf: &Pdf) -> Option<Editing> {
     let xref = pdf.xref();
     let catalog = xref.get::<Dict<'_>>(xref.root_id())?;
     let tree = catalog.get::<Dict<'_>>(b"Names")?.get::<Dict<'_>>(b"EmbeddedFiles")?;
-    let spec = editing_spec(&tree, 0)?;
+    let spec = editing_spec(&tree, 0, &mut HashSet::new())?;
     let files = spec.get::<Dict<'_>>(b"EF")?;
     let stream = files.get::<Stream<'_>>(b"F").or_else(|| files.get::<Stream<'_>>(b"UF"))?;
     let data = stream.decoded().map(|d| d.into_owned()).unwrap_or_default();
@@ -90,8 +91,9 @@ pub(crate) fn editing_in(pdf: &Pdf) -> Option<Editing> {
     Some(Editing { data, intact })
 }
 
-/// The file specification of the editing file in the name tree node `node`.
-fn editing_spec<'a>(node: &Dict<'a>, depth: usize) -> Option<Dict<'a>> {
+/// The file specification of the editing file in the name tree node `node`. `seen` holds the indirect
+/// nodes already searched, so a node shared by many `Kids` arrays is searched once, not once per path.
+fn editing_spec<'a>(node: &Dict<'a>, depth: usize, seen: &mut HashSet<ObjectIdentifier>) -> Option<Dict<'a>> {
     if let Some(names) = node.get::<Array<'a>>(b"Names") {
         let mut items = names.iter::<Object<'a>>();
         while let (Some(_), Some(spec)) = (items.next(), items.next()) {
@@ -105,7 +107,14 @@ fn editing_spec<'a>(node: &Dict<'a>, depth: usize) -> Option<Dict<'a>> {
     if depth >= MAX_DEPTH {
         return None;
     }
-    node.get::<Array<'a>>(b"Kids")?.iter::<Dict<'a>>().find_map(|kid| editing_spec(&kid, depth + 1))
+    let kids = node.get::<Array<'a>>(b"Kids")?;
+    // Only a referenced kid can be shared (a direct kid's `obj_id` is the object holding it).
+    kids.raw_iter().zip(kids.iter::<Dict<'a>>()).find_map(|(raw, kid)| {
+        if raw.as_obj_ref().is_some_and(|r| !seen.insert(r.into())) {
+            return None;
+        }
+        editing_spec(&kid, depth + 1, seen)
+    })
 }
 
 /// A file specification naming the editing file (current or legacy name).

crates/pdf/src/tests_editing.rs

@@ -88,3 +88,57 @@ fn pdf_a_carries_no_editing_data() {
     assert!(matches!(s.check(), Err(PdfError::BadSetting(m)) if m.contains("editing")));
     assert!(matches!(s.check_values(), Err(PdfError::BadSetting(_))));
 }
+
+/// A tiny PDF made of the given objects (numbered from 1), with the cross-reference table written.
+fn raw_pdf(objects: &[String]) -> Vec<u8> {
+    let mut out = b"%PDF-1.7\n".to_vec();
+    let mut offsets = Vec::new();
+    for (i, body) in objects.iter().enumerate() {
+        offsets.push(out.len());
+        out.extend_from_slice(format!("{} 0 obj\n{body}\nendobj\n", i + 1).as_bytes());
+    }
+    let xref = out.len();
+    out.extend_from_slice(format!("xref\n0 {}\n0000000000 65535 f \n", objects.len() + 1).as_bytes());
+    for at in offsets {
+        out.extend_from_slice(format!("{at:010} 00000 n \n").as_bytes());
+    }
+    out.extend_from_slice(format!("trailer\n<< /Size {} /Root 1 0 R >>\nstartxref\n{xref}\n%%EOF\n", objects.len() + 1).as_bytes());
+    out
+}
+
+#[test]
+fn a_name_tree_node_shared_by_many_kids_is_searched_once() {
+    // Objects 4..=15 are a chain of name tree nodes, each listing the next one ten times and
+    // holding no attachment. Walking every path would take 10^12 steps; the editing file sits in
+    // the leaf (object 16) listed after the chain.
+    let mut objects = vec![
+        "<< /Type /Catalog /Pages 2 0 R /Names << /EmbeddedFiles 3 0 R >> >>".to_string(),
+        "<< /Type /Pages /Kids [19 0 R] /Count 1 >>".to_string(),
+        "<< /Kids [4 0 R 16 0 R] >>".to_string(),
+    ];
+    for node in 4..=14 {
+        objects.push(format!("<< /Kids [{}] >>", format!("{} 0 R ", node + 1).repeat(10)));
+    }
+    objects.push("<< /Names [] >>".to_string());
+    objects.push("<< /Names [(a) 17 0 R] >>".to_string());
+    objects.push(format!("<< /Type /Filespec /UF ({EDITING_FILE}) /EF << /F 18 0 R >> >>"));
+    objects.push("<< /Length 5 >>\nstream\nhello\nendstream".to_string());
+    objects.push("<< /Type /Page /Parent 2 0 R /MediaBox [0 0 10 10] >>".to_string());
+    let e = editing(&raw_pdf(&objects)).unwrap();
+    assert_eq!((e.data.as_slice(), e.intact), (&b"hello"[..], true));
+}
+
+#[test]
+fn direct_name_tree_kids_are_all_searched() {
+    // Direct kids report the object holding them as their id; none of them may be taken as seen.
+    let objects = [
+        "<< /Type /Catalog /Pages 2 0 R /Names << /EmbeddedFiles 3 0 R >> >>".to_string(),
+        "<< /Type /Pages /Kids [6 0 R] /Count 1 >>".to_string(),
+        "<< /Kids [<< /Names [] >> << /Kids [<< /Names [] >> << /Names [(a) 4 0 R] >>] >>] >>".to_string(),
+        format!("<< /Type /Filespec /UF ({EDITING_FILE}) /EF << /F 5 0 R >> >>"),
+        "<< /Length 5 >>\nstream\nhello\nendstream".to_string(),
+        "<< /Type /Page /Parent 2 0 R /MediaBox [0 0 10 10] >>".to_string(),
+    ];
+    let e = editing(&raw_pdf(&objects)).unwrap();
+    assert_eq!(e.data, b"hello");
+}