Implement bytearray.take_bytes with stable export length snapshots (#9119)

Keep export counts and logical lengths in one synchronized epoch so bounds checks can precede export rejection without recursively borrowing payload data. Preserve shared buffer ownership and callback lifetimes.

Assisted-by: OpenAI Codex:gpt-6.1-sol

crates/vm/src/builtins/bytearray.rs

@@ -15,15 +15,12 @@ use crate::{
         ByteInnerTranslateOptions, ByteOr, BytesJoin, DecodeArgs, PyBytesInner, bytes_decode,
     },
     class::PyClassImpl,
-    common::{
-        atomic::{AtomicUsize, Ordering},
-        lock::{
-            PyDetachingRwLock, PyDetachingRwLockReadGuard, PyDetachingRwLockWriteGuard,
-            PyMappedDetachingRwLockReadGuard, PyMappedDetachingRwLockWriteGuard, PyMutex,
-        },
+    common::lock::{
+        PyDetachingRwLock, PyDetachingRwLockReadGuard, PyDetachingRwLockWriteGuard,
+        PyMappedDetachingRwLockReadGuard, PyMappedDetachingRwLockWriteGuard, PyMutex,
     },
     convert::{ToPyObject, ToPyResult},
-    function::{ArgBytesLike, FuncArgs, PyComparisonValue, PySsize},
+    function::{ArgBytesLike, FuncArgs, OptionalArg, PyComparisonValue, PySsize},
     protocol::{
         BufferDescriptor, BufferFlags, BufferMethods, BufferResizeGuard, PyBuffer, PyIterReturn,
         PyMappingMethods, PyNumberMethods, PySequenceMethods,
@@ -37,12 +34,62 @@ use crate::{
 };
 use bstr::ByteSlice;
 use core::mem::size_of;
+use num_traits::ToPrimitive;
 
 #[pyclass(module = false, name = "bytearray", unhashable = true)]
 #[derive(Debug, Default)]
 pub struct PyByteArray {
     inner: PyDetachingRwLock<PyBytesInner>,
-    exports: AtomicUsize,
+    exports: PyMutex<ExportState>,
+}
+
+// While count is nonzero, len is the stable logical length of the payload.
+// Publish both together; the length is meaningless between export epochs.
+#[derive(Debug, Default)]
+struct ExportState {
+    count: usize,
+    len: usize,
+}
+
+impl ExportState {
+    fn snapshot(&self) -> Option<usize> {
+        (self.count != 0).then_some(self.len)
+    }
+
+    fn retain_existing(&mut self) -> Option<usize> {
+        let len = self.snapshot()?;
+        self.count += 1;
+        Some(len)
+    }
+
+    // The caller holds a payload read guard, so first publication cannot
+    // overlap a resize and an existing epoch has the same length.
+    fn retain(&mut self, len: usize) {
+        if self.count == 0 {
+            self.len = len;
+        }
+        self.count += 1;
+    }
+
+    fn release(&mut self) {
+        self.count -= 1;
+    }
+}
+
+struct ByteArrayExportGuard<'a> {
+    owner: &'a PyByteArray,
+    len: usize,
+}
+
+impl Drop for ByteArrayExportGuard<'_> {
+    fn drop(&mut self) {
+        self.owner.exports.lock().release();
+    }
+}
+
+enum ResizeAccess<'a> {
+    Exclusive(PyDetachingRwLockWriteGuard<'a, PyBytesInner>),
+    Exported { len: usize },
 }
 
 pub(crate) type PyByteArrayRef = PyRef<PyByteArray>;
@@ -95,7 +142,7 @@ impl PyByteArray {
     const fn from_inner(inner: PyBytesInner) -> Self {
         Self {
             inner: PyDetachingRwLock::new(inner),
-            exports: AtomicUsize::new(0),
+            exports: PyMutex::new(ExportState { count: 0, len: 0 }),
         }
     }
 
@@ -103,15 +150,54 @@ impl PyByteArray {
         PyDetachingRwLockReadGuard::map(self.inner.read(), |inner| &*inner.elements)
     }
 
+    /// Borrow for in-place byte replacement only. Do not change the length,
+    /// replace or reallocate storage, or hold this guard across Python calls.
+    /// Operations that may resize or reallocate must use export-checked access.
     pub fn borrow_buf_mut(&self) -> PyMappedDetachingRwLockWriteGuard<'_, Vec<u8>> {
         PyDetachingRwLockWriteGuard::map(self.inner.write(), |inner| &mut inner.elements)
     }
 
+    fn begin_export(&self) -> ByteArrayExportGuard<'_> {
+        // Existing epochs need no payload lock, including in codec callbacks
+        // that already hold a payload read guard.
+        let existing_len = { self.exports.lock().retain_existing() };
+        let len = match existing_len {
+            Some(len) => len,
+            None => {
+                // Never wait for the payload with the metadata lock held.
+                // An earlier resize must finish before first publication.
+                let inner = self.inner.read();
+                let len = inner.elements.len();
+                self.exports.lock().retain(len);
+                len
+            }
+        };
+        ByteArrayExportGuard { owner: self, len }
+    }
+
+    fn resize_access(&self) -> ResizeAccess<'_> {
+        // An export is a borrow someone else still holds, so it is answered
+        // before the lock rather than by waiting on it.
+        let exported_len = { self.exports.lock().snapshot() };
+        if let Some(len) = exported_len {
+            return ResizeAccess::Exported { len };
+        }
+        let inner = self.inner.write();
+        // An export may have started while we waited. Capture its length in
+        // the same decision, without borrowing the payload a second time.
+        let exported_len = { self.exports.lock().snapshot() };
+        match exported_len {
+            Some(len) => {
+                drop(inner);
+                ResizeAccess::Exported { len }
+            }
+            None => ResizeAccess::Exclusive(inner),
+        }
+    }
+
     fn with_export<R>(&self, f: impl FnOnce(usize) -> R) -> R {
-        self.exports.fetch_add(1, Ordering::Release);
-        scopeguard::defer! { self.exports.fetch_sub(1, Ordering::Release); }
-        // Reading the length also waits for any earlier resize to finish.
-        f(self.__len__())
+        let export = self.begin_export();
+        f(export.len)
     }
 
     fn buffer_snapshot(buffer: &PyBuffer, vm: &VirtualMachine) -> PyResult<Vec<u8>> {
@@ -333,8 +419,7 @@ impl PyByteArray {
             return Ok(self.borrow_buf().contains(&byte));
         }
         // Converting the needle runs Python, which must not resize this bytearray.
-        self.exports.fetch_add(1, Ordering::Release);
-        scopeguard::defer! { self.exports.fetch_sub(1, Ordering::Release); }
+        let _export = self.begin_export();
         let needle = ByteInnerSub::from_contains_arg(needle, vm)?;
         self.inner().contains(needle, vm)
     }
