Fix UAF (address Nico Posada's review)

picnixz · picnixz · commit 22b21cc9aace · 2024-10-23T15:10:31.000+02:00
diff --git a/Lib/test/test_asyncio/test_futures.py b/Lib/test/test_asyncio/test_futures.py
@@ -970,6 +970,45 @@ def __eq__(self, other):
 
         self.assertIs(mem, cb_pad)
 
+    # Use-After-Free sanity checks.
+    # Credits to Nico-Posada for the PoCs.
+    # See https://github.com/python/cpython/pull/125833.
+
+    def test_use_after_free_fixed_1(self):
+        fut = self._new_future()
+
+        class setup:
+            def __eq__(self, other):
+                return False
+
+        # evil MUST be a subclass of setup for __eq__ to get called first
+        class evil(setup):
+            def __eq__(self, value):
+                fut._callbacks.clear()
+                return NotImplemented
+
+        cb_pad = lambda: ...
+        fut.add_done_callback(cb_pad)  # sets fut->fut_callback0
+        fut.add_done_callback(setup())  # sets fut->fut_callbacks[0]
+        # removes callback from fut->fut_callback0 setting it to NULL
+        fut.remove_done_callback(cb_pad)
+        fut.remove_done_callback(evil())
+
+    def test_use_after_free_fixed_2(self):
+        fut = self._new_future()
+
+        class cb_pad:
+            def __eq__(self, other):
+                return True
+
+        class evil(cb_pad):
+            def __eq__(self, other):
+                fut.remove_done_callback(None)
+                return NotImplemented
+
+        fut.add_done_callback(cb_pad())
+        fut.remove_done_callback(evil())
+
 
 @unittest.skipUnless(hasattr(futures, '_CFuture'),
                      'requires the C _asyncio module')
diff --git a/Modules/_asynciomodule.c b/Modules/_asynciomodule.c
@@ -1036,8 +1036,12 @@ _asyncio_Future_remove_done_callback_impl(FutureObj *self, PyTypeObject *cls,
 
     if (self->fut_callback0 != NULL) {
         // Beware: An evil PyObject_RichCompareBool could change fut_callback0
-        // (see https://github.com/python/cpython/issues/125789 for details).
-        int cmp = PyObject_RichCompareBool(self->fut_callback0, fn, Py_EQ);
+        // (see https://github.com/python/cpython/issues/125789 for details)
+        // In addition, the reference to self->fut_callback0 may be cleared,
+        // so we need to temporarily hold it explicitly.
+        PyObject *fut_callback0 = Py_NewRef(self->fut_callback0);
+        int cmp = PyObject_RichCompareBool(fut_callback0, fn, Py_EQ);
+        Py_DECREF(fut_callback0);
         if (cmp == -1) {
             return NULL;
         }
@@ -1076,7 +1080,7 @@ _asyncio_Future_remove_done_callback_impl(FutureObj *self, PyTypeObject *cls,
         Py_INCREF(cb_tup);
         PyObject *cb = PyTuple_GET_ITEM(cb_tup, 0);
         int cmp = PyObject_RichCompareBool(cb, fn, Py_EQ);
-        Py_DECREF(cb_tup);
+        Py_XDECREF(cb_tup);
         if (cmp == -1) {
             return NULL;
         }
