1 files changed, 96 insertions, 0 deletions
diff --git a/core/fpdfapi/parser/object_tree_traversal_util_embeddertest.cpp b/core/fpdfapi/parser/object_tree_traversal_util_embeddertest.cpp
new file mode 100644
index 000000000..c90c77e41
--- /dev/null
+++ b/core/fpdfapi/parser/object_tree_traversal_util_embeddertest.cpp
@@ -0,0 +1,96 @@
+// Copyright 2023 The PDFium Authors
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#include "core/fpdfapi/parser/object_tree_traversal_util.h"
+
+#include <stdint.h>
+
+#include <set>
+
+#include "core/fpdfapi/parser/cpdf_document.h"
+#include "testing/embedder_test.h"
+#include "testing/gmock/include/gmock/gmock.h"
+#include "testing/gtest/include/gtest/gtest.h"
+
+using testing::UnorderedElementsAreArray;
+using ObjectTreeTraversalUtilEmbedderTest = EmbedderTest;
+
+namespace {
+
+CPDF_Document* GetCPDFDocument(FPDF_DOCUMENT document) {
+  // This is cheating slightly to avoid a layering violation, since this file
+  // cannot include fpdfsdk/cpdfsdk_helpers.h to get access to
+  // CPDFDocumentFromFPDFDocument().
+  return reinterpret_cast<CPDF_Document*>((document));
+}
+
+}  // namespace
+
+TEST_F(ObjectTreeTraversalUtilEmbedderTest, GetObjectsWithReferencesBasic) {
+  ASSERT_TRUE(OpenDocument("hello_world.pdf"));
+  CPDF_Document* doc = GetCPDFDocument(document());
+  std::set<uint32_t> referenced_objects = GetObjectsWithReferences(doc);
+  EXPECT_THAT(referenced_objects,
+              UnorderedElementsAreArray({1, 2, 3, 4, 5, 6}));
+}
+
+TEST_F(ObjectTreeTraversalUtilEmbedderTest, GetObjectsWithReferencesNewDoc) {
+  ScopedFPDFDocument new_doc(FPDF_CreateNewDocument());
+  CPDF_Document* doc = GetCPDFDocument(new_doc.get());
+  std::set<uint32_t> referenced_objects = GetObjectsWithReferences(doc);
+  // Empty documents have a catalog and an empty pages object.
+  EXPECT_THAT(referenced_objects, UnorderedElementsAreArray({1, 2}));
+}
+
+TEST_F(ObjectTreeTraversalUtilEmbedderTest,
+       GetObjectsWithReferencesCircularRefs) {
+  ASSERT_TRUE(OpenDocument("circular_viewer_ref.pdf"));
+  CPDF_Document* doc = GetCPDFDocument(document());
+  std::set<uint32_t> referenced_objects = GetObjectsWithReferences(doc);
+  // The trailer points at a catalog, and the catalog only references itself.
+  EXPECT_THAT(referenced_objects, UnorderedElementsAreArray({1}));
+}
+
+TEST_F(ObjectTreeTraversalUtilEmbedderTest,
+       GetObjectsWithReferencesCrossRefStream) {
+  ASSERT_TRUE(OpenDocument("bug_1399.pdf"));
+  CPDF_Document* doc = GetCPDFDocument(document());
+  std::set<uint32_t> referenced_objects = GetObjectsWithReferences(doc);
+  // The trailer is the dictionary inside /XRef object 16 0. Note that it
+  // references object 3 0, but the rest of the document does not.
+  EXPECT_THAT(referenced_objects,
+              UnorderedElementsAreArray({1, 2, 3, 4, 5, 12, 13, 14, 16}));
+}
+
+TEST_F(ObjectTreeTraversalUtilEmbedderTest,
+       GetObjectsWithMultipleReferencesBasic) {
+  ASSERT_TRUE(OpenDocument("hello_world.pdf"));
+  CPDF_Document* doc = GetCPDFDocument(document());
+  std::set<uint32_t> referenced_objects = GetObjectsWithMultipleReferences(doc);
+  EXPECT_TRUE(referenced_objects.empty());
+}
+
+TEST_F(ObjectTreeTraversalUtilEmbedderTest,
+       GetObjectsWithMultipleReferencesNewDoc) {
+  ScopedFPDFDocument new_doc(FPDF_CreateNewDocument());
+  CPDF_Document* doc = GetCPDFDocument(new_doc.get());
+  std::set<uint32_t> referenced_objects = GetObjectsWithMultipleReferences(doc);
+  EXPECT_TRUE(referenced_objects.empty());
+}
+
+TEST_F(ObjectTreeTraversalUtilEmbedderTest,
+       GetObjectsWithMultipleReferencesCircularRefs) {
+  ASSERT_TRUE(OpenDocument("circular_viewer_ref.pdf"));
+  CPDF_Document* doc = GetCPDFDocument(document());
+  std::set<uint32_t> referenced_objects = GetObjectsWithMultipleReferences(doc);
+  EXPECT_TRUE(referenced_objects.empty());
+}
+
+TEST_F(ObjectTreeTraversalUtilEmbedderTest,
+       GetObjectsWithMultipleReferencesSharedObjects) {
+  ASSERT_TRUE(OpenDocument("hello_world_2_pages.pdf"));
+  CPDF_Document* doc = GetCPDFDocument(document());
+  std::set<uint32_t> referenced_objects = GetObjectsWithMultipleReferences(doc);
+  EXPECT_THAT(referenced_objects, UnorderedElementsAreArray({5, 6, 7}));
+}