Correct the way to count pages and to avoid infinite loop

core/src/fpdfapi/fpdf_parser/fpdf_parser_document.cpp

Index: core/src/fpdfapi/fpdf_parser/fpdf_parser_document.cpp
diff --git a/core/src/fpdfapi/fpdf_parser/fpdf_parser_document.cpp b/core/src/fpdfapi/fpdf_parser/fpdf_parser_document.cpp
index d8091e886193ad5069a95eb121a43fa683c620fb..72ab12d93b312c4b1fddab982539e7463813c8a9 100644
--- a/core/src/fpdfapi/fpdf_parser/fpdf_parser_document.cpp
+++ b/core/src/fpdfapi/fpdf_parser/fpdf_parser_document.cpp
@@ -6,7 +6,42 @@
 
 #include "core/include/fpdfapi/fpdf_parser.h"
 
+#include <set>
+
 #include "core/include/fpdfapi/fpdf_module.h"
+#include "third_party/base/stl_util.h"
+
+namespace {
+
+int CountPages(CPDF_Dictionary* pPages,
+               std::set<CPDF_Dictionary*>* counting_pages) {

[Lei Zhang, 2016/01/14 19:22:39]

counted_pages or encountered_pages?

[Wei Li, 2016/01/14 19:46:00]

Those pages are not counted yet. Chose a textbook classic name :)

+  int count = pPages->GetInteger("Count");
+  if (count > 0 && count < FPDF_PAGE_MAX_NUM) {
+    return count;
+  }
+  CPDF_Array* pKidList = pPages->GetArray("Kids");
+  if (!pKidList) {
+    return 0;
+  }
+  count = 0;
+  for (FX_DWORD i = 0; i < pKidList->GetCount(); i++) {
+    CPDF_Dictionary* pKid = pKidList->GetDict(i);
+    if (!pKid || pdfium::ContainsKey(*counting_pages, pKid)) {
+      continue;
+    }
+    if (!pKid->KeyExist("Kids")) {

[Lei Zhang, 2016/01/14 19:22:39]

Do you want to remove the '!' and flip the if/else blocks?

[Wei Li, 2016/01/14 19:45:59]

Done.

+      count++;
+    } else {
+      // Use |counting_pages| to help detect circular references of pages.
+      ScopedSetInsertion<CPDF_Dictionary*> local_add(counting_pages, pKid);
+      count += CountPages(pKid, counting_pages);
+    }
+  }
+  pPages->SetAtInteger("Count", count);
+  return count;
+}
+
+}  // namespace
 
 CPDF_Document::CPDF_Document(CPDF_Parser* pParser)
     : CPDF_IndirectObjectHolder(pParser) {
@@ -54,7 +89,7 @@ void CPDF_Document::LoadDoc() {
     m_ID1 = pIDArray->GetString(0);
     m_ID2 = pIDArray->GetString(1);
   }
-  m_PageList.SetSize(_GetPageCount());
+  m_PageList.SetSize(RetrievePageCount());
 }
 void CPDF_Document::LoadAsynDoc(CPDF_Dictionary* pLinearized) {
   m_bLinearized = TRUE;
@@ -87,7 +122,7 @@ void CPDF_Document::LoadAsynDoc(CPDF_Dictionary* pLinearized) {
     m_dwFirstPageObjNum = pObjNum->GetInteger();
 }
 void CPDF_Document::LoadPages() {
-  m_PageList.SetSize(_GetPageCount());
+  m_PageList.SetSize(RetrievePageCount());
 }
 CPDF_Document::~CPDF_Document() {
   if (m_pDocPage) {
@@ -256,34 +291,8 @@ int CPDF_Document::GetPageIndex(FX_DWORD objnum) {
 int CPDF_Document::GetPageCount() const {
   return m_PageList.GetSize();
 }
-static int _CountPages(CPDF_Dictionary* pPages, int level) {
-  if (level > 128) {
-    return 0;
-  }
-  int count = pPages->GetInteger("Count");
-  if (count > 0 && count < FPDF_PAGE_MAX_NUM) {
-    return count;
-  }
-  CPDF_Array* pKidList = pPages->GetArray("Kids");
-  if (!pKidList) {
-    return 0;
-  }
-  count = 0;
-  for (FX_DWORD i = 0; i < pKidList->GetCount(); i++) {
-    CPDF_Dictionary* pKid = pKidList->GetDict(i);
-    if (!pKid) {
-      continue;
-    }
-    if (!pKid->KeyExist("Kids")) {
-      count++;
-    } else {
-      count += _CountPages(pKid, level + 1);
-    }
-  }
-  pPages->SetAtInteger("Count", count);
-  return count;
-}
-int CPDF_Document::_GetPageCount() const {
+
+int CPDF_Document::RetrievePageCount() const {
   CPDF_Dictionary* pRoot = GetRoot();
   if (!pRoot) {
     return 0;
@@ -295,8 +304,11 @@ int CPDF_Document::_GetPageCount() const {
   if (!pPages->KeyExist("Kids")) {
     return 1;
   }
-  return _CountPages(pPages, 0);
+  std::set<CPDF_Dictionary*> counting_pages;
+  counting_pages.insert(pPages);
+  return CountPages(pPages, &counting_pages);
 }
+
 FX_BOOL CPDF_Document::IsContentUsedElsewhere(FX_DWORD objnum,
                                               CPDF_Dictionary* pThisPageDict) {
   for (int i = 0; i < m_PageList.GetSize(); i++) {