gh-157242: Fix PyBytesWriter_Resize() on MemoryError (#157243)

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.

Co-authored-by: Maurycy Pawłowski-Wieroński <maurycy@maurycy.com>

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_capi/test_bytes.py

@@ -446,6 +446,24 @@ 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.assertEqual(writer.finish(), self.result_type(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/C_API/2026-09-10-02-51-26.gh-issue-157242.tfXIsf.rst

@@ -0,0 +1,3 @@
+Fix :c:func:`PyBytesWriter_Resize` to handle properly memory allocation
+failure. Leave the writer unchanged on error, instead of leaving it in an
+inconsistent state. Patch by Victor Stinner.

Modules/_testcapi/mem.c

@@ -177,18 +177,29 @@ fm_remove_hooks(void)
     }
 }
 
+static void
+fm_set_nomemory(int start, int stop)
+{
+    /* Memory allocation fails after 'start' allocation requests, and until
+     * 'stop' allocation requests except when 'stop' is negative or equal
+     * to 0 (default) in which case allocation failures never stop. */
+    FmData.start = start;
+    FmData.stop = stop;
+    FmData.count = 0;
+    fm_setup_hooks();
+}
+
 static PyObject *
 set_nomemory(PyObject *self, PyObject *args)
 {
     /* Memory allocation fails after 'start' allocation requests, and until
      * 'stop' allocation requests except when 'stop' is negative or equal
      * to 0 (default) in which case allocation failures never stop. */
-    FmData.count = 0;
-    FmData.stop = 0;
-    if (!PyArg_ParseTuple(args, "i|i", &FmData.start, &FmData.stop)) {
+    int start, stop = 0;
+    if (!PyArg_ParseTuple(args, "i|i", &start, &stop)) {
         return NULL;
     }
-    fm_setup_hooks();
+    fm_set_nomemory(start, stop);
     Py_RETURN_NONE;
 }
 
@@ -821,6 +832,129 @@ get_process_memory_usage(PyObject *self, PyObject *args)
 #endif
 
 
+struct bytes_resize_tracer {
+    PyObject *create;
+    PyObject *destroy;
+};
+
+
+static int
+bytes_resize_tracer(PyObject *obj, PyRefTracerEvent event, void* data)
+{
+    if (event != PyRefTracer_CREATE && event != PyRefTracer_DESTROY) {
+        return 0;
+    }
+
+    struct bytes_resize_tracer *tracer = (struct bytes_resize_tracer*)data;
+    if (!PyBytes_Check(obj)) {
+        return 0;
+    }
+
+    switch (event) {
+        case PyRefTracer_CREATE:
+            tracer->create = obj;
+            break;
+        case PyRefTracer_DESTROY:
+            tracer->destroy = obj;
+            break;
+        default:
+            break;
+    }
+    return 0;
+}
+
+
+// When _PyBytes_Resize() resizes a bytes object in-place, check that
+// PyRefTracer_DESTROY and PyRefTracer_CREATE events are emitted.
+// If no_memory is non-zero, inject MemoryError.
+static int
+check_bytes_resize_tracer(int no_memory)
+{
+    PyObject *bytes = NULL;
+    PyRefTracer old_tracer = NULL;
+    void *old_tracer_data = NULL;
+    int restore_tracer = 0;
+
+    bytes = PyBytes_FromString("hello");
+    if (bytes == NULL) {
+        goto error;
+    }
+    assert(PyUnstable_Object_IsUniquelyReferenced(bytes));
+
+    old_tracer = PyRefTracer_GetTracer(&old_tracer_data);
+    restore_tracer = 1;
+
+    struct bytes_resize_tracer tracer = {0};
+    if (PyRefTracer_SetTracer(bytes_resize_tracer, &tracer) != 0) {
+        goto error;
+    }
+
+    PyObject *old_bytes = bytes;  // borrowed reference
+    if (no_memory) {
+        fm_set_nomemory(0, 0);
+        int res = _PyBytes_Resize(&bytes, 100);
+        assert(res < 0);
+        assert(bytes == NULL);
+        fm_remove_hooks();
+
+        assert(PyErr_ExceptionMatches(PyExc_MemoryError));
+        PyErr_Clear();
+    }
+    else {
+        if (_PyBytes_Resize(&bytes, 100) < 0) {
+            assert(bytes == NULL);
+            goto error;
+        }
+    }
+
+    if (tracer.destroy != old_bytes) {
+        PyErr_SetString(PyExc_AssertionError, "PyRefTracer_DESTROY not seen");
+        goto error;
+    }
+
+    int seen_create;
+    if (no_memory) {
+        seen_create = (tracer.create == old_bytes);
+    }
+    else {
+        seen_create = (tracer.create == bytes);
+    }
+    if (!seen_create) {
+        PyErr_SetString(PyExc_AssertionError, "PyRefTracer_CREATE not seen");
+        goto error;
+    }
+
+    Py_CLEAR(bytes);
+    if (PyRefTracer_SetTracer(old_tracer, old_tracer_data) != 0) {
+        restore_tracer = 0;
+        goto error;
+    }
+    return 0;
+
+error:
+    Py_XDECREF(bytes);
+    if (restore_tracer) {
+        if (PyRefTracer_SetTracer(old_tracer, old_tracer_data) != 0) {
+            return -1;
+        }
+    }
+    return -1;
+}
+
+
+static PyObject*
+test_bytes_resize_tracer(PyObject *self, PyObject *Py_UNUSED(ignored))
+{
+    if (check_bytes_resize_tracer(0) < 0) {
+        return NULL;
+    }
+    if (check_bytes_resize_tracer(1) < 0) {
+        return NULL;
+    }
+    Py_RETURN_NONE;
+}
+
+
 static PyMethodDef test_methods[] = {
     {"pymem_api_misuse",              pymem_api_misuse,              METH_NOARGS},
     {"pymem_buffer_overflow",         pymem_buffer_overflow,         METH_NOARGS},
@@ -838,6 +972,7 @@ static PyMethodDef test_methods[] = {
 #if TARGET_OS_OSX || defined(__FreeBSD__)
     {"get_process_memory_usage",      get_process_memory_usage,      METH_VARARGS},
 #endif
+    {"test_bytes_resize_tracer",      test_bytes_resize_tracer,      METH_NOARGS},
 
     // Tracemalloc tests
     {"tracemalloc_track",             tracemalloc_track,             METH_VARARGS},

Objects/bytesobject.c

@@ -3344,69 +3344,96 @@ PyBytes_ConcatAndDel(PyObject **pv, PyObject *w)
    does *not* include that), and a trailing \0 byte is stored.
 */
 
+// Similar to _PyBytes_Resize(), but leaves the object unchanged on error.
 int
-_PyBytes_Resize(PyObject **pv, Py_ssize_t newsize)
+_PyBytes_ResizeKeepOnError(PyObject **pv, Py_ssize_t newsize)
 {
-    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;
     }
+
     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) {
+            return -1;
+        }
+        *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)) {
         // Allocate and then copy so we don't get a shared immortal
         // one-character singleton!
-        *pv = _PyBytes_FromSize(newsize, 0);
-        if (*pv) {
-            memcpy(PyBytes_AS_STRING(*pv), PyBytes_AS_STRING(v),
-                   Py_MIN(oldsize, newsize));
+        result = _PyBytes_FromSize(newsize, 0);
+        if (!result) {
+            return -1;
         }
+
+        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());
 
 #ifdef Py_TRACE_REFS
     _Py_ForgetReference(v);
 #endif
     _PyReftracerTrack(v, PyRefTracer_DESTROY);
-    *pv = (PyObject *)
-        PyObject_Realloc(v, PyBytesObject_SIZE + newsize);
-    if (*pv == NULL) {
-#ifdef Py_REF_DEBUG
-        _Py_DecRefTotal(_PyThreadState_GET());
+
+    result = (PyObject *)PyObject_Realloc(v, PyBytesObject_SIZE + newsize);
+    if (result == NULL) {
+#ifdef Py_TRACE_REFS
+        _Py_AddToAllObjects(v);
 #endif
-        PyObject_Free(v);
+        _PyReftracerTrack(v, PyRefTracer_CREATE);
+
         PyErr_NoMemory();
         return -1;
     }
-    _Py_NewReferenceNoTotal(*pv);
-    sv = (PyBytesObject *) *pv;
+
+    *pv = result;
+    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 */
     return 0;
 }
 
 
+int
+_PyBytes_Resize(PyObject **pv, Py_ssize_t newsize)
+{
+    int res = _PyBytes_ResizeKeepOnError(pv, newsize);
+    if (res < 0) {
+        PyObject *v = *pv;
+        *pv = NULL;
+        Py_DECREF(v);
+    }
+    return res;
+}
+
+
 /*********************** Bytes Iterator ****************************/
 
 typedef struct {
@@ -3644,7 +3671,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;
             }
         }