From 89fc158c9136bfcfb4b0673e461998b79382fd62 Mon Sep 17 00:00:00 2001 From: youknowdot Date: Wed, 14 Oct 2026 21:16:12 +0900 Subject: [PATCH] Implement bytearray.take_bytes with stable export length snapshots 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 | 251 ++++++++++++++++++---- extra_tests/snippets/builtin_bytearray.py | 164 ++++++++++++++ 2 files changed, 376 insertions(+), 39 deletions(-) diff --git a/crates/vm/src/builtins/bytearray.rs b/crates/vm/src/builtins/bytearray.rs index fc046460b70..ac7e071f511 100644 --- a/crates/vm/src/builtins/bytearray.rs +++ b/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, - exports: AtomicUsize, + exports: PyMutex, +} + +// 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 { + (self.count != 0).then_some(self.len) + } + + fn retain_existing(&mut self) -> Option { + 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; @@ -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> { 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(&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> { @@ -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 { Ok(()) } + #[pymethod] + fn take_bytes(&self, n: OptionalArg, vm: &VirtualMachine) -> PyResult { + // __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 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 { #[cfg(debug_assertions)] #[pygetset] fn exports(&self) -> usize { - self.exports.load(Ordering::Relaxed) + self.exports.lock().count } #[pymethod] @@ -620,8 +763,7 @@ impl Py { // 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 { fn join(&self, iterable_of_bytes: PyObjectRef, vm: &VirtualMachine) -> PyResult { // 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 { fn split(&self, args: FuncArgs, vm: &VirtualMachine) -> PyResult> { // 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 { #[pymethod] fn rsplit(&self, args: FuncArgs, vm: &VirtualMachine) -> PyResult> { - 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::() - .exports - .fetch_sub(1, Ordering::Release); - }, + release: |buffer| buffer.obj_as::().exports.lock().release(), retain: |buffer| { + // as_buffer holds the payload read guard until PyBuffer::new returns. buffer .obj_as::() .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> { - // 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); + } +} diff --git a/extra_tests/snippets/builtin_bytearray.py b/extra_tests/snippets/builtin_bytearray.py index 1ea800b64db..dbee9a87a7a 100644 --- a/extra_tests/snippets/builtin_bytearray.py +++ b/extra_tests/snippets/builtin_bytearray.py @@ -1587,3 +1587,167 @@ def __release_buffer__(self, view): 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()