diff --git a/crates/vm/src/builtins/bytearray.rs b/crates/vm/src/builtins/bytearray.rs index 5afeec628..4904f96fb 100644 --- a/crates/vm/src/builtins/bytearray.rs +++ b/crates/vm/src/builtins/bytearray.rs @@ -173,11 +173,6 @@ impl PyByteArray { fn inner(&self) -> PyDetachingRwLockReadGuard<'_, PyBytesInner> { self.inner.read() } - #[inline] - fn inner_mut(&self) -> PyDetachingRwLockWriteGuard<'_, PyBytesInner> { - self.inner.write() - } - fn __len__(&self) -> usize { self.borrow_buf().len() } @@ -781,9 +776,18 @@ impl Initializer for PyByteArray { type Args = ByteInnerNewOptions; fn init(zelf: &Py, options: Self::Args, vm: &VirtualMachine) -> PyResult<()> { - // First unpack bytearray and *then* get a lock to set it. - let mut inner = options.get_inner(bytearray_from_object, vm)?; - core::mem::swap(&mut *zelf.inner_mut(), &mut inner); + // Reject exports and clear existing contents before converting source. + zelf.try_resizable(vm)?.elements.clear(); + // Conversion can call Python, so do not hold the lock across it and + // check again for exports acquired by a callback or another thread. + let (mut inner, zero_count) = + options.get_inner_with_zero_count(bytearray_from_object, vm)?; + // A zero integer count performs no further writes, including when its + // __index__ callback changed this bytearray or acquired an export. + if zero_count { + return Ok(()); + } + core::mem::swap(&mut *zelf.try_resizable(vm)?, &mut inner); Ok(()) } } @@ -841,9 +845,11 @@ impl AsBuffer for PyByteArray { } fn as_buffer(zelf: &Py, _vm: &VirtualMachine) -> PyResult { + // Keep the length stable until PyBuffer::new retains the export. + let inner = zelf.inner(); Ok(PyBuffer::new( zelf.to_owned().into(), - BufferDescriptor::simple(zelf.__len__(), false), + BufferDescriptor::simple(inner.elements.len(), false), &BUFFER_METHODS, )) } @@ -855,7 +861,12 @@ impl BufferResizeGuard for PyByteArray { 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. - (self.exports.load(Ordering::SeqCst) == 0).then(|| self.inner.write()) + if self.exports.load(Ordering::SeqCst) != 0 { + return 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) } } @@ -907,7 +918,7 @@ impl AsSequence for PyByteArray { if let Some(value) = value { zelf._setitem_by_index(i, &value, vm) } else { - zelf.borrow_buf_mut().delitem_by_index(vm, i) + zelf.try_resizable(vm)?.elements.delitem_by_index(vm, i) } }), contains: atomic_func!(|seq, other, vm| { diff --git a/crates/vm/src/bytes_inner.rs b/crates/vm/src/bytes_inner.rs index f577e3020..48f9bd1d4 100644 --- a/crates/vm/src/bytes_inner.rs +++ b/crates/vm/src/bytes_inner.rs @@ -225,12 +225,24 @@ impl ByteInnerNewOptions { /// read: `bytes()` and `bytearray()` differ in whether they ask it how long /// it is. pub fn get_inner(self, from_object: FromObject, vm: &VirtualMachine) -> PyResult { - match (self.source, self.encoding, self.errors) { + self.get_inner_with_zero_count(from_object, vm) + .map(|(inner, _)| inner) + } + + pub(crate) fn get_inner_with_zero_count( + self, + from_object: FromObject, + vm: &VirtualMachine, + ) -> PyResult<(PyBytesInner, bool)> { + let mut zero_count = false; + let inner = match (self.source, self.encoding, self.errors) { (OptionalArg::Present(obj), OptionalArg::Missing, OptionalArg::Missing) => { // Try __index__ first to handle int-like objects that might raise custom exceptions if let Some(index_result) = obj.try_index_opt(vm) { match index_result { - Ok(index) => Self::get_value_from_size(&index, vm), + Ok(index) => Self::get_value_from_size(&index, vm).inspect(|inner| { + zero_count = inner.elements.is_empty(); + }), Err(e) => { // Only propagate non-TypeError exceptions // TypeError means the object doesn't support __index__, so fall back @@ -266,7 +278,8 @@ impl ByteInnerNewOptions { (OptionalArg::Present(_), OptionalArg::Missing, OptionalArg::Present(_)) => { Err(vm.new_type_error(STRING_WITHOUT_ENCODING.to_owned())) } - } + }?; + Ok((inner, zero_count)) } } diff --git a/crates/vm/src/stdlib/_io.rs b/crates/vm/src/stdlib/_io.rs index 1d86717a2..86b68322a 100644 --- a/crates/vm/src/stdlib/_io.rs +++ b/crates/vm/src/stdlib/_io.rs @@ -762,10 +762,13 @@ mod _io { ))); } let n = n as usize; - let mut bytes = b.borrow_buf_mut(); - bytes.truncate(n); - // FIXME: try to use Arc::unwrap on the bytearray to get at the inner buffer - bytes.clone().to_pyobject(vm) + let mut inner = b.try_resizable(vm)?; + let bytes = &mut inner.elements; + bytes + .try_reserve_exact(n.saturating_sub(bytes.len())) + .map_err(|_| vm.no_memory_error())?; + bytes.resize(n, 0); + core::mem::take(bytes).to_pyobject(vm) } }) } else { diff --git a/extra_tests/snippets/builtin_bytearray_exports.py b/extra_tests/snippets/builtin_bytearray_exports.py new file mode 100644 index 000000000..06c102c32 --- /dev/null +++ b/extra_tests/snippets/builtin_bytearray_exports.py @@ -0,0 +1,136 @@ +"""Resizing paths must respect outstanding bytearray exports.""" + +import io + +from testutils import assert_raises + +for args in ((), (b"abc",), (b"x",), (b"",), (object(),)): + value = bytearray(b"abc") + with memoryview(value) as view: + with assert_raises(BufferError): + value.__init__(*args) + assert value == b"abc" + assert view.tobytes() == b"abc" + value.__init__(b"released") + assert value == b"released" + +value = bytearray(b"abc") +with assert_raises(TypeError): + value.__init__(object()) +assert value == b"" +value = bytearray(b"abc") +value.__init__(value) +assert value == b"" + + +class ExportDuringInit: + def __init__(self, count, payload=b""): + self.count = count + self.payload = payload + self.calls = 0 + + def __index__(self): + self.calls += 1 + value.extend(self.payload) + self.view = memoryview(value) + return self.count + + +value = bytearray(b"abc") +source = ExportDuringInit(2) +with assert_raises(BufferError): + value.__init__(source) +assert value == b"" +assert source.view.tobytes() == b"" +assert source.calls == 1 +source.view.release() +value.extend(b"released") +assert value == b"released" + +value = bytearray(b"abc") +source = ExportDuringInit(0) +value.__init__(source) +assert value == b"" +assert source.view.tobytes() == b"" +source.view.release() + +value = bytearray(b"abc") +source = ExportDuringInit(0, b"xy") +value.__init__(source) +assert value == b"xy" +assert source.view.tobytes() == b"xy" +assert source.calls == 1 +source.view.release() + + +class FailingIndex: + def __index__(self): + assert value == b"" + raise ValueError("conversion failed") + + +value = bytearray(b"abc") +with assert_raises(ValueError): + value.__init__(FailingIndex()) +assert value == b"" + +for key in (0, slice(0, 1)): + value = bytearray(b"abc") + with memoryview(value) as view: + with assert_raises(BufferError): + del value[key] + assert view.tobytes() == b"abc" + del value[key] + assert value == b"bc" + + +class Reader(io.RawIOBase): + def __init__(self, count, retain, payload=b"abcd"): + self.count = count + self.retain = retain + self.payload = payload + + def readinto(self, buffer): + buffer[:] = self.payload + self.buffer = buffer + if self.retain: + self.view = memoryview(buffer) + return self.count + + +for count in (0, 2, 4): + reader = Reader(count, False) + assert reader.read(4) == b"abcd"[:count] + assert reader.buffer == b"" + + reader = Reader(count, True) + with assert_raises(BufferError): + reader.read(4) + assert reader.view.tobytes() == b"abcd" + reader.view.release() + reader.buffer.clear() + +for count in (-1, 5): + reader = Reader(count, True) + with assert_raises(ValueError): + reader.read(4) + assert reader.view.tobytes() == b"abcd" + reader.view.release() + +reader = Reader(None, True) +assert reader.read(4) is None +assert reader.view.tobytes() == b"abcd" +reader.view.release() + +# The callback can shrink the temporary buffer, so slicing by the original +# requested size without checking the current buffer is not safe. +reader = Reader(2, False, b"xy") +assert reader.read(4) == b"xy" +assert reader.buffer == b"" + +reader = Reader(4, False, b"xy") +result = reader.read(4) +assert len(result) == 4 +assert result[:2] == b"xy" +# The contents of the newly grown region are unspecified by CPython. +assert reader.buffer == b"" diff --git a/extra_tests/snippets/builtin_bytearray_threading.py b/extra_tests/snippets/builtin_bytearray_threading.py new file mode 100644 index 000000000..c28bcb7b3 --- /dev/null +++ b/extra_tests/snippets/builtin_bytearray_threading.py @@ -0,0 +1,55 @@ +"""Acquiring a bytearray export must be atomic with respect to resizing.""" + +import threading + + +def exercise_exports(): + payload = b"x" * 65536 + value = bytearray(payload) + barrier = threading.Barrier(3, timeout=5) + errors = [] + + def worker(operation): + try: + for _ in range(256): + barrier.wait() + operation() + barrier.wait() + except BaseException as exc: + errors.append(exc) + barrier.abort() + + def resize(): + try: + value.clear() + value.extend(payload) + except BufferError: + pass + + def decode(): + result = value.decode("ascii") + assert result in ("", "x" * len(payload)) + + def export(): + with memoryview(value) as view: + size = len(view) + assert size in (0, len(payload)) + # The acquired export keeps its original length and contents even + # if the resizing thread runs between these two reads. + assert view.tobytes() == payload[:size] + assert len(view) == size + + threads = [ + threading.Thread(target=worker, args=(operation,), daemon=True) + for operation in (resize, decode, export) + ] + for thread in threads: + thread.start() + for thread in threads: + thread.join(timeout=10) + assert not any(thread.is_alive() for thread in threads), "worker did not finish" + if errors: + raise errors[0] + + +exercise_exports()