gh-154670: Fix use-after-free in itertools.count via re-entrant step

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('')), [])

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.

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