gh-154189: Fix use-after-free in functools.partial_vectorcall (GH-154508)

Author: brijkapadia

Lib/test/test_functools.py

@@ -579,6 +579,40 @@ def f(**kwargs):
         with self.assertRaises(RuntimeError):
             result = p(**{BadStr("poison"): "new_value"})
 
+    def test_call_safety_against_reentrant_mutation(self):
+        def old_function(*args, **kwargs):
+            return "old_function", args, kwargs
+
+        def new_function(*args, **kwargs):
+            return "new_function", args, kwargs
+
+        g_partial = None
+
+        class EvilKey(str):
+            armed = False
+            def __hash__(self):
+                if EvilKey.armed and g_partial is not None:
+                    EvilKey.armed = False
+                    new_args_tuple = ("new_arg",)
+                    new_keywords_dict = {"new_keyword": None}
+                    new_tuple_state = (new_function, new_args_tuple, new_keywords_dict, None)
+                    g_partial.__setstate__(new_tuple_state)
+                    gc.collect()
+                return str.__hash__(self)
+
+        g_partial = functools.partial(old_function, "old_arg", old_keyword=None)
+
+        kwargs = {EvilKey("evil_key"): None}
+        EvilKey.armed = True
+
+        result = g_partial(**kwargs)
+        expected = ("old_function", ("old_arg",), {"old_keyword": None, "evil_key": None})
+        self.assertEqual(result, expected)
+
+        result = g_partial()
+        expected = ("new_function", ("new_arg",), {"new_keyword": None})
+        self.assertEqual(result, expected)
+
 @unittest.skipUnless(c_functools, 'requires the C _functools module')
 class TestPartialC(TestPartial, unittest.TestCase):
     if c_functools:

Misc/NEWS.d/next/Library/2026-07-22-15-56-11.gh-issue-154189.7zWWNZ.rst

@@ -0,0 +1,4 @@
+Fixed a potential use-after-free when calling :func:`functools.partial`.
+Now, when invoking a :func:`~functools.partial` object, the stored function,
+positional arguments, and keyword arguments are preserved for the duration
+of the call in case of reentrancy.

Modules/_functoolsmodule.c

@@ -382,18 +382,24 @@ partial_vectorcall(PyObject *self, PyObject *const *args,
         return NULL;
     }
 
