Fix subscript access for free threading.

Subscript is a read-only operation and it should not try to do mutable operations.

Fix AssureWritable to also work on repeated sub messages. It now keeps an index hint to avoid an O(n) lookup if possible.

Fix BuildSubMessageFromPointer to not create MUTABLE messages. It is only taking a non-mutable one.

Fix GetItem to not perform mutable operations. This fixes the race in Subscript.

PiperOrigin-RevId: 967832766

python/google/protobuf/internal/descriptor_test.py

@@ -452,6 +452,67 @@ class DescriptorTest(unittest.TestCase):
     self.assertEqual(immutable_map.get('nonexistent_key'), None)
     self.assertEqual(immutable_map.get('nonexistent_key', 999), 999)
 
+  def testImmutableMessageMapLookup(self):
+    complex_opt1 = unittest_custom_options_pb2.complex_opt1
+    complex_options_msg = (
+        unittest_custom_options_pb2.VariousComplexOptions.DESCRIPTOR.GetOptions()
+    )
+    immutable_map = complex_options_msg.Extensions[complex_opt1].submsg_map
+
+    # Test lookups.
+    self.assertEqual(immutable_map['sub_key'].moo, 555)
+    self.assertIn('sub_key', immutable_map)
+    self.assertNotIn('nonexistent_key', immutable_map)
+    self.assertEqual(len(immutable_map), 1)
+
+    # Test lookups via bytes.
+    self.assertEqual(immutable_map[b'sub_key'].moo, 555)
+    self.assertIn(b'sub_key', immutable_map)
+    self.assertNotIn(b'nonexistent_key', immutable_map)
+
+    # Test iteration.
+    self.assertEqual(set(immutable_map.keys()), {'sub_key'})
+    self.assertEqual([item[1].moo for item in immutable_map.items()], [555])
+
+    # Test get().
+    self.assertEqual(immutable_map.get('sub_key').moo, 555)
+    self.assertIsNone(immutable_map.get('nonexistent_key'))
+    default_obj = object()
+    self.assertIs(
+        immutable_map.get('nonexistent_key', default_obj), default_obj
+    )
+
+    # Test get() with bytes.
+    self.assertEqual(immutable_map.get(b'sub_key').moo, 555)
+    self.assertIsNone(immutable_map.get(b'nonexistent_key'))
+
+    # Test text formatting on frozen message with message map.
+    text = text_format.MessageToString(complex_options_msg)
+    self.assertIn('sub_key', text)
+    self.assertIn('555', text)
+
+    # Verify the message, map, and elements are frozen (immutable).
+    with self.assertRaises(message.FrozenInstanceError):
+      immutable_map.clear()
+    with self.assertRaises(message.FrozenInstanceError):
+      del immutable_map['sub_key']
+    with self.assertRaises(message.FrozenInstanceError):
+      immutable_map.MergeFrom(immutable_map)
+    with self.assertRaises(message.FrozenInstanceError):
+      _ = immutable_map['missing_key']
+    with self.assertRaises(message.FrozenInstanceError):
+      immutable_map['sub_key'].moo = 999
+    with self.assertRaises(message.FrozenInstanceError):
+      immutable_map['sub_key'].Clear()
+    with self.assertRaises(message.FrozenInstanceError):
+      immutable_map['new_key'].moo = 999
+    with self.assertRaises((message.FrozenInstanceError, ValueError)):
+      immutable_map['new_key'] = (
+          unittest_custom_options_pb2.ComplexOptionType3()
+      )
+    with self.assertRaises(message.FrozenInstanceError):
+      complex_options_msg.Extensions[complex_opt1].ClearField('submsg_map')
+
   def testSimpleCustomOptions(self):
     file_descriptor = unittest_custom_options_pb2.DESCRIPTOR
     message_descriptor = (

python/google/protobuf/internal/message_test.py

@@ -1101,6 +1101,52 @@ class MessageTest(unittest.TestCase):
         [k.bb for k in message.repeated_nested_message], [6, 5, 4, 3, 2, 1]
     )
 
