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

@@ -195,15 +195,21 @@ class PROTOBUF_EXPORT PROTOBUF_ALIGNAS(8) Arena final {
         [arena](auto&&... args) {
           using Type = std::remove_const_t<T>;
 #ifdef __cpp_if_constexpr
-          constexpr auto construct_type = GetConstructType<T, Args&&...>();
-          // 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<Type*>(DefaultConstruct<Type>(arena));
-          } else if constexpr (construct_type == ConstructType::kCopy) {
-            return static_cast<Type*>(CopyConstruct<Type>(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<T>::value) {
+            constexpr auto construct_type = GetConstructType<T, Args&&...>();
+            // 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<Type*>(DefaultConstruct<Type>(arena));
+            } else if constexpr (construct_type == ConstructType::kCopy) {
+              return static_cast<Type*>(CopyConstruct<Type>(arena, &args...));
+            }
           }
 #endif
           return CreateArenaCompatible<Type>(arena,
@@ -651,6 +657,7 @@ class PROTOBUF_EXPORT PROTOBUF_ALIGNAS(8) Arena final {
 // keyword to make sure the `extern template` suppresses instantiations.
 template <typename T>
 PROTOBUF_NOINLINE void* Arena::DefaultConstruct(Arena* arena) {
+  static_assert(is_destructor_skippable<T>::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 <typename T>
 PROTOBUF_NOINLINE void* Arena::CopyConstruct(Arena* arena, const void* from) {
+  static_assert(is_destructor_skippable<T>::value, "");
   void* mem = arena != nullptr ? arena->AllocateAligned(sizeof(T))
                                : ::operator new(sizeof(T));
   return new (mem) T(arena, *static_cast<const T*>(from));

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<Arena>::type::value);
 }
 
+template <int>
+struct EmptyBase {};
+struct ArenaCtorBase {
+  using InternalArenaConstructable_ = void;
+};
+struct ArenaDtorBase {
+  using DestructorSkippable_ = void;
+};
+
+template <bool arena_ctor, bool arena_dtor>
+void TestCtorAndDtorTraits(std::vector<absl::string_view> def,
+                           std::vector<absl::string_view> copy,
+                           std::vector<absl::string_view> with_int) {
+  static auto& actions = *new std::vector<absl::string_view>;
+  struct TraitsProber
+      : std::conditional_t<arena_ctor, ArenaCtorBase, EmptyBase<0>>,
+        std::conditional_t<arena_dtor, ArenaDtorBase, EmptyBase<1>>,
+        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<TraitsProber>::value, "");
+  static_assert(
+      !arena_dtor || Arena::is_destructor_skippable<TraitsProber>::value, "");
+
+  {
+    actions.clear();
+    Arena arena;
+    Arena::Create<TraitsProber>(&arena);
+  }
+  EXPECT_THAT(actions, ElementsAreArray(def));
+
+  const TraitsProber p;
+  {
+    actions.clear();
+    Arena arena;
+    Arena::Create<TraitsProber>(&arena, p);
+  }
+  EXPECT_THAT(actions, ElementsAreArray(copy));
+
+  {
+    actions.clear();
+    Arena arena;
+    Arena::Create<TraitsProber>(&arena, 17);
+  }
+  EXPECT_THAT(actions, ElementsAreArray(with_int));
+}
+
+TEST(ArenaTest, AllConstructibleAndDestructibleCombinationsWorkCorrectly) {
+  TestCtorAndDtorTraits<false, false>({"()", "~()"}, {"(const T&)", "~()"},
+                                      {"(int)", "~()"});
+  // If the object is not arena constructible, then the destructor is always
+  // called even if marked as skippable.
+  TestCtorAndDtorTraits<false, true>({"()", "~()"}, {"(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<true, false>({"(Arena)", "~()"},
+                                     {"(Arena, const T&)", "~()"},
+                                     {"(Arena, int)", "~()"});
+  TestCtorAndDtorTraits<true, true>({"(Arena)"}, {"(Arena, const T&)"},
+                                    {"(Arena, int)"});
+}
+
 TEST(ArenaTest, BasicCreate) {
   Arena arena;
   EXPECT_TRUE(Arena::Create<int32_t>(&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); }