diff --git a/driver/callbacks.c b/driver/callbacks.c index f58ff01..8f52c31 100644 --- a/driver/callbacks.c +++ b/driver/callbacks.c @@ -5,46 +5,55 @@ #include "queue.h" -PQUEUE_HEAD report_queue = NULL; - -QUEUE_HEAD test_queue = { 0 }; +QUEUE_HEAD head = { 0 }; +/* +* This mutex is to prevent a new item being pushed to the queue +* while the HandlePeriodicCallbackReportQueue is iterating through +* the objects. This can be an issue because the spinlock is released +* after each report is placed in the IRP buffer which means a new report +* can be pushed into the queue before the next iteration can take ownership +* of the spinlock. +*/ KGUARDED_MUTEX mutex; -VOID InitCallbackReportQueue( PBOOLEAN Status ) +VOID InitCallbackReportQueue( + _In_ PBOOLEAN Status +) { - //report_queue = QueueCreate(); + head.start = NULL; + head.end = NULL; + head.entries = 0; - test_queue.start = NULL; - test_queue.end = NULL; - test_queue.entries = 0; - KeInitializeSpinLock( &test_queue.lock ); + KeInitializeSpinLock( &head.lock ); KeInitializeGuardedMutex( &mutex ); - //if ( report_queue == NULL ) - //{ - // *Status = FALSE; - // return; - //} - *Status = TRUE; } -//VOID DeleteCallbackReportQueueHead() -//{ -// ExFreePoolWithTag( report_queue, QUEUE_POOL_TAG ); -//} - -VOID InsertReportToQueue( - _In_ POPEN_HANDLE_FAILURE_REPORT Report +VOID InsertReportToQueue( + _In_ POPEN_HANDLE_FAILURE_REPORT Report ) { - QueuePush( &test_queue, Report ); + KeAcquireGuardedMutex( &mutex ); + QueuePush( &head, Report ); + KeReleaseGuardedMutex( &mutex ); } -POPEN_HANDLE_FAILURE_REPORT PopFirstReportFromQueue() +VOID FreeQueueObjectsAndCleanup() { - return QueuePop( &test_queue ); + KeAcquireGuardedMutex( &mutex ); + + PVOID report = QueuePop(&head ); + + if ( report == NULL ) + goto end; + + while ( report != NULL ) + report = QueuePop( &head ); + +end: + KeReleaseGuardedMutex( &mutex ); } NTSTATUS HandlePeriodicCallbackReportQueue( @@ -56,7 +65,7 @@ NTSTATUS HandlePeriodicCallbackReportQueue( OPEN_HANDLE_FAILURE_REPORT_HEADER header; KeAcquireGuardedMutex( &mutex ); - report = PopFirstReportFromQueue(); + report = QueuePop( &head ); if ( report == NULL ) { @@ -79,11 +88,12 @@ NTSTATUS HandlePeriodicCallbackReportQueue( sizeof( OPEN_HANDLE_FAILURE_REPORT ) ); - report = PopFirstReportFromQueue(); + report = QueuePop( &head ); count += 1; } end: + header.count = count; RtlCopyMemory( Irp->AssociatedIrp.SystemBuffer, &header, sizeof( OPEN_HANDLE_FAILURE_REPORT_HEADER )); KeReleaseGuardedMutex( &mutex ); diff --git a/driver/callbacks.h b/driver/callbacks.h index f1daa59..14ba8ec 100644 --- a/driver/callbacks.h +++ b/driver/callbacks.h @@ -55,8 +55,9 @@ OB_PREOP_CALLBACK_STATUS ObPreOpCallbackRoutine( _In_ POB_PRE_OPERATION_INFORMATION OperationInformation ); -VOID InitCallbackReportQueue(PBOOLEAN Status); -VOID DeleteCallbackReportQueueHead(); +VOID InitCallbackReportQueue( + _In_ PBOOLEAN Status +); NTSTATUS HandlePeriodicCallbackReportQueue( _In_ PIRP Irp @@ -68,4 +69,6 @@ VOID ProcessCreateNotifyRoutine( _In_ BOOLEAN Create ); +VOID FreeQueueObjectsAndCleanup(); + #endif diff --git a/driver/driver.c b/driver/driver.c index 17d5070..278b7fd 100644 --- a/driver/driver.c +++ b/driver/driver.c @@ -55,6 +55,7 @@ VOID DriverUnload( { PsSetCreateProcessNotifyRoutine( ProcessCreateNotifyRoutine, TRUE ); ObUnRegisterCallbacks( callback_registration_handle ); + FreeQueueObjectsAndCleanup(); IoDeleteSymbolicLink( &DEVICE_SYMBOLIC_LINK ); IoDeleteDevice( DriverObject->DeviceObject ); } @@ -105,7 +106,7 @@ NTSTATUS DriverEntry( { UNREFERENCED_PARAMETER( RegistryPath ); - BOOLEAN flag; + BOOLEAN flag = FALSE; NTSTATUS status; HANDLE handle; @@ -164,7 +165,6 @@ NTSTATUS DriverEntry( if ( !NT_SUCCESS( status ) ) { DEBUG_ERROR( "failed to launch thread to start tings" ); - //DeleteCallbackReportQueueHead(); IoDeleteSymbolicLink( &DEVICE_SYMBOLIC_LINK ); IoDeleteDevice( DriverObject->DeviceObject ); return STATUS_FAILED_DRIVER_ENTRY; diff --git a/driver/ioctl.c b/driver/ioctl.c index f848b39..f1fa2bb 100644 --- a/driver/ioctl.c +++ b/driver/ioctl.c @@ -82,9 +82,9 @@ NTSTATUS DeviceControl( /* KeWaitForSingleObject with infinite time must be called from IRQL <= APC_LEVEL */ PAGED_CODE(); - DEBUG_LOG( "waiting for thread to finish" ); + KeWaitForSingleObject( thread, Executive, KernelMode, FALSE, NULL ); - DEBUG_LOG( "THREAD FINISHED" ); + ZwClose( handle ); ObDereferenceObject( thread ); @@ -131,6 +131,7 @@ NTSTATUS DeviceCreate( _In_ PIRP Irp ) { + DEBUG_LOG( "Handle opened to DonnaAC" ); IoCompleteRequest( Irp, IO_NO_INCREMENT ); return Irp->IoStatus.Status; } \ No newline at end of file diff --git a/driver/nmi.c b/driver/nmi.c index 6029fec..253731f 100644 --- a/driver/nmi.c +++ b/driver/nmi.c @@ -3,6 +3,8 @@ #include "modules.h" #include "common.h" +#define NMI_DELAY 100 * 10000 + typedef struct _NMI_POOLS { PVOID thread_data_pool; @@ -208,7 +210,7 @@ NTSTATUS LaunchNonMaskableInterrupt( } LARGE_INTEGER delay = { 0 }; - delay.QuadPart -= 100 * 10000; + delay.QuadPart -= NMI_DELAY; for ( ULONG core = 0; core < NumCores; core++ ) { diff --git a/driver/queue.c b/driver/queue.c index dfa208b..512ac2d 100644 --- a/driver/queue.c +++ b/driver/queue.c @@ -1,24 +1,20 @@ #include "queue.h" -/* -* Basic thread-safe queue implementation that can be used to create queues for any data. -*/ - -PQUEUE_HEAD QueueCreate() -{ - PQUEUE_HEAD head = ExAllocatePool2( POOL_FLAG_NON_PAGED, sizeof( PQUEUE_HEAD ), QUEUE_POOL_TAG ); - - if ( !head ) - return NULL; - - head->end = NULL; - head->start = NULL; - head->entries = 0; - - KeInitializeSpinLock( &head->lock ); - - return head; -} +//PQUEUE_HEAD QueueCreate() +//{ +// PQUEUE_HEAD head = ExAllocatePool2( POOL_FLAG_NON_PAGED, sizeof( PQUEUE_HEAD ), QUEUE_POOL_TAG ); +// +// if ( !head ) +// return NULL; +// +// head->end = NULL; +// head->start = NULL; +// head->entries = 0; +// +// KeInitializeSpinLock( &head->lock ); +// +// return head; +//} VOID QueuePush( _In_ PQUEUE_HEAD Head, diff --git a/driver/queue.h b/driver/queue.h index d31b998..3cec39c 100644 --- a/driver/queue.h +++ b/driver/queue.h @@ -21,8 +21,6 @@ typedef struct QUEUE_HEAD }QUEUE_HEAD, *PQUEUE_HEAD; -PQUEUE_HEAD QueueCreate(); - VOID QueuePush( _In_ PQUEUE_HEAD Head, _In_ PVOID Data