diff --git a/Lib/test/test_capi/test_bytes.py b/Lib/test/test_capi/test_bytes.py index a500f2c702db0f..a0006ea35e21fe 100644 --- a/Lib/test/test_capi/test_bytes.py +++ b/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): diff --git a/Modules/_testcapi/bytes.c b/Modules/_testcapi/bytes.c index b4468dff0d0ba0..79effcad40090e 100644 --- a/Modules/_testcapi/bytes.c +++ b/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 */ }; diff --git a/Objects/bytesobject.c b/Objects/bytesobject.c index cf6b66d4dcd572..117d8b56017b64 100644 --- a/Objects/bytesobject.c +++ b/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,6 +3922,10 @@ PyBytesWriter_FinishWithPointer(PyBytesWriter *writer, void *buf) void* PyBytesWriter_GetData(PyBytesWriter *writer) { +#ifdef Py_DEBUG + byteswriter_check_canary_byte(writer); +#endif + return byteswriter_data(writer); } @@ -3908,6 +3933,10 @@ PyBytesWriter_GetData(PyBytesWriter *writer) Py_ssize_t PyBytesWriter_GetSize(PyBytesWriter *writer) { +#ifdef Py_DEBUG + byteswriter_check_canary_byte(writer); +#endif + return _PyBytesWriter_GetSize(writer); } @@ -3915,6 +3944,10 @@ PyBytesWriter_GetSize(PyBytesWriter *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