diff --git a/driver/callbacks.c b/driver/callbacks.c index 162113f..8a7ebfd 100644 --- a/driver/callbacks.c +++ b/driver/callbacks.c @@ -2,17 +2,16 @@ #include "driver.h" -#include "queue.h" #include "pool.h" #include "thread.h" #include "modules.h" #include "imports.h" -#include "list.h" #include "session.h" #include "crypt.h" -#include "map.h" #include "util.h" -#include "tree.h" + +#include "containers/tree.h" +#include "containers/map.h" #define PROCESS_HASHMAP_BUCKET_COUNT 101 @@ -74,32 +73,42 @@ CleanupThreadListOnDriverUnload() VOID CleanupDriverListOnDriverUnload() { - PDRIVER_LIST_HEAD list = GetDriverList(); - for (;;) { - if (!ListFreeFirstEntry(&list->start, &list->lock, NULL)) - return; + PDRIVER_LIST_HEAD head = GetDriverList(); + PLIST_ENTRY entry = NULL; + + ImpKeAcquireGuardedMutex(&head->lock); + + while (!IsListEmpty(&head->list_entry)) { + entry = RemoveHeadList(&head->list_entry); + PDRIVER_LIST_ENTRY driverEntry = + CONTAINING_RECORD(entry, DRIVER_LIST_ENTRY, list_entry); + ExFreePoolWithTag(entry, POOL_TAG_DRIVER_LIST); } + + ImpKeReleaseGuardedMutex(&head->lock); } VOID EnumerateDriverListWithCallbackRoutine( _In_ DRIVERLIST_CALLBACK_ROUTINE CallbackRoutine, _In_opt_ PVOID Context) { - PDRIVER_LIST_HEAD list = GetDriverList(); - ImpKeAcquireGuardedMutex(&list->lock); + PDRIVER_LIST_HEAD head = GetDriverList(); + PLIST_ENTRY list_entry = NULL; + PDRIVER_LIST_ENTRY driver_entry = NULL; - if (!CallbackRoutine) - goto unlock; + ImpKeAcquireGuardedMutex(&head->lock); - PDRIVER_LIST_ENTRY entry = list->start.Next; - - while (entry) { - CallbackRoutine(entry, Context); - entry = (PDRIVER_LIST_ENTRY)entry->list.Next; + if (CallbackRoutine) { + list_entry = head->list_entry.Flink; + while (list_entry != &head->list_entry) { + driver_entry = + CONTAINING_RECORD(list_entry, DRIVER_LIST_ENTRY, list_entry); + CallbackRoutine(driver_entry, Context); + list_entry = list_entry->Flink; + } } -unlock: - ImpKeReleaseGuardedMutex(&list->lock); + ImpKeReleaseGuardedMutex(&head->lock); } VOID @@ -121,16 +130,17 @@ InitialiseDriverList() SYSTEM_MODULES modules = {0}; PDRIVER_LIST_ENTRY entry = NULL; PRTL_MODULE_EXTENDED_INFO module_entry = NULL; - PDRIVER_LIST_HEAD list = GetDriverList(); + PDRIVER_LIST_HEAD head = GetDriverList(); - InterlockedExchange(&list->active, TRUE); - ListInit(&list->start, &list->lock); - InitializeListHead(&list->deferred_list); + InterlockedExchange(&head->active, TRUE); + InitializeListHead(&head->list_entry); + InitializeListHead(&head->deferred_list); + KeInitializeGuardedMutex(&head->lock); - list->can_hash_x86 = FALSE; - list->work_item = IoAllocateWorkItem(GetDriverDeviceObject()); + head->can_hash_x86 = FALSE; + head->work_item = IoAllocateWorkItem(GetDriverDeviceObject()); - if (!list->work_item) + if (!head->work_item) return STATUS_INSUFFICIENT_RESOURCES; status = GetSystemModuleInformation(&modules); @@ -141,7 +151,7 @@ InitialiseDriverList() } /* skip hal.dll and ntoskrnl.exe */ - for (INT index = 2; index < modules.module_count; index++) { + for (UINT32 index = 2; index < modules.module_count; index++) { entry = ImpExAllocatePool2(POOL_FLAG_NON_PAGED, sizeof(DRIVER_LIST_ENTRY), POOL_TAG_DRIVER_LIST); @@ -166,17 +176,19 @@ InitialiseDriverList() status); entry->hashed = FALSE; entry->x86 = TRUE; - InsertHeadList(&list->deferred_list, &entry->deferred_entry); + InsertHeadList(&head->deferred_list, &entry->deferred_entry); } else if (!NT_SUCCESS(status)) { DEBUG_ERROR("HashModule failed with status %x", status); entry->hashed = FALSE; } - ListInsert(&list->start, entry, &list->lock); + KeAcquireGuardedMutex(&head->lock); + InsertHeadList(&head->list_entry, &entry->list_entry); + KeReleaseGuardedMutex(&head->lock); } - list->active = TRUE; + head->active = TRUE; if (modules.address) ImpExFreePoolWithTag(modules.address, SYSTEM_MODULES_POOL); @@ -193,22 +205,29 @@ VOID FindDriverEntryByBaseAddress(_In_ PVOID ImageBase, _Out_ PDRIVER_LIST_ENTRY* Entry) { - PDRIVER_LIST_HEAD list = GetDriverList(); - ImpKeAcquireGuardedMutex(&list->lock); + PDRIVER_LIST_HEAD head = GetDriverList(); + PLIST_ENTRY list_entry = NULL; + PDRIVER_LIST_ENTRY driver_entry = NULL; + + ImpKeAcquireGuardedMutex(&head->lock); *Entry = NULL; - PDRIVER_LIST_ENTRY entry = (PDRIVER_LIST_ENTRY)list->start.Next; + list_entry = head->list_entry.Flink; - while (entry) { - if (entry->ImageBase == ImageBase) { - *Entry = entry; + while (list_entry != &head->list_entry) { + driver_entry = + CONTAINING_RECORD(list_entry, DRIVER_LIST_ENTRY, list_entry); + + if (driver_entry->ImageBase == ImageBase) { + *Entry = driver_entry; goto unlock; } - entry = entry->list.Next; + list_entry = list_entry->Flink; } + unlock: - ImpKeReleaseGuardedMutex(&list->lock); + ImpKeReleaseGuardedMutex(&head->lock); } STATIC @@ -299,10 +318,10 @@ ImageLoadNotifyRoutineCallback(_In_opt_ PUNICODE_STRING FullImageName, NTSTATUS status = STATUS_UNSUCCESSFUL; PDRIVER_LIST_ENTRY entry = NULL; RTL_MODULE_EXTENDED_INFO module = {0}; - PDRIVER_LIST_HEAD list = GetDriverList(); + PDRIVER_LIST_HEAD head = GetDriverList(); ANSI_STRING ansi_path = {0}; - if (InterlockedExchange(&list->active, list->active) == FALSE) + if (InterlockedExchange(&head->active, head->active) == FALSE) return; if (ImageInfo->SystemModeImage == FALSE) { @@ -353,7 +372,9 @@ hash: entry->hashed = FALSE; } - ListInsert(&list->start, entry, &list->lock); + KeAcquireGuardedMutex(&head->lock); + InsertHeadList(&head->list_entry, &entry->list_entry); + KeReleaseGuardedMutex(&head->lock); } /* assumes map lock is held */ diff --git a/driver/callbacks.h b/driver/callbacks.h index dc53bad..02e2902 100644 --- a/driver/callbacks.h +++ b/driver/callbacks.h @@ -13,13 +13,13 @@ typedef void (*THREADLIST_CALLBACK_ROUTINE)( #define SHA_256_HASH_LENGTH 32 typedef struct _DRIVER_LIST_ENTRY { - SINGLE_LIST_ENTRY list; - PVOID ImageBase; - ULONG ImageSize; - BOOLEAN hashed; - BOOLEAN x86; - CHAR path[DRIVER_PATH_LENGTH]; - CHAR text_hash[SHA_256_HASH_LENGTH]; + LIST_ENTRY list_entry; + PVOID ImageBase; + ULONG ImageSize; + BOOLEAN hashed; + BOOLEAN x86; + CHAR path[DRIVER_PATH_LENGTH]; + CHAR text_hash[SHA_256_HASH_LENGTH]; /* * This LIST_ENTRY is to be used for modules where the hashing needs to diff --git a/driver/common.h b/driver/common.h index 31a75c2..29f8d6f 100644 --- a/driver/common.h +++ b/driver/common.h @@ -70,10 +70,10 @@ typedef struct _THREAD_LIST_HEAD { } THREAD_LIST_HEAD, *PTHREAD_LIST_HEAD; typedef struct _DRIVER_LIST_HEAD { - SINGLE_LIST_ENTRY start; - volatile ULONG count; - volatile BOOLEAN active; - KGUARDED_MUTEX lock; + LIST_ENTRY list_entry; + volatile ULONG count; + volatile BOOLEAN active; + KGUARDED_MUTEX lock; /* modules that need to be hashed later. */ PIO_WORKITEM work_item; @@ -84,11 +84,11 @@ typedef struct _DRIVER_LIST_HEAD { } DRIVER_LIST_HEAD, *PDRIVER_LIST_HEAD; typedef struct _THREAD_LIST_ENTRY { - HANDLE thread_id; - PKTHREAD thread; - PKPROCESS owning_process; - BOOLEAN apc_queued; - PKAPC apc; + HANDLE thread_id; + PKTHREAD thread; + PKPROCESS owning_process; + BOOLEAN apc_queued; + PKAPC apc; } THREAD_LIST_ENTRY, *PTHREAD_LIST_ENTRY; diff --git a/driver/map.c b/driver/containers/map.c similarity index 100% rename from driver/map.c rename to driver/containers/map.c diff --git a/driver/map.h b/driver/containers/map.h similarity index 99% rename from driver/map.h rename to driver/containers/map.h index 9b6fc97..5a2b827 100644 --- a/driver/map.h +++ b/driver/containers/map.h @@ -1,7 +1,7 @@ #ifndef MAP_H #define MAP_H -#include "common.h" +#include "../common.h" /* To improve efficiency, each entry contains a common header * RTL_HASHMAP_ENTRY*, reducing the need to store a seperate pointer to the diff --git a/driver/tree.c b/driver/containers/tree.c similarity index 100% rename from driver/tree.c rename to driver/containers/tree.c diff --git a/driver/tree.h b/driver/containers/tree.h similarity index 98% rename from driver/tree.h rename to driver/containers/tree.h index 4ebda95..3469656 100644 --- a/driver/tree.h +++ b/driver/containers/tree.h @@ -1,7 +1,7 @@ #ifndef TREE_H #define TREE_H -#include "common.h" +#include "../common.h" #define RB_TREE_EQUAL 0 #define RB_TREE_LESS_THAN 1 diff --git a/driver/driver.h b/driver/driver.h index 350945c..0960e37 100644 --- a/driver/driver.h +++ b/driver/driver.h @@ -5,12 +5,12 @@ #include -#include "queue.h" #include "modules.h" #include "integrity.h" #include "callbacks.h" -#include "map.h" -#include "tree.h" + +#include "containers/map.h" +#include "containers/tree.h" BCRYPT_ALG_HANDLE* GetCryptHandle_AES(); diff --git a/driver/driver.vcxproj b/driver/driver.vcxproj index ff611e3..82e8032 100644 --- a/driver/driver.vcxproj +++ b/driver/driver.vcxproj @@ -252,16 +252,14 @@ - - + - - + @@ -276,15 +274,13 @@ - - + - - + diff --git a/driver/driver.vcxproj.filters b/driver/driver.vcxproj.filters index 3f9c5b8..05f8a97 100644 --- a/driver/driver.vcxproj.filters +++ b/driver/driver.vcxproj.filters @@ -33,9 +33,6 @@ Source Files - - Source Files - Source Files @@ -54,9 +51,6 @@ Source Files - - Source Files - Source Files @@ -75,10 +69,10 @@ Source Files - + Source Files - + Source Files @@ -95,9 +89,6 @@ Header Files - - Header Files - Header Files @@ -119,9 +110,6 @@ Header Files - - Header Files - Header Files @@ -152,10 +140,10 @@ Header Files - + Header Files - + Header Files diff --git a/driver/io.c b/driver/io.c index ff7f266..1ff47f3 100644 --- a/driver/io.c +++ b/driver/io.c @@ -6,13 +6,13 @@ #include "pool.h" #include "integrity.h" #include "thread.h" -#include "queue.h" + #include "hv.h" #include "imports.h" -#include "list.h" + #include "session.h" #include "hw.h" -#include "map.h" +#include "containers/map.h" STATIC NTSTATUS diff --git a/driver/list.c b/driver/list.c deleted file mode 100644 index 8d9481b..0000000 --- a/driver/list.c +++ /dev/null @@ -1,172 +0,0 @@ -#include "list.h" - -#include "imports.h" -#include "driver.h" - -/* - * Simple thread safe linked list implementation. All structures should begin - * with a SINGLE_LIST_ENTRY structure provided by the windows API. for example: - * - * typedef struct _LIST_ENTRY_STRUCTURE - * { - * SINGLE_LIST_ENTRY list; - * PVOID address; - * UINT32 data; - * ... - * }; - * - * This common structure layout allows us to pass in a callback routine when - *freeing allowing immense flexibility to ensure we can free and/or deference - *any objects that are referenced in said object. - * - * I've opted to use a mutex rather then a spinlock since there are many times - *we enumerate the list for extended periods aswell as queue up many insertions - *at once. - */ - -#define LIST_POOL_TAG 'list' - -VOID -ListInit(_Inout_ PSINGLE_LIST_ENTRY Head, _Inout_ PKGUARDED_MUTEX Lock) -{ - ImpKeInitializeGuardedMutex(Lock); - Head->Next = NULL; -} - -VOID -ListInsert(_Inout_ PSINGLE_LIST_ENTRY Head, - _Inout_ PSINGLE_LIST_ENTRY NewEntry, - _In_ PKGUARDED_MUTEX Lock) -{ - ImpKeAcquireGuardedMutex(Lock); - - PSINGLE_LIST_ENTRY old_entry = Head->Next; - - Head->Next = NewEntry; - NewEntry->Next = old_entry; - - ImpKeReleaseGuardedMutex(Lock); -} - -/* - * Assuming the SINGLE_LIST_ENTRY is the first item in the structure, we - * can pass a callback routine to be called before the free occurs. This - * allows us to dereference/free structure specific items whilst still allowing - * the list to remain flexible. - */ -BOOLEAN -ListFreeFirstEntry(_Inout_ PSINGLE_LIST_ENTRY Head, - _In_ PKGUARDED_MUTEX Lock, - _In_opt_ FREE_LIST_ITEM_CALLBACK CallbackRoutine) -{ - BOOLEAN result = FALSE; - ImpKeAcquireGuardedMutex(Lock); - - if (Head->Next) { - PSINGLE_LIST_ENTRY entry = Head->Next; - - if (CallbackRoutine) - CallbackRoutine(entry); - - Head->Next = Head->Next->Next; - ImpExFreePoolWithTag(entry, POOL_TAG_THREAD_LIST); - result = TRUE; - } - - ImpKeReleaseGuardedMutex(Lock); - return result; -} - -/* - * If we are removing a specific entry, its assumed we have freed and/or - * dereferenced any fields in the structure. - */ -VOID -ListRemoveEntry(_Inout_ PSINGLE_LIST_ENTRY Head, - _Inout_ PSINGLE_LIST_ENTRY Entry, - _In_ PKGUARDED_MUTEX Lock) -{ - ImpKeAcquireGuardedMutex(Lock); - - PSINGLE_LIST_ENTRY entry = Head->Next; - - if (!entry) - goto unlock; - - if (entry == Entry) { - Head->Next = entry->Next; - ImpExFreePoolWithTag(Entry, POOL_TAG_THREAD_LIST); - goto unlock; - } - - while (entry->Next) { - if (entry->Next == Entry) { - entry->Next = Entry->Next; - ImpExFreePoolWithTag(Entry, POOL_TAG_THREAD_LIST); - goto unlock; - } - - entry = entry->Next; - } - -unlock: - ImpKeReleaseGuardedMutex(Lock); -} - -VOID -LookasideListRemoveEntry(_Inout_ PSINGLE_LIST_ENTRY Head, - _Inout_ PSINGLE_LIST_ENTRY Entry, - _In_ PKGUARDED_MUTEX Lock) -{ - ImpKeAcquireGuardedMutex(Lock); - - PTHREAD_LIST_HEAD head = GetThreadTree(); - PSINGLE_LIST_ENTRY entry = Head->Next; - - if (!entry) - goto unlock; - - if (entry == Entry) { - Head->Next = entry->Next; - ExFreeToLookasideListEx(&head->lookaside_list, Entry); - goto unlock; - } - - while (entry->Next) { - if (entry->Next == Entry) { - entry->Next = Entry->Next; - ExFreeToLookasideListEx(&head->lookaside_list, Entry); - goto unlock; - } - - entry = entry->Next; - } - -unlock: - ImpKeReleaseGuardedMutex(Lock); -} - -BOOLEAN -LookasideListFreeFirstEntry(_Inout_ PSINGLE_LIST_ENTRY Head, - _In_ PKGUARDED_MUTEX Lock, - _In_opt_ FREE_LIST_ITEM_CALLBACK CallbackRoutine) -{ - ImpKeAcquireGuardedMutex(Lock); - - PTHREAD_LIST_HEAD head = GetThreadTree(); - BOOLEAN result = FALSE; - - if (Head->Next) { - PSINGLE_LIST_ENTRY entry = Head->Next; - - if (CallbackRoutine) - CallbackRoutine(entry); - - Head->Next = Head->Next->Next; - ExFreeToLookasideListEx(&head->lookaside_list, entry); - result = TRUE; - } - - ImpKeReleaseGuardedMutex(Lock); - return result; -} \ No newline at end of file diff --git a/driver/list.h b/driver/list.h deleted file mode 100644 index 3102acb..0000000 --- a/driver/list.h +++ /dev/null @@ -1,36 +0,0 @@ -#ifndef LIST_H -#define LIST_H - -#include "common.h" - -typedef void (*FREE_LIST_ITEM_CALLBACK)(_In_ PVOID Entry); - -VOID -LookasideListRemoveEntry(_Inout_ PSINGLE_LIST_ENTRY Head, - _Inout_ PSINGLE_LIST_ENTRY Entry, - _In_ PKGUARDED_MUTEX Lock); - -BOOLEAN -LookasideListFreeFirstEntry(_Inout_ PSINGLE_LIST_ENTRY Head, - _In_ PKGUARDED_MUTEX Lock, - _In_opt_ FREE_LIST_ITEM_CALLBACK CallbackRoutine); - -VOID -ListInit(_Inout_ PSINGLE_LIST_ENTRY Head, _Inout_ PKGUARDED_MUTEX Lock); - -VOID -ListInsert(_Inout_ PSINGLE_LIST_ENTRY Head, - _Inout_ PSINGLE_LIST_ENTRY NewEntry, - _In_ PKGUARDED_MUTEX Lock); - -BOOLEAN -ListFreeFirstEntry(_Inout_ PSINGLE_LIST_ENTRY Head, - _In_ PKGUARDED_MUTEX Lock, - _In_opt_ FREE_LIST_ITEM_CALLBACK CallbackRoutine); - -VOID -ListRemoveEntry(_Inout_ PSINGLE_LIST_ENTRY Head, - _Inout_ PSINGLE_LIST_ENTRY Entry, - _In_ PKGUARDED_MUTEX Lock); - -#endif \ No newline at end of file diff --git a/driver/modules.c b/driver/modules.c index 36ed4db..95157bf 100644 --- a/driver/modules.c +++ b/driver/modules.c @@ -9,7 +9,7 @@ #include "thread.h" #include "pe.h" #include "crypt.h" -#include "tree.h" +#include "containers/tree.h" #define WHITELISTED_MODULE_TAG 'whte' diff --git a/driver/modules.h b/driver/modules.h index 9364832..337f264 100644 --- a/driver/modules.h +++ b/driver/modules.h @@ -5,7 +5,6 @@ #include #include "common.h" -#include "queue.h" typedef struct _APC_OPERATION_ID { int operation_id; diff --git a/driver/pool.c b/driver/pool.c index 90a1a92..edbcfe8 100644 --- a/driver/pool.c +++ b/driver/pool.c @@ -3,7 +3,7 @@ #include #include "callbacks.h" -#include "queue.h" + #include "ia32.h" #include "imports.h" #include "crypt.h" diff --git a/driver/queue.c b/driver/queue.c deleted file mode 100644 index 31b1df9..0000000 --- a/driver/queue.c +++ /dev/null @@ -1,65 +0,0 @@ -#include "queue.h" - -#include "callbacks.h" - -#include "driver.h" - -#include "queue.h" -#include "pool.h" -#include "thread.h" -#include "io.h" -#include "common.h" -#include "imports.h" - -VOID -QueuePush(_Inout_ PQUEUE_HEAD Head, _In_ PVOID Data) -{ - ImpKeAcquireGuardedMutex(&Head->lock); - - PQUEUE_NODE temp = ExAllocatePool2( - POOL_FLAG_NON_PAGED, sizeof(QUEUE_NODE), QUEUE_POOL_TAG); - - if (!temp) - goto end; - - Head->entries += 1; - - temp->data = Data; - - if (Head->end != NULL) - Head->end->next = temp; - - Head->end = temp; - - if (Head->start == NULL) - Head->start = temp; - -end: - ImpKeReleaseGuardedMutex(&Head->lock); -} - -PVOID -QueuePop(_Inout_ PQUEUE_HEAD Head) -{ - ImpKeAcquireGuardedMutex(&Head->lock); - - PVOID data = NULL; - PQUEUE_NODE temp = Head->start; - - if (temp == NULL) - goto end; - - Head->entries = Head->entries - 1; - - data = temp->data; - Head->start = temp->next; - - if (Head->end == temp) - Head->end = NULL; - - ImpExFreePoolWithTag(temp, QUEUE_POOL_TAG); - -end: - ImpKeReleaseGuardedMutex(&Head->lock); - return data; -} \ No newline at end of file diff --git a/driver/queue.h b/driver/queue.h deleted file mode 100644 index 633507e..0000000 --- a/driver/queue.h +++ /dev/null @@ -1,54 +0,0 @@ -#ifndef QUEUE_H -#define QUEUE_H - -#include -#include "common.h" - -#define MAX_REPORTS_PER_IRP 20 - -typedef struct QUEUE_HEAD { - struct _QUEUE_NODE* start; - struct _QUEUE_NODE* end; - KGUARDED_MUTEX lock; - INT entries; - -} QUEUE_HEAD, *PQUEUE_HEAD; - -/* - * 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. - */ -typedef struct _REPORT_QUEUE_HEAD { - QUEUE_HEAD head; - volatile BOOLEAN is_driver_unloading; - KGUARDED_MUTEX lock; - -} REPORT_QUEUE_HEAD, *PREPORT_QUEUE_HEAD; - -typedef struct _QUEUE_NODE { - struct _QUEUE_NODE* next; - PVOID data; - -} QUEUE_NODE, *PQUEUE_NODE; - -typedef struct _GLOBAL_REPORT_QUEUE_HEADER { - INT count; - -} GLOBAL_REPORT_QUEUE_HEADER, *PGLOBAL_REPORT_QUEUE_HEADER; - -typedef struct _REPORT_HEADER { - INT report_id; - -} REPORT_HEADER, *PREPORT_HEADER; - -VOID -QueuePush(_Inout_ PQUEUE_HEAD Head, _In_ PVOID Data); - -PVOID -QueuePop(_Inout_ PQUEUE_HEAD Head); - -#endif \ No newline at end of file diff --git a/driver/thread.c b/driver/thread.c index 42d19e9..0e65f2b 100644 --- a/driver/thread.c +++ b/driver/thread.c @@ -5,10 +5,10 @@ #include "pool.h" #include "callbacks.h" #include "driver.h" -#include "queue.h" + #include "session.h" #include "imports.h" -#include "tree.h" +#include "containers/tree.h" #include "crypt.h" #ifdef ALLOC_PRAGMA