Sanitize CJBig2_SymbolDict's memory usage.chromium/2534chromium/2533chromium/2532

- Use std::vector<JBig2ArithCtx> instead of storing pointers to arrays.
- Make CJBig2_SymbolDict's members private with accessors.
- Use std::vector<JBig2ArithCtx> in related places.
- Steal Chromium's vector_as_array() and use it as an adaptor as needed.

BUG=514891
R=tsepez@chromium.org

Review URL: https://codereview.chromium.org/1388203003 .

8 files changed, 103 insertions, 73 deletions

diff --git a/core/src/fxcodec/jbig2/JBig2_Context.cpp b/core/src/fxcodec/jbig2/JBig2_Context.cpp
index bb01cb6be8..06863b0bc4 100644
--- a/core/src/fxcodec/jbig2/JBig2_Context.cpp
+++ b/core/src/fxcodec/jbig2/JBig2_Context.cpp
@@ -568,33 +568,32 @@ int32_t CJBig2_Context::parseSymbolDict(CJBig2_Segment* pSegment,
     }
   }
 
-  nonstd::unique_ptr<JBig2ArithCtx, FxFreeDeleter> gbContext;
-  nonstd::unique_ptr<JBig2ArithCtx, FxFreeDeleter> grContext;
-  if ((wFlags & 0x0100) && pLRSeg && pLRSeg->m_Result.sd->m_bContextRetained) {
-    if (pSymbolDictDecoder->SDHUFF == 0) {
-      const size_t size = GetHuffContextSize(pSymbolDictDecoder->SDTEMPLATE);
-      gbContext.reset(FX_Alloc(JBig2ArithCtx, size));
-      JBIG2_memcpy(gbContext.get(), pLRSeg->m_Result.sd->m_gbContext,
-                   sizeof(JBig2ArithCtx) * size);
+  const bool bUseGbContext = (pSymbolDictDecoder->SDHUFF == 0);
+  const bool bUseGrContext = (pSymbolDictDecoder->SDREFAGG == 1);
+  const size_t gbContextSize =
+      GetHuffContextSize(pSymbolDictDecoder->SDTEMPLATE);
+  const size_t grContextSize =
+      GetRefAggContextSize(pSymbolDictDecoder->SDRTEMPLATE);
+  std::vector<JBig2ArithCtx> gbContext;
+  std::vector<JBig2ArithCtx> grContext;
+  if ((wFlags & 0x0100) && pLRSeg) {
+    if (bUseGbContext) {
+      gbContext = pLRSeg->m_Result.sd->GbContext();
+      if (gbContext.size() != gbContextSize)
+        return JBIG2_ERROR_FATAL;
     }
-    if (pSymbolDictDecoder->SDREFAGG == 1) {
-      const size_t size = GetRefAggContextSize(pSymbolDictDecoder->SDRTEMPLATE);
-      grContext.reset(FX_Alloc(JBig2ArithCtx, size));
-      JBIG2_memcpy(grContext.get(), pLRSeg->m_Result.sd->m_grContext,
-                   sizeof(JBig2ArithCtx) * size);
+    if (bUseGrContext) {
+      grContext = pLRSeg->m_Result.sd->GrContext();
+      if (grContext.size() != grContextSize)
+        return JBIG2_ERROR_FATAL;
     }
   } else {
-    if (pSymbolDictDecoder->SDHUFF == 0) {
-      const size_t size = GetHuffContextSize(pSymbolDictDecoder->SDTEMPLATE);
-      gbContext.reset(FX_Alloc(JBig2ArithCtx, size));
-      JBIG2_memset(gbContext.get(), 0, sizeof(JBig2ArithCtx) * size);
-    }
-    if (pSymbolDictDecoder->SDREFAGG == 1) {
-      const size_t size = GetRefAggContextSize(pSymbolDictDecoder->SDRTEMPLATE);
-      grContext.reset(FX_Alloc(JBig2ArithCtx, size));
-      JBIG2_memset(grContext.get(), 0, sizeof(JBig2ArithCtx) * size);
-    }
+    if (bUseGbContext)
+      gbContext.resize(gbContextSize);
+    if (bUseGrContext)
+      grContext.resize(grContextSize);
   }
+
   CJBig2_CacheKey key =
       CJBig2_CacheKey(pSegment->m_dwObjNum, pSegment->m_dwDataOffset);
   FX_BOOL cache_hit = false;
@@ -613,11 +612,11 @@ int32_t CJBig2_Context::parseSymbolDict(CJBig2_Segment* pSegment,
     }
   }
   if (!cache_hit) {
-    if (pSymbolDictDecoder->SDHUFF == 0) {
+    if (bUseGbContext) {
       nonstd::unique_ptr<CJBig2_ArithDecoder> pArithDecoder(
           new CJBig2_ArithDecoder(m_pStream.get()));
       pSegment->m_Result.sd = pSymbolDictDecoder->decode_Arith(
-          pArithDecoder.get(), gbContext.get(), grContext.get());
+          pArithDecoder.get(), &gbContext, &grContext);
       if (!pSegment->m_Result.sd)
         return JBIG2_ERROR_FATAL;
 
@@ -625,7 +624,7 @@ int32_t CJBig2_Context::parseSymbolDict(CJBig2_Segment* pSegment,
       m_pStream->offset(2);
     } else {
       pSegment->m_Result.sd = pSymbolDictDecoder->decode_Huffman(
-          m_pStream.get(), gbContext.get(), grContext.get(), pPause);
+          m_pStream.get(), &gbContext, &grContext, pPause);
       if (!pSegment->m_Result.sd)
         return JBIG2_ERROR_FATAL;
       m_pStream->alignByte();
@@ -633,23 +632,18 @@ int32_t CJBig2_Context::parseSymbolDict(CJBig2_Segment* pSegment,
     if (m_bIsGlobal && kSymbolDictCacheMaxSize > 0) {
       nonstd::unique_ptr<CJBig2_SymbolDict> value =
           pSegment->m_Result.sd->DeepCopy();
-      if (value) {
-        while (m_pSymbolDictCache->size() >= kSymbolDictCacheMaxSize) {
-          delete m_pSymbolDictCache->back().second;
-          m_pSymbolDictCache->pop_back();
-        }
-        m_pSymbolDictCache->push_front(CJBig2_CachePair(key, value.release()));
+      while (m_pSymbolDictCache->size() >= kSymbolDictCacheMaxSize) {
+        delete m_pSymbolDictCache->back().second;
+        m_pSymbolDictCache->pop_back();
       }
+      m_pSymbolDictCache->push_front(CJBig2_CachePair(key, value.release()));
     }
   }
   if (wFlags & 0x0200) {
-    pSegment->m_Result.sd->m_bContextRetained = TRUE;
-    if (pSymbolDictDecoder->SDHUFF == 0) {
-      pSegment->m_Result.sd->m_gbContext = gbContext.release();
-    }
-    if (pSymbolDictDecoder->SDREFAGG == 1) {
-      pSegment->m_Result.sd->m_grContext = grContext.release();
-    }
+    if (bUseGbContext)
+      pSegment->m_Result.sd->SetGbContext(gbContext);
+    if (bUseGrContext)
+      pSegment->m_Result.sd->SetGrContext(grContext);
   }
   return JBIG2_SUCCESS;
 }
diff --git a/core/src/fxcodec/jbig2/JBig2_SddProc.cpp b/core/src/fxcodec/jbig2/JBig2_SddProc.cpp
index 1d0393a7e3..ab3b34f583 100644
--- a/core/src/fxcodec/jbig2/JBig2_SddProc.cpp
+++ b/core/src/fxcodec/jbig2/JBig2_SddProc.cpp
@@ -7,6 +7,7 @@
 #include "JBig2_SddProc.h"
 
 #include "../../../../third_party/base/nonstd_unique_ptr.h"
+#include "../../../../third_party/base/stl_util.h"
 #include "../../../include/fxcrt/fx_basic.h"
 #include "JBig2_ArithIntDecoder.h"
 #include "JBig2_GrdProc.h"
@@ -16,10 +17,12 @@
 #include "JBig2_SymbolDict.h"
 #include "JBig2_TrdProc.h"
 
+using pdfium::vector_as_array;
+
 CJBig2_SymbolDict* CJBig2_SDDProc::decode_Arith(
     CJBig2_ArithDecoder* pArithDecoder,
-    JBig2ArithCtx* gbContext,
-    JBig2ArithCtx* grContext) {
+    std::vector<JBig2ArithCtx>* gbContext,
+    std::vector<JBig2ArithCtx>* grContext) {
   CJBig2_Image** SDNEWSYMS;
   FX_DWORD HCHEIGHT, NSYMSDECODED;
   int32_t HCDH;
@@ -104,7 +107,7 @@ CJBig2_SymbolDict* CJBig2_SDDProc::decode_Arith(
         pGRD->GBAT[5] = SDAT[5];
         pGRD->GBAT[6] = SDAT[6];
         pGRD->GBAT[7] = SDAT[7];
-        BS = pGRD->decode_Arith(pArithDecoder, gbContext);
+        BS = pGRD->decode_Arith(pArithDecoder, vector_as_array(gbContext));
         if (!BS) {
           goto failed;
         }
@@ -192,7 +195,8 @@ CJBig2_SymbolDict* CJBig2_SDDProc::decode_Arith(
           ids.IARDX = IARDX.get();
           ids.IARDY = IARDY.get();
           ids.IAID = IAID.get();
-          BS = pDecoder->decode_Arith(pArithDecoder, grContext, &ids);
+          BS = pDecoder->decode_Arith(pArithDecoder, vector_as_array(grContext),
+                                      &ids);
           if (!BS) {
             FX_Free(SBSYMS);
             goto failed;
@@ -227,7 +231,7 @@ CJBig2_SymbolDict* CJBig2_SDDProc::decode_Arith(
           pGRRD->GRAT[1] = SDRAT[1];
           pGRRD->GRAT[2] = SDRAT[2];
           pGRRD->GRAT[3] = SDRAT[3];
-          BS = pGRRD->decode(pArithDecoder, grContext);
+          BS = pGRRD->decode(pArithDecoder, vector_as_array(grContext));
           if (!BS) {
             FX_Free(SBSYMS);
             goto failed;
@@ -285,10 +289,11 @@ failed:
   return nullptr;
 }
 
-CJBig2_SymbolDict* CJBig2_SDDProc::decode_Huffman(CJBig2_BitStream* pStream,
-                                                  JBig2ArithCtx* gbContext,
-                                                  JBig2ArithCtx* grContext,
-                                                  IFX_Pause* pPause) {
+CJBig2_SymbolDict* CJBig2_SDDProc::decode_Huffman(
+    CJBig2_BitStream* pStream,
+    std::vector<JBig2ArithCtx>* gbContext,
+    std::vector<JBig2ArithCtx>* grContext,
+    IFX_Pause* pPause) {
   CJBig2_Image** SDNEWSYMS;
   FX_DWORD* SDNEWSYMWIDTHS;
   FX_DWORD HCHEIGHT, NSYMSDECODED;
@@ -440,7 +445,7 @@ CJBig2_SymbolDict* CJBig2_SDDProc::decode_Huffman(CJBig2_BitStream* pStream,
           pDecoder->SBRAT[1] = SDRAT[1];
           pDecoder->SBRAT[2] = SDRAT[2];
           pDecoder->SBRAT[3] = SDRAT[3];
-          BS = pDecoder->decode_Huffman(pStream, grContext);
+          BS = pDecoder->decode_Huffman(pStream, vector_as_array(grContext));
           if (!BS) {
             FX_Free(SBSYMCODES);
             FX_Free(SBSYMS);
@@ -512,7 +517,7 @@ CJBig2_SymbolDict* CJBig2_SDDProc::decode_Huffman(CJBig2_BitStream* pStream,
           pGRRD->GRAT[3] = SDRAT[3];
           nonstd::unique_ptr<CJBig2_ArithDecoder> pArithDecoder(
               new CJBig2_ArithDecoder(pStream));
-          BS = pGRRD->decode(pArithDecoder.get(), grContext);
+          BS = pGRRD->decode(pArithDecoder.get(), vector_as_array(grContext));
           if (!BS) {
             FX_Free(SBSYMS);
             goto failed;
diff --git a/core/src/fxcodec/jbig2/JBig2_SddProc.h b/core/src/fxcodec/jbig2/JBig2_SddProc.h
index 2c55113dae..01f8014ccc 100644
--- a/core/src/fxcodec/jbig2/JBig2_SddProc.h
+++ b/core/src/fxcodec/jbig2/JBig2_SddProc.h
@@ -7,25 +7,26 @@
 #ifndef CORE_SRC_FXCODEC_JBIG2_JBIG2_SDDPROC_H_
 #define CORE_SRC_FXCODEC_JBIG2_JBIG2_SDDPROC_H_
 
+#include <vector>
+
 #include "../../../include/fxcrt/fx_system.h"
+#include "JBig2_ArithDecoder.h"
 
-class CJBig2_ArithDecoder;
 class CJBig2_BitStream;
 class CJBig2_HuffmanTable;
 class CJBig2_Image;
 class CJBig2_SymbolDict;
 class IFX_Pause;
-struct JBig2ArithCtx;
 
 class CJBig2_SDDProc {
  public:
   CJBig2_SymbolDict* decode_Arith(CJBig2_ArithDecoder* pArithDecoder,
-                                  JBig2ArithCtx* gbContext,
-                                  JBig2ArithCtx* grContext);
+                                  std::vector<JBig2ArithCtx>* gbContext,
+                                  std::vector<JBig2ArithCtx>* grContext);
 
   CJBig2_SymbolDict* decode_Huffman(CJBig2_BitStream* pStream,
-                                    JBig2ArithCtx* gbContext,
-                                    JBig2ArithCtx* grContext,
+                                    std::vector<JBig2ArithCtx>* gbContext,
+                                    std::vector<JBig2ArithCtx>* grContext,
                                     IFX_Pause* pPause);
 
  public:
diff --git a/core/src/fxcodec/jbig2/JBig2_SymbolDict.cpp b/core/src/fxcodec/jbig2/JBig2_SymbolDict.cpp
index a8f8a94529..351a8389c8 100644
--- a/core/src/fxcodec/jbig2/JBig2_SymbolDict.cpp
+++ b/core/src/fxcodec/jbig2/JBig2_SymbolDict.cpp
@@ -10,27 +10,19 @@
 #include "JBig2_Image.h"
 
 CJBig2_SymbolDict::CJBig2_SymbolDict() {
-  m_bContextRetained = FALSE;
-  m_gbContext = m_grContext = NULL;
 }
 
 CJBig2_SymbolDict::~CJBig2_SymbolDict() {
-  if (m_bContextRetained) {
-    FX_Free(m_gbContext);
-    FX_Free(m_grContext);
-  }
 }
 
 nonstd::unique_ptr<CJBig2_SymbolDict> CJBig2_SymbolDict::DeepCopy() const {
-  nonstd::unique_ptr<CJBig2_SymbolDict> dst;
   const CJBig2_SymbolDict* src = this;
-  if (src->m_bContextRetained || src->m_gbContext || src->m_grContext)
-    return dst;
-
-  dst.reset(new CJBig2_SymbolDict);
+  nonstd::unique_ptr<CJBig2_SymbolDict> dst(new CJBig2_SymbolDict);
   for (size_t i = 0; i < src->m_SDEXSYMS.size(); ++i) {
     CJBig2_Image* image = src->m_SDEXSYMS.get(i);
     dst->m_SDEXSYMS.push_back(image ? new CJBig2_Image(*image) : nullptr);
   }
+  dst->m_gbContext = src->m_gbContext;
+  dst->m_grContext = src->m_grContext;
   return dst;
 }
diff --git a/core/src/fxcodec/jbig2/JBig2_SymbolDict.h b/core/src/fxcodec/jbig2/JBig2_SymbolDict.h
index 577bfbc2f2..6ff4c2efe2 100644
--- a/core/src/fxcodec/jbig2/JBig2_SymbolDict.h
+++ b/core/src/fxcodec/jbig2/JBig2_SymbolDict.h
@@ -7,12 +7,14 @@
 #ifndef CORE_SRC_FXCODEC_JBIG2_JBIG2_SYMBOLDICT_H_
 #define CORE_SRC_FXCODEC_JBIG2_JBIG2_SYMBOLDICT_H_
 
+#include <vector>
+
 #include "../../../../third_party/base/nonstd_unique_ptr.h"
 #include "../../../include/fxcrt/fx_basic.h"
+#include "JBig2_ArithDecoder.h"
 #include "JBig2_List.h"
 
 class CJBig2_Image;
-struct JBig2ArithCtx;
 
 class CJBig2_SymbolDict {
  public:
@@ -27,12 +29,19 @@ class CJBig2_SymbolDict {
   size_t NumImages() const { return m_SDEXSYMS.size(); }
   CJBig2_Image* GetImage(size_t index) const { return m_SDEXSYMS.get(index); }
 
- public:
-  FX_BOOL m_bContextRetained;
-  JBig2ArithCtx* m_gbContext;
-  JBig2ArithCtx* m_grContext;
+  const std::vector<JBig2ArithCtx>& GbContext() const { return m_gbContext; }
+  const std::vector<JBig2ArithCtx>& GrContext() const { return m_grContext; }
+
+  void SetGbContext(const std::vector<JBig2ArithCtx>& gbContext) {
+    m_gbContext = gbContext;
+  }
+  void SetGrContext(const std::vector<JBig2ArithCtx>& grContext) {
+    m_grContext = grContext;
+  }
 
  private:
+  std::vector<JBig2ArithCtx> m_gbContext;
+  std::vector<JBig2ArithCtx> m_grContext;
   CJBig2_List<CJBig2_Image> m_SDEXSYMS;
 };
 
diff --git a/third_party/BUILD.gn b/third_party/BUILD.gn
index 21d3d5db32..f6358147be 100644
--- a/third_party/BUILD.gn
+++ b/third_party/BUILD.gn
@@ -300,6 +300,7 @@ source_set("pdfium_base") {
     "base/numerics/safe_conversions_impl.h",
     "base/numerics/safe_math.h",
     "base/numerics/safe_math_impl.h",
+    "base/stl_util.h",
     "base/template_util.h",
   ]
 }
diff --git a/third_party/base/stl_util.h b/third_party/base/stl_util.h
new file mode 100644
index 0000000000..fcbe5882a2
--- /dev/null
+++ b/third_party/base/stl_util.h
@@ -0,0 +1,27 @@
+// Copyright 2015 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#ifndef PDFIUM_THIRD_PARTY_BASE_STL_UTIL_H_
+#define PDFIUM_THIRD_PARTY_BASE_STL_UTIL_H_
+
+#include <vector>
+
+namespace pdfium {
+
+// To treat a possibly-empty vector as an array, use these functions.
+// If you know the array will never be empty, you can use &*v.begin()
+// directly, but that is undefined behaviour if |v| is empty.
+template <typename T>
+inline T* vector_as_array(std::vector<T>* v) {
+  return v->empty() ? nullptr : &*v->begin();
+}
+
+template <typename T>
+inline const T* vector_as_array(const std::vector<T>* v) {
+  return v->empty() ? nullptr : &*v->begin();
+}
+
+}  // namespace pdfium
+
+#endif  // PDFIUM_THIRD_PARTY_BASE_STL_UTIL_H_
diff --git a/third_party/third_party.gyp b/third_party/third_party.gyp
index 6ad8beb48c..1ec0509eb2 100644
--- a/third_party/third_party.gyp
+++ b/third_party/third_party.gyp
@@ -281,11 +281,12 @@
         'base/logging.h',
         'base/macros.h',
         'base/nonstd_unique_ptr.h',
-        'base/template_util.h',
         'base/numerics/safe_conversions.h',
         'base/numerics/safe_conversions_impl.h',
         'base/numerics/safe_math.h',
         'base/numerics/safe_math_impl.h',
+        'base/stl_util.h',
+        'base/template_util.h',
       ],
     },
   ],