gh-129813: Check size in PyBytesWriter_FinishWithSize()

Reject size larger than the allocated size.

Also, check negative size in PyBytesWriter_FinishWithSize() to always
raise ValueError. Previously, the function raised SystemError or
ValueError depending on the code path.

Replace _PyBytesWriter_GetAllocated() optimization with
_PyBytesWriter_ResizeToAllocated() to update the writer size to its
allocated size.

Author: vstinner

Lib/test/test_capi/test_bytes.py

@@ -317,14 +317,20 @@ def test_create(self):
 
     def test_finish_with_size(self):
         # Test PyBytesWriter_FinishWithSize()
-        writer = self.create_writer(10, b'abc')
+        writer = self.create_writer(10, b'abcdef')
         self.assertEqual(writer.get_size(), 10)
         self.assertEqual(writer.finish_with_size(3), self.result_type(b'abc'))
 
+        # Error if the size is negative
         writer = self.create_writer(3, b'abc')
-        with self.assertRaises(SystemError):
+        with self.assertRaises(ValueError):
             writer.finish_with_size(-3)
 
+        # Error if the requested size is larger than the allocated size
+        writer = self.create_writer(3, b'abc')
+        with self.assertRaises(ValueError):
+            writer.finish_with_size(4)
+
     def test_write_bytes(self):
          # Test PyBytesWriter_WriteBytes()
          writer = self.create_writer()

Objects/bytesobject.c

@@ -31,7 +31,7 @@ class bytes "PyBytesObject *" "&PyBytes_Type"
 /* Forward declaration */
 static void* _PyBytesWriter_ResizeAndUpdatePointer(PyBytesWriter *writer,
                                                    Py_ssize_t size, void *data);
-static Py_ssize_t _PyBytesWriter_GetAllocated(PyBytesWriter *writer);
+static Py_ssize_t _PyBytesWriter_ResizeToAllocated(PyBytesWriter *writer);
 
 
 #define CHARACTERS _Py_SINGLETON(bytes_characters)
@@ -2993,8 +2993,8 @@ _PyBytes_FromList(PyObject *x)
     if (writer == NULL) {
         return NULL;
     }
+    size = _PyBytesWriter_ResizeToAllocated(writer);
     char *str = PyBytesWriter_GetData(writer);
-    size = _PyBytesWriter_GetAllocated(writer);
 
     for (Py_ssize_t i = 0; i < PyList_GET_SIZE(x); i++) {
         PyObject *item = _PyList_GetItemRef((PyListObject *)x, i);
@@ -3017,7 +3017,9 @@ _PyBytes_FromList(PyObject *x)
             if (str == NULL) {
                 goto error;
             }
-            size = _PyBytesWriter_GetAllocated(writer);
+
+            // Set the writer size to its allocated size
+            size = _PyBytesWriter_ResizeToAllocated(writer);
         }
         *str++ = (char) value;
     }
@@ -3075,8 +3077,8 @@ _PyBytes_FromIterator(PyObject *it, PyObject *x)
     if (writer == NULL) {
         return NULL;
     }
+    size = _PyBytesWriter_ResizeToAllocated(writer);
     char *str = PyBytesWriter_GetData(writer);
-    size = _PyBytesWriter_GetAllocated(writer);
 
     /* Run the iterator to exhaustion */
     for (i = 0; ; i++) {
@@ -3110,7 +3112,9 @@ _PyBytes_FromIterator(PyObject *it, PyObject *x)
             if (str == NULL) {
                 goto error;
             }
-            size = _PyBytesWriter_GetAllocated(writer);
+
+            // Set the writer size to its allocated size
+            size = _PyBytesWriter_ResizeToAllocated(writer);
         }
         *str++ = (char) value;
     }
@@ -3747,6 +3751,19 @@ PyBytesWriter_Discard(PyBytesWriter *writer)
 PyObject*
 PyBytesWriter_FinishWithSize(PyBytesWriter *writer, Py_ssize_t size)
 {
+    // Check for negative size here to raise ValueError in all cases, rather
+    // than having a different exception depending on the code path. For
+    // example, _PyBytes_Resize() raises SystemError on negative size.
+    if (size < 0) {
+        PyErr_Format(PyExc_ValueError, "size must be positive");
+        return NULL;
+    }
+
+    if (size > writer->size) {
+        PyErr_SetString(PyExc_ValueError, "size larger than allocated size");
+        return NULL;
+    }
+
     PyObject *result;
     if (size == 0) {
         result = bytes_get_empty();
@@ -3828,13 +3845,6 @@ PyBytesWriter_GetSize(PyBytesWriter *writer)
 }
 
 
-static Py_ssize_t
-_PyBytesWriter_GetAllocated(PyBytesWriter *writer)
-{
-    return byteswriter_allocated(writer);
-}
-
-
 int
 PyBytesWriter_Resize(PyBytesWriter *writer, Py_ssize_t size)
 {
@@ -3934,3 +3944,15 @@ PyBytesWriter_Format(PyBytesWriter *writer, const char *format, ...)
     Py_ssize_t size = buf - byteswriter_data(writer);
     return PyBytesWriter_Resize(writer, size);
 }
+
+
+// Resize the writer to its allocated size.
+// Return the new size.
+// The function cannot fail.
+static Py_ssize_t
+_PyBytesWriter_ResizeToAllocated(PyBytesWriter *writer)
+{
+    Py_ssize_t allocated = byteswriter_allocated(writer);
+    writer->size = allocated;
+    return allocated;
+}