From fbaa55bc980b101051481155dcc1ff10e48cdaa2 Mon Sep 17 00:00:00 2001
From: Ilias Aberkane <286468824+iliasabk@users.noreply.github.com>
Date: Fri, 25 Sep 2026 02:03:44 +0200
Subject: [PATCH] gh-158031, gh-158033: Fix UAF in iter/aiter sentinel
 comparison

The sentinel comparison in calliter_iternext() and
acallawaitable_handle_error() ran arbitrary Python code (a custom
__eq__) while holding only a borrowed reference to it_sentinel.
A re-entrant __setstate__ or __anext__ could release the sentinel
mid-comparison, leaving the comparison machinery using freed memory.
Hold a strong reference to the sentinel across PyObject_RichCompareBool.
---
 Lib/test/test_asyncgen.py                     | 31 +++++++++++++++++++
 Lib/test/test_iter.py                         | 18 +++++++++++
 ...-23-23-45-00.gh-issue-158031.itSentUaf.rst |  5 +++
 Objects/iterobject.c                          | 14 +++++++--
 4 files changed, 66 insertions(+), 2 deletions(-)
 create mode 100644 Misc/NEWS.d/next/Core_and_Builtins/2026-09-23-23-45-00.gh-issue-158031.itSentUaf.rst

diff --git a/Lib/test/test_asyncgen.py b/Lib/test/test_asyncgen.py
index b5e0891feb794e..9f66a5d869dd9a 100644
--- a/Lib/test/test_asyncgen.py
+++ b/Lib/test/test_asyncgen.py
@@ -905,6 +905,37 @@ def test_aiter_callable_exhausted(self):
         with self.assertRaises(StopAsyncIteration):
             self.loop.run_until_complete(anext(it))
 
+    def test_aiter_callable_sentinel_reentrant_exhaustion(self):
+        # gh-158033: a sentinel __eq__ that exhausts the iterator
+        # re-entrantly must not leave the comparison using a freed
+        # sentinel.
+        state = {'stop': False}
+
+        async def produce():
+            return Result()
+
+        def spam():
+            if state['stop']:
+                raise StopAsyncIteration
+            return produce()
+
+        class Sentinel:
+            def __eq__(self, other):
+                state['stop'] = True
+                try:
+                    ait.__anext__().__await__().send(None)
+                except StopAsyncIteration:
+                    pass
+                return NotImplemented
+
+        class Result:
+            def __eq__(self, other):
+                return NotImplemented
+
+        ait = aiter(spam, Sentinel())
+        with self.assertRaises(StopIteration):
+            ait.__anext__().__await__().send(None)
+
     def test_aiter_callable_lazy(self):
         # The callable is only called when the awaitable is awaited
         calls = []
diff --git a/Lib/test/test_iter.py b/Lib/test/test_iter.py
index fe8617309da29d..4660b6cabbb4ce 100644
--- a/Lib/test/test_iter.py
+++ b/Lib/test/test_iter.py
@@ -507,6 +507,24 @@ def test_calliter_setstate(self):
         it.__setstate__(((10,), StopIteration))
         self.assertEqual(list(it), list(range(10)))
 
+    def test_calliter_sentinel_reentrant_setstate(self):
+        # gh-158031: a sentinel __eq__ that mutates the iterator
+        # re-entrantly must not leave the comparison using a freed
+        # sentinel.
+        class Sentinel:
+            def __eq__(self, other):
+                it.__setstate__(((), StopIteration))
+                return NotImplemented
+
+        class Result:
+            def __eq__(self, other):
+                return NotImplemented
+
+        it = iter(lambda: Result(), Sentinel())
+        self.assertIsInstance(next(it), Result)
+        # __setstate__ cleared the sentinel; iteration still works.
+        self.assertIsInstance(next(it), Result)
+
     def test_iter_function_concealing_reentrant_exhaustion(self):
         # gh-101892: Test two-argument iter() with a function that
         # exhausts its associated iterator but forgets to either return
diff --git a/Misc/NEWS.d/next/Core_and_Builtins/2026-09-23-23-45-00.gh-issue-158031.itSentUaf.rst b/Misc/NEWS.d/next/Core_and_Builtins/2026-09-23-23-45-00.gh-issue-158031.itSentUaf.rst
new file mode 100644
index 00000000000000..1532b8783a0250
--- /dev/null
+++ b/Misc/NEWS.d/next/Core_and_Builtins/2026-09-23-23-45-00.gh-issue-158031.itSentUaf.rst
@@ -0,0 +1,5 @@
+Fix use-after-free in :func:`iter` and :func:`aiter` callables with a
+sentinel: the sentinel comparison could run code that replaced or released
+the sentinel (for example via ``__setstate__`` or a re-entrant
+``__anext__``), leaving the comparison using freed memory.  A strong
+reference is now held for the duration of the comparison.
diff --git a/Objects/iterobject.c b/Objects/iterobject.c
index 2d5e3709a27dfb..7895bbe49b8476 100644
--- a/Objects/iterobject.c
+++ b/Objects/iterobject.c
@@ -282,7 +282,12 @@ calliter_iternext(PyObject *op)
         if (it->it_sentinel == NULL) {
             return result; /* Common case, fast path */
         }
-        int ok = PyObject_RichCompareBool(it->it_sentinel, result, Py_EQ);
+        /* The comparison can run code that mutates the iterator
+           (e.g. __setstate__), so hold a strong reference to the
+           sentinel while it is in use. */
+        PyObject *sentinel = Py_NewRef(it->it_sentinel);
+        int ok = PyObject_RichCompareBool(sentinel, result, Py_EQ);
+        Py_DECREF(sentinel);
         if (ok == 0) {
             return result; /* Common case, fast path */
         }
@@ -641,7 +646,12 @@ acallawaitable_handle_error(acallawaitableobject *aw)
         }
         int ok = 0;
         if (it->it_sentinel != NULL) {
-            ok = PyObject_RichCompareBool(it->it_sentinel, value, Py_EQ);
+            /* The comparison can run code that exhausts the iterator
+               re-entrantly, so hold a strong reference to the sentinel
+               while it is in use. */
+            PyObject *sentinel = Py_NewRef(it->it_sentinel);
+            ok = PyObject_RichCompareBool(sentinel, value, Py_EQ);
+            Py_DECREF(sentinel);
         }
         if (ok == 0) {
             (void)_PyGen_SetStopIterationValue(value);