Fix python upb crashes on map/repeated reference stub destructor

The original code removes map/repeated stub from the weak map without reifying in some situations. Changed _Reify() functions to decide if need to delete from WeakMap or SetConcreteSubobj.

PiperOrigin-RevId: 752524491

Author: anandolee

python/google/protobuf/internal/message_test.py

@@ -2699,6 +2699,20 @@ def testMapFindInitializationErrorsSmokeTest(self):
     msg.map_string_foreign_message['foo'].c = 5
     self.assertEqual(0, len(msg.FindInitializationErrors()))
 
+  def testMapStubReferenceSubMessageDestructor(self):
+    msg = map_unittest_pb2.TestMapSubmessage()
+    # A reference on map stub in sub message
+    map_ref = msg.test_map.map_int32_int32
+    # Make sure destructor after Clear the original message not crash
+    msg.Clear()
+
+  def testRepeatedStubReferenceSubMessageDestructor(self):
+    msg = unittest_pb2.NestedTestAllTypes()
+    # A reference on repeated stub in sub message
+    repeated_ref = msg.payload.repeated_int32
+    # Make sure destructor after Clear the original message not crash
+    msg.Clear()
+
   @unittest.skipIf(sys.maxunicode == UCS2_MAXUNICODE, 'Skip for ucs2')
   def testStrictUtf8Check(self):
     # Test u'\ud801' is rejected at parser in both python2 and python3.

python/map.c

@@ -90,7 +90,8 @@ PyObject* PyUpb_MapContainer_NewStub(PyObject* parent, const upb_FieldDef* f,
   return &map->ob_base;
 }
 
