RefCount CPDF_IccProfile all the time

Make the IccProfile track its stream so that it has a
proper key with which to purge the docpagedata map.

Change-Id: Ib05ebc1afb828f1f5e5df62a1a33a1bfdecf507d
Reviewed-on: https://pdfium-review.googlesource.com/3619
Commit-Queue: Tom Sepez <tsepez@chromium.org>
Reviewed-by: Lei Zhang <thestig@chromium.org>

5 files changed, 44 insertions, 49 deletions

diff --git a/core/fpdfapi/page/cpdf_colorspace.cpp b/core/fpdfapi/page/cpdf_colorspace.cpp
index 913c9c728d..500de26c56 100644
--- a/core/fpdfapi/page/cpdf_colorspace.cpp
+++ b/core/fpdfapi/page/cpdf_colorspace.cpp
@@ -166,7 +166,7 @@ class CPDF_ICCBasedCS : public CPDF_ColorSpace {
   void PopulateRanges(CPDF_Dictionary* pDict);
 
   CFX_MaybeOwned<CPDF_ColorSpace> m_pAlterCS;
-  CPDF_IccProfile* m_pProfile;
+  CFX_RetainPtr<CPDF_IccProfile> m_pProfile;
   uint8_t* m_pCache;
   float* m_pRanges;
 };
@@ -849,8 +849,11 @@ CPDF_ICCBasedCS::CPDF_ICCBasedCS(CPDF_Document* pDoc)
 CPDF_ICCBasedCS::~CPDF_ICCBasedCS() {
   FX_Free(m_pCache);
   FX_Free(m_pRanges);
-  if (m_pProfile && m_pDocument)
-    m_pDocument->GetPageData()->ReleaseIccProfile(m_pProfile);
+  if (m_pProfile && m_pDocument) {
+    CPDF_Stream* pStream = m_pProfile->GetStream();
+    m_pProfile.Reset();  // Give up our reference first.
+    m_pDocument->GetPageData()->MaybePurgeIccProfile(pStream);
+  }
 }
 
 bool CPDF_ICCBasedCS::v_Load(CPDF_Document* pDoc, CPDF_Array* pArray) {
diff --git a/core/fpdfapi/page/cpdf_docpagedata.cpp b/core/fpdfapi/page/cpdf_docpagedata.cpp
index 5789a531f0..a57c42d3cd 100644
--- a/core/fpdfapi/page/cpdf_docpagedata.cpp
+++ b/core/fpdfapi/page/cpdf_docpagedata.cpp
@@ -82,11 +82,7 @@ void CPDF_DocPageData::Clear(bool bForceRelease) {
 
   for (auto it = m_IccProfileMap.begin(); it != m_IccProfileMap.end();) {
     auto curr_it = it++;
-    CPDF_CountedIccProfile* ipData = curr_it->second;
-    if (!ipData->get())
-      continue;
-
-    if (bForceRelease || ipData->use_count() < 2) {
+    if (bForceRelease || curr_it->second->HasOneRef()) {
       for (auto hash_it = m_HashProfileMap.begin();
            hash_it != m_HashProfileMap.end(); ++hash_it) {
         if (curr_it->first == hash_it->second) {
@@ -94,8 +90,6 @@ void CPDF_DocPageData::Clear(bool bForceRelease) {
           break;
         }
       }
-      delete ipData->get();
-      delete ipData;
       m_IccProfileMap.erase(curr_it);
     }
   }
@@ -405,50 +399,40 @@ void CPDF_DocPageData::MaybePurgeImage(uint32_t dwStreamObjNum) {
     m_ImageMap.erase(it);
 }
 
-CPDF_IccProfile* CPDF_DocPageData::GetIccProfile(
-    CPDF_Stream* pIccProfileStream) {
-  if (!pIccProfileStream)
+CFX_RetainPtr<CPDF_IccProfile> CPDF_DocPageData::GetIccProfile(
+    CPDF_Stream* pProfileStream) {
+  if (!pProfileStream)
     return nullptr;
 
-  auto it = m_IccProfileMap.find(pIccProfileStream);
+  auto it = m_IccProfileMap.find(pProfileStream);
   if (it != m_IccProfileMap.end())
-    return it->second->AddRef();
+    return it->second;
+
+  CPDF_StreamAcc accessor;
+  accessor.LoadAllData(pProfileStream, false);
 
-  CPDF_StreamAcc stream;
-  stream.LoadAllData(pIccProfileStream, false);
   uint8_t digest[20];
-  CRYPT_SHA1Generate(stream.GetData(), stream.GetSize(), digest);
+  CRYPT_SHA1Generate(accessor.GetData(), accessor.GetSize(), digest);
+
   CFX_ByteString bsDigest(digest, 20);
   auto hash_it = m_HashProfileMap.find(bsDigest);
   if (hash_it != m_HashProfileMap.end()) {
     auto it_copied_stream = m_IccProfileMap.find(hash_it->second);
     if (it_copied_stream != m_IccProfileMap.end())
-      return it_copied_stream->second->AddRef();
+      return it_copied_stream->second;
   }
-  CPDF_CountedIccProfile* ipData = new CPDF_CountedIccProfile(
-      pdfium::MakeUnique<CPDF_IccProfile>(stream.GetData(), stream.GetSize()));
-  m_IccProfileMap[pIccProfileStream] = ipData;
-  m_HashProfileMap[bsDigest] = pIccProfileStream;
-  return ipData->AddRef();
+  auto pProfile = pdfium::MakeRetain<CPDF_IccProfile>(
+      pProfileStream, accessor.GetData(), accessor.GetSize());
+  m_IccProfileMap[pProfileStream] = pProfile;
+  m_HashProfileMap[bsDigest] = pProfileStream;
+  return pProfile;
 }
 
-void CPDF_DocPageData::ReleaseIccProfile(const CPDF_IccProfile* pIccProfile) {
-  ASSERT(pIccProfile);
-
-  for (auto it = m_IccProfileMap.begin(); it != m_IccProfileMap.end(); ++it) {
-    CPDF_CountedIccProfile* profile = it->second;
-    if (profile->get() != pIccProfile)
-      continue;
-
-    profile->RemoveRef();
-    if (profile->use_count() > 1)
-      continue;
-    // We have item only in m_IccProfileMap cache. Clean it.
-    delete profile->get();
-    delete profile;
+void CPDF_DocPageData::MaybePurgeIccProfile(CPDF_Stream* pProfileStream) {
+  ASSERT(pProfileStream);
+  auto it = m_IccProfileMap.find(pProfileStream);
+  if (it != m_IccProfileMap.end() && it->second->HasOneRef())
     m_IccProfileMap.erase(it);
-    return;
-  }
 }
 
 CPDF_StreamAcc* CPDF_DocPageData::GetFontFileStreamAcc(
diff --git a/core/fpdfapi/page/cpdf_docpagedata.h b/core/fpdfapi/page/cpdf_docpagedata.h
index d95c3c788d..24e0006efd 100644
--- a/core/fpdfapi/page/cpdf_docpagedata.h
+++ b/core/fpdfapi/page/cpdf_docpagedata.h
@@ -50,8 +50,8 @@ class CPDF_DocPageData {
   CFX_RetainPtr<CPDF_Image> GetImage(uint32_t dwStreamObjNum);
   void MaybePurgeImage(uint32_t dwStreamObjNum);
 
-  CPDF_IccProfile* GetIccProfile(CPDF_Stream* pIccProfileStream);
-  void ReleaseIccProfile(const CPDF_IccProfile* pIccProfile);
+  CFX_RetainPtr<CPDF_IccProfile> GetIccProfile(CPDF_Stream* pProfileStream);
+  void MaybePurgeIccProfile(CPDF_Stream* pProfileStream);
 
   CPDF_StreamAcc* GetFontFileStreamAcc(CPDF_Stream* pFontStream);
   void ReleaseFontFileStreamAcc(const CPDF_Stream* pFontStream);
@@ -61,7 +61,6 @@ class CPDF_DocPageData {
 
  private:
   using CPDF_CountedFont = CPDF_CountedObject<CPDF_Font>;
-  using CPDF_CountedIccProfile = CPDF_CountedObject<CPDF_IccProfile>;
   using CPDF_CountedStreamAcc = CPDF_CountedObject<CPDF_StreamAcc>;
 
   CPDF_ColorSpace* GetColorSpaceImpl(CPDF_Object* pCSObj,
@@ -74,7 +73,7 @@ class CPDF_DocPageData {
   std::map<const CPDF_Object*, CPDF_CountedColorSpace*> m_ColorSpaceMap;
   std::map<const CPDF_Stream*, CPDF_CountedStreamAcc*> m_FontFileMap;
   std::map<const CPDF_Dictionary*, CPDF_CountedFont*> m_FontMap;
-  std::map<const CPDF_Stream*, CPDF_CountedIccProfile*> m_IccProfileMap;
+  std::map<const CPDF_Stream*, CFX_RetainPtr<CPDF_IccProfile>> m_IccProfileMap;
   std::map<uint32_t, CFX_RetainPtr<CPDF_Image>> m_ImageMap;
   std::map<const CPDF_Object*, CPDF_CountedPattern*> m_PatternMap;
 };
diff --git a/core/fpdfapi/page/fpdf_page_colors.cpp b/core/fpdfapi/page/fpdf_page_colors.cpp
index ffb8da3a10..a59ca37d02 100644
--- a/core/fpdfapi/page/fpdf_page_colors.cpp
+++ b/core/fpdfapi/page/fpdf_page_colors.cpp
@@ -218,8 +218,10 @@ void CPDF_DeviceCS::TranslateImageLine(uint8_t* pDestBuf,
   }
 }
 
-CPDF_IccProfile::CPDF_IccProfile(const uint8_t* pData, uint32_t dwSize)
-    : m_bsRGB(DetectSRGB(pData, dwSize)) {
+CPDF_IccProfile::CPDF_IccProfile(CPDF_Stream* pStream,
+                                 const uint8_t* pData,
+                                 uint32_t dwSize)
+    : m_bsRGB(DetectSRGB(pData, dwSize)), m_pStream(pStream) {
   if (m_bsRGB) {
     m_nSrcComponents = 3;
     return;
diff --git a/core/fpdfapi/page/pageint.h b/core/fpdfapi/page/pageint.h
index 9637155343..a7a00f02ee 100644
--- a/core/fpdfapi/page/pageint.h
+++ b/core/fpdfapi/page/pageint.h
@@ -12,11 +12,13 @@
 
 #include "core/fpdfapi/page/cpdf_colorspace.h"
 #include "core/fpdfapi/page/cpdf_countedobject.h"
+#include "core/fxcrt/cfx_retain_ptr.h"
 
 class CPDF_ExpIntFunc;
 class CPDF_Pattern;
 class CPDF_SampledFunc;
 class CPDF_StitchFunc;
+class CPDF_Stream;
 class CPDF_StreamAcc;
 
 class CPDF_Function {
@@ -134,11 +136,12 @@ class CPDF_StitchFunc : public CPDF_Function {
   static const uint32_t kRequiredNumInputs = 1;
 };
 
-class CPDF_IccProfile {
+class CPDF_IccProfile : public CFX_Retainable {
  public:
-  CPDF_IccProfile(const uint8_t* pData, uint32_t dwSize);
-  ~CPDF_IccProfile();
+  template <typename T, typename... Args>
+  friend CFX_RetainPtr<T> pdfium::MakeRetain(Args&&... args);
 
+  CPDF_Stream* GetStream() const { return m_pStream; }
   bool IsValid() const { return IsSRGB() || IsSupported(); }
   bool IsSRGB() const { return m_bsRGB; }
   bool IsSupported() const { return !!m_pTransform; }
@@ -146,7 +149,11 @@ class CPDF_IccProfile {
   uint32_t GetComponents() const { return m_nSrcComponents; }
 
  private:
+  CPDF_IccProfile(CPDF_Stream* pStream, const uint8_t* pData, uint32_t dwSize);
+  ~CPDF_IccProfile() override;
+
   const bool m_bsRGB;
+  CPDF_Stream* const m_pStream;
   void* m_pTransform = nullptr;
   uint32_t m_nSrcComponents = 0;
 };