diff --git a/crates/stdlib/src/mmap.rs b/crates/stdlib/src/mmap.rs index a39c1278455..a242f8f1d29 100644 --- a/crates/stdlib/src/mmap.rs +++ b/crates/stdlib/src/mmap.rs @@ -26,7 +26,10 @@ mod mmap { use num_traits::Signed; use std::io::Write; #[cfg(windows)] - use {core::hint::cold_path, memchr::memchr, rustpython_vm::exceptions, std::io}; + use { + core::hint::cold_path, memchr::memchr, num_traits::ToPrimitive, rustpython_vm::exceptions, + std::io, + }; #[cfg(unix)] use rustpython_host_env::crt_fd; @@ -664,7 +667,9 @@ mod mmap { impl AsBuffer for PyMmap { const RELEASE_BUFFER: bool = true; - fn as_buffer(zelf: &Py, _vm: &VirtualMachine) -> PyResult { + fn as_buffer(zelf: &Py, vm: &VirtualMachine) -> PyResult { + // Publish the export while close/resize cannot invalidate the mapping. + let _mmap = zelf.check_valid(vm)?; let readonly = matches!(zelf.access, AccessMode::Read); let buf = PyBuffer::new( zelf.to_owned().into(), @@ -781,6 +786,7 @@ mod mmap { let m = self.mmap.lock(); if m.is_none() { + drop(m); return Err(vm.new_value_error("mmap closed or invalid")); } @@ -789,20 +795,31 @@ mod mmap { /// TODO: impl resize #[allow(dead_code)] - fn check_resizeable(&self, vm: &VirtualMachine) -> PyResult<()> { + fn check_resizeable( + &self, + vm: &VirtualMachine, + ) -> PyResult>> { + // Borrowed buffers can hold the data lock across blocking I/O. if self.exports.load() > 0 { return Err(vm.new_buffer_error("mmap can't resize with extant buffers exported.")); } + let mmap = self.check_valid(vm)?; + if self.exports.load() > 0 { + drop(mmap); + return Err(vm.new_buffer_error("mmap can't resize with extant buffers exported.")); + } #[cfg(unix)] if !self.trackfd { + drop(mmap); return Err(vm.new_value_error("mmap can't resize with trackfd=False.")); } if self.access == AccessMode::Write || self.access == AccessMode::Default { - return Ok(()); + return Ok(mmap); } + drop(mmap); Err(vm.new_type_error("mmap can't resize a readonly or copy-on-write memory map.")) } @@ -813,15 +830,20 @@ mod mmap { #[pymethod] fn close(zelf: &Py, vm: &VirtualMachine) -> PyResult<()> { + // Reject an existing export without waiting for its borrowed-data guard. + if zelf.exports.load() > 0 { + return Err(vm.new_buffer_error("cannot close exported pointers exist.")); + } + let mut mmap = zelf.mmap.lock(); if Self::closed(zelf) { return Ok(()); } if zelf.exports.load() > 0 { + drop(mmap); return Err(vm.new_buffer_error("cannot close exported pointers exist.")); } - let mut mmap = zelf.mmap.lock(); zelf.closed.store(true); *mmap = None; @@ -1051,7 +1073,7 @@ mod mmap { #[cfg(unix)] #[pymethod] fn resize(zelf: &Py, newsize: PyIntRef, vm: &VirtualMachine) -> PyResult<()> { - zelf.check_resizeable(vm)?; + drop(zelf.check_resizeable(vm)?); let new_size: isize = newsize.try_to_primitive(vm).map_err(|_| { vm.new_overflow_error("Python int too large to convert to C ssize_t") @@ -1076,23 +1098,24 @@ mod mmap { #[cfg(windows)] #[pymethod] fn resize(zelf: &Py, newsize: PyIntRef, vm: &VirtualMachine) -> PyResult<()> { - zelf.check_resizeable(vm)?; - - let newsize: usize = newsize - .try_to_primitive(vm) - .map_err(|_| vm.new_value_error("new size out of range"))?; + let newsize = newsize.as_bigint().to_usize(); + // Get the lock on mmap + // Keep the export check valid until resize ends. + let mut mmap_guard = zelf.check_resizeable(vm)?; + let Some(newsize) = newsize else { + drop(mmap_guard); + return Err(vm.new_value_error("new size out of range")); + }; if newsize == 0 { + drop(mmap_guard); return Err(vm.new_value_error("new size must be positive")); } - let handle = zelf.handle.load(); - // Get the lock on mmap - let mut mmap_guard = zelf.mmap.lock(); - // Check if this is a Named mmap - these cannot be resized if let Some(MmapObj::Named(_)) = mmap_guard.as_ref() { + drop(mmap_guard); return Err(vm.new_os_error("mmap: cannot resize a named memory mapping")); } @@ -1108,8 +1131,13 @@ mod mmap { let copy_size = core::cmp::min(old_size, newsize); // Create new anonymous mmap - let mut new_mmap = - host_mmap::map_anon(newsize).map_err(|e| e.to_pyexception(vm))?; + let mut new_mmap = match host_mmap::map_anon(newsize) { + Ok(mmap) => mmap, + Err(err) => { + drop(mmap_guard); + return Err(err.to_pyexception(vm)); + } + }; // Copy data from old mmap to new mmap if let Some(old_mmap) = mmap_guard.as_ref() { @@ -1134,17 +1162,23 @@ mod mmap { handle as host_mmap::Handle, zelf.size.load(), ); + drop(mmap_guard); return Err(err.to_pyexception(vm)); } // Create new mmap with the new size - let new_mmap = Self::create_mmap_windows( + let new_mmap = match Self::create_mmap_windows( handle as host_mmap::Handle, zelf.offset, newsize, &zelf.access, - ) - .map_err(|e| e.to_pyexception(vm))?; + ) { + Ok(mmap) => mmap, + Err(err) => { + drop(mmap_guard); + return Err(err.to_pyexception(vm)); + } + }; *mmap_guard = Some(new_mmap); zelf.size.store(newsize); @@ -1480,6 +1514,79 @@ mod mmap { } } + #[cfg(test)] + #[cfg(feature = "threading")] + mod tests { + use super::*; + use core::time::Duration; + use std::{sync::mpsc, thread}; + + fn rejects_borrowed_buffer(resize: bool) { + let builder = crate::vm::Interpreter::builder(Default::default()); + let defs = crate::stdlib_module_defs(&builder.ctx); + let interp = builder + .add_native_modules(&defs) + .add_frozen_modules(rustpython_pylib::FROZEN_STDLIB) + .build(); + interp.enter(|vm| { + let module = vm.import("mmap", 0).unwrap(); + let mapped = module + .get_attr("mmap", vm) + .unwrap() + .call((-1, 1), vm) + .unwrap() + .downcast::() + .unwrap(); + let buffer = PyMmap::as_buffer(&mapped, vm).unwrap(); + thread::scope(|scope| { + let (ready_tx, ready_rx) = mpsc::channel(); + let (done_tx, done_rx) = mpsc::channel(); + let buffer = &buffer; + let borrower = scope.spawn(move || { + // A native I/O operation can keep this guard while detached. + let bytes = buffer.as_contiguous().unwrap(); + assert!(bytes.is_locked()); + ready_tx.send(()).unwrap(); + // Release even on a regression, so the test fails rather than hangs. + let rejected_before_release = + done_rx.recv_timeout(Duration::from_secs(5)).is_ok(); + drop(bytes); + rejected_before_release + }); + ready_rx.recv().unwrap(); + let result = if resize { + mapped.check_resizeable(vm).map(drop) + } else { + PyMmap::close(&mapped, vm) + }; + let _ = done_tx.send(()); + assert!( + borrower.join().unwrap(), + "export rejection waited for the borrowed buffer" + ); + assert!( + result + .unwrap_err() + .class() + .is(vm.ctx.exceptions.buffer_error) + ); + }); + drop(buffer); + PyMmap::close(&mapped, vm).unwrap(); + }); + } + + #[test] + fn close_rejects_borrowed_buffer() { + rejects_borrowed_buffer(false); + } + + #[test] + fn resize_rejects_borrowed_buffer() { + rejects_borrowed_buffer(true); + } + } + impl Representable for PyMmap { #[inline] fn repr_str(zelf: &Py, _vm: &VirtualMachine) -> PyResult { diff --git a/extra_tests/snippets/stdlib_mmap.py b/extra_tests/snippets/stdlib_mmap.py index 9d4414e5552..2c4efdf58bf 100644 --- a/extra_tests/snippets/stdlib_mmap.py +++ b/extra_tests/snippets/stdlib_mmap.py @@ -39,3 +39,41 @@ mapped.find(b"") with assert_raises(ValueError): mapped.rfind(b"") + + +# Buffer acquisition must reject a closed mapping before returning a view. +with assert_raises(ValueError): + memoryview(mapped) + +mapped = mmap.mmap(-1, 4) +mapped[:] = b"data" +view = memoryview(mapped) +child = view[1:] +view.release() +with assert_raises(BufferError): + mapped.close() +assert not mapped.closed +assert child.tobytes() == b"ata" +child.release() +mapped.close() +mapped.close() + +# On Windows resize must neither invalidate an export nor revive a closed map. +import sys + +if sys.platform == "win32": + mapped = mmap.mmap(-1, 4) + mapped[:] = b"data" + with memoryview(mapped) as view: + with assert_raises(BufferError): + mapped.resize(8) + with assert_raises(BufferError): + mapped.resize(0) + assert view.tobytes() == b"data" + mapped.resize(8) + with memoryview(mapped) as view: + assert view.tobytes() == b"data" + b"\0" * 4 + mapped.close() + with assert_raises(ValueError): + mapped.resize(4) + assert mapped.closed