+  def testRepeatedCompositeSubscriptMutation(self, message_module):
+    """Check that accessing repeated composite items via subscript and mutating works."""
+    msg = message_module.TestAllTypes()
+    msg.repeated_nested_message.add(bb=1)
+    msg.repeated_nested_message.add(bb=2)
+    serialized = msg.SerializeToString()
+
+    msg2 = message_module.TestAllTypes()
+    msg2.ParseFromString(serialized)
+    item0 = msg2.repeated_nested_message[0]
+    item1 = msg2.repeated_nested_message[1]
+    # item0 and item1 start as default/lazy submessages.
+    # Mutating them promotes them in-place.
+    item0.bb = 10
+    item1.bb = 20
+    self.assertEqual(msg2.repeated_nested_message[0].bb, 10)
+    self.assertEqual(msg2.repeated_nested_message[1].bb, 20)
+    self.assertEqual(item0.bb, 10)
+    self.assertEqual(item1.bb, 20)
+
+  def testSortingRepeatedCompositeFieldsWithSubscriptReferences(
+      self, message_module
+  ):
+    """Check sorting repeated composite fields after retrieving elements via subscript."""
+    msg = message_module.TestAllTypes()
+    msg.repeated_nested_message.add(bb=30)
+    msg.repeated_nested_message.add(bb=10)
+    msg.repeated_nested_message.add(bb=20)
+    serialized = msg.SerializeToString()
+
+    msg2 = message_module.TestAllTypes()
+    msg2.ParseFromString(serialized)
+    ref0 = msg2.repeated_nested_message[0]
+    ref1 = msg2.repeated_nested_message[1]
+    ref2 = msg2.repeated_nested_message[2]
+
+    msg2.repeated_nested_message.sort(key=operator.attrgetter('bb'))
+    self.assertEqual([k.bb for k in msg2.repeated_nested_message], [10, 20, 30])
+
+    ref0.bb = 300
+    ref1.bb = 100
+    ref2.bb = 200
+    self.assertEqual(
+        [k.bb for k in msg2.repeated_nested_message], [100, 200, 300]
+    )
+
   def testRepeatedScalarFieldSortArguments(self, message_module):
     """Check sorting a scalar field using list.sort() arguments."""
     message = message_module.TestAllTypes()
@@ -2996,6 +3042,24 @@ class Proto3Test(unittest.TestCase):
         ('{-456: , 123: c: 1\n}', '{123: c: 1\n, -456: }'),
     )
 
+  def testMessageMapMutationAfterReprAndIteration(self):
+    msg = map_unittest_pb2.TestMap()
+    msg.map_int32_foreign_message[123].c = 1
+    msg.map_int32_foreign_message[456].c = 2
+
+    # Formatting/repr exercises const iteration over map elements.
+    _ = str(msg.map_int32_foreign_message)
+    _ = repr(msg.map_int32_foreign_message)
+    _ = str(msg)
+
+    # Mutating existing and new entries after repr/str must succeed.
+    msg.map_int32_foreign_message[123].c = 10
+    msg.map_int32_foreign_message[456].c = 20
+    msg.map_int32_foreign_message[789].c = 30
+    self.assertEqual(msg.map_int32_foreign_message[123].c, 10)
+    self.assertEqual(msg.map_int32_foreign_message[456].c, 20)
+    self.assertEqual(msg.map_int32_foreign_message[789].c, 30)
+
   def testNestedMessageMapItemDelete(self):
     msg = map_unittest_pb2.TestMap()
     msg.map_int32_all_types[1].optional_nested_message.bb = 1

python/google/protobuf/internal/test_proto2.proto

