From 566ab137c1cf8795d9632a97638731cfa38f6a53 Mon Sep 17 00:00:00 2001
From: joaoviictorti <joaovictorti08@gmail.com>
Date: Wed, 26 Feb 2025 09:49:18 -0300
Subject: [PATCH] Small fix

---
 driver/src/utils.rs  | 11 ++++++++---
 shadowx/src/error.rs |  7 +++++++
 2 files changed, 15 insertions(+), 3 deletions(-)

diff --git a/driver/src/utils.rs b/driver/src/utils.rs
index 5c2cc01..e5feda1 100644
--- a/driver/src/utils.rs
+++ b/driver/src/utils.rs
@@ -16,19 +16,24 @@ use wdk_sys::{
 /// * `Result<*mut T, ShadowError>` - A result containing the pointer to the input buffer or an NTSTATUS error code.
 pub unsafe fn get_input_buffer<T>(stack: *mut _IO_STACK_LOCATION) -> Result<*mut T, ShadowError> {
     // Retrieves the input buffer pointer from the I/O stack location.
-    let input_buffer= (*stack).Parameters.DeviceIoControl.Type3InputBuffer;
+    let input_buffer = (*stack).Parameters.DeviceIoControl.Type3InputBuffer;
     let input_length = (*stack).Parameters.DeviceIoControl.InputBufferLength;
-    
+
     // Validate that the input buffer is not null
     if input_buffer.is_null() {
         return Err(ShadowError::NullPointer("Type3InputBuffer"))
     } 
     
     // Validate that the input buffer size is sufficient
-    if input_length < size_of::<T>() as u32 || input_length % size_of::<T>() as u32 != 0 {
+    if input_length < size_of::<T>() as u32 {
         return Err(ShadowError::BufferTooSmall);
     }
 
+    // Alignment check
+    if (input_buffer as usize) % align_of::<T>() != 0 {
+        return Err(ShadowError::MisalignedBuffer);
+    }
+
     // Allocate a kernel-mode buffer in non-paged memory
     let buffer = ExAllocatePool2(POOL_FLAG_NON_PAGED, size_of::<T>() as u64, 0x1234) as *mut T;
     if buffer.is_null() {
diff --git a/shadowx/src/error.rs b/shadowx/src/error.rs
index eae0946..6efbea6 100644
--- a/shadowx/src/error.rs
+++ b/shadowx/src/error.rs
@@ -117,6 +117,13 @@ pub enum ShadowError {
     #[error("Small buffer")]
     BufferTooSmall,
 
+    /// Represents an error when a buffer is misaligned for the expected data structure.
+    ///
+    /// This error occurs when the provided buffer does not have the correct memory alignment
+    /// required for safe access.
+    #[error("Misaligned buffer")]
+    MisalignedBuffer,
+
     /// Error indicating that a callback could not be found.
     ///
     /// This occurs when the system is unable to locate the expected callback function.