From 4826d52338396758b2d6790a498c2a06eec19a86 Mon Sep 17 00:00:00 2001
From: Jamie Phan <jamie@ordinarylab.dev>
Date: Wed, 10 Jan 2024 16:21:00 +1100
Subject: [PATCH] gh-112182: Replace StopIteration with RuntimeError for future
 (#113220)

When an `StopIteration` raises into `asyncio.Future`, this will cause
a thread to hang. This commit address this by not raising an exception
and silently transforming the `StopIteration` with a `RuntimeError`,
which the caller can reconstruct from `fut.exception().__cause__`
---
 Lib/asyncio/futures.py                        | 10 +++++---
 Lib/test/test_asyncio/test_futures.py         | 23 ++++++++++++++---
 ...-12-17-10-22-55.gh-issue-112182.jLWGlr.rst |  3 +++
 Modules/_asynciomodule.c                      | 25 +++++++++++++++----
 4 files changed, 49 insertions(+), 12 deletions(-)
 create mode 100644 Misc/NEWS.d/next/Library/2023-12-17-10-22-55.gh-issue-112182.jLWGlr.rst

diff --git a/Lib/asyncio/futures.py b/Lib/asyncio/futures.py
index d19e5d8c919..5d35321db79 100644
--- a/Lib/asyncio/futures.py
+++ b/Lib/asyncio/futures.py
@@ -269,9 +269,13 @@ def set_exception(self, exception):
             raise exceptions.InvalidStateError(f'{self._state}: {self!r}')
         if isinstance(exception, type):
             exception = exception()
-        if type(exception) is StopIteration:
-            raise TypeError("StopIteration interacts badly with generators "
-                            "and cannot be raised into a Future")
+        if isinstance(exception, StopIteration):
+            new_exc = RuntimeError("StopIteration interacts badly with "
+                                   "generators and cannot be raised into a "
+                                   "Future")
+            new_exc.__cause__ = exception
+            new_exc.__context__ = exception
+            exception = new_exc
         self._exception = exception
         self._exception_tb = exception.__traceback__
         self._state = _FINISHED
diff --git a/Lib/test/test_asyncio/test_futures.py b/Lib/test/test_asyncio/test_futures.py
index 2184b2091f8..d3e8efec1c0 100644
--- a/Lib/test/test_asyncio/test_futures.py
+++ b/Lib/test/test_asyncio/test_futures.py
@@ -270,10 +270,6 @@ def test_exception(self):
         f = self._new_future(loop=self.loop)
         self.assertRaises(asyncio.InvalidStateError, f.exception)
 
-        # StopIteration cannot be raised into a Future - CPython issue26221
-        self.assertRaisesRegex(TypeError, "StopIteration .* cannot be raised",
-                               f.set_exception, StopIteration)
-
         f.set_exception(exc)
         self.assertFalse(f.cancelled())
         self.assertTrue(f.done())
@@ -283,6 +279,25 @@ def test_exception(self):
         self.assertRaises(asyncio.InvalidStateError, f.set_exception, None)
         self.assertFalse(f.cancel())
 
+    def test_stop_iteration_exception(self, stop_iteration_class=StopIteration):
+        exc = stop_iteration_class()
+        f = self._new_future(loop=self.loop)
+        f.set_exception(exc)
+        self.assertFalse(f.cancelled())
+        self.assertTrue(f.done())
+        self.assertRaises(RuntimeError, f.result)
+        exc = f.exception()
+        cause = exc.__cause__
+        self.assertIsInstance(exc, RuntimeError)
+        self.assertRegex(str(exc), 'StopIteration .* cannot be raised')
+        self.assertIsInstance(cause, stop_iteration_class)
+
+    def test_stop_iteration_subclass_exception(self):
+        class MyStopIteration(StopIteration):
+            pass
+
+        self.test_stop_iteration_exception(MyStopIteration)
+
     def test_exception_class(self):
         f = self._new_future(loop=self.loop)
         f.set_exception(RuntimeError)
diff --git a/Misc/NEWS.d/next/Library/2023-12-17-10-22-55.gh-issue-112182.jLWGlr.rst b/Misc/NEWS.d/next/Library/2023-12-17-10-22-55.gh-issue-112182.jLWGlr.rst
new file mode 100644
index 00000000000..dc5bb697aac
--- /dev/null
+++ b/Misc/NEWS.d/next/Library/2023-12-17-10-22-55.gh-issue-112182.jLWGlr.rst
@@ -0,0 +1,3 @@
+:meth:`asyncio.futures.Future.set_exception()` now transforms :exc:`StopIteration`
+into :exc:`RuntimeError` instead of hanging or other misbehavior. Patch
+contributed by Jamie Phan.
diff --git a/Modules/_asynciomodule.c b/Modules/_asynciomodule.c
index b929e6d9922..c1aa849ecf1 100644
--- a/Modules/_asynciomodule.c
+++ b/Modules/_asynciomodule.c
@@ -597,12 +597,27 @@ future_set_exception(asyncio_state *state, FutureObj *fut, PyObject *exc)
         PyErr_SetString(PyExc_TypeError, "invalid exception object");
         return NULL;
     }
-    if (Py_IS_TYPE(exc_val, (PyTypeObject *)PyExc_StopIteration)) {
+    if (PyErr_GivenExceptionMatches(exc_val, PyExc_StopIteration)) {
+        const char *msg = "StopIteration interacts badly with "
+                          "generators and cannot be raised into a "
+                          "Future";
+        PyObject *message = PyUnicode_FromString(msg);
+        if (message == NULL) {
+            Py_DECREF(exc_val);
+            return NULL;
+        }
+        PyObject *err = PyObject_CallOneArg(PyExc_RuntimeError, message);
+        Py_DECREF(message);
+        if (err == NULL) {
+            Py_DECREF(exc_val);
+            return NULL;
+        }
+        assert(PyExceptionInstance_Check(err));
+
+        PyException_SetCause(err, Py_NewRef(exc_val));
+        PyException_SetContext(err, Py_NewRef(exc_val));
         Py_DECREF(exc_val);
-        PyErr_SetString(PyExc_TypeError,
-                        "StopIteration interacts badly with generators "
-                        "and cannot be raised into a Future");
-        return NULL;
+        exc_val = err;
     }
 
     assert(!fut->fut_exception);