gh-156939: Detect PyBytesWriter buffer overflow earlier

Lib/test/test_capi/test_bytes.py

@@ -591,24 +591,42 @@ def test_canary_byte(self):
 
         # Test small buffer and large buffer
         for size in (0, self.SMALL_BUFFER, self.LARGE_BUFFER):
-            with self.subTest(size=size):
-                code = textwrap.dedent(f"""
-                    from test.support import SuppressCrashReport
-                    import _testcapi
-                    size = {size}
-                    # Add an extra '#' byte to trigger a buffer overflow
-                    data = b'x' * size + b'#'
-                    use_bytearray = {use_bytearray}
-                    writer = _testcapi.PyBytesWriter(size, use_bytearray)
-                    with SuppressCrashReport():
-                        writer.write(0, data, check=False)
-                        writer.finish()
-                """)
-                proc = assert_python_failure('-c', code)
-                self.assertIn(b'Buffer overflow detected in PyBytesWriter',
-                              proc.err)
-                self.assertIn(f'at position {size}'.encode(),
-                              proc.err)
+            for operation in (
+                'writer.get_data()',
+                'writer.get_size()',
+                f'writer.resize({size} * 2)',
+                f'writer.grow({size})',
+                'writer.discard()',
+                'writer.finish()',
+            ):
+                with self.subTest(size=size, operation=operation):
+                    code = textwrap.dedent(f"""
+                        from test.support import SuppressCrashReport
+                        import os
+                        import _testcapi
+                        size = {size}
+                        # Add an extra '#' byte to trigger a buffer overflow
+                        data = b'x' * size + b'#'
+                        use_bytearray = {use_bytearray}
+                        writer = _testcapi.PyBytesWriter(size, use_bytearray)
+                        with SuppressCrashReport():
+                            writer.write(0, data, check=False)
+                            try:
+                                {operation}
+                            except:
+                                # Ignore all exceptions
+                                pass
+                            # If we reached this line, the operation didn't
+                            # detect the overflow. Exit immediatetly without
+                            # calling the writer destructor since it can detect
+                            # the overflow.
+                            os._exit(0)
+                    """)
+                    proc = assert_python_failure('-c', code)
+                    self.assertIn(b'Buffer overflow detected in PyBytesWriter',
+                                  proc.err)
+                    self.assertIn(f'at position {size}'.encode(),
+                                  proc.err)
 
     @unittest.skipUnless(support.Py_DEBUG, 'need debug build')
     def test_get_data_canary(self):

Modules/_testcapi/bytes.c

@@ -315,6 +315,20 @@ writer_finish_with_size(PyObject *self_raw, PyObject *args)
 }
 
 
+static PyObject*
+writer_discard(PyObject *self_raw, PyObject *Py_UNUSED(args))
+{
+    WriterObject *self = (WriterObject *)self_raw;
+    if (writer_check(self) < 0) {
+        return NULL;
+    }
+
+    PyBytesWriter_Discard(self->writer);
+    self->writer = NULL;
+    Py_RETURN_NONE;
+}
+
+
 static PyMethodDef writer_methods[] = {
     {"write", _PyCFunction_CAST(writer_write), METH_VARARGS | METH_KEYWORDS},
     {"write_bytes", _PyCFunction_CAST(writer_write_bytes), METH_VARARGS},
@@ -325,6 +339,7 @@ static PyMethodDef writer_methods[] = {
     {"get_size", _PyCFunction_CAST(writer_get_size), METH_NOARGS},
     {"finish", _PyCFunction_CAST(writer_finish), METH_NOARGS},
     {"finish_with_size", _PyCFunction_CAST(writer_finish_with_size), METH_VARARGS},
+    {"discard", _PyCFunction_CAST(writer_discard), METH_VARARGS},
     {NULL,              NULL}           /* sentinel */
 };
 

Objects/bytesobject.c

@@ -3696,6 +3696,10 @@ byteswriter_resize(PyBytesWriter *writer, Py_ssize_t size, int resize)
     if (writer->obj != NULL) {
         if (writer->use_bytearray) {
             if (PyByteArray_Resize(writer->obj, size)) {
+#ifdef Py_DEBUG
+                // bytearray can override the canary byte on error
+                byteswriter_write_canary_byte(writer);
+#endif
                 return -1;
             }
         }
@@ -3770,6 +3774,11 @@ byteswriter_create(Py_ssize_t size, int use_bytearray)
 
     if (size >= 1) {
         if (byteswriter_resize(writer, size, 0) < 0) {
+#ifdef Py_DEBUG
+            // Write the canary byte so byteswriter_check_canary_byte()
+            // doesn't fail in PyBytesWriter_Discard()
+            byteswriter_write_canary_byte(writer);
+#endif
             PyBytesWriter_Discard(writer);
             return NULL;
         }
@@ -3803,6 +3812,10 @@ PyBytesWriter_Discard(PyBytesWriter *writer)
         return;
     }
 
+#ifdef Py_DEBUG
+    byteswriter_check_canary_byte(writer);
+#endif
+
     Py_XDECREF(writer->obj);
     _Py_FREELIST_FREE(bytes_writers, writer, PyMem_Free);
 }
@@ -3875,6 +3888,14 @@ PyBytesWriter_FinishWithSize(PyBytesWriter *writer, Py_ssize_t size)
         // The function returns single byte singleton if size equals 1
         result = PyBytes_FromStringAndSize(writer->small_buffer, size);
     }
+
+#ifdef Py_DEBUG
+    // Reset the writer, so byteswriter_check_canary_byte() doesn't fail
+    // in PyBytesWriter_Discard().
+    writer->size = 0;
+    byteswriter_write_canary_byte(writer);
+#endif
+
     PyBytesWriter_Discard(writer);
     return result;
 
@@ -3901,20 +3922,32 @@ PyBytesWriter_FinishWithPointer(PyBytesWriter *writer, void *buf)
 void*
 PyBytesWriter_GetData(PyBytesWriter *writer)
 {
+#ifdef Py_DEBUG
+    byteswriter_check_canary_byte(writer);
+#endif
+
     return byteswriter_data(writer);
 }
 
 
 Py_ssize_t
 PyBytesWriter_GetSize(PyBytesWriter *writer)
 {
+#ifdef Py_DEBUG
+    byteswriter_check_canary_byte(writer);
+#endif
+
     return _PyBytesWriter_GetSize(writer);
 }
 
 
 int
 PyBytesWriter_Resize(PyBytesWriter *writer, Py_ssize_t new_size)
 {
+#ifdef Py_DEBUG
+    byteswriter_check_canary_byte(writer);
+#endif
+
     if (new_size < 0) {
         PyErr_SetString(PyExc_ValueError, "size must be >= 0");
         return -1;
@@ -3950,6 +3983,10 @@ _PyBytesWriter_ResizeAndUpdatePointer(PyBytesWriter *writer, Py_ssize_t size,
 int
 PyBytesWriter_Grow(PyBytesWriter *writer, Py_ssize_t grow)
 {
+#ifdef Py_DEBUG
+    byteswriter_check_canary_byte(writer);
+#endif
+
     if (grow == 0) {
         // Nothing to do
         return 0;
@@ -4042,6 +4079,10 @@ PyBytesWriter_Format(PyBytesWriter *writer, const char *format, ...)
 static Py_ssize_t
 _PyBytesWriter_ResizeToAllocated(PyBytesWriter *writer)
 {
+#ifdef Py_DEBUG
+    byteswriter_check_canary_byte(writer);
+#endif
+
     Py_ssize_t allocated = byteswriter_allocated(writer);
     writer->size = allocated;
 #ifdef Py_DEBUG