gh-143543: Fix re-entrant use-after-free in itertools.groupby

Author: VanshAgarwal24036

Lib/test/test_itertools.py

@@ -733,6 +733,22 @@ def keyfunc(obj):
         keyfunc.skip = 1
         self.assertRaises(ExpectedError, gulp, [None, None], keyfunc)
 
+    def test_groupby_reentrant_eq_does_not_crash(self):
+
+        class Key(bytearray):
+            seen = False
+            def __eq__(self, other):
+                if not Key.seen:
+                    Key.seen = True
+                    next(g)
+                return False
+
+        data = [Key(b"a"), Key(b"b")]
+
+        g = itertools.groupby(data)
+        next(g)
+        next(g)  # must not segfault
+
     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])

Modules/itertoolsmodule.c

@@ -545,8 +545,17 @@ groupby_next(PyObject *op)
             break;
         else {
             int rcmp;
+            PyObject *tgtkey = gbo->tgtkey;
+            PyObject *currkey = gbo->currkey;
 
-            rcmp = PyObject_RichCompareBool(gbo->tgtkey, gbo->currkey, Py_EQ);
+            Py_INCREF(tgtkey);
+            Py_INCREF(currkey);
+
+            rcmp = PyObject_RichCompareBool(tgtkey, currkey, Py_EQ);
+
+            Py_DECREF(tgtkey);
+            Py_DECREF(currkey);
+            
             if (rcmp == -1)
                 return NULL;
             else if (rcmp == 0)