gah memory leak waaaaa

driver/driver.c

@@ -10,11 +10,140 @@
 #include "modules.h"
 #include "integrity.h"
 
-#include "queue.h"
-
 DRIVER_CONFIG driver_config = { 0 };
 PROCESS_CONFIG process_config = { 0 };
 
+VOID InitApcContextsList()
+{
+	KeAcquireGuardedMutex( &driver_config.lock );
+	driver_config.apc_contexts = ExAllocatePool2( POOL_FLAG_NON_PAGED, sizeof( LIST_HEAD ), POOL_TAG_APC );
+
+	if ( !driver_config.apc_contexts )
+		return;
+
+	ListInit( driver_config.apc_contexts );
+	KeReleaseGuardedMutex( &driver_config.lock );
+}
+
+VOID FreeApcContextStructures()
+{
+	KeAcquireGuardedMutex( &driver_config.lock );
+
+	PLIST_ITEM entry = driver_config.apc_contexts->start;
+
+	if ( !entry )
+		goto unlock;
+
+	while ( entry )
+	{
+		FreeApcContextStructure( entry->data );
+		entry = entry->next;
+		ListRemoveItem( driver_config.apc_contexts, entry );
+	}
+
+unlock:
+	KeReleaseGuardedMutex( &driver_config.lock );
+}
+
+VOID InsertApcIntoApcContextList(
+	_In_ PLIST_HEAD ListHead,
+	_In_ PAPC_STATUS ApcStatus
+)
+{
+	KeAcquireGuardedMutex( &driver_config.lock );
+	ListInsert( ListHead, ApcStatus );
+	KeReleaseGuardedMutex( &driver_config.lock );
+}
+
+VOID RemoveApcContext(
+	_In_ LONG ContextIdentifier,
+	_In_ PKAPC Apc
+)
+{
+	KeAcquireGuardedMutex( &driver_config.lock );
+
+	PLIST_ITEM entry = driver_config.apc_contexts->start;
+
+	while ( entry )
+	{
+		PAPC_CONTEXT_HEADER header = ( PAPC_CONTEXT_HEADER )entry->data;
+
+		if ( header->context_id == ContextIdentifier )
+		{
+			ListRemoveItem( driver_config.apc_contexts, entry );
+			ExFreePoolWithTag( entry->data, POOL_TAG_APC );
+			ExFreePoolWithTag( entry, LIST_POOL_TAG );
+			break;
+		}
+	}
+
+	KeReleaseGuardedMutex( &driver_config.lock );
+}
+
+VOID InsertApcContext(
+	_In_ PVOID Context
+)
+{
+	KeAcquireGuardedMutex( &driver_config.lock );
+
+	if ( Context )
+	{
+		PLIST_ITEM entry = ListInsert( driver_config.apc_contexts, Context );
+
+		if ( !entry )
+			goto end;
+
+		PAPC_CONTEXT_HEADER header = ( PAPC_CONTEXT_HEADER )entry->data;
+
+		header->head = ExAllocatePool2( POOL_FLAG_NON_PAGED, sizeof( LIST_HEAD ), POOL_TAG_APC );
+
+		if ( !header->head )
+			goto end;
+	}
+
+end:
+	KeReleaseGuardedMutex( &driver_config.lock );
+}
+
+VOID RemoveApcFromApcContextList(
+	_In_ PLIST_HEAD ListHead,
+	_Inout_ PLIST_ITEM ListEntry
+)
+{
+	KeAcquireGuardedMutex( &driver_config.lock );
+	ListRemoveItem( ListHead, ListEntry );
+	KeReleaseGuardedMutex( &driver_config.lock );
+}
+
+VOID GetApcContext(
+	_Inout_ PVOID* Context,
+	_In_ LONG ContextIdentifier
+)
+{
+	if ( !Context )
+		return;
+
+	KeAcquireGuardedMutex( &driver_config.lock );
+
+	PLIST_ITEM entry = driver_config.apc_contexts->start;
+
+	while ( entry )
+	{
+		PAPC_CONTEXT_HEADER header = ( PAPC_CONTEXT_HEADER )entry->data;
+
+		if ( header->context_id == ContextIdentifier )
+		{
+			*Context = header;
+			KeReleaseGuardedMutex( &driver_config.lock );
+			return;
+		}
+
+		entry = entry->next;
+	}
+
+	KeReleaseGuardedMutex( &driver_config.lock );
+}
+
 VOID ReadProcessInitialisedConfigFlag(
 	_Out_ PBOOLEAN Flag
 )
