diff --git a/src/google/protobuf/arena_unittest.cc b/src/google/protobuf/arena_unittest.cc index 5923cce62b59c..ca04704a56f43 100644 --- a/src/google/protobuf/arena_unittest.cc +++ b/src/google/protobuf/arena_unittest.cc @@ -1541,6 +1541,10 @@ TEST(ArenaTest, ClearOneofMessageOnArena) { #ifndef PROTOBUF_ASAN EXPECT_NE(child->moo_int(), 100); +#else +#if GTEST_HAS_DEATH_TEST + EXPECT_DEATH(EXPECT_EQ(child->moo_int(), 0), "use-after-poison"); +#endif // !GTEST_HAS_DEATH_TEST #endif // !PROTOBUF_ASAN } diff --git a/src/google/protobuf/generated_message_reflection.cc b/src/google/protobuf/generated_message_reflection.cc index 6857caf65fba2..9cea6662aa1d0 100644 --- a/src/google/protobuf/generated_message_reflection.cc +++ b/src/google/protobuf/generated_message_reflection.cc @@ -1330,7 +1330,53 @@ void Reflection::InternalSwap(Message* lhs, Message* rhs) const { } void Reflection::MaybePoisonAfterClear(Message& root) const { + struct MemBlock { + explicit MemBlock(Message& msg) + : ptr(static_cast(&msg)), size(GetSize(msg)) {} + + static uint32_t GetSize(const Message& msg) { + return msg.GetReflection()->schema_.GetObjectSize(); + } + + void* ptr; + uint32_t size; + }; + + bool heap_alloc = root.GetArena() == nullptr; + std::vector nodes; + +#ifdef __cpp_if_constexpr + nodes.emplace_back(root); + + std::queue queue; + queue.push(&root); + + while (!queue.empty() && !heap_alloc) { + Message* curr = queue.front(); + queue.pop(); + internal::VisitMutableMessageFields(*curr, [&](Message& msg) { + if (msg.GetArena() == nullptr) { + heap_alloc = true; + return; + } + + nodes.emplace_back(msg); + // Also visits child messages. + queue.push(&msg); + }); + } +#endif + root.Clear(); + + // Heap allocated oneof messages will be freed on clear. So, poisoning + // afterwards may cause use-after-free. Bailout. + if (heap_alloc) return; + + for (auto it : nodes) { + (void)it; + PROTOBUF_POISON_MEMORY_REGION(it.ptr, it.size); + } } int Reflection::FieldSize(const Message& message, diff --git a/src/google/protobuf/proto3_arena_unittest.cc b/src/google/protobuf/proto3_arena_unittest.cc index f1ee7ed122c47..c27e7cda85630 100644 --- a/src/google/protobuf/proto3_arena_unittest.cc +++ b/src/google/protobuf/proto3_arena_unittest.cc @@ -289,6 +289,10 @@ TEST(Proto3ArenaTest, CheckOneofMessageFieldIsCleared) { #ifndef PROTOBUF_ASAN EXPECT_EQ(child->bb(), 0); +#else +#if GTEST_HAS_DEATH_TEST + EXPECT_DEATH(EXPECT_EQ(child->bb(), 100), "use-after-poison"); +#endif // !GTEST_HAS_DEATH_TEST #endif // !PROTOBUF_ASAN }