Skip to content
22 changes: 22 additions & 0 deletions Lib/test/test_itertools.py
Original file line number Diff line number Diff line change
Expand Up @@ -754,6 +754,28 @@ def keys():
next(g)
next(g) # must pass with address sanitizer

def test_grouper_next_reentrant_eq_does_not_crash(self):
# regression test for gh-145678
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
next(g)
return NotImplemented
return self.val == other.val

def __hash__(self):
return hash(self.val)

keys_iter = iter([Key(1, True), Key(1, False), Key(2, False)])
g = itertools.groupby([1, 1, 2], lambda _: next(keys_iter))
k, grp = next(g)
list(grp) # must not crash with address sanitizer

def test_filter(self):
self.assertEqual(list(filter(isEven, range(6))), [0,2,4])
self.assertEqual(list(filter(None, [0,1,0,2,0])), [1,2])
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
Fix a use-after-free crash in :func:`itertools.groupby` when the
parent iterator is mutated while the iterator of groups was advanced.
6 changes: 5 additions & 1 deletion Modules/itertoolsmodule.c
Original file line number Diff line number Diff line change
Expand Up @@ -678,7 +678,11 @@ _grouper_next(PyObject *op)
}

assert(gbo->currkey != NULL);
rcmp = PyObject_RichCompareBool(igo->tgtkey, gbo->currkey, Py_EQ);
PyObject *tgtkey = Py_NewRef(igo->tgtkey);
PyObject *currkey = Py_NewRef(gbo->currkey);
rcmp = PyObject_RichCompareBool(tgtkey, currkey, Py_EQ);
Py_DECREF(tgtkey);
Py_DECREF(currkey);
if (rcmp <= 0)
/* got any error or current group is end */
return NULL;
Expand Down
Loading