@@ -56,3 +56,11 @@ message NestedMessageForFT {
 message ReproMessageForLazy {
   optional NestedMessageForFT lazy_field = 1 [lazy = true];
 }
+
+message MessageWithRepeatedComposite {
+  repeated NestedMessageForFT items = 1;
+}
+
+message ContainerForRepeatedComposite {
+  optional MessageWithRepeatedComposite submessage = 1;
+}

python/google/protobuf/internal/thread_safe_test.py

@@ -555,6 +555,40 @@ class FreeThreadingTest(unittest.TestCase):
     for _ in range(500):
       RunRace()
 
+  @unittest.skipIf(
+      api_implementation.Type() == 'upb',
+      'Upb has not been fixed to handle this case.',
+  )
+  def testConcurrentRepeatedCompositeSubscript(self):
+    msg = test_proto2_pb2.ContainerForRepeatedComposite()
+    msg.submessage.items.add(value='foo')
+    msg.submessage.items.add(value='bar')
+    serialized = msg.SerializeToString()
+
+    def RunRace():
+      shared_msg = test_proto2_pb2.ContainerForRepeatedComposite.FromString(
+          serialized
+      )
+      barrier = threading.Barrier(2)
+
+      def Thread1():
+        barrier.wait()
+        _ = shared_msg.submessage.items[0].value
+
+      def Thread2():
+        barrier.wait()
+        _ = shared_msg.submessage.items[1].value
+
+      t1 = threading.Thread(target=Thread1)
+      t2 = threading.Thread(target=Thread2)
+      t1.start()
+      t2.start()
+      t1.join()
+      t2.join()
+
+    for _ in range(500):
+      RunRace()
+
 
 if __name__ == '__main__':
   unittest.main()

python/google/protobuf/pyext/map_container.cc

@@ -577,12 +577,12 @@ static MessageMapContainer* GetMessageMap(PyObject* obj) {
   return reinterpret_cast<MessageMapContainer*>(obj);
 }
 
-static PyObject* GetCMessage(MessageMapContainer* self,
-                             const Message* message) {
-  // Get or create the CMessage object corresponding to this message.
+// Get or create the CMessage object corresponding to this message.
+static PyObject* GetCMessage(MessageMapContainer* self, const Message* message,
+                             MessageMutabilityState state) {
   return self->parent
       ->BuildSubMessageFromPointer(self->parent_field_descriptor, message,
-                                   self->message_class)
+                                   self->message_class, state)
       ->AsPyObject();
 }
 
@@ -662,26 +662,42 @@ int MapReflectionFriend::MessageMapSetItem(PyObject* _self, PyObject* key,
   }
 }
 
