diff --git a/driver/src/utils/handles.rs b/driver/src/utils/handles.rs index b645433..2c09663 100644 --- a/driver/src/utils/handles.rs +++ b/driver/src/utils/handles.rs @@ -1,15 +1,37 @@ use wdk_sys::{ntddk::ZwClose, HANDLE}; +/// A wrapper around a Windows `HANDLE` that automatically closes the handle when dropped. +/// +/// This struct provides a safe abstraction over raw Windows handles, ensuring that the +/// handle is properly closed when it goes out of scope, by calling `ZwClose` in its `Drop` +/// implementation. +/// pub struct Handle(HANDLE); impl Handle { - /// Create new instance `Handle`. + /// Creates a new `Handle` instance. + /// + /// This function wraps a raw Windows `HANDLE` inside the `Handle` struct. + /// + /// # Parameters + /// - `handle`: A raw Windows `HANDLE` to wrap. + /// + /// # Returns + /// - `Handle`: A new `Handle` instance that wraps the given `HANDLE`. + /// #[inline] pub fn new(handle: HANDLE) -> Self { Handle(handle) } - /// Return handle. + /// Returns the raw `HANDLE`. + /// + /// This function provides access to the underlying Windows handle + /// stored in the `Handle` struct. + /// + /// # Returns + /// - `HANDLE`: The raw Windows `HANDLE`. + /// #[inline] pub fn get(&self) -> HANDLE { self.0 @@ -17,6 +39,11 @@ impl Handle { } impl Drop for Handle { + /// Automatically closes the `HANDLE` when the `Handle` instance is dropped. + /// + /// When the `Handle` goes out of scope, this method is called to ensure that + /// the underlying Windows handle is closed using the `ZwClose` function, unless + /// the handle is null. #[inline] fn drop(&mut self) { if !self.0.is_null() { diff --git a/driver/src/utils/ioctls.rs b/driver/src/utils/ioctls.rs index c389ef4..c4ce09e 100644 --- a/driver/src/utils/ioctls.rs +++ b/driver/src/utils/ioctls.rs @@ -12,15 +12,39 @@ use { injection::ioctls::get_injection_ioctls, misc::ioctls::get_misc_ioctls, module::ioctls::get_module_ioctls, + port::ioctls::get_port_ioctls, }, }; +/// Type alias for an IOCTL handler function. +/// +/// This type represents a boxed function that handles IOCTL requests. Each handler takes +/// two parameters, `IRP` (I/O Request Packet) and `IO_STACK_LOCATION`, and returns +/// an `NTSTATUS` result, indicating the success or failure of the operation. +/// +/// # Parameters +/// - `*mut IRP`: Pointer to an IRP (I/O Request Packet), which represents an I/O request in Windows. +/// - `*mut IO_STACK_LOCATION`: Pointer to the current I/O stack location. +/// +/// # Returns +/// - `NTSTATUS`: A status code indicating the success or failure of the operation. pub type IoctlHandler = Box NTSTATUS + Send + Sync>; lazy_static! { + /// A static map that holds the mapping of IOCTL codes to their corresponding handlers. pub static ref IOCTL_MAP: HashMap = load_ioctls(); } +/// Loads the IOCTL handlers into a `HashMap`. +/// +/// This function collects IOCTL handlers from various modules and inserts them +/// into a `HashMap`, which maps IOCTL codes (`u32`) to their respective handler functions (`IoctlHandler`). +/// +/// # Returns +/// - `HashMap`: A map containing IOCTL handlers for process, thread, driver, +/// callback, injection, miscellaneous, module, and port operations. +/// If the "mapper" feature is disabled, registry-related IOCTLs are also included. +/// fn load_ioctls() -> HashMap { let mut ioctls = HashMap::new(); @@ -31,6 +55,7 @@ fn load_ioctls() -> HashMap { get_injection_ioctls(&mut ioctls); get_misc_ioctls(&mut ioctls); get_module_ioctls(&mut ioctls); + get_port_ioctls(&mut ioctls); #[cfg(not(feature = "mapper"))] { get_registry_ioctls(&mut ioctls); diff --git a/driver/src/utils/macros.rs b/driver/src/utils/macros.rs index ab41173..57edb5a 100644 --- a/driver/src/utils/macros.rs +++ b/driver/src/utils/macros.rs @@ -1,8 +1,12 @@ -/// Macro to handle process-related IRP (I/O Request Packet) operations. +/// A macro to handle I/O operations with input and output buffers. /// -/// Matches the input buffer type and executes the given action, returning the status. +/// This macro abstracts common patterns in handling IOCTL operations, where input and/or output +/// buffers are required. It fetches the buffers from the provided IRP (I/O Request Packet) and +/// passes them to the specified action. If fetching the buffers fails, it immediately returns +/// the failure status. +/// #[macro_export] -macro_rules! handle_process { +macro_rules! handle { ($irp:expr, $stack:expr, $action:expr, $input_type:ty, $output_type:ty, $information:expr) => {{ let output_buffer = match crate::utils::get_output_buffer::<$output_type>($irp) { Ok(buffer) => buffer, @@ -26,87 +30,23 @@ macro_rules! handle_process { $action(input_buffer) }}; - ($irp:expr, $action:expr, $type_:ty, $information:expr) => { - let output_buffer = match crate::utils::get_output_buffer::<$type_>($irp) { - Ok(buffer) => buffer, - Err(status) => return status, - } - - $action(output_buffer, $information) - }; -} - -/// Macro to handle thread-related IRP (I/O Request Packet) operations. -/// -/// Matches the input buffer type and executes the given action, returning the status. -#[macro_export] -macro_rules! handle_thread { - ($irp:expr, $stack:expr, $action:expr, $input_type:ty, $output_type:ty, $information:expr) => {{ - let output_buffer = match crate::utils::get_output_buffer::<$output_type>($irp) { - Ok(buffer) => buffer, - Err(status) => return status, - }; - - let input_buffer = match crate::utils::get_input_buffer::<$input_type>($stack) { - Ok(buffer) => buffer, - Err(status) => return status, - }; - - $action(input_buffer, output_buffer, $information) - }}; - - ($irp:expr, $action:expr, $type_:ty) => {{ - let input_buffer = match crate::utils::get_input_buffer::<$type_>($irp) { - Ok(buffer) => buffer, - Err(status) => return status, - }; - - $action(input_buffer) - }}; -} - -/// Macro to handle driver-related operations. -/// -/// Executes the given action based on the provided parameters and returns the status. -#[macro_export] -macro_rules! handle_driver { ($irp:expr, $action:expr, $type_:ty, $information:expr) => {{ let output_buffer = match crate::utils::get_output_buffer::<$type_>($irp) { - Ok(buffer) => buffer, + Ok(buffer) => buffer, Err(status) => return status, }; $action(output_buffer, $information) }}; - - ($stack:expr, $action:expr, $type_:ty) => {{ - let input_buffer = match crate::utils::get_input_buffer::<$type_>($stack) { - Ok(buffer) => buffer, - Err(status) => return status, - }; - - $action(input_buffer) - }}; -} - -/// Macro to handle injection-related operations. -/// -/// Executes the given action based on the provided parameters and returns the status. -#[macro_export] -macro_rules! handle_injection { - ($stack:expr, $action:expr, $type_:ty) => {{ - let input_buffer = match crate::utils::get_input_buffer::<$type_>($stack) { - Ok(buffer) => buffer, - Err(status) => return status, - }; - - $action(input_buffer) - }}; } /// Macro to handle registry-related operations. /// -/// Executes the given action based on the provided parameters and returns the status. +/// This macro abstracts common patterns in handling IOCTL operations, where input and/or output +/// buffers are required. It fetches the buffers from the provided IRP (I/O Request Packet) and +/// passes them to the specified action. If fetching the buffers fails, it immediately returns +/// the failure status. +/// #[cfg(not(feature = "mapper"))] #[macro_export] macro_rules! handle_registry { @@ -128,38 +68,12 @@ macro_rules! handle_registry { }}; } -/// Macro to handle module-related operations. -/// -/// Executes the given action based on the provided parameters and returns the status. -#[macro_export] -macro_rules! handle_module { - ($irp:expr, $stack:expr, $action:expr, $input_type:ty, $output_type:ty, $information:expr) => {{ - let output_buffer = match crate::utils::get_output_buffer::<$output_type>($irp) { - Ok(buffer) => buffer, - Err(status) => return status, - }; - - let input_buffer = match crate::utils::get_input_buffer::<$input_type>($stack) { - Ok(buffer) => buffer, - Err(status) => return status, - }; - - $action(input_buffer, output_buffer, $information) - }}; - - ($stack:expr, $action:expr, $input_type:ty) => {{ - let input_buffer = match crate::utils::get_input_buffer::<$input_type>($stack) { - Ok(buffer) => buffer, - Err(status) => return status, - }; - - $action(input_buffer) - }}; -} - /// Macro to handle callback-related operations. /// -/// Executes the given action based on the provided parameters and returns the status. +/// This macro abstracts common patterns in handling IOCTL operations, where input and/or output +/// buffers are required. It fetches the buffers from the provided IRP (I/O Request Packet) and +/// passes them to the specified action. If fetching the buffers fails, it immediately returns +/// the failure status. /// #[macro_export] macro_rules! handle_callback { diff --git a/driver/src/utils/mod.rs b/driver/src/utils/mod.rs index 1b227e1..26155f8 100644 --- a/driver/src/utils/mod.rs +++ b/driver/src/utils/mod.rs @@ -419,4 +419,8 @@ pub fn return_module() -> Option<(*mut LDR_DATA_TABLE_ENTRY, i32)> { #[allow(dead_code)] pub fn valid_kernel_memory(addr: u64) -> bool { addr > 0x8000000000000000 && addr < 0xFFFFFFFFFFFFFFFF +} + +pub fn valid_user_memory(addr: u64) -> bool { + addr > 0 && addr < 0x7FFFFFFFFFFFFFFF } \ No newline at end of file diff --git a/driver/src/utils/patterns.rs b/driver/src/utils/patterns.rs index 2cefa37..f913620 100644 --- a/driver/src/utils/patterns.rs +++ b/driver/src/utils/patterns.rs @@ -22,6 +22,10 @@ use { }, }; +/// +/// +/// +/// fn slice_to_number(slice: &[u8], func: fn([u8; N]) -> T) -> Result { if slice.len() != N { return Err("Slice length mismatch"); @@ -35,6 +39,7 @@ fn slice_to_number(slice: &[u8], func: fn([u8; N]) -> T) -> R } /// Scans memory for a specific pattern of bytes in a specific section. +/// /// # Parameters /// - `base_addr`: The base address (in `usize` format) from which the scan should start. /// - `section_name`: The name of the section to scan. This string must match the name of the section you want to scan. diff --git a/driver/src/utils/pool.rs b/driver/src/utils/pool.rs index e89a243..b7a4f3c 100644 --- a/driver/src/utils/pool.rs +++ b/driver/src/utils/pool.rs @@ -1,11 +1,31 @@ use core::ffi::c_void; use wdk_sys::{ntddk::{ExAllocatePool2, ExFreePool}, POOL_FLAGS}; +/// A wrapper around memory allocated from the pool in the Windows kernel. +/// +/// This struct provides a safe abstraction over memory allocated from the kernel pool. +/// It ensures that the allocated memory is properly freed when the `PoolMemory` goes out +/// of scope, by calling `ExFreePool` in its `Drop` implementation. +/// pub struct PoolMemory { + /// A raw pointer to the allocated pool memory. pub ptr: *mut c_void, } impl PoolMemory { + /// Allocates memory from the Windows kernel pool. + /// + /// This function uses `ExAllocatePool2` to allocate a block of memory from the Windows kernel + /// pool. It returns `None` if the allocation fails, or `Some(PoolMemory)` if successful. + /// + /// # Parameters + /// - `flag`: Flags controlling the behavior of the memory allocation, of type `POOL_FLAGS`. + /// - `number_of_bytes`: The size of the memory block to allocate, in bytes. + /// - `tag`: A tag (typically a 4-character identifier) used to identify the allocation. + /// + /// # Returns + /// - `Option`: `Some(PoolMemory)` if the memory is successfully allocated, or `None` if the allocation fails. + /// #[inline] pub fn new(flag: POOL_FLAGS, number_of_bytes: u64, tag: u32) -> Option { let ptr = unsafe { ExAllocatePool2(flag, number_of_bytes, tag) }; @@ -20,6 +40,11 @@ impl PoolMemory { } impl Drop for PoolMemory { + /// Frees the allocated pool memory when the `PoolMemory` instance is dropped. + /// + /// This method is automatically called when the `PoolMemory` goes out of scope. It ensures that + /// the memory allocated with `ExAllocatePool2` is properly freed using `ExFreePool`, unless + /// the pointer is null. fn drop(&mut self) { if !self.ptr.is_null() { unsafe { ExFreePool(self.ptr) }; diff --git a/driver/src/utils/process_attach.rs b/driver/src/utils/process_attach.rs index 8e5d113..08f661f 100644 --- a/driver/src/utils/process_attach.rs +++ b/driver/src/utils/process_attach.rs @@ -1,12 +1,34 @@ use wdk_sys::{ntddk::{KeStackAttachProcess, KeUnstackDetachProcess}, KAPC_STATE, PRKPROCESS}; +/// A wrapper for managing the attachment to a process context in the Windows kernel. +/// +/// This struct provides a safe abstraction for attaching to the context of a target process using +/// `KeStackAttachProcess` and ensures that the process is properly detached when no longer needed +/// (either manually or automatically when it goes out of scope). +/// +/// When a `ProcessAttach` instance is dropped, it will automatically detach from the process +/// if still attached. +/// pub struct ProcessAttach { + /// The APC (Asynchronous Procedure Call) state used to manage process attachment. apc_state: KAPC_STATE, + /// Indicates whether the process is currently attached. attached: bool, } impl ProcessAttach { - // Function for attaching the context of a process + /// Attaches to the context of a target process. + /// + /// This function attaches the current thread to the address space of the specified + /// process using `KeStackAttachProcess`. This allows the current thread to operate within + /// the target process context. + /// + /// # Parameters + /// - `target_process`: A pointer to the target process (`PRKPROCESS`) to attach to. + /// + /// # Returns + /// - `ProcessAttach`: A new `ProcessAttach` instance representing the attached process context. + /// #[inline] pub fn new(target_process: PRKPROCESS) -> Self { let mut apc_state: KAPC_STATE = unsafe { core::mem::zeroed() }; @@ -21,7 +43,11 @@ impl ProcessAttach { } } - // Method for manually detaching the process + /// Manually detaches from the process context. + /// + /// This method can be called to explicitly detach the current thread from the target process's + /// address space, if it was previously attached. + /// #[inline] pub fn detach(&mut self) { if self.attached { @@ -34,8 +60,13 @@ impl ProcessAttach { } impl Drop for ProcessAttach { + /// Automatically detaches from the process context when the `ProcessAttach` instance is dropped. + /// + /// This method ensures that the current thread is detached from the target process's address space + /// when the `ProcessAttach` object goes out of scope. If the process is still attached when `drop` + /// is called, it will be safely detached using `KeUnstackDetachProcess`. fn drop(&mut self) { - // If it is still attached, it unattaches when it leaves the scope + // If it is still attached, it unattaches when it leaves the scope. if self.attached { unsafe { KeUnstackDetachProcess(&mut self.apc_state); diff --git a/driver/src/utils/uni.rs b/driver/src/utils/uni.rs index 3a6373b..94de7d9 100644 --- a/driver/src/utils/uni.rs +++ b/driver/src/utils/uni.rs @@ -2,8 +2,11 @@ use alloc::vec::Vec; use wdk_sys::UNICODE_STRING; /// A wrapper around a Vec that represents a unicode string +#[derive(Default)] pub(crate) struct OwnedUnicodeString { + /// buffer: Vec, + /// _phantompinned: core::marker::PhantomPinned, } @@ -21,6 +24,9 @@ impl OwnedUnicodeString { } /// Creates a new OwnedUnicodeString from a rust string. The string is converted to a wide string and null-terminated. +/// +/// +/// pub(crate) fn str_to_unicode(s: &str) -> OwnedUnicodeString { // Convert the rust string to a wide string let mut wide_string: Vec = s.encode_utf16().collect();