From 4bdd07858d7d0c03c10f18692ee5e270fb12c409 Mon Sep 17 00:00:00 2001 From: Samuel Benzaquen Date: Thu, 20 Aug 2026 08:14:10 -0700 Subject: [PATCH] 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 --- .../protobuf/internal/descriptor_test.py | 61 ++++++++++++++++++ .../google/protobuf/internal/message_test.py | 64 +++++++++++++++++++ .../protobuf/internal/test_proto2.proto | 8 +++ .../protobuf/internal/thread_safe_test.py | 34 ++++++++++ python/google/protobuf/pyext/map_container.cc | 49 ++++++++++---- python/google/protobuf/pyext/message.cc | 7 +- python/google/protobuf/pyext/message.h | 3 +- .../pyext/repeated_composite_container.cc | 48 +++++++------- .../protobuf/unittest_custom_options.proto | 5 ++ 9 files changed, 237 insertions(+), 42 deletions(-) diff --git a/python/google/protobuf/internal/descriptor_test.py b/python/google/protobuf/internal/descriptor_test.py index a91970c946..098e40d8a8 100755 --- a/python/google/protobuf/internal/descriptor_test.py +++ b/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 = ( diff --git a/python/google/protobuf/internal/message_test.py b/python/google/protobuf/internal/message_test.py index 94895a00ae..03f43be814 100755 --- a/python/google/protobuf/internal/message_test.py +++ b/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 diff --git a/python/google/protobuf/internal/test_proto2.proto b/python/google/protobuf/internal/test_proto2.proto index bb95dd614a..82444225ad 100644 --- a/python/google/protobuf/internal/test_proto2.proto +++ b/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; +} diff --git a/python/google/protobuf/internal/thread_safe_test.py b/python/google/protobuf/internal/thread_safe_test.py index a3b904bd8f..dc55bbabde 100644 --- a/python/google/protobuf/internal/thread_safe_test.py +++ b/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() diff --git a/python/google/protobuf/pyext/map_container.cc b/python/google/protobuf/pyext/map_container.cc index 4b6b71ef92..ca21add1ec 100644 --- a/python/google/protobuf/pyext/map_container.cc +++ b/python/google/protobuf/pyext/map_container.cc @@ -577,12 +577,12 @@ static MessageMapContainer* GetMessageMap(PyObject* obj) { return reinterpret_cast(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; } diff --git a/python/google/protobuf/pyext/message.cc b/python/google/protobuf/pyext/message.cc index a52bc9821b..a4139b1b22 100644 --- a/python/google/protobuf/pyext/message.cc +++ b/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(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; } diff --git a/python/google/protobuf/pyext/message.h b/python/google/protobuf/pyext/message.h index bc993e2c17..89339f075f 100644 --- a/python/google/protobuf/pyext/message.h +++ b/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; diff --git a/python/google/protobuf/pyext/repeated_composite_container.cc b/python/google/protobuf/pyext/repeated_composite_container.cc index 044016cff6..42fc0c05ba 100644 --- a/python/google/protobuf/pyext/repeated_composite_container.cc +++ b/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(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(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(child_cmsg->message)); } } diff --git a/src/google/protobuf/unittest_custom_options.proto b/src/google/protobuf/unittest_custom_options.proto index 173aa4ea82..454b63bdbf 100644 --- a/src/google/protobuf/unittest_custom_options.proto +++ b/src/google/protobuf/unittest_custom_options.proto @@ -283,6 +283,7 @@ message ComplexOptionType1 { optional int32 foo3 = 3; repeated int32 foo4 = 4; map my_map = 5; + map 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;