Deal with __release_buffer__ saving references to the buffer

python · JelleZijlstra · May 8, 2023 · May 6, 2023 · May 6, 2023 · May 6, 2023
commit 80a4b0b9b3dc7c3bc440272c2079f17e032da64a
diff --git a/Include/cpython/memoryobject.h b/Include/cpython/memoryobject.h
@@ -24,6 +24,7 @@ typedef struct {
 #define _Py_MEMORYVIEW_FORTRAN     0x004  /* Fortran contiguous layout */
 #define _Py_MEMORYVIEW_SCALAR      0x008  /* scalar: ndim = 0 */
 #define _Py_MEMORYVIEW_PIL         0x010  /* PIL-style layout */
+#define _Py_MEMORYVIEW_RESTRICTED  0x020  /* Disallow additional references */
 
 typedef struct {
     PyObject_VAR_HEAD

diff --git a/Lib/test/test_buffer.py b/Lib/test/test_buffer.py
@@ -4634,6 +4634,42 @@ def __release_buffer__(self, view):
         self.assertEqual(rb_call_count, 1)
         self.assertIs(rb_raised, True)
 
+    def test_override_only_release(self):
+        class C(bytearray):
+            def __release_buffer__(self, buffer):
+                super().__release_buffer__(buffer)
+
+        c = C(b"hello")
+        with memoryview(c) as mv:
+            self.assertEqual(mv.tobytes(), b"hello")
+
+    def test_release_saves_reference(self):
+        smuggled_buffer = None
+
+        class C(bytearray):
+            def __release_buffer__(s, buffer: memoryview):
+                with self.assertRaises(ValueError):
+                    memoryview(buffer)
+                with self.assertRaises(ValueError):
+                    buffer.cast("b")
+                with self.assertRaises(ValueError):
+                    buffer.toreadonly()
+                with self.assertRaises(ValueError):
+                    buffer[:1]
+                with self.assertRaises(ValueError):
+                    buffer.__buffer__(0)
+                nonlocal smuggled_buffer
+                smuggled_buffer = buffer
+                self.assertEqual(buffer.tobytes(), b"hello")
+                super().__release_buffer__(buffer)
+
+        c = C(b"hello")
+        with memoryview(c) as mv:
+            self.assertEqual(mv.tobytes(), b"hello")
+        c.clear()
+        with self.assertRaises(ValueError):
+            smuggled_buffer.tobytes()
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/Objects/memoryobject.c b/Objects/memoryobject.c
@@ -193,6 +193,20 @@ PyTypeObject _PyManagedBuffer_Type = {
         return -1;                                                \
     }
 
+#define CHECK_RESTRICTED(mv) \
+    if (((PyMemoryViewObject *)(mv))->flags & _Py_MEMORYVIEW_RESTRICTED) { \
+        PyErr_SetString(PyExc_ValueError,                                  \
+            "cannot create new view on restricted memoryview");            \
+        return NULL;                                                       \
+    }
+
+#define CHECK_RESTRICTED_INT(mv) \
+    if (((PyMemoryViewObject *)(mv))->flags & _Py_MEMORYVIEW_RESTRICTED) { \
+        PyErr_SetString(PyExc_ValueError,                                  \
+            "cannot create new view on restricted memoryview");            \
+        return -1;                                                       \
+    }
+
 /* See gh-92888. These macros signal that we need to check the memoryview
    again due to possible read after frees. */
 #define CHECK_RELEASED_AGAIN(mv) CHECK_RELEASED(mv)
@@ -789,6 +803,7 @@ PyMemoryView_FromObjectAndFlags(PyObject *v, int flags)
     if (PyMemoryView_Check(v)) {
         PyMemoryViewObject *mv = (PyMemoryViewObject *)v;
         CHECK_RELEASED(mv);
+        CHECK_RESTRICTED(mv);
         return mbuf_add_view(mv->mbuf, &mv->view);
     }
     else if (PyObject_CheckBuffer(v)) {
@@ -1421,6 +1436,7 @@ memoryview_cast_impl(PyMemoryViewObject *self, PyObject *format,
     Py_ssize_t ndim = 1;
 
     CHECK_RELEASED(self);
+    CHECK_RESTRICTED(self);
 
     if (!MV_C_CONTIGUOUS(self->flags)) {
         PyErr_SetString(PyExc_TypeError,
@@ -1476,6 +1492,7 @@ memoryview_toreadonly_impl(PyMemoryViewObject *self)
 /*[clinic end generated code: output=2c7e056f04c99e62 input=dc06d20f19ba236f]*/
 {
     CHECK_RELEASED(self);
+    CHECK_RESTRICTED(self);
     /* Even if self is already readonly, we still need to create a new
      * object for .release() to work correctly.
      */
@@ -1498,6 +1515,7 @@ memory_getbuf(PyMemoryViewObject *self, Py_buffer *view, int flags)
     int baseflags = self->flags;
 
     CHECK_RELEASED_INT(self);
+    CHECK_RESTRICTED_INT(self);
 
     /* start with complete information */
     *view = *base;
@@ -2559,6 +2577,7 @@ memory_subscript(PyMemoryViewObject *self, PyObject *key)
         return memory_item(self, index);
     }
     else if (PySlice_Check(key)) {
+        CHECK_RESTRICTED(self);
         PyMemoryViewObject *sliced;
 
         sliced = (PyMemoryViewObject *)mbuf_add_view(self->mbuf, view);

@@ -9107,27 +9107,39 @@ static void
 releasebuffer_call_python(PyObject *self, Py_buffer *buffer)
 {
     PyObject *mv;
+    bool is_buffer_wrapper = Py_TYPE(buffer->obj) == &_PyBufferWrapper_Type;
     if (Py_TYPE(buffer->obj) == &_PyBufferWrapper_Type) {
         // Make sure we pass the same memoryview to
         // __release_buffer__() that __buffer__() returned.
         mv = Py_NewRef(((PyBufferWrapper *)buffer->obj)->mv);
     }
     else {
+        // This means we are not dealing with a memoryview returned
+        // from a Python __buffer__ function.
         mv = PyMemoryView_FromBuffer(buffer);
         if (mv == NULL) {
             PyErr_WriteUnraisable(self);
             return;
         }
+        // Set the memoryview to restricted mode, which forbids
+        // users from saving any reference to the underlying buffer
+        // (e.g., by doing .cast()). This is necessary to ensure
+        // no Python code retains a reference to the to-be-released
+        // buffer.
+        ((PyMemoryViewObject *)mv)->flags |= _Py_MEMORYVIEW_RESTRICTED;
     }
     PyObject *stack[2] = {self, mv};
     PyObject *ret = vectorcall_method(&_Py_ID(__release_buffer__), stack, 2);
-    Py_DECREF(mv);
     if (ret == NULL) {
         PyErr_WriteUnraisable(self);
     }
     else {
         Py_DECREF(ret);
     }
+    if (!is_buffer_wrapper) {
+        PyObject_CallMethodNoArgs(mv, &_Py_ID(release));
+    }
+    Py_DECREF(mv);
 }
 
 /*
@@ -9144,12 +9156,12 @@ releasebuffer_call_python(PyObject *self, Py_buffer *buffer)
 static void
 slot_bf_releasebuffer(PyObject *self, Py_buffer *buffer)
 {
+    releasebuffer_call_python(self, buffer);
     if (releasebuffer_maybe_call_super(self, buffer) < 0) {
         if (PyErr_Occurred()) {
             PyErr_WriteUnraisable(self);
         }
     }
-    releasebuffer_call_python(self, buffer);
 }
 
 static PyObject *