rpdfium_parser/

object_walker.rs

// Derived from PDFium's cpdf_object_walker.cpp
// Original: Copyright 2014 The PDFium Authors
// Licensed under BSD-3-Clause / Apache-2.0
// See pdfium-upstream/LICENSE for the original license.

//! Object graph walker -- iterative BFS traversal with cycle detection.
//!
//! Provides a visitor pattern for traversing the PDF object graph starting
//! from a root object. Uses BFS (breadth-first search) with a `VecDeque`
//! work queue and `HashSet` for cycle detection. Never recurses on the
//! call stack (WASM-safe).

use std::collections::{HashMap, HashSet, VecDeque};

use rpdfium_core::fx_system::MAX_OBJECT_NUMBER;
use rpdfium_core::{Name, PdfSource};

use crate::object::{Object, ObjectId};
use crate::store::ObjectStore;

/// Visitor trait for walking the PDF object graph.
///
/// Implement this trait to receive callbacks for each object type
/// encountered during the walk.
pub trait ObjectVisitor {
    /// Called for each dictionary encountered.
    fn visit_dict(&mut self, id: Option<ObjectId>, dict: &HashMap<Name, Object>);

    /// Called for each array encountered.
    fn visit_array(&mut self, id: Option<ObjectId>, arr: &[Object]);

    /// Called for each stream encountered.
    fn visit_stream(&mut self, id: Option<ObjectId>);

    /// Called for each reference encountered.
    fn visit_reference(&mut self, target: ObjectId);

    /// Called for each primitive value (Integer, Real, Boolean, String, Name, Null).
    fn visit_primitive(&mut self, id: Option<ObjectId>, obj: &Object);
}

/// Iterative BFS walker for the PDF object graph.
///
/// Traverses objects starting from a root, tracking visited objects to
/// detect cycles. Enforces a safety limit of `MAX_OBJECT_NUMBER` visited
/// objects to prevent denial of service.
pub struct ObjectWalker {
    visited: HashSet<ObjectId>,
    max_objects: usize,
}

impl ObjectWalker {
    /// Create a new walker with the default safety limit.
    pub fn new() -> Self {
        Self {
            visited: HashSet::new(),
            max_objects: MAX_OBJECT_NUMBER as usize,
        }
    }

    /// Create a new walker with a custom maximum objects limit.
    pub fn with_max_objects(max_objects: usize) -> Self {
        Self {
            visited: HashSet::new(),
            max_objects,
        }
    }

    /// Walk the object graph starting from `root_id`, calling visitor methods
    /// for each object encountered.
    ///
    /// Uses iterative BFS with cycle detection. Silently skips objects that
    /// cannot be resolved from the store.
    pub fn walk<S: PdfSource>(
        &mut self,
        store: &ObjectStore<S>,
        root_id: ObjectId,
        visitor: &mut dyn ObjectVisitor,
    ) {
        let mut queue: VecDeque<ObjectId> = VecDeque::new();
        queue.push_back(root_id);

        while let Some(id) = queue.pop_front() {
            // Cycle detection
            if self.visited.contains(&id) {
                continue;
            }

            // Safety limit
            if self.visited.len() >= self.max_objects {
                break;
            }

            self.visited.insert(id);

            // Resolve the object from the store
            let obj = match store.resolve(id) {
                Ok(obj) => obj,
                Err(_) => continue,
            };

            self.visit_object(obj, Some(id), visitor, &mut queue);
        }
    }

    /// Visit a single object, dispatching to the appropriate visitor method
    /// and enqueueing any references found.
    fn visit_object(
        &self,
        obj: &Object,
        id: Option<ObjectId>,
        visitor: &mut dyn ObjectVisitor,
        queue: &mut VecDeque<ObjectId>,
    ) {
        match obj {
            Object::Dictionary(dict) => {
                visitor.visit_dict(id, dict);
                self.enqueue_dict_refs(dict, queue);
            }
            Object::Stream { dict, .. } => {
                visitor.visit_stream(id);
                visitor.visit_dict(id, dict);
                self.enqueue_dict_refs(dict, queue);
            }
            Object::Array(arr) => {
                visitor.visit_array(id, arr);
                self.enqueue_array_refs(arr, queue);
            }
            Object::Reference(target) => {
                visitor.visit_reference(*target);
                if !self.visited.contains(target) {
                    queue.push_back(*target);
                }
            }
            _ => {
                // Null, Boolean, Integer, Real, String, Name
                visitor.visit_primitive(id, obj);
            }
        }
    }

