itertools.groupby: use-after-free in _grouper_next() when __eq__ re-enters the iterator #145678
Description
Description
_grouper_next() in Modules/itertoolsmodule.c passes igo->tgtkey and gbo->currkey directly to PyObject_RichCompareBool() without holding strong references to them first.
A user-defined __eq__ method can re-enter the parent groupby iterator during that comparison. Re-entry calls groupby_step(), which contains:
Py_XSETREF(gbo->currkey, newkey);This frees gbo->currkey while it is still being compared — a use-after-free.
Relationship to gh-143543
gh-143543 fixed the same bug in groupby_next() by taking INCREF'd local snapshots before calling PyObject_RichCompareBool(). The fix was not applied to the analogous code in _grouper_next().
Fix
Apply the same INCREF snapshot pattern used in the groupby_next() fix:
PyObject *tgtkey = igo->tgtkey;
PyObject *currkey = gbo->currkey;
Py_INCREF(tgtkey);
Py_INCREF(currkey);
rcmp = PyObject_RichCompareBool(tgtkey, currkey, Py_EQ);
Py_DECREF(tgtkey);
Py_DECREF(currkey);Reproducer
import itertools
outer_grouper = None
class Key:
def __init__(self, val, do_advance):
self.val = val
self.do_advance = do_advance
def __eq__(self, other):
if self.do_advance:
self.do_advance = False
try:
next(outer_grouper)
except StopIteration:
pass
return NotImplemented
return self.val == other.val
def __hash__(self):
return hash(self.val)
values = [1, 1, 2]
keys_iter = iter([Key(1, True), Key(1, False), Key(2, False)])
g = itertools.groupby(values, lambda _: next(keys_iter))
outer_grouper = g
k, grp = next(g)
list(grp) # use-after-free / crash under ASAN