From 533a674a1832086d8417ec72be64a0cdd4da4b6a Mon Sep 17 00:00:00 2001 From: tonghuaroot Date: Sat, 25 Jul 2026 13:39:42 +0800 Subject: [PATCH 1/4] gh-154670: Fix use-after-free in itertools.count via re-entrant step --- Lib/test/test_itertools.py | 9 +++++++++ .../2026-07-25-12-00-00.gh-issue-154670.CtUAF1.rst | 2 ++ Modules/itertoolsmodule.c | 7 +++---- 3 files changed, 14 insertions(+), 4 deletions(-) create mode 100644 Misc/NEWS.d/next/Library/2026-07-25-12-00-00.gh-issue-154670.CtUAF1.rst diff --git a/Lib/test/test_itertools.py b/Lib/test/test_itertools.py index cf579d4da4e0dfb..89afd62be5305af 100644 --- a/Lib/test/test_itertools.py +++ b/Lib/test/test_itertools.py @@ -624,6 +624,15 @@ def counting_thread(): def test_count_with_step_threading(self): self.test_count_threading(step=5) + def test_count_reentrant_step(self): + c = count(1 << 100, Step()) + class Step: + def __radd__(self, other): + 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; } From 7742e77e7f672fab7f5befb8c2d405de816ed53e Mon Sep 17 00:00:00 2001 From: tonghuaroot Date: Sat, 25 Jul 2026 13:50:34 +0800 Subject: [PATCH 2/4] Fix test: define Step class before using it --- Lib/test/test_itertools.py | 1 - 1 file changed, 1 deletion(-) diff --git a/Lib/test/test_itertools.py b/Lib/test/test_itertools.py index 89afd62be5305af..d56f606c76de8e1 100644 --- a/Lib/test/test_itertools.py +++ b/Lib/test/test_itertools.py @@ -625,7 +625,6 @@ def test_count_with_step_threading(self): self.test_count_threading(step=5) def test_count_reentrant_step(self): - c = count(1 << 100, Step()) class Step: def __radd__(self, other): next(c) From 05c5bd39f9cd73ba9d0f81aeb7a3a6f123ca204c Mon Sep 17 00:00:00 2001 From: tonghuaroot Date: Sat, 25 Jul 2026 14:02:52 +0800 Subject: [PATCH 3/4] Make Step subclass int to pass PyNumber_Check on all platforms --- Lib/test/test_itertools.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Lib/test/test_itertools.py b/Lib/test/test_itertools.py index d56f606c76de8e1..c9009644852b993 100644 --- a/Lib/test/test_itertools.py +++ b/Lib/test/test_itertools.py @@ -625,7 +625,7 @@ def test_count_with_step_threading(self): self.test_count_threading(step=5) def test_count_reentrant_step(self): - class Step: + class Step(int): def __radd__(self, other): next(c) return other + 1 From 1e195e26c1c55c108b9671b0bb30ed21a999283b Mon Sep 17 00:00:00 2001 From: tonghuaroot Date: Sat, 25 Jul 2026 14:21:14 +0800 Subject: [PATCH 4/4] Guard re-entrant test against infinite recursion --- Lib/test/test_itertools.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/Lib/test/test_itertools.py b/Lib/test/test_itertools.py index c9009644852b993..a7cea6dde47ea9e 100644 --- a/Lib/test/test_itertools.py +++ b/Lib/test/test_itertools.py @@ -625,9 +625,13 @@ 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): - next(c) + nonlocal entered + if not entered: + entered = True + next(c) return other + 1 c = count(1 << 100, Step()) next(c)