gh-154672: Fix use-after-free in itertools.zip_longest via re-entrant iterator

Author: tonghuaroot

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.

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.

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);
         }