    /// Enqueue all indirect references found in a dictionary's values.
    fn enqueue_dict_refs(&self, dict: &HashMap<Name, Object>, queue: &mut VecDeque<ObjectId>) {
        for value in dict.values() {
            if let Object::Reference(target) = value {
                if !self.visited.contains(target) {
                    queue.push_back(*target);
                }
            }
        }
    }

    /// Enqueue all indirect references found in an array's elements.
    fn enqueue_array_refs(&self, arr: &[Object], queue: &mut VecDeque<ObjectId>) {
        for elem in arr {
            if let Object::Reference(target) = elem {
                if !self.visited.contains(target) {
                    queue.push_back(*target);
                }
            }
        }
    }

    /// Returns the number of unique objects visited so far.
    pub fn visited_count(&self) -> usize {
        self.visited.len()
    }

    /// Returns a reference to the set of visited object IDs.
    pub fn visited(&self) -> &HashSet<ObjectId> {
        &self.visited
    }
}

impl Default for ObjectWalker {
    fn default() -> Self {
        Self::new()
    }
}

/// A built-in visitor that collects object type statistics.
#[derive(Debug, Clone, Default)]
pub struct ObjectStats {
    /// Number of dictionaries visited.
    pub dict_count: usize,
    /// Number of arrays visited.
    pub array_count: usize,
    /// Number of streams visited.
    pub stream_count: usize,
    /// Number of references visited.
    pub reference_count: usize,
    /// Number of primitive values visited.
    pub primitive_count: usize,
    /// Total number of visitor callbacks.
    pub total_visited: usize,
}

impl ObjectVisitor for ObjectStats {
    fn visit_dict(&mut self, _id: Option<ObjectId>, _dict: &HashMap<Name, Object>) {
        self.dict_count += 1;
        self.total_visited += 1;
    }

    fn visit_array(&mut self, _id: Option<ObjectId>, _arr: &[Object]) {
        self.array_count += 1;
        self.total_visited += 1;
    }

    fn visit_stream(&mut self, _id: Option<ObjectId>) {
        self.stream_count += 1;
        self.total_visited += 1;
    }

    fn visit_reference(&mut self, _target: ObjectId) {
        self.reference_count += 1;
        self.total_visited += 1;
    }

    fn visit_primitive(&mut self, _id: Option<ObjectId>, _obj: &Object) {
        self.primitive_count += 1;
        self.total_visited += 1;
    }
}

#[cfg(test)]
mod tests {
    use rpdfium_core::ParsingMode;

    use super::*;

    /// Build a minimal valid PDF for testing.
    fn build_minimal_pdf() -> Vec<u8> {
        let mut pdf = Vec::new();
        pdf.extend_from_slice(b"%PDF-1.4\n");

        let obj1_offset = pdf.len();
        pdf.extend_from_slice(b"1 0 obj\n<< /Type /Catalog /Pages 2 0 R >>\nendobj\n");

        let obj2_offset = pdf.len();
        pdf.extend_from_slice(b"2 0 obj\n<< /Type /Pages /Kids [] /Count 0 >>\nendobj\n");

        let xref_offset = pdf.len();
        pdf.extend_from_slice(b"xref\n");
        pdf.extend_from_slice(b"0 3\n");
        pdf.extend_from_slice(b"0000000000 65535 f \r\n");
        pdf.extend_from_slice(format!("{:010} 00000 n \r\n", obj1_offset).as_bytes());
        pdf.extend_from_slice(format!("{:010} 00000 n \r\n", obj2_offset).as_bytes());
        pdf.extend_from_slice(b"trailer\n");
        pdf.extend_from_slice(b"<< /Size 3 /Root 1 0 R >>\n");
        pdf.extend_from_slice(format!("startxref\n{}\n%%EOF", xref_offset).as_bytes());

        pdf
    }

    #[test]
    fn test_walk_minimal_pdf() {
        let pdf = build_minimal_pdf();
        let store = ObjectStore::open(pdf, ParsingMode::Strict).unwrap();

        let mut stats = ObjectStats::default();
        let mut walker = ObjectWalker::new();
        walker.walk(&store, ObjectId::new(1, 0), &mut stats);

        // Object 1 (Catalog dict) and Object 2 (Pages dict) should be visited
        assert_eq!(walker.visited_count(), 2);
        assert!(walker.visited().contains(&ObjectId::new(1, 0)));
        assert!(walker.visited().contains(&ObjectId::new(2, 0)));
    }

