From 8aba995fb03f20c07874a304050ee747773183c3 Mon Sep 17 00:00:00 2001 From: CNWei Date: Tue, 25 Aug 2026 15:12:20 +0800 Subject: [PATCH] =?UTF-8?q?refactor(main):=20=E4=BC=98=E5=8C=96=20Lua=20?= =?UTF-8?q?=E9=85=8D=E7=BD=AE=E8=A7=A3=E6=9E=90=E4=B8=8E=E6=A0=A1=E9=AA=8C?= =?UTF-8?q?=E9=80=BB=E8=BE=91=EF=BC=8C?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 优化 Lua 配置解析与校验逻辑 - 清理废弃的注释代码 --- src/lib.rs | 15 ++-- src/main.rs | 200 +++---------------------------------------------- src/sys.rs | 5 +- src/sys/win.rs | 30 ++++---- 4 files changed, 39 insertions(+), 211 deletions(-) diff --git a/src/lib.rs b/src/lib.rs index abbaf9a..cefb0ac 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,17 +1,20 @@ extern crate core; -mod spec; mod layout; +mod loader; mod logger; mod runtime; mod shim; -mod loader; -mod validators; +mod spec; +pub mod sys; mod utils; -mod sys; +mod validators; -pub use spec::ShimSpec; pub use layout::Layout; pub use runtime::LuaRuntime; pub use shim::Shim; -pub use sys::execute_elevated; +pub use spec::ShimSpec; +// pub use sys::{ +// ERROR_ELEVATION_REQUIRED, EXIT_FAILED_LOAD_SHIM, EXIT_FAILED_SPAWN_PROG, EXIT_FAILED_WAIT_PROG, +// EXIT_PROG_TERMINATED, execute_elevated, set_console_ctrl_handler, +// }; diff --git a/src/main.rs b/src/main.rs index 04acf7d..2def279 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,68 +1,19 @@ +use rshim::sys::*; 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 rshim::execute_elevated; -use std::ffi::{OsStr, OsString}; +use std::{env, process::exit}; +use tracing_subscriber::{fmt, EnvFilter}; -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 + //初始化日志：输出到 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!("警告: 注册控制台中断事件处理器失败。"); - } - + // 2. 注册 Windows 控制台信号 + set_console_ctrl_handler(); + // 3. 解析调用参数与代理 Shim 配置 let calling_args: Vec<_> = env::args_os().skip(1).collect(); let shim = match Shim::load() { Ok(v) => v, @@ -71,17 +22,16 @@ fn main() { 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) => { + // 提权回退时需要完整参数：配置默认参数 + 调用方透传参数 + let mut args = shim.args.clone(); + args.extend_from_slice(&calling_args); + exit(execute_elevated(&shim.target, &args, Some(&shim.env))) } Err(e) => { @@ -93,6 +43,7 @@ fn main() { exit(EXIT_FAILED_SPAWN_PROG); } }; + // 5. 等待子进程退出 let status = match child.wait() { Ok(v) => v, Err(e) => { @@ -106,130 +57,3 @@ fn main() { }; exit(status.code().unwrap_or(EXIT_PROG_TERMINATED)) } -// 辅助函数：将任意 OsStr 转换为以 \0 结尾的 UTF-16 宽字符向量 (Vec) -// fn to_wide_null(s: impl AsRef) -> Vec { -// s.as_ref().encode_wide().chain(std::iter::once(0)).collect() -// } -// -// fn execute_elevated( -// program: &Path, -// args: &[OsString], -// env_vars: Option<&std::collections::HashMap>, -// ) -> 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::() 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::() as u32; -// // info.fMask = SEE_MASK_NOASYNC | SEE_MASK_NOCLOSEPROCESS; -// // info.lpVerb = runas.as_ptr().cast::(); -// // info.lpFile = program.as_ptr().cast::(); -// // info.lpParameters = arguments.as_ptr().cast::(); -// // 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 -// } diff --git a/src/sys.rs b/src/sys.rs index d8da9c7..da797ec 100644 --- a/src/sys.rs +++ b/src/sys.rs @@ -1,3 +1,6 @@ mod win; -pub use win::{set_console_ctrl_handler, execute_elevated}; \ No newline at end of file +pub use win::{ + ERROR_ELEVATION_REQUIRED, EXIT_FAILED_LOAD_SHIM, EXIT_FAILED_SPAWN_PROG, EXIT_FAILED_WAIT_PROG, + EXIT_PROG_TERMINATED, execute_elevated, set_console_ctrl_handler, +}; diff --git a/src/sys/win.rs b/src/sys/win.rs index 14c5f73..bd22d82 100644 --- a/src/sys/win.rs +++ b/src/sys/win.rs @@ -1,40 +1,38 @@ -use crate::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::{env, mem::size_of, path::Path, ptr::null_mut}; use std::ffi::{OsStr, OsString}; -use windows_sys::Win32::UI::Shell::{SHELLEXECUTEINFOW, ShellExecuteExW}; +use windows_sys::Win32::UI::Shell::{ShellExecuteExW, SHELLEXECUTEINFOW}; use windows_sys::Win32::Foundation::CloseHandle; use windows_sys::{ + core::BOOL, Win32::{ Foundation::{FALSE, TRUE}, System::{ - Com::{COINIT_APARTMENTTHREADED, COINIT_DISABLE_OLE1DDE, CoInitializeEx}, + Com::{CoInitializeEx, COINIT_APARTMENTTHREADED, COINIT_DISABLE_OLE1DDE}, Console::{ - CTRL_BREAK_EVENT, CTRL_C_EVENT, CTRL_CLOSE_EVENT, CTRL_LOGOFF_EVENT, - CTRL_SHUTDOWN_EVENT, SetConsoleCtrlHandler, + SetConsoleCtrlHandler, CTRL_BREAK_EVENT, CTRL_CLOSE_EVENT, CTRL_C_EVENT, + CTRL_LOGOFF_EVENT, CTRL_SHUTDOWN_EVENT, }, - Threading::{GetExitCodeProcess, INFINITE, WaitForSingleObject}, + Threading::{GetExitCodeProcess, WaitForSingleObject, INFINITE}, }, UI::{ Shell::{ - SEE_MASK_NOASYNC, SEE_MASK_NOCLOSEPROCESS, SHELLEXECUTEINFOA, ShellExecuteExA, + SEE_MASK_NOASYNC, SEE_MASK_NOCLOSEPROCESS, }, WindowsAndMessaging::SW_NORMAL, }, }, - core::BOOL, }; -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; +pub const EXIT_FAILED_LOAD_SHIM: i32 = 1; +pub const EXIT_FAILED_SPAWN_PROG: i32 = 2; +pub const EXIT_FAILED_WAIT_PROG: i32 = 3; +pub const EXIT_PROG_TERMINATED: i32 = 4; -const ERROR_ELEVATION_REQUIRED: i32 = 740; +pub const ERROR_ELEVATION_REQUIRED: i32 = 740; unsafe extern "system" fn console_ctrl_handler(evt: u32) -> BOOL { match evt { @@ -153,5 +151,5 @@ pub fn execute_elevated( } exit_code as i32 - + }