-    PyObject **pto_args = _PyTuple_ITEMS(pto->args);
-    Py_ssize_t pto_nargs = PyTuple_GET_SIZE(pto->args);
-    Py_ssize_t pto_nkwds = PyDict_GET_SIZE(pto->kw);
+    PyObject *result = NULL;
+    PyObject *partial_function = Py_NewRef(pto->fn);
+    PyObject *partial_args = Py_NewRef(pto->args);
+    PyObject *partial_keywords = Py_NewRef(pto->kw);
+
+    PyObject **pto_args = _PyTuple_ITEMS(partial_args);
+    Py_ssize_t pto_nargs = PyTuple_GET_SIZE(partial_args);
+    Py_ssize_t pto_nkwds = PyDict_GET_SIZE(partial_keywords);
     Py_ssize_t nkwds = kwnames == NULL ? 0 : PyTuple_GET_SIZE(kwnames);
     Py_ssize_t nargskw = nargs + nkwds;
 
     /* Special cases */
     if (!pto_nkwds) {
         /* Fast path if we're called without arguments */
         if (nargskw == 0) {
-            return _PyObject_VectorcallTstate(tstate, pto->fn, pto_args,
-                                              pto_nargs, NULL);
+            result = _PyObject_VectorcallTstate(tstate, partial_function, pto_args,
+                                                pto_nargs, NULL);
+            goto done;
         }
 
         /* Use PY_VECTORCALL_ARGUMENTS_OFFSET to prepend a single
@@ -402,10 +408,10 @@ partial_vectorcall(PyObject *self, PyObject *const *args,
             PyObject **newargs = (PyObject **)args - 1;
             PyObject *tmp = newargs[0];
             newargs[0] = pto_args[0];
-            PyObject *ret = _PyObject_VectorcallTstate(tstate, pto->fn, newargs,
-                                                       nargs + 1, kwnames);
+            result = _PyObject_VectorcallTstate(tstate, partial_function, newargs,
+                                                nargs + 1, kwnames);
             newargs[0] = tmp;
-            return ret;
+            goto done;
         }
     }
 
@@ -435,7 +441,8 @@ partial_vectorcall(PyObject *self, PyObject *const *args,
     else {
         stack = PyMem_Malloc(init_stack_size * sizeof(PyObject *));
         if (stack == NULL) {
-            return PyErr_NoMemory();
+            PyErr_NoMemory();
+            goto done;
         }
     }
 
@@ -457,20 +464,20 @@ partial_vectorcall(PyObject *self, PyObject *const *args,
         for (Py_ssize_t i = 0; i < nkwds; ++i) {
             key = PyTuple_GET_ITEM(kwnames, i);
             val = args[nargs + i];
-            int contains = PyDict_Contains(pto->kw, key);
+            int contains = PyDict_Contains(partial_keywords, key);
             if (contains < 0) {
-                goto error;
+                goto clean_stack;
             }
             else if (contains == 1) {
                 if (pto_kw_merged == NULL) {
-                    pto_kw_merged = PyDict_Copy(pto->kw);
+                    pto_kw_merged = PyDict_Copy(partial_keywords);
                     if (pto_kw_merged == NULL) {
-                        goto error;
+                        goto clean_stack;
                     }
                 }
                 if (PyDict_SetItem(pto_kw_merged, key, val) < 0) {
                     Py_DECREF(pto_kw_merged);
-                    goto error;
+                    goto clean_stack;
                 }
             }
             else {
@@ -486,7 +493,7 @@ partial_vectorcall(PyObject *self, PyObject *const *args,
         tot_kwnames = PyTuple_New(tot_nkwds - n_merges);
         if (tot_kwnames == NULL) {
             Py_XDECREF(pto_kw_merged);
-            goto error;
+            goto clean_stack;
         }
         for (Py_ssize_t i = 0; i < n_tail; ++i) {
             key = Py_NewRef(stack[tot_nargskw + i]);
@@ -496,7 +503,7 @@ partial_vectorcall(PyObject *self, PyObject *const *args,
         /* Copy pto_keywords with overlapping call keywords merged
          * Note, tail is already coppied. */
         Py_ssize_t pos = 0, i = 0;
-        PyObject *keyword_dict = n_merges ? pto_kw_merged : pto->kw;
+        PyObject *keyword_dict = n_merges ? pto_kw_merged : partial_keywords;
         Py_BEGIN_CRITICAL_SECTION(keyword_dict);
         while (PyDict_Next(keyword_dict, &pos, &key, &val)) {
             assert(i < pto_nkwds);
@@ -515,10 +522,8 @@ partial_vectorcall(PyObject *self, PyObject *const *args,
             tmp_stack = PyMem_Realloc(stack, (tot_nargskw - n_merges) * sizeof(PyObject *));
             if (tmp_stack == NULL) {
                 Py_DECREF(tot_kwnames);
-                if (stack != small_stack) {
-                    PyMem_Free(stack);
-                }
-                return PyErr_NoMemory();
+                PyErr_NoMemory();
+                goto clean_stack;
             }
             stack = tmp_stack;
         }
@@ -547,21 +552,22 @@ partial_vectorcall(PyObject *self, PyObject *const *args,
         memcpy(stack + pto_nargs, args, nargs * sizeof(PyObject*));
     }
 
-    PyObject *ret = _PyObject_VectorcallTstate(tstate, pto->fn, stack,
-                                               tot_nargs, tot_kwnames);
-    if (stack != small_stack) {
-        PyMem_Free(stack);
-    }
+    result = _PyObject_VectorcallTstate(tstate, partial_function, stack,
+                                        tot_nargs, tot_kwnames);
     if (pto_nkwds) {
         Py_DECREF(tot_kwnames);
     }
-    return ret;
 
- error:
+ clean_stack:
     if (stack != small_stack) {
         PyMem_Free(stack);
     }
-    return NULL;
+
+ done:
+    Py_DECREF(partial_function);
+    Py_DECREF(partial_args);
+    Py_DECREF(partial_keywords);
+    return result;
 }
 
 /* Set pto->vectorcall depending on the parameters of the partial object */