e

driver/callbacks.c

@@ -5,6 +5,7 @@
 
 #include "queue.h"
 
+CALLBACK_CONFIGURATION configuration;
 QUEUE_HEAD head = { 0 };
 
 /*
@@ -117,6 +118,8 @@ OB_PREOP_CALLBACK_STATUS ObPreOpCallbackRoutine(
 	_In_ POB_PRE_OPERATION_INFORMATION OperationInformation
 )
 {
+	KeAcquireGuardedMutex( &configuration.mutex );
+
 	UNREFERENCED_PARAMETER( RegistrationContext );
 
 	/* access mask to completely strip permissions */
@@ -178,6 +181,7 @@ OB_PREOP_CALLBACK_STATUS ObPreOpCallbackRoutine(
 
 end:
 
+	KeReleaseGuardedMutex( &configuration.mutex );
 	return OB_PREOP_SUCCESS;
 }
 
@@ -459,3 +463,57 @@ VOID EnumerateProcessListWithCallbackFunction(
 
 	} while ( current_process != base_process || !current_process );
 }
+
+NTSTATUS InitiateDriverCallbacks()
+{
+	NTSTATUS status;
+
+	/*
+	* This mutex ensures we don't  unregister our ObRegisterCallbacks while
+	* the callback function is running since this might cause some funny stuff
+	* to happen. Better to be safe then sorry :)
+	*/
+	KeInitializeGuardedMutex( &configuration.mutex );
+
+	OB_CALLBACK_REGISTRATION callback_registration = { 0 };
+	OB_OPERATION_REGISTRATION operation_registration = { 0 };
+
+	operation_registration.ObjectType = PsProcessType;
+	operation_registration.Operations = OB_OPERATION_HANDLE_CREATE | OB_OPERATION_HANDLE_DUPLICATE;
+	operation_registration.PreOperation = ObPreOpCallbackRoutine;
+	operation_registration.PostOperation = ObPostOpCallbackRoutine;
+
+	callback_registration.Version = OB_FLT_REGISTRATION_VERSION;
+	callback_registration.OperationRegistration = &operation_registration;
+	callback_registration.OperationRegistrationCount = 1;
+	callback_registration.RegistrationContext = NULL;
+
+	status = ObRegisterCallbacks(
+		&callback_registration,
+		&configuration.registration_handle
+	);
+
+	if ( !NT_SUCCESS( status ) )
+	{
+		DEBUG_ERROR( "failed to launch obregisters with status %x", status );
+		return status;
+	}
+
+	//status = PsSetCreateProcessNotifyRoutine(
+	//	ProcessCreateNotifyRoutine,
+	//	FALSE
+	//);
+
+	//if ( !NT_SUCCESS( status ) )
+	//	DEBUG_ERROR( "Failed to launch ps create notif routines with status %x", status );
+
+	return status;
+}
+
+VOID UnregisterCallbacksOnProcessTermination()
+{
+	KeAcquireGuardedMutex( &configuration.mutex );
+	ObUnRegisterCallbacks( configuration.registration_handle );
+	configuration.registration_handle = NULL;
+	KeReleaseGuardedMutex( &configuration.mutex );
+}

driver/callbacks.h

@@ -30,6 +30,13 @@ typedef struct _OPEN_HANDLE_FAILURE_REPORT
 
 }OPEN_HANDLE_FAILURE_REPORT, *POPEN_HANDLE_FAILURE_REPORT;
 
+typedef struct _CALLBACKS_CONFIGURATION
+{
+	PVOID registration_handle;
+	KGUARDED_MUTEX mutex;
+
+}CALLBACK_CONFIGURATION, *PCALLBACK_CONFIGURATION;
+
 //handle access masks
 //https://learn.microsoft.com/en-us/windows/win32/procthread/process-security-and-access-rights
 #define PROCESS_CREATE_PROCESS 0x0080
