commit 993811ffe75c2573f97fb3fd1414b34609b8c8db
author Miss Islington (bot) <31488909+miss-islington@users.noreply.github.com> Thu Jan 23 07:22:25 2020 -0800
committer GitHub <noreply@github.com> Thu Jan 23 07:22:25 2020 -0800
tree ea16d4b5c8c75cb7cc1c0f9e33082b6b016f837e
parent 545fc51d950558ecec9ff64cb2f9c11469051524

bpo-39421: Fix posible crash in heapq with custom comparison operators (GH-18118)


* bpo-39421: Fix posible crash in heapq with custom comparison operators

* fixup! bpo-39421: Fix posible crash in heapq with custom comparison operators

* fixup! fixup! bpo-39421: Fix posible crash in heapq with custom comparison operators
(cherry picked from commit 79f89e6e5a659846d1068e8b1bd8e491ccdef861)

Co-authored-by: Pablo Galindo <Pablogsal@gmail.com>

Modules/_heapqmodule.c

diff --git a/Modules/_heapqmodule.c b/Modules/_heapqmodule.c
index a84cade..6bc18b5 100644
--- a/Modules/_heapqmodule.c
+++ b/Modules/_heapqmodule.c
@@ -36,7 +36,11 @@
     while (pos > startpos) {
         parentpos = (pos - 1) >> 1;
         parent = arr[parentpos];
+        Py_INCREF(newitem);
+        Py_INCREF(parent);
         cmp = PyObject_RichCompareBool(newitem, parent, Py_LT);
+        Py_DECREF(parent);
+        Py_DECREF(newitem);
         if (cmp < 0)
             return -1;
         if (size != PyList_GET_SIZE(heap)) {
@@ -78,10 +82,13 @@
         /* Set childpos to index of smaller child.   */
         childpos = 2*pos + 1;    /* leftmost child position  */
         if (childpos + 1 < endpos) {
-            cmp = PyObject_RichCompareBool(
-                arr[childpos],
-                arr[childpos + 1],
-                Py_LT);
+            PyObject* a = arr[childpos];
+            PyObject* b = arr[childpos + 1];
+            Py_INCREF(a);
+            Py_INCREF(b);
+            cmp = PyObject_RichCompareBool(a, b, Py_LT);
+            Py_DECREF(a);
+            Py_DECREF(b);
             if (cmp < 0)
                 return -1;
             childpos += ((unsigned)cmp ^ 1);   /* increment when cmp==0 */
@@ -264,7 +271,10 @@
         return item;
     }
 
-    cmp = PyObject_RichCompareBool(PyList_GET_ITEM(heap, 0), item, Py_LT);
+    PyObject* top = PyList_GET_ITEM(heap, 0);
+    Py_INCREF(top);
+    cmp = PyObject_RichCompareBool(top, item, Py_LT);
+    Py_DECREF(top);
     if (cmp < 0)
         return NULL;
     if (cmp == 0) {
@@ -420,7 +430,11 @@
     while (pos > startpos) {
         parentpos = (pos - 1) >> 1;
         parent = arr[parentpos];
+        Py_INCREF(parent);
+        Py_INCREF(newitem);
         cmp = PyObject_RichCompareBool(parent, newitem, Py_LT);
+        Py_DECREF(parent);
+        Py_DECREF(newitem);
         if (cmp < 0)
             return -1;
         if (size != PyList_GET_SIZE(heap)) {
@@ -462,10 +476,13 @@
         /* Set childpos to index of smaller child.   */
         childpos = 2*pos + 1;    /* leftmost child position  */
         if (childpos + 1 < endpos) {
-            cmp = PyObject_RichCompareBool(
-                arr[childpos + 1],
-                arr[childpos],
-                Py_LT);
+            PyObject* a = arr[childpos + 1];
+            PyObject* b = arr[childpos];
+            Py_INCREF(a);
+            Py_INCREF(b);
+            cmp = PyObject_RichCompareBool(a, b, Py_LT);
+            Py_DECREF(a);
+            Py_DECREF(b);
             if (cmp < 0)
                 return -1;
             childpos += ((unsigned)cmp ^ 1);   /* increment when cmp==0 */