itertools.count use-after-free via re-entrant step.__radd__

Bug description

itertools.count has a use-after-free when the step object's __radd__ re-enters the iterator.

count_nextlong() borrows lz->long_cnt (the running total) without Py_INCREF, then calls PyNumber_Add(result, step). If step.__radd__ calls next() on the same count, the inner call moves lz->long_cnt to the new value and hands the old total's only reference back, which is then dropped and freed. The outer call still returns that freed total as a dangling pointer.

The step needs __index__ so count() accepts it as a number, and the returned value has to be used to hit the freed memory:

from itertools import count

class Step:
    armed = True
    def __index__(self):        # so count() accepts it as a number
        return 1
    def __radd__(self, other):
        if Step.armed:
            Step.armed = False
            inner = next(c)     # re-enter; steals the running total's ref
            f"{inner!r}"        # churn the heap so the freed slot is reused
        return other + 1

c = count(1 << 100, Step())
val = next(c)
val + 1                         # use the returned (dangling) total

Run with PYTHONMALLOC=debug python repro.py -> SIGSEGV. With the fix it prints the correct value.

CPython versions tested on

main

Operating systems tested on

macOS

Bắt đầu bằng cách lần theo đường đi của count_nextlong() được mô tả trong báo cáo và chạy reproducer được cung cấp với PYTHONMALLOC=debug. Điều tra lệnh gọi next() tái nhập và thời gian tồn tại của tham chiếu xung quanh tổng đang chạy; công việc hoàn tất khi reproducer in ra giá trị chính xác mà không bị crash.