yay

driver/ioctl.c

@@ -15,6 +15,7 @@ NTSTATUS DeviceControl(
 	NTSTATUS status = STATUS_SUCCESS;
 	PIO_STACK_LOCATION stack_location = IoGetCurrentIrpStackLocation( Irp );
 	HANDLE handle;
+	PKTHREAD thread;
 
 	switch ( stack_location->Parameters.DeviceIoControl.IoControlCode )
 	{
@@ -29,6 +30,8 @@ NTSTATUS DeviceControl(
 
 	case IOCTL_VALIDATE_DRIVER_OBJECTS:
 
+		PAGED_CODE();
+
 		/*
 		* The reason this function is run in a new thread and not the thread
 		* issuing the IOCTL is because ZwOpenDirectoryObject issues a
@@ -36,7 +39,6 @@ NTSTATUS DeviceControl(
 		* This is a problem because when we pass said handle to ObReferenceObjectByHandle
 		* it will issue a bug check under windows driver verifier.
 		*/
-
 		status = PsCreateSystemThread(
 			&handle,
 			PROCESS_ALL_ACCESS,
@@ -48,10 +50,38 @@ NTSTATUS DeviceControl(
 		);
 
 		if ( !NT_SUCCESS( status ) )
+		{
 			DEBUG_ERROR( "Failed to start thread to validate system drivers" );
+			goto end;
+		}
+
+		/*
+		* Thread objects are a type of dispatcher object, meaning when they are freed
+		* its set to the signal state and any waiters will be signalled. This allows
+		* us to wait til our threads terminated and the IRP buffer has been either filled
+		* or left empty and then from there we can complete the IRP and return.
+		*/
+		status = ObReferenceObjectByHandle(
+			handle,
+			THREAD_ALL_ACCESS,
+			*PsThreadType,
+			KernelMode,
+			&thread,
+			NULL
+		);
+
+		if ( !NT_SUCCESS( status ) )
+		{
+			DEBUG_ERROR( "ObReferenceObjectbyhandle failed with status %lx", status );
+			ZwClose( handle );
+			goto end;
+		}
+
+		KeWaitForSingleObject( thread, Executive, KernelMode, FALSE, NULL );
 
-		/* return early as IRP completion was handled inside the function */
 		ZwClose( handle );
+		ObDereferenceObject( thread );
+
 		break;
 
 	default:
@@ -59,6 +89,8 @@ NTSTATUS DeviceControl(
 		break;
 	}
 
+end:
+
 	IoCompleteRequest( Irp, IO_NO_INCREMENT );
 	Irp->IoStatus.Status = status;
 	return status;

driver/modules.c

@@ -297,9 +297,9 @@ NTSTATUS HandleValidateDriversIOCTL(
 		Irp->IoStatus.Information = sizeof( MODULE_VALIDATION_FAILURE_HEADER ) +
 			MODULE_VALIDATION_FAILURE_MAX_REPORT_COUNT * sizeof( MODULE_VALIDATION_FAILURE );
 
-		RtlCopyMemory( 
+		RtlCopyMemory(
-			Irp->AssociatedIrp.SystemBuffer, 
+			Irp->AssociatedIrp.SystemBuffer,
-			&header, 
+			&header,
 			sizeof( MODULE_VALIDATION_FAILURE_HEADER ) );
 
 		for ( INT i = 0; i < head->count; i++ )
@@ -312,16 +312,16 @@ NTSTATUS HandleValidateDriversIOCTL(
 			report.report_code = REPORT_MODULE_VALIDATION_FAILURE;
 			report.driver_base_address = head->first_entry->driver->DriverStart;
 			report.driver_size = head->first_entry->driver->Size;
-			
+
 			RtlCopyMemory(
 				&report.driver_name,
 				head->first_entry->driver->DriverName.Buffer,
 				MODULE_REPORT_DRIVER_NAME_BUFFER_SIZE );
 
 			RtlCopyMemory(
-				(UINT64)Irp->AssociatedIrp.SystemBuffer + sizeof( MODULE_VALIDATION_FAILURE_HEADER ) + i * sizeof( MODULE_VALIDATION_FAILURE ),
+				( UINT64 )Irp->AssociatedIrp.SystemBuffer + sizeof( MODULE_VALIDATION_FAILURE_HEADER ) + i * sizeof( MODULE_VALIDATION_FAILURE ),
 				&report,
-				sizeof( MODULE_VALIDATION_FAILURE ));
+				sizeof( MODULE_VALIDATION_FAILURE ) );
 
 			RemoveInvalidDriverFromList( head );
 		}
@@ -338,8 +338,8 @@ NTSTATUS HandleValidateDriversIOCTL(
 	* Complete the IRP here so we don't have to implement a waiting mechanism
 	* to prevent an early completion of the IRP.
 	*/
-	IoCompleteRequest( Irp, IO_NO_INCREMENT );
+	//IoCompleteRequest( Irp, IO_NO_INCREMENT );
-	Irp->IoStatus.Status = status;
+	//Irp->IoStatus.Status = status;
 
 	return status;
 }

user/km/driver.cpp

@@ -2,7 +2,7 @@
 
 #include "../common.h"
 
-kernelmode::Driver::Driver(LPCWSTR DriverName, std::shared_ptr<global::Report> ReportInterface )
+kernelmode::Driver::Driver( LPCWSTR DriverName, std::shared_ptr<global::Report> ReportInterface )
 {
 	this->driver_name = DriverName;
 	this->report_interface = ReportInterface;
@@ -56,7 +56,7 @@ void kernelmode::Driver::RunNmiCallbacks()
 void kernelmode::Driver::VerifySystemModules()
 {
 	BOOLEAN status;
-	DWORD bytes_returned = 0;
+	DWORD bytes_returned;
 	PVOID buffer;
 	SIZE_T buffer_size;
 	SIZE_T header_size;
@@ -70,7 +70,7 @@ void kernelmode::Driver::VerifySystemModules()
 	*/
 	header_size = sizeof( global::report_structures::MODULE_VALIDATION_FAILURE_HEADER );
 
-	buffer_size = sizeof( global::report_structures::MODULE_VALIDATION_FAILURE ) * 
+	buffer_size = sizeof( global::report_structures::MODULE_VALIDATION_FAILURE ) *
 		MODULE_VALIDATION_FAILURE_MAX_REPORT_COUNT +
 		header_size;
 
@@ -87,7 +87,7 @@ void kernelmode::Driver::VerifySystemModules()
 		buffer,
 		buffer_size,
 		&bytes_returned,
-		( LPOVERLAPPED )NULL
+		NULL
 	);
 
 	if ( status == NULL )
@@ -102,10 +102,10 @@ void kernelmode::Driver::VerifySystemModules()
 		LOG_INFO( "All modules valid :)" );
 		free( buffer );
 		return;
-	}	
+	}
 
 	/*
-	* We are splitting up each packet here and passing them on one by one since 
+	* We are splitting up each packet here and passing them on one by one since
 	* if I am being honest it is just easier in c++ and that way the process
 	* is streamlined just like all other report packets.
 	*/
@@ -115,10 +115,10 @@ void kernelmode::Driver::VerifySystemModules()
 
 	for ( int i = 0; i < header.module_count; i++ )
 	{
-		memcpy( 
+		memcpy(
-			&report, 
+			&report,
 			PVOID( base + i * sizeof( global::report_structures::MODULE_VALIDATION_FAILURE ) ),
-			sizeof( global::report_structures::MODULE_VALIDATION_FAILURE ) 
+			sizeof( global::report_structures::MODULE_VALIDATION_FAILURE )
 		);
 
 		this->report_interface->ReportViolation( &report );