gh-143308: fix UAF when PickleBuffer is concurrently mutated in a callback (#143312)

Co-authored-by: Aaron Wieczorek <woz@Aarons-MacBook-Pro.local>
Co-authored-by: Bénédikt Tran <10796600+picnixz@users.noreply.github.com>

Author: 3 people

Lib/test/pickletester.py

@@ -4261,6 +4261,35 @@ def check_array(arr):
         # 2-D, non-contiguous
         check_array(arr[::2])
 
+    def test_concurrent_mutation_in_buffer_with_bytearray(self):
+        def factory():
+            s = b"a" * 16
+            return bytearray(s), s
+        self.do_test_concurrent_mutation_in_buffer_callback(factory)
+
+    def test_concurrent_mutation_in_buffer_with_memoryview(self):
+        def factory():
+            obj = memoryview(b"a" * 32)[10:26]
+            sub = b"a" * len(obj)
+            return obj, sub
+        self.do_test_concurrent_mutation_in_buffer_callback(factory)
+
+    def do_test_concurrent_mutation_in_buffer_callback(self, factory):
+        # See: https://github.com/python/cpython/issues/143308.
+        class R:
+            def __bool__(self):
+                buf.release()
+                return True
+
+        for proto in range(5, pickle.HIGHEST_PROTOCOL + 1):
+            obj, sub = factory()
+            buf = pickle.PickleBuffer(obj)
+            buffer_callback = lambda _: R()
+
+            with self.subTest(proto=proto, obj=obj, sub=sub):
+                res = self.dumps(buf, proto, buffer_callback=buffer_callback)
+                self.assertIn(sub, res)
+
     def test_evil_class_mutating_dict(self):
         # https://github.com/python/cpython/issues/92930
         from random import getrandbits

Misc/NEWS.d/next/Library/2025-12-31-17-38-33.gh-issue-143308.lY8UCR.rst

@@ -0,0 +1,3 @@
+:mod:`pickle`: fix use-after-free crashes when a :class:`~pickle.PickleBuffer`
+is concurrently mutated by a custom buffer callback during pickling.
+Patch by Bénédikt Tran and Aaron Wieczorek.

Modules/_pickle.c

@@ -2636,53 +2636,61 @@ save_picklebuffer(PickleState *st, PicklerObject *self, PyObject *obj)
                         "PickleBuffer can only be pickled with protocol >= 5");
         return -1;
     }
-    const Py_buffer* view = PyPickleBuffer_GetBuffer(obj);
-    if (view == NULL) {
+    Py_buffer view;
+    if (PyObject_GetBuffer(obj, &view, PyBUF_FULL_RO) != 0) {
         return -1;
     }
-    if (view->suboffsets != NULL || !PyBuffer_IsContiguous(view, 'A')) {
+    if (view.suboffsets != NULL || !PyBuffer_IsContiguous(&view, 'A')) {
         PyErr_SetString(st->PicklingError,
                         "PickleBuffer can not be pickled when "
                         "pointing to a non-contiguous buffer");
-        return -1;
+        goto error;
     }
+
+    int rc = 0;
     int in_band = 1;
     if (self->buffer_callback != NULL) {
         PyObject *ret = PyObject_CallOneArg(self->buffer_callback, obj);
         if (ret == NULL) {
-            return -1;
+            goto error;
         }
         in_band = PyObject_IsTrue(ret);
         Py_DECREF(ret);
         if (in_band == -1) {
-            return -1;
+            goto error;
         }
     }
     if (in_band) {
         /* Write data in-band */
-        if (view->readonly) {
-            return _save_bytes_data(st, self, obj, (const char *)view->buf,
-                                    view->len);
+        if (view.readonly) {
+            rc = _save_bytes_data(st, self, obj, (const char *)view.buf,
+                                  view.len);
         }
         else {
-            return _save_bytearray_data(st, self, obj, (const char *)view->buf,
-                                        view->len);
+            rc = _save_bytearray_data(st, self, obj, (const char *)view.buf,
+                                      view.len);
         }
     }
     else {
         /* Write data out-of-band */
         const char next_buffer_op = NEXT_BUFFER;
         if (_Pickler_Write(self, &next_buffer_op, 1) < 0) {
-            return -1;
+            goto error;
         }
-        if (view->readonly) {
+        if (view.readonly) {
             const char readonly_buffer_op = READONLY_BUFFER;
             if (_Pickler_Write(self, &readonly_buffer_op, 1) < 0) {
-                return -1;
+                goto error;
             }
         }
     }
-    return 0;
+
+    PyBuffer_Release(&view);
+    return rc;
+
+error:
+    PyBuffer_Release(&view);
+    return -1;
 }
 
 /* A copy of PyUnicode_AsRawUnicodeEscapeString() that also translates