Skip to content

Commit fbaa55b

Browse files
committed
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.
1 parent 09bf4c5 commit fbaa55b

4 files changed

Lines changed: 66 additions & 2 deletions

File tree

‎Lib/test/test_asyncgen.py‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -905,6 +905,37 @@ def test_aiter_callable_exhausted(self):
905905
with self.assertRaises(StopAsyncIteration):
906906
self.loop.run_until_complete(anext(it))
907907

908+
def test_aiter_callable_sentinel_reentrant_exhaustion(self):
909+
# gh-158033: a sentinel __eq__ that exhausts the iterator
910+
# re-entrantly must not leave the comparison using a freed
911+
# sentinel.
912+
state = {'stop': False}
913+
914+
async def produce():
915+
return Result()
916+
917+
def spam():
918+
if state['stop']:
919+
raise StopAsyncIteration
920+
return produce()
921+
922+
class Sentinel:
923+
def __eq__(self, other):
924+
state['stop'] = True
925+
try:
926+
ait.__anext__().__await__().send(None)
927+
except StopAsyncIteration:
928+
pass
929+
return NotImplemented
930+
931+
class Result:
932+
def __eq__(self, other):
933+
return NotImplemented
934+
935+
ait = aiter(spam, Sentinel())
936+
with self.assertRaises(StopIteration):
937+
ait.__anext__().__await__().send(None)
938+
908939
def test_aiter_callable_lazy(self):
909940
# The callable is only called when the awaitable is awaited
910941
calls = []

‎Lib/test/test_iter.py‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -507,6 +507,24 @@ def test_calliter_setstate(self):
507507
it.__setstate__(((10,), StopIteration))
508508
self.assertEqual(list(it), list(range(10)))
509509

510+
def test_calliter_sentinel_reentrant_setstate(self):
511+
# gh-158031: a sentinel __eq__ that mutates the iterator
512+
# re-entrantly must not leave the comparison using a freed
513+
# sentinel.
514+
class Sentinel:
515+
def __eq__(self, other):
516+
it.__setstate__(((), StopIteration))
517+
return NotImplemented
518+
519+
class Result:
520+
def __eq__(self, other):
521+
return NotImplemented
522+
523+
it = iter(lambda: Result(), Sentinel())
524+
self.assertIsInstance(next(it), Result)
525+
# __setstate__ cleared the sentinel; iteration still works.
526+
self.assertIsInstance(next(it), Result)
527+
510528
def test_iter_function_concealing_reentrant_exhaustion(self):
511529
# gh-101892: Test two-argument iter() with a function that
512530
# exhausts its associated iterator but forgets to either return
Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
Fix use-after-free in :func:`iter` and :func:`aiter` callables with a
2+
sentinel: the sentinel comparison could run code that replaced or released
3+
the sentinel (for example via ``__setstate__`` or a re-entrant
4+
``__anext__``), leaving the comparison using freed memory. A strong
5+
reference is now held for the duration of the comparison.

‎Objects/iterobject.c‎

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -282,7 +282,12 @@ calliter_iternext(PyObject *op)
282282
if (it->it_sentinel == NULL) {
283283
return result; /* Common case, fast path */
284284
}
285-
int ok = PyObject_RichCompareBool(it->it_sentinel, result, Py_EQ);
285+
/* The comparison can run code that mutates the iterator
286+
(e.g. __setstate__), so hold a strong reference to the
287+
sentinel while it is in use. */
288+
PyObject *sentinel = Py_NewRef(it->it_sentinel);
289+
int ok = PyObject_RichCompareBool(sentinel, result, Py_EQ);
290+
Py_DECREF(sentinel);
286291
if (ok == 0) {
287292
return result; /* Common case, fast path */
288293
}
@@ -641,7 +646,12 @@ acallawaitable_handle_error(acallawaitableobject *aw)
641646
}
642647
int ok = 0;
643648
if (it->it_sentinel != NULL) {
644-
ok = PyObject_RichCompareBool(it->it_sentinel, value, Py_EQ);
649+
/* The comparison can run code that exhausts the iterator
650+
re-entrantly, so hold a strong reference to the sentinel
651+
while it is in use. */
652+
PyObject *sentinel = Py_NewRef(it->it_sentinel);
653+
ok = PyObject_RichCompareBool(sentinel, value, Py_EQ);
654+
Py_DECREF(sentinel);
645655
}
646656
if (ok == 0) {
647657
(void)_PyGen_SetStopIterationValue(value);

0 commit comments

Comments
 (0)