diff --git a/Cargo.toml b/Cargo.toml index bcfb181..393ac81 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -16,4 +16,16 @@ unicode-bom = "1" [dependencies.winapi] version = "0.3" -features = ["wincon","consoleapi","minwindef"] +features = [ + "wincon", + "consoleapi", + "minwindef", + "shellapi", + "winuser", + "synchapi", + "combaseapi", + "winbase", + "processthreadsapi", + "objbase", + "impl-default" +] diff --git a/src/main.rs b/src/main.rs index ea8ab2a..fc22cfa 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,11 +1,25 @@ use std::{ env, + ffi::CString, + mem::size_of, + path::Path, process::{exit, Command}, + ptr::null_mut, }; use winapi::{ shared::minwindef::{BOOL, DWORD, FALSE, TRUE}, - um::{consoleapi, wincon}, + um::{ + combaseapi::CoInitializeEx, + consoleapi, + objbase::{COINIT_APARTMENTTHREADED, COINIT_DISABLE_OLE1DDE}, + processthreadsapi::GetExitCodeProcess, + shellapi::{ShellExecuteExA, SEE_MASK_NOASYNC, SEE_MASK_NOCLOSEPROCESS, SHELLEXECUTEINFOA}, + synchapi::WaitForSingleObject, + winbase::INFINITE, + wincon, + winuser::SW_NORMAL, + }, }; unsafe extern "system" fn routine_handler(evt: DWORD) -> BOOL { @@ -30,6 +44,7 @@ const EXIT_FAILED_SPAWN_PROG: i32 = 2; const EXIT_FAILED_WAIT_PROG: i32 = 3; const EXIT_PROG_TERMINATED: i32 = 4; +const ERROR_ELEVATION_REQUIRED: i32 = 740; fn main() { let res: BOOL = unsafe { consoleapi::SetConsoleCtrlHandler(Some(routine_handler), TRUE) }; if res == FALSE { @@ -52,6 +67,9 @@ fn main() { }; let mut cmd = match Command::new(&shim.target_path).args(&args).spawn() { Ok(v) => v, + Err(e) if e.raw_os_error() == Some(ERROR_ELEVATION_REQUIRED) => { + exit(execute_elevated(&shim.target_path, &args)) + } Err(e) => { eprintln!( "Error while spawning target program `{}`: {}", @@ -74,3 +92,54 @@ fn main() { }; exit(status.code().unwrap_or(EXIT_PROG_TERMINATED)) } + +fn execute_elevated(program: &Path, args: &[String]) -> i32 { + let runas = CString::new("runas").unwrap(); + let program = CString::new(program.to_str().unwrap()).unwrap(); + let mut params = String::new(); + for arg in args.iter() { + params.push(' '); + if arg.len() == 0 { + params.push_str("\"\""); + } else if arg.find(&[' ', '\t', '"'][..]).is_none() { + params.push_str(&arg); + } else { + params.push('"'); + for c in arg.chars() { + match c { + '\\' => params.push_str("\\\\"), + '"' => params.push_str("\\\""), + c => params.push(c), + } + } + params.push('"'); + } + } + + let params = CString::new(¶ms[..]).unwrap(); + let mut info = SHELLEXECUTEINFOA::default(); + info.cbSize = size_of::() as DWORD; + info.fMask = SEE_MASK_NOASYNC | SEE_MASK_NOCLOSEPROCESS; + info.lpVerb = runas.as_ptr(); + info.lpFile = program.as_ptr(); + info.lpParameters = params.as_ptr(); + info.nShow = SW_NORMAL; + let res = unsafe { + CoInitializeEx( + null_mut(), + COINIT_APARTMENTTHREADED | COINIT_DISABLE_OLE1DDE, + ); + ShellExecuteExA(&mut info as *mut _) + }; + if res == FALSE || info.hProcess == null_mut() { + return EXIT_FAILED_SPAWN_PROG; + } + let mut code: DWORD = 0; + unsafe { + WaitForSingleObject(info.hProcess, INFINITE); + if GetExitCodeProcess(info.hProcess, &mut code as *mut _) == FALSE { + return EXIT_FAILED_WAIT_PROG; + } + } + return code as i32; +}