@@ -275,6 +404,8 @@ NTSTATUS InitialiseDriverConfigOnDriverEntry(
 		return status;
 	}
 
+	InitApcContextsList();
+
 	DEBUG_LOG( "Motherboard serial: %s", driver_config.system_information.motherboard_serial );
 	DEBUG_LOG( "Drive 0 serial: %s", driver_config.system_information.drive_0_serial );
 
@@ -330,6 +461,7 @@ VOID DriverUnload(
 {
 	//PsSetCreateProcessNotifyRoutine( ProcessCreateNotifyRoutine, TRUE );
 	CleanupDriverConfigOnUnload();
+	FreeApcContextStructures();
 	IoDeleteDevice( DriverObject->DeviceObject );
 }
 

driver/driver.h

@@ -4,7 +4,10 @@
 #include <ntifs.h>
 #include <wdftypes.h>
 #include <wdf.h>
+
 #include "common.h"
+#include "queue.h"
+#include "modules.h"
 
 #define DRIVER_PATH_MAX_LENGTH 512
 #define MOTHERBOARD_SERIAL_CODE_LENGTH 64
@@ -30,6 +33,7 @@ typedef struct _SYSTEM_INFORMATION
 * Note that the lock isnt really needed here but Im using one
 * just in case c:
 */
+
 typedef struct _DRIVER_CONFIG
 {
 	UNICODE_STRING unicode_driver_name;
@@ -39,6 +43,7 @@ typedef struct _DRIVER_CONFIG
 	UNICODE_STRING driver_path;
 	UNICODE_STRING registry_path;
 	SYSTEM_INFORMATION system_information;
+	PLIST_HEAD apc_contexts;
 	KGUARDED_MUTEX lock;
 
 }DRIVER_CONFIG, *PDRIVER_CONFIG;
@@ -82,6 +87,25 @@ VOID GetDriverConfigSystemInformation(
 	_Out_ PSYSTEM_INFORMATION* SystemInformation
 );
 
+VOID GetApcContext(
+	_Inout_ PVOID* Context,
+	_In_ LONG ContextIdentifier
+);
+
+VOID InsertApcContext(
+	_In_ PVOID Context
+);
+
+VOID RemoveApcFromApcContextList(
+	_In_ PLIST_HEAD ListHead,
+	_Inout_ PLIST_ITEM ListEntry
+);
+
+VOID InsertApcIntoApcContextList(
+	_In_ PLIST_HEAD ListHead,
+	_In_ PAPC_STATUS ApcStatus
+);
+
 VOID TerminateProtectedProcessOnViolation();
 
 VOID ClearProcessConfigOnProcessTermination();

driver/modules.c

@@ -1,6 +1,7 @@
 #include "modules.h"
 
 #include "callbacks.h"
+#include "driver.h"
 
 #define WHITELISTED_MODULE_TAG 'whte'
 
@@ -873,22 +874,173 @@ VOID ApcRundownRoutine(
 	_In_ PRKAPC Apc
 )
 {
-	DEBUG_LOG( "Thread discarding APC queue, freeing APC object" );
-	ExFreePoolWithTag( Apc, POOL_TAG_APC );
+
+}
+
+VOID InsertApcToContextList(
+	_In_ PKAPC Apc,
+	_In_ LONG ContextId
+)
+{
+	PAPC_CONTEXT_HEADER header = NULL;
+
+	GetApcContext( &header, ContextId );
+
+	if ( !header )
+		return;
+
+	KeAcquireGuardedMutex( &header->lock );
+
+	switch ( header->context_id )
+	{
+	case APC_CONTEXT_ID_STACKWALK:;
+
+		PAPC_STACKWALK_CONTEXT context = ( PAPC_STACKWALK_CONTEXT )header;
+
+		PAPC_STATUS apc_status = ExAllocatePool2( POOL_FLAG_NON_PAGED, sizeof( APC_STATUS ), POOL_TAG_APC );
+
+		if ( !apc_status )
+			goto unlock;
+
+		apc_status->apc = Apc;
+		apc_status->apc_complete = FALSE;
+
+		InsertApcIntoApcContextList( context->head, apc_status );
+		break;
+	}
+
+unlock:
+
+	KeReleaseGuardedMutex( &header->lock );
+}
+
+VOID FreeApcContextStructure(
+	_Inout_ PVOID Context
+)
+{
+	PAPC_CONTEXT_HEADER header = ( PAPC_CONTEXT_HEADER )Context;
+
+	KeAcquireGuardedMutex( &header->lock );
+
+	switch ( header->context_id )
+	{
+	case APC_CONTEXT_ID_STACKWALK:;
+
+		PAPC_STACKWALK_CONTEXT context = ( PAPC_STACKWALK_CONTEXT )header;
+
+		ExFreePoolWithTag( context->head, POOL_TAG_APC );
+		ExFreePoolWithTag( context->modules, POOL_TAG_APC );
+		break;
+	}
+
+	KeReleaseGuardedMutex( &header->lock );
+}
+
+VOID FreeApcAndStatusStructure(
+	_Inout_ PKAPC Apc,
+	_In_ LONG ContextId
+)
+{
+	PAPC_CONTEXT_HEADER header = NULL;
+
+	GetApcContext( &header, ContextId );
+
+	if ( !header )
+		return;
+
+	KeAcquireGuardedMutex( &header->lock );
+
+	switch ( header->context_id )
+	{
+	case APC_CONTEXT_ID_STACKWALK:;
+
+		PAPC_STACKWALK_CONTEXT context = ( PAPC_STACKWALK_CONTEXT )header;
+
+		PLIST_ITEM entry = ( PLIST_ITEM )context->head->start;
+
+		while ( entry != NULL )
+		{
+			PAPC_STATUS apc_entry = ( PAPC_STATUS )entry->data;
+
+			if ( apc_entry->apc == Apc )
+			{
+				ExFreePoolWithTag( apc_entry->apc, POOL_TAG_APC );
+				ExFreePoolWithTag( entry->data, POOL_TAG_APC );
+				RemoveApcFromApcContextList( context->head, entry );
+				KeReleaseGuardedMutex( &header->lock );
+				return;
+			}
+
+			entry = entry->next;
+		}
+
+		break;
+	}
+
+	KeReleaseGuardedMutex( &header->lock );
 }
 
 /*
-* The KernelRoutine is executed in kernel mode at APC_LEVEL before the APC is delivered.
+* The KernelRoutine is executed in kernel mode at APC_LEVEL before the APC is delivered. This
+* is also where we want to free our APC object.
 */
 VOID ApcKernelRoutine(
 	_In_ PRKAPC Apc,
-	_Inout_ _Deref_pre_maybenull_ PKNORMAL_ROUTINE* NormalRoutine,
+	_Inout_ _Deref_pre_maybenull_ PKNORMAL_ROUTINE* NormalRoutine, 
 	_Inout_ _Deref_pre_maybenull_ PVOID* NormalContext,
 	_Inout_ _Deref_pre_maybenull_ PVOID* SystemArgument1,
 	_Inout_ _Deref_pre_maybenull_ PVOID* SystemArgument2
 )
 {
-	DEBUG_LOG( "Hello from apc routine! ThreadId: %lx", (LONG)PsGetCurrentThreadId() );
+	PVOID buffer = NULL;
+	INT frames_captured = 0;
+	UINT64 stack_frame = 0;
+	NTSTATUS status;
+	BOOLEAN flag = FALSE;
+	PAPC_STACKWALK_CONTEXT context;
+
+	context = ( PAPC_STACKWALK_CONTEXT )Apc->NormalContext;
+
+	buffer = ExAllocatePool2( POOL_FLAG_NON_PAGED, 0x200, POOL_TAG_APC );
+
+	if ( !buffer )
+		return;
+
+	frames_captured = RtlCaptureStackBackTrace(
+		NULL,
+		STACK_FRAME_POOL_SIZE / sizeof(UINT64),
+		buffer,
+		NULL
+	);
+
+	if ( frames_captured == NULL )
+		goto free;
+
+	for ( INT index = 0; index < frames_captured; index++ )
+	{
+		stack_frame = *( UINT64* )( ( UINT64 )buffer + index * sizeof(UINT64) );
+
+		/*
+		* Apc->NormalContext holds the address of our context data structure that we passed into
+		* KeInitializeApc as the last argument.
+		*/
+		status = IsInstructionPointerInInvalidRegion(
+			stack_frame,
+			context->modules,
+			&flag
+		);
+
+		if ( !NT_SUCCESS( status ) )
+		{
+			DEBUG_ERROR( "IsInstructionPointerInInvalidRegion failed with status %x", status );
+			goto free;
+		}
+	}
+
+free:
+
+	ExFreePoolWithTag( buffer, POOL_TAG_APC );
+	FreeApcAndStatusStructure( Apc, APC_CONTEXT_ID_STACKWALK );
 }
 
 /*
@@ -908,8 +1060,6 @@ VOID ValidateThreadViaKernelApcCallback(
 	_In_ PVOID Context
 )
 {
-	UNREFERENCED_PARAMETER( Context );
-
 	NTSTATUS status;
 	PLIST_ENTRY thread_list_head;
 	PLIST_ENTRY thread_list_entry;
@@ -924,10 +1074,13 @@ VOID ValidateThreadViaKernelApcCallback(
 	{
 		current_thread = ( PETHREAD )( ( UINT64 )thread_list_entry - KTHREAD_THREADLIST_OFFSET );
 
-		/* sanity check */
+		/* ensure thread has a valid cid entry */
 		if ( PsGetThreadId( current_thread ) == NULL )
 			goto increment;
 
+		if (current_thread == KeGetCurrentThread())
+			goto increment;
+
 		apc = ( PKAPC )ExAllocatePool2( POOL_FLAG_NON_PAGED, sizeof( KAPC ), POOL_TAG_APC );
 
 		if ( !apc )
@@ -941,7 +1094,7 @@ VOID ValidateThreadViaKernelApcCallback(
 			ApcRundownRoutine,
 			ApcNormalRoutine,
 			KernelMode,
-			NULL
+			Context
 		);
 
 		apc_status = KeInsertQueueApc(
@@ -952,7 +1105,12 @@ VOID ValidateThreadViaKernelApcCallback(
 		);
 
 		if ( !apc_status )
+		{
 			DEBUG_ERROR( "KeInsertQueueApc failed" );
+			goto increment;
+		}
+
+		InsertApcToContextList( apc, APC_CONTEXT_ID_STACKWALK );
 
 	increment:
 		thread_list_entry = thread_list_entry->Flink;
@@ -967,8 +1125,43 @@ VOID ValidateThreadViaKernelApcCallback(
 */
 NTSTATUS ValidateThreadsViaKernelApc()
 {
+	NTSTATUS status;
+	PAPC_STACKWALK_CONTEXT context = NULL;
+
+	context = ExAllocatePool2( POOL_FLAG_NON_PAGED, sizeof( APC_STACKWALK_CONTEXT ), POOL_TAG_APC );
+
+	if ( !context )
+		return STATUS_MEMORY_NOT_ALLOCATED;
+
+	context->context_id = APC_CONTEXT_ID_STACKWALK;
+	context->modules = ExAllocatePool2( POOL_FLAG_NON_PAGED, sizeof( SYSTEM_MODULES ), POOL_TAG_APC );
+
+	KeInitializeGuardedMutex( &context->lock );
+
+	if ( !context->modules )
+	{
+		ExFreePoolWithTag( context, POOL_TAG_APC );
+		return STATUS_MEMORY_NOT_ALLOCATED;
+	}
+
+	status = GetSystemModuleInformation( context->modules );
+
+	if ( !NT_SUCCESS( status ) )
+	{
+		ExFreePoolWithTag( context->modules, POOL_TAG_APC );
+		ExFreePoolWithTag( context, POOL_TAG_APC );
+		return STATUS_MEMORY_NOT_ALLOCATED;
+	}
+
+	InsertApcContext( context );
+
 	EnumerateProcessListWithCallbackFunction(
 		ValidateThreadViaKernelApcCallback,
-		NULL
+		context
 	);
+}
+
+NTSTATUS QueryApcListToCheckForCompletion()
+{
+
 }

driver/modules.h

@@ -3,7 +3,9 @@
 
 #include <ntifs.h>
 #include <intrin.h>
+
 #include "common.h"
+#include "queue.h"
 
 #define MODULE_REPORT_DRIVER_NAME_BUFFER_SIZE 128
 
@@ -87,6 +89,32 @@ typedef struct _SYSTEM_MODULES
 
 }SYSTEM_MODULES, * PSYSTEM_MODULES;
 
+typedef struct _APC_STATUS
+{
+	PKAPC apc;
+	BOOLEAN apc_complete;
+
+}APC_STATUS, * PAPC_STATUS;
+
+typedef struct _APC_STACKWALK_CONTEXT
+{
+	LONG context_id;
+	PLIST_HEAD head;
+	KGUARDED_MUTEX lock;
+	PSYSTEM_MODULES modules;
+
+}APC_STACKWALK_CONTEXT, * PAPC_STACKWALK_CONTEXT;
+
+#define APC_CONTEXT_ID_STACKWALK 0x1
+
+typedef struct _APC_CONTEXT_HEADER
+{
+	LONG context_id;
+	PLIST_HEAD head;
+	KGUARDED_MUTEX lock;
+
+}APC_CONTEXT_HEADER, *PAPC_CONTEXT_HEADER;
+
 NTSTATUS GetSystemModuleInformation(
 	_Out_ PSYSTEM_MODULES ModuleInformation
 );
@@ -104,6 +132,10 @@ NTSTATUS HandleNmiIOCTL(
 	_In_ PIRP Irp
 );
 
+VOID FreeApcContextStructure(
+	_Inout_ PVOID Context
+);
+
 NTSTATUS ValidateThreadsViaKernelApc();
 
 #endif

driver/queue.c

@@ -253,3 +253,86 @@ end:
 	DEBUG_LOG( "Moved all reports into the IRP, sending !" );
 	return STATUS_SUCCESS;
 }
+
+VOID ListInit(
+	_In_ PLIST_HEAD ListHead
+)
+{
+	KeInitializeSpinLock( &ListHead->lock );
+	ListHead->start = NULL;
+}
+
+PLIST_ITEM ListInsert(
+	_In_ PLIST_HEAD ListHead,
+	_In_ PVOID Data
+)
+{
+	KIRQL irql = KeGetCurrentIrql();
+	KeAcquireSpinLock( &ListHead->lock, &irql );
+
+	PLIST_ITEM entry = ExAllocatePool2( POOL_FLAG_NON_PAGED, sizeof( LIST_ITEM ), POOL_TAG_APC );
+
+	if ( !entry )
+		return;
+
+	entry->data = Data;
+	entry->next = ListHead->start;
+	ListHead->start = entry;
+
+	KeReleaseSpinLock( &ListHead->lock, irql );
+
+	return entry;
+}
+
+PVOID ListRemoveFirst(
+	_In_ PLIST_HEAD ListHead
+)
+{
+	KIRQL irql = KeGetCurrentIrql();
+	KeAcquireSpinLock( &ListHead->lock, &irql );
+
+	if ( ListHead->start )
+	{
+		PLIST_ITEM entry = ListHead->start;
+		ListHead->start = ListHead->start->next;
+		ExFreePoolWithTag( entry, POOL_TAG_APC );
+	}
+
+	KeReleaseSpinLock( &ListHead->lock, irql );
+}
+
+PVOID ListRemoveItem(
+	_In_ PLIST_HEAD ListHead,
+	_Inout_ PLIST_ITEM ListItem
+)
+{
+	KIRQL irql = KeGetCurrentIrql();
+	KeAcquireSpinLock( &ListHead->lock, &irql );
+
+	PLIST_ITEM entry = ListHead->start;
+
+	if ( !entry )
+		goto unlock;
+
+	if ( entry == ListItem )
+	{
+		ListHead->start = entry->next;
+		ExFreePoolWithTag( ListItem, POOL_TAG_APC );
+		goto unlock;
+	}
+
+	while ( entry->next )
+	{
+		if ( entry->next == ListItem )
+		{
+			entry->next = entry->next->next;
+			ExFreePoolWithTag( ListItem, POOL_TAG_APC );
+			goto unlock;
+		}
+
+		entry = entry->next;
+	}
+
+unlock:
+	KeReleaseSpinLock( &ListHead->lock, irql);
+}

driver/queue.h

@@ -34,6 +34,22 @@ typedef struct _REPORT_HEADER
 
 }REPORT_HEADER, * PREPORT_HEADER;
 
+typedef struct _LIST_ITEM
+{
+	struct _LIST_ITEM* next;
+	PVOID data;
+
+}LIST_ITEM, * PLIST_ITEM;
+
+typedef struct _LIST_HEAD
+{
+	PLIST_ITEM start;
+	KSPIN_LOCK lock;
+
+}LIST_HEAD, * PLIST_HEAD;
+
+#define LIST_POOL_TAG 'list'
+
 VOID QueuePush(
 	_In_ PQUEUE_HEAD Head,
 	_In_ PVOID Data
@@ -57,4 +73,22 @@ NTSTATUS HandlePeriodicGlobalReportQueueQuery(
 
 VOID FreeGlobalReportQueueObjects();
 
+VOID ListInit(
+	_In_ PLIST_HEAD ListHead
+);
+
+PLIST_ITEM ListInsert(
+	_In_ PLIST_HEAD ListHead,
+	_In_ PVOID Data
+);
+
+PVOID ListRemoveFirst(
+	_In_ PLIST_HEAD ListHead
+);
+
+PVOID ListRemoveItem(
+	_In_ PLIST_HEAD ListHead,
+	_Inout_ PLIST_ITEM ListItem
+);
+
 #endif