diff --git a/Cargo.toml b/Cargo.toml index bd19522..b14f8f5 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -22,4 +22,4 @@ libc = "0.2" mach = "0.3" [target.'cfg(windows)'.dependencies] -winapi = { version = "0.3", features = ["winnt", "memoryapi", "minwindef", "processthreadsapi"] } +winapi = { version = "0.3", features = ["winnt", "memoryapi", "minwindef", "processthreadsapi", "handleapi"] } diff --git a/examples/fastyboy.rs b/examples/fastyboy.rs index 31002f1..ce3d3fc 100644 --- a/examples/fastyboy.rs +++ b/examples/fastyboy.rs @@ -15,64 +15,19 @@ fn main() { println!("FastyBoy can only be run on systems supporting the game Mirror's Edge Catalyst, which as of writing is only Windows.") } -/// A helper function to get a Pid from the name of a process -#[cfg(windows)] -pub fn get_pid(process_name: &str) -> process_memory::Pid { - /// A helper function to turn a c_char array to a String - fn utf8_to_string(bytes: &[i8]) -> String { - use std::ffi::CStr; - unsafe { - CStr::from_ptr(bytes.as_ptr()) - .to_string_lossy() - .into_owned() - } - } - let mut entry = winapi::um::tlhelp32::PROCESSENTRY32 { - dwSize: std::mem::size_of::() as u32, - cntUsage: 0, - th32ProcessID: 0, - th32DefaultHeapID: 0, - th32ModuleID: 0, - cntThreads: 0, - th32ParentProcessID: 0, - pcPriClassBase: 0, - dwFlags: 0, - szExeFile: [0; winapi::shared::minwindef::MAX_PATH], - }; - let snapshot: winapi::um::winnt::HANDLE; - unsafe { - snapshot = winapi::um::tlhelp32::CreateToolhelp32Snapshot( - winapi::um::tlhelp32::TH32CS_SNAPPROCESS, - 0, - ); - if winapi::um::tlhelp32::Process32First(snapshot, &mut entry) - == winapi::shared::minwindef::TRUE - { - while winapi::um::tlhelp32::Process32Next(snapshot, &mut entry) - == winapi::shared::minwindef::TRUE - { - if utf8_to_string(&entry.szExeFile) == process_name { - return entry.th32ProcessID; - } - } - } - } - 0 -} - #[cfg(windows)] fn main() -> std::io::Result<()> { use process_memory::*; - let process_handle = get_pid("MirrorsEdgeCatalyst.exe").try_into_process_handle()?; + let process_handle = get_pid("MirrorsEdgeCatalyst.exe")?.try_into_process_handle()?; let mut spawn_timer = DataMember::::new(process_handle); - spawn_timer.set_offset(vec![0x1_42_14_2a_d8, 0xac]); + spawn_timer.set_offset(0x1_42_14_2a_d8, vec![0xac]); let mut level_warmup = DataMember::::new(process_handle); - level_warmup.set_offset(vec![0x1_42_14_2a_d8, 0x9c]); + level_warmup.set_offset(0x1_42_14_2a_d8, vec![0x9c]); let mut emitters_enabled = DataMember::::new(process_handle); - emitters_enabled.set_offset(vec![0x1_42_3e_44_78, 0xac]); + emitters_enabled.set_offset(0x1_42_3e_44_78, vec![0xac]); spawn_timer.write(&1.0)?; level_warmup.write(&1.0)?; diff --git a/src/data_member.rs b/src/data_member.rs index 11ae376..ccd4109 100644 --- a/src/data_member.rs +++ b/src/data_member.rs @@ -14,7 +14,7 @@ use crate::{CopyAddress, Memory, ProcessHandle, PutAddress}; /// // We need to make sure that we get a handle to a process, in this case, ourselves /// let handle = (std::process::id() as Pid).try_into_process_handle().unwrap(); /// // We make a `DataMember` that has an offset referring to its location in memory -/// let member = DataMember::new_offset(handle, vec![&x as *const _ as usize]); +/// let member = DataMember::::new_addr(handle, &x as *const _ as usize); /// // The memory refered to is now the same /// println!("Memory location: &x: {}, member: {}", &x as *const _ as usize, /// member.get_offset().unwrap()); @@ -29,7 +29,8 @@ use crate::{CopyAddress, Memory, ProcessHandle, PutAddress}; /// ``` #[derive(Clone, Debug)] pub struct DataMember { - offsets: Vec, + base: usize, + offsets: Vec, process: ProcessHandle, _phantom: std::marker::PhantomData<*mut T>, } @@ -51,13 +52,31 @@ impl DataMember { #[must_use] pub fn new(handle: ProcessHandle) -> Self { Self { + base: 0, offsets: Vec::new(), process: handle, _phantom: std::marker::PhantomData, } } - /// Create a new `DataMember` from a [`ProcessHandle`] and some number of offsets. You must + // /// Create a new `DataMember` from a [`ProcessHandle`] and some number of offsets. You must + // /// remember to call [`try_into_process_handle`] on a [`Pid`] as sometimes the `Pid` can have + // /// the same backing type as a [`ProcessHandle`], resulting in an error. + // /// + // /// [`try_into_process_handle`]: trait.TryIntoProcessHandle.html#tymethod.try_into_process_handle + // /// [`ProcessHandle`]: type.ProcessHandle.html + // /// [`Pid`]: type.Pid.html + // #[must_use] + // pub fn new_offset(handle: ProcessHandle, offsets: Vec) -> Self { + // Self { + // base: 0, + // offsets, + // process: handle, + // _phantom: std::marker::PhantomData, + // } + // } + + /// Create a new `DataMember` from a [`ProcessHandle`], a base address, and some number of offsets. You must /// remember to call [`try_into_process_handle`] on a [`Pid`] as sometimes the `Pid` can have /// the same backing type as a [`ProcessHandle`], resulting in an error. /// @@ -65,26 +84,58 @@ impl DataMember { /// [`ProcessHandle`]: type.ProcessHandle.html /// [`Pid`]: type.Pid.html #[must_use] - pub fn new_offset(handle: ProcessHandle, offsets: Vec) -> Self { + pub fn new_addr_offset(handle: ProcessHandle, base: usize, offsets: Vec) -> Self { Self { + base, offsets, process: handle, _phantom: std::marker::PhantomData, } } + + /// Create a new `DataMember` from a [`ProcessHandle`], and an absolute address (just one offset, so to speak). You must + /// remember to call [`try_into_process_handle`] on a [`Pid`] as sometimes the `Pid` can have + /// the same backing type as a [`ProcessHandle`], resulting in an error. + /// + /// [`try_into_process_handle`]: trait.TryIntoProcessHandle.html#tymethod.try_into_process_handle + /// [`ProcessHandle`]: type.ProcessHandle.html + /// [`Pid`]: type.Pid.html + #[must_use] + pub fn new_addr(handle: ProcessHandle, base: usize) -> Self { + Self { + base, + offsets: Vec::new(), + process: handle, + _phantom: std::marker::PhantomData, + } + } + + /// Clones an existing `DataMember`, and appends some offsets. Mostly provided for convenience. + #[must_use] + pub fn extend(&self, more_offsets: Vec) -> DataMember { + let mut clone = DataMember { + base: self.base, + offsets: self.offsets.clone(), + process: self.process, + _phantom: std::marker::PhantomData, + }; + clone.offsets.extend(more_offsets); + clone + } } impl Memory for DataMember { - fn set_offset(&mut self, new_offsets: Vec) { + fn set_offset(&mut self, new_base: usize, new_offsets: Vec) { + self.base = new_base; self.offsets = new_offsets; } fn get_offset(&self) -> std::io::Result { - self.process.get_offset(&self.offsets) + self.process.get_offset(self.base, &self.offsets) } fn read(&self) -> std::io::Result { - let offset = self.process.get_offset(&self.offsets)?; + let offset = self.process.get_offset(self.base, &self.offsets)?; // This can't be [0_u8;size_of::()] because no const generics. // It will be freed at the end of the function because no references are held to it. let mut buffer = vec![0_u8; std::mem::size_of::()]; @@ -94,7 +145,7 @@ impl Memory for DataMember { fn write(&self, value: &T) -> std::io::Result<()> { use std::slice; - let offset = self.process.get_offset(&self.offsets)?; + let offset = self.process.get_offset(self.base, &self.offsets)?; let buffer: &[u8] = unsafe { slice::from_raw_parts(value as *const _ as _, std::mem::size_of::()) }; self.process.put_address(offset, &buffer) @@ -114,7 +165,7 @@ mod test { .unwrap(); println!("Process Handle: {:?}", handle); let mut member = DataMember::::new(handle); - member.set_offset(vec![&test as *const _ as usize]); + member.set_offset(&test as *const _ as usize, vec![]); assert_eq!(test, member.read().unwrap()); member.write(&5_i32).unwrap(); assert_eq!(test, 5_i32); @@ -128,7 +179,7 @@ mod test { .unwrap(); println!("Process Handle: {:?}", handle); let mut member = DataMember::::new(handle); - member.set_offset(vec![&test as *const _ as usize]); + member.set_offset(&test as *const _ as usize, vec![]); assert_eq!(test, member.read().unwrap()); member.write(&-1_i64).unwrap(); assert_eq!(test, -1); @@ -141,8 +192,9 @@ mod test { .try_into_process_handle() .unwrap(); println!("Process Handle: {:?}", handle); - let mut member = DataMember::::new(handle); - member.set_offset(vec![&test as *const _ as usize]); + // let mut member = DataMember::::new(handle); + // member.set_offset(&test as *const _ as usize, vec![]); + let member = DataMember::::new_addr(handle, &test as *const _ as usize); assert_eq!(test, member.read().unwrap()); member.write(&0xffff).unwrap(); assert_eq!(test, 0xffff); diff --git a/src/lib.rs b/src/lib.rs index 27024cb..d7c8028 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -15,9 +15,9 @@ //! // We need to make sure that we get a handle to a process, in this case, ourselves //! let handle = (std::process::id() as Pid).try_into_process_handle().unwrap(); //! // We make a `DataMember` that has an offset referring to its location in memory -//! let member = DataMember::new_offset(handle, vec![&x as *const _ as usize]); +//! let member = DataMember::new_addr(handle, &x as *const _ as usize); //! // The memory refered to is now the same -//! println!("Memory location: &x: {}, member: {}", &x as *const _ as usize, +//! println!("Memory location: &x: {:X}, member: {}", &x as *const _ as usize, //! member.get_offset().unwrap()); //! assert_eq!(&x as *const _ as usize, member.get_offset().unwrap()); //! // The value of the member is the same as the variable @@ -35,9 +35,9 @@ //! println!("Original x-value: {}", x); //! //! // We make a `LocalMember` that has an offset referring to its location in memory -//! let member = LocalMember::new_offset(vec![&x as *const _ as usize]); +//! let member = LocalMember::new_addr(&x as *const _ as usize); //! // The memory refered to is now the same -//! println!("Memory location: &x: {}, member: {}", &x as *const _ as usize, +//! println!("Memory location: &x: {:X}, member: {}", &x as *const _ as usize, //! member.get_offset().unwrap()); //! assert_eq!(&x as *const _ as usize, member.get_offset().unwrap()); //! // The value of the member is the same as the variable @@ -58,7 +58,7 @@ //! .set_arch(Architecture::Arch32Bit); //! // We make a `DataMember` that has a series of offsets refering to a known value in //! // the target processes memory -//! let member = DataMember::new_offset(handle, vec![0x01_02_03_04, 0x04, 0x08, 0x10]); +//! let member = DataMember::new_addr_offset(handle, 0x01020304, vec![0x04, 0x08, 0x10]); //! // The memory offset can now be correctly calculated: //! println!("Target memory location: {}", member.get_offset().unwrap()); //! // The memory offset can now be used to retrieve and modify values: @@ -92,6 +92,9 @@ mod platform; #[path = "windows.rs"] mod platform; +#[cfg(windows)] +pub use platform::get_pid; + /// A trait that defines that it is possible to copy some memory from something represented by a /// type into a buffer. pub trait CopyAddress { @@ -109,19 +112,26 @@ pub trait CopyAddress { /// /// # Errors /// `std::io::Error` if an error occurs copying the address. - fn get_offset(&self, offsets: &[usize]) -> std::io::Result { + #[allow(clippy::cast_sign_loss)] + fn get_offset(&self, base: usize, offsets: &[isize]) -> std::io::Result { // Look ma! No unsafes! - let mut offset: usize = 0; + let noffsets: usize = offsets.len(); - let mut copy = vec![0_u8; self.get_pointer_width() as usize]; - for next_offset in offsets.iter().take(noffsets - 1) { - offset += next_offset; - self.copy_address(offset, &mut copy)?; - offset = self.get_pointer_width().pointer_from_ne_bytes(©); - } + if noffsets > 0 { + let mut offset: usize = base; + let mut copy = vec![0_u8; self.get_pointer_width() as usize]; + for next_offset in offsets.iter().take(noffsets - 1) { + // should work because of 2s-complement. + offset = offset.wrapping_add(*next_offset as usize); + self.copy_address(offset, &mut copy)?; + offset = self.get_pointer_width().pointer_from_ne_bytes(©); + } - offset += offsets[noffsets - 1]; - Ok(offset) + // should work because of 2s-complement. + Ok(offset.wrapping_add(offsets[noffsets - 1] as usize)) + } else { + Ok(base) + } } /// Get the the pointer width of the underlying process. @@ -190,6 +200,22 @@ pub trait ProcessHandleExt { fn set_arch(self, arch: Architecture) -> Self; } +/// Handling modules (e.g. DLLs) in a process. +pub trait ModuleInfo { + /// Gets the base address of a module in a process. For example, "GameAssembly.dll" when on Windows. + /// You can then use the address in the `base` parameter of [`set_offset`] for example. + /// + /// # Errors + /// Returns `std::io::ErrorKind::NotFound` when no such module name exists. + /// Returns OS Error when something else went wrong. + /// + /// # Panics + /// Panics when closing the handle fails (e.g. double close). + /// + /// [`set_offset`]: trait.Memory.html#tymethod.set_offset + fn get_module_base(&self, name: &str) -> std::io::Result; +} + /// A trait that refers to and allows writing to a region of memory in a running program. pub trait Memory { /// Set the offsets to the location in memory. This is used for things such as multi-level @@ -198,7 +224,7 @@ pub trait Memory { /// For those sorts of data structures, to access data you need to go via multiple pointers, so /// that if an inner region reallocates its size, the variable that is being modified will be /// correctly modified. - fn set_offset(&mut self, new_offsets: Vec); + fn set_offset(&mut self, new_base: usize, new_offsets: Vec); /// Gets the actual total offset from the offsets given by [`Memory::set_offset`]. /// @@ -259,3 +285,53 @@ where source.copy_address(addr, &mut copy)?; Ok(copy) } + +#[cfg(test)] +mod tests { + use super::*; + + #[repr(C)] + #[derive(Clone, Copy)] + struct Player { + x: u32, + y: u32, + } + + #[repr(C)] + struct GameState { + garbage: u32, + garbage2: u32, + players: [Box; 2], // note that this array is in-place, since it's fixed size. + } + + #[test] + fn multilevel_pointers() { + let game = GameState { + garbage: 42, + garbage2: 1337, + players: [ + Box::new(Player { x: 1, y: 2 }), + Box::new(Player { x: 3, y: 4 }), + ], + }; + let handle = (std::process::id() as Pid) + .try_into_process_handle() + .unwrap(); + + let garbage2 = + DataMember::::new_addr_offset(handle, &game as *const _ as usize, vec![4]); + assert_eq!(1337, garbage2.read().unwrap()); + + let second_player = DataMember::::new_addr_offset( + handle, + &game as *const _ as usize, + vec![8 + (handle.get_pointer_width() as u8 as isize) * 1], + ); // skip u32 + i64 + first player. + + let second_player_x = second_player.extend::(vec![0]); + let second_player_y = second_player.extend::(vec![4]); // sizeof u32 + + assert_eq!(3, second_player_x.read().unwrap()); + assert_eq!(4, second_player_y.read().unwrap()); + } +} diff --git a/src/local_member.rs b/src/local_member.rs index 06b5316..0f56474 100644 --- a/src/local_member.rs +++ b/src/local_member.rs @@ -10,7 +10,7 @@ use crate::Memory; /// let x = 4u32; /// /// // We make a `LocalMember` that has an offset referring to its location in memory -/// let member = LocalMember::new_offset(vec![&x as *const _ as usize]); +/// let member = LocalMember::new_addr(&x as *const _ as usize); /// // The memory refered to is now the same /// assert_eq!(&x as *const _ as usize, member.get_offset().unwrap()); /// // The value of the member is the same as the variable @@ -34,7 +34,8 @@ use crate::Memory; /// to mess something up really badly in your program. #[derive(Clone, Debug, Default)] pub struct LocalMember { - offsets: Vec, + base: usize, + offsets: Vec, _phantom: std::marker::PhantomData<*mut T>, } @@ -51,6 +52,7 @@ impl LocalMember { #[must_use] pub fn new() -> Self { Self { + base: 0, offsets: Vec::new(), _phantom: std::marker::PhantomData, } @@ -58,8 +60,19 @@ impl LocalMember { /// Create a new `LocalMember` with a given set of offsets. #[must_use] - pub fn new_offset(offsets: Vec) -> Self { + pub fn new_addr(base: usize) -> Self { Self { + base, + offsets: Vec::new(), + _phantom: std::marker::PhantomData, + } + } + + /// Create a new `LocalMember` with a given set of offsets. + #[must_use] + pub fn new_addr_offset(base: usize, offsets: Vec) -> Self { + Self { + base, offsets, _phantom: std::marker::PhantomData, } @@ -67,28 +80,34 @@ impl LocalMember { } impl Memory for LocalMember { - fn set_offset(&mut self, new_offsets: Vec) { + fn set_offset(&mut self, new_base: usize, new_offsets: Vec) { + self.base = new_base; self.offsets = new_offsets; } + #[allow(clippy::cast_sign_loss)] fn get_offset(&self) -> std::io::Result { - let mut offset = 0_usize; - for i in 0..self.offsets.len() - 1 { - offset = offset.wrapping_add(self.offsets[i]); - if offset == 0 { - return Err(std::io::Error::new( - std::io::ErrorKind::Other, - "Would be a null dereference!", - )); - } - // We can't guarantee alignment, so we must use `read_unaligned()` - // to ensure that its ok to read from, as `read()` requires that - // our source pointer is properly aligned. - unsafe { - offset = (offset as *const usize).read_unaligned(); + if self.offsets.is_empty() { + Ok(self.base) + } else { + let mut offset = self.base; + for i in 0..self.offsets.len() - 1 { + offset = offset.wrapping_add(self.offsets[i] as usize); // should work because of 2s-complement + if offset == 0 { + return Err(std::io::Error::new( + std::io::ErrorKind::Other, + "Would be a null dereference!", + )); + } + // We can't guarantee alignment, so we must use `read_unaligned()` + // to ensure that its ok to read from, as `read()` requires that + // our source pointer is properly aligned. + unsafe { + offset = (offset as *const usize).read_unaligned(); + } } + Ok(offset.wrapping_add(self.offsets[self.offsets.len() - 1] as usize)) } - Ok(offset.wrapping_add(self.offsets[self.offsets.len() - 1])) } /// This will only return a error if one of the offsets gives a null pointer. or give a @@ -120,7 +139,7 @@ mod test { fn modify_local_i32() { let test = 4_i32; let mut member = LocalMember::::new(); - member.set_offset(vec![&test as *const _ as usize]); + member.set_offset(&test as *const _ as usize, vec![]); assert_eq!(test, member.read().unwrap()); member.write(&5_i32).unwrap(); assert_eq!(test, 5_i32); @@ -129,7 +148,7 @@ mod test { fn modify_local_i64() { let test = 3_i64; let mut member = LocalMember::::new(); - member.set_offset(vec![&test as *const _ as usize]); + member.set_offset(&test as *const _ as usize, vec![]); assert_eq!(test, member.read().unwrap()); member.write(&-1_i64).unwrap(); assert_eq!(test, -1); @@ -137,8 +156,7 @@ mod test { #[test] fn modify_local_usize() { let test = 0_usize; - let mut member = LocalMember::::new(); - member.set_offset(vec![&test as *const _ as usize]); + let member = LocalMember::::new_addr(&test as *const _ as usize); assert_eq!(test, member.read().unwrap()); member.write(&0xffff).unwrap(); assert_eq!(test, 0xffff); diff --git a/src/windows.rs b/src/windows.rs index 1ed2a39..c16c516 100644 --- a/src/windows.rs +++ b/src/windows.rs @@ -1,10 +1,13 @@ -use winapi::shared::minwindef; +use winapi::um::{handleapi::CloseHandle, tlhelp32}; +use winapi::{shared::minwindef, um::handleapi::INVALID_HANDLE_VALUE}; use std::os::windows::io::AsRawHandle; use std::process::Child; use std::ptr; -use super::{Architecture, CopyAddress, ProcessHandleExt, PutAddress, TryIntoProcessHandle}; +use super::{ + Architecture, CopyAddress, ModuleInfo, ProcessHandleExt, PutAddress, TryIntoProcessHandle, +}; /// On Windows a `Pid` is a `DWORD`. pub type Pid = minwindef::DWORD; @@ -107,3 +110,172 @@ impl PutAddress for ProcessHandle { } } } + +/// Use `CreateToolhelp32Snapshot` to get and filter list of loaded modules (called DLLs in Windows) +/// of this process, returning the base address of it. +#[allow(clippy::clippy::cast_possible_truncation)] // for size_of as u32 +impl ModuleInfo for Pid { + fn get_module_base(&self, name: &str) -> std::io::Result { + // taken from https://stackoverflow.com/questions/41552466/how-do-i-get-the-physical-baseaddress-of-an-dll-used-in-a-process + let mut module_entry = tlhelp32::MODULEENTRY32 { + dwSize: 0, + th32ModuleID: 0, + th32ProcessID: 0, + GlblcntUsage: 0, + ProccntUsage: 0, + modBaseAddr: std::ptr::null_mut(), // yikes + modBaseSize: 0, + hModule: std::ptr::null_mut(), // yikes + szModule: [0; tlhelp32::MAX_MODULE_NAME32 + 1], + szExePath: [0; winapi::shared::minwindef::MAX_PATH], + }; + + unsafe { + module_entry.dwSize = std::mem::size_of::() as u32; + + let snapshot = tlhelp32::CreateToolhelp32Snapshot( + tlhelp32::TH32CS_SNAPMODULE | tlhelp32::TH32CS_SNAPMODULE32, + *self, + ); + if snapshot == INVALID_HANDLE_VALUE { + return Err(std::io::Error::last_os_error()); + } + + if tlhelp32::Module32First(snapshot, &mut module_entry) == minwindef::TRUE { + // makeshift do-while + loop { + if utf8_to_string(&module_entry.szModule) == name { + let addr = module_entry.modBaseAddr as usize; + if CloseHandle(snapshot) == minwindef::FALSE { + panic!("Could not close handle") + }; + return Ok(addr); + } + + if tlhelp32::Module32Next(snapshot, &mut module_entry) == minwindef::FALSE { + break; + } + } + } + + // We searched everything, nothing found + if CloseHandle(snapshot) == minwindef::FALSE { + panic!("Could not close handle") + }; + Err(std::io::Error::new( + std::io::ErrorKind::NotFound, + format!( + "Process PID#{} contains no module named \"{}\".", + *self, name + ), + )) + } + } +} + +/// Get the process ID of some process by name. For example,"MyGame.exe". +/// If you want to get the PID of your own process, use `std::process:id() as Pid` instead. +/// +/// # Errors +/// If no process exists of that name, returns an `std::io::Error` with kind `std::io::ErrorKind::NotFound`. +/// If something went very wrong with the windows API, returns last OS error. +#[allow(clippy::cast_possible_truncation)] +pub fn get_pid(process_name: &str) -> std::io::Result { + let mut entry = winapi::um::tlhelp32::PROCESSENTRY32 { + dwSize: std::mem::size_of::() as u32, + cntUsage: 0, + th32ProcessID: 0, + th32DefaultHeapID: 0, + th32ModuleID: 0, + cntThreads: 0, + th32ParentProcessID: 0, + pcPriClassBase: 0, + dwFlags: 0, + szExeFile: [0; winapi::shared::minwindef::MAX_PATH], + }; + + let snapshot: winapi::um::winnt::HANDLE; + unsafe { + snapshot = winapi::um::tlhelp32::CreateToolhelp32Snapshot( + winapi::um::tlhelp32::TH32CS_SNAPPROCESS, + 0, + ); + if snapshot == INVALID_HANDLE_VALUE { + return Err(std::io::Error::last_os_error()); + } + + if winapi::um::tlhelp32::Process32First(snapshot, &mut entry) + == winapi::shared::minwindef::TRUE + { + // makeshift do-while + loop { + // println!("Have process: {}", utf8_to_string(&entry.szExeFile)); + if utf8_to_string(&entry.szExeFile) == process_name { + let pid = entry.th32ProcessID; + if CloseHandle(snapshot) == minwindef::FALSE { + panic!("Could not close handle") + }; + return Ok(pid); + } + + if winapi::um::tlhelp32::Process32Next(snapshot, &mut entry) + == winapi::shared::minwindef::FALSE + { + break; + } + } + } + + if CloseHandle(snapshot) == minwindef::FALSE { + panic!("Could not close handle") + }; + Err(std::io::Error::new( + std::io::ErrorKind::NotFound, + format!("Could not find Process ID of \"{}\".", process_name), + )) + } +} + +/// A helper function to turn a `c_char` array to a String +fn utf8_to_string(bytes: &[i8]) -> String { + use std::ffi::CStr; + unsafe { + CStr::from_ptr(bytes.as_ptr()) + .to_string_lossy() + .into_owned() + } +} + +#[cfg(test)] +mod tests { + use std::io::ErrorKind; + + use super::*; + + #[test] + fn module_info() { + let pid = std::process::id() as Pid; + let base = pid.get_module_base("ntdll.dll").unwrap(); + assert_ne!(0, base); + // println!("ntdll.exe address: 0x{:X}", base); + + match pid.get_module_base("this_dll_doesnt_exist.dll") { + Ok(_) => panic!(), + Err(e) => { + assert_eq!(ErrorKind::NotFound, e.kind()); + } + } + } + + #[test] + fn getpid() { + let _pid = get_pid("svchost.exe").unwrap(); + + match get_pid("this_process_doesnt_exist.exe") { + Ok(_) => panic!(), + Err(e) => { + assert_eq!(ErrorKind::NotFound, e.kind()); + } + } + } +}