Skip to content

Commit 79bcea0

Browse files
sampsoncclaude
andcommitted
gh-145678: Fix use-after-free in itertools.groupby _grouper_next()
_grouper_next() passed igo->tgtkey and gbo->currkey directly to PyObject_RichCompareBool() without first holding strong references. A re-entrant __eq__ that advanced the parent groupby iterator would call groupby_step(), which executes Py_XSETREF(gbo->currkey, newkey), freeing currkey while it was still under comparison. Fix by taking INCREF'd local snapshots before the comparison, mirroring the protection added to groupby_next() in gh-143543. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
1 parent 0dfe649 commit 79bcea0

File tree

3 files changed

+52
-1
lines changed

3 files changed

+52
-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_next_reentrant_eq_does_not_crash(self):
758+
# regression test for gh-145678: _grouper_next() did not protect
759+
# gbo->currkey / igo->tgtkey before calling PyObject_RichCompareBool,
760+
# so a reentrant __eq__ that advanced the parent groupby could free
761+
# those objects while they were still being compared (use-after-free).
762+
outer_grouper = None
763+
764+
class Key:
765+
def __init__(self, val, do_advance):
766+
self.val = val
767+
self.do_advance = do_advance
768+
769+
def __eq__(self, other):
770+
if self.do_advance:
771+
self.do_advance = False
772+
# Advance the parent groupby iterator from inside __eq__,
773+
# which calls groupby_step() and frees the old currkey.
774+
try:
775+
next(outer_grouper)
776+
except StopIteration:
777+
pass
778+
return NotImplemented
779+
return self.val == other.val
780+
781+
def __hash__(self):
782+
return hash(self.val)
783+
784+
values = [1, 1, 2]
785+
keys_iter = iter([Key(1, True), Key(1, False), Key(2, False)])
786+
g = itertools.groupby(values, lambda _: next(keys_iter))
787+
outer_grouper = g
788+
k, grp = next(g)
789+
list(grp) # must not crash with address sanitizer
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])
Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
Fix a use-after-free in :func:`itertools.groupby`: the internal
2+
``_grouper_next()`` function did not hold strong references to
3+
``gbo->currkey`` and ``igo->tgtkey`` before calling
4+
:c:func:`PyObject_RichCompareBool`, so a re-entrant ``__eq__`` that advanced
5+
the parent iterator (triggering ``groupby_step()`` and ``Py_XSETREF`` on
6+
``currkey``) could free those objects while they were still being compared.
7+
The fix mirrors the protection added in :gh:`143543` for ``groupby_next()``.

Modules/itertoolsmodule.c

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -678,7 +678,17 @@ _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 groupby (via the parent iterator)
682+
and call groupby_step(), which frees gbo->currkey via Py_XSETREF while
683+
we are still comparing it. Take local snapshots with strong references
684+
so INCREF/DECREF apply to the same objects even under re-entrancy. */
685+
PyObject *tgtkey = igo->tgtkey;
686+
PyObject *currkey = gbo->currkey;
687+
Py_INCREF(tgtkey);
688+
Py_INCREF(currkey);
689+
rcmp = PyObject_RichCompareBool(tgtkey, currkey, Py_EQ);
690+
Py_DECREF(tgtkey);
691+
Py_DECREF(currkey);
682692
if (rcmp <= 0)
683693
/* got any error or current group is end */
684694
return NULL;

0 commit comments

Comments
 (0)