Skip to content

Commit a91b5c3

Browse files
VanshAgarwal24036skirpichevencukou
authored
gh-143543: Fix re-entrant use-after-free in itertools.groupby (GH-143738)
Co-authored-by: Sergey B Kirpichev <skirpichev@gmail.com> Co-authored-by: Petr Viktorin <encukou@gmail.com>
1 parent 988286e commit a91b5c3

File tree

3 files changed

+35
-2
lines changed

3 files changed

+35
-2
lines changed

Lib/test/test_itertools.py

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -733,6 +733,27 @@ def keyfunc(obj):
733733
keyfunc.skip = 1
734734
self.assertRaises(ExpectedError, gulp, [None, None], keyfunc)
735735

736+
def test_groupby_reentrant_eq_does_not_crash(self):
737+
# regression test for gh-143543
738+
class Key:
739+
def __init__(self, do_advance):
740+
self.do_advance = do_advance
741+
742+
def __eq__(self, other):
743+
if self.do_advance:
744+
self.do_advance = False
745+
next(g)
746+
return NotImplemented
747+
return False
748+
749+
def keys():
750+
yield Key(True)
751+
yield Key(False)
752+
753+
g = itertools.groupby([None, None], keys().send)
754+
next(g)
755+
next(g) # must pass with address sanitizer
756+
736757
def test_filter(self):
737758
self.assertEqual(list(filter(isEven, range(6))), [0,2,4])
738759
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+
Fix a crash in itertools.groupby that could occur when a user-defined
2+
:meth:`~object.__eq__` method re-enters the iterator during key comparison.

Modules/itertoolsmodule.c

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -544,9 +544,19 @@ groupby_next(PyObject *op)
544544
else if (gbo->tgtkey == NULL)
545545
break;
546546
else {
547-
int rcmp;
547+
/* A user-defined __eq__ can re-enter groupby and advance the iterator,
548+
mutating gbo->tgtkey / gbo->currkey while we are comparing them.
549+
Take local snapshots and hold strong references so INCREF/DECREF
550+
apply to the same objects even under re-entrancy. */
551+
PyObject *tgtkey = gbo->tgtkey;
552+
PyObject *currkey = gbo->currkey;
553+
554+
Py_INCREF(tgtkey);
555+
Py_INCREF(currkey);
556+
int rcmp = PyObject_RichCompareBool(tgtkey, currkey, Py_EQ);
557+
Py_DECREF(tgtkey);
558+
Py_DECREF(currkey);
548559

549-
rcmp = PyObject_RichCompareBool(gbo->tgtkey, gbo->currkey, Py_EQ);
550560
if (rcmp == -1)
551561
return NULL;
552562
else if (rcmp == 0)

0 commit comments

Comments
 (0)