In debug mode, after clearing oneof messages on arenas, poison them if ASAN.

PiperOrigin-RevId: 621886314

src/google/protobuf/arena.cc

@@ -92,9 +92,6 @@ class GetDeallocator {
         space_allocated_(space_allocated) {}
 
   void operator()(SizedPtr mem) const {
-    // This memory was provided by the underlying allocator as unpoisoned,
-    // so return it in an unpoisoned state.
-    PROTOBUF_UNPOISON_MEMORY_REGION(mem.p, mem.n);
     if (dealloc_) {
       dealloc_(mem.p, mem.n);
     } else {
@@ -666,6 +663,15 @@ void ThreadSafeArena::AddSerialArena(void* id, SerialArena* serial) {
   head_.store(new_head, std::memory_order_release);
 }
 
+void ThreadSafeArena::UnpoisonAllArenaBlocks() const {
+  VisitSerialArena([](const SerialArena* serial) {
+    for (const auto* b = serial->head(); b != nullptr && !b->IsSentry();
+         b = b->next) {
+      PROTOBUF_UNPOISON_MEMORY_REGION(b, b->size);
+    }
+  });
+}
+
 void ThreadSafeArena::Init() {
   tag_and_id_ = GetNextLifeCycleId();
   arena_stats_ = Sample();
@@ -868,6 +874,10 @@ template void*
     ThreadSafeArena::AllocateAlignedFallback<AllocationClient::kArray>(size_t);
 
 void ThreadSafeArena::CleanupList() {
+#ifdef PROTOBUF_ASAN
+  UnpoisonAllArenaBlocks();
+#endif
+
   WalkSerialArenaChunk([](SerialArenaChunk* chunk) {
     absl::Span<std::atomic<SerialArena*>> span = chunk->arenas();
     // Walks arenas backward to handle the first serial arena the last.

src/google/protobuf/arena_unittest.cc

@@ -1506,6 +1506,45 @@ TEST(ArenaTest, MutableMessageReflection) {
 #endif  // PROTOBUF_RTTI
 
 
+TEST(ArenaTest, ClearOneofMessageOnArena) {
+  if (!internal::DebugHardenClearOneofMessageOnArena()) {
+    GTEST_SKIP() << "arena allocated oneof message fields are not hardened.";
+  }
+
+  Arena arena;
+  auto* message = Arena::Create<unittest::TestOneof2>(&arena);
+  // Intentionally nested to force poisoning recursively to catch the access.
+  auto* child =
+      message->mutable_foo_message()->mutable_child()->mutable_child();
+  child->set_moo_int(100);
+  message->clear_foo_message();
+
+#ifndef PROTOBUF_ASAN
+  EXPECT_NE(child->moo_int(), 100);
+#else
+#if GTEST_HAS_DEATH_TEST && defined(__cpp_if_constexpr)
+  EXPECT_DEATH(EXPECT_EQ(child->moo_int(), 0), "use-after-poison");
+#endif
+#endif
+}
+
+TEST(ArenaTest, CopyValuesWithinOneof) {
+  if (!internal::DebugHardenClearOneofMessageOnArena()) {
+    GTEST_SKIP() << "arena allocated oneof message fields are not hardened.";
+  }
+
+  Arena arena;
+  auto* message = Arena::Create<unittest::TestOneof>(&arena);
+  auto* foo = message->mutable_foogroup();
+  foo->set_a(100);
+  foo->set_b("hello world");
+  message->set_foo_string(message->foogroup().b());
+
+  // As a debug hardening measure, `set_foo_string` would clear `foo` in
+  // (!NDEBUG && !ASAN) and the copy wouldn't work.
+  EXPECT_TRUE(message->foo_string().empty()) << message->foo_string();
+}
+
 void FillArenaAwareFields(TestAllTypes* message) {
   std::string test_string = "hello world";
   message->set_optional_int32(42);
@@ -1526,6 +1565,10 @@ void FillArenaAwareFields(TestAllTypes* message) {
 
 // Test: no allocations occur on heap while touching all supported field types.
 TEST(ArenaTest, NoHeapAllocationsTest) {
+  if (internal::DebugHardenClearOneofMessageOnArena()) {
+    GTEST_SKIP() << "debug hardening may cause heap allocation.";
+  }
+
   // Allocate a large initial block to avoid mallocs during hooked test.
   std::vector<char> arena_block(128 * 1024);
   ArenaOptions options;

src/google/protobuf/compiler/cpp/field_generators/message_field.cc

@@ -593,13 +593,25 @@ void OneofMessage::GenerateInlineAccessorDefinitions(io::Printer* p) const {
 }
 
 void OneofMessage::GenerateClearingCode(io::Printer* p) const {
+  p->Emit({{"poison_or_clear",
+            [&] {
+              if (HasDescriptorMethods(field_->file(), options_)) {
+                p->Emit(R"cc(
+                  $pbi$::MaybePoisonAfterClear($field_$);
+                )cc");
+              } else {
                 p->Emit(R"cc(
-    if (GetArena() == nullptr) {
-      delete $field_$;
-    } else if ($pbi$::DebugHardenClearOneofMessageOnArena()) {
                   if ($field_$ != nullptr) {
                     $field_$->Clear();
                   }
+                )cc");
+              }
+            }}},
+          R"cc(
+            if (GetArena() == nullptr) {
+              delete $field_$;
+            } else if ($pbi$::DebugHardenClearOneofMessageOnArena()) {
+              $poison_or_clear$;
             }
           )cc");
 }

src/google/protobuf/compiler/cpp/unittest.inc

@@ -552,6 +552,7 @@ TEST(GENERATED_MESSAGE_TEST_NAME, CopyConstructor) {
   }
 }
 
+#ifndef PROTOBUF_TEST_NO_DESCRIPTORS
 TEST(GENERATED_MESSAGE_TEST_NAME, CopyConstructorWithArenas) {
   Arena arena;
   UNITTEST::TestAllTypes* message1 =
@@ -572,7 +573,6 @@ TEST(GENERATED_MESSAGE_TEST_NAME, CopyConstructorWithArenas) {
   TestUtil::ExpectAllFieldsSet(*message2_heap);
 }
 
-#ifndef PROTOBUF_TEST_NO_DESCRIPTORS
 TEST(GENERATED_MESSAGE_TEST_NAME, UpcastCopyFrom) {
   // Test the CopyFrom method that takes in the generic const Message&
   // parameter.

src/google/protobuf/generated_message_reflection.cc

@@ -16,6 +16,7 @@
 #include <cstdint>
 #include <cstring>
 #include <new>  // IWYU pragma: keep for operator delete
+#include <queue>
 #include <string>
 #include <type_traits>
 #include <utility>
@@ -43,6 +44,7 @@
 #include "google/protobuf/message_lite.h"
 #include "google/protobuf/port.h"
 #include "google/protobuf/raw_ptr.h"
+#include "google/protobuf/reflection_visit_fields.h"
 #include "google/protobuf/repeated_field.h"
 #include "google/protobuf/repeated_ptr_field.h"
 #include "google/protobuf/unknown_field_set.h"
@@ -1307,6 +1309,10 @@ void Reflection::InternalSwap(Message* lhs, Message* rhs) const {
   }
 }
 
+void Reflection::MaybePoisonAfterClear(Message& root) const {
+  root.Clear();
+}
+
 int Reflection::FieldSize(const Message& message,
                           const FieldDescriptor* field) const {
   USAGE_CHECK_MESSAGE(FieldSize, &message);

src/google/protobuf/message.h

@@ -226,6 +226,8 @@ bool CreateUnknownEnumValues(const FieldDescriptor* field);
 
 // Returns true if "message" is a descendant of "root".
 PROTOBUF_EXPORT bool IsDescendant(Message& root, const Message& message);
+
+inline void MaybePoisonAfterClear(Message* root);
 }  // namespace internal
 
 // Abstract interface for protocol messages.
@@ -1030,10 +1032,16 @@ class PROTOBUF_EXPORT Reflection final {
     return schema_.IsSplit(field);
   }
 
+  // Walks the message tree from "root" and poisons (under ASAN) the memory to
+  // force subsequent accesses to fail. Always calls Clear beforehand to clear
+  // strings, etc.
+  void MaybePoisonAfterClear(Message& root) const;
+
   friend class FastReflectionBase;
   friend class FastReflectionMessageMutator;
   friend class internal::ReflectionVisit;
   friend bool internal::IsDescendant(Message& root, const Message& message);
+  friend void internal::MaybePoisonAfterClear(Message* root);
 
   const Descriptor* const descriptor_;
   const internal::ReflectionSchema schema_;
@@ -1626,6 +1634,16 @@ class RawMessageBase : public Message {
   virtual size_t SpaceUsedLong() const = 0;
 };
 
+inline void MaybePoisonAfterClear(Message* root) {
+  if (root == nullptr) return;
+#ifndef PROTOBUF_ASAN
+  root->Clear();
+#else
+  const Reflection* reflection = root->GetReflection();
+  reflection->MaybePoisonAfterClear(*root);
+#endif
+}
+
 }  // namespace internal
 
 template <typename Type>

src/google/protobuf/port.h

@@ -239,7 +239,8 @@ inline constexpr bool DebugHardenStringValues() {
 #endif
 }
 
-// Returns true if force clearing oneof message on arena is enabled.
+// Returns true if debug hardening for clearing oneof message on arenas is
+// enabled.
 inline constexpr bool DebugHardenClearOneofMessageOnArena() {
 #ifdef NDEBUG
   return false;

src/google/protobuf/proto3_arena_unittest.cc

@@ -273,7 +273,7 @@ TEST(Proto3ArenaTest, CheckMessageFieldIsCleared) {
 
 TEST(Proto3ArenaTest, CheckOneofMessageFieldIsCleared) {
   if (!internal::DebugHardenClearOneofMessageOnArena()) {
-    GTEST_SKIP() << "arena allocated oneof message fields are not cleared.";
+    GTEST_SKIP() << "arena allocated oneof message fields are not hardened.";
   }
 
   Arena arena;
@@ -286,7 +286,13 @@ TEST(Proto3ArenaTest, CheckOneofMessageFieldIsCleared) {
   child->set_bb(100);
   msg->Clear();
 
+#ifndef PROTOBUF_ASAN
   EXPECT_EQ(child->bb(), 0);
+#else
+#if GTEST_HAS_DEATH_TEST && defined(__cpp_if_constexpr)
+  EXPECT_DEATH(EXPECT_EQ(child->bb(), 100), "use-after-poison");
+#endif
+#endif
 }
 
 TEST(Proto3OptionalTest, OptionalFieldDescriptor) {

src/google/protobuf/reflection_visit_fields.h

@@ -302,7 +302,6 @@ void ReflectionVisit::VisitFields(MessageT& message, CallbackFn&& func,
   auto& set = ExtensionSet(reflection, message);
   auto* extendee = reflection->descriptor_;
   auto* pool = reflection->descriptor_pool_;
-  auto* arena = message.GetArena();
 
   set.ForEach([&](int number, auto& ext) {
     ABSL_DCHECK_GT(ext.type, 0);

src/google/protobuf/thread_safe_arena.h

@@ -129,6 +129,8 @@ class PROTOBUF_EXPORT ThreadSafeArena {
   // Adds SerialArena to the chunked list. May create a new chunk.
   void AddSerialArena(void* id, SerialArena* serial);
 
+  void UnpoisonAllArenaBlocks() const;
+
   // Members are declared here to track sizeof(ThreadSafeArena) and hotness
   // centrally.