    #[test]
    fn test_stats_collection_from_minimal_pdf() {
        let pdf = build_minimal_pdf();
        let store = ObjectStore::open(pdf, ParsingMode::Strict).unwrap();

        let mut stats = ObjectStats::default();
        let mut walker = ObjectWalker::new();
        walker.walk(&store, ObjectId::new(1, 0), &mut stats);

        // Catalog: 1 dict. Pages: 1 dict.
        assert_eq!(stats.dict_count, 2);
        // Pages /Kids is an empty array
        assert_eq!(stats.array_count, 0); // Arrays inside dicts are not separately walked
        assert_eq!(stats.stream_count, 0);
        assert!(stats.total_visited > 0);
    }

    #[test]
    fn test_cycle_detection() {
        // Build a PDF where object 3 references itself via a dict value
        let mut pdf = Vec::new();
        pdf.extend_from_slice(b"%PDF-1.4\n");

        let obj1_offset = pdf.len();
        pdf.extend_from_slice(b"1 0 obj\n<< /Type /Catalog /Pages 2 0 R >>\nendobj\n");

        let obj2_offset = pdf.len();
        pdf.extend_from_slice(b"2 0 obj\n<< /Type /Pages /Kids [] /Count 0 >>\nendobj\n");

        let obj3_offset = pdf.len();
        pdf.extend_from_slice(b"3 0 obj\n<< /Self 3 0 R >>\nendobj\n");

        let xref_offset = pdf.len();
        pdf.extend_from_slice(b"xref\n");
        pdf.extend_from_slice(b"0 4\n");
        pdf.extend_from_slice(b"0000000000 65535 f \r\n");
        pdf.extend_from_slice(format!("{:010} 00000 n \r\n", obj1_offset).as_bytes());
        pdf.extend_from_slice(format!("{:010} 00000 n \r\n", obj2_offset).as_bytes());
        pdf.extend_from_slice(format!("{:010} 00000 n \r\n", obj3_offset).as_bytes());
        pdf.extend_from_slice(b"trailer\n");
        pdf.extend_from_slice(b"<< /Size 4 /Root 1 0 R >>\n");
        pdf.extend_from_slice(format!("startxref\n{}\n%%EOF", xref_offset).as_bytes());

        let store = ObjectStore::open(pdf, ParsingMode::Strict).unwrap();

        let mut stats = ObjectStats::default();
        let mut walker = ObjectWalker::new();
        walker.walk(&store, ObjectId::new(3, 0), &mut stats);

        // Should visit object 3 exactly once, then detect the cycle
        assert_eq!(walker.visited_count(), 1);
        assert!(walker.visited().contains(&ObjectId::new(3, 0)));
    }

    #[test]
    fn test_max_objects_safety_limit() {
        // Build a PDF with 3 objects but set max_objects to 1
        let pdf = build_minimal_pdf();
        let store = ObjectStore::open(pdf, ParsingMode::Strict).unwrap();

        let mut stats = ObjectStats::default();
        let mut walker = ObjectWalker::with_max_objects(1);
        walker.walk(&store, ObjectId::new(1, 0), &mut stats);

        // Should stop after visiting 1 object
        assert_eq!(walker.visited_count(), 1);
    }

    #[test]
    fn test_walk_with_no_references() {
        // Build a PDF with a single object containing no references
        let mut pdf = Vec::new();
        pdf.extend_from_slice(b"%PDF-1.4\n");

        let obj1_offset = pdf.len();
        pdf.extend_from_slice(b"1 0 obj\n<< /Type /Catalog >>\nendobj\n");

        let xref_offset = pdf.len();
        pdf.extend_from_slice(b"xref\n");
        pdf.extend_from_slice(b"0 2\n");
        pdf.extend_from_slice(b"0000000000 65535 f \r\n");
        pdf.extend_from_slice(format!("{:010} 00000 n \r\n", obj1_offset).as_bytes());
        pdf.extend_from_slice(b"trailer\n");
        pdf.extend_from_slice(b"<< /Size 2 /Root 1 0 R >>\n");
        pdf.extend_from_slice(format!("startxref\n{}\n%%EOF", xref_offset).as_bytes());

        let store = ObjectStore::open(pdf, ParsingMode::Strict).unwrap();

        let mut stats = ObjectStats::default();
        let mut walker = ObjectWalker::new();
        walker.walk(&store, ObjectId::new(1, 0), &mut stats);

        assert_eq!(walker.visited_count(), 1);
        assert_eq!(stats.dict_count, 1);
        assert_eq!(stats.reference_count, 0);
    }

