From a0dc7c2e13e78d11c264981b5ec29cf52ef2a4aa Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" <69878+youknowone@users.noreply.github.com> Date: Sun, 13 Sep 2026 09:41:57 +0900 Subject: [PATCH] host_env: share file identity, dll dirs, link and startfile (#8696) * host_env: share file identity, dll dirs, link and startfile Add GetFileInformationByHandle, AddDllDirectory/RemoveDllDirectory, CreateHardLinkW, ShellExecuteW, and a wide ExpandEnvironmentStringsW reader that keeps unpaired surrogates. Assisted-by: Claude * vm: use host_env for startfile, dll dirs and link Wire nt.startfile, nt._add_dll_directory/_remove_dll_directory and Windows os.link through the host_env wrappers. ExpandEnvironmentStrings keeps unpaired surrogates via the wide reader. Assisted-by: Claude:grok-4.6 * host_env: publish Windows =X: drive cwd for spawn Extract publish/ensure helpers from set_current_dir and call ensure before _wspawnv/_wexecv so CRT env walks have a first =X: entry. Assisted-by: Claude:grok-4.6 --- crates/host_env/src/nt.rs | 20 ++++++++ crates/host_env/src/os.rs | 90 +++++++++++++++++++++++++--------- crates/host_env/src/winapi.rs | 60 +++++++++++++++++++++++ crates/host_env/src/winreg.rs | 22 +++++++++ crates/vm/src/stdlib/nt.rs | 82 ++++++++++++++++++++++++++++++- crates/vm/src/stdlib/os.rs | 22 ++++++++- crates/vm/src/stdlib/winreg.rs | 10 ++-- 7 files changed, 277 insertions(+), 29 deletions(-) diff --git a/crates/host_env/src/nt.rs b/crates/host_env/src/nt.rs index 51efb6d47..b0bf5f57a 100644 --- a/crates/host_env/src/nt.rs +++ b/crates/host_env/src/nt.rs @@ -1634,6 +1634,24 @@ pub fn handle_from_fd(fd: i32) -> HANDLE { unsafe { crate::suppress_iph!(libc::get_osfhandle(fd)) as HANDLE } } +/// The triple `os._getfileinformation` / `ntpath.samefile` identify a file by. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub struct FileInformation { + pub volume_serial_number: u32, + pub file_index_high: u32, + pub file_index_low: u32, +} + +pub fn get_file_information(handle: HANDLE) -> io::Result { + let mut info: BY_HANDLE_FILE_INFORMATION = unsafe { core::mem::zeroed() }; + unsafe { GetFileInformationByHandle(handle, &mut info) }.check_win32_bool()?; + Ok(FileInformation { + volume_serial_number: info.dwVolumeSerialNumber, + file_index_high: info.nFileIndexHigh, + file_index_low: info.nFileIndexLow, + }) +} + pub fn console_type(handle: HANDLE) -> char { if is_invalid_handle(handle) { return '\0'; @@ -2106,6 +2124,7 @@ pub fn spawnv( path: &widestring::WideCStr, argv: &[&widestring::WideCStr], ) -> io::Result { + crate::os::ensure_drive_current_directory(); let argv_ptrs = null_terminated_ptrs(argv); let result = unsafe { crate::suppress_iph!(_wspawnv(mode, path.as_ptr(), argv_ptrs.as_ptr())) }; if result == -1 { @@ -2141,6 +2160,7 @@ pub fn spawnve( #[cfg(target_env = "msvc")] pub fn execv(path: &widestring::WideCStr, argv: &[&widestring::WideCStr]) -> io::Result<()> { + crate::os::ensure_drive_current_directory(); let argv_ptrs = null_terminated_ptrs(argv); let result = unsafe { crate::suppress_iph!(_wexecv(path.as_ptr(), argv_ptrs.as_ptr())) }; if result == -1 { diff --git a/crates/host_env/src/os.rs b/crates/host_env/src/os.rs index 6c2521ca1..c5c2a7a34 100644 --- a/crates/host_env/src/os.rs +++ b/crates/host_env/src/os.rs @@ -107,32 +107,76 @@ pub unsafe fn remove_var(key: impl AsRef) { unsafe { env::remove_var(key) }; } +/// Publish the working directory as `=X:` for the drive it is on. +/// +/// CRT `_wspawnve` / `_wexecve` walk the environment for the first `=X:` +/// entry with no bound; a process that was not started by cmd.exe has none +/// and that walk runs off the block. A UNC-like path is on no drive and +/// publishes nothing. +#[cfg(windows)] +pub fn publish_drive_current_directory() -> io::Result<()> { + use std::os::windows::ffi::OsStrExt; + use windows_sys::Win32::System::Environment::SetEnvironmentVariableW; + + let mut cwd_wide: Vec = env::current_dir()?.as_os_str().encode_wide().collect(); + let unc_like = cwd_wide.len() >= 2 + && ((cwd_wide[0] == b'\\' as u16 && cwd_wide[1] == b'\\' as u16) + || (cwd_wide[0] == b'/' as u16 && cwd_wide[1] == b'/' as u16)); + if unc_like || cwd_wide.is_empty() { + return Ok(()); + } + let env_name = [b'=' as u16, cwd_wide[0], b':' as u16, 0]; + cwd_wide.push(0); + let ok = unsafe { SetEnvironmentVariableW(env_name.as_ptr(), cwd_wide.as_ptr()) }; + if ok == 0 { + Err(io::Error::last_os_error()) + } else { + Ok(()) + } +} + +/// Whether the process environment already has an `=X:` drive-cwd entry. +#[cfg(windows)] +pub fn environment_has_drive_current_directory() -> bool { + use windows_sys::Win32::System::Environment::{ + FreeEnvironmentStringsW, GetEnvironmentStringsW, + }; + + let block = unsafe { GetEnvironmentStringsW() }; + if block.is_null() { + return false; + } + let mut present = false; + let mut entry = block; + unsafe { + while *entry != 0 { + if *entry == b'=' as u16 { + present = true; + break; + } + while *entry != 0 { + entry = entry.add(1); + } + entry = entry.add(1); + } + FreeEnvironmentStringsW(block); + } + present +} + +/// Publish `=X:` when the process environment has none, so a later +/// `_wspawnv` / `_wexecv` walk has a first entry. +#[cfg(windows)] +pub fn ensure_drive_current_directory() { + if !environment_has_drive_current_directory() { + let _ = publish_drive_current_directory(); + } +} + pub fn set_current_dir(path: impl AsRef) -> io::Result<()> { env::set_current_dir(&path)?; - #[cfg(windows)] - { - use std::os::windows::ffi::OsStrExt; - use windows_sys::Win32::System::Environment::SetEnvironmentVariableW; - - if let Ok(cwd) = env::current_dir() { - let cwd_str = cwd.as_os_str(); - let mut cwd_wide: Vec = cwd_str.encode_wide().collect(); - - let is_unc_like_path = cwd_wide.len() >= 2 - && ((cwd_wide[0] == b'\\' as u16 && cwd_wide[1] == b'\\' as u16) - || (cwd_wide[0] == b'/' as u16 && cwd_wide[1] == b'/' as u16)); - - if !is_unc_like_path { - let env_name: [u16; 4] = [b'=' as u16, cwd_wide[0], b':' as u16, 0]; - cwd_wide.push(0); - unsafe { - SetEnvironmentVariableW(env_name.as_ptr(), cwd_wide.as_ptr()); - } - } - } - } - + publish_drive_current_directory()?; Ok(()) } diff --git a/crates/host_env/src/winapi.rs b/crates/host_env/src/winapi.rs index 67e133878..d281b1733 100644 --- a/crates/host_env/src/winapi.rs +++ b/crates/host_env/src/winapi.rs @@ -1366,3 +1366,63 @@ pub fn need_current_directory_for_exe_path_w(exe_name: &widestring::WideCStr) -> != 0 } } + +pub type DllDirectoryCookie = *mut core::ffi::c_void; + +pub fn add_dll_directory(path: &widestring::WideCStr) -> io::Result { + let cookie = + unsafe { windows_sys::Win32::System::LibraryLoader::AddDllDirectory(path.as_ptr()) }; + if cookie.is_null() { + Err(io::Error::last_os_error()) + } else { + Ok(cookie) + } +} + +pub fn remove_dll_directory(cookie: DllDirectoryCookie) -> io::Result<()> { + unsafe { windows_sys::Win32::System::LibraryLoader::RemoveDllDirectory(cookie) } + .check_win32_bool() +} + +/// `CreateHardLinkW(new, existing)` — `dst` is the new name, `src` the file it names. +pub fn create_hard_link(dst: &widestring::WideCStr, src: &widestring::WideCStr) -> io::Result<()> { + unsafe { + windows_sys::Win32::Storage::FileSystem::CreateHardLinkW( + dst.as_ptr(), + src.as_ptr(), + core::ptr::null(), + ) + } + .check_win32_bool() +} + +/// `SW_SHOWNORMAL` — the default `os.startfile` show command. +pub const SW_SHOWNORMAL: i32 = 1; + +/// `ShellExecuteW` with a null owner window. Failure is a return of 32 or less. +pub fn shell_execute_w( + file: &widestring::WideCStr, + operation: Option<&widestring::WideCStr>, + arguments: Option<&widestring::WideCStr>, + directory: Option<&widestring::WideCStr>, + show_cmd: i32, +) -> io::Result<()> { + let as_ptr = |wide: Option<&widestring::WideCStr>| { + wide.map_or(core::ptr::null(), widestring::WideCStr::as_ptr) + }; + let rc = unsafe { + windows_sys::Win32::UI::Shell::ShellExecuteW( + core::ptr::null_mut(), + as_ptr(operation), + file.as_ptr(), + as_ptr(arguments), + as_ptr(directory), + show_cmd, + ) + }; + if rc as isize <= 32 { + Err(io::Error::last_os_error()) + } else { + Ok(()) + } +} diff --git a/crates/host_env/src/winreg.rs b/crates/host_env/src/winreg.rs index 5a67e4e2b..a8bc56b14 100644 --- a/crates/host_env/src/winreg.rs +++ b/crates/host_env/src/winreg.rs @@ -517,3 +517,25 @@ pub fn expand_environment_strings( let len = out.iter().position(|&c| c == 0).unwrap_or(out.len()); String::from_utf16(&out[..len]).map_err(ExpandEnvironmentStringsError::Utf16) } + +/// Same walk as [`expand_environment_strings`], keeping unpaired surrogates. +pub fn expand_environment_strings_wide( + wide_input: &widestring::WideCStr, +) -> Result, ExpandEnvironmentStringsError> { + let required_size = unsafe { + Environment::ExpandEnvironmentStringsW(wide_input.as_ptr(), core::ptr::null_mut(), 0) + }; + if required_size == 0 { + return Err(ExpandEnvironmentStringsError::Os); + } + let mut out = vec![0u16; required_size as usize]; + let written = unsafe { + Environment::ExpandEnvironmentStringsW(wide_input.as_ptr(), out.as_mut_ptr(), required_size) + }; + if written == 0 { + return Err(ExpandEnvironmentStringsError::Os); + } + let end = out.iter().position(|&unit| unit == 0).unwrap_or(out.len()); + out.truncate(end); + Ok(out) +} diff --git a/crates/vm/src/stdlib/nt.rs b/crates/vm/src/stdlib/nt.rs index 6a27906f2..a269646e5 100644 --- a/crates/vm/src/stdlib/nt.rs +++ b/crates/vm/src/stdlib/nt.rs @@ -6,8 +6,8 @@ pub use module::raw_set_handle_inheritable; #[pymodule(name = "nt", with(super::os::_os))] pub(crate) mod module { use crate::{ - Py, PyResult, TryFromObject, VirtualMachine, - builtins::{PyBytes, PyDictRef, PyListRef, PyStr, PyStrRef, PyTupleRef}, + AsObject, Py, PyObjectRef, PyResult, TryFromObject, VirtualMachine, + builtins::{PyBytes, PyCapsule, PyDictRef, PyListRef, PyStr, PyStrRef, PyTupleRef}, convert::ToPyException, exceptions::{self, OSErrorBuilder, ToOSErrorBuilder}, function::{ArgMapping, Either, OptionalArg}, @@ -19,6 +19,7 @@ pub(crate) mod module { use libc::intptr_t; use rustpython_common::wtf8::Wtf8Buf; use rustpython_host_env::nt as host_nt; + use rustpython_host_env::winapi as host_winapi; use std::os::windows::ffi::{OsStrExt, OsStringExt}; use std::os::windows::io::AsRawHandle; @@ -1090,4 +1091,81 @@ pub(crate) mod module { // TODO: Implement the actual logic here false } + + const DLL_DIRECTORY_COOKIE: &core::ffi::CStr = c"DLL directory cookie"; + + fn pystr_to_wide(s: &PyStrRef, vm: &VirtualMachine) -> PyResult { + widestring::WideCString::from_vec(s.as_wtf8().encode_wide().collect::>()) + .map_err(|_| vm.new_value_error("embedded null character")) + } + + #[pyfunction] + fn _add_dll_directory(path: OsPath, vm: &VirtualMachine) -> PyResult { + let wide = path.to_wide_cstring(vm)?; + let cookie = host_winapi::add_dll_directory(&wide) + .map_err(|err| OSErrorBuilder::with_filename(&err, path, vm))?; + Ok(vm + .ctx + .new_capsule(cookie, Some(DLL_DIRECTORY_COOKIE), None) + .into()) + } + + #[pyfunction] + fn _remove_dll_directory(cookie: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> { + let not_a_cookie = + || vm.new_type_error("Provided cookie was not returned from os.add_dll_directory"); + let capsule = cookie + .downcast_ref::() + .ok_or_else(not_a_cookie)?; + if capsule.name() != Some(DLL_DIRECTORY_COOKIE) || capsule.pointer().is_null() { + return Err(not_a_cookie()); + } + host_winapi::remove_dll_directory(capsule.pointer()) + .map_err(|err| err.to_pyexception(vm))?; + capsule.set_pointer(core::ptr::null_mut()); + Ok(()) + } + + #[derive(FromArgs)] + struct StartfileArgs { + #[pyarg(any)] + filepath: OsPath, + #[pyarg(any, default)] + operation: Option, + #[pyarg(any, default)] + arguments: Option, + #[pyarg(any, default)] + cwd: Option, + #[pyarg(any, default)] + show_cmd: Option, + } + + #[pyfunction] + fn startfile(args: StartfileArgs, vm: &VirtualMachine) -> PyResult<()> { + let file = args.filepath.to_wide_cstring(vm)?; + let operation = args + .operation + .as_ref() + .map(|s| pystr_to_wide(s, vm)) + .transpose()?; + let arguments = args + .arguments + .as_ref() + .map(|s| pystr_to_wide(s, vm)) + .transpose()?; + let directory = args + .cwd + .as_ref() + .map(|p| p.to_wide_cstring(vm)) + .transpose()?; + let show_cmd = args.show_cmd.unwrap_or(host_winapi::SW_SHOWNORMAL); + host_winapi::shell_execute_w( + &file, + operation.as_deref(), + arguments.as_deref(), + directory.as_deref(), + show_cmd, + ) + .map_err(|err| OSErrorBuilder::with_filename(&err, args.filepath, vm)) + } } diff --git a/crates/vm/src/stdlib/os.rs b/crates/vm/src/stdlib/os.rs index f63ed3974..ee38693b9 100644 --- a/crates/vm/src/stdlib/os.rs +++ b/crates/vm/src/stdlib/os.rs @@ -1634,7 +1634,7 @@ pub(super) mod _os { Ok(()) } - #[cfg(not(unix))] + #[cfg(windows)] { let src_path = match follow_symlinks.into_option() { Some(true) => { @@ -1647,6 +1647,26 @@ pub(super) mod _os { PathBuf::from(src.path.clone()) } }; + // CreateHardLinkW(new, existing) + let src_wide = src_path + .to_wide_cstring() + .map_err(|err| err.into_pyexception(vm))?; + let dst_wide = dst.to_wide_cstring(vm)?; + rustpython_host_env::winapi::create_hard_link(&dst_wide, &src_wide).map_err(|err| { + let builder = err.to_os_error_builder(vm); + let builder = builder.filename(src.filename(vm)); + let builder = builder.filename2(dst.filename(vm)); + builder.build(vm).upcast() + }) + } + + #[cfg(not(any(unix, windows)))] + { + let src_path = match follow_symlinks.into_option() { + Some(true) => crate::host_env::fs::canonicalize(&src.path) + .unwrap_or_else(|_| PathBuf::from(src.path.clone())), + Some(false) | None => PathBuf::from(src.path.clone()), + }; fs::hard_link(&src_path, &dst.path).map_err(|err| { let builder = err.to_os_error_builder(vm); diff --git a/crates/vm/src/stdlib/winreg.rs b/crates/vm/src/stdlib/winreg.rs index df3959acd..179143759 100644 --- a/crates/vm/src/stdlib/winreg.rs +++ b/crates/vm/src/stdlib/winreg.rs @@ -18,6 +18,7 @@ mod winreg { use crossbeam_utils::atomic::AtomicCell; use malachite_bigint::Sign; use num_traits::ToPrimitive; + use rustpython_common::wtf8::Wtf8Buf; use rustpython_host_env::windows::ToWideString; use rustpython_host_env::winreg as host_winreg; use widestring::{WideCString, WideString}; @@ -869,8 +870,11 @@ mod winreg { } #[pyfunction] - fn ExpandEnvironmentStrings(i: String, vm: &VirtualMachine) -> PyResult { - let i = WideCString::from_str(i).map_err(|err| err.to_pyexception(vm))?; - host_winreg::expand_environment_strings(&i).map_err(|err| err.to_pyexception(vm)) + fn ExpandEnvironmentStrings(i: PyRef, vm: &VirtualMachine) -> PyResult> { + let i = WideCString::from_vec(i.as_wtf8().encode_wide().collect::>()) + .map_err(|_| nul_char_error(vm))?; + let expanded = host_winreg::expand_environment_strings_wide(&i) + .map_err(|err| err.to_pyexception(vm))?; + Ok(vm.ctx.new_str(Wtf8Buf::from_wide(&expanded))) } }