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

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

If PyBytesWriter_Resize() fails, leave the writer unchanged.

Add a new internal _PyBytes_ResizeKeepOnError() function similar to
_PyBytes_Resize() but leaves the bytes object unchanged on error.

Author: vstinner

Include/internal/pycore_bytesobject.h

@@ -75,6 +75,8 @@ PyAPI_FUNC(PyObject *) _PyBytes_Repeat(PyObject *self, Py_ssize_t n);
 */
 #define _PyBytesObject_SIZE (offsetof(PyBytesObject, ob_sval) + 1)
 
+extern int _PyBytes_ResizeKeepOnError(PyObject **pv, Py_ssize_t newsize);
+
 /* --- PyBytesWriter ------------------------------------------------------ */
 
 struct PyBytesWriter {

Lib/test/test_bytes.py

@@ -1555,6 +1555,21 @@ 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 (memory allocation failure),
+        # the bytearray must be left unchanged.
+        _testcapi = import_helper.import_module('_testcapi')
+
+        data = b'some data'
+        ba = bytearray(data)
+        try:
+            with self.assertRaises(MemoryError):
+                _testcapi.set_nomemory(0)
+                ba.resize(1024)
+        finally:
+            _testcapi.remove_mem_hooks()
+        self.assertEqual(ba, data)
+
     def test_take_bytes(self):
         ba = bytearray(b'ab')
         self.assertEqual(ba.take_bytes(), b'ab')

Lib/test/test_capi/test_bytes.py

@@ -373,6 +373,21 @@ def test_resize(self):
         writer.resize(len(b'number=123456'), b'456')
         self.assertEqual(writer.finish(), self.result_type(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.assertEqual(writer.finish(), self.result_type(init + suffix))
+
     def test_format_i(self):
         # Test PyBytesWriter_Format()
         writer = self.create_writer()

Objects/bytearrayobject.c

@@ -288,18 +288,16 @@ bytearray_resize_lock_held(PyObject *self, Py_ssize_t requested_size)
                 Py_MIN(requested_size, Py_SIZE(self)));
     }
 
-    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;
+    if (_PyBytes_ResizeKeepOnError(&obj->ob_bytes_object, alloc) < 0) {
+        return -1;
     }
+
     bytearray_reinit_from_bytes(obj, size, alloc);
     if (alloc != size) {
         /* Add mid-buffer null; end provided by bytes. */
         obj->ob_bytes[size] = '\0';
     }
-
-    return ret;
+    return 0;
 }
 
 int
@@ -1632,10 +1630,7 @@ bytearray_take_bytes_impl(PyByteArrayObject *self, PyObject *n)
         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 (_PyBytes_ResizeKeepOnError(&self->ob_bytes_object, to_take) == -1) {
         Py_DECREF(remaining);
         return NULL;
     }

Objects/bytesobject.c

@@ -3344,68 +3344,99 @@ PyBytes_ConcatAndDel(PyObject **pv, PyObject *w)
    does *not* include that), and a trailing \0 byte is stored.
 */
 
-int
-_PyBytes_Resize(PyObject **pv, Py_ssize_t newsize)
+static int
+bytes_resize(PyObject **pv, Py_ssize_t newsize, int clear_obj_on_error)
 {
-    PyObject *v;
-    PyBytesObject *sv;
-    v = *pv;
+    PyObject *v = *pv;
+    PyObject *result;
+
     if (!PyBytes_Check(v) || newsize < 0) {
-        *pv = 0;
-        Py_DECREF(v);
         PyErr_BadInternalCall();
-        return -1;
+        goto error;
     }
+
     Py_ssize_t oldsize = PyBytes_GET_SIZE(v);
     if (oldsize == newsize) {
         /* return early if newsize equals to v->ob_size */
         return 0;
     }
+
     if (oldsize == 0) {
-        *pv = _PyBytes_FromSize(newsize, 0);
+        result = _PyBytes_FromSize(newsize, 0);
+        if (result == NULL) {
+            goto error;
+        }
+        *pv = result;
         Py_DECREF(v);
-        return (*pv == NULL) ? -1 : 0;
+        return 0;
     }
+
     if (newsize == 0) {
-        *pv = bytes_get_empty();
+        *pv = bytes_get_empty();  // cannot fail
         Py_DECREF(v);
         return 0;
     }
+
     if (!_PyObject_IsUniquelyReferenced(v)) {
-        if (oldsize < newsize) {
-            *pv = _PyBytes_FromSize(newsize, 0);
-            if (*pv) {
-                memcpy(PyBytes_AS_STRING(*pv), PyBytes_AS_STRING(v), oldsize);
-            }
-        }
-        else {
-            *pv = PyBytes_FromStringAndSize(PyBytes_AS_STRING(v), newsize);
+        result = _PyBytes_FromSize(newsize, 0);
+        if (!result) {
+            goto error;
         }
+
+        memcpy(PyBytes_AS_STRING(result), PyBytes_AS_STRING(v), Py_MIN(oldsize, newsize));
+        *pv = result;
         Py_DECREF(v);
-        return (*pv == NULL) ? -1 : 0;
+        return 0;
     }
-    assert(v != bytes_get_empty());
 
+    assert(v != bytes_get_empty());
 #ifdef Py_TRACE_REFS
     _Py_ForgetReference(v);
 #endif
     _PyReftracerTrack(v, PyRefTracer_DESTROY);
-    *pv = (PyObject *)
-        PyObject_Realloc(v, PyBytesObject_SIZE + newsize);
-    if (*pv == NULL) {
+    result = (PyObject *)PyObject_Realloc(v, PyBytesObject_SIZE + newsize);
+    if (result == NULL) {
+        if (clear_obj_on_error) {
+            *pv = NULL;
 #ifdef Py_REF_DEBUG
-        _Py_DecRefTotal(_PyThreadState_GET());
+            _Py_DecRefTotal(_PyThreadState_GET());
 #endif
-        PyObject_Free(v);
+            PyObject_Free(v);
+        }
         PyErr_NoMemory();
         return -1;
     }
-    _Py_NewReferenceNoTotal(*pv);
-    sv = (PyBytesObject *) *pv;
+
+    v = result;
+    _Py_NewReferenceNoTotal(v);
+    PyBytesObject *sv = (PyBytesObject *)v;
     Py_SET_SIZE(sv, newsize);
     sv->ob_sval[newsize] = '\0';
     set_ob_shash(sv, -1);          /* invalidate cached hash value */
+    *pv = v;
     return 0;
+
+error:
+    if (clear_obj_on_error) {
+        *pv = NULL;
+        Py_DECREF(v);
+    }
+    return -1;
+}
+
+
+int
+_PyBytes_Resize(PyObject **pv, Py_ssize_t newsize)
+{
+    return bytes_resize(pv, newsize, 1);
+}
+
+
+// Similar to _PyBytes_Resize(), but leaves the object unchanged on error.
+int
+_PyBytes_ResizeKeepOnError(PyObject **pv, Py_ssize_t newsize)
+{
+    return bytes_resize(pv, newsize, 0);
 }
 
 
@@ -3646,7 +3677,8 @@ byteswriter_resize(PyBytesWriter *writer, Py_ssize_t size, int resize)
             }
         }
         else {
-            if (_PyBytes_Resize(&writer->obj, size)) {
+            if (_PyBytes_ResizeKeepOnError(&writer->obj, size)) {
+                assert(writer->obj != NULL);
                 return -1;
             }
         }