diff --git a/Lib/test/test_itertools.py b/Lib/test/test_itertools.py index cf579d4da4e0dfb..a7cea6dde47ea9e 100644 --- a/Lib/test/test_itertools.py +++ b/Lib/test/test_itertools.py @@ -624,6 +624,18 @@ def counting_thread(): def test_count_with_step_threading(self): self.test_count_threading(step=5) + def test_count_reentrant_step(self): + entered = False + class Step(int): + def __radd__(self, other): + nonlocal entered + if not entered: + entered = True + next(c) + return other + 1 + c = count(1 << 100, Step()) + next(c) + def test_cycle(self): self.assertEqual(take(10, cycle('abc')), list('abcabcabca')) self.assertEqual(list(cycle('')), []) diff --git a/Misc/NEWS.d/next/Library/2026-07-25-12-00-00.gh-issue-154670.CtUAF1.rst b/Misc/NEWS.d/next/Library/2026-07-25-12-00-00.gh-issue-154670.CtUAF1.rst new file mode 100644 index 000000000000000..dbb03b10c69fdd4 --- /dev/null +++ b/Misc/NEWS.d/next/Library/2026-07-25-12-00-00.gh-issue-154670.CtUAF1.rst @@ -0,0 +1,2 @@ +Fix use-after-free in :func:`itertools.count` when the step object's +``__radd__`` re-enters the iterator. Patch by tonghuaroot. diff --git a/Modules/itertoolsmodule.c b/Modules/itertoolsmodule.c index c0023c839ca7fe3..19a690025b0190d 100644 --- a/Modules/itertoolsmodule.c +++ b/Modules/itertoolsmodule.c @@ -3628,15 +3628,14 @@ count_nextlong(countobject *lz) } assert(lz->cnt == PY_SSIZE_T_MAX && lz->long_cnt != NULL); - // We hold one reference to "result" (a.k.a. the old value of - // lz->long_cnt); we'll either return it or keep it in lz->long_cnt. - PyObject *result = lz->long_cnt; + PyObject *result = Py_NewRef(lz->long_cnt); PyObject *stepped_up = PyNumber_Add(result, lz->long_step); if (stepped_up == NULL) { + Py_DECREF(result); return NULL; } - lz->long_cnt = stepped_up; + Py_SETREF(lz->long_cnt, stepped_up); return result; }