diff --git a/driver/src/callback.rs b/driver/src/callback.rs
index 854b49b..ded059c 100644
--- a/driver/src/callback.rs
+++ b/driver/src/callback.rs
@@ -159,11 +159,18 @@ impl<'a> Callback<'a> {
     pub fn unload() {
         unsafe {
             // Unregister process and thread creation callbacks
-            ObUnRegisterCallbacks(CALLBACK_REGISTRATION_HANDLE_PROCESS);
-            ObUnRegisterCallbacks(CALLBACK_REGISTRATION_HANDLE_THREAD);
-    
+            if !CALLBACK_REGISTRATION_HANDLE_PROCESS.is_null() {
+                ObUnRegisterCallbacks(CALLBACK_REGISTRATION_HANDLE_PROCESS);
+            }
+
+            if !CALLBACK_REGISTRATION_HANDLE_THREAD.is_null() {
+                ObUnRegisterCallbacks(CALLBACK_REGISTRATION_HANDLE_THREAD);
+            }
+
             // Unregister registry modification callback
-            CmUnRegisterCallback(CALLBACK_REGISTRY);
+            if CALLBACK_REGISTRY.QuadPart != 0 {
+                CmUnRegisterCallback(CALLBACK_REGISTRY);
+            }
     
             // Unregister bug check (crash dump) callback
             KeDeregisterBugCheckReasonCallback(&mut BUG_CHECK);
diff --git a/driver/src/ioctls.rs b/driver/src/ioctls.rs
index 0787821..fe15553 100644
--- a/driver/src/ioctls.rs
+++ b/driver/src/ioctls.rs
@@ -269,10 +269,13 @@ impl IoctlManager {
         
                 // Write the USER_ADDRESS to the output buffer provided by the IRP.
                 let output_buffer = (*irp).AssociatedIrp.SystemBuffer;
-                if !output_buffer.is_null() {
-                    *(output_buffer as *mut usize) = USER_ADDRESS;
+                if output_buffer.is_null() {
+                    log::error!("IRP SystemBuffer is null");
+                    return Ok(STATUS_UNSUCCESSFUL);
                 }
-        
+
+                *(output_buffer as *mut usize) = USER_ADDRESS;
+
                 // Set the number of bytes returned to the size of a `usize`.
                 (*irp).IoStatus.Information = size_of::<usize>() as u64;
                 Ok(STATUS_SUCCESS)
diff --git a/driver/src/utils.rs b/driver/src/utils.rs
index eed3d49..5c2cc01 100644
--- a/driver/src/utils.rs
+++ b/driver/src/utils.rs
@@ -29,7 +29,7 @@ pub unsafe fn get_input_buffer<T>(stack: *mut _IO_STACK_LOCATION) -> Result<*mut
         return Err(ShadowError::BufferTooSmall);
     }
 
-    //  Allocate a kernel-mode buffer in non-paged memory
+    // 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() {
         return Err(ShadowError::NullPointer("buffer"));
diff --git a/shadowx/src/utils/address.rs b/shadowx/src/utils/address.rs
index 5805406..b19862b 100644
--- a/shadowx/src/utils/address.rs
+++ b/shadowx/src/utils/address.rs
@@ -56,10 +56,7 @@ pub unsafe fn get_module_base_address(module_name: &str) -> Result<*mut c_void>
     );
 
     if !NT_SUCCESS(status) {
-        return Err(ShadowError::ApiCallFailed(
-            "ZwQuerySystemInformation",
-            status,
-        ));
+        return Err(ShadowError::ApiCallFailed("ZwQuerySystemInformation", status));
     }
 
     // Iterates over the list of modules to find the one that matches the provided name