noury/mirror-ac
1
0
Fork 0
mirror of https://github.com/donnaskiez/ac.git synced 2024-11-21 22:24:08 +01:00
Code Issues Projects Releases Packages Wiki Activity

implement new thread list

Browse source
This commit is contained in:
lhodges1 2023-10-09 18:34:30 +11:00
parent 7e18820fcc
commit b4e8b7b576
7 changed files with 245 additions and 253 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

29
driver/callbacks.c
View file

@ -8,9 +8,9 @@
/* /*
* Interlocked intrinsics are only atomic with respect to other InterlockedXxx functions, * 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 typedef struct _THREAD_LIST
{ {
SINGLE_LIST_ENTRY start; SINGLE_LIST_ENTRY start;
@ -19,13 +19,6 @@ typedef struct _THREAD_LIST
}THREAD_LIST, * PTHREAD_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.. */ /* todo: maybe put this in the global config? hmm.. I kinda like how its encapsulated here tho hm.. */
PTHREAD_LIST thread_list = NULL; 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 VOID
EnumerateThreadListWithCallbackRoutine( EnumerateThreadListWithCallbackRoutine(
@ -81,8 +75,8 @@ EnumerateThreadListWithCallbackRoutine(
while (entry) while (entry)
{ {
VOID(*callback_function_ptr)(PKTHREAD, PVOID) = CallbackRoutine; VOID(*callback_function_ptr)(PTHREAD_LIST_ENTRY, PVOID) = CallbackRoutine;
(*callback_function_ptr)(entry->thread, Context); (*callback_function_ptr)(entry, Context);
entry = entry->list.Next; entry = entry->list.Next;
} }
@ -104,7 +98,6 @@ InitialiseThreadList()
return STATUS_SUCCESS; return STATUS_SUCCESS;
} }
STATIC
VOID VOID
FindThreadListEntryByThreadAddress( FindThreadListEntryByThreadAddress(
_In_ PKTHREAD Thread, _In_ PKTHREAD Thread,
@ -140,14 +133,16 @@ ThreadCreateNotifyRoutine(
{ {
PTHREAD_LIST_ENTRY entry = NULL; PTHREAD_LIST_ENTRY entry = NULL;
PKTHREAD thread = NULL; PKTHREAD thread = NULL;
PKPROCESS process = NULL;
/* ensure we don't insert new entries if we are unloading */ /* ensure we don't insert new entries if we are unloading */
if (InterlockedExchange(&thread_list->active, thread_list->active) == FALSE) if (InterlockedExchange(&thread_list->active, thread_list->active) == FALSE)
return; return;
PsLookupThreadByThreadId(ThreadId, &thread); PsLookupThreadByThreadId(ThreadId, &thread);
PsLookupProcessByProcessId(ProcessId, &process);
if (!thread) if (!thread || !process)
return; return;
if (Create) if (Create)
@ -158,8 +153,11 @@ ThreadCreateNotifyRoutine(
return; return;
entry->thread = thread; entry->thread = thread;
entry->owning_process = process;
entry->apc = NULL;
entry->apc_queued = FALSE;
ListInsert(&thread_list->start, &entry->list, &thread_list->lock); ListInsert(&thread_list->start, &entry->list, &thread_list->lock);
DEBUG_LOG("Thread inserted: %llx", (UINT64)thread);
} }
else else
{ {
@ -169,7 +167,6 @@ ThreadCreateNotifyRoutine(
return; return;
ListRemoveEntry(&thread_list->start, entry, &thread_list->lock); ListRemoveEntry(&thread_list->start, entry, &thread_list->lock);
DEBUG_LOG("Thread removed: %llx", (UINT64)thread);
} }
} }

16
driver/callbacks.h
View file

@ -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_PROCESS = RTL_CONSTANT_STRING(L"Process");
static UNICODE_STRING OBJECT_TYPE_THREAD = RTL_CONSTANT_STRING(L"Thread"); 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 VOID
NTAPI NTAPI
ExUnlockHandleTableEntry( ExUnlockHandleTableEntry(
@ -93,4 +103,10 @@ ThreadCreateNotifyRoutine(
VOID VOID
CleanupThreadListOnDriverUnload(); CleanupThreadListOnDriverUnload();
VOID
FindThreadListEntryByThreadAddress(
_In_ PKTHREAD Thread,
_Inout_ PTHREAD_LIST_ENTRY* Entry
);
#endif #endif

66
driver/driver.c
View file

@ -65,9 +65,6 @@ NTSTATUS RegistryPathQueryCallbackRoutine(
#pragma alloc_text(PAGE, UnregisterCallbacksOnProcessTermination) #pragma alloc_text(PAGE, UnregisterCallbacksOnProcessTermination)
#pragma alloc_text(PAGE, CleanupDriverCallbacksOnDriverUnload) #pragma alloc_text(PAGE, CleanupDriverCallbacksOnDriverUnload)
#pragma alloc_text(PAGE, GetCallbackConfigStructure) #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, ReadProcessInitialisedConfigFlag)
#pragma alloc_text(PAGE, GetProtectedProcessEProcess) #pragma alloc_text(PAGE, GetProtectedProcessEProcess)
#pragma alloc_text(PAGE, GetProtectedProcessId) #pragma alloc_text(PAGE, GetProtectedProcessId)
@ -101,7 +98,9 @@ typedef struct _DRIVER_CONFIG
SYSTEM_INFORMATION system_information; SYSTEM_INFORMATION system_information;
PVOID apc_contexts[MAXIMUM_APC_CONTEXTS]; PVOID apc_contexts[MAXIMUM_APC_CONTEXTS];
CALLBACK_CONFIGURATION callback_config; CALLBACK_CONFIGURATION callback_config;
volatile BOOLEAN unload_in_progress;
KGUARDED_MUTEX lock; KGUARDED_MUTEX lock;
KSPIN_LOCK spin_lock;
}DRIVER_CONFIG, * PDRIVER_CONFIG; }DRIVER_CONFIG, * PDRIVER_CONFIG;
@ -288,8 +287,9 @@ BOOLEAN
FreeAllApcContextStructures() FreeAllApcContextStructures()
{ {
BOOLEAN flag = TRUE; 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++) for (INT index = 0; index < MAXIMUM_APC_CONTEXTS; index++)
{ {
@ -310,7 +310,7 @@ FreeAllApcContextStructures()
} }
unlock: unlock:
KeReleaseGuardedMutex(&driver_config.lock); KeReleaseSpinLock(&driver_config.spin_lock, irql);
return flag; return flag;
} }
@ -356,14 +356,15 @@ IncrementApcCount(
) )
{ {
PAPC_CONTEXT_HEADER header = NULL; PAPC_CONTEXT_HEADER header = NULL;
KIRQL irql = KeGetCurrentIrql();
GetApcContext(&header, ContextId); GetApcContext(&header, ContextId);
if (!header) if (!header)
return; return;
KeAcquireGuardedMutex(&driver_config.lock); KeAcquireSpinLock(&driver_config.spin_lock, &irql);
header->count += 1; header->count += 1;
KeReleaseGuardedMutex(&driver_config.lock); KeReleaseSpinLock(&driver_config.spin_lock, irql);
} }
VOID VOID
@ -373,16 +374,18 @@ FreeApcAndDecrementApcCount(
) )
{ {
PAPC_CONTEXT_HEADER context = NULL; PAPC_CONTEXT_HEADER context = NULL;
KIRQL irql = KeGetCurrentIrql();
ExFreePoolWithTag(Apc, POOL_TAG_APC); ExFreePoolWithTag(Apc, POOL_TAG_APC);
GetApcContext(&context, ContextId); GetApcContext(&context, ContextId);
if (!context) if (!context)
goto end; goto end;
KeAcquireGuardedMutex(&driver_config.lock); KeAcquireSpinLock(&driver_config.spin_lock, &irql);
context->count -= 1; context->count -= 1;
end: end:
KeReleaseGuardedMutex(&driver_config.lock); KeReleaseSpinLock(&driver_config.spin_lock, irql);
} }
/* /*
@ -411,17 +414,18 @@ end:
NTSTATUS NTSTATUS
QueryActiveApcContextsForCompletion() QueryActiveApcContextsForCompletion()
{ {
KIRQL irql = KeGetCurrentIrql();
for (INT index = 0; index < MAXIMUM_APC_CONTEXTS; index++) for (INT index = 0; index < MAXIMUM_APC_CONTEXTS; index++)
{ {
PAPC_CONTEXT_HEADER entry = NULL; PAPC_CONTEXT_HEADER entry = NULL;
GetApcContextByIndex(&entry, index); GetApcContextByIndex(&entry, index);
/* acquire mutex after we get the context to prevent thread deadlock */ /* acquire mutex after we get the context to prevent thread deadlock */
KeAcquireGuardedMutex(&driver_config.lock); KeAcquireSpinLock(&driver_config.spin_lock, &irql);
if (entry == NULL) if (entry == NULL)
{ {
KeReleaseGuardedMutex(&driver_config.lock); KeReleaseSpinLock(&driver_config.spin_lock, irql);
continue; continue;
} }
@ -430,7 +434,7 @@ QueryActiveApcContextsForCompletion()
if (entry->count > 0 || entry->allocation_in_progress == TRUE) if (entry->count > 0 || entry->allocation_in_progress == TRUE)
{ {
KeReleaseGuardedMutex(&driver_config.lock); KeReleaseSpinLock(&driver_config.spin_lock, irql);
continue; continue;
} }
@ -442,21 +446,33 @@ QueryActiveApcContextsForCompletion()
break; break;
} }
KeReleaseGuardedMutex(&driver_config.lock); KeReleaseSpinLock(&driver_config.spin_lock, irql);
} }
return STATUS_SUCCESS; return STATUS_SUCCESS;
} }
VOID NTSTATUS
InsertApcContext( InsertApcContext(
_In_ PVOID Context _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; 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++) for (INT index = 0; index < MAXIMUM_APC_CONTEXTS; index++)
{ {
PUINT64 entry = driver_config.apc_contexts; PUINT64 entry = driver_config.apc_contexts;
@ -468,7 +484,8 @@ InsertApcContext(
} }
} }
end: end:
KeReleaseGuardedMutex(&driver_config.lock); KeReleaseSpinLock(&driver_config.spin_lock, irql);
return status;
} }
VOID VOID
@ -477,7 +494,8 @@ GetApcContext(
_In_ LONG ContextIdentifier _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++) for (INT index = 0; index < MAXIMUM_APC_CONTEXTS; index++)
{ {
@ -494,7 +512,7 @@ GetApcContext(
} }
unlock: unlock:
KeReleaseGuardedMutex(&driver_config.lock); KeReleaseSpinLock(&driver_config.spin_lock, irql);
} }
VOID VOID
@ -503,13 +521,15 @@ GetApcContextByIndex(
_In_ INT Index _In_ INT Index
) )
{ {
KIRQL irql = KeGetCurrentIrql();
if (!Context) if (!Context)
return; return;
*Context = NULL; *Context = NULL;
KeAcquireGuardedMutex(&driver_config.lock); KeAcquireSpinLock(&driver_config.spin_lock, &irql);
*Context = driver_config.apc_contexts[Index]; *Context = driver_config.apc_contexts[Index];
KeReleaseGuardedMutex(&driver_config.lock); KeReleaseSpinLock(&driver_config.spin_lock, irql);
} }
VOID VOID
@ -718,6 +738,8 @@ InitialiseDriverConfigOnDriverEntry(
RTL_QUERY_REGISTRY_TABLE query_table[3] = { 0 }; RTL_QUERY_REGISTRY_TABLE query_table[3] = { 0 };
KeInitializeGuardedMutex(&driver_config.lock); 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_name, L"\\Device\\DonnaAC");
RtlInitUnicodeString(&driver_config.device_symbolic_link, L"\\??\\DonnaAC"); RtlInitUnicodeString(&driver_config.device_symbolic_link, L"\\??\\DonnaAC");
@ -856,6 +878,8 @@ DriverUnload(
_In_ PDRIVER_OBJECT DriverObject _In_ PDRIVER_OBJECT DriverObject
) )
{ {
InterlockedExchange(&driver_config.unload_in_progress, TRUE);
DEBUG_LOG("Unloading driver..."); DEBUG_LOG("Unloading driver...");
//PsSetCreateProcessNotifyRoutine( ProcessCreateNotifyRoutine, TRUE ); //PsSetCreateProcessNotifyRoutine( ProcessCreateNotifyRoutine, TRUE );
//QueryActiveApcContextsForCompletion(); //QueryActiveApcContextsForCompletion();
@ -878,7 +902,7 @@ VOID
TerminateProtectedProcessOnViolation() TerminateProtectedProcessOnViolation()
{ {
NTSTATUS status; NTSTATUS status;
ULONG process_id; ULONG process_id = 0;
GetProtectedProcessId(&process_id); GetProtectedProcessId(&process_id);

3
driver/driver.h
View file

@ -63,7 +63,8 @@ VOID GetApcContext(
_In_ LONG ContextIdentifier _In_ LONG ContextIdentifier
); );
VOID InsertApcContext( NTSTATUS
InsertApcContext(
_In_ PVOID Context _In_ PVOID Context
); );

97
driver/modules.c
View file

@ -194,7 +194,7 @@ ApcNormalRoutine(
STATIC STATIC
VOID VOID
ValidateThreadViaKernelApcCallback( ValidateThreadViaKernelApcCallback(
_In_ PEPROCESS Process, _In_ PTHREAD_LIST_ENTRY ThreadListEntry,
_Inout_opt_ PVOID Context); _Inout_opt_ PVOID Context);
#ifdef ALLOC_PRAGMA #ifdef ALLOC_PRAGMA
@ -217,9 +217,7 @@ ValidateThreadViaKernelApcCallback(
#pragma alloc_text(PAGE, ApcKernelRoutine) #pragma alloc_text(PAGE, ApcKernelRoutine)
#pragma alloc_text(PAGE, ApcNormalRoutine) #pragma alloc_text(PAGE, ApcNormalRoutine)
#pragma alloc_text(PAGE, FlipKThreadMiscFlagsFlag) #pragma alloc_text(PAGE, FlipKThreadMiscFlagsFlag)
#pragma alloc_text(PAGE, ValidateThreadViaKernelApcCallback)
#pragma alloc_text(PAGE, ValidateThreadsViaKernelApc) #pragma alloc_text(PAGE, ValidateThreadsViaKernelApc)
#pragma alloc_text(PAGE, FreeApcStackwalkApcContextInformation)
#endif #endif
/* /*
@ -1125,9 +1123,15 @@ ApcKernelRoutine(
NTSTATUS status; NTSTATUS status;
BOOLEAN flag = FALSE; BOOLEAN flag = FALSE;
PAPC_STACKWALK_CONTEXT context; PAPC_STACKWALK_CONTEXT context;
PTHREAD_LIST_ENTRY thread_list_entry = NULL;
context = (PAPC_STACKWALK_CONTEXT)Apc->NormalContext; 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); buffer = ExAllocatePool2(POOL_FLAG_NON_PAGED, STACK_FRAME_POOL_SIZE, POOL_TAG_APC);
if (!buffer) if (!buffer)
@ -1191,6 +1195,9 @@ free:
ExFreePoolWithTag(buffer, POOL_TAG_APC); ExFreePoolWithTag(buffer, POOL_TAG_APC);
FreeApcAndDecrementApcCount(Apc, APC_CONTEXT_ID_STACKWALK); FreeApcAndDecrementApcCount(Apc, APC_CONTEXT_ID_STACKWALK);
thread_list_entry->apc = NULL;
thread_list_entry->apc_queued = FALSE;
} }
/* /*
@ -1230,25 +1237,21 @@ FlipKThreadMiscFlagsFlag(
STATIC STATIC
VOID VOID
ValidateThreadViaKernelApcCallback( ValidateThreadViaKernelApcCallback(
_In_ PEPROCESS Process, _In_ PTHREAD_LIST_ENTRY ThreadListEntry,
_Inout_opt_ PVOID Context _Inout_opt_ PVOID Context
) )
{ {
NTSTATUS status;
PLIST_ENTRY thread_list_head;
PLIST_ENTRY thread_list_entry;
PETHREAD current_thread;
PKAPC apc = NULL; PKAPC apc = NULL;
BOOLEAN apc_status; BOOLEAN apc_status = FALSE;
PLONG misc_flags = NULL; PLONG misc_flags = NULL;
PCHAR previous_mode = NULL; PCHAR previous_mode = NULL;
PUCHAR state = NULL; PUCHAR state = NULL;
BOOLEAN apc_queueable = FALSE; BOOLEAN apc_queueable = FALSE;
PAPC_STACKWALK_CONTEXT context = (PAPC_STACKWALK_CONTEXT)Context; 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 */ /* we dont want to schedule an apc to threads owned by the kernel */
if (Process == PsInitialSystemProcess || !Context) if (ThreadListEntry->owning_process == PsInitialSystemProcess || !Context)
return; return;
/* We are not interested in these processess.. for now lol */ /* We are not interested in these processess.. for now lol */
@ -1264,46 +1267,34 @@ ValidateThreadViaKernelApcCallback(
DEBUG_LOG("Process: %s", process_name); DEBUG_LOG("Process: %s", process_name);
thread_list_head = (PLIST_ENTRY)((UINT64)Process + KPROCESS_THREADLIST_OFFSET); if (ThreadListEntry->thread == KeGetCurrentThread() || !ThreadListEntry->thread)
thread_list_entry = thread_list_head->Flink; return;
context->header.allocation_in_progress = TRUE;
while (thread_list_entry != thread_list_head)
{
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 * 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 * 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 * lets make sure we flip this before before queueing ours. Since we filter out any system
* threads this should be fine... c: * threads this should be fine... c:
*/ */
misc_flags = (PLONG)((UINT64)current_thread + KTHREAD_MISC_FLAGS_OFFSET); misc_flags = (PLONG)((UINT64)ThreadListEntry->thread + KTHREAD_MISC_FLAGS_OFFSET);
previous_mode = (PCHAR)((UINT64)current_thread + KTHREAD_PREVIOUS_MODE_OFFSET); previous_mode = (PCHAR)((UINT64)ThreadListEntry->thread + KTHREAD_PREVIOUS_MODE_OFFSET);
state = (PUCHAR)((UINT64)current_thread + KTHREAD_STATE_OFFSET); state = (PUCHAR)((UINT64)ThreadListEntry->thread + KTHREAD_STATE_OFFSET);
if (!MmIsAddressValid(current_thread))
goto increment;
/* we dont care about user mode threads */ /* we dont care about user mode threads */
//if (*previous_mode == UserMode) //if (*previous_mode == UserMode)
// goto increment; // return;
/* todo: We should also flag all threads that have the flag set to false */ /* todo: We should also flag all threads that have the flag set to false */
if (*misc_flags >> KTHREAD_MISC_FLAGS_APC_QUEUEABLE == FALSE) if (*misc_flags >> KTHREAD_MISC_FLAGS_APC_QUEUEABLE == FALSE)
FlipKThreadMiscFlagsFlag(current_thread, KTHREAD_MISC_FLAGS_APC_QUEUEABLE, TRUE); FlipKThreadMiscFlagsFlag(ThreadListEntry->thread, KTHREAD_MISC_FLAGS_APC_QUEUEABLE, TRUE);
/* force thread into an alertable state */ /* force thread into an alertable state */
if (*misc_flags >> KTHREAD_MISC_FLAGS_ALERTABLE == FALSE) if (*misc_flags >> KTHREAD_MISC_FLAGS_ALERTABLE == FALSE)
FlipKThreadMiscFlagsFlag(current_thread, KTHREAD_MISC_FLAGS_ALERTABLE, TRUE); FlipKThreadMiscFlagsFlag(ThreadListEntry->thread, KTHREAD_MISC_FLAGS_ALERTABLE, TRUE);
apc = (PKAPC)ExAllocatePool2(POOL_FLAG_NON_PAGED, sizeof(KAPC), POOL_TAG_APC); apc = (PKAPC)ExAllocatePool2(POOL_FLAG_NON_PAGED, sizeof(KAPC), POOL_TAG_APC);
if (!apc) if (!apc)
goto increment; return;
/* /*
* KTHREAD->State values: * KTHREAD->State values:
@ -1322,18 +1313,17 @@ ValidateThreadViaKernelApcCallback(
* and checking the thread state is as small as possible. * and checking the thread state is as small as possible.
*/ */
if (*state == THREAD_STATE_TERMINATED || THREAD_STATE_INIT) //if (*state == THREAD_STATE_TERMINATED || THREAD_STATE_INIT)
{ //{
ExFreePoolWithTag(apc, POOL_TAG_APC); // ExFreePoolWithTag(apc, POOL_TAG_APC);
apc = NULL; // return;
goto increment; //}
}
DEBUG_LOG("Apc: %llx", (UINT64)apc); DEBUG_LOG("Apc: %llx", (UINT64)apc);
KeInitializeApc( KeInitializeApc(
apc, apc,
current_thread, ThreadListEntry->thread,
OriginalApcEnvironment, OriginalApcEnvironment,
ApcKernelRoutine, ApcKernelRoutine,
ApcRundownRoutine, ApcRundownRoutine,
@ -1353,17 +1343,13 @@ ValidateThreadViaKernelApcCallback(
{ {
DEBUG_ERROR("KeInsertQueueApc failed"); DEBUG_ERROR("KeInsertQueueApc failed");
ExFreePoolWithTag(apc, POOL_TAG_APC); ExFreePoolWithTag(apc, POOL_TAG_APC);
apc = NULL; return;
goto increment;
} }
ThreadListEntry->apc = apc;
ThreadListEntry->apc_queued = TRUE;
IncrementApcCount(APC_CONTEXT_ID_STACKWALK); IncrementApcCount(APC_CONTEXT_ID_STACKWALK);
increment:
thread_list_entry = thread_list_entry->Flink;
}
context->header.allocation_in_progress = FALSE;
} }
/* /*
@ -1410,12 +1396,25 @@ ValidateThreadsViaKernelApc()
return STATUS_MEMORY_NOT_ALLOCATED; 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, ValidateThreadViaKernelApcCallback,
context context
); );
context->header.allocation_in_progress = FALSE;
return status;
} }
VOID VOID

77
driver/thread.c
View file

@ -7,6 +7,11 @@
#include "driver.h" #include "driver.h"
#include "queue.h" #include "queue.h"
#ifdef ALLOC_PRAGMA
#pragma alloc_text(PAGE, ValidateKPCRBThreads)
#pragma alloc_text(PAGE, DetectThreadsAttachedToProtectedProcess)
#endif
typedef struct _KPRCB_THREAD_VALIDATION_CTX typedef struct _KPRCB_THREAD_VALIDATION_CTX
{ {
UINT64 current_kpcrb_thread; UINT64 current_kpcrb_thread;
@ -19,51 +24,21 @@ typedef struct _KPRCB_THREAD_VALIDATION_CTX
STATIC STATIC
VOID VOID
KPRCBThreadValidationProcessCallback( KPRCBThreadValidationProcessCallback(
_In_ PEPROCESS Process, _In_ PTHREAD_LIST_ENTRY ThreadListEntry,
_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,
_Inout_opt_ PVOID Context _Inout_opt_ PVOID Context
) )
{ {
NTSTATUS status;
PLIST_ENTRY thread_list_head;
PLIST_ENTRY thread_list_entry;
PETHREAD current_thread;
UINT32 thread_id; UINT32 thread_id;
PKPRCB_THREAD_VALIDATION_CTX context = (PKPRCB_THREAD_VALIDATION_CTX)Context; PKPRCB_THREAD_VALIDATION_CTX context = (PKPRCB_THREAD_VALIDATION_CTX)Context;
if (!Context || context->finished == TRUE) if (!Context || context->finished == TRUE)
return; return;
thread_list_head = (PLIST_ENTRY)((UINT64)Process + KPROCESS_THREADLIST_OFFSET); if (ThreadListEntry->thread == context->current_kpcrb_thread)
thread_list_entry = thread_list_head->Flink;
while (thread_list_entry != thread_list_head)
{
current_thread = (PETHREAD)((UINT64)thread_list_entry - KTHREAD_THREADLIST_OFFSET);
if (current_thread == context->current_kpcrb_thread)
{ {
context->thread_found_in_kthreadlist = TRUE; context->thread_found_in_kthreadlist = TRUE;
thread_id = PsGetThreadId(current_thread); thread_id = PsGetThreadId(ThreadListEntry->thread);
if (thread_id != NULL) if (thread_id != NULL)
{ {
@ -71,9 +46,6 @@ KPRCBThreadValidationProcessCallback(
context->finished = 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. * process' thread list , we know it's been removed from the KTHREAD linked list.
* *
*/ */
VOID VOID
ValidateKPCRBThreads( ValidateKPCRBThreads(
_Inout_ PIRP Irp _Inout_ PIRP Irp
@ -130,7 +101,7 @@ ValidateKPCRBThreads(
if (!context.current_kpcrb_thread) if (!context.current_kpcrb_thread)
continue; continue;
EnumerateProcessListWithCallbackFunction( EnumerateThreadListWithCallbackRoutine(
KPRCBThreadValidationProcessCallback, KPRCBThreadValidationProcessCallback,
&context &context
); );
@ -167,17 +138,12 @@ ValidateKPCRBThreads(
STATIC STATIC
VOID VOID
DetectAttachedThreadsProcessCallback( DetectAttachedThreadsProcessCallback(
_In_ PEPROCESS Process, _In_ PTHREAD_LIST_ENTRY ThreadListEntry,
_Inout_opt_ PVOID Context _Inout_opt_ PVOID Context
) )
{ {
UNREFERENCED_PARAMETER(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; PKAPC_STATE apc_state;
PEPROCESS protected_process = NULL; PEPROCESS protected_process = NULL;
@ -186,35 +152,24 @@ DetectAttachedThreadsProcessCallback(
if (protected_process == NULL) if (protected_process == NULL)
return; return;
thread_list_head = (PLIST_ENTRY)((UINT64)Process + KPROCESS_THREADLIST_OFFSET); apc_state = (PKAPC_STATE)((UINT64)ThreadListEntry->thread + KTHREAD_APC_STATE_OFFSET);
thread_list_entry = thread_list_head->Flink;
while (thread_list_entry != thread_list_head)
{
current_thread = (PETHREAD)((UINT64)thread_list_entry - KTHREAD_THREADLIST_OFFSET);
apc_state = (PKAPC_STATE)((UINT64)current_thread + KTHREAD_APC_STATE_OFFSET);
if (apc_state->Process == protected_process) if (apc_state->Process == protected_process)
{ {
DEBUG_LOG("Program attached to notepad: %llx", (UINT64)current_thread); DEBUG_LOG("Program attached to notepad: %llx", (UINT64)ThreadListEntry->thread);
PATTACH_PROCESS_REPORT report = PATTACH_PROCESS_REPORT report =
ExAllocatePool2(POOL_FLAG_NON_PAGED, sizeof(ATTACH_PROCESS_REPORT), REPORT_POOL_TAG); ExAllocatePool2(POOL_FLAG_NON_PAGED, sizeof(ATTACH_PROCESS_REPORT), REPORT_POOL_TAG);
if (!report) if (!report)
goto increment; return;
report->report_code = REPORT_ILLEGAL_ATTACH_PROCESS; report->report_code = REPORT_ILLEGAL_ATTACH_PROCESS;
report->thread_id = PsGetThreadId(current_thread); report->thread_id = PsGetThreadId(ThreadListEntry->thread);
report->thread_address = current_thread; report->thread_address = ThreadListEntry->thread;
InsertReportToQueue(report); InsertReportToQueue(report);
} }
increment:
thread_list_entry = thread_list_entry->Flink;
}
} }
/* /*
@ -232,7 +187,7 @@ DetectAttachedThreadsProcessCallback(
*/ */
VOID DetectThreadsAttachedToProtectedProcess() VOID DetectThreadsAttachedToProtectedProcess()
{ {
EnumerateProcessListWithCallbackFunction( EnumerateThreadListWithCallbackRoutine(
DetectAttachedThreadsProcessCallback, DetectAttachedThreadsProcessCallback,
NULL NULL
); );

66
user/main.cpp
View file

@ -47,39 +47,39 @@ DWORD WINAPI Init(HINSTANCE hinstDLL)
std::cout << "Seed: " << seed << std::endl; std::cout << "Seed: " << seed << std::endl;
switch (seed) //switch (seed)
{ //{
case 0: //case 0:
kmanager.EnumerateHandleTables(); // kmanager.EnumerateHandleTables();
break; // break;
case 1: //case 1:
kmanager.PerformIntegrityCheck(); // kmanager.PerformIntegrityCheck();
break; // break;
case 2: //case 2:
kmanager.ScanPoolsForUnlinkedProcesses(); // kmanager.ScanPoolsForUnlinkedProcesses();
break; // break;
case 3: //case 3:
kmanager.VerifySystemModules(); // kmanager.VerifySystemModules();
break; // break;
case 4: //case 4:
kmanager.ValidateProcessModules(); // kmanager.ValidateProcessModules();
break; // break;
case 5: //case 5:
kmanager.RunNmiCallbacks(); // kmanager.RunNmiCallbacks();
break; // break;
case 6: //case 6:
kmanager.CheckForAttachedThreads(); // kmanager.CheckForAttachedThreads();
break; // break;
case 7: //case 7:
kmanager.InitiateApcStackwalkOperation(); // kmanager.InitiateApcStackwalkOperation();
break; // break;
case 8: //case 8:
kmanager.CheckForHiddenThreads(); // kmanager.CheckForHiddenThreads();
break; // break;
case 9: //case 9:
kmanager.CheckForEptHooks(); // kmanager.CheckForEptHooks();
break; // break;
} //}
kmanager.InitiateApcStackwalkOperation(); kmanager.InitiateApcStackwalkOperation();
kmanager.MonitorCallbackReports(); kmanager.MonitorCallbackReports();