-void PyUpb_MapContainer_Reify(PyObject* _self, upb_Map* map) {
+upb_Map* PyUpb_MapContainer_Reify(PyObject* _self, upb_Map* map,
+                                  PyUpb_WeakMap* subobj_map, intptr_t iter) {
   PyUpb_MapContainer* self = (PyUpb_MapContainer*)_self;
   if (!map) {
     const upb_FieldDef* f = PyUpb_MapContainer_GetField(self);
@@ -101,11 +102,19 @@ void PyUpb_MapContainer_Reify(PyObject* _self, upb_Map* map) {
     map = upb_Map_New(arena, upb_FieldDef_CType(key_f),
                       upb_FieldDef_CType(val_f));
   }
+  if (subobj_map) {
+    PyUpb_WeakMap_DeleteIter(subobj_map, &iter);
+  } else {
+    const upb_FieldDef* f = PyUpb_MapContainer_GetField(self);
+    upb_MessageValue msgval = {.map_val = map};
+    PyUpb_Message_SetConcreteSubobj(self->ptr.parent, f, msgval);
+  }
   PyUpb_ObjCache_Add(map, &self->ob_base);
   Py_DECREF(self->ptr.parent);
   self->ptr.map = map;  // Overwrites self->ptr.parent.
   self->field &= ~(uintptr_t)1;
   assert(!PyUpb_MapContainer_IsStub(self));
+  return map;
 }
 
 void PyUpb_MapContainer_Invalidate(PyObject* obj) {
@@ -119,17 +128,7 @@ upb_Map* PyUpb_MapContainer_EnsureReified(PyObject* _self) {
   upb_Map* map = PyUpb_MapContainer_GetIfReified(self);
   if (map) return map;  // Already writable.
 
-  const upb_FieldDef* f = PyUpb_MapContainer_GetField(self);
-  upb_Arena* arena = PyUpb_Arena_Get(self->arena);
-  const upb_MessageDef* entry_m = upb_FieldDef_MessageSubDef(f);
-  const upb_FieldDef* key_f = upb_MessageDef_Field(entry_m, 0);
-  const upb_FieldDef* val_f = upb_MessageDef_Field(entry_m, 1);
-  map =
-      upb_Map_New(arena, upb_FieldDef_CType(key_f), upb_FieldDef_CType(val_f));
-  upb_MessageValue msgval = {.map_val = map};
-  PyUpb_Message_SetConcreteSubobj(self->ptr.parent, f, msgval);
-  PyUpb_MapContainer_Reify((PyObject*)self, map);
-  return map;
+  return PyUpb_MapContainer_Reify((PyObject*)self, NULL, NULL, 0);
 }
 
 static bool PyUpb_MapContainer_Set(PyUpb_MapContainer* self, upb_Map* map,

python/map.h

@@ -10,6 +10,7 @@
 
 #include <stdbool.h>
 
+#include "python/protobuf.h"
 #include "python/python_api.h"
 #include "upb/reflection/def.h"
 
@@ -28,7 +29,8 @@ PyObject* PyUpb_MapContainer_GetOrCreateWrapper(upb_Map* map,
 
 // Reifies a map stub to point to the concrete data in `map`.
 // If `map` is NULL, an appropriate empty map will be constructed.
-void PyUpb_MapContainer_Reify(PyObject* self, upb_Map* map);
+upb_Map* PyUpb_MapContainer_Reify(PyObject* self, upb_Map* map,
+                                  PyUpb_WeakMap* subobj_map, intptr_t iter);
 
 // Reifies this map object if it is not already reified.
 upb_Map* PyUpb_MapContainer_EnsureReified(PyObject* self);

python/message.c

@@ -685,8 +685,10 @@ static void PyUpb_Message_SyncSubobjs(PyUpb_Message* self);
  * the set state (having a non-owning pointer to self->ptr.msg).
  */
 static void PyUpb_Message_Reify(PyUpb_Message* self, const upb_FieldDef* f,
-                                upb_Message* msg) {
+                                upb_Message* msg, PyUpb_WeakMap* subobj_map,
+                                intptr_t iter) {
   assert(f == PyUpb_Message_GetFieldDef(self));
+  PyUpb_WeakMap_DeleteIter(subobj_map, &iter);
   if (!msg) {
     const upb_MessageDef* msgdef = PyUpb_Message_GetMsgdef((PyObject*)self);
     const upb_MiniTable* layout = upb_MessageDef_MiniTable(msgdef);
@@ -738,17 +740,18 @@ static void PyUpb_Message_SyncSubobjs(PyUpb_Message* self) {
     if (upb_FieldDef_HasPresence(f) && !upb_Message_HasFieldByDef(msg, f))
       continue;
     upb_MessageValue msgval = upb_Message_GetFieldByDef(msg, f);
-    PyUpb_WeakMap_DeleteIter(subobj_map, &iter);
     if (upb_FieldDef_IsMap(f)) {
       if (!msgval.map_val) continue;
-      PyUpb_MapContainer_Reify(obj, (upb_Map*)msgval.map_val);
+      PyUpb_MapContainer_Reify(obj, (upb_Map*)msgval.map_val, subobj_map, iter);
     } else if (upb_FieldDef_IsRepeated(f)) {
       if (!msgval.array_val) continue;
-      PyUpb_RepeatedContainer_Reify(obj, (upb_Array*)msgval.array_val);
+      PyUpb_RepeatedContainer_Reify(obj, (upb_Array*)msgval.array_val,
+                                    subobj_map, iter);
     } else {
       PyUpb_Message* sub = (void*)obj;
       assert(self == sub->ptr.parent);
-      PyUpb_Message_Reify(sub, f, (upb_Message*)msgval.msg_val);
+      PyUpb_Message_Reify(sub, f, (upb_Message*)msgval.msg_val, subobj_map,
+                          iter);
     }
   }
 
@@ -1404,18 +1407,17 @@ static PyObject* PyUpb_Message_Clear(PyUpb_Message* self) {
 
     while (PyUpb_WeakMap_Next(subobj_map, &key, &obj, &iter)) {
       const upb_FieldDef* f = key;
-      PyUpb_WeakMap_DeleteIter(subobj_map, &iter);
       if (upb_FieldDef_IsMap(f)) {
         assert(upb_Message_GetFieldByDef(msg, f).map_val == NULL);
-        PyUpb_MapContainer_Reify(obj, NULL);
+        PyUpb_MapContainer_Reify(obj, NULL, subobj_map, iter);
       } else if (upb_FieldDef_IsRepeated(f)) {
         assert(upb_Message_GetFieldByDef(msg, f).array_val == NULL);
-        PyUpb_RepeatedContainer_Reify(obj, NULL);
+        PyUpb_RepeatedContainer_Reify(obj, NULL, subobj_map, iter);
       } else {
         assert(!upb_Message_HasFieldByDef(msg, f));
         PyUpb_Message* sub = (void*)obj;
         assert(self == sub->ptr.parent);
-        PyUpb_Message_Reify(sub, f, NULL);
+        PyUpb_Message_Reify(sub, f, NULL, subobj_map, iter);
       }
     }
   }

python/repeated.c

@@ -55,33 +55,36 @@ static upb_Array* PyUpb_RepeatedContainer_GetIfReified(
   return PyUpb_RepeatedContainer_IsStub(self) ? NULL : self->ptr.arr;
 }
 
-void PyUpb_RepeatedContainer_Reify(PyObject* _self, upb_Array* arr) {
+upb_Array* PyUpb_RepeatedContainer_Reify(PyObject* _self, upb_Array* arr,
+                                         PyUpb_WeakMap* subobj_map,
+                                         intptr_t iter) {
   PyUpb_RepeatedContainer* self = (PyUpb_RepeatedContainer*)_self;
   assert(PyUpb_RepeatedContainer_IsStub(self));
+  const upb_FieldDef* f = PyUpb_RepeatedContainer_GetField(self);
   if (!arr) {
-    const upb_FieldDef* f = PyUpb_RepeatedContainer_GetField(self);
     upb_Arena* arena = PyUpb_Arena_Get(self->arena);
     arr = upb_Array_New(arena, upb_FieldDef_CType(f));
   }
+  if (subobj_map) {
+    PyUpb_WeakMap_DeleteIter(subobj_map, &iter);
+  } else {
+    PyUpb_Message_SetConcreteSubobj(self->ptr.parent, f,
+                                    (upb_MessageValue){.array_val = arr});
+  }
   PyUpb_ObjCache_Add(arr, &self->ob_base);
   Py_DECREF(self->ptr.parent);
   self->ptr.arr = arr;  // Overwrites self->ptr.parent.
   self->field &= ~(uintptr_t)1;
   assert(!PyUpb_RepeatedContainer_IsStub(self));
+  return arr;
 }
 
 upb_Array* PyUpb_RepeatedContainer_EnsureReified(PyObject* _self) {
   PyUpb_RepeatedContainer* self = (PyUpb_RepeatedContainer*)_self;
   upb_Array* arr = PyUpb_RepeatedContainer_GetIfReified(self);
   if (arr) return arr;  // Already writable.
 
-  const upb_FieldDef* f = PyUpb_RepeatedContainer_GetField(self);
-  upb_Arena* arena = PyUpb_Arena_Get(self->arena);
-  arr = upb_Array_New(arena, upb_FieldDef_CType(f));
-  PyUpb_Message_SetConcreteSubobj(self->ptr.parent, f,
-                                  (upb_MessageValue){.array_val = arr});
-  PyUpb_RepeatedContainer_Reify((PyObject*)self, arr);
-  return arr;
+  return PyUpb_RepeatedContainer_Reify((PyObject*)self, NULL, NULL, 0);
 }
 
 static void PyUpb_RepeatedContainer_Dealloc(PyObject* _self) {

python/repeated.h

@@ -10,6 +10,7 @@
 
 #include <stdbool.h>
 
+#include "python/protobuf.h"
 #include "python/python_api.h"
 #include "upb/reflection/def.h"
 
@@ -29,7 +30,9 @@ PyObject* PyUpb_RepeatedContainer_GetOrCreateWrapper(upb_Array* arr,
 
 // Reifies a repeated field stub to point to the concrete data in `arr`.
 // If `arr` is NULL, an appropriate empty array will be constructed.
-void PyUpb_RepeatedContainer_Reify(PyObject* self, upb_Array* arr);
+upb_Array* PyUpb_RepeatedContainer_Reify(PyObject* self, upb_Array* arr,
+                                         PyUpb_WeakMap* subobj_map,
+                                         intptr_t iter);
 
 // Reifies this repeated object if it is not already reified.
 upb_Array* PyUpb_RepeatedContainer_EnsureReified(PyObject* self);