GH-91153: Handle mutating __index__ methods in bytearray item assignment (GH-94891)

1 files changed, 25 insertions, 11 deletions

diff --git a/Objects/bytearrayobject.c b/Objects/bytearrayobject.c
index d017962..b2962fd 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;
         }