Correct the way to count pages and to avoid infinite loop

core/src/fpdfapi/fpdf_parser/fpdf_parser_document.cpp

 // Copyright 2014 PDFium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
 // Original code copyright 2014 Foxit Software Inc. http://www.foxitsoftware.com
 
 #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*>* visited_pages) {
+  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(*visited_pages, pKid)) {
+      continue;
+    }
+    if (pKid->KeyExist("Kids")) {
+      // Use |visited_pages| to help detect circular references of pages.
+      ScopedSetInsertion<CPDF_Dictionary*> local_add(visited_pages, pKid);
+      count += CountPages(pKid, visited_pages);
+    } else {
+      // This page is a leaf node.
+      count++;
+    }
+  }
+  pPages->SetAtInteger("Count", count);
+  return count;
+}
+
+}  // namespace
 
 CPDF_Document::CPDF_Document(CPDF_Parser* pParser)
     : CPDF_IndirectObjectHolder(pParser) {
   ASSERT(pParser);
   m_pRootDict = NULL;
   m_pInfoDict = NULL;
   m_bLinearized = FALSE;
   m_dwFirstPageNo = 0;
   m_dwFirstPageObjNum = 0;
   m_pDocPage = CPDF_ModuleMgr::Get()->GetPageModule()->CreateDocData(this);
@@ skipping 27 matching lines @@
   CPDF_Object* pInfoObj =
       GetIndirectObject(m_pParser->GetInfoObjNum(), nullptr);
   if (pInfoObj) {
     m_pInfoDict = pInfoObj->GetDict();
   }
   CPDF_Array* pIDArray = m_pParser->GetIDArray();
   if (pIDArray) {
     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;
   m_LastObjNum = m_pParser->GetLastObjNum();
   CPDF_Object* pIndirectObj =
       GetIndirectObject(m_pParser->GetRootObjNum(), nullptr);
   m_pRootDict = pIndirectObj ? pIndirectObj->GetDict() : nullptr;
   if (!m_pRootDict) {
     return;
   }
@@ skipping 12 matching lines @@
   m_PageList.SetSize(dwPageCount);
   CPDF_Object* pNo = pLinearized->GetElement("P");
   if (ToNumber(pNo))
     m_dwFirstPageNo = pNo->GetInteger();
 
   CPDF_Object* pObjNum = pLinearized->GetElement("O");
   if (ToNumber(pObjNum))
     m_dwFirstPageObjNum = pObjNum->GetInteger();
 }
 void CPDF_Document::LoadPages() {
-  m_PageList.SetSize(_GetPageCount());
+  m_PageList.SetSize(RetrievePageCount());
 }
 CPDF_Document::~CPDF_Document() {
   if (m_pDocPage) {
     CPDF_ModuleMgr::Get()->GetPageModule()->ReleaseDoc(this);
     CPDF_ModuleMgr::Get()->GetPageModule()->ClearStockFont(this);
   }
   if (m_pDocRender) {
     CPDF_ModuleMgr::Get()->GetRenderModule()->DestroyDocData(m_pDocRender);
   }
 }
@@ skipping 148 matching lines @@
   CPDF_Dictionary* pPages = pRoot->GetDict("Pages");
   if (!pPages) {
     return -1;
   }
   int index = 0;
   return _FindPageIndex(pPages, skip_count, objnum, index);
 }
 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;
   }
   CPDF_Dictionary* pPages = pRoot->GetDict("Pages");
   if (!pPages) {
     return 0;
   }
   if (!pPages->KeyExist("Kids")) {
     return 1;
   }
-  return _CountPages(pPages, 0);
+  std::set<CPDF_Dictionary*> visited_pages;
+  visited_pages.insert(pPages);
+  return CountPages(pPages, &visited_pages);
 }
+
 FX_BOOL CPDF_Document::IsContentUsedElsewhere(FX_DWORD objnum,
                                               CPDF_Dictionary* pThisPageDict) {
   for (int i = 0; i < m_PageList.GetSize(); i++) {
     CPDF_Dictionary* pPageDict = GetPage(i);
     if (pPageDict == pThisPageDict) {
       continue;
     }
     CPDF_Object* pContents =
         pPageDict ? pPageDict->GetElement("Contents") : NULL;
     if (!pContents) {
@@ skipping 37 matching lines @@
 void CPDF_Document::ClearPageData() {
   if (m_pDocPage) {
     CPDF_ModuleMgr::Get()->GetPageModule()->ClearDoc(this);
   }
 }
 void CPDF_Document::ClearRenderData() {
   if (m_pDocRender) {
     CPDF_ModuleMgr::Get()->GetRenderModule()->ClearDocData(m_pDocRender);
   }
 }