diff --git a/src/butil/third_party/rapidjson/allocators.h b/src/butil/third_party/rapidjson/allocators.h index 7ac0007ca9..9a3092873e 100644 --- a/src/butil/third_party/rapidjson/allocators.h +++ b/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(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(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. diff --git a/test/BUILD.bazel b/test/BUILD.bazel index a2c291eb7e..8555968f0f 100644 --- a/test/BUILD.bazel +++ b/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", diff --git a/test/CMakeLists.txt b/test/CMakeLists.txt index 4c3a09f529..71dbd1f69f 100644 --- a/test/CMakeLists.txt +++ b/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 diff --git a/test/rapidjson_allocator_unittest.cc b/test/rapidjson_allocator_unittest.cc new file mode 100644 index 0000000000..1e3752ff92 --- /dev/null +++ b/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 +#include +#include +#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 pool( + 32, &allocator); + EXPECT_EQ(nullptr, pool.Malloc(8)); + EXPECT_EQ(0u, pool.Size()); + allocator.fail = false; + char* saved = static_cast(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(pool.Realloc(saved, 8, 64)); + ASSERT_NE(nullptr, resized); + EXPECT_STREQ("stored", resized); +} +} // namespace