From 54d068e11c77ed387b97a60f435998b384e36e34 Mon Sep 17 00:00:00 2001
From: Tony Liao <tonyliaoss@google.com>
Date: Thu, 24 Oct 2024 17:16:43 -0700
Subject: [PATCH] Breaking Change: Add ASAN poisoning after clearing oneof
 messages on arena.

Note:
  This change primarily affects debug + ASAN builds using protobuf arenas.
  If this change causes a crash in your debug build, it probably means that
  there is a use-after-free bug in your program. This change has already been
  implemented and battle-tested within Google for some time.

  Oneof messages on the regular heap should not be affected because the memory
  they hold are already deleted. Users will already see use-after-free errors
  if they attempt to access heap-allocated oneof messages after calling Clear().

When a protobuf message is cleared, all raw pointers should be invalidated
because undefined things may happen to any of the fields pointed to by
mutable_foo() APIs. While destructors may not necessarily be invoked, Clear()
should be considered a pointer invalidation event.

#test-continuous

PiperOrigin-RevId: 689569669
---
 src/google/protobuf/arena_unittest.cc         |  4 ++
 .../protobuf/generated_message_reflection.cc  | 46 +++++++++++++++++++
 src/google/protobuf/proto3_arena_unittest.cc  |  4 ++
 3 files changed, 54 insertions(+)

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<void*>(&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<MemBlock> nodes;
+
+#ifdef __cpp_if_constexpr
+  nodes.emplace_back(root);
+
+  std::queue<Message*> 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
 }