Skip to content

Commit

Permalink
In debug mode, after clearing oneof messages on arenas, poison them i…
Browse files Browse the repository at this point in the history
…f ASAN.

PiperOrigin-RevId: 621886314
  • Loading branch information
protobuf-github-bot authored and copybara-github committed Apr 4, 2024
1 parent b4bf6b2 commit 8826baf
Show file tree
Hide file tree
Showing 10 changed files with 113 additions and 16 deletions.
16 changes: 13 additions & 3 deletions src/google/protobuf/arena.cc
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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();
Expand Down Expand Up @@ -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.
Expand Down
43 changes: 43 additions & 0 deletions src/google/protobuf/arena_unittest.cc
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand All @@ -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;
Expand Down
30 changes: 21 additions & 9 deletions src/google/protobuf/compiler/cpp/field_generators/message_field.cc
Original file line number Diff line number Diff line change
Expand Up @@ -593,15 +593,27 @@ void OneofMessage::GenerateInlineAccessorDefinitions(io::Printer* p) const {
}

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

void OneofMessage::GenerateMessageClearingCode(io::Printer* p) const {
Expand Down
2 changes: 1 addition & 1 deletion src/google/protobuf/compiler/cpp/unittest.inc
Original file line number Diff line number Diff line change
Expand Up @@ -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 =
Expand All @@ -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.
Expand Down
6 changes: 6 additions & 0 deletions src/google/protobuf/generated_message_reflection.cc
Original file line number Diff line number Diff line change
Expand Up @@ -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>
Expand Down Expand Up @@ -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"
Expand Down Expand Up @@ -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);
Expand Down
18 changes: 18 additions & 0 deletions src/google/protobuf/message.h
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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_;
Expand Down Expand Up @@ -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>
Expand Down
3 changes: 2 additions & 1 deletion src/google/protobuf/port.h
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
8 changes: 7 additions & 1 deletion src/google/protobuf/proto3_arena_unittest.cc
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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) {
Expand Down
1 change: 0 additions & 1 deletion src/google/protobuf/reflection_visit_fields.h
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down
2 changes: 2 additions & 0 deletions src/google/protobuf/thread_safe_arena.h
Original file line number Diff line number Diff line change
Expand Up @@ -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.

Expand Down

0 comments on commit 8826baf

Please sign in to comment.