From a6a20ba8b355359735262d0b3c7977e451634956 Mon Sep 17 00:00:00 2001 From: tonghuaroot Date: Sat, 25 Jul 2026 13:39:42 +0800 Subject: [PATCH 1/5] gh-000000: 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-000000.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-000000.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-000000.CtUAF1.rst b/Misc/NEWS.d/next/Library/2026-07-25-12-00-00.gh-issue-000000.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-000000.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 cd78520336d009f468b420ad97b1399d448bf45c Mon Sep 17 00:00:00 2001 From: tonghuaroot Date: Sat, 25 Jul 2026 13:41:06 +0800 Subject: [PATCH 2/5] gh-154672: Fix use-after-free in itertools.zip_longest via re-entrant iterator --- Lib/test/test_itertools.py | 25 ++++++++++++------- ...-07-25-12-30-00.gh-issue-154672.ZlUAF2.rst | 2 ++ Modules/itertoolsmodule.c | 25 +++++++++++++------ 3 files changed, 36 insertions(+), 16 deletions(-) create mode 100644 Misc/NEWS.d/next/Library/2026-07-25-12-30-00.gh-issue-154672.ZlUAF2.rst diff --git a/Lib/test/test_itertools.py b/Lib/test/test_itertools.py index 89afd62be5305af..ade95dc9644dadd 100644 --- a/Lib/test/test_itertools.py +++ b/Lib/test/test_itertools.py @@ -624,15 +624,6 @@ 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('')), []) @@ -1514,6 +1505,22 @@ def test_zip_longest_result_gc(self): gc.collect() self.assertTrue(gc.is_tracked(next(it))) + def test_zip_longest_reentrant_iterator(self): + # A re-entrant iterator exhaust must not cause use-after-free. + zl = None + class ReentrantFilter: + def __init__(self, it): + self.it = it + def __iter__(self): + return self + def __next__(self): + if zl is not None: + for _ in zl: + pass + return next(self.it) + zl = zip_longest(ReentrantFilter(iter([1])), iter([2, 3])) + list(zl) + @support.cpython_only def test_pairwise_result_gc(self): # Ditto for pairwise. diff --git a/Misc/NEWS.d/next/Library/2026-07-25-12-30-00.gh-issue-154672.ZlUAF2.rst b/Misc/NEWS.d/next/Library/2026-07-25-12-30-00.gh-issue-154672.ZlUAF2.rst new file mode 100644 index 000000000000000..0aa2aa006033db7 --- /dev/null +++ b/Misc/NEWS.d/next/Library/2026-07-25-12-30-00.gh-issue-154672.ZlUAF2.rst @@ -0,0 +1,2 @@ +Fix use-after-free in :func:`itertools.zip_longest` when a re-entrant +iterator callback exhausts a sibling iterator. Patch by tonghuaroot. diff --git a/Modules/itertoolsmodule.c b/Modules/itertoolsmodule.c index 19a690025b0190d..5b71b2db09b55bf 100644 --- a/Modules/itertoolsmodule.c +++ b/Modules/itertoolsmodule.c @@ -3628,14 +3628,15 @@ count_nextlong(countobject *lz) } assert(lz->cnt == PY_SSIZE_T_MAX && lz->long_cnt != NULL); - PyObject *result = Py_NewRef(lz->long_cnt); + // 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 *stepped_up = PyNumber_Add(result, lz->long_step); if (stepped_up == NULL) { - Py_DECREF(result); return NULL; } - Py_SETREF(lz->long_cnt, stepped_up); + lz->long_cnt = stepped_up; return result; } @@ -3976,19 +3977,24 @@ zip_longest_next_lock_held(PyObject *op) if (it == NULL) { item = Py_NewRef(lz->fillvalue); } else { + Py_INCREF(it); item = PyIter_Next(it); if (item == NULL) { lz->numactive -= 1; if (lz->numactive == 0 || PyErr_Occurred()) { lz->numactive = 0; + Py_DECREF(it); Py_DECREF(result); return NULL; } else { item = Py_NewRef(lz->fillvalue); - PyTuple_SET_ITEM(lz->ittuple, i, NULL); - Py_DECREF(it); + if (PyTuple_GET_ITEM(lz->ittuple, i) != NULL) { + PyTuple_SET_ITEM(lz->ittuple, i, NULL); + Py_DECREF(it); + } } } + Py_DECREF(it); } olditem = PyTuple_GET_ITEM(result, i); PyTuple_SET_ITEM(result, i, item); @@ -4006,19 +4012,24 @@ zip_longest_next_lock_held(PyObject *op) if (it == NULL) { item = Py_NewRef(lz->fillvalue); } else { + Py_INCREF(it); item = PyIter_Next(it); if (item == NULL) { lz->numactive -= 1; if (lz->numactive == 0 || PyErr_Occurred()) { lz->numactive = 0; + Py_DECREF(it); Py_DECREF(result); return NULL; } else { item = Py_NewRef(lz->fillvalue); - PyTuple_SET_ITEM(lz->ittuple, i, NULL); - Py_DECREF(it); + if (PyTuple_GET_ITEM(lz->ittuple, i) != NULL) { + PyTuple_SET_ITEM(lz->ittuple, i, NULL); + Py_DECREF(it); + } } } + Py_DECREF(it); } PyTuple_SET_ITEM(result, i, item); } From 1f51f66f73071c2ad1be0624b4a91d8555910a6c Mon Sep 17 00:00:00 2001 From: tonghuaroot Date: Sat, 25 Jul 2026 13:52:50 +0800 Subject: [PATCH 3/5] Remove stale NEWS file from count fix --- .../next/Library/2026-07-25-12-00-00.gh-issue-000000.CtUAF1.rst | 2 -- 1 file changed, 2 deletions(-) delete mode 100644 Misc/NEWS.d/next/Library/2026-07-25-12-00-00.gh-issue-000000.CtUAF1.rst diff --git a/Misc/NEWS.d/next/Library/2026-07-25-12-00-00.gh-issue-000000.CtUAF1.rst b/Misc/NEWS.d/next/Library/2026-07-25-12-00-00.gh-issue-000000.CtUAF1.rst deleted file mode 100644 index dbb03b10c69fdd4..000000000000000 --- a/Misc/NEWS.d/next/Library/2026-07-25-12-00-00.gh-issue-000000.CtUAF1.rst +++ /dev/null @@ -1,2 +0,0 @@ -Fix use-after-free in :func:`itertools.count` when the step object's -``__radd__`` re-enters the iterator. Patch by tonghuaroot. From f21835bd09864fab3b9979dbd0a8e76bafab58c3 Mon Sep 17 00:00:00 2001 From: tonghuaroot Date: Sat, 25 Jul 2026 14:08:38 +0800 Subject: [PATCH 4/5] Fix test: guard against infinite recursion in re-entrant iterator --- Lib/test/test_itertools.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/Lib/test/test_itertools.py b/Lib/test/test_itertools.py index ade95dc9644dadd..eef899219c57d3f 100644 --- a/Lib/test/test_itertools.py +++ b/Lib/test/test_itertools.py @@ -1507,18 +1507,20 @@ def test_zip_longest_result_gc(self): def test_zip_longest_reentrant_iterator(self): # A re-entrant iterator exhaust must not cause use-after-free. - zl = None - class ReentrantFilter: + entered = False + class ReentrantIter: def __init__(self, it): self.it = it def __iter__(self): return self def __next__(self): - if zl is not None: + nonlocal entered + if not entered: + entered = True for _ in zl: pass return next(self.it) - zl = zip_longest(ReentrantFilter(iter([1])), iter([2, 3])) + zl = zip_longest(ReentrantIter(iter([1])), iter([2, 3])) list(zl) @support.cpython_only From 54a489ad410f4c60bedbb8b072d032a97fdccb68 Mon Sep 17 00:00:00 2001 From: tonghuaroot Date: Sat, 25 Jul 2026 14:21:01 +0800 Subject: [PATCH 5/5] Rewrite test to exhaust sibling iterator without self-recursion --- Lib/test/test_itertools.py | 20 +++++++++++--------- 1 file changed, 11 insertions(+), 9 deletions(-) diff --git a/Lib/test/test_itertools.py b/Lib/test/test_itertools.py index eef899219c57d3f..784b8ba7ea8e2e2 100644 --- a/Lib/test/test_itertools.py +++ b/Lib/test/test_itertools.py @@ -1507,20 +1507,22 @@ def test_zip_longest_result_gc(self): def test_zip_longest_reentrant_iterator(self): # A re-entrant iterator exhaust must not cause use-after-free. - entered = False - class ReentrantIter: - def __init__(self, it): + class ExhaustingIter: + def __init__(self, it, sibling): self.it = it + self.sibling = sibling + self.first = True def __iter__(self): return self def __next__(self): - nonlocal entered - if not entered: - entered = True - for _ in zl: + val = next(self.it) + if self.first: + self.first = False + for _ in self.sibling: pass - return next(self.it) - zl = zip_longest(ReentrantIter(iter([1])), iter([2, 3])) + return val + sib = iter([2, 3]) + zl = zip_longest(ExhaustingIter(iter([1, 4]), sib), sib) list(zl) @support.cpython_only