gh-125984: fix use-after-free on `fut->fut_{callback,context}0` due to an evil `loop.__getattribute__` (#126003)

1 files changed, 71 insertions, 2 deletions

diff --git a/Lib/test/test_asyncio/test_futures.py b/Lib/test/test_asyncio/test_futures.py
index 4e3efa7..bac76b0 100644
--- a/Lib/test/test_asyncio/test_futures.py
+++ b/Lib/test/test_asyncio/test_futures.py
@@ -32,6 +32,25 @@ def last_cb():
     pass
 
 
+class ReachableCode(Exception):
+    """Exception to raise to indicate that some code was reached.
+
+    Use this exception if using mocks is not a good alternative.
+    """
+
+
+class SimpleEvilEventLoop(asyncio.base_events.BaseEventLoop):
+    """Base class for UAF and other evil stuff requiring an evil event loop."""
+
+    def get_debug(self):  # to suppress tracebacks
+        return False
+
+    def __del__(self):
+        # Automatically close the evil event loop to avoid warnings.
+        if not self.is_closed() and not self.is_running():
+            self.close()
+
+
 class DuckFuture:
     # Class that does not inherit from Future but aims to be duck-type
     # compatible with it.
@@ -948,6 +967,7 @@ class BaseFutureDoneCallbackTests():
         fut.remove_done_callback(evil())
 
     def test_evil_call_soon_list_mutation(self):
+        # see: https://github.com/python/cpython/issues/125969
         called_on_fut_callback0 = False
 
         pad = lambda: ...
@@ -962,9 +982,8 @@ class BaseFutureDoneCallbackTests():
             else:
                 called_on_fut_callback0 = True
 
-        fake_event_loop = lambda: ...
+        fake_event_loop = SimpleEvilEventLoop()
         fake_event_loop.call_soon = evil_call_soon
-        fake_event_loop.get_debug = lambda: False  # suppress traceback
 
         with mock.patch.object(self, 'loop', fake_event_loop):
             fut = self._new_future()
@@ -980,6 +999,56 @@ class BaseFutureDoneCallbackTests():
             # returns an empty list but the C implementation returns None.
             self.assertIn(fut._callbacks, (None, []))
 
+    def test_use_after_free_on_fut_callback_0_with_evil__getattribute__(self):
+        # see: https://github.com/python/cpython/issues/125984
+
+        class EvilEventLoop(SimpleEvilEventLoop):
+            def call_soon(self, *args, **kwargs):
+                super().call_soon(*args, **kwargs)
+                raise ReachableCode
+
+            def __getattribute__(self, name):
+                nonlocal fut_callback_0
+                if name == 'call_soon':
+                    fut.remove_done_callback(fut_callback_0)
+                    del fut_callback_0
+                return object.__getattribute__(self, name)
+
+        evil_loop = EvilEventLoop()
+        with mock.patch.object(self, 'loop', evil_loop):
+            fut = self._new_future()
+            self.assertIs(fut.get_loop(), evil_loop)
+
+            fut_callback_0 = lambda: ...
+            fut.add_done_callback(fut_callback_0)
+            self.assertRaises(ReachableCode, fut.set_result, "boom")
+
+    def test_use_after_free_on_fut_context_0_with_evil__getattribute__(self):
+        # see: https://github.com/python/cpython/issues/125984
+
+        class EvilEventLoop(SimpleEvilEventLoop):
+            def call_soon(self, *args, **kwargs):
+                super().call_soon(*args, **kwargs)
+                raise ReachableCode
+
+            def __getattribute__(self, name):
+                if name == 'call_soon':
+                    # resets the future's event loop
+                    fut.__init__(loop=SimpleEvilEventLoop())
+                return object.__getattribute__(self, name)
+
+        evil_loop = EvilEventLoop()
+        with mock.patch.object(self, 'loop', evil_loop):
+            fut = self._new_future()
+            self.assertIs(fut.get_loop(), evil_loop)
+
+            fut_callback_0 = mock.Mock()
+            fut_context_0 = mock.Mock()
+            fut.add_done_callback(fut_callback_0, context=fut_context_0)
+            del fut_context_0
+            del fut_callback_0
+            self.assertRaises(ReachableCode, fut.set_result, "boom")
+
 
 @unittest.skipUnless(hasattr(futures, '_CFuture'),
                      'requires the C _asyncio module')