Skip to content

Commit fc7a188

Browse files
authored
gh-146613: Fix re-entrant use-after-free in itertools._grouper (#147962)
1 parent c1a4112 commit fc7a188

File tree

3 files changed

+44
-1
lines changed

3 files changed

+44
-1
lines changed

Lib/test/test_itertools.py

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -754,6 +754,38 @@ 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+
761+
class Key:
762+
__hash__ = None
763+
764+
def __init__(self, do_advance):
765+
self.do_advance = do_advance
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+
return NotImplemented
777+
return True
778+
779+
def keyfunc(element):
780+
if element == 0:
781+
return Key(do_advance=True)
782+
return Key(do_advance=False)
783+
784+
g = itertools.groupby(range(4), keyfunc)
785+
key, grouper_iter = next(g)
786+
items = list(grouper_iter)
787+
self.assertEqual(len(items), 1)
788+
757789
def test_filter(self):
758790
self.assertEqual(list(filter(isEven, range(6))), [0,2,4])
759791
self.assertEqual(list(filter(None, [0,1,0,2,0])), [1,2])
Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
:mod:`itertools`: Fix a crash in :func:`itertools.groupby` when
2+
the grouper iterator is concurrently mutated.

Modules/itertoolsmodule.c

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -678,7 +678,16 @@ _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 = Py_NewRef(igo->tgtkey);
686+
PyObject *currkey = Py_NewRef(gbo->currkey);
687+
rcmp = PyObject_RichCompareBool(tgtkey, currkey, Py_EQ);
688+
Py_DECREF(tgtkey);
689+
Py_DECREF(currkey);
690+
682691
if (rcmp <= 0)
683692
/* got any error or current group is end */
684693
return NULL;

0 commit comments

Comments
 (0)