diff --git a/Lib/test/test_itertools.py b/Lib/test/test_itertools.py index cf579d4da4e0df..784b8ba7ea8e2e 100644 --- a/Lib/test/test_itertools.py +++ b/Lib/test/test_itertools.py @@ -1505,6 +1505,26 @@ 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. + class ExhaustingIter: + def __init__(self, it, sibling): + self.it = it + self.sibling = sibling + self.first = True + def __iter__(self): + return self + def __next__(self): + val = next(self.it) + if self.first: + self.first = False + for _ in self.sibling: + pass + return val + sib = iter([2, 3]) + zl = zip_longest(ExhaustingIter(iter([1, 4]), sib), sib) + 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 00000000000000..0aa2aa006033db --- /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 c0023c839ca7fe..5b71b2db09b55b 100644 --- a/Modules/itertoolsmodule.c +++ b/Modules/itertoolsmodule.c @@ -3977,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); @@ -4007,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); }