CNWei
e3cb065b35
- 使用 `OsString`/`PathBuf` 替代 `String`，实现跨平台无损编码与命令行参数透传 - 完善 `FromLua` 转换与边界校验，拦截空洞 (`nil`) 及非整数键 - 修复环境变量拼接时由于内置分隔符导致的 `join_paths` 报错 - 优化 Windows 下 UAC 提权执行逻辑，改用宽字符 API (`ShellExecuteExW`)
236 lines
8.1 KiB
Rust
236 lines
8.1 KiB
Rust
use rshim::Shim;
|
||
use std::os::windows::ffi::OsStrExt;
|
||
use std::{env, ffi::CString, mem::size_of, path::Path, process::exit, ptr::null_mut};
|
||
use tracing_subscriber::{EnvFilter, fmt};
|
||
|
||
use std::ffi::{OsStr, OsString};
|
||
|
||
use windows_sys::Win32::UI::Shell::{SHELLEXECUTEINFOW, ShellExecuteExW};
|
||
|
||
use windows_sys::Win32::Foundation::CloseHandle;
|
||
use windows_sys::{
|
||
Win32::{
|
||
Foundation::{FALSE, TRUE},
|
||
System::{
|
||
Com::{COINIT_APARTMENTTHREADED, COINIT_DISABLE_OLE1DDE, CoInitializeEx},
|
||
Console::{
|
||
CTRL_BREAK_EVENT, CTRL_C_EVENT, CTRL_CLOSE_EVENT, CTRL_LOGOFF_EVENT,
|
||
CTRL_SHUTDOWN_EVENT, SetConsoleCtrlHandler,
|
||
},
|
||
Threading::{GetExitCodeProcess, INFINITE, WaitForSingleObject},
|
||
},
|
||
UI::{
|
||
Shell::{
|
||
SEE_MASK_NOASYNC, SEE_MASK_NOCLOSEPROCESS, SHELLEXECUTEINFOA, ShellExecuteExA,
|
||
},
|
||
WindowsAndMessaging::SW_NORMAL,
|
||
},
|
||
},
|
||
core::BOOL,
|
||
};
|
||
|
||
unsafe extern "system" fn console_ctrl_handler(evt: u32) -> BOOL {
|
||
match evt {
|
||
CTRL_C_EVENT => TRUE, //eprintln!("ctrl_c handled!"),
|
||
CTRL_BREAK_EVENT => TRUE, //eprintln!("ctrl_break handled!"),
|
||
CTRL_CLOSE_EVENT => TRUE, //eprintln!("ctrl_close handled!"),
|
||
CTRL_LOGOFF_EVENT => TRUE, //eprintln!("ctrl_logoff handled!"),
|
||
CTRL_SHUTDOWN_EVENT => TRUE, //eprintln!("ctrl_shutdown handled!"),
|
||
other => {
|
||
eprintln!("未知的系统事件编号: {}，未处理！", other);
|
||
FALSE
|
||
}
|
||
}
|
||
}
|
||
|
||
const EXIT_FAILED_LOAD_SHIM: i32 = 1;
|
||
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() {
|
||
// 初始化日志：输出到 stderr，避免污染 shim 子进程的 stdout
|
||
fmt()
|
||
.with_writer(std::io::stderr)
|
||
.with_env_filter(
|
||
EnvFilter::try_from_env("SHIM_LOG").unwrap_or_else(|_| EnvFilter::new("warn")),
|
||
)
|
||
.init();
|
||
|
||
let res: BOOL = unsafe { SetConsoleCtrlHandler(Some(console_ctrl_handler), TRUE) };
|
||
if res == FALSE {
|
||
eprintln!("警告: 注册控制台中断事件处理器失败。");
|
||
}
|
||
|
||
let calling_args: Vec<_> = env::args_os().skip(1).collect();
|
||
let shim = match Shim::load() {
|
||
Ok(v) => v,
|
||
Err(e) => {
|
||
eprintln!("加载代理(shim)配置时发生错误: {}", e);
|
||
exit(EXIT_FAILED_LOAD_SHIM);
|
||
}
|
||
};
|
||
|
||
// 构建 Command：复用 ShimConfig::to_command（含 args/env 注入），避免重复逻辑
|
||
let mut cmd = shim.to_command(&calling_args);
|
||
|
||
// 提权回退时需要完整参数：配置默认参数 + 调用方透传参数
|
||
let mut args = shim.args.clone();
|
||
args.extend_from_slice(&calling_args);
|
||
|
||
let mut child = match cmd.spawn() {
|
||
Ok(v) => v,
|
||
Err(e) if e.raw_os_error() == Some(ERROR_ELEVATION_REQUIRED) => {
|
||
exit(execute_elevated(&shim.target, &args, Some(&shim.env)))
|
||
}
|
||
Err(e) => {
|
||
eprintln!(
|
||
"启动目标程序 [{}] 时发生错误: {}",
|
||
shim.target.to_string_lossy(),
|
||
e
|
||
);
|
||
exit(EXIT_FAILED_SPAWN_PROG);
|
||
}
|
||
};
|
||
let status = match child.wait() {
|
||
Ok(v) => v,
|
||
Err(e) => {
|
||
eprintln!(
|
||
"等待目标程序 [{}] 执行完毕时发生错误: {}",
|
||
shim.target.to_string_lossy(),
|
||
e
|
||
);
|
||
exit(EXIT_FAILED_WAIT_PROG);
|
||
}
|
||
};
|
||
exit(status.code().unwrap_or(EXIT_PROG_TERMINATED))
|
||
}
|
||
// 辅助函数：将任意 OsStr 转换为以 \0 结尾的 UTF-16 宽字符向量 (Vec<u16>)
|
||
fn to_wide_null(s: impl AsRef<OsStr>) -> Vec<u16> {
|
||
s.as_ref().encode_wide().chain(std::iter::once(0)).collect()
|
||
}
|
||
|
||
fn execute_elevated(
|
||
program: &Path,
|
||
args: &[OsString],
|
||
env_vars: Option<&std::collections::HashMap<String, OsString>>,
|
||
) -> i32 {
|
||
// 若提权启动，在此处将环境变量设置给当前进程（即将弹窗 UAC 的进程，随后会被子进程继承）
|
||
if let Some(env_map) = env_vars {
|
||
for (k, v) in env_map {
|
||
unsafe {
|
||
env::set_var(k, v);
|
||
}
|
||
}
|
||
}
|
||
// 2. 将参数列表按 Windows 命令行规则拼装为单个命令行字符串
|
||
let mut arguments_os = OsString::new();
|
||
for (i, arg) in args.iter().enumerate() {
|
||
if i > 0 {
|
||
arguments_os.push(" ");
|
||
}
|
||
let arg_str = arg.to_string_lossy();
|
||
if arg_str.is_empty() {
|
||
arguments_os.push("\"\"");
|
||
} else if !arg_str.contains([' ', '\t', '"']) {
|
||
arguments_os.push(arg);
|
||
} else {
|
||
// 包含空格或引号时进行标准转义包裹
|
||
arguments_os.push("\"");
|
||
for c in arg_str.chars() {
|
||
match c {
|
||
'\\' => arguments_os.push("\\\\"),
|
||
'"' => arguments_os.push("\\\""),
|
||
_ => arguments_os.push(c.to_string()),
|
||
}
|
||
}
|
||
arguments_os.push("\"");
|
||
}
|
||
}
|
||
// let runas = CString::new("runas").unwrap();
|
||
// let program = CString::new(program.to_str().unwrap()).unwrap();
|
||
// let mut arguments = String::new();
|
||
// for arg in args.iter() {
|
||
// arguments.push(' ');
|
||
// if arg.len() == 0 {
|
||
// arguments.push_str("\"\"");
|
||
// } else if arg.find(&[' ', '\t', '"'][..]).is_none() {
|
||
// arguments.push_str(&arg);
|
||
// } else {
|
||
// arguments.push('"');
|
||
// for c in arg.chars() {
|
||
// match c {
|
||
// '\\' => arguments.push_str("\\\\"),
|
||
// '"' => arguments.push_str("\\\""),
|
||
// c => arguments.push(c),
|
||
// }
|
||
// }
|
||
// arguments.push('"');
|
||
// }
|
||
// }
|
||
// 3. 准备 Windows 宽字符参数
|
||
let runas = to_wide_null("runas");
|
||
let program_wide = to_wide_null(program.as_os_str());
|
||
let arguments_wide = to_wide_null(&arguments_os);
|
||
|
||
let mut info = SHELLEXECUTEINFOW {
|
||
cbSize: size_of::<SHELLEXECUTEINFOW>() as u32,
|
||
fMask: SEE_MASK_NOASYNC | SEE_MASK_NOCLOSEPROCESS,
|
||
hwnd: null_mut(),
|
||
lpVerb: runas.as_ptr(),
|
||
lpFile: program_wide.as_ptr(),
|
||
lpParameters: arguments_wide.as_ptr(),
|
||
lpDirectory: null_mut(),
|
||
nShow: SW_NORMAL as i32,
|
||
hInstApp: null_mut(),
|
||
lpIDList: null_mut(),
|
||
lpClass: null_mut(),
|
||
hkeyClass: null_mut(),
|
||
dwHotKey: 0,
|
||
Anonymous: unsafe { std::mem::zeroed() },
|
||
hProcess: null_mut(),
|
||
};
|
||
// let arguments = CString::new(&arguments[..]).unwrap();
|
||
// let mut info = SHELLEXECUTEINFOA::default();
|
||
// info.cbSize = size_of::<SHELLEXECUTEINFOA>() as u32;
|
||
// info.fMask = SEE_MASK_NOASYNC | SEE_MASK_NOCLOSEPROCESS;
|
||
// info.lpVerb = runas.as_ptr().cast::<u8>();
|
||
// info.lpFile = program.as_ptr().cast::<u8>();
|
||
// info.lpParameters = arguments.as_ptr().cast::<u8>();
|
||
// info.nShow = SW_NORMAL;
|
||
let res = unsafe {
|
||
CoInitializeEx(
|
||
null_mut(),
|
||
(COINIT_APARTMENTTHREADED | COINIT_DISABLE_OLE1DDE) as u32,
|
||
);
|
||
// ShellExecuteExA(&mut info as *mut _)
|
||
ShellExecuteExW(&mut info)
|
||
};
|
||
if res == FALSE || info.hProcess == null_mut() {
|
||
return EXIT_FAILED_SPAWN_PROG;
|
||
}
|
||
// 5. 等待提权子进程执行结束并获取退出状态码
|
||
let mut exit_code: u32 = 0;
|
||
unsafe {
|
||
WaitForSingleObject(info.hProcess, INFINITE);
|
||
let ok = GetExitCodeProcess(info.hProcess, &mut exit_code);
|
||
CloseHandle(info.hProcess); // 释放进程句柄，防止资源泄露
|
||
|
||
if ok == FALSE {
|
||
return EXIT_FAILED_WAIT_PROG;
|
||
}
|
||
}
|
||
|
||
exit_code as i32
|
||
|
||
// let mut code: u32 = 0;
|
||
// unsafe {
|
||
// WaitForSingleObject(info.hProcess, INFINITE);
|
||
// if GetExitCodeProcess(info.hProcess, &mut code as *mut _) == FALSE {
|
||
// return EXIT_FAILED_WAIT_PROG;
|
||
// }
|
||
// }
|
||
// code as i32
|
||
}
|