+// For mutable messages, subscript access has get-or-create semantics: looking
+// up an existing entry or inserting a new default entry, returning a
+// MESSAGE_MUTABLE wrapper. For frozen messages, mutation is disallowed, so
+// subscript access is a read-only lookup returning a MESSAGE_FROZEN wrapper for
+// existing keys, or raising a FrozenInstanceError if the key is not present.
 PyObject* MapReflectionFriend::MessageMapGetItem(PyObject* _self,
                                                  PyObject* key) {
   MessageMapContainer* self = GetMessageMap(_self);
 
-  Message* message = self->GetMutableMessage();
-  if (message == nullptr) return nullptr;
-  const Reflection* reflection = message->GetReflection();
   MapKey map_key;
-  MapValueRef value;
 
   if (!PythonToMapKey(self, key, &map_key)) {
     return nullptr;
   }
 
-  if (reflection->InsertOrLookupMapValue(message, self->parent_field_descriptor,
-                                         map_key, &value)) {
-    self->version++;
+  // Mutable path: insert-or-lookup and return mutable submessage wrapper.
+  if (Message* message = self->GetMutableMessage(); message != nullptr) {
+    const Reflection* reflection = message->GetReflection();
+    MapValueRef value;
+    if (reflection->InsertOrLookupMapValue(
+            message, self->parent_field_descriptor, map_key, &value)) {
+      self->version++;
+    }
+    return GetCMessage(self, value.MutableMessageValue(), MESSAGE_MUTABLE);
   }
 
-  return GetCMessage(self, value.MutableMessageValue());
+  // Frozen path: read-only lookup without mutation.
+  PyErr_Clear();
+  const Message* message = self->GetReadOnlyMessage();
+  const Reflection* reflection = message->GetReflection();
+  MapValueConstRef value;
+  if (!reflection->LookupMapValue(*message, self->parent_field_descriptor,
+                                  map_key, &value)) {
+    return SetMessageFrozenError();
+  }
+  return GetCMessage(self, &value.GetMessageValue(), MESSAGE_FROZEN);
 }
 
 PyObject* MapReflectionFriend::MessageMapToStr(PyObject* _self) {
@@ -695,6 +711,13 @@ PyObject* MapReflectionFriend::MessageMapToStr(PyObject* _self) {
   MessageMapContainer* self = GetMessageMap(_self);
   const Message* message = self->GetReadOnlyMessage();
   const Reflection* reflection = message->GetReflection();
+  // Map elements only exist if the parent map contains allocated nodes (default
+  // instances are empty). If the parent is not frozen, the underlying C++
+  // message in the map is already fully mutable in memory, so creating or
+  // reusing CMessage wrappers with MESSAGE_MUTABLE is safe and allows
+  // subsequent mutations without requiring lazy promotion.
+  MessageMutabilityState state =
+      self->parent->state == MESSAGE_FROZEN ? MESSAGE_FROZEN : MESSAGE_MUTABLE;
   for (google::protobuf::ConstMapIterator it =
            reflection->ConstMapBegin(message, self->parent_field_descriptor);
        it != reflection->ConstMapEnd(message, self->parent_field_descriptor);
@@ -703,7 +726,7 @@ PyObject* MapReflectionFriend::MessageMapToStr(PyObject* _self) {
     if (key == nullptr) {
       return nullptr;
     }
-    value.reset(GetCMessage(self, &it.GetValueRef().GetMessageValue()));
+    value.reset(GetCMessage(self, &it.GetValueRef().GetMessageValue(), state));
     if (value == nullptr) {
       return nullptr;
     }

python/google/protobuf/pyext/message.cc

@@ -873,7 +873,6 @@ Message* AssureWritable(CMessage* self) {
     return nullptr;
   }
 
-  // Make self->message writable.
   const Reflection* reflection = parent_message->GetReflection();
   Message* mutable_message = reflection->MutableMessage(
       parent_message, self->parent_field_descriptor,
@@ -2910,7 +2909,7 @@ void ContainerBase::RemoveFromParentCache() {
 
 CMessage* CMessage::BuildSubMessageFromPointer(
     const FieldDescriptor* field_descriptor, const Message* sub_message,
-    CMessageClass* message_class) {
+    CMessageClass* message_class, MessageMutabilityState state) {
   if (PyObject* value =
           this->child_submessages.Get()->Get(sub_message, nullptr)) {
     return reinterpret_cast<CMessage*>(value);
@@ -2925,9 +2924,7 @@ CMessage* CMessage::BuildSubMessageFromPointer(
   Py_INCREF(this);
   cmsg->parent = this;
   cmsg->parent_field_descriptor = field_descriptor;
-  if (this->state == MESSAGE_FROZEN) {
-    cmsg->state = MESSAGE_FROZEN;
-  }
+  cmsg->state = this->state == MESSAGE_FROZEN ? MESSAGE_FROZEN : state;
   cmessage::SetSubmessage(this, cmsg);
   return cmsg;
 }

python/google/protobuf/pyext/message.h

@@ -132,7 +132,8 @@ typedef struct CMessage : public ContainerBase {
   // pointer to a message.
   CMessage* BuildSubMessageFromPointer(const FieldDescriptor* field_descriptor,
                                        const Message* sub_message,
-                                       CMessageClass* message_class);
+                                       CMessageClass* message_class,
+                                       MessageMutabilityState state);
   CMessage* MaybeReleaseSubMessage(const Message* sub_message);
 } CMessage;
 

python/google/protobuf/pyext/repeated_composite_container.cc

@@ -50,11 +50,14 @@ PyObject* Add(RepeatedCompositeContainer* self, PyObject* args,
   Message* message = cmessage::AssureWritable(self->parent);
   if (message == nullptr) return nullptr;
 
-  Message* sub_message = message->GetReflection()->AddMessage(
+  const Reflection* reflection = message->GetReflection();
+  Message* sub_message = reflection->AddMessage(
       message, self->parent_field_descriptor,
       self->child_message_class->py_message_factory->message_factory);
   CMessage* cmsg = self->parent->BuildSubMessageFromPointer(
-      self->parent_field_descriptor, sub_message, self->child_message_class);
+      self->parent_field_descriptor, sub_message, self->child_message_class,
+      MESSAGE_MUTABLE);
+  if (cmsg == nullptr) return nullptr;
 
   if (cmessage::InitAttributes(cmsg, args, kwargs) < 0) {
     message->GetReflection()->RemoveLast(message,
@@ -180,33 +183,31 @@ static PyObject* MergeFromMethod(PyObject* self, PyObject* other) {
 // This function does not check the bounds.
 static PyObject* GetItem(RepeatedCompositeContainer* self, Py_ssize_t index,
                          Py_ssize_t length = -1) {
+  const Message* message = self->parent->message;
+  const Reflection* reflection = message->GetReflection();
   if (length == -1) {
-    const Message* message = self->parent->message;
-    const Reflection* reflection = message->GetReflection();
     length = reflection->FieldSize(*message, self->parent_field_descriptor);
   }
   if (index < 0 || index >= length) {
     PyErr_Format(PyExc_IndexError, "list index (%zd) out of range", index);
     return nullptr;
   }
-  const Message* message = self->parent->message;
-  const Reflection* reflection = message->GetReflection();
-  const Message* sub_message = nullptr;
   const int int_index = static_cast<int>(index);
-  if (self->parent->state == python::MESSAGE_FROZEN) {
-    sub_message = &reflection->GetRepeatedMessage(
-        *message, self->parent_field_descriptor, int_index);
-  } else {
-    Message* mutable_parent = cmessage::AssureWritable(self->parent);
-    if (mutable_parent == nullptr) {
-      return nullptr;
-    }
-    sub_message = mutable_parent->GetReflection()->MutableRepeatedMessage(
-        mutable_parent, self->parent_field_descriptor, int_index);
-  }
+  const Message* sub_message = &reflection->GetRepeatedMessage(
+      *message, self->parent_field_descriptor, int_index);
+  // Elements in a repeated field only exist if the parent container is
+  // non-empty (default instances have size 0). If the parent is not frozen,
+  // the underlying C++ message in RepeatedPtrField is already an allocated,
+  // fully mutable instance in memory. Using GetRepeatedMessage performs a
+  // const, thread-safe read without mutating parent state during subscript
+  // access, and wrapping with MESSAGE_MUTABLE allows subsequent mutations
+  // without requiring lazy promotion.
+  MessageMutabilityState state = self->parent->state == python::MESSAGE_FROZEN
+                                     ? MESSAGE_FROZEN
+                                     : MESSAGE_MUTABLE;
   return self->parent
       ->BuildSubMessageFromPointer(self->parent_field_descriptor, sub_message,
-                                   self->child_message_class)
+                                   self->child_message_class, state)
       ->AsPyObject();
 }
 
@@ -376,10 +377,11 @@ static void ReorderAttached(RepeatedCompositeContainer* self,
   for (Py_ssize_t i = 0; i < length; ++i) {
     CMessage* child_cmsg =
         reinterpret_cast<CMessage*>(PyList_GET_ITEM(child_list, i));
-    Message* child_message = cmessage::AssureWritable(child_cmsg);
-    if (child_message == nullptr) return;
-    reflection->UnsafeArenaAddAllocatedMessage(message, descriptor,
-                                               child_message);
+    // const_cast is safe because each child_cmsg originated from this mutable
+    // parent's repeated field (released above) and is already an allocated,
+    // mutable Message object in memory.
+    reflection->UnsafeArenaAddAllocatedMessage(
+        message, descriptor, const_cast<Message*>(child_cmsg->message));
   }
 }
 

src/google/protobuf/unittest_custom_options.proto

@@ -283,6 +283,7 @@ message ComplexOptionType1 {
   optional int32 foo3 = 3;
   repeated int32 foo4 = 4;
   map<string, int32> my_map = 5;
+  map<string, ComplexOptionType3> submsg_map = 6;
 
   extensions 100 to max;
 }
@@ -347,6 +348,10 @@ message VariousComplexOptions {
     key: "other_key"
     value: 456
   };
+  option (complex_opt1).submsg_map = {
+    key: "sub_key"
+    value: { moo: 555 }
+  };
   option (complex_opt2).baz = 987;
   option (complex_opt2).(grault) = 654;
   option (complex_opt2).bar.foo = 743;