From 31a0df5630ae225857c5a08977e86359b96fd615 Mon Sep 17 00:00:00 2001
From: Wang Xiaofeng <wasphin@gmail.com>
Date: Sun, 25 Oct 2026 23:44:48 +0800
Subject: [PATCH] Handle Arena block allocation failures

---
 src/butil/arena.cpp    |  9 +++++++++
 test/BUILD.bazel       |  1 +
 test/CMakeLists.txt    |  1 +
 test/arena_unittest.cc | 43 ++++++++++++++++++++++++++++++++++++++++++
 4 files changed, 54 insertions(+)
 create mode 100644 test/arena_unittest.cc

diff --git a/src/butil/arena.cpp b/src/butil/arena.cpp
index dc07d8d75a..43d59d860b 100644
--- a/src/butil/arena.cpp
+++ b/src/butil/arena.cpp
@@ -64,7 +64,13 @@ void Arena::clear() {
 }
 
 void* Arena::allocate_new_block(size_t n) {
+    if (n > UINT32_MAX || n > SIZE_MAX - offsetof(Block, data)) {
+        return nullptr;
+    }
     Block* b = (Block*)malloc(offsetof(Block, data) + n);
+    if (b == nullptr) {
+        return nullptr;
+    }
     b->next = _isolated_blocks;
     b->alloc_size = n;
     b->size = n;
@@ -86,6 +92,9 @@ void* Arena::allocate_in_other_blocks(size_t n) {
     if (new_size < n) {
         new_size = n;
     }
+    if (new_size > UINT32_MAX || new_size > SIZE_MAX - offsetof(Block, data)) {
+        return nullptr;
+    }
     Block* b = (Block*)malloc(offsetof(Block, data) + new_size);
     if (nullptr == b) {
         return nullptr;
diff --git a/test/BUILD.bazel b/test/BUILD.bazel
index a2c291eb7e..408a8d2510 100644
--- a/test/BUILD.bazel
+++ b/test/BUILD.bazel
@@ -27,6 +27,7 @@ COPTS = [
 ]
 
 TEST_BUTIL_SOURCES = [
+    "arena_unittest.cc",
     "at_exit_unittest.cc",
     "atomicops_unittest.cc",
     "base64_unittest.cc",
diff --git a/test/CMakeLists.txt b/test/CMakeLists.txt
index 4c3a09f529..2120139927 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}/arena_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/arena_unittest.cc b/test/arena_unittest.cc
new file mode 100644
index 0000000000..ff0f73aa7e
--- /dev/null
+++ b/test/arena_unittest.cc
@@ -0,0 +1,43 @@
+// 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 <cstring>
+#include <limits>
+#include "butil/arena.h"
+
+TEST(ArenaTest, rejects_unrepresentable_blocks_and_remains_usable) {
+    butil::Arena arena;
+    char* saved = static_cast<char*>(arena.allocate(16));
+    ASSERT_NE(nullptr, saved);
+    memcpy(saved, "retained", 9);
+    EXPECT_EQ(nullptr, arena.allocate(std::numeric_limits<size_t>::max()));
+    if (sizeof(size_t) > sizeof(uint32_t)) {
+        EXPECT_EQ(nullptr, arena.allocate(static_cast<size_t>(UINT32_MAX) + 1));
+    }
+    EXPECT_STREQ("retained", saved);
+    ASSERT_NE(nullptr, arena.allocate(16));
+    arena.clear();
+    ASSERT_NE(nullptr, arena.allocate(16));
+}
+
+TEST(ArenaTest, rejects_unrepresentable_initial_block) {
+    butil::ArenaOptions options;
+    options.initial_block_size = std::numeric_limits<size_t>::max();
+    butil::Arena arena(options);
+    EXPECT_EQ(nullptr, arena.allocate(1));
+}