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

[tonghuaroot]

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

[terryjreedy]

On Windows main and older and macOS 10.5 13.5,

    c = itertools.count(1 << 100, Step())
TypeError: a number is required

the step argument has to be a number, and likely an int, not a Step(). How did that run?

[tonghuaroot]

You're right, sorry -- the step has to pass the number check, so a bare class fails. It needs __index__. I have updated the issue with a working reproducer.

The short version: the step's __radd__ calls next() on the same count. That inner call frees the running total, and the outer call still returns it as a dangling pointer. With PYTHONMALLOC=debug it segfaults on unfixed main; with the fix it prints the right value.