3 files changed, 31 insertions, 35 deletions

diff --git a/core/fpdfapi/page/cpdf_generalstate.h b/core/fpdfapi/page/cpdf_generalstate.h
index ddc33ac800..1c5ddee4e0 100644
--- a/core/fpdfapi/page/cpdf_generalstate.h
+++ b/core/fpdfapi/page/cpdf_generalstate.h
@@ -7,6 +7,7 @@
 #ifndef CORE_FPDFAPI_PAGE_CPDF_GENERALSTATE_H_
 #define CORE_FPDFAPI_PAGE_CPDF_GENERALSTATE_H_
 
+#include "core/fxcrt/cfx_shared_copy_on_write.h"
 #include "core/fxcrt/fx_basic.h"
 #include "core/fxcrt/fx_coordinates.h"
 #include "core/fxge/fx_dib.h"
diff --git a/core/fpdfapi/page/cpdf_image.cpp b/core/fpdfapi/page/cpdf_image.cpp
index 3de0511771..4fc1c344ef 100644
--- a/core/fpdfapi/page/cpdf_image.cpp
+++ b/core/fpdfapi/page/cpdf_image.cpp
@@ -157,7 +157,7 @@ void CPDF_Image::SetJpegImageInline(
   m_pStream->InitStream(&(data[0]), size, std::move(pDict));
 }
 
