From 79688b5b0ea761183193ffb0859415f3b02fa44d Mon Sep 17 00:00:00 2001
From: Sam Gross <colesbury@gmail.com>
Date: Mon, 29 Apr 2024 15:49:01 -0400
Subject: [PATCH] gh-118331: Handle errors in _PyObject_SetManagedDict
 (#118334)

When detaching a dict, the `copy_values` call may fail due to
out-of-memory errors. This can be triggered by test_no_memory in
test_repl.
---
 Include/cpython/object.h |  2 +-
 Objects/dictobject.c     | 29 ++++++++++++++++++-----------
 Objects/typeobject.c     |  2 +-
 3 files changed, 20 insertions(+), 13 deletions(-)

diff --git a/Include/cpython/object.h b/Include/cpython/object.h
index a8f57827a96..a6b93b93ab0 100644
--- a/Include/cpython/object.h
+++ b/Include/cpython/object.h
@@ -493,7 +493,7 @@ do { \
 PyAPI_FUNC(void *) PyObject_GetItemData(PyObject *obj);
 
 PyAPI_FUNC(int) PyObject_VisitManagedDict(PyObject *obj, visitproc visit, void *arg);
-PyAPI_FUNC(void) _PyObject_SetManagedDict(PyObject *obj, PyObject *new_dict);
+PyAPI_FUNC(int) _PyObject_SetManagedDict(PyObject *obj, PyObject *new_dict);
 PyAPI_FUNC(void) PyObject_ClearManagedDict(PyObject *obj);
 
 #define TYPE_MAX_WATCHERS 8
diff --git a/Objects/dictobject.c b/Objects/dictobject.c
index 43cb2350b7f..1f21f70f149 100644
--- a/Objects/dictobject.c
+++ b/Objects/dictobject.c
@@ -7056,11 +7056,12 @@ set_dict_inline_values(PyObject *obj, PyDictObject *new_dict)
     }
 }
 
-void
+int
 _PyObject_SetManagedDict(PyObject *obj, PyObject *new_dict)
 {
     assert(Py_TYPE(obj)->tp_flags & Py_TPFLAGS_MANAGED_DICT);
     assert(_PyObject_InlineValuesConsistencyCheck(obj));
+    int err = 0;
     PyTypeObject *tp = Py_TYPE(obj);
     if (tp->tp_flags & Py_TPFLAGS_INLINE_VALUES) {
         PyDictObject *dict = _PyObject_GetManagedDict(obj);
@@ -7076,11 +7077,11 @@ _PyObject_SetManagedDict(PyObject *obj, PyObject *new_dict)
             Py_END_CRITICAL_SECTION();
 
             if (dict == NULL) {
-                return;
+                return 0;
             }
 #else
             set_dict_inline_values(obj, (PyDictObject *)new_dict);
-            return;
+            return 0;
 #endif
         }
 
@@ -7089,15 +7090,16 @@ _PyObject_SetManagedDict(PyObject *obj, PyObject *new_dict)
         // We've locked dict, but the actual dict could have changed
         // since we locked it.
         dict = _PyObject_ManagedDictPointer(obj)->dict;
-
-        FT_ATOMIC_STORE_PTR(_PyObject_ManagedDictPointer(obj)->dict,
-                            (PyDictObject *)Py_XNewRef(new_dict));
-
-        _PyDict_DetachFromObject(dict, obj);
-
+        err = _PyDict_DetachFromObject(dict, obj);
+        if (err == 0) {
+            FT_ATOMIC_STORE_PTR(_PyObject_ManagedDictPointer(obj)->dict,
+                                (PyDictObject *)Py_XNewRef(new_dict));
+        }
         Py_END_CRITICAL_SECTION2();
 
-        Py_XDECREF(dict);
+        if (err == 0) {
+            Py_XDECREF(dict);
+        }
     }
     else {
         PyDictObject *dict;
@@ -7114,18 +7116,23 @@ _PyObject_SetManagedDict(PyObject *obj, PyObject *new_dict)
         Py_XDECREF(dict);
     }
     assert(_PyObject_InlineValuesConsistencyCheck(obj));
+    return err;
 }
 
 void
 PyObject_ClearManagedDict(PyObject *obj)
 {
-    _PyObject_SetManagedDict(obj, NULL);
+    if (_PyObject_SetManagedDict(obj, NULL) < 0) {
+        PyErr_WriteUnraisable(NULL);
+    }
 }
 
 int
 _PyDict_DetachFromObject(PyDictObject *mp, PyObject *obj)
 {
     _Py_CRITICAL_SECTION_ASSERT_OBJECT_LOCKED(obj);
+    assert(_PyObject_ManagedDictPointer(obj)->dict == mp);
+    assert(_PyObject_InlineValuesConsistencyCheck(obj));
 
     if (FT_ATOMIC_LOAD_PTR_RELAXED(mp->ma_values) != _PyObject_InlineValues(obj)) {
         return 0;
diff --git a/Objects/typeobject.c b/Objects/typeobject.c
index 07e0a5a02da..50efbb6e182 100644
--- a/Objects/typeobject.c
+++ b/Objects/typeobject.c
@@ -3167,7 +3167,7 @@ subtype_setdict(PyObject *obj, PyObject *value, void *context)
     }
 
     if (Py_TYPE(obj)->tp_flags & Py_TPFLAGS_MANAGED_DICT) {
-        _PyObject_SetManagedDict(obj, value);
+        return _PyObject_SetManagedDict(obj, value);
     }
     else {
         dictptr = _PyObject_ComputedDictPointer(obj);