Properly refcount CPDF_StructElement

Change-Id: Idc9921fe7389abea8e23f07a58fb6e7bfd1c09eb
Reviewed-on: https://pdfium-review.googlesource.com/2433
Reviewed-by: dsinclair <dsinclair@chromium.org>
Commit-Queue: Tom Sepez <tsepez@chromium.org>

4 files changed, 88 insertions, 103 deletions

diff --git a/core/fpdfdoc/doc_tagged.cpp b/core/fpdfdoc/doc_tagged.cpp
index 5feb205542..af5cf85938 100644
--- a/core/fpdfdoc/doc_tagged.cpp
+++ b/core/fpdfdoc/doc_tagged.cpp
@@ -32,6 +32,17 @@ bool IsTagged(const CPDF_Document* pDoc) {
 
 }  // namespace
 
+CPDF_StructKid::CPDF_StructKid()
+    : m_Type(Invalid),
+      m_pDict(nullptr),
+      m_PageObjNum(0),
+      m_RefObjNum(0),
+      m_ContentId(0) {}
+
+CPDF_StructKid::CPDF_StructKid(const CPDF_StructKid& that) = default;
+
+CPDF_StructKid::~CPDF_StructKid() {}
+
 // static
 std::unique_ptr<IPDF_StructTree> IPDF_StructTree::LoadPage(
     const CPDF_Document* pDoc,
@@ -91,52 +102,50 @@ void CPDF_StructTree::LoadPageTree(const CPDF_Dictionary* pPageDict) {
   if (!pParentArray)
     return;
 
-  std::map<CPDF_Dictionary*, CPDF_StructElement*> element_map;
+  std::map<CPDF_Dictionary*, CFX_RetainPtr<CPDF_StructElement>> element_map;
   for (size_t i = 0; i < pParentArray->GetCount(); i++) {
     if (CPDF_Dictionary* pParent = pParentArray->GetDictAt(i))
-      AddPageNode(pParent, element_map);
+      AddPageNode(pParent, &element_map);
   }
 }
 
-CPDF_StructElement* CPDF_StructTree::AddPageNode(
+CFX_RetainPtr<CPDF_StructElement> CPDF_StructTree::AddPageNode(
     CPDF_Dictionary* pDict,
-    std::map<CPDF_Dictionary*, CPDF_StructElement*>& map,
+    std::map<CPDF_Dictionary*, CFX_RetainPtr<CPDF_StructElement>>* map,
     int nLevel) {
   if (nLevel > nMaxRecursion)
     return nullptr;
 
-  auto it = map.find(pDict);
-  if (it != map.end())
+  auto it = map->find(pDict);
+  if (it != map->end())
     return it->second;
 
-  CPDF_StructElement* pElement = new CPDF_StructElement(this, nullptr, pDict);
-  map[pDict] = pElement;
+  auto pElement = pdfium::MakeRetain<CPDF_StructElement>(this, nullptr, pDict);
+  (*map)[pDict] = pElement;
   CPDF_Dictionary* pParent = pDict->GetDictFor("P");
   if (!pParent || pParent->GetStringFor("Type") == "StructTreeRoot") {
-    if (!AddTopLevelNode(pDict, pElement)) {
-      pElement->Release();
-      map.erase(pDict);
-    }
-  } else {
-    CPDF_StructElement* pParentElement = AddPageNode(pParent, map, nLevel + 1);
-    bool bSave = false;
-    for (CPDF_StructKid& kid : *pParentElement->GetKids()) {
-      if (kid.m_Type != CPDF_StructKid::Element)
-        continue;
-      if (kid.m_Element.m_pDict != pDict)
-        continue;
-      kid.m_Element.m_pElement = pElement->Retain();
+    if (!AddTopLevelNode(pDict, pElement))
+      map->erase(pDict);
+    return pElement;
+  }
+
+  CFX_RetainPtr<CPDF_StructElement> pParentElement =
+      AddPageNode(pParent, map, nLevel + 1);
+  bool bSave = false;
+  for (CPDF_StructKid& kid : *pParentElement->GetKids()) {
+    if (kid.m_Type == CPDF_StructKid::Element && kid.m_pDict == pDict) {
+      kid.m_pElement = pElement;
       bSave = true;
     }
-    if (!bSave) {
-      pElement->Release();
-      map.erase(pDict);
-    }
   }
+  if (!bSave)
+    map->erase(pDict);
   return pElement;
 }
-bool CPDF_StructTree::AddTopLevelNode(CPDF_Dictionary* pDict,
-                                      CPDF_StructElement* pElement) {
+
+bool CPDF_StructTree::AddTopLevelNode(
+    CPDF_Dictionary* pDict,
+    const CFX_RetainPtr<CPDF_StructElement>& pElement) {
   CPDF_Object* pObj = m_pTreeRoot->GetDirectObjectFor("K");
   if (!pObj)
     return false;
@@ -144,14 +153,14 @@ bool CPDF_StructTree::AddTopLevelNode(CPDF_Dictionary* pDict,
   if (pObj->IsDictionary()) {
     if (pObj->GetObjNum() != pDict->GetObjNum())
       return false;
-    m_Kids[0].Reset(pElement);
+    m_Kids[0] = pElement;
   }
   if (CPDF_Array* pTopKids = pObj->AsArray()) {
     bool bSave = false;
     for (size_t i = 0; i < pTopKids->GetCount(); i++) {
       CPDF_Reference* pKidRef = ToReference(pTopKids->GetObjectAt(i));
       if (pKidRef && pKidRef->GetRefObjNum() == pDict->GetObjNum()) {
-        m_Kids[i].Reset(pElement);
+        m_Kids[i] = pElement;
         bSave = true;
       }
     }
@@ -164,8 +173,7 @@ bool CPDF_StructTree::AddTopLevelNode(CPDF_Dictionary* pDict,
 CPDF_StructElement::CPDF_StructElement(CPDF_StructTree* pTree,
                                        CPDF_StructElement* pParent,
                                        CPDF_Dictionary* pDict)
-    : m_RefCount(0),
-      m_pTree(pTree),
+    : m_pTree(pTree),
       m_pParent(pParent),
       m_pDict(pDict),
       m_Type(pDict->GetStringFor("S")) {
@@ -198,27 +206,14 @@ int CPDF_StructElement::CountKids() const {
 }
 
 IPDF_StructElement* CPDF_StructElement::GetKidIfElement(int index) const {
-  return m_Kids[index].m_Type == CPDF_StructKid::Element
-             ? m_Kids[index].m_Element.m_pElement
-             : nullptr;
-}
+  if (m_Kids[index].m_Type != CPDF_StructKid::Element)
+    return nullptr;
 
-CPDF_StructElement::~CPDF_StructElement() {
-  for (CPDF_StructKid& kid : m_Kids) {
-    if (kid.m_Type == CPDF_StructKid::Element && kid.m_Element.m_pElement)
-      static_cast<CPDF_StructElement*>(kid.m_Element.m_pElement)->Release();
-  }
+  return m_Kids[index].m_pElement.Get();
 }
 
-CPDF_StructElement* CPDF_StructElement::Retain() {
-  m_RefCount++;
-  return this;
-}
-void CPDF_StructElement::Release() {
-  if (--m_RefCount < 1) {
-    delete this;
-  }
-}
+CPDF_StructElement::~CPDF_StructElement() {}
+
 void CPDF_StructElement::LoadKids(CPDF_Dictionary* pDict) {
   CPDF_Object* pObj = pDict->GetObjectFor("Pg");
   uint32_t PageObjNum = 0;
@@ -253,8 +248,8 @@ void CPDF_StructElement::LoadKid(uint32_t PageObjNum,
       return;
     }
     pKid->m_Type = CPDF_StructKid::PageContent;
-    pKid->m_PageContent.m_ContentId = pKidObj->GetInteger();
-    pKid->m_PageContent.m_PageObjNum = PageObjNum;
+    pKid->m_ContentId = pKidObj->GetInteger();
+    pKid->m_PageObjNum = PageObjNum;
     return;
   }
 
@@ -271,32 +266,26 @@ void CPDF_StructElement::LoadKid(uint32_t PageObjNum,
       return;
     }
     pKid->m_Type = CPDF_StructKid::StreamContent;
-    if (CPDF_Reference* pRef = ToReference(pKidDict->GetObjectFor("Stm"))) {
-      pKid->m_StreamContent.m_RefObjNum = pRef->GetRefObjNum();
-    } else {
-      pKid->m_StreamContent.m_RefObjNum = 0;
-    }
-    pKid->m_StreamContent.m_PageObjNum = PageObjNum;
-    pKid->m_StreamContent.m_ContentId = pKidDict->GetIntegerFor("MCID");
+    CPDF_Reference* pRef = ToReference(pKidDict->GetObjectFor("Stm"));
+    pKid->m_RefObjNum = pRef ? pRef->GetRefObjNum() : 0;
+    pKid->m_PageObjNum = PageObjNum;
+    pKid->m_ContentId = pKidDict->GetIntegerFor("MCID");
   } else if (type == "OBJR") {
     if (m_pTree->m_pPage && m_pTree->m_pPage->GetObjNum() != PageObjNum) {
       return;
     }
     pKid->m_Type = CPDF_StructKid::Object;
-    if (CPDF_Reference* pObj = ToReference(pKidDict->GetObjectFor("Obj"))) {
-      pKid->m_Object.m_RefObjNum = pObj->GetRefObjNum();
-    } else {
-      pKid->m_Object.m_RefObjNum = 0;
-    }
-    pKid->m_Object.m_PageObjNum = PageObjNum;
+    CPDF_Reference* pObj = ToReference(pKidDict->GetObjectFor("Obj"));
+    pKid->m_RefObjNum = pObj ? pObj->GetRefObjNum() : 0;
+    pKid->m_PageObjNum = PageObjNum;
   } else {
     pKid->m_Type = CPDF_StructKid::Element;
-    pKid->m_Element.m_pDict = pKidDict;
+    pKid->m_pDict = pKidDict;
     if (!m_pTree->m_pPage) {
-      pKid->m_Element.m_pElement =
-          new CPDF_StructElement(m_pTree, this, pKidDict);
+      pKid->m_pElement =
+          pdfium::MakeRetain<CPDF_StructElement>(m_pTree, this, pKidDict);
     } else {
-      pKid->m_Element.m_pElement = nullptr;
+      pKid->m_pElement = nullptr;
     }
   }
 }
diff --git a/core/fpdfdoc/fpdf_tagged.h b/core/fpdfdoc/fpdf_tagged.h
index 3a1e2231fc..fbbb49f856 100644
--- a/core/fpdfdoc/fpdf_tagged.h
+++ b/core/fpdfdoc/fpdf_tagged.h
@@ -21,16 +21,17 @@ class IPDF_StructTree {
       const CPDF_Document* pDoc,
       const CPDF_Dictionary* pPageDict);
 
-  virtual ~IPDF_StructTree() {}
 
   virtual int CountTopElements() const = 0;
   virtual IPDF_StructElement* GetTopElement(int i) const = 0;
+
+ protected:
+  friend std::default_delete<IPDF_StructTree>;
+  virtual ~IPDF_StructTree() {}
 };
 
 class IPDF_StructElement {
  public:
-  virtual ~IPDF_StructElement() {}
-
   virtual IPDF_StructTree* GetTree() const = 0;
   virtual const CFX_ByteString& GetType() const = 0;
   virtual IPDF_StructElement* GetParent() const = 0;
@@ -66,6 +67,9 @@ class IPDF_StructElement {
                          int default_value,
                          bool bInheritable = false,
                          int subindex = -1) = 0;
+
+ protected:
+  virtual ~IPDF_StructElement() {}
 };
 
 #endif  // CORE_FPDFDOC_FPDF_TAGGED_H_
diff --git a/core/fpdfdoc/tagged_int.h b/core/fpdfdoc/tagged_int.h
index 6a24700502..ce246023df 100644
--- a/core/fpdfdoc/tagged_int.h
+++ b/core/fpdfdoc/tagged_int.h
@@ -18,27 +18,17 @@
 class CPDF_StructElement;
 
 struct CPDF_StructKid {
+  CPDF_StructKid();
+  CPDF_StructKid(const CPDF_StructKid& that);
+  ~CPDF_StructKid();
+
   enum { Invalid, Element, PageContent, StreamContent, Object } m_Type;
 
-  union {
-    struct {
-      CPDF_StructElement* m_pElement;
-      CPDF_Dictionary* m_pDict;
-    } m_Element;
-    struct {
-      uint32_t m_PageObjNum;
-      uint32_t m_ContentId;
-    } m_PageContent;
-    struct {
-      uint32_t m_PageObjNum;
-      uint32_t m_ContentId;
-      uint32_t m_RefObjNum;
-    } m_StreamContent;
-    struct {
-      uint32_t m_PageObjNum;
-      uint32_t m_RefObjNum;
-    } m_Object;
-  };
+  CFX_RetainPtr<CPDF_StructElement> m_pElement;  // For Element.
+  CPDF_Dictionary* m_pDict;                      // For Element.
+  uint32_t m_PageObjNum;  // For PageContent, StreamContent, Object.
+  uint32_t m_RefObjNum;   // For StreamContent, Object.
+  uint32_t m_ContentId;   // For PageContent, StreamContent.
 };
 
 class CPDF_StructTree final : public IPDF_StructTree {
@@ -51,11 +41,12 @@ class CPDF_StructTree final : public IPDF_StructTree {
   IPDF_StructElement* GetTopElement(int i) const override;
 
   void LoadPageTree(const CPDF_Dictionary* pPageDict);
-  CPDF_StructElement* AddPageNode(
+  CFX_RetainPtr<CPDF_StructElement> AddPageNode(
       CPDF_Dictionary* pElement,
-      std::map<CPDF_Dictionary*, CPDF_StructElement*>& map,
+      std::map<CPDF_Dictionary*, CFX_RetainPtr<CPDF_StructElement>>* map,
       int nLevel = 0);
-  bool AddTopLevelNode(CPDF_Dictionary* pDict, CPDF_StructElement* pElement);
+  bool AddTopLevelNode(CPDF_Dictionary* pDict,
+                       const CFX_RetainPtr<CPDF_StructElement>& pElement);
 
  protected:
   const CPDF_Dictionary* const m_pTreeRoot;
@@ -66,11 +57,11 @@ class CPDF_StructTree final : public IPDF_StructTree {
   friend class CPDF_StructElement;
 };
 
-class CPDF_StructElement final : public IPDF_StructElement {
+class CPDF_StructElement final : public CFX_Retainable,
+                                 public IPDF_StructElement {
  public:
-  CPDF_StructElement(CPDF_StructTree* pTree,
-                     CPDF_StructElement* pParent,
-                     CPDF_Dictionary* pDict);
+  template <typename T, typename... Args>
+  friend CFX_RetainPtr<T> pdfium::MakeRetain(Args&&... args);
 
   // IPDF_StructElement
   IPDF_StructTree* GetTree() const override;
@@ -112,13 +103,12 @@ class CPDF_StructElement final : public IPDF_StructElement {
                        bool bInheritable,
                        int subindex);
 
-  CPDF_StructElement* Retain();
-  void Release();
-
- protected:
+ private:
+  CPDF_StructElement(CPDF_StructTree* pTree,
+                     CPDF_StructElement* pParent,
+                     CPDF_Dictionary* pDict);
   ~CPDF_StructElement() override;
 
-  int m_RefCount;
   CPDF_StructTree* const m_pTree;
   CPDF_StructElement* const m_pParent;
   CPDF_Dictionary* const m_pDict;
diff --git a/fpdfsdk/fpdf_structtree.cpp b/fpdfsdk/fpdf_structtree.cpp
index 26ded3d460..5a922a1c3b 100644
--- a/fpdfsdk/fpdf_structtree.cpp
+++ b/fpdfsdk/fpdf_structtree.cpp
@@ -4,6 +4,8 @@
 
 #include "public/fpdf_structtree.h"
 
+#include <memory>
+
 #include "core/fpdfapi/page/cpdf_page.h"
 #include "core/fpdfapi/parser/cpdf_dictionary.h"
 #include "core/fpdfdoc/fpdf_tagged.h"
@@ -30,7 +32,7 @@ DLLEXPORT FPDF_STRUCTTREE STDCALL FPDF_StructTree_GetForPage(FPDF_PAGE page) {
 }
 
 DLLEXPORT void STDCALL FPDF_StructTree_Close(FPDF_STRUCTTREE struct_tree) {
-  delete ToStructTree(struct_tree);
+  std::unique_ptr<IPDF_StructTree>(ToStructTree(struct_tree));
 }
 
 DLLEXPORT int STDCALL