-void CPDF_Image::SetImage(const CFX_DIBitmap* pBitmap) {
+void CPDF_Image::SetImage(const CFX_RetainPtr<CFX_DIBitmap>& pBitmap) {
   int32_t BitmapWidth = pBitmap->GetWidth();
   int32_t BitmapHeight = pBitmap->GetHeight();
   if (BitmapWidth < 1 || BitmapHeight < 1)
@@ -248,7 +248,7 @@ void CPDF_Image::SetImage(const CFX_DIBitmap* pBitmap) {
     bCopyWithoutAlpha = false;
   }
 
-  std::unique_ptr<CFX_DIBitmap> pMaskBitmap;
+  CFX_RetainPtr<CFX_DIBitmap> pMaskBitmap;
   if (pBitmap->HasAlpha())
     pMaskBitmap = pBitmap->CloneAlphaMask();
 
@@ -320,28 +320,25 @@ void CPDF_Image::SetImage(const CFX_DIBitmap* pBitmap) {
   FX_Free(dest_buf);
 }
 
-void CPDF_Image::ResetCache(CPDF_Page* pPage, const CFX_DIBitmap* pBitmap) {
+void CPDF_Image::ResetCache(CPDF_Page* pPage,
+                            const CFX_RetainPtr<CFX_DIBitmap>& pBitmap) {
   pPage->GetRenderCache()->ResetBitmap(m_pStream.Get(), pBitmap);
 }
 
-std::unique_ptr<CFX_DIBSource> CPDF_Image::LoadDIBSource() const {
-  auto source = pdfium::MakeUnique<CPDF_DIBSource>();
+CFX_RetainPtr<CFX_DIBSource> CPDF_Image::LoadDIBSource() const {
+  auto source = pdfium::MakeRetain<CPDF_DIBSource>();
   if (!source->Load(m_pDocument, m_pStream.Get()))
     return nullptr;
 
-  return std::move(source);
+  return source;
 }
 
-CFX_DIBSource* CPDF_Image::DetachBitmap() {
-  CFX_DIBSource* pBitmap = m_pDIBSource;
-  m_pDIBSource = nullptr;
-  return pBitmap;
+CFX_RetainPtr<CFX_DIBSource> CPDF_Image::DetachBitmap() {
+  return std::move(m_pDIBSource);
 }
 
-CFX_DIBSource* CPDF_Image::DetachMask() {
-  CFX_DIBSource* pBitmap = m_pMask;
-  m_pMask = nullptr;
-  return pBitmap;
+CFX_RetainPtr<CFX_DIBSource> CPDF_Image::DetachMask() {
+  return std::move(m_pMask);
 }
 
 bool CPDF_Image::StartLoadDIBSource(CPDF_Dictionary* pFormResource,
@@ -349,35 +346,33 @@ bool CPDF_Image::StartLoadDIBSource(CPDF_Dictionary* pFormResource,
                                     bool bStdCS,
                                     uint32_t GroupFamily,
                                     bool bLoadMask) {
-  auto source = pdfium::MakeUnique<CPDF_DIBSource>();
+  auto source = pdfium::MakeRetain<CPDF_DIBSource>();
   int ret = source->StartLoadDIBSource(m_pDocument, m_pStream.Get(), true,
                                        pFormResource, pPageResource, bStdCS,
                                        GroupFamily, bLoadMask);
-  if (ret == 2) {
-    m_pDIBSource = source.release();
-    return true;
-  }
   if (!ret) {
-    m_pDIBSource = nullptr;
+    m_pDIBSource.Reset();
     return false;
   }
+  m_pDIBSource = source;
+  if (ret == 2)
+    return true;
+
   m_pMask = source->DetachMask();
   m_MatteColor = source->GetMatteColor();
-  m_pDIBSource = source.release();
   return false;
 }
 
 bool CPDF_Image::Continue(IFX_Pause* pPause) {
-  CPDF_DIBSource* pSource = static_cast<CPDF_DIBSource*>(m_pDIBSource);
+  CFX_RetainPtr<CPDF_DIBSource> pSource = m_pDIBSource.As<CPDF_DIBSource>();
   int ret = pSource->ContinueLoadDIBSource(pPause);
-  if (ret == 2) {
-    return true;
-  }
   if (!ret) {
-    delete m_pDIBSource;
-    m_pDIBSource = nullptr;
+    m_pDIBSource.Reset();
     return false;
   }
+  if (ret == 2)
+    return true;
+
   m_pMask = pSource->DetachMask();
   m_MatteColor = pSource->GetMatteColor();
   return false;
diff --git a/core/fpdfapi/page/cpdf_image.h b/core/fpdfapi/page/cpdf_image.h
index 1909498102..5b737e67b9 100644
--- a/core/fpdfapi/page/cpdf_image.h
+++ b/core/fpdfapi/page/cpdf_image.h
@@ -44,25 +44,25 @@ class CPDF_Image {
   bool IsMask() const { return m_bIsMask; }
   bool IsInterpol() const { return m_bInterpolate; }
 
-  std::unique_ptr<CFX_DIBSource> LoadDIBSource() const;
+  CFX_RetainPtr<CFX_DIBSource> LoadDIBSource() const;
 
-  void SetImage(const CFX_DIBitmap* pDIBitmap);
+  void SetImage(const CFX_RetainPtr<CFX_DIBitmap>& pDIBitmap);
   void SetJpegImage(const CFX_RetainPtr<IFX_SeekableReadStream>& pFile);
   void SetJpegImageInline(const CFX_RetainPtr<IFX_SeekableReadStream>& pFile);
 
-  void ResetCache(CPDF_Page* pPage, const CFX_DIBitmap* pDIBitmap);
-
+  void ResetCache(CPDF_Page* pPage,
+                  const CFX_RetainPtr<CFX_DIBitmap>& pDIBitmap);
   bool StartLoadDIBSource(CPDF_Dictionary* pFormResource,
                           CPDF_Dictionary* pPageResource,
                           bool bStdCS = false,
                           uint32_t GroupFamily = 0,
                           bool bLoadMask = false);
   bool Continue(IFX_Pause* pPause);
-  CFX_DIBSource* DetachBitmap();
-  CFX_DIBSource* DetachMask();
+  CFX_RetainPtr<CFX_DIBSource> DetachBitmap();
+  CFX_RetainPtr<CFX_DIBSource> DetachMask();
 
-  CFX_DIBSource* m_pDIBSource = nullptr;
-  CFX_DIBSource* m_pMask = nullptr;
+  CFX_RetainPtr<CFX_DIBSource> m_pDIBSource;
+  CFX_RetainPtr<CFX_DIBSource> m_pMask;
   uint32_t m_MatteColor = 0;
 
  private: