[3.15] gh-157242: Leave bytearray unchanged if resize() fails (#157340) (#159274)

gh-157242: Leave bytearray unchanged if resize() fails (#157340)

If bytearray.resize() or bytearray.take_bytes() fails, leave the
bytearray unchanged.

Add PyBytesWriter_Resize() error test on bytearray.

(cherry picked from commit 67e6be7)

Author: vstinner

Lib/test/test_bytes.py

@@ -6,6 +6,7 @@
 
 import array
 import codecs
+import contextlib
 import copy
 import functools
 import operator
@@ -49,6 +50,19 @@ def __index__(self):
         return self.value
 
 
+@contextlib.contextmanager
+def inject_memory_error(testcase, start):
+    # Raise SkipTest if _testcapi extension module is missing
+    _testcapi = import_helper.import_module('_testcapi')
+
+    with testcase.assertRaises(MemoryError):
+        try:
+            _testcapi.set_nomemory(start)
+            yield
+        finally:
+            _testcapi.remove_mem_hooks()
+
+
 class BaseBytesTest:
 
     def assertTypedEqual(self, actual, expected):
@@ -1565,6 +1579,35 @@ def test_resize(self):
         self.assertRaises(MemoryError, bytearray().resize, sys.maxsize)
         self.assertRaises(MemoryError, bytearray(1000).resize, sys.maxsize)
 
+    def test_resize_error(self):
+        # gh-157242: If bytearray.resize() fails (MemoryError),
+        # the bytearray must be left unchanged.
+
+        offset = 3
+        for logical_offset in (False, True):
+            with self.subTest(logical_offset=logical_offset):
+                # grow bytearray
+                ba = bytearray(b'0123456789')
+                if logical_offset:
+                    expected = ba[offset:]
+                    del ba[:offset]
+                else:
+                    expected = ba.copy()
+                with inject_memory_error(self, 0):
+                    ba.resize(1024)
+                self.assertEqual(ba, expected)
+
+                # shrink bytearray
+                ba = bytearray(b'0123456789')
+                if logical_offset:
+                    expected = ba[offset:]
+                    del ba[:offset]
+                else:
+                    expected = ba.copy()
+                with inject_memory_error(self, 0):
+                    ba.resize(1)
+                self.assertEqual(ba, expected)
+
     def test_take_bytes(self):
         ba = bytearray(b'ab')
         self.assertEqual(ba.take_bytes(), b'ab')
@@ -1629,6 +1672,28 @@ def test_take_bytes(self):
         self.assertEqual(ba, bytearray(b'A'))
         self.assertEqual(ord(b'c'), ord('c'))
 
+    def test_take_bytes_error(self):
+        # gh-157242: If bytearray.take_bytes() fails (MemoryError),
+        # the bytearray must be left unchanged.
+
+        for logical_offset, to_take, mem_errors in (
+            (True, 5, (0, 1)),
+            (False, 5, (0, 1)),
+            (True, None, (0,)),
+        ):
+            for mem_error in mem_errors:
+                with self.subTest(logical_offset=logical_offset,
+                                  to_take=to_take, mem_error=mem_error):
+                    ba = bytearray(b'0123456789')
+                    if logical_offset:
+                        expected = ba[3:]
+                        del ba[:3]
+                    else:
+                        expected = ba.copy()
+                    with inject_memory_error(self, mem_error):
+                        ba.take_bytes(to_take)
+                    self.assertEqual(ba, expected)
+
     @support.cpython_only  # tests an implementation detail
     def test_take_bytes_optimization(self):
         # Validate optimization around taking lots of little chunks out of a

Lib/test/test_capi/test_bytes.py

@@ -396,6 +396,24 @@ def test_resize(self):
         writer.resize(len(b'number=123456'), b'456')
         self.bytes_equal(writer.finish(), b'number=123456')
 
+    def test_resize_error(self):
+        small_buffer = _testcapi.PyBytesWriter_small_buffer
+        init = b'x' * (small_buffer * 2)
+        writer = self.create_writer(len(init), init)
+        size = len(init) + 100
+        try:
+            with self.assertRaises(MemoryError):
+                _testcapi.set_nomemory(0)
+                writer.resize(size, b'')
+        finally:
+            _testcapi.remove_mem_hooks()
+        suffix = b'still working'
+        writer.write_bytes(suffix, -1)
+        self.bytes_equal(writer.finish(), init + suffix)
+
+        # Note: PyBytesWriter_Resize() leaves the buffer unchanged (no resize)
+        # if the new size is smaller than the allocated size
+
     def test_format_i(self):
         # Test PyBytesWriter_Format()
         writer = self.create_writer()
@@ -465,24 +483,6 @@ def test_example_resize(self):
     def test_example_highlevel(self):
         self.assertEqual(_testcapi.byteswriter_highlevel(), b'Hello World!')
 
-    def test_resize_error(self):
-        small_buffer = _testcapi.PyBytesWriter_small_buffer
-        init = b'x' * (small_buffer * 2)
-        writer = self.create_writer(len(init), init)
-        size = len(init) + 100
-        try:
-            with self.assertRaises(MemoryError):
-                _testcapi.set_nomemory(0)
-                writer.resize(size, b'')
-        finally:
-            _testcapi.remove_mem_hooks()
-        suffix = b'still working'
-        writer.write_bytes(suffix, -1)
-        self.bytes_equal(writer.finish(), init + suffix)
-
-        # Note: PyBytesWriter_Resize() leaves the buffer unchanged (no resize)
-        # if the new size is smaller than the allocated size
-
 
 class ByteArrayWriterTest(BaseWriterTest, unittest.TestCase):
     RESULT_TYPE = bytearray

Misc/NEWS.d/next/Core_and_Builtins/2026-09-10-02-50-14.gh-issue-157242.LsqOUJ.rst

@@ -0,0 +1,3 @@
+If :meth:`bytearray.resize` or :meth:`bytearray.take_bytes` fails, leave the
+:class:`bytearray` unchanged, instead of clearing it. Patch by Victor
+Stinner.

Objects/bytearrayobject.c

@@ -43,12 +43,24 @@ _getbytevalue(PyObject* arg, int *value)
     return 1;
 }
 
+static inline void
+bytearray_write_trailing_null_byte(PyByteArrayObject *self)
+{
+    char *data = PyByteArray_AS_STRING(self);
+    Py_ssize_t size = PyByteArray_GET_SIZE(self);
+    data[size] = '\0';
+}
+
+
 static void
-bytearray_reinit_from_bytes(PyByteArrayObject *self, Py_ssize_t size,
-                            Py_ssize_t alloc)
+bytearray_reinit_from_bytes(PyByteArrayObject *self, Py_ssize_t size)
 {
+    Py_ssize_t alloc = PyBytes_GET_SIZE(self->ob_bytes_object);
+    assert(0 <= size && size <= alloc);
+
     /* Only the empty bytes may be immortal. */
     assert((alloc == 0) == _Py_IsImmortal(self->ob_bytes_object));
+
     self->ob_bytes = self->ob_start = PyBytes_AS_STRING(self->ob_bytes_object);
     Py_SET_SIZE(self, size);
     FT_ATOMIC_STORE_SSIZE_RELAXED(self->ob_alloc, alloc);
@@ -185,7 +197,7 @@ PyByteArray_FromStringAndSize(const char *bytes, Py_ssize_t size)
         Py_DECREF(new);
         return NULL;
     }
-    bytearray_reinit_from_bytes(new, size, size);
+    bytearray_reinit_from_bytes(new, size);
     if (bytes != NULL && size > 0) {
         memcpy(new->ob_bytes, bytes, size);
     }
@@ -211,6 +223,43 @@ PyByteArray_AsString(PyObject *self)
     return PyByteArray_AS_STRING(self);
 }
 
+
+static int
+bytearray_resize_storage(PyByteArrayObject *self,
+                         Py_ssize_t new_size, Py_ssize_t alloc)
+{
+    _Py_CRITICAL_SECTION_ASSERT_OBJECT_LOCKED(self);
+    assert(1 <= new_size && new_size <= alloc);
+
+    Py_ssize_t size = Py_SIZE(self);
+
+    /* Re-align data to the start of the allocation. */
+    char *old_start = self->ob_start;
+    if (self->ob_start != self->ob_bytes) {
+        /* optimization tradeoff: This is faster than a new allocation when
+           the number of bytes being removed in a resize is small; for
+           large size changes it may be better to just make a new bytes
+           object as _PyBytes_Resize will do a malloc + memcpy internally.
+           */
+        Py_ssize_t move = Py_MIN(new_size, size);
+        memmove(self->ob_bytes, self->ob_start, move);
+        self->ob_start = self->ob_bytes;
+    }
+
+    if (_PyBytes_ResizeKeepOnError(&self->ob_bytes_object, alloc) < 0) {
+        if (old_start != self->ob_bytes && new_size < size) {
+            // Move remaining bytes
+            Py_ssize_t moved = new_size;
+            Py_ssize_t remaining = size - moved;
+            memmove(self->ob_bytes + moved, old_start + moved, remaining);
+        }
+        bytearray_write_trailing_null_byte(self);
+        return -1;
+    }
+    return 0;
+}
+
+
 static int
 bytearray_resize_lock_held(PyObject *self, Py_ssize_t requested_size)
 {
@@ -246,7 +295,7 @@ bytearray_resize_lock_held(PyObject *self, Py_ssize_t requested_size)
     if (requested_size == 0) {
         Py_SETREF(obj->ob_bytes_object,
                    Py_GetConstant(Py_CONSTANT_EMPTY_BYTES));
-        bytearray_reinit_from_bytes(obj, 0, 0);
+        bytearray_reinit_from_bytes(obj, 0);
         return 0;
     }
 
@@ -261,7 +310,7 @@ bytearray_resize_lock_held(PyObject *self, Py_ssize_t requested_size)
             /* Minor downsize; quick exit */
             Py_SET_SIZE(self, size);
             /* Add mid-buffer null; end provided by bytes. */
-            PyByteArray_AS_STRING(self)[size] = '\0'; /* Trailing null */
+            bytearray_write_trailing_null_byte(_PyByteArray_CAST(self));
             return 0;
         }
     }
@@ -281,28 +330,16 @@ bytearray_resize_lock_held(PyObject *self, Py_ssize_t requested_size)
         return -1;
     }
 
-    /* Re-align data to the start of the allocation. */
-    if (logical_offset > 0) {
-        /* optimization tradeoff: This is faster than a new allocation when
-           the number of bytes being removed in a resize is small; for large
-           size changes it may be better to just make a new bytes object as
-           _PyBytes_Resize will do a malloc + memcpy internally. */
-        memmove(obj->ob_bytes, obj->ob_start,
-                Py_MIN(requested_size, Py_SIZE(self)));
+    if (bytearray_resize_storage(obj, requested_size, (Py_ssize_t)alloc) < 0) {
+        return -1;
     }
 
-    int ret = _PyBytes_Resize(&obj->ob_bytes_object, alloc);
-    if (ret == -1) {
-        obj->ob_bytes_object = Py_GetConstant(Py_CONSTANT_EMPTY_BYTES);
-        size = alloc = 0;
-    }
-    bytearray_reinit_from_bytes(obj, size, alloc);
+    bytearray_reinit_from_bytes(obj, size);
     if (alloc != size) {
         /* Add mid-buffer null; end provided by bytes. */
-        obj->ob_bytes[size] = '\0';
+        bytearray_write_trailing_null_byte(obj);
     }
-
-    return ret;
+    return 0;
 }
 
 int
@@ -928,7 +965,7 @@ bytearray_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
     }
     PyByteArrayObject *self = _PyByteArray_CAST(op);
     self->ob_bytes_object = Py_GetConstant(Py_CONSTANT_EMPTY_BYTES);
-    bytearray_reinit_from_bytes(self, 0, 0);
+    bytearray_reinit_from_bytes(self, 0);
     self->ob_exports = 0;
     return op;
 }
@@ -995,9 +1032,9 @@ bytearray___init___impl(PyByteArrayObject *self, PyObject *arg,
             && PyBytes_CheckExact(encoded)
             && _PyBytes_GET_CACHED_HASH((PyBytesObject*)encoded) == -1)
         {
-            Py_ssize_t size = Py_SIZE(encoded);
+            Py_ssize_t size = PyBytes_GET_SIZE(encoded);
             self->ob_bytes_object = encoded;
-            bytearray_reinit_from_bytes(self, size, size);
+            bytearray_reinit_from_bytes(self, size);
             return 0;
         }
         new = bytearray_iconcat((PyObject*)self, encoded);
@@ -1121,7 +1158,7 @@ bytearray___init___impl(PyByteArrayObject *self, PyObject *arg,
         /* Append the byte */
         if (Py_SIZE(self) + 1 < self->ob_alloc) {
             Py_SET_SIZE(self, Py_SIZE(self) + 1);
-            PyByteArray_AS_STRING(self)[Py_SIZE(self)] = '\0';
+            bytearray_write_trailing_null_byte(self);
         }
         else if (PyByteArray_Resize((PyObject *)self, Py_SIZE(self)+1) < 0)
             goto error;
@@ -1611,6 +1648,7 @@ bytearray_take_bytes_impl(PyByteArrayObject *self, PyObject *n)
     }
 
     Py_ssize_t remaining_length = size - to_take;
+
     // optimization: If taking less than leaving, just copy the small to_take
     // portion out and move ob_start.
     if (to_take < remaining_length) {
@@ -1632,24 +1670,15 @@ bytearray_take_bytes_impl(PyByteArrayObject *self, PyObject *n)
     memcpy(PyBytes_AS_STRING(remaining), self->ob_start + to_take,
            remaining_length);
 
-    // If the bytes are offset inside the buffer must first align.
-    if (self->ob_start != self->ob_bytes) {
-        memmove(self->ob_bytes, self->ob_start, to_take);
-        self->ob_start = self->ob_bytes;
-    }
-
-    if (_PyBytes_Resize(&self->ob_bytes_object, to_take) == -1) {
-        assert(self->ob_bytes_object == NULL);
-        self->ob_bytes_object = Py_GetConstant(Py_CONSTANT_EMPTY_BYTES);
-        bytearray_reinit_from_bytes(self, 0, 0);
+    if (bytearray_resize_storage(self, to_take, to_take) < 0) {
         Py_DECREF(remaining);
         return NULL;
     }
 
     // Point the bytearray towards the buffer with the remaining data.
     PyObject *result = self->ob_bytes_object;
     self->ob_bytes_object = remaining;
-    bytearray_reinit_from_bytes(self, remaining_length, remaining_length);
+    bytearray_reinit_from_bytes(self, remaining_length);
     return result;
 }