10 files changed, 65 insertions, 9 deletions
diff --git a/core/fpdfapi/edit/cpdf_pagecontentgenerator_unittest.cpp b/core/fpdfapi/edit/cpdf_pagecontentgenerator_unittest.cpp
index 331cec3813..fdf31243b7 100644
--- a/core/fpdfapi/edit/cpdf_pagecontentgenerator_unittest.cpp
+++ b/core/fpdfapi/edit/cpdf_pagecontentgenerator_unittest.cpp
@@ -14,13 +14,20 @@
 #include "core/fpdfapi/parser/cpdf_parser.h"
 #include "core/fpdfapi/parser/cpdf_reference.h"
 #include "testing/gtest/include/gtest/gtest.h"
+#include "testing/test_support.h"
 #include "third_party/base/ptr_util.h"
 
-class CPDF_PageContentGeneratorTest : public testing::Test {
+class CPDF_PageContentGeneratorTest : public pdfium::FPDF_Test {
  protected:
-  void SetUp() override { CPDF_ModuleMgr::Get()->InitPageModule(); }
+  void SetUp() override {
+    FPDF_Test::SetUp();
+    CPDF_ModuleMgr::Get()->InitPageModule();
+  }
 
-  void TearDown() override { CPDF_ModuleMgr::Destroy(); }
+  void TearDown() override {
+    CPDF_ModuleMgr::Destroy();
+    FPDF_Test::TearDown();
+  }
 
   void TestProcessPath(CPDF_PageContentGenerator* pGen,
                        CFX_ByteTextBuf* buf,
diff --git a/core/fxcrt/cfx_string_data_template.h b/core/fxcrt/cfx_string_data_template.h
index 2e87a811a5..33013fde33 100644
--- a/core/fxcrt/cfx_string_data_template.h
+++ b/core/fxcrt/cfx_string_data_template.h
@@ -35,7 +35,8 @@ class CFX_StringDataTemplate {
     int usableLen = (totalSize - overhead) / sizeof(CharType);
     ASSERT(usableLen >= nLen);
 
-    void* pData = FX_Alloc(uint8_t, totalSize);
+    void* pData = pdfium::base::PartitionAllocGeneric(
+        gStringPartitionAllocator.root(), totalSize, "CFX_StringDataTemplate");
     return new (pData) CFX_StringDataTemplate(nLen, usableLen);
   }
 
@@ -54,7 +55,7 @@ class CFX_StringDataTemplate {
   void Retain() { ++m_nRefs; }
   void Release() {
     if (--m_nRefs <= 0)
-      FX_Free(this);
+      pdfium::base::PartitionFree(this);
   }
 
   bool CanOperateInPlace(FX_STRSIZE nTotalLen) const {
diff --git a/core/fxcrt/fx_basic_memmgr.cpp b/core/fxcrt/fx_basic_memmgr.cpp
index 927b994b5f..06568c04ad 100644
--- a/core/fxcrt/fx_basic_memmgr.cpp
+++ b/core/fxcrt/fx_basic_memmgr.cpp
@@ -8,6 +8,9 @@
 
 #include "core/fxcrt/fx_memory.h"
 
+pdfium::base::PartitionAllocatorGeneric gArrayBufferPartitionAllocator;
+pdfium::base::PartitionAllocatorGeneric gStringPartitionAllocator;
+
 void* FXMEM_DefaultAlloc(size_t byte_size, int flags) {
   return (void*)malloc(byte_size);
 }
diff --git a/core/fxcrt/fx_memory.h b/core/fxcrt/fx_memory.h
index 0ad28ce896..c4d619efab 100644
--- a/core/fxcrt/fx_memory.h
+++ b/core/fxcrt/fx_memory.h
@@ -26,6 +26,11 @@ void FXMEM_DefaultFree(void* pointer, int flags);
 #include <memory>
 #include <new>
 
+#include "third_party/base/allocator/partition_allocator/partition_alloc.h"
+
+extern pdfium::base::PartitionAllocatorGeneric gArrayBufferPartitionAllocator;
+extern pdfium::base::PartitionAllocatorGeneric gStringPartitionAllocator;
+
 NEVER_INLINE void FX_OutOfMemoryTerminate();
 
 inline void* FX_SafeRealloc(void* ptr, size_t num_members, size_t member_size) {
diff --git a/fpdfsdk/fpdfview.cpp b/fpdfsdk/fpdfview.cpp
index 4c3631fc88..d4284a3217 100644
--- a/fpdfsdk/fpdfview.cpp
+++ b/fpdfsdk/fpdfview.cpp
@@ -33,6 +33,7 @@
 #include "fpdfsdk/javascript/ijs_runtime.h"
 #include "public/fpdf_ext.h"
 #include "public/fpdf_progressive.h"
+#include "third_party/base/allocator/partition_allocator/partition_alloc.h"
 #include "third_party/base/numerics/safe_conversions_impl.h"
 #include "third_party/base/ptr_util.h"
 
@@ -367,6 +368,14 @@ FPDF_InitLibraryWithConfig(const FPDF_LIBRARY_CONFIG* cfg) {
   if (g_pCodecModule)
     return;
 
+  static bool s_gPartitionAllocatorsInitialized = false;
+  if (!s_gPartitionAllocatorsInitialized) {
+    pdfium::base::PartitionAllocGlobalInit(FX_OutOfMemoryTerminate);
+    gArrayBufferPartitionAllocator.init();
+    gStringPartitionAllocator.init();
+    s_gPartitionAllocatorsInitialized = true;
+  }
+
   g_pCodecModule = new CCodec_ModuleMgr();
 
   CFX_GEModule* pModule = CFX_GEModule::Get();
diff --git a/fxjs/fxjs_v8.cpp b/fxjs/fxjs_v8.cpp
index 5f9426b643..9c1c3e84f3 100644
--- a/fxjs/fxjs_v8.cpp
+++ b/fxjs/fxjs_v8.cpp
@@ -9,6 +9,7 @@
 #include <vector>
 
 #include "core/fxcrt/fx_basic.h"
+#include "third_party/base/allocator/partition_allocator/partition_alloc.h"
 
 // Keep this consistent with the values defined in gin/public/context_holder.h
 // (without actually requiring a dependency on gin itself for the standalone
@@ -144,15 +145,23 @@ static v8::Local<v8::ObjectTemplate> GetGlobalObjectTemplate(
 }
 
 void* FXJS_ArrayBufferAllocator::Allocate(size_t length) {
-  return length <= kMaxAllowedBytes ? calloc(1, length) : nullptr;
+  if (length > kMaxAllowedBytes)
+    return nullptr;
+  void* p = AllocateUninitialized(length);
+  if (p)
+    memset(p, 0, length);
+  return p;
 }
 
 void* FXJS_ArrayBufferAllocator::AllocateUninitialized(size_t length) {
-  return length < kMaxAllowedBytes ? malloc(length) : nullptr;
+  if (length > kMaxAllowedBytes)
+    return nullptr;
+  return pdfium::base::PartitionAllocGeneric(
+      gArrayBufferPartitionAllocator.root(), length, "FXJS_ArrayBuffer");
 }
 
 void FXJS_ArrayBufferAllocator::Free(void* data, size_t length) {
-  free(data);
+  pdfium::base::PartitionFree(data);
 }
 
 void V8TemplateMapTraits::Dispose(v8::Isolate* isolate,
diff --git a/fxjs/fxjs_v8.h b/fxjs/fxjs_v8.h
index bdcf425f53..d44af2f8fa 100644
--- a/fxjs/fxjs_v8.h
+++ b/fxjs/fxjs_v8.h
@@ -22,6 +22,7 @@
 #include <vector>
 
 #include "core/fxcrt/fx_string.h"
+
 #ifdef PDF_ENABLE_XFA
 // Header for CFXJSE_RuntimeData. FXJS_V8 doesn't interpret this class,
 // it is just passed along to XFA.
diff --git a/testing/test_support.cpp b/testing/test_support.cpp
index a530b3ea4c..6e1bb64bc0 100644
--- a/testing/test_support.cpp
+++ b/testing/test_support.cpp
@@ -210,3 +210,15 @@ int TestSaver::WriteBlockCallback(FPDF_FILEWRITE* pFileWrite,
   pThis->m_String.append(static_cast<const char*>(data), size);
   return 1;
 }
+
+namespace pdfium {
+
+void FPDF_Test::SetUp() {
+  FPDF_InitLibrary();
+}
+
+void FPDF_Test::TearDown() {
+  FPDF_DestroyLibrary();
+}
+
+}  // namespace pdfium
diff --git a/testing/test_support.h b/testing/test_support.h
index 4111ca61ed..b734bc31f0 100644
--- a/testing/test_support.h
+++ b/testing/test_support.h
@@ -12,6 +12,7 @@
 
 #include "public/fpdf_save.h"
 #include "public/fpdfview.h"
+#include "testing/gtest/include/gtest/gtest.h"
 
 #ifdef PDF_ENABLE_V8
 #include "v8/include/v8.h"
@@ -58,6 +59,12 @@ struct FreeDeleter {
   inline void operator()(void* ptr) const { free(ptr); }
 };
 
+class FPDF_Test : public ::testing::Test {
+ public:
+  void SetUp() override;
+  void TearDown() override;
+};
+
 }  // namespace pdfium
 
 // Reads the entire contents of a file into a newly alloc'd buffer.
diff --git a/xfa/fde/cfde_txtedtbuf_unittest.cpp b/xfa/fde/cfde_txtedtbuf_unittest.cpp
index 9c564093e5..aa80a66b60 100644
--- a/xfa/fde/cfde_txtedtbuf_unittest.cpp
+++ b/xfa/fde/cfde_txtedtbuf_unittest.cpp
@@ -8,12 +8,14 @@
 #include "testing/test_support.h"
 #include "third_party/base/ptr_util.h"
 
-class CFDE_TxtEdtBufTest : public testing::Test {
+class CFDE_TxtEdtBufTest : public pdfium::FPDF_Test {
  public:
   void SetUp() override {
+    FPDF_Test::SetUp();
     buf_ = pdfium::MakeUnique<CFDE_TxtEdtBuf>();
     buf_->SetChunkSizeForTesting(5);
   }
+
   size_t ChunkCount() const { return buf_->m_chunks.size(); }
 
   std::unique_ptr<CFDE_TxtEdtBuf> buf_;