diff --git a/Lib/test/test_os.py b/Lib/test/test_os.py index 28b0874658d..25e2c30e4d7 100644 --- a/Lib/test/test_os.py +++ b/Lib/test/test_os.py @@ -1152,7 +1152,6 @@ def test_putenv_unsetenv(self): self.assertEqual(proc.stdout.rstrip(), repr(None)) # On OS X < 10.6, unsetenv() doesn't return a value (bpo-13415). - @unittest.expectedFailureIfWindows('TODO: RUSTPYTHON; (AssertionError: ValueError not raised by putenv)') @support.requires_mac_ver(10, 6) def test_putenv_unsetenv_error(self): # Empty variable name is invalid. diff --git a/crates/common/src/windows.rs b/crates/common/src/windows.rs index 4a922ce4354..dd4ea82c874 100644 --- a/crates/common/src/windows.rs +++ b/crates/common/src/windows.rs @@ -3,6 +3,9 @@ use std::{ os::windows::ffi::{OsStrExt, OsStringExt}, }; +/// _MAX_ENV from Windows CRT stdlib.h - maximum environment variable size +pub const _MAX_ENV: usize = 32767; + pub trait ToWideString { fn to_wide(&self) -> Vec; fn to_wide_with_nul(&self) -> Vec; diff --git a/crates/vm/src/stdlib/nt.rs b/crates/vm/src/stdlib/nt.rs index c041dab9aa5..0bbb5881a9f 100644 --- a/crates/vm/src/stdlib/nt.rs +++ b/crates/vm/src/stdlib/nt.rs @@ -682,6 +682,87 @@ pub(crate) mod module { Ok(vm.ctx.new_list(drives)) } + #[pyfunction] + fn listvolumes(vm: &VirtualMachine) -> PyResult { + use windows_sys::Win32::Foundation::ERROR_NO_MORE_FILES; + + let mut result = Vec::new(); + let mut buffer = [0u16; Foundation::MAX_PATH as usize + 1]; + + let find = unsafe { FileSystem::FindFirstVolumeW(buffer.as_mut_ptr(), buffer.len() as _) }; + if find == INVALID_HANDLE_VALUE { + return Err(errno_err(vm)); + } + + loop { + // Find the null terminator + let len = buffer.iter().position(|&c| c == 0).unwrap_or(buffer.len()); + let volume = String::from_utf16_lossy(&buffer[..len]); + result.push(vm.new_pyobj(volume)); + + let ret = unsafe { + FileSystem::FindNextVolumeW(find, buffer.as_mut_ptr(), buffer.len() as _) + }; + if ret == 0 { + let err = io::Error::last_os_error(); + unsafe { FileSystem::FindVolumeClose(find) }; + if err.raw_os_error() == Some(ERROR_NO_MORE_FILES as i32) { + break; + } + return Err(err.to_pyexception(vm)); + } + } + + Ok(vm.ctx.new_list(result)) + } + + #[pyfunction] + fn listmounts(volume: OsPath, vm: &VirtualMachine) -> PyResult { + use windows_sys::Win32::Foundation::ERROR_MORE_DATA; + + let wide = volume.to_wide_cstring(vm)?; + let mut buflen: u32 = Foundation::MAX_PATH + 1; + let mut buffer: Vec = vec![0; buflen as usize]; + + loop { + let success = unsafe { + FileSystem::GetVolumePathNamesForVolumeNameW( + wide.as_ptr(), + buffer.as_mut_ptr(), + buflen, + &mut buflen, + ) + }; + if success != 0 { + break; + } + let err = io::Error::last_os_error(); + if err.raw_os_error() == Some(ERROR_MORE_DATA as i32) { + buffer.resize(buflen as usize, 0); + continue; + } + return Err(err.to_pyexception(vm)); + } + + // Parse null-separated strings + let mut result = Vec::new(); + let mut start = 0; + for (i, &c) in buffer.iter().enumerate() { + if c == 0 { + if i > start { + let mount = String::from_utf16_lossy(&buffer[start..i]); + result.push(vm.new_pyobj(mount)); + } + start = i + 1; + if start < buffer.len() && buffer[start] == 0 { + break; // Double null = end + } + } + } + + Ok(vm.ctx.new_list(result)) + } + #[pyfunction] fn set_handle_inheritable( handle: intptr_t, diff --git a/crates/vm/src/stdlib/os.rs b/crates/vm/src/stdlib/os.rs index f484f69c06d..e9e1337235b 100644 --- a/crates/vm/src/stdlib/os.rs +++ b/crates/vm/src/stdlib/os.rs @@ -428,6 +428,20 @@ pub(super) mod _os { } } + /// Check if environment variable length exceeds Windows limit. + /// size should be key.len() + value.len() + 2 (for '=' and null terminator) + #[cfg(windows)] + fn check_env_var_len(size: usize, vm: &VirtualMachine) -> PyResult<()> { + use crate::common::windows::_MAX_ENV; + if size > _MAX_ENV { + return Err(vm.new_value_error(format!( + "the environment variable is longer than {} characters", + _MAX_ENV + ))); + } + Ok(()) + } + #[pyfunction] fn putenv( key: Either, @@ -442,6 +456,8 @@ pub(super) mod _os { if key.is_empty() || key.contains(&b'=') { return Err(vm.new_value_error("illegal environment variable name")); } + #[cfg(windows)] + check_env_var_len(key.len() + value.len() + 2, vm)?; let key = super::bytes_as_os_str(key, vm)?; let value = super::bytes_as_os_str(value, vm)?; // SAFETY: requirements forwarded from the caller @@ -464,6 +480,9 @@ pub(super) mod _os { ), )); } + // For unsetenv, size is key + '=' (no value, just clearing) + #[cfg(windows)] + check_env_var_len(key.len() + 1, vm)?; let key = super::bytes_as_os_str(key, vm)?; // SAFETY: requirements forwarded from the caller unsafe { env::remove_var(key) }; @@ -507,17 +526,17 @@ pub(super) mod _os { Ok(self.mode.process_path(&self.pathval, vm)) } - fn perform_on_metadata( - &self, - follow_symlinks: FollowSymlinks, - action: fn(fs::Metadata) -> bool, - vm: &VirtualMachine, - ) -> PyResult { + #[pymethod] + fn is_dir(&self, follow_symlinks: FollowSymlinks, vm: &VirtualMachine) -> PyResult { match super::fs_metadata(&self.pathval, follow_symlinks.0) { - Ok(meta) => Ok(action(meta)), + Ok(meta) => Ok(meta.is_dir()), Err(e) => { - // FileNotFoundError is caught and not raised if e.kind() == io::ErrorKind::NotFound { + // On Windows, use cached file_type when file is removed + #[cfg(windows)] + if let Ok(file_type) = &self.file_type { + return Ok(file_type.is_dir()); + } Ok(false) } else { Err(e.into_pyexception(vm)) @@ -526,22 +545,23 @@ pub(super) mod _os { } } - #[pymethod] - fn is_dir(&self, follow_symlinks: FollowSymlinks, vm: &VirtualMachine) -> PyResult { - self.perform_on_metadata( - follow_symlinks, - |meta: fs::Metadata| -> bool { meta.is_dir() }, - vm, - ) - } - #[pymethod] fn is_file(&self, follow_symlinks: FollowSymlinks, vm: &VirtualMachine) -> PyResult { - self.perform_on_metadata( - follow_symlinks, - |meta: fs::Metadata| -> bool { meta.is_file() }, - vm, - ) + match super::fs_metadata(&self.pathval, follow_symlinks.0) { + Ok(meta) => Ok(meta.is_file()), + Err(e) => { + if e.kind() == io::ErrorKind::NotFound { + // On Windows, use cached file_type when file is removed + #[cfg(windows)] + if let Ok(file_type) = &self.file_type { + return Ok(file_type.is_file()); + } + Ok(false) + } else { + Err(e.into_pyexception(vm)) + } + } + } } #[pymethod]