WIP sorting new process stuff

This commit is contained in:
Caelan Sayler
2024-12-03 21:10:41 +00:00
committed by Kevin Bost
parent 1079074178
commit 14b7119637
5 changed files with 222 additions and 116 deletions

1
Cargo.lock generated
View File

@@ -2277,7 +2277,6 @@ dependencies = [
"tempfile",
"time 0.3.41",
"velopack",
"wait-timeout",
"waitpid-any",
"webview2-com-sys",
"windows",

View File

@@ -59,7 +59,6 @@ regex.workspace = true
normpath.workspace = true
simple-stopwatch.workspace = true
file-rotate.workspace = true
wait-timeout.workspace = true
pretty-bytes-rust.workspace = true
enum-flags.workspace = true
log-panics.workspace = true

View File

@@ -1,5 +1,4 @@
use ::windows::Win32::System::ProcessStatus::EnumProcesses;
use ::windows::Win32::UI::WindowsAndMessaging::AllowSetForegroundWindow;
use anyhow::{anyhow, bail, Result};
use regex::Regex;
use semver::Version;
@@ -7,23 +6,14 @@ use std::{
collections::HashMap,
fs,
path::{Path, PathBuf},
process::Command as Process,
};
use windows::{Wdk::System::Threading::{NtQueryInformationProcess, ProcessBasicInformation}, Win32::Foundation::HANDLE};
use windows::Wdk::System::Threading::{NtQueryInformationProcess, ProcessBasicInformation};
use windows::Win32::System::Threading::{GetCurrentProcess, PROCESS_BASIC_INFORMATION};
use winsafe::{self as w, co, prelude::*};
use velopack::locator::VelopackLocator;
pub fn wait_for_handle_to_exit(handle: HANDLE, ms_to_wait: Option<u32>) -> Result<()> {
let handle = unsafe { w::HPROCESS::from_ptr(handle.0) };
info!("Waiting {:?}ms for handle to exit.", ms_to_wait);
match handle.WaitForSingleObject(ms_to_wait) {
Ok(co::WAIT::OBJECT_0) => Ok(()),
// Ok(co::WAIT::TIMEOUT) => Ok(()),
_ => Err(anyhow!("WaitForSingleObject Failed.")),
}
}
use crate::windows::process;
pub fn wait_for_pid_to_exit(pid: u32, ms_to_wait: u32) -> Result<()> {
info!("Waiting {}ms for process ({}) to exit.", ms_to_wait, pid);
@@ -180,28 +170,18 @@ fn _force_stop_package<P: AsRef<Path>>(root_dir: P) -> Result<()> {
}
pub fn start_package(locator: &VelopackLocator, exe_args: Option<Vec<&str>>, set_env: Option<&str>) -> Result<()> {
let current = locator.get_current_bin_dir();
let current = locator.get_current_bin_dir_as_string();
let exe_to_execute = locator.get_main_exe_path();
if !exe_to_execute.exists() {
bail!("Unable to find executable to start: '{}'", exe_to_execute.to_string_lossy());
}
let mut psi = Process::new(&exe_to_execute);
psi.current_dir(&current);
if let Some(args) = exe_args {
psi.args(args);
}
if let Some(env) = set_env {
debug!("Setting environment variable: {}={}", env, "true");
psi.env(env, "true");
}
info!("About to launch: '{:?}' in dir '{:?}'", exe_to_execute, current);
info!("Args: {:?}", psi.get_args());
let child = psi.spawn().map_err(|z| anyhow!("Failed to start application ({}).", z))?;
let _ = unsafe { AllowSetForegroundWindow(child.id()) };
let envp = set_env.map(|f| HashMap::from([(f.to_string(), "true".to_string())]));
let args = exe_args.map(|opt| opt.iter().map(|f| f.to_string()).collect()).unwrap_or(Vec::new());
info!("About to launch: '{:?}' in dir '{:?}' with args: {:?}", exe_to_execute, current, args);
process::run_process(exe_to_execute.to_string_lossy().to_string(), args, Some(current), envp, true)?;
Ok(())
}
@@ -258,10 +238,10 @@ pub fn get_latest_app_version_folder<P: AsRef<Path>>(parent_path: P) -> Result<O
pub fn has_app_prefixed_folder<P: AsRef<Path>>(parent_path: P) -> bool {
match get_app_prefixed_folders(parent_path) {
Ok(folders) => !folders.is_empty(),
Err(e) => {
Err(e) => {
warn!("Failed to check for app-prefixed folders: {}", e);
false
},
}
}
}

View File

@@ -1,15 +1,24 @@
use std::{
collections::HashMap,
ffi::{OsStr, OsString},
os::windows::ffi::OsStrExt,
os::{raw::c_void, windows::ffi::OsStrExt},
time::Duration,
};
use anyhow::{bail, Result};
use windows::{
core::PCWSTR,
core::{PCWSTR, PWSTR},
Win32::{
Foundation::HANDLE,
System::Threading::CreateProcessW,
UI::Shell::{ShellExecuteExW, SEE_MASK_NOCLOSEPROCESS, SHELLEXECUTEINFOW},
Foundation::{CloseHandle, HANDLE, WAIT_OBJECT_0, WAIT_TIMEOUT},
Security::{GetTokenInformation, TokenElevation, TOKEN_ELEVATION},
System::Threading::{
CreateProcessW, GetCurrentProcess, GetExitCodeProcess, GetProcessId, OpenProcessToken, WaitForSingleObject, CREATE_NO_WINDOW,
PROCESS_CREATION_FLAGS, STARTUPINFOW, STARTUPINFOW_FLAGS,
},
UI::{
Shell::{ShellExecuteExW, SEE_MASK_NOCLOSEPROCESS, SHELLEXECUTEINFOW},
WindowsAndMessaging::AllowSetForegroundWindow,
},
},
};
@@ -104,10 +113,108 @@ fn make_command_line(argv0: Option<&OsStr>, args: &[Arg], force_quotes: bool) ->
cmd.push(' ' as u16);
}
cmd.push(0);
Ok(cmd)
}
pub fn run_process_as_admin(exe_path: String, args: Vec<String>, work_dir: Option<String>) -> Result<HANDLE> {
fn make_envp(maybe_env: Option<HashMap<String, String>>) -> Result<(Option<*const c_void>, Vec<u16>)> {
// On Windows we pass an "environment block" which is not a char**, but
// rather a concatenation of null-terminated k=v\0 sequences, with a final
// \0 to terminate.
if let Some(env) = maybe_env {
let mut blk = Vec::new();
// If there are no environment variables to set then signal this by
// pushing a null.
if env.is_empty() {
blk.push(0);
}
for (k, v) in env {
let os_key = OsString::from(k);
let os_value = OsString::from(v);
blk.extend(ensure_no_nuls(os_key)?.encode_wide());
blk.push('=' as u16);
blk.extend(ensure_no_nuls(os_value)?.encode_wide());
blk.push(0);
}
blk.push(0);
Ok((Some(blk.as_ptr() as *mut c_void), blk))
} else {
Ok((None, Vec::new()))
}
}
fn make_dirp(d: Option<String>) -> Result<(PCWSTR, Vec<u16>)> {
match d {
Some(dir) => {
let dir = OsString::from(dir);
let mut dir_str: Vec<u16> = ensure_no_nuls(dir)?.encode_wide().collect();
dir_str.push(0);
Ok((PCWSTR(dir_str.as_ptr()), dir_str))
}
None => Ok((PCWSTR::null(), Vec::new())),
}
}
pub fn is_process_elevated() -> bool {
// Get the current process handle
let process = unsafe { GetCurrentProcess() };
// Variable to hold the process token
let mut token: HANDLE = HANDLE::default();
// Open the process token with the TOKEN_QUERY access rights
unsafe {
if OpenProcessToken(process, windows::Win32::Security::TOKEN_QUERY, &mut token).is_ok() {
// Allocate a buffer for the TOKEN_ELEVATION structure
let mut elevation = TOKEN_ELEVATION::default();
let mut size: u32 = 0;
let elevation_ptr: *mut core::ffi::c_void = &mut elevation as *mut _ as *mut _;
// Query the token information to check if it is elevated
if GetTokenInformation(token, TokenElevation, Some(elevation_ptr), std::mem::size_of::<TOKEN_ELEVATION>() as u32, &mut size)
.is_ok()
{
// Return whether the token is elevated
CloseHandle(token);
return elevation.TokenIsElevated != 0;
}
}
}
// Clean up the token handle
if !token.is_invalid() {
unsafe { CloseHandle(token) };
}
false
}
struct SafeProcessHandle(HANDLE);
impl Drop for SafeProcessHandle {
fn drop(&mut self) {
if !self.0.is_invalid() {
let _ = unsafe { CloseHandle(self.0) };
}
}
}
impl SafeProcessHandle {
pub fn handle(&self) -> HANDLE {
self.0
}
}
// impl Into<u32> for SafeProcessHandle {
// fn into(self) -> u32 {
// self.1
// }
// }
pub fn run_process_as_admin(exe_path: String, args: Vec<String>, work_dir: Option<String>) -> Result<SafeProcessHandle> {
let verb = string_to_u16("runas");
let verb = PCWSTR(verb.as_ptr());
@@ -134,21 +241,87 @@ pub fn run_process_as_admin(exe_path: String, args: Vec<String>, work_dir: Optio
unsafe {
ShellExecuteExW(&mut exe_info as *mut SHELLEXECUTEINFOW)?;
let process_id = GetProcessId(exe_info.hProcess);
let _ = AllowSetForegroundWindow(process_id);
Ok(SafeProcessHandle(exe_info.hProcess))
}
Ok(exe_info.hProcess)
}
// pub fn run_process_no_window(exe_path: String, args: Vec<String>, work_dir: Option<String>) -> Result<HANDLE> {
// CreateProcessW(
// lpapplicationname,
// lpcommandline,
// lpprocessattributes,
// lpthreadattributes,
// binherithandles,
// dwcreationflags,
// lpenvironment,
// lpcurrentdirectory,
// lpstartupinfo,
// lpprocessinformation,
// )
// }
pub fn run_process(
exe_path: String,
args: Vec<String>,
work_dir: Option<String>,
set_env: Option<HashMap<String, String>>,
show_window: bool,
) -> Result<SafeProcessHandle> {
let exe_path = OsString::from(exe_path);
let exe_name = PCWSTR(exe_path.encode_wide().chain(Some(0)).collect::<Vec<_>>().as_mut_ptr());
let args: Vec<Arg> = args.iter().map(|a| Arg::Regular(a.into())).collect();
let mut params = make_command_line(Some(&exe_path), &args, false)?;
let params = PWSTR(params.as_mut_ptr());
let mut pi = windows::Win32::System::Threading::PROCESS_INFORMATION::default();
let si = STARTUPINFOW {
cb: std::mem::size_of::<STARTUPINFOW>() as u32,
lpReserved: PWSTR::null(),
lpDesktop: PWSTR::null(),
lpTitle: PWSTR::null(),
dwX: 0,
dwY: 0,
dwXSize: 0,
dwYSize: 0,
dwXCountChars: 0,
dwYCountChars: 0,
dwFillAttribute: 0,
dwFlags: STARTUPINFOW_FLAGS(0),
wShowWindow: 0,
cbReserved2: 0,
lpReserved2: std::ptr::null_mut(),
hStdInput: HANDLE(std::ptr::null_mut()),
hStdOutput: HANDLE(std::ptr::null_mut()),
hStdError: HANDLE(std::ptr::null_mut()),
};
let envp = make_envp(set_env)?;
let dirp = make_dirp(work_dir)?;
let flags = if show_window { PROCESS_CREATION_FLAGS(0) } else { CREATE_NO_WINDOW };
unsafe {
CreateProcessW(exe_name, params, None, None, false, flags, envp.0, dirp.0, &si, &mut pi)?;
let _ = AllowSetForegroundWindow(pi.dwProcessId);
let _ = CloseHandle(pi.hThread);
}
Ok(SafeProcessHandle(pi.hProcess))
}
pub fn wait_process_timeout(process: HANDLE, dur: Duration) -> std::io::Result<Option<u32>> {
let ms = dur
.as_secs()
.checked_mul(1000)
.and_then(|amt| amt.checked_add((dur.subsec_nanos() / 1_000_000) as u64))
.expect("failed to convert duration to milliseconds");
let ms: u32 = if ms > (u32::max_value() as u64) { u32::max_value() } else { ms as u32 };
unsafe {
match WaitForSingleObject(process, ms) {
WAIT_OBJECT_0 => {}
WAIT_TIMEOUT => return Ok(None),
_ => return Err(std::io::Error::last_os_error()),
}
let mut exit_code = 0;
GetExitCodeProcess(process, &mut exit_code)?;
Ok(Some(exit_code))
}
}
pub fn kill_process(process: HANDLE) -> std::io::Result<()> {
unsafe {
let _ = CloseHandle(process);
}
Ok(())
}

View File

@@ -1,8 +1,5 @@
use std::{
ffi::{c_void, OsStr},
os::windows::{ffi::OsStrExt, process::CommandExt},
path::{Path, PathBuf},
process::Command as Process,
time::Duration,
};
@@ -10,56 +7,52 @@ use velopack::locator::VelopackLocator;
use anyhow::{anyhow, Result};
use normpath::PathExt;
use wait_timeout::ChildExt;
use windows::core::PCWSTR;
use windows::Win32::Storage::FileSystem::GetLongPathNameW;
use windows::Win32::System::SystemInformation::{VerSetConditionMask, VerifyVersionInfoW, OSVERSIONINFOEXW, VER_FLAGS};
use windows::Win32::UI::WindowsAndMessaging::AllowSetForegroundWindow;
use windows::Win32::{
Foundation::{CloseHandle, HANDLE},
Security::{GetTokenInformation, TokenElevation, TOKEN_ELEVATION},
Storage::FileSystem::GetLongPathNameW,
System::Threading::{GetCurrentProcess, OpenProcessToken},
UI::Shell::{ShellExecuteExA, ShellExecuteExW, SEE_MASK_NOCLOSEPROCESS, SHELLEXECUTEINFOW},
};
use windows::{
core::PCWSTR,
Win32::{Foundation::HWND, UI::Shell::ShellExecuteW},
};
use crate::shared::{self, runtime_arch::RuntimeArch};
use crate::windows::strings::{string_to_u16, u16_to_string};
use crate::{
shared::{self, runtime_arch::RuntimeArch},
windows::process,
};
pub fn run_hook(locator: &VelopackLocator, hook_name: &str, timeout_secs: u64) -> bool {
let sw = simple_stopwatch::Stopwatch::start_new();
let root_dir = locator.get_root_dir();
let current_path = locator.get_current_bin_dir();
let main_exe_path = locator.get_main_exe_path();
let main_exe_path = locator.get_main_exe_path_as_string();
let ver_string = locator.get_manifest_version_full_string();
let args = vec![hook_name, &ver_string];
let mut success = false;
info!("Running {} hook...", hook_name);
const CREATE_NO_WINDOW: u32 = 0x08000000;
let cmd = Process::new(&main_exe_path).args(args).current_dir(&current_path).creation_flags(CREATE_NO_WINDOW).spawn();
let cmd = process::run_process(
main_exe_path,
args.iter().map(|f| f.to_string()).collect(),
Some(current_path.to_string_lossy().to_string()),
None,
false,
);
if let Err(e) = cmd {
warn!("Failed to start hook {}: {}", hook_name, e);
return false;
}
let mut cmd = cmd.unwrap();
let _ = unsafe { AllowSetForegroundWindow(cmd.id()) };
let cmd = cmd.unwrap();
match cmd.wait_timeout(Duration::from_secs(timeout_secs)) {
match process::wait_process_timeout(cmd.handle(), Duration::from_secs(timeout_secs)) {
Ok(Some(status)) => {
if status.success() {
if status == 0 {
info!("Hook executed successfully (took {}ms)", sw.ms());
success = true;
} else {
warn!("Hook exited with non-zero exit code: {}", status.code().unwrap_or(0));
warn!("Hook exited with non-zero exit code: {}", status);
}
}
Ok(None) => {
let _ = cmd.kill();
let _ = process::kill_process(cmd.handle());
error!("Process timed out after {}s", timeout_secs);
}
Err(e) => {
@@ -72,44 +65,6 @@ pub fn run_hook(locator: &VelopackLocator, hook_name: &str, timeout_secs: u64) -
success
}
pub fn is_process_elevated() -> bool {
// Get the current process handle
let process = unsafe { GetCurrentProcess() };
// Variable to hold the process token
let mut token: HANDLE = HANDLE::default();
// Open the process token with the TOKEN_QUERY access rights
unsafe {
if OpenProcessToken(process, windows::Win32::Security::TOKEN_QUERY, &mut token).is_ok() {
// Allocate a buffer for the TOKEN_ELEVATION structure
let mut elevation = TOKEN_ELEVATION::default();
let mut size: u32 = 0;
let elevation_ptr: *mut core::ffi::c_void = &mut elevation as *mut _ as *mut _;
// Query the token information to check if it is elevated
if GetTokenInformation(token, TokenElevation, Some(elevation_ptr), std::mem::size_of::<TOKEN_ELEVATION>() as u32, &mut size)
.is_ok()
{
// Return whether the token is elevated
CloseHandle(token);
return elevation.TokenIsElevated != 0;
}
}
}
// Clean up the token handle
if !token.is_invalid() {
unsafe { CloseHandle(token) };
}
false
}
pub fn expand_environment_strings<P: AsRef<str>>(input: P) -> Result<String> {
use windows::Win32::System::Environment::ExpandEnvironmentStringsW;
let encoded_u16 = super::strings::string_to_u16(input);