gh-153981: Fix use after free bug itertools.count()

Obtain a reference to lz->long_cnt, lz->long_step
such that we don't hit a use after free

Author: johng

Lib/test/test_free_threading/test_itertools.py

@@ -4,6 +4,7 @@
     batched,
     chain,
     combinations_with_replacement,
+    count,
     cycle,
     permutations,
     tee,
@@ -142,5 +143,32 @@ def consume():
         self.assertEqual(errors, [], msg=f"unexpected errors: {errors}")
 
 
+class TestCountConcurrent(unittest.TestCase):
+    @staticmethod
+    def _spin_next(it, n=2000):
+        for _ in range(n):
+            next(it)
+
+    @staticmethod
+    def _spin_repr(it, n=2000):
+        for _ in range(n):
+            repr(it)
+
+    @threading_helper.reap_threads
+    def test_repr_racing_next_fast_mode(self):
+        for _ in range(10):
+            it = count()
+            workers = [self._spin_next] * 2 + [self._spin_repr] * 4
+            threading_helper.run_concurrently(workers, args=(it,))
+
+    @threading_helper.reap_threads
+    def test_repr_racing_next_slow_mode(self):
+        for _ in range(10):
+            # Large count to trigger "slow mode"
+            it = count(10**18, 2)
+            workers = [self._spin_next] * 2 + [self._spin_repr] * 4
+            threading_helper.run_concurrently(workers, args=(it,))
+
+
 if __name__ == "__main__":
     unittest.main()

Modules/itertoolsmodule.c

@@ -3675,27 +3675,45 @@ static PyObject *
 count_repr(PyObject *op)
 {
     countobject *lz = countobject_CAST(op);
-    if (lz->long_cnt == NULL) {
+    PyObject *long_cnt, *long_step, *result;
+
+    Py_BEGIN_CRITICAL_SECTION(lz);
+    long_cnt = Py_XNewRef(lz->long_cnt);
+    long_step = Py_XNewRef(lz->long_step);
+    Py_END_CRITICAL_SECTION();
+
+    if (long_cnt == NULL) {
+        /* Fast mode: cnt is advanced by count_next()'s lock-free atomic CAS,
+           which never takes this critical section, so read it atomically. */
         Py_ssize_t cnt = FT_ATOMIC_LOAD_SSIZE_RELAXED(lz->cnt);
-        return PyUnicode_FromFormat("%s(%zd)",
-                                    _PyType_Name(Py_TYPE(lz)), cnt);
+        result = PyUnicode_FromFormat("%s(%zd)",
+                                      _PyType_Name(Py_TYPE(lz)), cnt);
     }
-
-    if (PyLong_Check(lz->long_step)) {
-        long step = PyLong_AsLong(lz->long_step);
-        if (step == -1 && PyErr_Occurred()) {
-            PyErr_Clear();
-        }
-        if (step == 1) {
+    else {
+        int hide_step = 0;
+        if (PyLong_Check(long_step)) {
+            long step = PyLong_AsLong(long_step);
+            if (step == -1 && PyErr_Occurred()) {
+                PyErr_Clear();
+            }
             /* Don't display step when it is an integer equal to 1 */
-            return PyUnicode_FromFormat("%s(%R)",
-                                        _PyType_Name(Py_TYPE(lz)),
-                                        lz->long_cnt);
+            hide_step = (step == 1);
+        }
+        if (hide_step) {
+            result = PyUnicode_FromFormat("%s(%R)",
+                                          _PyType_Name(Py_TYPE(lz)),
+                                          long_cnt);
+        }
+        else {
+            result = PyUnicode_FromFormat("%s(%R, %R)",
+                                          _PyType_Name(Py_TYPE(lz)),
+                                          long_cnt, long_step);
         }
     }
-    return PyUnicode_FromFormat("%s(%R, %R)",
-                                _PyType_Name(Py_TYPE(lz)),
-                                lz->long_cnt, lz->long_step);
+
+    Py_XDECREF(long_cnt);
+    Py_XDECREF(long_step);
+    return result;
 }
 
 static PyType_Slot count_slots[] = {