BLAH

driver/callbacks.c

@@ -512,6 +512,7 @@ NTSTATUS InitiateDriverCallbacks()
 
 VOID UnregisterCallbacksOnProcessTermination()
 {
+	DEBUG_LOG( "Process closed, unregistering callbacks" );
 	KeAcquireGuardedMutex( &configuration.mutex );
 	ObUnRegisterCallbacks( configuration.registration_handle );
 	configuration.registration_handle = NULL;

driver/driver.c

@@ -13,8 +13,17 @@ DRIVER_CONFIG config = { 0 };
 UNICODE_STRING DEVICE_NAME = RTL_CONSTANT_STRING( L"\\Device\\DonnaAC" );
 UNICODE_STRING DEVICE_SYMBOLIC_LINK = RTL_CONSTANT_STRING( L"\\??\\DonnaAC" );
 
+VOID ReadInitialisedConfigFlag(
+	_Out_ PBOOLEAN Flag
+)
+{
+	KeAcquireGuardedMutex( &config.lock );
+	*Flag = config.initialised;
+	KeReleaseGuardedMutex( &config.lock );
+}
+
 VOID GetProtectedProcessEProcess( 
-	_In_ PEPROCESS Process 
+	_Out_ PEPROCESS Process 
 )
 {
 	KeAcquireGuardedMutex( &config.lock );
@@ -23,7 +32,7 @@ VOID GetProtectedProcessEProcess(
 }
 
 VOID GetProtectedProcessId( 
-	_In_ PLONG ProcessId 
+	_Out_ PLONG ProcessId 
 )
 {
 	KeAcquireGuardedMutex( &config.lock );
@@ -33,6 +42,7 @@ VOID GetProtectedProcessId(
 
 VOID ClearDriverConfigOnProcessTermination()
 {
+	DEBUG_LOG( "Process closed, clearing driver configuration" );
 	KeAcquireGuardedMutex( &config.lock );
 	config.protected_process_id = NULL;
 	config.protected_process_eprocess = NULL;
@@ -55,10 +65,18 @@ NTSTATUS InitialiseDriverConfigOnProcessLaunch(
 	if ( !NT_SUCCESS( status ) )
 		return status;
 
+	/*
+	* acquire the mutex here to prevent a race condition if an unknown party trys 
+	* to fuzz our IOCTL codes whilst the target process launches.
+	*/
+	KeAcquireGuardedMutex( &config.lock );
+
 	config.protected_process_eprocess = eprocess;
 	config.protected_process_id = information->protected_process_id;
 	config.initialised = TRUE;
 
+	KeReleaseGuardedMutex( &config.lock );
+
 	Irp->IoStatus.Status = status;
 
 	return status;
@@ -86,6 +104,10 @@ NTSTATUS DriverEntry(
 
 	KeInitializeGuardedMutex( &config.lock );
 
+	config.initialised = FALSE;
+	config.protected_process_eprocess = NULL;
+	config.protected_process_id = NULL;
+
 	status = IoCreateDevice(
 		DriverObject,
 		NULL,

driver/driver.h

@@ -19,12 +19,16 @@ NTSTATUS InitialiseDriverConfigOnProcessLaunch(
 );
 
 VOID GetProtectedProcessEProcess(
-	_In_ PEPROCESS Process
+	_Out_ PEPROCESS Process
 );
 
 
 VOID GetProtectedProcessId(
-	_In_ PLONG ProcessId
+	_Out_ PLONG ProcessId
+);
+
+VOID ReadInitialisedConfigFlag(
+	_Out_ PBOOLEAN Flag
 );
 
 

driver/integrity.c

@@ -43,6 +43,8 @@ NTSTATUS GetDriverImageSize(
 * analyse the executable sections from there. Until I find a better way to enumerate
 * kernel memory without having to walk the pages tables to check the EDB bit this 
 * is how I will be doing it. c:
+* 
+* TODO: We will hash this based on timestamp sent from the server.
 */
 NTSTATUS CopyDriverExecutableRegions(
 	_In_ PIRP Irp

driver/ioctl.c

@@ -21,6 +21,17 @@ NTSTATUS DeviceControl(
 	PIO_STACK_LOCATION stack_location = IoGetCurrentIrpStackLocation( Irp );
 	HANDLE handle;
 	PKTHREAD thread;
+	BOOLEAN security_flag = FALSE;
+
+	/*
+	* The purpose of this is to prevent programs from opening a handle to our driver
+	* and trying to fuzz the IOCTL access or codes. This definitely isnt a perfect 
+	* solution though... xD
+	*/
+	ReadInitialisedConfigFlag( &security_flag );
+
+	if ( security_flag == FALSE )
+		goto end;
 
 	switch ( stack_location->Parameters.DeviceIoControl.IoControlCode )
 	{
@@ -190,7 +201,7 @@ NTSTATUS DeviceControl(
 
 		break;
 
-	case IOCTL_CLEAR_CONFIG_ON_PROCESS_CLOSE:
+	case IOCTL_NOTIFY_DRIVER_ON_PROCESS_TERMINATION:
 
 		ClearDriverConfigOnProcessTermination();
 		UnregisterCallbacksOnProcessTermination();

driver/ioctl.h

@@ -15,7 +15,7 @@
 #define IOCTL_ENUMERATE_HANDLE_TABLES CTL_CODE(FILE_DEVICE_UNKNOWN, 0x2007, METHOD_BUFFERED, FILE_ANY_ACCESS)
 #define IOCTL_RETRIEVE_MODULE_EXECUTABLE_REGIONS CTL_CODE(FILE_DEVICE_UNKNOWN, 0x2008, METHOD_BUFFERED, FILE_ANY_ACCESS)
 #define IOCTL_REQUEST_TOTAL_MODULE_SIZE CTL_CODE(FILE_DEVICE_UNKNOWN, 0x2009, METHOD_BUFFERED, FILE_ANY_ACCESS)
-#define IOCTL_CLEAR_CONFIG_ON_PROCESS_CLOSE CTL_CODE(FILE_DEVICE_UNKNOWN, 0x2010, METHOD_BUFFERED, FILE_ANY_ACCESS)
+#define IOCTL_NOTIFY_DRIVER_ON_PROCESS_TERMINATION CTL_CODE(FILE_DEVICE_UNKNOWN, 0x2010, METHOD_BUFFERED, FILE_ANY_ACCESS)
 
 typedef struct _DRIVER_INITIATION_INFORMATION
 {

user/km/driver.cpp

@@ -24,6 +24,11 @@ kernelmode::Driver::Driver( LPCWSTR DriverName, std::shared_ptr<global::Client>
 	this->NotifyDriverOnProcessLaunch();
 }
 
+kernelmode::Driver::~Driver()
+{
+	this->NotifyDriverOnProcessTermination();
+}
+
 VOID kernelmode::Driver::RunNmiCallbacks()
 {
 	BOOLEAN status;
@@ -356,6 +361,25 @@ ULONG kernelmode::Driver::RequestTotalModuleSize()
 	return module_size;
 }
 
+VOID kernelmode::Driver::NotifyDriverOnProcessTermination()
+{
+	BOOLEAN status;
+	
+	status = DeviceIoControl(
+		this->driver_handle,
+		IOCTL_NOTIFY_DRIVER_ON_PROCESS_TERMINATION,
+		NULL,
+		NULL,
+		NULL,
+		NULL,
+		NULL,
+		NULL
+	);
+
+	if ( status == NULL )
+		LOG_ERROR( "NotifyDriverOnProcessTermination failed with status %x", status );
+}
+
 VOID kernelmode::Driver::ValidateKPRCBThreads()
 {
 

user/km/driver.h

@@ -8,14 +8,13 @@
 
 #define IOCCTL_RUN_NMI_CALLBACKS CTL_CODE(FILE_DEVICE_UNKNOWN, 0x2001, METHOD_BUFFERED, FILE_ANY_ACCESS)
 #define IOCTL_VALIDATE_DRIVER_OBJECTS CTL_CODE(FILE_DEVICE_UNKNOWN, 0x2002, METHOD_BUFFERED, FILE_ANY_ACCESS)
-#define IOCTL_MONITOR_CALLBACKS_FOR_REPORTS CTL_CODE(FILE_DEVICE_UNKNOWN, 0x2003, METHOD_BUFFERED, FILE_ANY_ACCESS)
 #define IOCTL_NOTIFY_DRIVER_ON_PROCESS_LAUNCH CTL_CODE(FILE_DEVICE_UNKNOWN, 0x2004, METHOD_BUFFERED, FILE_ANY_ACCESS)
 #define IOCTL_HANDLE_REPORTS_IN_CALLBACK_QUEUE CTL_CODE(FILE_DEVICE_UNKNOWN, 0x2005, METHOD_BUFFERED, FILE_ANY_ACCESS)
 #define IOCTL_PERFORM_VIRTUALIZATION_CHECK CTL_CODE(FILE_DEVICE_UNKNOWN, 0x2006, METHOD_BUFFERED, FILE_ANY_ACCESS)
 #define IOCTL_ENUMERATE_HANDLE_TABLES CTL_CODE(FILE_DEVICE_UNKNOWN, 0x2007, METHOD_BUFFERED, FILE_ANY_ACCESS)
 #define IOCTL_RETRIEVE_MODULE_EXECUTABLE_REGIONS CTL_CODE(FILE_DEVICE_UNKNOWN, 0x2008, METHOD_BUFFERED, FILE_ANY_ACCESS)
 #define IOCTL_REQUEST_TOTAL_MODULE_SIZE CTL_CODE(FILE_DEVICE_UNKNOWN, 0x2009, METHOD_BUFFERED, FILE_ANY_ACCESS)
-
+#define IOCTL_NOTIFY_DRIVER_ON_PROCESS_TERMINATION CTL_CODE(FILE_DEVICE_UNKNOWN, 0x2010, METHOD_BUFFERED, FILE_ANY_ACCESS)
 
 #define MAX_HANDLE_REPORTS_PER_IRP 10
 
@@ -29,18 +28,20 @@ namespace kernelmode
 
 		VOID QueryReportQueue();
 		ULONG RequestTotalModuleSize();
+		VOID NotifyDriverOnProcessLaunch();
+		VOID CheckDriverHeartbeat();
+		VOID NotifyDriverOnProcessTermination();
 
 	public:
 
 		Driver(LPCWSTR DriverName, std::shared_ptr<global::Client> ReportInterface );
+		~Driver();
 
 		VOID RunNmiCallbacks();
 		VOID VerifySystemModules();
 		VOID RunCallbackReportQueue();
-		VOID NotifyDriverOnProcessLaunch();
 		VOID DetectSystemVirtualization();
 		VOID ValidateKPRCBThreads();
-		VOID CheckDriverHeartbeat();
 		VOID CheckHandleTableEntries();
 		VOID RequestModuleExecutableRegions();
 		/* todo: driver integrity check */

user/main.cpp

@@ -31,7 +31,7 @@ DWORD WINAPI Init(HINSTANCE hinstDLL)
     //kmanager.MonitorCallbackReports();
     //kmanager.RunNmiCallbacks();
     //kmanager.VerifySystemModules();
-    kmanager.RequestModuleExecutableRegionsForIntegrityCheck();
+    //kmanager.RequestModuleExecutableRegionsForIntegrityCheck();
     //kmanager.MonitorCallbackReports();
 
     //umanager.ValidateProcessModules();