Handle JSON memory pool allocation failures

src/butil/third_party/rapidjson/allocators.h

@@ -178,8 +178,9 @@ class MemoryPoolAllocator {
             return NULL;
 
         size = RAPIDJSON_ALIGN(size);
-        if (chunkHead_ == 0 || chunkHead_->size + size > chunkHead_->capacity)
-            AddChunk(chunk_capacity_ > size ? chunk_capacity_ : size);
+        if ((chunkHead_ == 0 || chunkHead_->size + size > chunkHead_->capacity)
+            && !AddChunk(chunk_capacity_ > size ? chunk_capacity_ : size))
+            return NULL;
 
         void *buffer = reinterpret_cast<char *>(chunkHead_) + RAPIDJSON_ALIGN(sizeof(ChunkHeader)) + chunkHead_->size;
         chunkHead_->size += size;
@@ -210,7 +211,8 @@ class MemoryPoolAllocator {
 
         // Realloc process: allocate and copy memory, do not free original buffer.
         void* newBuffer = Malloc(newSize);
-        RAPIDJSON_ASSERT(newBuffer != 0);   // Do not handle out-of-memory explicitly.
+        if (newBuffer == 0)
+            return NULL;
         if (originalSize)
             std::memcpy(newBuffer, originalPtr, originalSize);
         return newBuffer;
@@ -228,14 +230,17 @@ class MemoryPoolAllocator {
     //! Creates a new chunk.
     /*! \param capacity Capacity of the chunk in bytes.
     */
-    void AddChunk(size_t capacity) {
+    bool AddChunk(size_t capacity) {
         if (!baseAllocator_)
             ownBaseAllocator_ = baseAllocator_ = RAPIDJSON_NEW(BaseAllocator());
         ChunkHeader* chunk = reinterpret_cast<ChunkHeader*>(baseAllocator_->Malloc(RAPIDJSON_ALIGN(sizeof(ChunkHeader)) + capacity));
+        if (chunk == 0)
+            return false;
         chunk->capacity = capacity;
         chunk->size = 0;
         chunk->next = chunkHead_;
         chunkHead_ =  chunk;
+        return true;
     }
 
     static const int kDefaultChunkCapacity = 64 * 1024; //!< Default chunk capacity.

test/BUILD.bazel

@@ -27,6 +27,7 @@ COPTS = [
 ]
 
 TEST_BUTIL_SOURCES = [
+    "rapidjson_allocator_unittest.cc",
     "at_exit_unittest.cc",
     "atomicops_unittest.cc",
     "base64_unittest.cc",

test/CMakeLists.txt

@@ -88,6 +88,7 @@ file(COPY ${CMAKE_CURRENT_SOURCE_DIR}/run_tests.sh DESTINATION ${CMAKE_CURRENT_B
 SET(TEST_BUTIL_SOURCES
     ${CMAKE_CURRENT_SOURCE_DIR}/abalist_unittest.cc
     ${CMAKE_CURRENT_SOURCE_DIR}/aligned_memory_unittest.cc
+    ${CMAKE_CURRENT_SOURCE_DIR}/rapidjson_allocator_unittest.cc
     ${CMAKE_CURRENT_SOURCE_DIR}/at_exit_unittest.cc
     ${CMAKE_CURRENT_SOURCE_DIR}/atomicops_unittest.cc
     ${CMAKE_CURRENT_SOURCE_DIR}/baidu_thread_local_unittest.cpp

test/rapidjson_allocator_unittest.cc

@@ -0,0 +1,53 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include <gtest/gtest.h>
+#include <cstdlib>
+#include <cstring>
+#include "butil/third_party/rapidjson/allocators.h"
+
+namespace {
+struct FailingAllocator {
+    bool fail = true;
+    void* Malloc(size_t size) { return fail ? nullptr : std::malloc(size); }
+    static void Free(void* ptr) { std::free(ptr); }
+};
+
+TEST(RapidJsonAllocatorTest, allocation_failure_preserves_existing_chunks) {
+    FailingAllocator allocator;
+    BUTIL_RAPIDJSON_NAMESPACE::MemoryPoolAllocator<FailingAllocator> pool(
+        32, &allocator);
+    EXPECT_EQ(nullptr, pool.Malloc(8));
+    EXPECT_EQ(0u, pool.Size());
+    allocator.fail = false;
+    char* saved = static_cast<char*>(pool.Malloc(8));
+    ASSERT_NE(nullptr, saved);
+    memcpy(saved, "stored", 7);
+    const size_t size = pool.Size();
+    const size_t capacity = pool.Capacity();
+    allocator.fail = true;
+    EXPECT_EQ(nullptr, pool.Malloc(64));
+    EXPECT_EQ(nullptr, pool.Realloc(saved, 8, 64));
+    EXPECT_EQ(size, pool.Size());
+    EXPECT_EQ(capacity, pool.Capacity());
+    EXPECT_STREQ("stored", saved);
+    allocator.fail = false;
+    char* resized = static_cast<char*>(pool.Realloc(saved, 8, 64));
+    ASSERT_NE(nullptr, resized);
+    EXPECT_STREQ("stored", resized);
+}
+}  // namespace