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>

4 files changed, 60 insertions, 11 deletions

diff --git a/Lib/test/test_bytes.py b/Lib/test/test_bytes.py
index b457ff6..521e391 100644
--- a/Lib/test/test_bytes.py
+++ b/Lib/test/test_bytes.py
@@ -1710,6 +1710,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 3467c04..9aa7fca 100644
--- a/Modules/_testcapimodule.c
+++ b/Modules/_testcapimodule.c
@@ -5489,6 +5489,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.
  *
@@ -6272,6 +6287,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_CAST(meth_varargs_keywords), METH_VARARGS|METH_KEYWORDS},
     {"meth_o", meth_o, METH_O},
diff --git a/Objects/bytearrayobject.c b/Objects/bytearrayobject.c
index b9436d9..dfeed68 100644
--- a/Objects/bytearrayobject.c
+++ b/Objects/bytearrayobject.c
@@ -563,22 +563,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;
 }
@@ -593,11 +599,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");
@@ -612,9 +628,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;
         }