🎨 使用 widestring 替代现有处理

2026-03-15 03:53:16 +08:00 · 2025-12-29 21:28:49 +08:00
parent cc3655a700
commit 41987a9a58
2 changed files with 48 additions and 53 deletions
--- a/src-tauri/src/watchdog.rs
+++ b/src-tauri/src/watchdog.rs
@@ -1,35 +1,37 @@
 //! 重启提权相关处理
-//! @since Beta v0.8.7
+//! @since Beta v0.9.1
 #![cfg(target_os = "windows")]

-use std::ffi::OsStr;
-use std::iter::once;
-use std::os::windows::ffi::OsStrExt;
 use std::ptr::null_mut;
 use std::time::Duration;
 use tauri::AppHandle;
+use widestring::U16CString;
 use windows_sys::Win32::Foundation::{HANDLE, HWND};
 use windows_sys::Win32::Storage::FileSystem::SYNCHRONIZE;
 use windows_sys::Win32::System::Threading::{OpenProcess, WaitForSingleObject, INFINITE};
 use windows_sys::Win32::UI::Shell::ShellExecuteW;
 use windows_sys::Win32::UI::WindowsAndMessaging::SW_SHOWNORMAL;

-// 带参数启动
+/// 带参数启动（使用 ShellExecuteW + runas）
 fn shell_runas_with_args(args: &str) -> Result<(), String> {
-  fn to_wide(s: &OsStr) -> Vec<u16> {
-    s.encode_wide().chain(once(0)).collect()
-  }
-
  let exe_path = std::env::current_exe().map_err(|e| e.to_string())?;
-  let exe_wide = to_wide(exe_path.as_os_str());
-  let args_wide = to_wide(OsStr::new(args));
-  let cwd_wide =
-    exe_path.parent().map(|p| to_wide(p.as_os_str())).unwrap_or_else(|| to_wide(OsStr::new("")));
+
+  let exe_wide =
+    U16CString::from_os_str(exe_path.as_os_str()).map_err(|e| format!("路径转换失败: {e}"))?;
+  let args_wide = U16CString::from_str(args).map_err(|e| format!("参数转换失败: {e}"))?;
+  let cwd_wide = exe_path
+    .parent()
+    .map(|p| U16CString::from_os_str(p.as_os_str()))
+    .transpose()
+    .map_err(|e| format!("路径转换失败: {e}"))?
+    .unwrap_or_else(|| U16CString::from_str("").unwrap());
+
+  let verb = U16CString::from_str("runas").unwrap();

  unsafe {
    let result = ShellExecuteW(
      0 as HWND,
-      to_wide(OsStr::new("runas")).as_ptr(),
+      verb.as_ptr(),
      exe_wide.as_ptr(),
      args_wide.as_ptr(),
      if cwd_wide.len() > 1 { cwd_wide.as_ptr() } else { null_mut() },
@@ -38,17 +40,15 @@ fn shell_runas_with_args(args: &str) -> Result<(), String> {
    if (result as usize) > 32 {
      Ok(())
    } else {
-      Err("Failed to ShellExecuteW runas".into())
+      Err("ShellExecuteW runas 启动失败".into())
    }
  }
 }

-// 等待父进程退出（释放单例锁）后，再以管理员身份启动新实例
+/// 等待父进程退出（释放单例锁）后，再以管理员身份启动新实例
 pub fn run_watchdog(parent_pid: u32, args_to_pass: &str) -> Result<(), String> {
-  // 打开父进程句柄用于等待
  let handle: HANDLE = unsafe { OpenProcess(SYNCHRONIZE, 0, parent_pid) };
-  if handle == std::ptr::null_mut() {
-    // 如果拿不到句柄，可能父进程已退出，稍作等待后继续
+  if handle.is_null() {
    std::thread::sleep(Duration::from_millis(300));
  } else {
    unsafe {
@@ -56,11 +56,10 @@ pub fn run_watchdog(parent_pid: u32, args_to_pass: &str) -> Result<(), String> {
    }
  }

-  // 父进程已退出 → 触发 UAC 提权启动新实例
  shell_runas_with_args(args_to_pass)
 }

-// 以管理员权限重启应用
+/// 以管理员权限重启应用
 #[tauri::command]
 pub fn run_with_admin(app_handle: AppHandle) -> Result<(), String> {
  let parent_pid = std::process::id();
@@ -69,11 +68,9 @@ pub fn run_with_admin(app_handle: AppHandle) -> Result<(), String> {
  cmd
    .arg("--watchdog")
    .arg(format!("--ppid={}", parent_pid))
-    // 看门狗不加载单例插件（通过参数决定 main 的初始化）
    .spawn()
    .map_err(|e| format!("spawn watchdog failed: {e}"))?;

-  // 立即退出：单例锁释放
  app_handle.exit(0);
  Ok(())
 }
--- a/src-tauri/src/yae/inject.rs
+++ b/src-tauri/src/yae/inject.rs
@@ -2,10 +2,8 @@
 //! @since Beta v0.9.1
 #![cfg(target_os = "windows")]

-use std::ffi::OsStr;
-use std::iter::once;
-use std::os::windows::ffi::OsStrExt;
 use std::ptr;
+use widestring::U16CString;
 use windows_sys::Win32::Foundation::{CloseHandle, FreeLibrary, HANDLE, INVALID_HANDLE_VALUE};
 use windows_sys::Win32::Storage::FileSystem::PIPE_ACCESS_DUPLEX;
 use windows_sys::Win32::System::Diagnostics::Debug::WriteProcessMemory;
@@ -24,15 +22,10 @@ use windows_sys::Win32::System::Threading::{
  STARTUPINFOW,
 };

-/// 转为宽字符串
-pub fn to_wide_string(s: &str) -> Vec<u16> {
-  OsStr::new(s).encode_wide().chain(once(0)).collect()
-}
-
 /// 创建命名管道
 pub fn create_named_pipe(pipe_name: &str) -> HANDLE {
  let full_pipe_name = format!(r"\\.\pipe\{}", pipe_name);
-  let wide: Vec<u16> = to_wide_string(&full_pipe_name);
+  let wide = U16CString::from_str(&full_pipe_name).expect("invalid pipe name");

  unsafe {
    let handle = CreateNamedPipeW(
@@ -52,46 +45,52 @@ pub fn create_named_pipe(pipe_name: &str) -> HANDLE {
  }
 }

-/// 启动目标进程，附加cwd
+/// 启动目标进程，附加 cwd 或 ticket
 pub fn spawn_process(path: &str, ticket: Option<String>) -> PROCESS_INFORMATION {
-  let wide_path: Vec<u16> = to_wide_string(path); // 如果有 ticket，构造命令行
-  let mut wide_cmd: Option<Vec<u16>> = None;
-  if let Some(ref t) = ticket {
+  let wide_path = U16CString::from_str(path).expect("invalid path");
+
+  let wide_cmd = ticket.as_ref().map(|t| {
    let cmdline = format!("{} login_auth_ticket={}", path, t);
-    wide_cmd = Some(to_wide_string(&cmdline));
-  } // 如果没有 ticket，使用 cwd
-  let wide_cwd: Option<Vec<u16>> = if ticket.is_none() {
-    let cwd = std::path::Path::new(path).parent().unwrap().to_str().unwrap();
-    Some(to_wide_string(cwd))
+    U16CString::from_str(&cmdline).expect("invalid cmdline")
+  });
+
+  let wide_cwd: Option<U16CString> = if ticket.is_none() {
+    std::path::Path::new(path)
+      .parent()
+      .and_then(|p| Some(U16CString::from_os_str(p.as_os_str()).unwrap()))
  } else {
    None
  };
+
  unsafe {
    let mut si: STARTUPINFOW = std::mem::zeroed();
    si.cb = std::mem::size_of::<STARTUPINFOW>() as u32;
    let mut pi: PROCESS_INFORMATION = std::mem::zeroed();
+
    let success = CreateProcessW(
      wide_path.as_ptr(),
-      wide_cmd.as_mut().map(|v| v.as_mut_ptr()).unwrap_or(ptr::null_mut()), // 有 ticket 时传命令行
+      wide_cmd.as_ref().map(|s| s.as_ptr() as *mut u16).unwrap_or(ptr::null_mut()),
      ptr::null_mut(),
      ptr::null_mut(),
      0,
      0,
      ptr::null_mut(),
-      wide_cwd.as_ref().map(|v| v.as_ptr()).unwrap_or(ptr::null()), // 无 ticket 时传 cwd
+      wide_cwd.as_ref().map(|s: &widestring::U16CString| s.as_ptr()).unwrap_or(ptr::null()),
      &mut si,
      &mut pi,
    );
+
    if success == 0 {
      panic!("CreateProcessW failed");
    }
+
    pi
  }
 }

 /// 注入 DLL
 pub fn inject_dll(pi: &PROCESS_INFORMATION, dll_path: &str) {
-  let dll_utf16: Vec<u16> = to_wide_string(dll_path);
+  let dll_utf16 = U16CString::from_str(dll_path).expect("invalid dll path");
  let size = dll_utf16.len() * 2;

  unsafe {
@@ -107,7 +106,7 @@ pub fn inject_dll(pi: &PROCESS_INFORMATION, dll_path: &str) {
    }

    let k32 = GetModuleHandleA(b"kernel32.dll\0".as_ptr());
-    if k32 == std::ptr::null_mut() {
+    if k32.is_null() {
      panic!("GetModuleHandleA failed");
    }

@@ -164,33 +163,32 @@ pub fn find_module_base(pid: u32, dll_name: &str) -> Option<usize> {

 /// 执行 YaeMain
 pub fn call_yaemain(pi: &PROCESS_INFORMATION, base: usize, dll_path: &str) {
-  let dll_path_wide: Vec<u16> = to_wide_string(dll_path);
+  let dll_path_wide = U16CString::from_str(dll_path).expect("invalid dll path");
+
  unsafe {
    let local =
-      LoadLibraryExW(dll_path_wide.as_ptr(), std::ptr::null_mut(), DONT_RESOLVE_DLL_REFERENCES);
-    if local == std::ptr::null_mut() {
+      LoadLibraryExW(dll_path_wide.as_ptr(), ptr::null_mut(), DONT_RESOLVE_DLL_REFERENCES);
+    if local.is_null() {
      panic!("LoadLibraryExW failed");
    }

    let proc = GetProcAddress(local, b"YaeMain\0".as_ptr()).expect("无法找到 YaeMain");

-    // Option<unsafe extern "system" fn() -> isize>
-    let proc_addr = proc as *const () as usize;
+    let proc_addr = proc as usize;
    let rva = proc_addr - local as usize;
    println!("YaeMain RVA: {:#x}", rva);

    FreeLibrary(local);

    let remote_yaemain = (base + rva) as *mut std::ffi::c_void;
-    // 在远程进程里调用 YaeMain(hModule)
    CreateRemoteThread(
      pi.hProcess,
-      std::ptr::null_mut(),
+      ptr::null_mut(),
      0,
      Some(std::mem::transmute(remote_yaemain)),
      base as *mut _,
      0,
-      std::ptr::null_mut(),
+      ptr::null_mut(),
    );
  }
 }