diff --git a/driver/src/utils.rs b/driver/src/utils.rs
index f97230f..eed3d49 100644
--- a/driver/src/utils.rs
+++ b/driver/src/utils.rs
@@ -30,7 +30,7 @@ pub unsafe fn get_input_buffer<T>(stack: *mut _IO_STACK_LOCATION) -> Result<*mut
     }
 
     //  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;
+    let buffer = ExAllocatePool2(POOL_FLAG_NON_PAGED, size_of::<T>() as u64, 0x1234) as *mut T;
     if buffer.is_null() {
         return Err(ShadowError::NullPointer("buffer"));
     }
@@ -79,6 +79,5 @@ pub unsafe fn get_output_buffer<T>(irp: *mut IRP, stack: *mut _IO_STACK_LOCATION
     }
 
     let count = output_length as usize / size_of::<T>();
-
     Ok((buffer as *mut T, count))
 }
diff --git a/shadowx/src/utils/mod.rs b/shadowx/src/utils/mod.rs
index aa21ad7..15e06f7 100644
--- a/shadowx/src/utils/mod.rs
+++ b/shadowx/src/utils/mod.rs
@@ -334,7 +334,7 @@ pub unsafe fn get_process_by_name(process_name: &str) -> Result<usize> {
 ///
 /// * True if the address is within the kernel memory range, False otherwise.
 pub fn valid_kernel_memory(addr: u64) -> bool {
-    (addr >> 48) == 0xFFFF
+    unsafe { addr >= wdk_sys::MmSystemRangeStart as u64 }
 }
 
 /// Validates if the given address is within the user memory range.
@@ -347,7 +347,7 @@ pub fn valid_kernel_memory(addr: u64) -> bool {
 ///
 /// * True if the address is within the user memory range, False otherwise.
 pub fn valid_user_memory(addr: u64) -> bool {
-    (addr >> 48) == 0x0000
+    unsafe { addr > 0 && addr <= wdk_sys::MmHighestUserAddress as u64 }
 }
 
 /// Responsible for returning information on the modules loaded.