From a5e5e02bdf87af3cacc29f839ee0b153eab13c95 Mon Sep 17 00:00:00 2001 From: Protobuf Team Bot Date: Tue, 12 Mar 2024 10:48:02 -0700 Subject: [PATCH] Fix memory leak in the arena when interacting with a type that derives from MessageLite and is marked as arena constructible but not as destructor skippable. The object was not being registered for destruction. This currently only affects ImplicitWeakMessage. PiperOrigin-RevId: 615099201 --- src/google/protobuf/arena.h | 26 +++++--- src/google/protobuf/arena_unittest.cc | 85 +++++++++++++++++++++++++++ 2 files changed, 102 insertions(+), 9 deletions(-) diff --git a/src/google/protobuf/arena.h b/src/google/protobuf/arena.h index 868a90f185..a13fccf93e 100644 --- a/src/google/protobuf/arena.h +++ b/src/google/protobuf/arena.h @@ -195,15 +195,21 @@ class PROTOBUF_EXPORT PROTOBUF_ALIGNAS(8) Arena final { [arena](auto&&... args) { using Type = std::remove_const_t; #ifdef __cpp_if_constexpr - constexpr auto construct_type = GetConstructType(); - // We delegate to DefaultConstruct/CopyConstruct where appropriate - // because protobuf generated classes have external templates for - // these functions for code size reasons. When `if constexpr` is not - // available always use the fallback. - if constexpr (construct_type == ConstructType::kDefault) { - return static_cast(DefaultConstruct(arena)); - } else if constexpr (construct_type == ConstructType::kCopy) { - return static_cast(CopyConstruct(arena, &args...)); + // DefaultConstruct/CopyConstruct are optimized for messages, which + // are both arena constructible and destructor skippable and they + // assume much. Don't use these functions unless the invariants + // hold. + if constexpr (is_destructor_skippable::value) { + constexpr auto construct_type = GetConstructType(); + // We delegate to DefaultConstruct/CopyConstruct where appropriate + // because protobuf generated classes have external templates for + // these functions for code size reasons. When `if constexpr` is not + // available always use the fallback. + if constexpr (construct_type == ConstructType::kDefault) { + return static_cast(DefaultConstruct(arena)); + } else if constexpr (construct_type == ConstructType::kCopy) { + return static_cast(CopyConstruct(arena, &args...)); + } } #endif return CreateArenaCompatible(arena, @@ -651,6 +657,7 @@ class PROTOBUF_EXPORT PROTOBUF_ALIGNAS(8) Arena final { // keyword to make sure the `extern template` suppresses instantiations. template PROTOBUF_NOINLINE void* Arena::DefaultConstruct(Arena* arena) { + static_assert(is_destructor_skippable::value, ""); void* mem = arena != nullptr ? arena->AllocateAligned(sizeof(T)) : ::operator new(sizeof(T)); return new (mem) T(arena); @@ -658,6 +665,7 @@ PROTOBUF_NOINLINE void* Arena::DefaultConstruct(Arena* arena) { template PROTOBUF_NOINLINE void* Arena::CopyConstruct(Arena* arena, const void* from) { + static_assert(is_destructor_skippable::value, ""); void* mem = arena != nullptr ? arena->AllocateAligned(sizeof(T)) : ::operator new(sizeof(T)); return new (mem) T(arena, *static_cast(from)); diff --git a/src/google/protobuf/arena_unittest.cc b/src/google/protobuf/arena_unittest.cc index 99815d9cfc..7e2f0ad75a 100644 --- a/src/google/protobuf/arena_unittest.cc +++ b/src/google/protobuf/arena_unittest.cc @@ -54,6 +54,7 @@ using protobuf_unittest::TestAllTypes; using protobuf_unittest::TestEmptyMessage; using protobuf_unittest::TestOneof2; using protobuf_unittest::TestRepeatedString; +using ::testing::ElementsAreArray; namespace google { namespace protobuf { @@ -158,6 +159,89 @@ TEST(ArenaTest, DestructorSkippable) { EXPECT_FALSE(Arena::is_destructor_skippable::type::value); } +template +struct EmptyBase {}; +struct ArenaCtorBase { + using InternalArenaConstructable_ = void; +}; +struct ArenaDtorBase { + using DestructorSkippable_ = void; +}; + +template +void TestCtorAndDtorTraits(std::vector def, + std::vector copy, + std::vector with_int) { + static auto& actions = *new std::vector; + struct TraitsProber + : std::conditional_t>, + std::conditional_t>, + Message { + TraitsProber() { actions.push_back("()"); } + TraitsProber(const TraitsProber&) { actions.push_back("(const T&)"); } + explicit TraitsProber(int) { actions.push_back("(int)"); } + explicit TraitsProber(Arena* arena) { actions.push_back("(Arena)"); } + TraitsProber(Arena* arena, const TraitsProber&) { + actions.push_back("(Arena, const T&)"); + } + TraitsProber(Arena* arena, int) { actions.push_back("(Arena, int)"); } + ~TraitsProber() override { actions.push_back("~()"); } + + TraitsProber* New(Arena*) const final { + ABSL_LOG(FATAL); + return nullptr; + } + const ClassData* GetClassData() const final { + ABSL_LOG(FATAL); + return nullptr; + } + }; + + static_assert( + !arena_ctor || Arena::is_arena_constructable::value, ""); + static_assert( + !arena_dtor || Arena::is_destructor_skippable::value, ""); + + { + actions.clear(); + Arena arena; + Arena::Create(&arena); + } + EXPECT_THAT(actions, ElementsAreArray(def)); + + const TraitsProber p; + { + actions.clear(); + Arena arena; + Arena::Create(&arena, p); + } + EXPECT_THAT(actions, ElementsAreArray(copy)); + + { + actions.clear(); + Arena arena; + Arena::Create(&arena, 17); + } + EXPECT_THAT(actions, ElementsAreArray(with_int)); +} + +TEST(ArenaTest, AllConstructibleAndDestructibleCombinationsWorkCorrectly) { + TestCtorAndDtorTraits({"()", "~()"}, {"(const T&)", "~()"}, + {"(int)", "~()"}); + // If the object is not arena constructible, then the destructor is always + // called even if marked as skippable. + TestCtorAndDtorTraits({"()", "~()"}, {"(const T&)", "~()"}, + {"(int)", "~()"}); + + // Some types are arena constructible but we can't skip the destructor. Those + // are constructed with an arena but still destroyed. + TestCtorAndDtorTraits({"(Arena)", "~()"}, + {"(Arena, const T&)", "~()"}, + {"(Arena, int)", "~()"}); + TestCtorAndDtorTraits({"(Arena)"}, {"(Arena, const T&)"}, + {"(Arena, int)"}); +} + TEST(ArenaTest, BasicCreate) { Arena arena; EXPECT_TRUE(Arena::Create(&arena) != nullptr); @@ -420,6 +504,7 @@ TEST(ArenaTest, GetConstructTypeWorks) { class DispatcherTestProto : public Message { public: using InternalArenaConstructable_ = void; + using DestructorSkippable_ = void; // For the test below to construct. explicit DispatcherTestProto(absl::in_place_t) {} explicit DispatcherTestProto(Arena*) { ABSL_LOG(FATAL); }