Free-threaded (no-GIL) use-after-free: borrowed dict key from PyDict_Next used across __hash__/__eq__ in assert_same_structure #143

"""dm-tree `tree.cc:552` — borrowed key from PyDict_Next used across arbitrary Python.

    while (PyDict_Next(o1, &pos, &key, nullptr)) {    /* key is borrowed */
      if (PyDict_GetItem(o2, key) == nullptr) {       /* runs __hash__/__eq__ */

`o1` and `o2` are the two structures the caller passed to
`tree.assert_same_structure`. Under Py_GIL_DISABLED another thread mutating `o1`
can free the borrowed `key` while the lookup in `o2` is running user code.

Two hazards stacked: PyDict_Next itself is documented as unsafe against
concurrent mutation, and the key it hands back is borrowed.

Keys here are a str subclass whose __hash__ is real Python, so the lookup in `o2`
genuinely leaves C.
"""
import gc
import sys
import threading

import tree

N_CMP, N_MUT, ROUNDS, SIZE = 8, 3, 4000, 32


class Key(str):
    """A str subclass: __hash__ goes through Python, so PyDict_GetItem leaves C."""
    __slots__ = ()

    def __hash__(self):
        return str.__hash__(self)


def main():
    print(f"py={sys.version.split()[0]}{sys.abiflags} "
          f"gil={getattr(sys, '_is_gil_enabled', lambda: True)()}", flush=True)
    a = {Key(f"k{i}"): i for i in range(SIZE)}
    b = {Key(f"k{i}"): i for i in range(SIZE)}
    stop = threading.Event()
    barrier = threading.Barrier(N_CMP + N_MUT + 1)
    done = [0] * N_CMP

    def comparer(t):
        barrier.wait()
        n = 0
        for _ in range(ROUNDS):
            try:
                tree.assert_same_structure(a, b)
            except Exception:
                pass
            n += 1
        done[t] = n

    def mutator():
        barrier.wait()
        i = 0
        while not stop.is_set():
            k = Key(f"k{i % SIZE}")
            a[k] = i               # replace: the old Key object loses its last ref
            i += 1

    cs = [threading.Thread(target=comparer, args=(t,)) for t in range(N_CMP)]
    ms = [threading.Thread(target=mutator) for _ in range(N_MUT)]
    gc.disable()
    for t in cs + ms:
        t.start()
    barrier.wait()
    for t in cs:
        t.join()
    stop.set()
    for t in ms:
        t.join()
    gc.enable()
    print(f"clean: {sum(done)} comparisons", flush=True)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())