refactor: 迁移 winapi 到 windows-sys,修复配置解析漏洞
- 依赖替换为 windows-sys 0.61,main.rs 全面适配新 API - 配置解析错误显式传播:env/args 类型错误、数组空洞、非法键不再静默 吞错 - 修正空环境变量与空参数语义,补充 UTF-8 与路径拼接校验 - require 容错移至 Rust 侧,模块加载失败记录日志并跳过 - 新增配置解析与运行时单元测试(19 个)
This commit is contained in:
133
src/main.rs
133
src/main.rs
@@ -1,36 +1,36 @@
|
||||
use std::{
|
||||
env,
|
||||
ffi::CString,
|
||||
mem::size_of,
|
||||
path::Path,
|
||||
process::{Command, exit},
|
||||
ptr::null_mut,
|
||||
};
|
||||
use std::{env, ffi::CString, mem::size_of, path::Path, process::exit, ptr::null_mut};
|
||||
|
||||
use rshim::Shim;
|
||||
use tracing_subscriber::{EnvFilter, fmt};
|
||||
|
||||
use winapi::{
|
||||
shared::minwindef::{BOOL, DWORD, FALSE, TRUE},
|
||||
um::{
|
||||
combaseapi::CoInitializeEx,
|
||||
consoleapi,
|
||||
objbase::{COINIT_APARTMENTTHREADED, COINIT_DISABLE_OLE1DDE},
|
||||
processthreadsapi::GetExitCodeProcess,
|
||||
shellapi::{SEE_MASK_NOASYNC, SEE_MASK_NOCLOSEPROCESS, SHELLEXECUTEINFOA, ShellExecuteExA},
|
||||
synchapi::WaitForSingleObject,
|
||||
winbase::INFINITE,
|
||||
wincon,
|
||||
winuser::SW_NORMAL,
|
||||
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 routine_handler(evt: DWORD) -> BOOL {
|
||||
unsafe extern "system" fn console_ctrl_handler(evt: u32) -> BOOL {
|
||||
match evt {
|
||||
wincon::CTRL_C_EVENT => TRUE, //eprintln!("ctrl_c handled!"),
|
||||
wincon::CTRL_BREAK_EVENT => TRUE, //eprintln!("ctrl_break handled!"),
|
||||
wincon::CTRL_CLOSE_EVENT => TRUE, //eprintln!("ctrl_close handled!"),
|
||||
wincon::CTRL_LOGOFF_EVENT => TRUE, //eprintln!("ctrl_logoff handled!"),
|
||||
wincon::CTRL_SHUTDOWN_EVENT => TRUE, //eprintln!("ctrl_shutdown handled!"),
|
||||
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
|
||||
@@ -45,13 +45,21 @@ const EXIT_PROG_TERMINATED: i32 = 4;
|
||||
|
||||
const ERROR_ELEVATION_REQUIRED: i32 = 740;
|
||||
fn main() {
|
||||
let res: BOOL = unsafe { consoleapi::SetConsoleCtrlHandler(Some(routine_handler), TRUE) };
|
||||
// 初始化日志:输出到 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().skip(1).collect();
|
||||
let shim = match Shim::init() {
|
||||
let shim = match Shim::load() {
|
||||
Ok(v) => v,
|
||||
Err(e) => {
|
||||
eprintln!("加载代理(shim)配置时发生错误: {}", e);
|
||||
@@ -59,31 +67,22 @@ fn main() {
|
||||
}
|
||||
};
|
||||
|
||||
let args = if let Some(mut shim_args) = shim.args {
|
||||
shim_args.extend_from_slice(calling_args.as_slice());
|
||||
shim_args
|
||||
} else {
|
||||
calling_args
|
||||
};
|
||||
// ======= 【修改位置 1:构建 Command 并注入环境变量】 =======
|
||||
let mut cmd_builder = Command::new(&shim.target_path);
|
||||
cmd_builder.args(&args);
|
||||
// 构建 Command:复用 ShimConfig::to_command(含 args/env 注入),避免重复逻辑
|
||||
let mut cmd = shim.to_command(&calling_args);
|
||||
|
||||
// 仅作用于目标子进程,完全 Safe 且隔离
|
||||
if let Some(ref envs) = shim.envs {
|
||||
cmd_builder.envs(envs);
|
||||
}
|
||||
let mut cmd = match cmd_builder.spawn() {
|
||||
// 提权回退时需要完整参数:配置默认参数 + 调用方透传参数
|
||||
let mut args = shim.args.clone();
|
||||
args.extend_from_slice(&calling_args);
|
||||
|
||||
let mut cmd = match cmd.spawn() {
|
||||
Ok(v) => v,
|
||||
Err(e) if e.raw_os_error() == Some(ERROR_ELEVATION_REQUIRED) => exit(execute_elevated(
|
||||
&shim.target_path,
|
||||
&args,
|
||||
shim.envs.as_ref(),
|
||||
)),
|
||||
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_path.to_string_lossy(),
|
||||
shim.target.to_string_lossy(),
|
||||
e
|
||||
);
|
||||
exit(EXIT_FAILED_SPAWN_PROG);
|
||||
@@ -94,7 +93,7 @@ fn main() {
|
||||
Err(e) => {
|
||||
eprintln!(
|
||||
"等待目标程序 [{}] 执行完毕时发生错误: {}",
|
||||
shim.target_path.to_string_lossy(),
|
||||
shim.target.to_string_lossy(),
|
||||
e
|
||||
);
|
||||
exit(EXIT_FAILED_WAIT_PROG);
|
||||
@@ -106,10 +105,10 @@ fn main() {
|
||||
fn execute_elevated(
|
||||
program: &Path,
|
||||
args: &[String],
|
||||
envs: Option<&std::collections::HashMap<String, String>>,
|
||||
env_vars: Option<&std::collections::HashMap<String, String>>,
|
||||
) -> i32 {
|
||||
// 若提权启动,在此处将环境变量设置给当前进程(即将弹窗 UAC 的进程,随后会被子进程继承)
|
||||
if let Some(env_map) = envs {
|
||||
if let Some(env_map) = env_vars {
|
||||
for (k, v) in env_map {
|
||||
unsafe {
|
||||
env::set_var(k, v);
|
||||
@@ -119,45 +118,45 @@ fn execute_elevated(
|
||||
|
||||
let runas = CString::new("runas").unwrap();
|
||||
let program = CString::new(program.to_str().unwrap()).unwrap();
|
||||
let mut params = String::new();
|
||||
let mut arguments = String::new();
|
||||
for arg in args.iter() {
|
||||
params.push(' ');
|
||||
arguments.push(' ');
|
||||
if arg.len() == 0 {
|
||||
params.push_str("\"\"");
|
||||
arguments.push_str("\"\"");
|
||||
} else if arg.find(&[' ', '\t', '"'][..]).is_none() {
|
||||
params.push_str(&arg);
|
||||
arguments.push_str(&arg);
|
||||
} else {
|
||||
params.push('"');
|
||||
arguments.push('"');
|
||||
for c in arg.chars() {
|
||||
match c {
|
||||
'\\' => params.push_str("\\\\"),
|
||||
'"' => params.push_str("\\\""),
|
||||
c => params.push(c),
|
||||
'\\' => arguments.push_str("\\\\"),
|
||||
'"' => arguments.push_str("\\\""),
|
||||
c => arguments.push(c),
|
||||
}
|
||||
}
|
||||
params.push('"');
|
||||
arguments.push('"');
|
||||
}
|
||||
}
|
||||
|
||||
let params = CString::new(¶ms[..]).unwrap();
|
||||
let arguments = CString::new(&arguments[..]).unwrap();
|
||||
let mut info = SHELLEXECUTEINFOA::default();
|
||||
info.cbSize = size_of::<SHELLEXECUTEINFOA>() as DWORD;
|
||||
info.cbSize = size_of::<SHELLEXECUTEINFOA>() as u32;
|
||||
info.fMask = SEE_MASK_NOASYNC | SEE_MASK_NOCLOSEPROCESS;
|
||||
info.lpVerb = runas.as_ptr();
|
||||
info.lpFile = program.as_ptr();
|
||||
info.lpParameters = params.as_ptr();
|
||||
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,
|
||||
(COINIT_APARTMENTTHREADED | COINIT_DISABLE_OLE1DDE) as u32,
|
||||
);
|
||||
ShellExecuteExA(&mut info as *mut _)
|
||||
};
|
||||
if res == FALSE || info.hProcess == null_mut() {
|
||||
return EXIT_FAILED_SPAWN_PROG;
|
||||
}
|
||||
let mut code: DWORD = 0;
|
||||
let mut code: u32 = 0;
|
||||
unsafe {
|
||||
WaitForSingleObject(info.hProcess, INFINITE);
|
||||
if GetExitCodeProcess(info.hProcess, &mut code as *mut _) == FALSE {
|
||||
|
||||
Reference in New Issue
Block a user