From 2d79804bb00908d6563522c86f81128123752c5e Mon Sep 17 00:00:00 2001
From: "Miss Islington (bot)"
 <31488909+miss-islington@users.noreply.github.com>
Date: Tue, 19 Jul 2022 10:04:47 -0700
Subject: GH-91153: Handle mutating __index__ methods in bytearray item
 assignment (GH-94891)

(cherry picked from commit f36589510b8708fa224d799d5b328deab558aa4e)

Co-authored-by: Brandt Bucher <brandtbucher@microsoft.com>
---
 Lib/test/test_bytes.py                             | 17 ++++++++++
 .../2022-07-15-16-15-04.gh-issue-91153.HiBmtt.rst  |  2 ++
 Modules/_testcapimodule.c                          | 16 ++++++++++
 Objects/bytearrayobject.c                          | 36 +++++++++++++++-------
 4 files changed, 60 insertions(+), 11 deletions(-)
 create mode 100644 Misc/NEWS.d/next/Core and Builtins/2022-07-15-16-15-04.gh-issue-91153.HiBmtt.rst

diff --git a/Lib/test/test_bytes.py b/Lib/test/test_bytes.py
index 13ad238..72c5e18 100644
--- a/Lib/test/test_bytes.py
+++ b/Lib/test/test_bytes.py
@@ -1698,6 +1698,23 @@ class ByteArrayTest(BaseBytesTest, unittest.TestCase):
         self.assertEqual(b1, b)
         self.assertEqual(b3, b'xcxcxc')
 
+    def test_mutating_index(self):
+        class Boom:
+            def __index__(self):
+                b.clear()
+                return 0
+
+        with self.subTest("tp_as_mapping"):
+            b = bytearray(b'Now you see me...')
+            with self.assertRaises(IndexError):
+                b[0] = Boom()
+
+        with self.subTest("tp_as_sequence"):
+            _testcapi = import_helper.import_module('_testcapi')
+            b = bytearray(b'Now you see me...')
+            with self.assertRaises(IndexError):
+                _testcapi.sequence_setitem(b, 0, Boom())
+
 
 class AssortedBytesTest(unittest.TestCase):
     #
diff --git a/Misc/NEWS.d/next/Core and Builtins/2022-07-15-16-15-04.gh-issue-91153.HiBmtt.rst b/Misc/NEWS.d/next/Core and Builtins/2022-07-15-16-15-04.gh-issue-91153.HiBmtt.rst
new file mode 100644
index 0000000..2caa017
--- /dev/null
+++ b/Misc/NEWS.d/next/Core and Builtins/2022-07-15-16-15-04.gh-issue-91153.HiBmtt.rst	
@@ -0,0 +1,2 @@
+Fix an issue where a :class:`bytearray` item assignment could crash if it's
+resized by the new value's :meth:`__index__` method.
diff --git a/Modules/_testcapimodule.c b/Modules/_testcapimodule.c
index aaf29ad..f4826c2 100644
--- a/Modules/_testcapimodule.c
+++ b/Modules/_testcapimodule.c
@@ -5388,6 +5388,21 @@ sequence_getitem(PyObject *self, PyObject *args)
 }
 
 
+static PyObject *
+sequence_setitem(PyObject *self, PyObject *args)
+{
+    Py_ssize_t i;
+    PyObject *seq, *val;
+    if (!PyArg_ParseTuple(args, "OnO", &seq, &i, &val)) {
+        return NULL;
+    }
+    if (PySequence_SetItem(seq, i, val)) {
+        return NULL;
+    }
+    Py_RETURN_NONE;
+}
+
+
 /* Functions for testing C calling conventions (METH_*) are named meth_*,
  * e.g. "meth_varargs" for METH_VARARGS.
  *
@@ -5884,6 +5899,7 @@ static PyMethodDef TestMethods[] = {
 #endif
     {"write_unraisable_exc", test_write_unraisable_exc, METH_VARARGS},
     {"sequence_getitem", sequence_getitem, METH_VARARGS},
+    {"sequence_setitem", sequence_setitem, METH_VARARGS},
     {"meth_varargs", meth_varargs, METH_VARARGS},
     {"meth_varargs_keywords", (PyCFunction)(void(*)(void))meth_varargs_keywords, METH_VARARGS|METH_KEYWORDS},
     {"meth_o", meth_o, METH_O},
diff --git a/Objects/bytearrayobject.c b/Objects/bytearrayobject.c
index 1ab9621..2a675d5 100644
--- a/Objects/bytearrayobject.c
+++ b/Objects/bytearrayobject.c
@@ -577,22 +577,28 @@ bytearray_setslice(PyByteArrayObject *self, Py_ssize_t lo, Py_ssize_t hi,
 static int
 bytearray_setitem(PyByteArrayObject *self, Py_ssize_t i, PyObject *value)
 {
-    int ival;
+    int ival = -1;
 
-    if (i < 0)
+    // GH-91153: We need to do this *before* the size check, in case value has a
+    // nasty __index__ method that changes the size of the bytearray:
+    if (value && !_getbytevalue(value, &ival)) {
+        return -1;
+    }
+
+    if (i < 0) {
         i += Py_SIZE(self);
+    }
 
     if (i < 0 || i >= Py_SIZE(self)) {
         PyErr_SetString(PyExc_IndexError, "bytearray index out of range");
         return -1;
     }
 
-    if (value == NULL)
+    if (value == NULL) {
         return bytearray_setslice(self, i, i+1, NULL);
+    }
 
-    if (!_getbytevalue(value, &ival))
-        return -1;
-
+    assert(0 <= ival && ival < 256);
     PyByteArray_AS_STRING(self)[i] = ival;
     return 0;
 }
@@ -607,11 +613,21 @@ bytearray_ass_subscript(PyByteArrayObject *self, PyObject *index, PyObject *valu
     if (_PyIndex_Check(index)) {
         Py_ssize_t i = PyNumber_AsSsize_t(index, PyExc_IndexError);
 
-        if (i == -1 && PyErr_Occurred())
+        if (i == -1 && PyErr_Occurred()) {
             return -1;
+        }
 
-        if (i < 0)
+        int ival = -1;
+
+        // GH-91153: We need to do this *before* the size check, in case values
+        // has a nasty __index__ method that changes the size of the bytearray:
+        if (values && !_getbytevalue(values, &ival)) {
+            return -1;
+        }
+
+        if (i < 0) {
             i += PyByteArray_GET_SIZE(self);
+        }
 
         if (i < 0 || i >= Py_SIZE(self)) {
             PyErr_SetString(PyExc_IndexError, "bytearray index out of range");
@@ -626,9 +642,7 @@ bytearray_ass_subscript(PyByteArrayObject *self, PyObject *index, PyObject *valu
             slicelen = 1;
         }
         else {
-            int ival;
-            if (!_getbytevalue(values, &ival))
-                return -1;
+            assert(0 <= ival && ival < 256);
             buf[i] = (char)ival;
             return 0;
         }
-- 
cgit v0.12