From e282e7967c48c86ced4c511be96831f0b9e67ed7 Mon Sep 17 00:00:00 2001 From: Andrea Venuta Date: Wed, 31 Jan 2024 11:37:48 +0100 Subject: [PATCH 1/2] Make string usage in `inject` more idiomatic --- src/inject.rs | 53 +++++++++++++++++++++------------------------------ 1 file changed, 22 insertions(+), 31 deletions(-) diff --git a/src/inject.rs b/src/inject.rs index 76b992a2..6ae99795 100644 --- a/src/inject.rs +++ b/src/inject.rs @@ -1,19 +1,17 @@ //! Facilities for injecting compiled DLLs into target processes. -use std::ffi::{CStr, CString}; use std::mem::{self, size_of}; use std::path::PathBuf; -use std::ptr::null; use tracing::debug; -use windows::core::{Error, Result, HRESULT, HSTRING, PCSTR, PCWSTR}; +use windows::core::{s, w, Error, Result, HRESULT, HSTRING, PCSTR, PCWSTR}; use windows::Win32::Foundation::{CloseHandle, GetLastError, BOOL, HANDLE, MAX_PATH}; use windows::Win32::System::Diagnostics::Debug::WriteProcessMemory; use windows::Win32::System::Diagnostics::ToolHelp::{ CreateToolhelp32Snapshot, Process32First, Process32FirstW, Process32Next, Process32NextW, PROCESSENTRY32, PROCESSENTRY32W, TH32CS_SNAPPROCESS, }; -use windows::Win32::System::LibraryLoader::{GetModuleHandleA, GetProcAddress}; +use windows::Win32::System::LibraryLoader::{GetModuleHandleW, GetProcAddress}; use windows::Win32::System::Memory::{ VirtualAllocEx, VirtualFreeEx, MEM_COMMIT, MEM_RELEASE, MEM_RESERVE, PAGE_READWRITE, }; @@ -41,17 +39,10 @@ impl Process { /// Inject the DLL in the process. pub fn inject(&self, dll_path: PathBuf) -> Result<()> { - let kernel32 = CString::new("Kernel32").unwrap(); - let loadlibraryw = CString::new("LoadLibraryW").unwrap(); + let proc_addr = + unsafe { GetProcAddress(GetModuleHandleW(w!("Kernel32"))?, s!("LoadLibraryW")) }; - let proc_addr = unsafe { - GetProcAddress( - GetModuleHandleA(PCSTR(kernel32.as_ptr() as _))?, - PCSTR(loadlibraryw.as_ptr() as _), - ) - }; - - let dll_path = HSTRING::from(dll_path.canonicalize().unwrap().as_os_str()); + let dll_path = HSTRING::from(dll_path.canonicalize().unwrap().as_path()); let dll_path_buf = unsafe { VirtualAllocEx( self.0, @@ -97,6 +88,11 @@ impl Process { Ok(()) } } + + /// Retrieve the process handle. + pub fn handle(&self) -> HANDLE { + self.0 + } } impl Drop for Process { @@ -118,9 +114,8 @@ fn get_process_by_title(title: &str) -> Result { // 32-bit implementation. Uses [`std::ffi::CString`] and `FindWindowA`. unsafe fn get_process_by_title32(title: &str) -> Result { - let title = title.split_once('\0').map(|(t, _)| t).unwrap_or(title); - let title = CString::new(title).unwrap(); - let hwnd = FindWindowA(PCSTR(null()), PCSTR(title.as_bytes_with_nul().as_ptr())); + let title = HSTRING::from(title).to_os_string(); + let hwnd = FindWindowA(None, PCSTR(title.as_encoded_bytes().as_ptr())); if hwnd.0 == 0 { let last_error = match GetLastError() { @@ -142,7 +137,7 @@ unsafe fn get_process_by_title32(title: &str) -> Result { // 64-bit implementation. Uses [`widestring::U16CString`] and `FindWindowW`. unsafe fn get_process_by_title64(title: &str) -> Result { let title = HSTRING::from(title); - let hwnd = FindWindowW(PCWSTR(null()), PCWSTR(title.as_ptr())); + let hwnd = FindWindowW(None, PCWSTR(title.as_ptr())); if hwnd.0 == 0 { let last_error = match GetLastError() { @@ -173,9 +168,8 @@ fn get_process_by_name(name: &str) -> Result { } // 32-bit implementation. Uses [`PROCESSENTRY32`]. -unsafe fn get_process_by_name32(name: &str) -> Result { - let name = name.split_once('\0').map(|(t, _)| t).unwrap_or(name); - let name = CString::new(name).unwrap(); +unsafe fn get_process_by_name32(name_str: &str) -> Result { + let name = PCSTR(HSTRING::from(name_str).to_os_string().as_encoded_bytes().as_ptr()); let snapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0)?; let mut pe32 = @@ -193,18 +187,15 @@ unsafe fn get_process_by_name32(name: &str) -> Result { } let pid = loop { - let proc_name = CStr::from_ptr(pe32.szExeFile.as_ptr() as *const i8); + let proc_name = PCSTR(pe32.szExeFile.as_ptr()); - if proc_name == name.as_ref() { + if proc_name == name { break Ok(pe32.th32ProcessID); } if Process32Next(snapshot, &mut pe32).is_err() { CloseHandle(snapshot)?; - break Err(Error::new( - HRESULT(0), - format!("Process {} not found", name.to_string_lossy()).into(), - )); + break Err(Error::new(HRESULT(0), format!("Process {name_str} not found").into())); } }?; @@ -215,8 +206,8 @@ unsafe fn get_process_by_name32(name: &str) -> Result { // 64-bit implementation. Uses [`PROCESSENTRY32W`] and // [`widestring::U16CString`]. -unsafe fn get_process_by_name64(name: &str) -> Result { - let name = PCWSTR::from_raw(HSTRING::from(name).as_ptr()).display().to_string(); +unsafe fn get_process_by_name64(name_str: &str) -> Result { + let name = PCWSTR(HSTRING::from(name_str).as_ptr()); let snapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0)?; let mut pe32 = @@ -234,14 +225,14 @@ unsafe fn get_process_by_name64(name: &str) -> Result { } let pid = loop { - let proc_name = PCWSTR::from_raw(pe32.szExeFile.as_ptr()).display().to_string(); + let proc_name = PCWSTR(pe32.szExeFile.as_ptr()); if proc_name == name { break Ok(pe32.th32ProcessID); } if Process32NextW(snapshot, &mut pe32).is_err() { CloseHandle(snapshot)?; - break Err(Error::new(HRESULT(0), format!("Process {} not found", name).into())); + break Err(Error::new(HRESULT(0), format!("Process {name_str} not found").into())); } }?; From 689803d5b752fabeb2d6718cb3e7ddf0f1883dbc Mon Sep 17 00:00:00 2001 From: Andrea Venuta Date: Wed, 31 Jan 2024 12:02:06 +0100 Subject: [PATCH 2/2] Change things --- src/inject.rs | 16 +++++++++------- tests/inject.rs | 2 +- 2 files changed, 10 insertions(+), 8 deletions(-) diff --git a/src/inject.rs b/src/inject.rs index 6ae99795..d6e09a37 100644 --- a/src/inject.rs +++ b/src/inject.rs @@ -169,7 +169,8 @@ fn get_process_by_name(name: &str) -> Result { // 32-bit implementation. Uses [`PROCESSENTRY32`]. unsafe fn get_process_by_name32(name_str: &str) -> Result { - let name = PCSTR(HSTRING::from(name_str).to_os_string().as_encoded_bytes().as_ptr()); + let name = HSTRING::from(name_str).to_os_string(); + let name = name.as_encoded_bytes(); let snapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0)?; let mut pe32 = @@ -187,7 +188,8 @@ unsafe fn get_process_by_name32(name_str: &str) -> Result { } let pid = loop { - let proc_name = PCSTR(pe32.szExeFile.as_ptr()); + let zero_idx = pe32.szExeFile.iter().position(|&x| x == 0).unwrap_or(pe32.szExeFile.len()); + let proc_name = &pe32.szExeFile[..zero_idx]; if proc_name == name { break Ok(pe32.th32ProcessID); @@ -204,10 +206,9 @@ unsafe fn get_process_by_name32(name_str: &str) -> Result { OpenProcess(PROCESS_ALL_ACCESS, BOOL(0), pid) } -// 64-bit implementation. Uses [`PROCESSENTRY32W`] and -// [`widestring::U16CString`]. +// 64-bit implementation. Uses [`PROCESSENTRY32W`]. unsafe fn get_process_by_name64(name_str: &str) -> Result { - let name = PCWSTR(HSTRING::from(name_str).as_ptr()); + let name = HSTRING::from(name_str); let snapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0)?; let mut pe32 = @@ -225,8 +226,9 @@ unsafe fn get_process_by_name64(name_str: &str) -> Result { } let pid = loop { - let proc_name = PCWSTR(pe32.szExeFile.as_ptr()); - if proc_name == name { + let zero_idx = pe32.szExeFile.iter().position(|&x| x == 0).unwrap_or(pe32.szExeFile.len()); + let proc_name = HSTRING::from_wide(&pe32.szExeFile[..zero_idx])?; + if name == proc_name { break Ok(pe32.th32ProcessID); } diff --git a/tests/inject.rs b/tests/inject.rs index da22ee7d..cac281f8 100644 --- a/tests/inject.rs +++ b/tests/inject.rs @@ -29,7 +29,7 @@ fn test_inject_by_name() { std::thread::sleep(Duration::from_millis(500)); println!("Should show a message box that says \"Hello\"."); - Process::by_name("Notepad.exe") + Process::by_name("notepad.exe") .unwrap() .inject(examples_path().join("dummy_hook.dll")) .unwrap();