🚨 修复内存分配和句柄关闭错误

src-tauri/src/commands.rs

@@ -89,30 +89,39 @@ pub fn is_in_admin() -> bool {
 
   unsafe {
     let mut token_handle: HANDLE = std::ptr::null_mut();
-    if OpenProcessToken(GetCurrentProcess(), TOKEN_QUERY, &mut token_handle) != 0 {
-      let nt_authority = SID_IDENTIFIER_AUTHORITY { Value: [0, 0, 0, 0, 0, 5] };
-      let mut admin_group = std::ptr::null_mut();
-      if AllocateAndInitializeSid(
-        &nt_authority,
-        2,
-        SECURITY_BUILTIN_DOMAIN_RID.try_into().unwrap(),
-        DOMAIN_ALIAS_RID_ADMINS.try_into().unwrap(),
-        0,
-        0,
-        0,
-        0,
-        0,
-        0,
-        &mut admin_group,
-      ) != 0
-      {
-        let mut is_admin = 0i32;
-        CheckTokenMembership(std::ptr::null_mut(), admin_group, &mut is_admin);
-        FreeSid(admin_group);
-        return is_admin != 0;
-      }
+    if OpenProcessToken(GetCurrentProcess(), TOKEN_QUERY, &mut token_handle) == 0 {
+      return false;
     }
-    false
+
+    let nt_authority = SID_IDENTIFIER_AUTHORITY { Value: [0, 0, 0, 0, 0, 5] };
+    let mut admin_group = std::ptr::null_mut();
+
+    let success = AllocateAndInitializeSid(
+      &nt_authority,
+      2,
+      SECURITY_BUILTIN_DOMAIN_RID.try_into().unwrap(),
+      DOMAIN_ALIAS_RID_ADMINS.try_into().unwrap(),
+      0,
+      0,
+      0,
+      0,
+      0,
+      0,
+      &mut admin_group,
+    );
+
+    if success == 0 {
+      CloseHandle(token_handle);
+      return false;
+    }
+
+    let mut is_admin = 0i32;
+    let result = CheckTokenMembership(std::ptr::null_mut(), admin_group, &mut is_admin);
+
+    FreeSid(admin_group);
+    CloseHandle(token_handle);
+
+    result != 0 && is_admin != 0
   }
 }
 

src-tauri/src/yae/inject.rs

@@ -5,7 +5,7 @@ use std::ffi::{c_void, OsStr};
 use std::iter::once;
 use std::os::windows::ffi::OsStrExt;
 use std::ptr;
-use windows_sys::Win32::Foundation::{HANDLE, INVALID_HANDLE_VALUE};
+use windows_sys::Win32::Foundation::{CloseHandle, HANDLE, INVALID_HANDLE_VALUE};
 use windows_sys::Win32::Storage::FileSystem::PIPE_ACCESS_DUPLEX;
 use windows_sys::Win32::System::Diagnostics::Debug::WriteProcessMemory;
 use windows_sys::Win32::System::Diagnostics::ToolHelp::{
@@ -89,15 +89,27 @@ pub fn inject_dll(pi: &PROCESS_INFORMATION, dll_path: &str) {
   let size = dll_utf16.len() * 2;
 
   unsafe {
-    // 在远程进程分配内存并写入 DLL 路径
     let addr = VirtualAllocEx(pi.hProcess, ptr::null_mut(), size, MEM_COMMIT, PAGE_READWRITE);
-    WriteProcessMemory(pi.hProcess, addr, dll_utf16.as_ptr() as *const _, size, ptr::null_mut());
+    if addr.is_null() {
+      panic!("VirtualAllocEx failed");
+    }
+
+    let success =
+      WriteProcessMemory(pi.hProcess, addr, dll_utf16.as_ptr() as *const _, size, ptr::null_mut());
+    if success == 0 {
+      panic!("WriteProcessMemory failed");
+    }
 
-    // 获取 kernel32!LoadLibraryW 地址
     let k32 = GetModuleHandleA(b"kernel32.dll\0".as_ptr());
-    let loadlib = GetProcAddress(k32, b"LoadLibraryW\0".as_ptr());
+    if k32 == 0 {
+      panic!("GetModuleHandleA failed");
+    }
+
+    let loadlib = GetProcAddress(k32, b"LoadLibraryW\0".as_ptr());
+    if loadlib.is_null() {
+      panic!("GetProcAddress failed");
+    }
 
-    // 在远程进程里调用 LoadLibraryW
     let thread = CreateRemoteThread(
       pi.hProcess,
       ptr::null_mut(),
@@ -107,6 +119,10 @@ pub fn inject_dll(pi: &PROCESS_INFORMATION, dll_path: &str) {
       0,
       ptr::null_mut(),
     );
+    if thread.is_null() {
+      panic!("CreateRemoteThread failed");
+    }
+
     WaitForSingleObject(thread, INFINITE);
   }
 }
@@ -115,6 +131,10 @@ pub fn inject_dll(pi: &PROCESS_INFORMATION, dll_path: &str) {
 pub fn find_module_base(pid: u32, dll_name: &str) -> Option<usize> {
   unsafe {
     let snapshot = CreateToolhelp32Snapshot(TH32CS_SNAPMODULE, pid);
+    if snapshot == INVALID_HANDLE_VALUE {
+      return None;
+    }
+
     let mut me32 =
       MODULEENTRY32W { dwSize: std::mem::size_of::<MODULEENTRY32W>() as u32, ..Default::default() };
 
@@ -122,6 +142,7 @@ pub fn find_module_base(pid: u32, dll_name: &str) -> Option<usize> {
       loop {
         let name = String::from_utf16_lossy(&me32.szModule);
         if name.contains(dll_name) {
+          CloseHandle(snapshot);
           return Some(me32.modBaseAddr as usize);
         }
         if Module32NextW(snapshot, &mut me32) == 0 {
@@ -129,6 +150,8 @@ pub fn find_module_base(pid: u32, dll_name: &str) -> Option<usize> {
         }
       }
     }
+
+    CloseHandle(snapshot);
   }
   None
 }
@@ -137,29 +160,34 @@ pub fn find_module_base(pid: u32, dll_name: &str) -> Option<usize> {
 pub fn call_yaemain(pi: &PROCESS_INFORMATION, base: usize, dll_path: &str) {
   let dll_path_wide: Vec<u16> = to_wide_string(dll_path);
   unsafe {
-    // 本地解析 YaeMain 地址
     let local =
-      LoadLibraryExW(dll_path_wide.as_ptr(), ptr::null_mut(), DONT_RESOLVE_DLL_REFERENCES);
-    let proc = GetProcAddress(local, b"YaeMain\0".as_ptr()).expect("无法找到 YaeMain");
+      LoadLibraryExW(dll_path_wide.as_ptr(), std::ptr::null_mut(), DONT_RESOLVE_DLL_REFERENCES);
+    if local == 0 {
+      panic!("LoadLibraryExW failed");
+    }
 
-    // 把函数指针转成裸地址
-    let proc_addr = proc as *const () as usize;
+    let proc = GetProcAddress(local, b"YaeMain\0".as_ptr());
+    if proc.is_null() {
+      FreeLibrary(local);
+      panic!("无法找到 YaeMain");
+    }
 
-    // 计算 RVA
+    let proc_addr = proc as usize;
     let rva = proc_addr - local as usize;
     println!("YaeMain RVA: {:#x}", rva);
 
-    let remote_yaemain = (base + rva) as *mut c_void;
+    FreeLibrary(local);
 
+    let remote_yaemain = (base + rva) as *mut std::ffi::c_void;
     // 在远程进程里调用 YaeMain(hModule)
     CreateRemoteThread(
       pi.hProcess,
-      ptr::null_mut(),
+      std::ptr::null_mut(),
       0,
       Some(std::mem::transmute(remote_yaemain)),
-      base as *mut _, // hModule 参数
+      base as *mut _,
       0,
-      ptr::null_mut(),
+      std::ptr::null_mut(),
     );
   }
 }

src-tauri/src/yae/mod.rs

@@ -199,4 +199,8 @@ pub fn call_yae_dll(app_handle: AppHandle, game_path: String) -> () {
       }
     }
   });
+  unsafe {
+    CloseHandle(pi.hProcess);
+    CloseHandle(pi.hThread);
+  }
 }