Skip to content

Commit 9aa672e

Browse files
committed
gh-146613: Fix re-entrant use-after-free in itertools._grouper
The same pattern was fixed in groupby.__next__ (gh-143543 / a91b5c3), but _grouper_next (the inner group iterator returned by groupby) was missed. A user-defined __eq__ can re-enter the grouper during PyObject_RichCompareBool, causing Py_XSETREF to free currkey while it is still being used. Fix by taking local snapshots of tgtkey/currkey + INCREF/DECREF protection, exactly as done in groupby_next. Added regression test in test_itertools.py.
1 parent 4497cf3 commit 9aa672e

File tree

2 files changed

+47
-1
lines changed

2 files changed

+47
-1
lines changed

Lib/test/test_itertools.py

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -754,6 +754,40 @@ def keys():
754754
next(g)
755755
next(g) # must pass with address sanitizer
756756

757+
def test_grouper_reentrant_eq_does_not_crash(self):
758+
# regression test for gh-146613
759+
grouper_iter = None
760+
class Key:
761+
__hash__ = None
762+
763+
def __init__(self, do_advance):
764+
self.do_advance = do_advance
765+
self.payload = bytearray(256)
766+
767+
def __eq__(self, other):
768+
nonlocal grouper_iter
769+
if self.do_advance:
770+
self.do_advance = False
771+
if grouper_iter is not None:
772+
try:
773+
next(grouper_iter)
774+
except StopIteration:
775+
pass
776+
for _ in range(50):
777+
bytearray(256)
778+
return NotImplemented
779+
return True
780+
781+
def keyfunc(element):
782+
if element == 0:
783+
return Key(do_advance=True)
784+
return Key(do_advance=False)
785+
786+
g = itertools.groupby(range(4), keyfunc)
787+
key, grouper_iter = next(g)
788+
items = list(grouper_iter)
789+
self.assertEqual(len(items), 1)
790+
757791
def test_filter(self):
758792
self.assertEqual(list(filter(isEven, range(6))), [0,2,4])
759793
self.assertEqual(list(filter(None, [0,1,0,2,0])), [1,2])

Modules/itertoolsmodule.c

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -678,7 +678,19 @@ _grouper_next(PyObject *op)
678678
}
679679

680680
assert(gbo->currkey != NULL);
681-
rcmp = PyObject_RichCompareBool(igo->tgtkey, gbo->currkey, Py_EQ);
681+
/* A user-defined __eq__ can re-enter the grouper and advance the iterator,
682+
mutating gbo->currkey while we are comparing them.
683+
Take local snapshots and hold strong references so INCREF/DECREF
684+
apply to the same objects even under re-entrancy. */
685+
PyObject *tgtkey = igo->tgtkey;
686+
PyObject *currkey = gbo->currkey;
687+
688+
Py_INCREF(tgtkey);
689+
Py_INCREF(currkey);
690+
rcmp = PyObject_RichCompareBool(tgtkey, currkey, Py_EQ);
691+
Py_DECREF(tgtkey);
692+
Py_DECREF(currkey);
693+
682694
if (rcmp <= 0)
683695
/* got any error or current group is end */
684696
return NULL;

0 commit comments

Comments
 (0)