From c563f409ea30bcb0623d785428c9257917371b76 Mon Sep 17 00:00:00 2001 From: "Miss Islington (bot)" <31488909+miss-islington@users.noreply.github.com> Date: Thu, 23 Jan 2020 06:49:19 -0800 Subject: [PATCH] bpo-39421: Fix posible crash in heapq with custom comparison operators (GH-18118) (GH-18146) (cherry picked from commit 79f89e6e5a659846d1068e8b1bd8e491ccdef861) Co-authored-by: Pablo Galindo --- Lib/test/test_heapq.py | 31 ++++++++++++++++ .../2020-01-22-15-53-37.bpo-39421.O3nG7u.rst | 2 ++ Modules/_heapqmodule.c | 35 ++++++++++++++----- 3 files changed, 59 insertions(+), 9 deletions(-) create mode 100644 Misc/NEWS.d/next/Core and Builtins/2020-01-22-15-53-37.bpo-39421.O3nG7u.rst diff --git a/Lib/test/test_heapq.py b/Lib/test/test_heapq.py index 2f8c648d84a58..7c3fb0210f69b 100644 --- a/Lib/test/test_heapq.py +++ b/Lib/test/test_heapq.py @@ -414,6 +414,37 @@ def test_heappop_mutating_heap(self): with self.assertRaises((IndexError, RuntimeError)): self.module.heappop(heap) + def test_comparison_operator_modifiying_heap(self): + # See bpo-39421: Strong references need to be taken + # when comparing objects as they can alter the heap + class EvilClass(int): + def __lt__(self, o): + heap.clear() + return NotImplemented + + heap = [] + self.module.heappush(heap, EvilClass(0)) + self.assertRaises(IndexError, self.module.heappushpop, heap, 1) + + def test_comparison_operator_modifiying_heap_two_heaps(self): + + class h(int): + def __lt__(self, o): + list2.clear() + return NotImplemented + + class g(int): + def __lt__(self, o): + list1.clear() + return NotImplemented + + list1, list2 = [], [] + + self.module.heappush(list1, h(0)) + self.module.heappush(list2, g(0)) + + self.assertRaises((IndexError, RuntimeError), self.module.heappush, list1, g(1)) + self.assertRaises((IndexError, RuntimeError), self.module.heappush, list2, h(1)) class TestErrorHandlingPython(TestErrorHandling, TestCase): module = py_heapq diff --git a/Misc/NEWS.d/next/Core and Builtins/2020-01-22-15-53-37.bpo-39421.O3nG7u.rst b/Misc/NEWS.d/next/Core and Builtins/2020-01-22-15-53-37.bpo-39421.O3nG7u.rst new file mode 100644 index 0000000000000..bae008150ee12 --- /dev/null +++ b/Misc/NEWS.d/next/Core and Builtins/2020-01-22-15-53-37.bpo-39421.O3nG7u.rst @@ -0,0 +1,2 @@ +Fix possible crashes when operating with the functions in the :mod:`heapq` +module and custom comparison operators. diff --git a/Modules/_heapqmodule.c b/Modules/_heapqmodule.c index b499e1f668aae..0fb35ffe5ec48 100644 --- a/Modules/_heapqmodule.c +++ b/Modules/_heapqmodule.c @@ -29,7 +29,11 @@ siftdown(PyListObject *heap, Py_ssize_t startpos, Py_ssize_t pos) 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)) { @@ -71,10 +75,13 @@ siftup(PyListObject *heap, Py_ssize_t pos) /* 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 */ @@ -229,7 +236,10 @@ heappushpop(PyObject *self, PyObject *args) 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) { @@ -383,7 +393,11 @@ siftdown_max(PyListObject *heap, Py_ssize_t startpos, Py_ssize_t pos) 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)) { @@ -425,10 +439,13 @@ siftup_max(PyListObject *heap, Py_ssize_t pos) /* 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 */