@@ -87,4 +94,8 @@ NTSTATUS EnumerateProcessHandles(
 	_In_ PEPROCESS Process
 );
 
+NTSTATUS InitiateDriverCallbacks();
+
+VOID UnregisterCallbacksOnProcessTermination();
+
 #endif

driver/driver.c

@@ -8,9 +8,6 @@
 
 #include "integrity.h"
 
-
-PVOID callback_registration_handle;
-
 DRIVER_CONFIG config = { 0 };
 
 UNICODE_STRING DEVICE_NAME = RTL_CONSTANT_STRING( L"\\Device\\DonnaAC" );
@@ -34,9 +31,7 @@ VOID GetProtectedProcessId(
 	KeReleaseGuardedMutex( &config.lock );
 }
 
-VOID ClearDriverConfigOnProcessTermination(
+VOID ClearDriverConfigOnProcessTermination()
-	_In_ PIRP Irp
-)
 {
 	KeAcquireGuardedMutex( &config.lock );
 	config.protected_process_id = NULL;
@@ -74,51 +69,11 @@ VOID DriverUnload(
 )
 {
 	//PsSetCreateProcessNotifyRoutine( ProcessCreateNotifyRoutine, TRUE );
-	ObUnRegisterCallbacks( callback_registration_handle );
 	FreeQueueObjectsAndCleanup();
 	IoDeleteSymbolicLink( &DEVICE_SYMBOLIC_LINK );
 	IoDeleteDevice( DriverObject->DeviceObject );
 }
 
-NTSTATUS InitiateDriverCallbacks()
-{
-	NTSTATUS status;
-
-	OB_CALLBACK_REGISTRATION callback_registration = { 0 };
-	OB_OPERATION_REGISTRATION operation_registration = { 0 };
-
-	operation_registration.ObjectType = PsProcessType;
-	operation_registration.Operations = OB_OPERATION_HANDLE_CREATE | OB_OPERATION_HANDLE_DUPLICATE;
-	operation_registration.PreOperation = ObPreOpCallbackRoutine;
-	operation_registration.PostOperation = ObPostOpCallbackRoutine;
-
-	callback_registration.Version = OB_FLT_REGISTRATION_VERSION;
-	callback_registration.OperationRegistration = &operation_registration;
-	callback_registration.OperationRegistrationCount = 1;
-	callback_registration.RegistrationContext = NULL;
-
-	status = ObRegisterCallbacks(
-		&callback_registration,
-		&callback_registration_handle
-	);
-
-	if ( !NT_SUCCESS( status ) )
-	{
-		DEBUG_ERROR( "failed to launch obregisters with status %x", status );
-		return status;
-	}
-
-	//status = PsSetCreateProcessNotifyRoutine(
-	//	ProcessCreateNotifyRoutine,
-	//	FALSE
-	//);
-
-	//if ( !NT_SUCCESS( status ) )
-	//	DEBUG_ERROR( "Failed to launch ps create notif routines with status %x", status );
-
-	return status;
-}
-
 NTSTATUS DriverEntry(
 	_In_ PDRIVER_OBJECT DriverObject,
 	_In_ PUNICODE_STRING RegistryPath

driver/driver.h

@@ -28,10 +28,5 @@ VOID GetProtectedProcessId(
 );
 
 
-VOID ClearDriverConfigOnProcessTermination(
+VOID ClearDriverConfigOnProcessTermination();
-	_In_ PIRP Irp
-);
-
-NTSTATUS InitiateDriverCallbacks();
-
 #endif

driver/ioctl.c

@@ -192,7 +192,8 @@ NTSTATUS DeviceControl(
 
 	case IOCTL_CLEAR_CONFIG_ON_PROCESS_CLOSE:
 
-		ClearDriverConfigOnProcessTermination( Irp );
+		ClearDriverConfigOnProcessTermination();
+		UnregisterCallbacksOnProcessTermination();
 
 		break;