gh-143378: refactor the buffer check and unit tests for C and Python implementation

Author: superboy-zjc

Lib/_pyio.py

@@ -955,6 +955,8 @@ def write(self, b):
             raise TypeError("can't write str to binary stream")
         with memoryview(b) as view:
             n = view.nbytes  # Size of any bytes-like object
+            if self.closed:
+                raise ValueError("write to closed file")
         if n == 0:
             return 0
 

Lib/test/test_io/test_memoryio.py

@@ -587,6 +587,52 @@ def test_issue5449(self):
         self.ioclass(initial_bytes=buf)
         self.assertRaises(TypeError, self.ioclass, buf, foo=None)
 
+    def test_write_concurrent_close(self):
+        # Prevent crashes when memio.write() concurrently closes 'memio'.
+        # See: https://github.com/python/cpython/issues/143378.
+        class B:
+            def __buffer__(self, flags):
+                memio.close()
+                return memoryview(b"A")
+
+        memio = self.ioclass()
+        self.assertRaises(ValueError, memio.write, B())
+
+    def test_writelines_concurrent_close(self):
+        # Prevent crashes when memio.writelines() concurrently closes 'memio'.
+        # See: https://github.com/python/cpython/issues/143378.
+        class B:
+            def __buffer__(self, flags):
+                memio.close()
+                return memoryview(b"A")
+
+        memio = self.ioclass()
+        self.assertRaises(ValueError, memio.writelines, [B()])
+
+    def test_write_concurrent_export(self):
+        # Prevent crashes when memio.write() concurrently exports 'memio'.
+        # See: https://github.com/python/cpython/issues/143378.
+        class B:
+            buf = None
+            def __buffer__(self, flags):
+                self.buf = memio.getbuffer()
+                return memoryview(b"A")
+
+        memio = self.ioclass()
+        self.assertRaises(BufferError, memio.write, B())
+
+    def test_writelines_concurrent_export(self):
+        # Prevent crashes when memio.writelines() concurrently exports 'memio'.
+        # See: https://github.com/python/cpython/issues/143378.
+        class B:
+            buf = None
+            def __buffer__(self, flags):
+                self.buf = memio.getbuffer()
+                return memoryview(b"A")
+
+        memio = self.ioclass()
+        self.assertRaises(BufferError, memio.writelines, [B()])
+
 
 class TextIOTestMixin:
 
@@ -856,31 +902,6 @@ def test_cow_mutable(self):
         memio = self.ioclass(ba)
         self.assertEqual(sys.getrefcount(ba), old_rc)
 
-    @support.cpython_only
-    def test_write_concurrent_mutation(self):
-        # Prevent crashes when buf.write() concurrently mutates 'buf'.
-        # See: https://github.com/python/cpython/issues/143378.
-        class B:
-            def __buffer__(self, flags):
-                memio.close()
-                return memoryview(b"A")
-
-        memio = self.ioclass()
-        self.assertRaises(BufferError, memio.write, B())
-
-    @support.cpython_only
-    def test_writelines_concurrent_mutation(self):
-        # Prevent crashes when buf.writelines() concurrently mutates 'buf'.
-        # See: https://github.com/python/cpython/issues/143378.
-        class B:
-            def __buffer__(self, flags):
-                memio.close()
-                return memoryview(b"A")
-
-        memio = self.ioclass()
-        self.assertRaises(BufferError, memio.writelines, [B()])
-
-
 class CStringIOTest(PyStringIOTest):
     ioclass = io.StringIO
     UnsupportedOperation = io.UnsupportedOperation

Modules/_io/bytesio.c

@@ -202,14 +202,17 @@ write_bytes_lock_held(bytesio *self, PyObject *b)
     }
 
     Py_buffer buf;
-    self->exports++;
+    Py_ssize_t len;
     if (PyObject_GetBuffer(b, &buf, PyBUF_CONTIG_RO) < 0) {
-        self->exports--;
         return -1;
     }
-    self->exports--;
 
-    Py_ssize_t len = buf.len;
+    if (check_closed(self) || check_exports(self)) {
+        len = -1;
+        goto done;
+    }
+
+    len = buf.len;
     if (len == 0) {
         goto done;
     }