Add support for elevation

This commit is contained in:
Zoritle
2021-01-21 23:05:50 +08:00
parent 60e3190020
commit 7934dde74d
2 changed files with 83 additions and 2 deletions

View File

@@ -16,4 +16,16 @@ unicode-bom = "1"
[dependencies.winapi] [dependencies.winapi]
version = "0.3" version = "0.3"
features = ["wincon","consoleapi","minwindef"] features = [
"wincon",
"consoleapi",
"minwindef",
"shellapi",
"winuser",
"synchapi",
"combaseapi",
"winbase",
"processthreadsapi",
"objbase",
"impl-default"
]

View File

@@ -1,11 +1,25 @@
use std::{ use std::{
env, env,
ffi::CString,
mem::size_of,
path::Path,
process::{exit, Command}, process::{exit, Command},
ptr::null_mut,
}; };
use winapi::{ use winapi::{
shared::minwindef::{BOOL, DWORD, FALSE, TRUE}, 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 { 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_FAILED_WAIT_PROG: i32 = 3;
const EXIT_PROG_TERMINATED: i32 = 4; const EXIT_PROG_TERMINATED: i32 = 4;
const ERROR_ELEVATION_REQUIRED: i32 = 740;
fn main() { fn main() {
let res: BOOL = unsafe { consoleapi::SetConsoleCtrlHandler(Some(routine_handler), TRUE) }; let res: BOOL = unsafe { consoleapi::SetConsoleCtrlHandler(Some(routine_handler), TRUE) };
if res == FALSE { if res == FALSE {
@@ -52,6 +67,9 @@ fn main() {
}; };
let mut cmd = match Command::new(&shim.target_path).args(&args).spawn() { let mut cmd = match Command::new(&shim.target_path).args(&args).spawn() {
Ok(v) => v, Ok(v) => v,
Err(e) if e.raw_os_error() == Some(ERROR_ELEVATION_REQUIRED) => {
exit(execute_elevated(&shim.target_path, &args))
}
Err(e) => { Err(e) => {
eprintln!( eprintln!(
"Error while spawning target program `{}`: {}", "Error while spawning target program `{}`: {}",
@@ -74,3 +92,54 @@ fn main() {
}; };
exit(status.code().unwrap_or(EXIT_PROG_TERMINATED)) 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(&params[..]).unwrap();
let mut info = SHELLEXECUTEINFOA::default();
info.cbSize = size_of::<SHELLEXECUTEINFOA>() 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;
}