From b4e8b7b576689918d3cec3983d35d1a86ba7304b Mon Sep 17 00:00:00 2001 From: lhodges1 Date: Mon, 9 Oct 2023 18:34:30 +1100 Subject: [PATCH] implement new thread list --- driver/callbacks.c | 29 +++--- driver/callbacks.h | 16 ++++ driver/driver.c | 66 +++++++++----- driver/driver.h | 3 +- driver/modules.c | 219 ++++++++++++++++++++++----------------------- driver/thread.c | 99 ++++++-------------- user/main.cpp | 66 +++++++------- 7 files changed, 245 insertions(+), 253 deletions(-) diff --git a/driver/callbacks.c b/driver/callbacks.c index 9fb219b..1925614 100644 --- a/driver/callbacks.c +++ b/driver/callbacks.c @@ -8,9 +8,9 @@ /* * Interlocked intrinsics are only atomic with respect to other InterlockedXxx functions, -* so all reads and writes to the THREAD_LIST->active flag must be with Interlocked instrinsics. +* so all reads and writes to the THREAD_LIST->active flag must be with Interlocked instrinsics +* to ensure atomicity. */ - typedef struct _THREAD_LIST { SINGLE_LIST_ENTRY start; @@ -19,13 +19,6 @@ typedef struct _THREAD_LIST }THREAD_LIST, * PTHREAD_LIST; -typedef struct _THREAD_LIST_ENTRY -{ - SINGLE_LIST_ENTRY list; - PKTHREAD thread; - -}THREAD_LIST_ENTRY, *PTHREAD_LIST_ENTRY; - /* todo: maybe put this in the global config? hmm.. I kinda like how its encapsulated here tho hm.. */ PTHREAD_LIST thread_list = NULL; @@ -63,7 +56,8 @@ CleanupThreadListOnDriverUnload() } /* -* Safely enumerate the threads list. +* Important to remember the callback function will run at irql = DISPATCH_LEVEL since +* we hold the spinlock during enumeration. */ VOID EnumerateThreadListWithCallbackRoutine( @@ -81,8 +75,8 @@ EnumerateThreadListWithCallbackRoutine( while (entry) { - VOID(*callback_function_ptr)(PKTHREAD, PVOID) = CallbackRoutine; - (*callback_function_ptr)(entry->thread, Context); + VOID(*callback_function_ptr)(PTHREAD_LIST_ENTRY, PVOID) = CallbackRoutine; + (*callback_function_ptr)(entry, Context); entry = entry->list.Next; } @@ -104,7 +98,6 @@ InitialiseThreadList() return STATUS_SUCCESS; } -STATIC VOID FindThreadListEntryByThreadAddress( _In_ PKTHREAD Thread, @@ -140,14 +133,16 @@ ThreadCreateNotifyRoutine( { PTHREAD_LIST_ENTRY entry = NULL; PKTHREAD thread = NULL; + PKPROCESS process = NULL; /* ensure we don't insert new entries if we are unloading */ if (InterlockedExchange(&thread_list->active, thread_list->active) == FALSE) return; PsLookupThreadByThreadId(ThreadId, &thread); + PsLookupProcessByProcessId(ProcessId, &process); - if (!thread) + if (!thread || !process) return; if (Create) @@ -158,8 +153,11 @@ ThreadCreateNotifyRoutine( return; entry->thread = thread; + entry->owning_process = process; + entry->apc = NULL; + entry->apc_queued = FALSE; + ListInsert(&thread_list->start, &entry->list, &thread_list->lock); - DEBUG_LOG("Thread inserted: %llx", (UINT64)thread); } else { @@ -169,7 +167,6 @@ ThreadCreateNotifyRoutine( return; ListRemoveEntry(&thread_list->start, entry, &thread_list->lock); - DEBUG_LOG("Thread removed: %llx", (UINT64)thread); } } diff --git a/driver/callbacks.h b/driver/callbacks.h index 646602f..eb5b6ae 100644 --- a/driver/callbacks.h +++ b/driver/callbacks.h @@ -44,6 +44,16 @@ static const uintptr_t EPROCESS_PLIST_ENTRY_OFFSET = 0x448; static UNICODE_STRING OBJECT_TYPE_PROCESS = RTL_CONSTANT_STRING(L"Process"); static UNICODE_STRING OBJECT_TYPE_THREAD = RTL_CONSTANT_STRING(L"Thread"); +typedef struct _THREAD_LIST_ENTRY +{ + SINGLE_LIST_ENTRY list; + PKTHREAD thread; + PKPROCESS owning_process; + BOOLEAN apc_queued; + PKAPC apc; + +}THREAD_LIST_ENTRY, * PTHREAD_LIST_ENTRY; + VOID NTAPI ExUnlockHandleTableEntry( @@ -93,4 +103,10 @@ ThreadCreateNotifyRoutine( VOID CleanupThreadListOnDriverUnload(); +VOID +FindThreadListEntryByThreadAddress( + _In_ PKTHREAD Thread, + _Inout_ PTHREAD_LIST_ENTRY* Entry +); + #endif diff --git a/driver/driver.c b/driver/driver.c index 1383d87..5031dd4 100644 --- a/driver/driver.c +++ b/driver/driver.c @@ -65,9 +65,6 @@ NTSTATUS RegistryPathQueryCallbackRoutine( #pragma alloc_text(PAGE, UnregisterCallbacksOnProcessTermination) #pragma alloc_text(PAGE, CleanupDriverCallbacksOnDriverUnload) #pragma alloc_text(PAGE, GetCallbackConfigStructure) -#pragma alloc_text(PAGE, InsertApcContext) -#pragma alloc_text(PAGE, GetApcContext) -#pragma alloc_text(PAGE, GetApcContextByIndex) #pragma alloc_text(PAGE, ReadProcessInitialisedConfigFlag) #pragma alloc_text(PAGE, GetProtectedProcessEProcess) #pragma alloc_text(PAGE, GetProtectedProcessId) @@ -101,7 +98,9 @@ typedef struct _DRIVER_CONFIG SYSTEM_INFORMATION system_information; PVOID apc_contexts[MAXIMUM_APC_CONTEXTS]; CALLBACK_CONFIGURATION callback_config; + volatile BOOLEAN unload_in_progress; KGUARDED_MUTEX lock; + KSPIN_LOCK spin_lock; }DRIVER_CONFIG, * PDRIVER_CONFIG; @@ -288,8 +287,9 @@ BOOLEAN FreeAllApcContextStructures() { BOOLEAN flag = TRUE; + KIRQL irql = KeGetCurrentIrql(); - KeAcquireGuardedMutex(&driver_config.lock); + KeAcquireSpinLock(&driver_config.spin_lock, &irql); for (INT index = 0; index < MAXIMUM_APC_CONTEXTS; index++) { @@ -310,7 +310,7 @@ FreeAllApcContextStructures() } unlock: - KeReleaseGuardedMutex(&driver_config.lock); + KeReleaseSpinLock(&driver_config.spin_lock, irql); return flag; } @@ -356,14 +356,15 @@ IncrementApcCount( ) { PAPC_CONTEXT_HEADER header = NULL; + KIRQL irql = KeGetCurrentIrql(); GetApcContext(&header, ContextId); if (!header) return; - KeAcquireGuardedMutex(&driver_config.lock); + KeAcquireSpinLock(&driver_config.spin_lock, &irql); header->count += 1; - KeReleaseGuardedMutex(&driver_config.lock); + KeReleaseSpinLock(&driver_config.spin_lock, irql); } VOID @@ -373,16 +374,18 @@ FreeApcAndDecrementApcCount( ) { PAPC_CONTEXT_HEADER context = NULL; + KIRQL irql = KeGetCurrentIrql(); + ExFreePoolWithTag(Apc, POOL_TAG_APC); GetApcContext(&context, ContextId); if (!context) goto end; - KeAcquireGuardedMutex(&driver_config.lock); + KeAcquireSpinLock(&driver_config.spin_lock, &irql); context->count -= 1; end: - KeReleaseGuardedMutex(&driver_config.lock); + KeReleaseSpinLock(&driver_config.spin_lock, irql); } /* @@ -411,17 +414,18 @@ end: NTSTATUS QueryActiveApcContextsForCompletion() { + KIRQL irql = KeGetCurrentIrql(); for (INT index = 0; index < MAXIMUM_APC_CONTEXTS; index++) { PAPC_CONTEXT_HEADER entry = NULL; GetApcContextByIndex(&entry, index); /* acquire mutex after we get the context to prevent thread deadlock */ - KeAcquireGuardedMutex(&driver_config.lock); + KeAcquireSpinLock(&driver_config.spin_lock, &irql); if (entry == NULL) { - KeReleaseGuardedMutex(&driver_config.lock); + KeReleaseSpinLock(&driver_config.spin_lock, irql); continue; } @@ -430,7 +434,7 @@ QueryActiveApcContextsForCompletion() if (entry->count > 0 || entry->allocation_in_progress == TRUE) { - KeReleaseGuardedMutex(&driver_config.lock); + KeReleaseSpinLock(&driver_config.spin_lock, irql); continue; } @@ -442,21 +446,33 @@ QueryActiveApcContextsForCompletion() break; } - KeReleaseGuardedMutex(&driver_config.lock); + KeReleaseSpinLock(&driver_config.spin_lock, irql); } return STATUS_SUCCESS; } -VOID +NTSTATUS InsertApcContext( _In_ PVOID Context ) { - KeAcquireGuardedMutex(&driver_config.lock); + NTSTATUS status = STATUS_SUCCESS; + KIRQL irql = KeGetCurrentIrql(); + KeAcquireSpinLock(&driver_config.spin_lock, &irql); PAPC_CONTEXT_HEADER header = Context; + /* + * prevents the race condition where the driver is unloaded whilst a new apc operation + * is attempted to start, ensuring that even if it holds + */ + if (InterlockedExchange(&driver_config.unload_in_progress, driver_config.unload_in_progress)) + { + status = STATUS_ABANDONED; + goto end; + } + for (INT index = 0; index < MAXIMUM_APC_CONTEXTS; index++) { PUINT64 entry = driver_config.apc_contexts; @@ -468,7 +484,8 @@ InsertApcContext( } } end: - KeReleaseGuardedMutex(&driver_config.lock); + KeReleaseSpinLock(&driver_config.spin_lock, irql); + return status; } VOID @@ -477,7 +494,8 @@ GetApcContext( _In_ LONG ContextIdentifier ) { - KeAcquireGuardedMutex(&driver_config.lock); + KIRQL irql = KeGetCurrentIrql(); + KeAcquireSpinLock(&driver_config.spin_lock, &irql); for (INT index = 0; index < MAXIMUM_APC_CONTEXTS; index++) { @@ -494,7 +512,7 @@ GetApcContext( } unlock: - KeReleaseGuardedMutex(&driver_config.lock); + KeReleaseSpinLock(&driver_config.spin_lock, irql); } VOID @@ -503,13 +521,15 @@ GetApcContextByIndex( _In_ INT Index ) { + KIRQL irql = KeGetCurrentIrql(); + if (!Context) return; *Context = NULL; - KeAcquireGuardedMutex(&driver_config.lock); + KeAcquireSpinLock(&driver_config.spin_lock, &irql); *Context = driver_config.apc_contexts[Index]; - KeReleaseGuardedMutex(&driver_config.lock); + KeReleaseSpinLock(&driver_config.spin_lock, irql); } VOID @@ -718,6 +738,8 @@ InitialiseDriverConfigOnDriverEntry( RTL_QUERY_REGISTRY_TABLE query_table[3] = { 0 }; KeInitializeGuardedMutex(&driver_config.lock); + KeInitializeSpinLock(&driver_config.spin_lock); + driver_config.unload_in_progress = FALSE; RtlInitUnicodeString(&driver_config.device_name, L"\\Device\\DonnaAC"); RtlInitUnicodeString(&driver_config.device_symbolic_link, L"\\??\\DonnaAC"); @@ -856,6 +878,8 @@ DriverUnload( _In_ PDRIVER_OBJECT DriverObject ) { + InterlockedExchange(&driver_config.unload_in_progress, TRUE); + DEBUG_LOG("Unloading driver..."); //PsSetCreateProcessNotifyRoutine( ProcessCreateNotifyRoutine, TRUE ); //QueryActiveApcContextsForCompletion(); @@ -878,7 +902,7 @@ VOID TerminateProtectedProcessOnViolation() { NTSTATUS status; - ULONG process_id; + ULONG process_id = 0; GetProtectedProcessId(&process_id); diff --git a/driver/driver.h b/driver/driver.h index 537957e..662463a 100644 --- a/driver/driver.h +++ b/driver/driver.h @@ -63,7 +63,8 @@ VOID GetApcContext( _In_ LONG ContextIdentifier ); -VOID InsertApcContext( +NTSTATUS +InsertApcContext( _In_ PVOID Context ); diff --git a/driver/modules.c b/driver/modules.c index 7cefcda..e5d10c5 100644 --- a/driver/modules.c +++ b/driver/modules.c @@ -191,10 +191,10 @@ ApcNormalRoutine( _In_opt_ PVOID SystemArgument1, _In_opt_ PVOID SystemArgument2); -STATIC -VOID +STATIC +VOID ValidateThreadViaKernelApcCallback( - _In_ PEPROCESS Process, + _In_ PTHREAD_LIST_ENTRY ThreadListEntry, _Inout_opt_ PVOID Context); #ifdef ALLOC_PRAGMA @@ -217,9 +217,7 @@ ValidateThreadViaKernelApcCallback( #pragma alloc_text(PAGE, ApcKernelRoutine) #pragma alloc_text(PAGE, ApcNormalRoutine) #pragma alloc_text(PAGE, FlipKThreadMiscFlagsFlag) -#pragma alloc_text(PAGE, ValidateThreadViaKernelApcCallback) #pragma alloc_text(PAGE, ValidateThreadsViaKernelApc) -#pragma alloc_text(PAGE, FreeApcStackwalkApcContextInformation) #endif /* @@ -1125,9 +1123,15 @@ ApcKernelRoutine( NTSTATUS status; BOOLEAN flag = FALSE; PAPC_STACKWALK_CONTEXT context; + PTHREAD_LIST_ENTRY thread_list_entry = NULL; context = (PAPC_STACKWALK_CONTEXT)Apc->NormalContext; + FindThreadListEntryByThreadAddress(KeGetCurrentThread(), &thread_list_entry); + + if (!thread_list_entry) + return; + buffer = ExAllocatePool2(POOL_FLAG_NON_PAGED, STACK_FRAME_POOL_SIZE, POOL_TAG_APC); if (!buffer) @@ -1191,6 +1195,9 @@ free: ExFreePoolWithTag(buffer, POOL_TAG_APC); FreeApcAndDecrementApcCount(Apc, APC_CONTEXT_ID_STACKWALK); + + thread_list_entry->apc = NULL; + thread_list_entry->apc_queued = FALSE; } /* @@ -1230,25 +1237,21 @@ FlipKThreadMiscFlagsFlag( STATIC VOID ValidateThreadViaKernelApcCallback( - _In_ PEPROCESS Process, + _In_ PTHREAD_LIST_ENTRY ThreadListEntry, _Inout_opt_ PVOID Context ) { - NTSTATUS status; - PLIST_ENTRY thread_list_head; - PLIST_ENTRY thread_list_entry; - PETHREAD current_thread; PKAPC apc = NULL; - BOOLEAN apc_status; + BOOLEAN apc_status = FALSE; PLONG misc_flags = NULL; PCHAR previous_mode = NULL; PUCHAR state = NULL; BOOLEAN apc_queueable = FALSE; PAPC_STACKWALK_CONTEXT context = (PAPC_STACKWALK_CONTEXT)Context; - LPCSTR process_name = PsGetProcessImageFileName(Process); + LPCSTR process_name = PsGetProcessImageFileName(ThreadListEntry->owning_process); /* we dont want to schedule an apc to threads owned by the kernel */ - if (Process == PsInitialSystemProcess || !Context) + if (ThreadListEntry->owning_process == PsInitialSystemProcess || !Context) return; /* We are not interested in these processess.. for now lol */ @@ -1264,106 +1267,89 @@ ValidateThreadViaKernelApcCallback( DEBUG_LOG("Process: %s", process_name); - thread_list_head = (PLIST_ENTRY)((UINT64)Process + KPROCESS_THREADLIST_OFFSET); - thread_list_entry = thread_list_head->Flink; + if (ThreadListEntry->thread == KeGetCurrentThread() || !ThreadListEntry->thread) + return; + /* + * Its possible to set the KThread->ApcQueueable flag to false ensuring that no APCs can be + * queued to the thread, as KeInsertQueueApc will check this flag before queueing an APC so + * lets make sure we flip this before before queueing ours. Since we filter out any system + * threads this should be fine... c: + */ + misc_flags = (PLONG)((UINT64)ThreadListEntry->thread + KTHREAD_MISC_FLAGS_OFFSET); + previous_mode = (PCHAR)((UINT64)ThreadListEntry->thread + KTHREAD_PREVIOUS_MODE_OFFSET); + state = (PUCHAR)((UINT64)ThreadListEntry->thread + KTHREAD_STATE_OFFSET); - context->header.allocation_in_progress = TRUE; + /* we dont care about user mode threads */ + //if (*previous_mode == UserMode) + // return; - while (thread_list_entry != thread_list_head) + /* todo: We should also flag all threads that have the flag set to false */ + if (*misc_flags >> KTHREAD_MISC_FLAGS_APC_QUEUEABLE == FALSE) + FlipKThreadMiscFlagsFlag(ThreadListEntry->thread, KTHREAD_MISC_FLAGS_APC_QUEUEABLE, TRUE); + + /* force thread into an alertable state */ + if (*misc_flags >> KTHREAD_MISC_FLAGS_ALERTABLE == FALSE) + FlipKThreadMiscFlagsFlag(ThreadListEntry->thread, KTHREAD_MISC_FLAGS_ALERTABLE, TRUE); + + apc = (PKAPC)ExAllocatePool2(POOL_FLAG_NON_PAGED, sizeof(KAPC), POOL_TAG_APC); + + if (!apc) + return; + + /* + * KTHREAD->State values: + * + * 0 is INITIALIZED; + * 1 is READY; + * 2 is RUNNING; + * 3 is STANDBY; + * 4 is TERMINATED; + * 5 is WAIT; + * 6 is TRANSITION. + * + * Since we are unsafely enumerating the threads linked list, it's best just + * to make sure we don't queue an APC to a terminated thread. We also check after + * we've allocated memory for the apc to ensure the window between queuing our APC + * and checking the thread state is as small as possible. + */ + + //if (*state == THREAD_STATE_TERMINATED || THREAD_STATE_INIT) + //{ + // ExFreePoolWithTag(apc, POOL_TAG_APC); + // return; + //} + + DEBUG_LOG("Apc: %llx", (UINT64)apc); + + KeInitializeApc( + apc, + ThreadListEntry->thread, + OriginalApcEnvironment, + ApcKernelRoutine, + ApcRundownRoutine, + ApcNormalRoutine, + KernelMode, + Context + ); + + apc_status = KeInsertQueueApc( + apc, + NULL, + NULL, + IO_NO_INCREMENT + ); + + if (!apc_status) { - current_thread = (PETHREAD)((UINT64)thread_list_entry - KTHREAD_THREADLIST_OFFSET); - - if (current_thread == KeGetCurrentThread() || !current_thread) - goto increment; - /* - * Its possible to set the KThread->ApcQueueable flag to false ensuring that no APCs can be - * queued to the thread, as KeInsertQueueApc will check this flag before queueing an APC so - * lets make sure we flip this before before queueing ours. Since we filter out any system - * threads this should be fine... c: - */ - misc_flags = (PLONG)((UINT64)current_thread + KTHREAD_MISC_FLAGS_OFFSET); - previous_mode = (PCHAR)((UINT64)current_thread + KTHREAD_PREVIOUS_MODE_OFFSET); - state = (PUCHAR)((UINT64)current_thread + KTHREAD_STATE_OFFSET); - - if (!MmIsAddressValid(current_thread)) - goto increment; - - /* we dont care about user mode threads */ - //if (*previous_mode == UserMode) - // goto increment; - - /* todo: We should also flag all threads that have the flag set to false */ - if (*misc_flags >> KTHREAD_MISC_FLAGS_APC_QUEUEABLE == FALSE) - FlipKThreadMiscFlagsFlag(current_thread, KTHREAD_MISC_FLAGS_APC_QUEUEABLE, TRUE); - - /* force thread into an alertable state */ - if (*misc_flags >> KTHREAD_MISC_FLAGS_ALERTABLE == FALSE) - FlipKThreadMiscFlagsFlag(current_thread, KTHREAD_MISC_FLAGS_ALERTABLE, TRUE); - - apc = (PKAPC)ExAllocatePool2(POOL_FLAG_NON_PAGED, sizeof(KAPC), POOL_TAG_APC); - - if (!apc) - goto increment; - - /* - * KTHREAD->State values: - * - * 0 is INITIALIZED; - * 1 is READY; - * 2 is RUNNING; - * 3 is STANDBY; - * 4 is TERMINATED; - * 5 is WAIT; - * 6 is TRANSITION. - * - * Since we are unsafely enumerating the threads linked list, it's best just - * to make sure we don't queue an APC to a terminated thread. We also check after - * we've allocated memory for the apc to ensure the window between queuing our APC - * and checking the thread state is as small as possible. - */ - - if (*state == THREAD_STATE_TERMINATED || THREAD_STATE_INIT) - { - ExFreePoolWithTag(apc, POOL_TAG_APC); - apc = NULL; - goto increment; - } - - DEBUG_LOG("Apc: %llx", (UINT64)apc); - - KeInitializeApc( - apc, - current_thread, - OriginalApcEnvironment, - ApcKernelRoutine, - ApcRundownRoutine, - ApcNormalRoutine, - KernelMode, - Context - ); - - apc_status = KeInsertQueueApc( - apc, - NULL, - NULL, - IO_NO_INCREMENT - ); - - if (!apc_status) - { - DEBUG_ERROR("KeInsertQueueApc failed"); - ExFreePoolWithTag(apc, POOL_TAG_APC); - apc = NULL; - goto increment; - } - - IncrementApcCount(APC_CONTEXT_ID_STACKWALK); - - increment: - thread_list_entry = thread_list_entry->Flink; + DEBUG_ERROR("KeInsertQueueApc failed"); + ExFreePoolWithTag(apc, POOL_TAG_APC); + return; } - context->header.allocation_in_progress = FALSE; + ThreadListEntry->apc = apc; + ThreadListEntry->apc_queued = TRUE; + + IncrementApcCount(APC_CONTEXT_ID_STACKWALK); } /* @@ -1410,12 +1396,25 @@ ValidateThreadsViaKernelApc() return STATUS_MEMORY_NOT_ALLOCATED; } - InsertApcContext(context); + status = InsertApcContext(context); - EnumerateProcessListWithCallbackFunction( + if (!NT_SUCCESS(status)) + { + ExFreePoolWithTag(context->modules, POOL_TAG_APC); + ExFreePoolWithTag(context, POOL_TAG_APC); + return status; + } + + context->header.allocation_in_progress = TRUE; + + EnumerateThreadListWithCallbackRoutine( ValidateThreadViaKernelApcCallback, context ); + + context->header.allocation_in_progress = FALSE; + + return status; } VOID diff --git a/driver/thread.c b/driver/thread.c index 275b5d2..478f0ee 100644 --- a/driver/thread.c +++ b/driver/thread.c @@ -7,6 +7,11 @@ #include "driver.h" #include "queue.h" +#ifdef ALLOC_PRAGMA +#pragma alloc_text(PAGE, ValidateKPCRBThreads) +#pragma alloc_text(PAGE, DetectThreadsAttachedToProtectedProcess) +#endif + typedef struct _KPRCB_THREAD_VALIDATION_CTX { UINT64 current_kpcrb_thread; @@ -16,63 +21,30 @@ typedef struct _KPRCB_THREAD_VALIDATION_CTX }KPRCB_THREAD_VALIDATION_CTX, * PKPRCB_THREAD_VALIDATION_CTX; -STATIC -VOID -KPRCBThreadValidationProcessCallback( - _In_ PEPROCESS Process, - _Inout_opt_ PVOID Context); - -STATIC -VOID -DetectAttachedThreadsProcessCallback( - _In_ PEPROCESS Process, - _Inout_opt_ PVOID Context); - -#ifdef ALLOC_PRAGMA -#pragma alloc_text(PAGE, KPRCBThreadValidationProcessCallback) -#pragma alloc_text(PAGE, ValidateKPCRBThreads) -#pragma alloc_text(PAGE, DetectAttachedThreadsProcessCallback) -#pragma alloc_text(PAGE, DetectThreadsAttachedToProtectedProcess) -#endif - STATIC VOID KPRCBThreadValidationProcessCallback( - _In_ PEPROCESS Process, + _In_ PTHREAD_LIST_ENTRY ThreadListEntry, _Inout_opt_ PVOID Context ) { - NTSTATUS status; - PLIST_ENTRY thread_list_head; - PLIST_ENTRY thread_list_entry; - PETHREAD current_thread; UINT32 thread_id; PKPRCB_THREAD_VALIDATION_CTX context = (PKPRCB_THREAD_VALIDATION_CTX)Context; if (!Context || context->finished == TRUE) return; - thread_list_head = (PLIST_ENTRY)((UINT64)Process + KPROCESS_THREADLIST_OFFSET); - thread_list_entry = thread_list_head->Flink; - - while (thread_list_entry != thread_list_head) + if (ThreadListEntry->thread == context->current_kpcrb_thread) { - current_thread = (PETHREAD)((UINT64)thread_list_entry - KTHREAD_THREADLIST_OFFSET); + context->thread_found_in_kthreadlist = TRUE; - if (current_thread == context->current_kpcrb_thread) + thread_id = PsGetThreadId(ThreadListEntry->thread); + + if (thread_id != NULL) { - context->thread_found_in_kthreadlist = TRUE; - - thread_id = PsGetThreadId(current_thread); - - if (thread_id != NULL) - { - context->thread_found_in_pspcidtable = TRUE; - context->finished = TRUE; - } + context->thread_found_in_pspcidtable = TRUE; + context->finished = TRUE; } - - thread_list_entry = thread_list_entry->Flink; } } @@ -97,7 +69,6 @@ KPRCBThreadValidationProcessCallback( * process' thread list , we know it's been removed from the KTHREAD linked list. * */ - VOID ValidateKPCRBThreads( _Inout_ PIRP Irp @@ -130,7 +101,7 @@ ValidateKPCRBThreads( if (!context.current_kpcrb_thread) continue; - EnumerateProcessListWithCallbackFunction( + EnumerateThreadListWithCallbackRoutine( KPRCBThreadValidationProcessCallback, &context ); @@ -167,17 +138,12 @@ ValidateKPCRBThreads( STATIC VOID DetectAttachedThreadsProcessCallback( - _In_ PEPROCESS Process, + _In_ PTHREAD_LIST_ENTRY ThreadListEntry, _Inout_opt_ PVOID Context ) { UNREFERENCED_PARAMETER(Context); - NTSTATUS status; - PLIST_ENTRY thread_list_head; - PLIST_ENTRY thread_list_entry; - PETHREAD current_thread; - UINT32 thread_id; PKAPC_STATE apc_state; PEPROCESS protected_process = NULL; @@ -186,34 +152,23 @@ DetectAttachedThreadsProcessCallback( if (protected_process == NULL) return; - thread_list_head = (PLIST_ENTRY)((UINT64)Process + KPROCESS_THREADLIST_OFFSET); - thread_list_entry = thread_list_head->Flink; + apc_state = (PKAPC_STATE)((UINT64)ThreadListEntry->thread + KTHREAD_APC_STATE_OFFSET); - while (thread_list_entry != thread_list_head) + if (apc_state->Process == protected_process) { - current_thread = (PETHREAD)((UINT64)thread_list_entry - KTHREAD_THREADLIST_OFFSET); + DEBUG_LOG("Program attached to notepad: %llx", (UINT64)ThreadListEntry->thread); - apc_state = (PKAPC_STATE)((UINT64)current_thread + KTHREAD_APC_STATE_OFFSET); + PATTACH_PROCESS_REPORT report = + ExAllocatePool2(POOL_FLAG_NON_PAGED, sizeof(ATTACH_PROCESS_REPORT), REPORT_POOL_TAG); - if (apc_state->Process == protected_process) - { - DEBUG_LOG("Program attached to notepad: %llx", (UINT64)current_thread); + if (!report) + return; - PATTACH_PROCESS_REPORT report = - ExAllocatePool2(POOL_FLAG_NON_PAGED, sizeof(ATTACH_PROCESS_REPORT), REPORT_POOL_TAG); + report->report_code = REPORT_ILLEGAL_ATTACH_PROCESS; + report->thread_id = PsGetThreadId(ThreadListEntry->thread); + report->thread_address = ThreadListEntry->thread; - if (!report) - goto increment; - - report->report_code = REPORT_ILLEGAL_ATTACH_PROCESS; - report->thread_id = PsGetThreadId(current_thread); - report->thread_address = current_thread; - - InsertReportToQueue(report); - } - - increment: - thread_list_entry = thread_list_entry->Flink; + InsertReportToQueue(report); } } @@ -232,7 +187,7 @@ DetectAttachedThreadsProcessCallback( */ VOID DetectThreadsAttachedToProtectedProcess() { - EnumerateProcessListWithCallbackFunction( + EnumerateThreadListWithCallbackRoutine( DetectAttachedThreadsProcessCallback, NULL ); diff --git a/user/main.cpp b/user/main.cpp index 6e44380..6ff4358 100644 --- a/user/main.cpp +++ b/user/main.cpp @@ -47,39 +47,39 @@ DWORD WINAPI Init(HINSTANCE hinstDLL) std::cout << "Seed: " << seed << std::endl; - switch (seed) - { - case 0: - kmanager.EnumerateHandleTables(); - break; - case 1: - kmanager.PerformIntegrityCheck(); - break; - case 2: - kmanager.ScanPoolsForUnlinkedProcesses(); - break; - case 3: - kmanager.VerifySystemModules(); - break; - case 4: - kmanager.ValidateProcessModules(); - break; - case 5: - kmanager.RunNmiCallbacks(); - break; - case 6: - kmanager.CheckForAttachedThreads(); - break; - case 7: - kmanager.InitiateApcStackwalkOperation(); - break; - case 8: - kmanager.CheckForHiddenThreads(); - break; - case 9: - kmanager.CheckForEptHooks(); - break; - } + //switch (seed) + //{ + //case 0: + // kmanager.EnumerateHandleTables(); + // break; + //case 1: + // kmanager.PerformIntegrityCheck(); + // break; + //case 2: + // kmanager.ScanPoolsForUnlinkedProcesses(); + // break; + //case 3: + // kmanager.VerifySystemModules(); + // break; + //case 4: + // kmanager.ValidateProcessModules(); + // break; + //case 5: + // kmanager.RunNmiCallbacks(); + // break; + //case 6: + // kmanager.CheckForAttachedThreads(); + // break; + //case 7: + // kmanager.InitiateApcStackwalkOperation(); + // break; + //case 8: + // kmanager.CheckForHiddenThreads(); + // break; + //case 9: + // kmanager.CheckForEptHooks(); + // break; + //} kmanager.InitiateApcStackwalkOperation(); kmanager.MonitorCallbackReports();