diff options
Diffstat (limited to 'src')
| -rw-r--r-- | src/platform/mod.rs | 5 | ||||
| -rw-r--r-- | src/platform/windows.rs | 248 | ||||
| -rw-r--r-- | src/server/autodetect.rs | 9 |
3 files changed, 253 insertions, 9 deletions
diff --git a/src/platform/mod.rs b/src/platform/mod.rs index ba92ca63..613bff1a 100644 --- a/src/platform/mod.rs +++ b/src/platform/mod.rs @@ -59,6 +59,11 @@ pub(crate) const fn capabilities() -> PlatformCapabilities { } } +#[cfg(not(windows))] +pub fn launch_server_daemon_command(command: &mut std::process::Command) -> std::io::Result<u32> { + command.spawn().map(|child| child.id()) +} + #[cfg(any(target_os = "linux", target_os = "macos"))] pub fn detach_server_daemon_command(command: &mut std::process::Command) { use std::os::unix::process::CommandExt; diff --git a/src/platform/windows.rs b/src/platform/windows.rs index 48b31c89..35ecfd82 100644 --- a/src/platform/windows.rs +++ b/src/platform/windows.rs @@ -1,6 +1,7 @@ use std::{ + cmp::Ordering, collections::{HashMap, HashSet, VecDeque}, - ffi::c_void, + ffi::{c_void, OsStr}, mem::{size_of, MaybeUninit}, path::PathBuf, ptr::{copy_nonoverlapping, null_mut}, @@ -15,6 +16,7 @@ use windows_sys::{ CloseHandle, GlobalFree, LocalFree, HANDLE, HWND, INVALID_HANDLE_VALUE, NTSTATUS, STATUS_SUCCESS, UNICODE_STRING, }, + Globalization::{CompareStringOrdinal, CSTR_EQUAL, CSTR_GREATER_THAN, CSTR_LESS_THAN}, System::{ Console::GetConsoleWindow, DataExchange::{CloseClipboard, EmptyClipboard, OpenClipboard, SetClipboardData}, @@ -25,7 +27,10 @@ use windows_sys::{ TH32CS_SNAPPROCESS, }, }, - JobObjects::IsProcessInJob, + JobObjects::{ + IsProcessInJob, JobObjectExtendedLimitInformation, QueryInformationJobObject, + JOBOBJECT_EXTENDED_LIMIT_INFORMATION, JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE, + }, Memory::{GlobalAlloc, GlobalLock, GlobalUnlock, GMEM_MOVEABLE}, Ole::CF_UNICODETEXT, Threading::{ @@ -248,6 +253,178 @@ pub(crate) fn configure_background_command_platform(command: &mut std::process:: command.creation_flags(CREATE_NO_WINDOW); } +pub fn launch_server_daemon_command(command: &mut std::process::Command) -> std::io::Result<u32> { + if current_job_kills_processes_on_close()? { + launch_server_daemon_with_wmi(command) + } else { + command.spawn().map(|child| child.id()) + } +} + +fn launch_server_daemon_with_wmi(command: &std::process::Command) -> std::io::Result<u32> { + // WMI resolves the class from this Rust type name, including CIM casing. + #[allow(non_camel_case_types)] + #[derive(serde::Deserialize)] + struct Win32_Process; + + // WMI serializes this embedded object using the matching CIM class name. + #[allow(non_camel_case_types)] + #[derive(serde::Serialize)] + struct Win32_ProcessStartup { + #[serde(rename = "CreateFlags")] + create_flags: u32, + #[serde(rename = "EnvironmentVariables")] + environment_variables: Vec<String>, + } + + #[derive(serde::Serialize)] + struct CreateInput { + #[serde(rename = "CommandLine")] + command_line: String, + #[serde(rename = "CurrentDirectory")] + current_directory: String, + #[serde(rename = "ProcessStartupInformation")] + process_startup_information: Win32_ProcessStartup, + } + + #[derive(serde::Deserialize)] + struct CreateOutput { + #[serde(rename = "ProcessId")] + process_id: Option<u32>, + #[serde(rename = "ReturnValue")] + return_value: u32, + } + + let current_directory = command + .get_current_dir() + .map(std::path::Path::to_path_buf) + .map(Ok) + .unwrap_or_else(std::env::current_dir)?; + let input = CreateInput { + command_line: windows_command_line(command)?, + current_directory: unicode_windows_value( + ¤t_directory.into_os_string(), + "working directory", + )?, + process_startup_information: Win32_ProcessStartup { + create_flags: DETACHED_PROCESS, + environment_variables: effective_command_environment(command)?, + }, + }; + + let connection = wmi::WMIConnection::new() + .map_err(|err| std::io::Error::other(format!("failed to connect to WMI: {err}")))?; + let output: CreateOutput = connection + .exec_class_method::<Win32_Process, _>("Create", &input) + .map_err(|err| std::io::Error::other(format!("WMI Win32_Process.Create failed: {err}")))?; + if output.return_value != 0 { + return Err(std::io::Error::other(format!( + "WMI Win32_Process.Create returned error {}", + output.return_value + ))); + } + output.process_id.ok_or_else(|| { + std::io::Error::other("WMI Win32_Process.Create succeeded without a process id") + }) +} + +fn windows_command_line(command: &std::process::Command) -> std::io::Result<String> { + std::iter::once(command.get_program()) + .chain(command.get_args()) + .map(|value| { + unicode_windows_value(value, "server command argument") + .map(|value| quote_windows_command_line_arg(&value)) + }) + .collect::<std::io::Result<Vec<_>>>() + .map(|parts| parts.join(" ")) +} + +fn effective_command_environment(command: &std::process::Command) -> std::io::Result<Vec<String>> { + let mut environment = std::env::vars_os() + .map(|(key, value)| { + Ok(( + unicode_windows_value(&key, "inherited environment variable name")?, + unicode_windows_value(&value, "inherited environment variable value")?, + )) + }) + .collect::<std::io::Result<Vec<(String, String)>>>()?; + for (key, value) in command.get_envs() { + let key = unicode_windows_value(key, "environment variable name")?; + environment.retain(|(inherited, _)| windows_environment_key_cmp(inherited, &key).is_ne()); + if let Some(value) = value { + environment.push(( + key, + unicode_windows_value(value, "environment variable value")?, + )); + } + } + environment.sort_unstable_by(|(left, _), (right, _)| windows_environment_key_cmp(left, right)); + Ok(environment + .into_iter() + .map(|(key, value)| format!("{key}={value}")) + .collect()) +} + +fn windows_environment_key_cmp(left: &str, right: &str) -> Ordering { + let left_wide: Vec<u16> = left.encode_utf16().collect(); + let right_wide: Vec<u16> = right.encode_utf16().collect(); + // SAFETY: both pointers remain valid for the call and lengths count UTF-16 units. + match unsafe { + CompareStringOrdinal( + left_wide.as_ptr(), + left_wide.len() as i32, + right_wide.as_ptr(), + right_wide.len() as i32, + 1, + ) + } { + CSTR_LESS_THAN => Ordering::Less, + CSTR_EQUAL => Ordering::Equal, + CSTR_GREATER_THAN => Ordering::Greater, + _ => left.cmp(right), + } +} + +fn unicode_windows_value(value: &OsStr, label: &str) -> std::io::Result<String> { + value.to_str().map(str::to_owned).ok_or_else(|| { + std::io::Error::new( + std::io::ErrorKind::InvalidData, + format!("{label} is not valid Unicode"), + ) + }) +} + +fn current_process_is_in_job() -> std::io::Result<bool> { + let mut in_job = 0; + // SAFETY: `in_job` is a valid writable BOOL for the duration of the call. + if unsafe { IsProcessInJob(GetCurrentProcess(), null_mut(), &mut in_job) } == 0 { + return Err(std::io::Error::last_os_error()); + } + Ok(in_job != 0) +} + +fn current_job_kills_processes_on_close() -> std::io::Result<bool> { + if !current_process_is_in_job()? { + return Ok(false); + } + + let mut limits = JOBOBJECT_EXTENDED_LIMIT_INFORMATION::default(); + // SAFETY: `limits` is writable and its exact buffer size is supplied. + if unsafe { + QueryInformationJobObject( + null_mut(), + JobObjectExtendedLimitInformation, + &mut limits as *mut _ as *mut c_void, + size_of::<JOBOBJECT_EXTENDED_LIMIT_INFORMATION>() as u32, + null_mut(), + ) + } == 0 + { + return Err(std::io::Error::last_os_error()); + } + Ok(limits.BasicLimitInformation.LimitFlags & JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE != 0) +} + pub fn detach_server_daemon_command(command: &mut std::process::Command) { use std::os::windows::process::CommandExt; @@ -259,8 +436,7 @@ pub fn current_process_is_detached_server_daemon() -> bool { return false; } - let mut in_job = 0; - unsafe { IsProcessInJob(GetCurrentProcess(), null_mut(), &mut in_job) != 0 && in_job == 0 } + matches!(current_process_is_in_job(), Ok(false)) } pub fn foreground_job(child_pid: u32) -> Option<ForegroundJob> { @@ -1333,6 +1509,70 @@ mod tests { const CONSOLE_TEST_CHILD_ENV: &str = "HERDR_TEST_CONSOLE_CHILD_MODE"; const CONSOLE_TEST_PARENT_PID_ENV: &str = "HERDR_TEST_CONSOLE_PARENT_PID"; + const WMI_DAEMON_TEST_CHILD_ENV: &str = "HERDR_TEST_WMI_DAEMON_CHILD"; + + #[test] + fn windows_environment_keys_use_unicode_case_insensitive_ordering() { + assert_eq!( + super::windows_environment_key_cmp("hérdr", "HÉRDR"), + std::cmp::Ordering::Equal + ); + } + + #[test] + fn windows_wmi_daemon_preserves_environment_and_working_directory() { + if let Some(capture) = std::env::var_os(WMI_DAEMON_TEST_CHILD_ENV) { + let cwd = std::env::current_dir().expect("WMI daemon test working directory"); + fs::write( + capture, + format!( + "{}\n{}", + cwd.display(), + super::current_process_is_detached_server_daemon() + ), + ) + .expect("write WMI daemon test capture"); + return; + } + + let base = std::env::temp_dir().join(format!( + "herdr-wmi-daemon-test-{}-{}", + std::process::id(), + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_millis() + )); + fs::create_dir_all(&base).unwrap(); + let capture = base.join("capture.txt"); + let test_exe = std::env::current_exe().expect("resolve test executable"); + let mut child = Command::new(test_exe); + child + .arg("windows_wmi_daemon_preserves_environment_and_working_directory") + .current_dir(&base) + .env(WMI_DAEMON_TEST_CHILD_ENV, &capture) + .stdin(Stdio::null()) + .stdout(Stdio::null()) + .stderr(Stdio::null()); + + let pid = super::launch_server_daemon_with_wmi(&child) + .expect("launch detached process through WMI"); + assert_ne!(pid, 0, "WMI returned an invalid process id"); + + let expected = format!("{}\ntrue", base.display()); + let deadline = Instant::now() + Duration::from_secs(10); + loop { + if fs::read_to_string(&capture).is_ok_and(|captured| captured == expected) { + break; + } + assert!( + Instant::now() < deadline, + "WMI daemon child did not write the expected capture" + ); + thread::sleep(Duration::from_millis(50)); + } + let _ = fs::remove_dir_all(base); + } fn console_process_ids() -> Vec<u32> { let mut process_ids = vec![0; 8]; diff --git a/src/server/autodetect.rs b/src/server/autodetect.rs index 4b58196e..337c388f 100644 --- a/src/server/autodetect.rs +++ b/src/server/autodetect.rs @@ -198,11 +198,10 @@ pub fn spawn_server_daemon() -> io::Result<u32> { let mut command = build_server_daemon_command(exe); - let child = command.spawn().map_err(|err: io::Error| { - io::Error::new(err.kind(), format!("failed to spawn herdr server: {err}")) - })?; - - let pid = child.id(); + let pid = + crate::platform::launch_server_daemon_command(&mut command).map_err(|err: io::Error| { + io::Error::new(err.kind(), format!("failed to spawn herdr server: {err}")) + })?; info!(pid, "server daemon spawned"); Ok(pid) |