Fix Python cpp memory crash when MergeFrom Oneof.

Maybe release oneof messages before MergeFrom in python cpp extension

PiperOrigin-RevId: 819802492

Author: anandolee

python/google/protobuf/internal/message_test.py

@@ -1001,6 +1001,50 @@ def testOneofMessageMergeFrom(self, message_module):
     self.assertEqual('oneof_nested_message',
                      m2.child.payload.WhichOneof('oneof_field'))
 
+  def testOneofReleaseMergeFrom(self, message_module):
+    m = unittest_pb2.TestOneof2()
+    m.foo_message.moo_int = 123
+    reference = m.foo_message
+    self.assertEqual(m.foo_message.moo_int, 123)
+    m2 = unittest_pb2.TestOneof2()
+    m2.foo_lazy_message.moo_int = 456
+    m.MergeFrom(m2)
+    self.assertEqual(reference.moo_int, 123)
+    self.assertEqual(m.foo_message.moo_int, 0)
+    m.foo_message.CopyFrom(reference)
+    self.assertEqual(m.foo_message.moo_int, 123)
+
+  def testNestedOneofRleaseMergeFrom(self, message_module):
+    m = message_module.NestedTestAllTypes()
+    m.payload.oneof_nested_message.bb = 1
+    m.child.payload.oneof_nested_message.bb = 2
+    ref1 = m.payload.oneof_nested_message
+    ref2 = m.child.payload.oneof_nested_message
+    other = message_module.NestedTestAllTypes()
+    other.payload.oneof_uint32 = 22
+    other.child.payload.oneof_string = 'hi'
+    self.assertEqual(ref1.bb, 1)
+    self.assertEqual(ref2.bb, 2)
+    m.MergeFrom(other)
+    # oneof messages are released
+    self.assertEqual(ref1.bb, 1)
+    self.assertEqual(ref2.bb, 2)
+    self.assertEqual(m.payload.oneof_nested_message.bb, 0)
+    self.assertEqual(m.child.payload.oneof_nested_message.bb, 0)
+    self.assertEqual(m.payload.oneof_uint32, 22)
+    self.assertEqual(m.child.payload.oneof_string, 'hi')
+
+  def testOneofNotReleaseMergeFrom(self, message_module):
+    m = message_module.NestedTestAllTypes()
+    m.payload.oneof_nested_message.bb = 1
+    ref = m.payload.oneof_nested_message
+    other = message_module.NestedTestAllTypes()
+    other.payload.oneof_nested_message.bb = 2
+    self.assertEqual(ref.bb, 1)
+    m.MergeFrom(other)
+    # oneof message is not released
+    self.assertEqual(ref.bb, 2)
+
   def testOneofNestedMessageInit(self, message_module):
     m = message_module.TestAllTypes(
         oneof_nested_message=message_module.TestAllTypes.NestedMessage())

python/google/protobuf/pyext/message.cc

@@ -764,6 +764,50 @@ static int MaybeReleaseOverlappingOneofField(CMessage* cmessage,
   return 0;
 }
 
+int MaybeReleaseOneofBeforeMerge(CMessage* self, const Message& other) {
+  if (!self->composite_fields) {
+    return 0;
+  }
+
+  Message* message = self->message;
+  const Reflection* reflection = message->GetReflection();
+  PyMessageFactory* factory = GetFactoryForMessage(self);
+  std::vector<const FieldDescriptor*> fields_to_release;
+  std::vector<const FieldDescriptor*> nested_message_fields;
+  for (const auto& item : *self->composite_fields) {
+    const FieldDescriptor* descriptor = item.first;
+    if (descriptor->cpp_type() == FieldDescriptor::CPPTYPE_MESSAGE &&
+        // For normal repeated message, MergeFrom will append the messages.
+        // For message map with same keys, it is overwrite
+        !descriptor->is_repeated() &&
+        reflection->HasField(*message, descriptor)) {
+      if (reflection->HasField(other, descriptor)) {
+        nested_message_fields.push_back(descriptor);
+      } else {
+        // Release oneof message if the other message has set a different oneof
+        const OneofDescriptor* oneof = descriptor->containing_oneof();
+        if (oneof && reflection->HasOneof(other, oneof)) {
+          fields_to_release.push_back(descriptor);
+        }
+      }
+    }
+  }
+  for (const FieldDescriptor* field : nested_message_fields) {
+    if (MaybeReleaseOneofBeforeMerge(
+            reinterpret_cast<CMessage*>(
+                self->composite_fields->find(field)->second),
+            reflection->GetMessage(other, field, factory->message_factory)) <
+        0) {
+      return -1;
+    }
+  }
+  for (const FieldDescriptor* field : fields_to_release) {
+    if (InternalReleaseFieldByDescriptor(self, field) < 0) {
+      return -1;
+    }
+  }
+  return 0;
+}
 // After a Merge, visit every sub-message that was read-only, and
 // eventually update their pointer if the Merge operation modified them.
 int FixupMessageAfterMerge(CMessage* self) {
@@ -1846,6 +1890,10 @@ PyObject* MergeFrom(CMessage* self, PyObject* arg) {
   }
   AssureWritable(self);
 
+  if (MaybeReleaseOneofBeforeMerge(self, *other_message->message) < 0) {
+    return nullptr;
+  }
+
   self->message->MergeFrom(*other_message->message);
   // Child message might be lazily created before MergeFrom. Make sure they
   // are mutable at this point if child messages are really created.