Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 0 additions & 1 deletion Lib/test/test_os.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
3 changes: 3 additions & 0 deletions crates/common/src/windows.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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<u16>;
fn to_wide_with_nul(&self) -> Vec<u16>;
Expand Down
81 changes: 81 additions & 0 deletions crates/vm/src/stdlib/nt.rs
Original file line number Diff line number Diff line change
Expand Up @@ -682,6 +682,87 @@ pub(crate) mod module {
Ok(vm.ctx.new_list(drives))
}

#[pyfunction]
fn listvolumes(vm: &VirtualMachine) -> PyResult<PyListRef> {
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<PyListRef> {
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<u16> = 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,
Expand Down
64 changes: 42 additions & 22 deletions crates/vm/src/stdlib/os.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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<PyStrRef, PyBytesRef>,
Expand All @@ -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
Expand All @@ -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)?;
Comment on lines +483 to +485
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟡 Minor

Potential off-by-one: missing null terminator in size calculation.

The comment states "size is key + '='" but Windows _MAX_ENV includes the null terminator. For unsetenv, Windows internally sets "KEY=\0", so the size should be key.len() + 2 (key + '=' + '\0'), not key.len() + 1.

         // For unsetenv, size is key + '=' (no value, just clearing)
         #[cfg(windows)]
-        check_env_var_len(key.len() + 1, vm)?;
+        check_env_var_len(key.len() + 2, vm)?;  // key + '=' + '\0'
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
// For unsetenv, size is key + '=' (no value, just clearing)
#[cfg(windows)]
check_env_var_len(key.len() + 1, vm)?;
// For unsetenv, size is key + '=' (no value, just clearing)
#[cfg(windows)]
check_env_var_len(key.len() + 2, vm)?; // key + '=' + '\0'
🤖 Prompt for AI Agents
In crates/vm/src/stdlib/os.rs around lines 483 to 485, the size computed for
unsetenv on Windows omits the null terminator; change the Windows branch to pass
key.len() + 2 (key + '=' + '\0') into check_env_var_len instead of key.len() + 1
so the length accounts for the terminating NUL when validating against _MAX_ENV.

let key = super::bytes_as_os_str(key, vm)?;
// SAFETY: requirements forwarded from the caller
unsafe { env::remove_var(key) };
Expand Down Expand Up @@ -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<bool> {
#[pymethod]
fn is_dir(&self, follow_symlinks: FollowSymlinks, vm: &VirtualMachine) -> PyResult<bool> {
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))
Expand All @@ -526,22 +545,23 @@ pub(super) mod _os {
}
}

#[pymethod]
fn is_dir(&self, follow_symlinks: FollowSymlinks, vm: &VirtualMachine) -> PyResult<bool> {
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<bool> {
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]
Expand Down
Loading