@@ -502,6 +587,64 @@ impl Py<PyByteArray> {
         Ok(())
     }
 
+    #[pymethod]
+    fn take_bytes(&self, n: OptionalArg, vm: &VirtualMachine) -> PyResult<PyBytes> {
+        // __index__ can resize self or acquire/release exports. Convert before
+        // observing either the payload or the export epoch.
+        let n = match n.into_option().filter(|n| !vm.is_none(n)) {
+            None => None,
+            Some(n) => {
+                let index = n
+                    .try_index_opt(vm)
+                    .ok_or_else(|| vm.new_type_error("n must be an integer or None"))??;
+                Some(index.as_bigint().to_isize().ok_or_else(|| {
+                    vm.new_index_error(format!(
+                        "cannot fit '{}' into an index-sized integer",
+                        n.class().slot_name()
+                    ))
+                })?)
+            }
+        };
+        let access = self.resize_access();
+        let len = match &access {
+            ResizeAccess::Exclusive(inner) => inner.elements.len(),
+            ResizeAccess::Exported { len } => *len,
+        };
+        // Vec<u8> lengths fit isize. Adding a nonnegative length to a negative
+        // index cannot overflow; preserve the adjusted index in the error.
+        let to_take = match n {
+            None => len as isize,
+            Some(n) if n < 0 => n + len as isize,
+            Some(n) => n,
+        };
+        if to_take < 0 || to_take as usize > len {
+            drop(access);
+            return Err(
+                vm.new_index_error(format!("can't take {to_take} bytes outside size {len}"))
+            );
+        }
+        let ResizeAccess::Exclusive(mut inner) = access else {
+            return Err(vm.new_buffer_error("Existing exports of data: object cannot be re-sized"));
+        };
+        let to_take = to_take as usize;
+        let result = if to_take == 0 {
+            Vec::new()
+        } else if to_take == len {
+            core::mem::take(&mut inner.elements)
+        } else {
+            let mut result = Vec::new();
+            if result.try_reserve_exact(to_take).is_err() {
+                drop(inner);
+                return Err(vm.no_memory_error());
+            }
+            result.extend_from_slice(&inner.elements[..to_take]);
+            inner.elements.drain(..to_take);
+            result
+        };
+        drop(inner);
+        Ok(PyBytes::from(result))
+    }
+
     #[pymethod]
     fn __reduce_ex__(
         &self,
@@ -525,7 +668,7 @@ impl Py<PyByteArray> {
     #[cfg(debug_assertions)]
     #[pygetset]
     fn exports(&self) -> usize {
-        self.exports.load(Ordering::Relaxed)
+        self.exports.lock().count
     }
 
     #[pymethod]
@@ -620,8 +763,7 @@ impl Py<PyByteArray> {
         // is borrowed, with the buffer exported so that resizing this
         // bytearray from there raises BufferError.
         let (sep, bytes_per_sep) = {
-            self.exports.fetch_add(1, Ordering::Release);
-            scopeguard::defer! { self.exports.fetch_sub(1, Ordering::Release); }
+            let _export = self.begin_export();
             options.resolve(vm)?
         };
         Ok(self.inner().hex(sep, bytes_per_sep))
@@ -687,9 +829,8 @@ impl Py<PyByteArray> {
     fn join(&self, iterable_of_bytes: PyObjectRef, vm: &VirtualMachine) -> PyResult<PyByteArray> {
         // Export before driving the iterable: resizing is forbidden, but
         // same-size mutations must be reflected in the joined result.
-        self.exports.fetch_add(1, Ordering::Release);
-        scopeguard::defer! { self.exports.fetch_sub(1, Ordering::Release); }
-        let separator_len = self.__len__();
+        let _export = self.begin_export();
+        let separator_len = _export.len;
         let items = BytesJoin::new(iterable_of_bytes, vm)?;
         Ok(items
             .join(separator_len, || self.borrow_buf().into(), vm)?
@@ -801,8 +942,7 @@ impl Py<PyByteArray> {
     fn split(&self, args: FuncArgs, vm: &VirtualMachine) -> PyResult<Vec<PyObjectRef>> {
         // Converting the separator runs Python, which must not resize this
         // bytearray, so the arguments are bound with the buffer exported.
-        self.exports.fetch_add(1, Ordering::Release);
-        scopeguard::defer! { self.exports.fetch_sub(1, Ordering::Release); }
+        let _export = self.begin_export();
         let options: ByteInnerSplitOptions = args.bind_for(vm, "split")?;
         self.inner()
             .split(options, |s, vm| vm.ctx.new_bytearray(s.to_vec()).into(), vm)
@@ -810,8 +950,7 @@ impl Py<PyByteArray> {
 
     #[pymethod]
     fn rsplit(&self, args: FuncArgs, vm: &VirtualMachine) -> PyResult<Vec<PyObjectRef>> {
-        self.exports.fetch_add(1, Ordering::Release);
-        scopeguard::defer! { self.exports.fetch_sub(1, Ordering::Release); }
+        let _export = self.begin_export();
         let options: ByteInnerSplitOptions = args.bind_for(vm, "rsplit")?;
         self.inner()
             .rsplit(options, |s, vm| vm.ctx.new_bytearray(s.to_vec()).into(), vm)
@@ -1001,17 +1140,14 @@ static BUFFER_METHODS: BufferMethods = BufferMethods {
         )
         .into()
     },
-    release: |buffer| {
-        buffer
-            .obj_as::<PyByteArray>()
-            .exports
-            .fetch_sub(1, Ordering::Release);
-    },
+    release: |buffer| buffer.obj_as::<PyByteArray>().exports.lock().release(),
     retain: |buffer| {
+        // as_buffer holds the payload read guard until PyBuffer::new returns.
         buffer
             .obj_as::<PyByteArray>()
             .exports
-            .fetch_add(1, Ordering::Release);
+            .lock()
+            .retain(buffer.desc.len);
     },
 };
 
@@ -1045,14 +1181,10 @@ impl BufferResizeGuard for PyByteArray {
     type Resizable<'a> = PyDetachingRwLockWriteGuard<'a, PyBytesInner>;
 
     fn try_resizable_opt(&self) -> Option<Self::Resizable<'_>> {
-        // An export is a borrow someone else still holds, so it is answered
-        // before the lock rather than by waiting on it.
-        if self.exports.load(Ordering::SeqCst) != 0 {
-            return None;
+        match self.resize_access() {
+            ResizeAccess::Exclusive(inner) => Some(inner),
+            ResizeAccess::Exported { .. } => None,
         }
-        let inner = self.inner.write();
-        // An export may have been acquired while we waited for the lock.
-        (self.exports.load(Ordering::SeqCst) == 0).then_some(inner)
     }
 }
 
@@ -1241,3 +1373,44 @@ impl IterNext for PyByteArrayIterator {
         })
     }
 }
+
+#[cfg(test)]
+mod export_state_tests {
+    use super::ExportState;
+
+    #[test]
+    fn export_epoch_length_is_stable_until_last_release() {
+        let mut state = ExportState::default();
+        assert_eq!(state.snapshot(), None);
+        assert_eq!(state.retain_existing(), None);
+        state.retain(3);
+        assert_eq!(state.snapshot(), Some(3));
+        assert_eq!(state.retain_existing(), Some(3));
+        state.retain(3);
+        assert_eq!(state.count, 3);
+        state.release();
+        state.release();
+        assert_eq!(state.count, 1);
+        assert_eq!(state.snapshot(), Some(3));
+        let rejected_length = state.snapshot();
+        state.release();
+        assert_eq!(state.snapshot(), None);
+        state.retain(7);
+        assert_eq!(state.snapshot(), Some(7));
+        assert_eq!(rejected_length, Some(3));
+        state.release();
+        assert_eq!(state.snapshot(), None);
+    }
+
+    #[test]
+    fn empty_export_is_distinct_from_no_export() {
+        let mut state = ExportState::default();
+        state.retain(0);
+        assert_eq!(state.snapshot(), Some(0));
+        assert_eq!(state.retain_existing(), Some(0));
+        state.release();
+        assert_eq!(state.snapshot(), Some(0));
+        state.release();
+        assert_eq!(state.snapshot(), None);
+    }
+}

extra_tests/snippets/builtin_bytearray.py

@@ -1587,3 +1587,167 @@ def test_bytearray_contains_builtin_and_callback_exports():
 
 
 test_bytearray_contains_builtin_and_callback_exports()
+
+
+def test_bytearray_take_bytes_bounded():
+    for n, expected, remaining in (
+        (None, b"abcd", b""),
+        (0, b"", b"abcd"),
+        (1, b"a", b"bcd"),
+        (2, b"ab", b"cd"),
+        (4, b"abcd", b""),
+        (-1, b"abc", b"d"),
+        (-4, b"", b"abcd"),
+        (True, b"a", b"bcd"),
+    ):
+        target = bytearray(b"abcd")
+        result = target.take_bytes(n)
+        assert type(result) is bytes
+        assert result == expected
+        assert target == remaining
+        target[:] = b"z"
+        assert result == expected
+    target = bytearray(b"abc")
+    assert target.take_bytes() == b"abc"
+    assert target == b""
+    for n in (None, 0):
+        assert bytearray().take_bytes(n) == b""
+    for n in (5, -5, sys.maxsize + 1, -sys.maxsize - 2):
+        target = bytearray(b"abcd")
+        with assert_raises(IndexError):
+            target.take_bytes(n)
+        assert target == b"abcd"
+    for n in (1.0, "1", b"1", object()):
+        with assert_raises(TypeError):
+            bytearray(b"a").take_bytes(n)
+    with assert_raises(TypeError):
+        bytearray(b"a").take_bytes(n=1)
+    with assert_raises(TypeError):
+        bytearray(b"a").take_bytes(1, 2)
+    for content in (b"", b"abcd"):
+        target = bytearray(content)
+        view = memoryview(target)
+        try:
+            for n in (None, 0, len(content)):
+                with assert_raises(BufferError):
+                    target.take_bytes(n)
+                assert target == content
+            for n in (len(content) + 1, -len(content) - 1):
+                with assert_raises(IndexError):
+                    target.take_bytes(n)
+                assert target == content
+        finally:
+            view.release()
+        assert target.take_bytes() == content
+    target = bytearray(b"abcd")
+    parent = memoryview(target)
+    child = parent[1:3]
+    clone = memoryview(parent)
+    parent.release()
+    try:
+        with assert_raises(BufferError):
+            target.take_bytes(0)
+        child.release()
+        with assert_raises(BufferError):
+            target.take_bytes(1)
+    finally:
+        child.release()
+        clone.release()
+    assert target.take_bytes() == b"abcd"
+
+
+def test_bytearray_take_bytes_reentrant_bounded():
+    class Index:
+        def __init__(self, action, result):
+            self.action = action
+            self.result = result
+
+        def __index__(self):
+            self.action()
+            return self.result
+
+    for replacement, n, expected, remaining in (
+        (b"ab", 2, b"ab", b""),
+        (b"", 0, b"", b""),
+        (b"abcdef", 5, b"abcde", b"f"),
+        (b"ab", -1, b"a", b"b"),
+    ):
+        target = bytearray(b"abcd")
+        index = Index(lambda: target.__setitem__(slice(None), replacement), n)
+        assert target.take_bytes(index) == expected
+        assert target == remaining
+    target = bytearray(b"abcd")
+    with assert_raises(IndexError):
+        target.take_bytes(Index(target.clear, 1))
+    assert target == b""
+    for n, error in ((1, BufferError), (5, IndexError)):
+        target = bytearray(b"abcd")
+        views = []
+        try:
+            with assert_raises(error):
+                target.take_bytes(Index(lambda: views.append(memoryview(target)), n))
+            assert target == b"abcd"
+        finally:
+            for view in views:
+                view.release()
+    for n, error in ((2, None), (5, IndexError)):
+        target = bytearray(b"abcd")
+        view = memoryview(target)
+        try:
+            index = Index(view.release, n)
+            if error is None:
+                assert target.take_bytes(index) == b"ab"
+                assert target == b"cd"
+            else:
+                with assert_raises(error):
+                    target.take_bytes(index)
+                assert target == b"abcd"
+        finally:
+            view.release()
+    target = bytearray(b"abcd")
+    with assert_raises(TypeError):
+        target.take_bytes(Index(lambda: None, "wrong"))
+    assert target == b"abcd"
+
+    def raises():
+        raise ValueError("index callback")
+
+    with assert_raises(ValueError):
+        target.take_bytes(Index(raises, 1))
+    assert target == b"abcd"
+
+    class MembershipIndex:
+        def __index__(self):
+            with assert_raises(IndexError):
+                target.take_bytes(5)
+            with assert_raises(BufferError):
+                target.take_bytes(0)
+            return ord("a")
+
+    assert MembershipIndex() in target
+    assert target.take_bytes() == b"abcd"
+
+
+def test_bytearray_take_bytes_codec_callback_bounded():
+    import codecs
+
+    target = bytearray(b"\xff")
+    calls = []
+
+    def handler(error):
+        with assert_raises(IndexError):
+            target.take_bytes(2)
+        with assert_raises(BufferError):
+            target.take_bytes(0)
+        calls.append(error.start)
+        return ("?", error.end)
+
+    codecs.register_error("bytearray_take_bytes_bounded", handler)
+    assert target.decode("ascii", "bytearray_take_bytes_bounded") == "?"
+    assert calls == [0]
+    assert target.take_bytes() == b"\xff"
+
+
+test_bytearray_take_bytes_bounded()
+test_bytearray_take_bytes_reentrant_bounded()
+test_bytearray_take_bytes_codec_callback_bounded()