    #[test]
    fn test_object_stats_counters() {
        let mut stats = ObjectStats::default();
        assert_eq!(stats.dict_count, 0);
        assert_eq!(stats.array_count, 0);
        assert_eq!(stats.stream_count, 0);
        assert_eq!(stats.reference_count, 0);
        assert_eq!(stats.primitive_count, 0);
        assert_eq!(stats.total_visited, 0);

        // Manually call visitor methods
        stats.visit_dict(None, &HashMap::new());
        stats.visit_array(None, &[]);
        stats.visit_stream(None);
        stats.visit_reference(ObjectId::new(1, 0));
        stats.visit_primitive(None, &Object::Integer(42));

        assert_eq!(stats.dict_count, 1);
        assert_eq!(stats.array_count, 1);
        assert_eq!(stats.stream_count, 1);
        assert_eq!(stats.reference_count, 1);
        assert_eq!(stats.primitive_count, 1);
        assert_eq!(stats.total_visited, 5);
    }

    #[test]
    fn test_visit_callbacks_fire_correctly() {
        // Track which callbacks fired and in what order
        struct CallbackTracker {
            events: Vec<String>,
        }
        impl ObjectVisitor for CallbackTracker {
            fn visit_dict(&mut self, id: Option<ObjectId>, _dict: &HashMap<Name, Object>) {
                self.events
                    .push(format!("dict:{}", id.map_or(0, |i| i.number)));
            }
            fn visit_array(&mut self, id: Option<ObjectId>, _arr: &[Object]) {
                self.events
                    .push(format!("array:{}", id.map_or(0, |i| i.number)));
            }
            fn visit_stream(&mut self, id: Option<ObjectId>) {
                self.events
                    .push(format!("stream:{}", id.map_or(0, |i| i.number)));
            }
            fn visit_reference(&mut self, target: ObjectId) {
                self.events.push(format!("ref:{}", target.number));
            }
            fn visit_primitive(&mut self, id: Option<ObjectId>, _obj: &Object) {
                self.events
                    .push(format!("prim:{}", id.map_or(0, |i| i.number)));
            }
        }

        let pdf = build_minimal_pdf();
        let store = ObjectStore::open(pdf, ParsingMode::Strict).unwrap();

        let mut tracker = CallbackTracker { events: Vec::new() };
        let mut walker = ObjectWalker::new();
        walker.walk(&store, ObjectId::new(1, 0), &mut tracker);

        // Object 1 is a dict (Catalog) containing /Pages 2 0 R
        assert!(tracker.events.contains(&"dict:1".to_string()));
        // Object 2 is a dict (Pages)
        assert!(tracker.events.contains(&"dict:2".to_string()));
    }

    #[test]
    fn test_walk_nonexistent_root() {
        let pdf = build_minimal_pdf();
        let store = ObjectStore::open(pdf, ParsingMode::Lenient).unwrap();

        let mut stats = ObjectStats::default();
        let mut walker = ObjectWalker::new();
        walker.walk(&store, ObjectId::new(999, 0), &mut stats);

        // Nonexistent root should result in 0 visited objects
        // (it gets inserted into visited set but fails to resolve)
        assert_eq!(stats.total_visited, 0);
    }

    // -----------------------------------------------------------------------
    // Upstream-derived object walker tests (cpdf_object_walker_unittest.cpp)
    // -----------------------------------------------------------------------

    /// Upstream: TEST(ObjectWalkerTest, Simple)
    ///
    /// Walking a single simple object should visit exactly that object.
    /// In our BFS walker design, we test with store-based objects.
    #[test]
    fn test_walker_simple_dict() {
        // Build a PDF with a single dict object (no references)
        let mut pdf = Vec::new();
        pdf.extend_from_slice(b"%PDF-1.4\n");
        let obj1_offset = pdf.len();
        pdf.extend_from_slice(b"1 0 obj\n<< /Type /Catalog >>\nendobj\n");
        let xref_offset = pdf.len();
        pdf.extend_from_slice(b"xref\n0 2\n");
        pdf.extend_from_slice(b"0000000000 65535 f \r\n");
        pdf.extend_from_slice(format!("{:010} 00000 n \r\n", obj1_offset).as_bytes());
        pdf.extend_from_slice(b"trailer\n<< /Size 2 /Root 1 0 R >>\n");
        pdf.extend_from_slice(format!("startxref\n{}\n%%EOF", xref_offset).as_bytes());

        let store = ObjectStore::open(pdf, ParsingMode::Strict).unwrap();
        let mut stats = ObjectStats::default();
        let mut walker = ObjectWalker::new();
        walker.walk(&store, ObjectId::new(1, 0), &mut stats);

        // Single dict object, no references to follow
        assert_eq!(walker.visited_count(), 1);
        assert_eq!(stats.dict_count, 1);
    }

