Fix SetExtension to fuse the right arenas.

Also, make sure the operation has no side effects if the arenas can't be fused.

PiperOrigin-RevId: 561378595

upb/protos/BUILD

@@ -75,13 +75,12 @@ cc_library(
         ":protos_extension_lock",
         "//:base",
         "//:mem",
-        "//:message",
+        "//:message_accessors_internal",
         "//:message_copy",
         "//:message_internal",
         "//:message_promote",
         "//:mini_table",
         "//:wire",
-        "@com_google_absl//absl/base:core_headers",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
         "@com_google_absl//absl/strings",

upb/protos/protos.h

@@ -39,6 +39,7 @@
 #include "upb/base/status.hpp"
 #include "upb/mem/arena.hpp"
 #include "upb/message/copy.h"
+#include "upb/message/internal/accessors.h"
 #include "upb/message/internal/extension.h"
 #include "upb/wire/decode.h"
 #include "upb/wire/encode.h"
@@ -371,15 +372,24 @@ absl::Status SetExtension(
     const ::protos::internal::ExtensionIdentifier<Extendee, Extension>& id,
     Extension& value) {
   static_assert(!std::is_const_v<T>);
-  auto* message_arena = static_cast<upb_Arena*>(message->GetInternalArena());
+  auto* message_arena = internal::GetArena(message);
   upb_Message_Extension* msg_ext = _upb_Message_GetOrCreateExtension(
       internal::GetInternalMsg(message), id.mini_table_ext(), message_arena);
   if (!msg_ext) {
     return MessageAllocationError();
   }
-  auto* extension_arena = static_cast<upb_Arena*>(message->GetInternalArena());
+  auto* extension_arena = internal::GetArena(&value);
   if (message_arena != extension_arena) {
-    upb_Arena_Fuse(message_arena, extension_arena);
+    if (!upb_Arena_Fuse(message_arena, extension_arena)) {
+      // We have to undo the Create part. Otherwise ,we end up with a broken
+      // extension. We do fuse last because we can undo Create, but we can't
+      // undo Fuse.
+      if (msg_ext->data.ptr == nullptr) {
+        _upb_Message_ClearExtensionField(internal::GetInternalMsg(message),
+                                         id.mini_table_ext());
+      }
+      return absl::InvalidArgumentError("Unable to fuse arenas.");
+    }
   }
   msg_ext->data.ptr = internal::GetInternalMsg(&value);
   return absl::OkStatus();

upb/protos_generator/tests/test_generated.cc

@@ -57,6 +57,7 @@ using ::protos_generator::test::protos::TestModel_Category_VIDEO;
 using ::protos_generator::test::protos::theme;
 using ::protos_generator::test::protos::ThemeExtension;
 using ::testing::ElementsAre;
+using ::testing::HasSubstr;
 
 TEST(CppGeneratedCode, Constructor) { TestModel test_model; }
 
@@ -682,11 +683,34 @@ TEST(CppGeneratedCode, ClearExtensionWithEmptyExtensionPtr) {
 
 TEST(CppGeneratedCode, SetExtension) {
   TestModel model;
+  {
+    // Use a nested scope to make sure the arenas are fused correctly.
+    ThemeExtension extension1;
+    extension1.set_ext_name("Hello World");
+    EXPECT_EQ(false, ::protos::HasExtension(&model, theme));
+    EXPECT_EQ(true, ::protos::SetExtension(&model, theme, extension1).ok());
+  }
+  EXPECT_EQ(true, ::protos::HasExtension(&model, theme));
+  auto ext = ::protos::GetExtension(&model, theme);
+  EXPECT_TRUE(ext.ok());
+  EXPECT_EQ((*ext)->ext_name(), "Hello World");
+}
+
+TEST(CppGeneratedCode, SetExtensionFailsFusing) {
+  // Use an initial block to disallow fusing.
+  char initial_block[1000];
+  protos::Arena arena(initial_block, sizeof(initial_block));
+
+  protos::Ptr<TestModel> model = protos::CreateMessage<TestModel>(arena);
+
   ThemeExtension extension1;
   extension1.set_ext_name("Hello World");
-  EXPECT_EQ(false, ::protos::HasExtension(&model, theme));
-  EXPECT_EQ(true, ::protos::SetExtension(&model, theme, extension1).ok());
-  EXPECT_EQ(true, ::protos::HasExtension(&model, theme));
+  EXPECT_FALSE(::protos::HasExtension(model, theme));
+  auto status = ::protos::SetExtension(model, theme, extension1);
+  EXPECT_FALSE(status.ok());
+  EXPECT_THAT(status.message(), HasSubstr("Unable to fuse arenas."));
+  EXPECT_FALSE(::protos::HasExtension(model, theme));
+  EXPECT_FALSE(::protos::GetExtension(model, theme).ok());
 }
 
 TEST(CppGeneratedCode, SetExtensionOnMutableChild) {