    /// Upstream: TEST(ObjectWalkerTest, CombinedObject)
    ///
    /// Walking a dict that contains references to other objects should visit
    /// all reachable objects.
    #[test]
    fn test_walker_combined_object() {
        let pdf = build_minimal_pdf();
        let store = ObjectStore::open(pdf, ParsingMode::Strict).unwrap();

        let mut stats = ObjectStats::default();
        let mut walker = ObjectWalker::new();
        walker.walk(&store, ObjectId::new(1, 0), &mut stats);

        // Object 1 (Catalog dict with /Pages 2 0 R) → follows ref to Object 2
        assert_eq!(walker.visited_count(), 2);
        assert!(walker.visited().contains(&ObjectId::new(1, 0)));
        assert!(walker.visited().contains(&ObjectId::new(2, 0)));
        // Both are dicts
        assert_eq!(stats.dict_count, 2);
    }

    /// Upstream: TEST(ObjectWalkerTest, GetParent)
    ///
    /// In the C++ implementation, GetParent returns the parent object.
    /// In our BFS walker, parent tracking happens through the visitor pattern.
    /// Our walker follows indirect references from dict values but does not
    /// descend into inline arrays. Here we verify that direct dict-to-dict
    /// reference chains are walked correctly.
    #[test]
    fn test_walker_get_parent_tracking() {
        // Build a 2-level deep PDF where references are direct dict values:
        // Obj 1 (Catalog) → /Pages 2 0 R → Obj 2 (Pages)
        let pdf = build_minimal_pdf();
        let store = ObjectStore::open(pdf, ParsingMode::Strict).unwrap();

        // Track visit order
        struct OrderTracker {
            order: Vec<u32>,
        }
        impl ObjectVisitor for OrderTracker {
            fn visit_dict(&mut self, id: Option<ObjectId>, _dict: &HashMap<Name, Object>) {
                if let Some(id) = id {
                    self.order.push(id.number);
                }
            }
            fn visit_array(&mut self, _id: Option<ObjectId>, _arr: &[Object]) {}
            fn visit_stream(&mut self, _id: Option<ObjectId>) {}
            fn visit_reference(&mut self, _target: ObjectId) {}
            fn visit_primitive(&mut self, _id: Option<ObjectId>, _obj: &Object) {}
        }

        let mut tracker = OrderTracker { order: Vec::new() };
        let mut walker = ObjectWalker::new();
        walker.walk(&store, ObjectId::new(1, 0), &mut tracker);

        // Both objects visited, Object 1 first (root), then Object 2
        assert_eq!(walker.visited_count(), 2);
        assert!(tracker.order.contains(&1));
        assert!(tracker.order.contains(&2));
    }

    /// Upstream: TEST(ObjectWalkerTest, SkipWalkIntoCurrentObject)
    ///
    /// In the C++ walker, SkipWalkIntoCurrentObject prevents descending into
    /// the current object's children. Our BFS walker doesn't have this
    /// mechanism, but we test that the max_objects limit achieves similar control.
    #[test]
    fn test_walker_skip_via_max_objects() {
        let pdf = build_minimal_pdf();
        let store = ObjectStore::open(pdf, ParsingMode::Strict).unwrap();

        // Limit to 1 object — should only visit the root
        let mut stats = ObjectStats::default();
        let mut walker = ObjectWalker::with_max_objects(1);
        walker.walk(&store, ObjectId::new(1, 0), &mut stats);

        assert_eq!(walker.visited_count(), 1);
        assert!(walker.visited().contains(&ObjectId::new(1, 0)));
        assert!(!walker.visited().contains(&ObjectId::new(2, 0)));
    }

    /// Upstream: TEST(ObjectWalkerTest, DictionaryKey)
    ///
    /// In the C++ walker, dictionary_key() returns the key name for the current
    /// child being iterated. We verify that dict keys are preserved when
    /// objects are resolved from the store.
    #[test]
    fn test_walker_dictionary_key_tracking() {
        let pdf = build_minimal_pdf();
        let store = ObjectStore::open(pdf, ParsingMode::Strict).unwrap();

        // Resolve the Catalog dict and verify its keys are correct
        let catalog = store.resolve(ObjectId::new(1, 0)).unwrap();
        let dict = catalog.as_dict().unwrap();

        // The Catalog should have /Type and /Pages keys
        assert!(dict.contains_key(&Name::r#type()));
        assert!(dict.contains_key(&Name::pages()));

        // Verify the values are accessible
        let type_val = dict.get(&Name::r#type()).unwrap();
        assert!(type_val.is_name());

        let pages_val = dict.get(&Name::pages()).unwrap();
        assert!(pages_val.is_reference());
        assert_eq!(pages_val.as_reference(), Some(ObjectId::new(2, 0)));
    }
}