diff --git a/.clang-format b/.clang-format index f33dc92..09eafac 100644 --- a/.clang-format +++ b/.clang-format @@ -12,11 +12,10 @@ AlignOperands: true AlignTrailingComments: true -AllowAllArgumentsOnNextLine: true +AllowAllArgumentsOnNextLine: false AllowShortBlocksOnASingleLine: true AllowShortCaseLabelsOnASingleLine: true -AllowShortFunctionsOnASingleLine: false AllowShortIfStatementsOnASingleLine: false AllowShortLoopsOnASingleLine: false AlwaysBreakAfterReturnType: TopLevel @@ -26,8 +25,9 @@ AlwaysBreakTemplateDeclarations: true #false BinPackArguments: false BinPackParameters: false - +AllowShortFunctionsOnASingleLine: false AllowAllParametersOfDeclarationOnNextLine: true +PenaltyBreakBeforeFirstCallParameter: 0 BreakBeforeBraces: Stroustrup BraceWrapping: diff --git a/driver/common.h b/driver/common.h index 29f8d6f..682c8bd 100644 --- a/driver/common.h +++ b/driver/common.h @@ -53,9 +53,11 @@ #define MAX_MODULE_PATH 260 -#define CONVERT_RELATIVE_ADDRESS(Cast, Base, Rel) \ +#define RVA(Cast, Base, Rel) \ ((Cast)((DWORD_PTR)(Base) + (DWORD_PTR)(Rel))) +#define ARRAYLEN(len, type) ((len) / sizeof(type)) + /* * Interlocked intrinsics are only atomic with respect to other InterlockedXxx * functions, so all reads and writes to the THREAD_LIST->active flag must be diff --git a/driver/containers/tree.h b/driver/containers/tree.h index d5b4409..c71ebe1 100644 --- a/driver/containers/tree.h +++ b/driver/containers/tree.h @@ -3,8 +3,8 @@ #include "../common.h" -#define RB_TREE_EQUAL 0 -#define RB_TREE_LESS_THAN 1 +#define RB_TREE_EQUAL 0 +#define RB_TREE_LESS_THAN 1 #define RB_TREE_GREATER_THAN 2 typedef enum _COLOUR { red, black } COLOUR; @@ -55,6 +55,9 @@ RtlRbTreeEnumerate(_In_ PRB_TREE Tree, _In_ RB_ENUM_CALLBACK Callback, _In_opt_ PVOID Context); +#define ENUMERATE_THREADS(callback, context) \ + RtlRbTreeEnumerate(GetThreadTree(), callback, context) + VOID RtlRbTreeDeleteTree(_In_ PRB_TREE Tree); diff --git a/driver/crypt.c b/driver/crypt.c index b4607e3..1c1e0f6 100644 --- a/driver/crypt.c +++ b/driver/crypt.c @@ -58,14 +58,16 @@ CryptEncryptImportsArray(_In_ PUINT64 Array, _In_ UINT32 Entries) __m256i load_block = {0}; __m256i xored_block = {0}; - RtlCopyMemory( - ¤t_block, &Array[block_index * block_size], sizeof(__m256i)); + RtlCopyMemory(¤t_block, + &Array[block_index * block_size], + sizeof(__m256i)); load_block = _mm256_loadu_si256(¤t_block); xored_block = _mm256_xor_si256(load_block, *imports_key); - RtlCopyMemory( - &Array[block_index * block_size], &xored_block, sizeof(__m256i)); + RtlCopyMemory(&Array[block_index * block_size], + &xored_block, + sizeof(__m256i)); } } @@ -78,8 +80,9 @@ CryptDecryptImportBlock(_In_ PUINT64 Array, _In_ UINT32 BlockIndex) __m256i* imports_key = GetDriverImportsKey(); UINT32 block_size = sizeof(__m256i) / sizeof(UINT64); - RtlCopyMemory( - &load_block, &Array[BlockIndex * block_size], sizeof(__m256i)); + RtlCopyMemory(&load_block, + &Array[BlockIndex * block_size], + sizeof(__m256i)); return _mm256_xor_si256(load_block, *imports_key); } @@ -128,8 +131,10 @@ CryptDecryptImportsArrayEntry(_In_ PUINT64 Array, UINT32 block_sub_index = 0; UINT64 pointer = 0; - CryptFindContainingBlockForArrayIndex( - EntryIndex, block_size, &containing_block_index, &block_sub_index); + CryptFindContainingBlockForArrayIndex(EntryIndex, + block_size, + &containing_block_index, + &block_sub_index); original_block = CryptDecryptImportBlock(Array, containing_block_index); @@ -281,8 +286,9 @@ CryptInitialiseSessionCryptObjects() goto end; } - session->key_object = ExAllocatePool2( - POOL_FLAG_NON_PAGED, session->key_object_length, POOL_TAG_CRYPT); + session->key_object = ExAllocatePool2(POOL_FLAG_NON_PAGED, + session->key_object_length, + POOL_TAG_CRYPT); if (!session->key_object) { status = STATUS_INSUFFICIENT_RESOURCES; @@ -323,8 +329,10 @@ CryptInitialiseProvider() NTSTATUS status = STATUS_UNSUCCESSFUL; BCRYPT_ALG_HANDLE* handle = GetCryptHandle_AES(); - status = BCryptOpenAlgorithmProvider( - handle, BCRYPT_AES_ALGORITHM, NULL, BCRYPT_PROV_DISPATCH); + status = BCryptOpenAlgorithmProvider(handle, + BCRYPT_AES_ALGORITHM, + NULL, + BCRYPT_PROV_DISPATCH); if (!NT_SUCCESS(status)) DEBUG_ERROR("BCryptOpenAlgorithmProvider: %x", status); @@ -499,4 +507,132 @@ TpmExtractEndorsementKey() DEBUG_INFO("TPM2.0 PTP Interface Type: %x", (UINT32)type); return status; +} + +NTSTATUS +CryptHashBuffer_sha256(_In_ PVOID Buffer, + _In_ ULONG BufferSize, + _Out_ PVOID* HashResult, + _Out_ PULONG HashResultSize) +{ + PAGED_CODE(); + + NTSTATUS status = STATUS_UNSUCCESSFUL; + BCRYPT_ALG_HANDLE* algo_handle = GetCryptHandle_Sha256(); + BCRYPT_HASH_HANDLE hash_handle = NULL; + ULONG bytes_copied = 0; + ULONG resulting_hash_size = 0; + ULONG hash_object_size = 0; + PCHAR hash_object = NULL; + PCHAR resulting_hash = NULL; + + *HashResult = NULL; + *HashResultSize = 0; + + /* + * Request the size of the hash object buffer, this is different then + * the buffer that will store the resulting hash, instead this will be + * used to store the hash object used to create the hash. + */ + status = BCryptGetProperty(*algo_handle, + BCRYPT_OBJECT_LENGTH, + (PCHAR)&hash_object_size, + sizeof(ULONG), + &bytes_copied, + NULL); + + if (!NT_SUCCESS(status)) { + DEBUG_ERROR("BCryptGetProperty failed with status %x", status); + goto end; + } + + hash_object = ImpExAllocatePool2(POOL_FLAG_NON_PAGED, + hash_object_size, + POOL_TAG_INTEGRITY); + + if (!hash_object) { + status = STATUS_MEMORY_NOT_ALLOCATED; + goto end; + } + + /* + * This call gets the size of the resulting hash, which we will use to + * allocate the resulting hash buffer. + */ + status = BCryptGetProperty(*algo_handle, + BCRYPT_HASH_LENGTH, + (PCHAR)&resulting_hash_size, + sizeof(ULONG), + &bytes_copied, + NULL); + + if (!NT_SUCCESS(status)) { + DEBUG_ERROR("BCryptGetProperty failed with status %x", status); + goto end; + } + + resulting_hash = ImpExAllocatePool2(POOL_FLAG_NON_PAGED, + resulting_hash_size, + POOL_TAG_INTEGRITY); + + if (!resulting_hash) { + status = STATUS_MEMORY_NOT_ALLOCATED; + goto end; + } + + /* + * Here we create our hash object and store it in the hash_object + * buffer. + */ + status = BCryptCreateHash(*algo_handle, + &hash_handle, + hash_object, + hash_object_size, + NULL, + NULL, + NULL); + + if (!NT_SUCCESS(status)) { + DEBUG_ERROR("BCryptCreateHash failed with status %x", status); + goto end; + } + + /* + * This function hashes the buffer, but does NOT store it in our + * resulting buffer yet, we need to call BCryptFinishHash to retrieve + * the final hash. + */ + status = BCryptHashData(hash_handle, Buffer, BufferSize, NULL); + + if (!NT_SUCCESS(status)) { + DEBUG_ERROR("BCryptHashData failed with status %x", status); + goto end; + } + + /* + * As said in the previous comment, this is where we retrieve the final + * hash and store it in our output buffer. + */ + status = BCryptFinishHash(hash_handle, + resulting_hash, + resulting_hash_size, + NULL); + + if (!NT_SUCCESS(status)) { + DEBUG_ERROR("BCryptFinishHash failed with status %x", status); + goto end; + } + + *HashResult = resulting_hash; + *HashResultSize = resulting_hash_size; + +end: + + if (hash_handle) + BCryptDestroyHash(hash_handle); + + if (hash_object) + ImpExFreePoolWithTag(hash_object, POOL_TAG_INTEGRITY); + + return status; } \ No newline at end of file diff --git a/driver/crypt.h b/driver/crypt.h index 6636618..5218ab3 100644 --- a/driver/crypt.h +++ b/driver/crypt.h @@ -69,4 +69,10 @@ CryptDecryptPointer64(_Inout_ PUINT64 Pointer, _In_ UINT64 Key); UINT64 CryptDecryptPointerOutOfPlace64(_In_ PUINT64 Pointer, _In_ UINT64 Key); +NTSTATUS +CryptHashBuffer_sha256(_In_ PVOID Buffer, + _In_ ULONG BufferSize, + _Out_ PVOID* HashResult, + _Out_ PULONG HashResultSize); + #endif \ No newline at end of file diff --git a/driver/integrity.c b/driver/integrity.c index 4b0a6f6..97fc9b4 100644 --- a/driver/integrity.c +++ b/driver/integrity.c @@ -15,9 +15,11 @@ #include #include +/* Header for a buffer that contains an array of sections copied from a module + */ typedef struct _INTEGRITY_CHECK_HEADER { - INT executable_section_count; - LONG total_packet_size; + UINT32 section_count; + UINT32 total_size; } INTEGRITY_CHECK_HEADER, *PINTEGRITY_CHECK_HEADER; @@ -28,72 +30,89 @@ typedef struct _PROCESS_MODULE_INFORMATION { } PROCESS_MODULE_INFORMATION, *PPROCESS_MODULE_INFORMATION; +/* Structure representing the data passed back to user-mode after validating a + * process module sections*/ typedef struct _PROCESS_MODULE_VALIDATION_RESULT { - INT is_module_valid; + UINT32 is_module_valid; } PROCESS_MODULE_VALIDATION_RESULT, *PPROCESS_MODULE_VALIDATION_RESULT; +typedef struct _VAL_INTEGRITY_HEADER { + INTEGRITY_CHECK_HEADER integrity_check_header; + IMAGE_SECTION_HEADER section_header; + CHAR section_base[]; + +} VAL_INTEGRITY_HEADER, *PVAL_INTEGRITY_HEADER; + +// clang-format off + STATIC NTSTATUS InitiateEptFunctionAddressArrays(); STATIC NTSTATUS -GetModuleInformationByName(_Out_ PRTL_MODULE_EXTENDED_INFO ModuleInfo, - _In_ LPCSTR ModuleName); +GetModuleInformationByName( + _Out_ PRTL_MODULE_EXTENDED_INFO ModuleInfo, + _In_ LPCSTR ModuleName); STATIC NTSTATUS -StoreModuleExecutableRegionsInBuffer(_Out_ PVOID* Buffer, - _In_ PVOID ModuleBase, - _In_ SIZE_T ModuleSize, - _Out_ PSIZE_T BytesWritten, - _In_ BOOLEAN IsModulex86); +StoreModuleExecutableRegionsInBuffer( + _Out_ PVOID* Buffer, + _In_ PVOID ModuleBase, + _In_ SIZE_T ModuleSize, + _Out_ PSIZE_T BytesWritten, + _In_ BOOLEAN IsModulex86 +); STATIC NTSTATUS -MapDiskImageIntoVirtualAddressSpace(_Inout_ PHANDLE SectionHandle, - _Out_ PVOID* Section, - _In_ PUNICODE_STRING Path, - _Out_ PSIZE_T Size); - -STATIC -NTSTATUS -ComputeHashOfBuffer(_In_ PVOID Buffer, - _In_ ULONG BufferSize, - _Out_ PVOID* HashResult, - _Out_ PULONG HashResultSize); +MapDiskImageIntoVirtualAddressSpace( + _Inout_ PHANDLE SectionHandle, + _Out_ PVOID* Section, + _In_ PUNICODE_STRING Path, + _Out_ PSIZE_T Size +); STATIC VOID -GetNextSMBIOSStructureInTable(_Inout_ PSMBIOS_TABLE_HEADER* CurrentStructure); +GetNextSMBIOSStructureInTable( + _Inout_ PSMBIOS_TABLE_HEADER* CurrentStructure +); STATIC NTSTATUS -GetStringAtIndexFromSMBIOSTable(_In_ PSMBIOS_TABLE_HEADER Table, - _In_ INT Index, - _In_ PVOID Buffer, - _In_ SIZE_T BufferSize); +GetStringAtIndexFromSMBIOSTable( + _In_ PSMBIOS_TABLE_HEADER Table, + _In_ UINT32 Index, + _In_ PVOID Buffer, + _In_ SIZE_T BufferSize +); STATIC NTSTATUS -GetAverageReadTimeAtRoutine(_In_ PVOID RoutineAddress, - _Out_ PUINT64 AverageTime); +GetAverageReadTimeAtRoutine( + _In_ PVOID RoutineAddress, + _Out_ PUINT64 AverageTime +); STATIC VOID -HeartbeatDpcRoutine(_In_ PKDPC Dpc, - _In_opt_ PVOID DeferredContext, - _In_opt_ PVOID SystemArgument1, - _In_opt_ PVOID SystemArgument2); +HeartbeatDpcRoutine( + _In_ PKDPC Dpc, + _In_opt_ PVOID DeferredContext, + _In_opt_ PVOID SystemArgument1, + _In_opt_ PVOID SystemArgument2 +); + +// clang-format on #ifdef ALLOC_PRAGMA # pragma alloc_text(PAGE, GetDriverImageSize) # pragma alloc_text(PAGE, GetModuleInformationByName) # pragma alloc_text(PAGE, StoreModuleExecutableRegionsInBuffer) # pragma alloc_text(PAGE, MapDiskImageIntoVirtualAddressSpace) -# pragma alloc_text(PAGE, ComputeHashOfBuffer) -// # pragma alloc_text(PAGE, VerifyInMemoryImageVsDiskImage) # pragma alloc_text(PAGE, RetrieveInMemoryModuleExecutableSections) # pragma alloc_text(PAGE, GetNextSMBIOSStructureInTable) # pragma alloc_text(PAGE, GetStringAtIndexFromSMBIOSTable) @@ -103,7 +122,6 @@ HeartbeatDpcRoutine(_In_ PKDPC Dpc, # pragma alloc_text(PAGE, ScanForSignature) # pragma alloc_text(PAGE, InitiateEptFunctionAddressArrays) # pragma alloc_text(PAGE, DetectEptHooksInKeyFunctions) -// #pragma alloc_text(PAGE, DetermineIfTestSigningIsEnabled) #endif /* @@ -144,6 +162,7 @@ GetDriverImageSize(_Inout_ PIRP Irp) } Irp->IoStatus.Information = sizeof(ULONG); + RtlCopyMemory(Irp->AssociatedIrp.SystemBuffer, &driver_info->ImageSize, sizeof(ULONG)); @@ -232,8 +251,8 @@ InitIntegrityCheckHeader(_Out_ PINTEGRITY_CHECK_HEADER Header, _In_ UINT32 SectionCount, _In_ UINT32 TotalSize) { - Header->executable_section_count = SectionCount; - Header->total_packet_size = TotalSize + sizeof(INTEGRITY_CHECK_HEADER); + Header->section_count = SectionCount; + Header->total_size = TotalSize + sizeof(INTEGRITY_CHECK_HEADER); } STATIC @@ -246,21 +265,17 @@ StoreModuleExecutableRegionsInBuffer(_Out_ PVOID* Buffer, { PAGED_CODE(); + UINT32 total_packet_size = 0; + UINT32 num_sections = 0; + UINT32 num_executable_sections = 0; + UINT64 buffer_base = 0; + UINT32 bytes_returned = 0; NTSTATUS status = STATUS_UNSUCCESSFUL; PNT_HEADER_64 nt_header = NULL; PIMAGE_SECTION_HEADER section = NULL; - ULONG total_packet_size = 0; - ULONG num_sections = 0; - ULONG num_executable_sections = 0; - UINT64 buffer_base = 0; - ULONG bytes_returned = 0; MM_COPY_ADDRESS address = {0}; INTEGRITY_CHECK_HEADER header = {0}; - // DEBUG_VERBOSE("Storing x regions -> x86 module: %lx", - // (UINT32)IsModulex86); DEBUG_VERBOSE("MmIsAddressValid: %lx", - // MmIsAddressValid(ModuleBase)); - if (!ModuleBase || !ModuleSize) return STATUS_INVALID_PARAMETER; @@ -300,7 +315,7 @@ StoreModuleExecutableRegionsInBuffer(_Out_ PVOID* Buffer, section = IMAGE_FIRST_SECTION(nt_header); buffer_base = (UINT64)*Buffer + sizeof(INTEGRITY_CHECK_HEADER); - for (ULONG index = 0; index < num_sections - 1; index++) { + for (UINT32 index = 0; index < num_sections - 1; index++) { if (!IsSectionExecutable(section)) { section++; continue; @@ -338,8 +353,9 @@ StoreModuleExecutableRegionsInBuffer(_Out_ PVOID* Buffer, section++; } - InitIntegrityCheckHeader( - &header, num_executable_sections, total_packet_size); + InitIntegrityCheckHeader(&header, + num_executable_sections, + total_packet_size); RtlCopyMemory(*Buffer, &header, sizeof(INTEGRITY_CHECK_HEADER)); *BytesWritten = total_packet_size + sizeof(INTEGRITY_CHECK_HEADER); @@ -355,29 +371,27 @@ MapDiskImageIntoVirtualAddressSpace(_Inout_ PHANDLE SectionHandle, { PAGED_CODE(); - NTSTATUS status = STATUS_UNSUCCESSFUL; - HANDLE file_handle = NULL; - OBJECT_ATTRIBUTES object_attributes = {0}; - PIO_STATUS_BLOCK pio_block = NULL; - UNICODE_STRING path = {0}; + NTSTATUS status = STATUS_UNSUCCESSFUL; + HANDLE handle = NULL; + OBJECT_ATTRIBUTES oa = {0}; + PIO_STATUS_BLOCK io = NULL; + UNICODE_STRING path = {0}; *Section = NULL; *Size = 0; ImpRtlInitUnicodeString(&path, Path->Buffer); - InitializeObjectAttributes( - &object_attributes, &path, OBJ_KERNEL_HANDLE, NULL, NULL); + InitializeObjectAttributes(&oa, &path, OBJ_KERNEL_HANDLE, NULL, NULL); - status = ImpZwOpenFile( - &file_handle, GENERIC_READ, &object_attributes, &pio_block, NULL, NULL); + status = ImpZwOpenFile(&handle, GENERIC_READ, &oa, &io, NULL, NULL); if (!NT_SUCCESS(status)) { DEBUG_ERROR("ZwOpenFile failed with status %x", status); return status; } - object_attributes.ObjectName = NULL; + oa.ObjectName = NULL; /* * Its important that we set the SEC_IMAGE flag with the PAGE_READONLY @@ -385,15 +399,15 @@ MapDiskImageIntoVirtualAddressSpace(_Inout_ PHANDLE SectionHandle, */ status = ImpZwCreateSection(SectionHandle, SECTION_ALL_ACCESS, - &object_attributes, + &oa, NULL, PAGE_READONLY, SEC_IMAGE, - file_handle); + handle); if (!NT_SUCCESS(status)) { DEBUG_ERROR("ZwCreateSection failed with status %x", status); - ImpZwClose(file_handle); + ImpZwClose(handle); *SectionHandle = NULL; return status; } @@ -422,138 +436,13 @@ MapDiskImageIntoVirtualAddressSpace(_Inout_ PHANDLE SectionHandle, if (!NT_SUCCESS(status)) { DEBUG_ERROR("ZwMapViewOfSection failed with status %x", status); - ImpZwClose(file_handle); + ImpZwClose(handle); ImpZwClose(*SectionHandle); *SectionHandle = NULL; return status; } - ImpZwClose(file_handle); - return status; -} - -STATIC -NTSTATUS -ComputeHashOfBuffer(_In_ PVOID Buffer, - _In_ ULONG BufferSize, - _Out_ PVOID* HashResult, - _Out_ PULONG HashResultSize) -{ - PAGED_CODE(); - - NTSTATUS status = STATUS_UNSUCCESSFUL; - BCRYPT_ALG_HANDLE* algo_handle = GetCryptHandle_Sha256(); - BCRYPT_HASH_HANDLE hash_handle = NULL; - ULONG bytes_copied = 0; - ULONG resulting_hash_size = 0; - ULONG hash_object_size = 0; - PCHAR hash_object = NULL; - PCHAR resulting_hash = NULL; - - *HashResult = NULL; - *HashResultSize = 0; - - /* - * Request the size of the hash object buffer, this is different then - * the buffer that will store the resulting hash, instead this will be - * used to store the hash object used to create the hash. - */ - status = BCryptGetProperty(*algo_handle, - BCRYPT_OBJECT_LENGTH, - (PCHAR)&hash_object_size, - sizeof(ULONG), - &bytes_copied, - NULL); - - if (!NT_SUCCESS(status)) { - DEBUG_ERROR("BCryptGetProperty failed with status %x", status); - goto end; - } - - hash_object = ImpExAllocatePool2( - POOL_FLAG_NON_PAGED, hash_object_size, POOL_TAG_INTEGRITY); - - if (!hash_object) { - status = STATUS_MEMORY_NOT_ALLOCATED; - goto end; - } - - /* - * This call gets the size of the resulting hash, which we will use to - * allocate the resulting hash buffer. - */ - status = BCryptGetProperty(*algo_handle, - BCRYPT_HASH_LENGTH, - (PCHAR)&resulting_hash_size, - sizeof(ULONG), - &bytes_copied, - NULL); - - if (!NT_SUCCESS(status)) { - DEBUG_ERROR("BCryptGetProperty failed with status %x", status); - goto end; - } - - resulting_hash = ImpExAllocatePool2( - POOL_FLAG_NON_PAGED, resulting_hash_size, POOL_TAG_INTEGRITY); - - if (!resulting_hash) { - status = STATUS_MEMORY_NOT_ALLOCATED; - goto end; - } - - /* - * Here we create our hash object and store it in the hash_object - * buffer. - */ - status = BCryptCreateHash(*algo_handle, - &hash_handle, - hash_object, - hash_object_size, - NULL, - NULL, - NULL); - - if (!NT_SUCCESS(status)) { - DEBUG_ERROR("BCryptCreateHash failed with status %x", status); - goto end; - } - - /* - * This function hashes the buffer, but does NOT store it in our - * resulting buffer yet, we need to call BCryptFinishHash to retrieve - * the final hash. - */ - status = BCryptHashData(hash_handle, Buffer, BufferSize, NULL); - - if (!NT_SUCCESS(status)) { - DEBUG_ERROR("BCryptHashData failed with status %x", status); - goto end; - } - - /* - * As said in the previous comment, this is where we retrieve the final - * hash and store it in our output buffer. - */ - status = BCryptFinishHash( - hash_handle, resulting_hash, resulting_hash_size, NULL); - - if (!NT_SUCCESS(status)) { - DEBUG_ERROR("BCryptFinishHash failed with status %x", status); - goto end; - } - - *HashResult = resulting_hash; - *HashResultSize = resulting_hash_size; - -end: - - if (hash_handle) - BCryptDestroyHash(hash_handle); - - if (hash_object) - ImpExFreePoolWithTag(hash_object, POOL_TAG_INTEGRITY); - + ImpZwClose(handle); return status; } @@ -663,16 +552,16 @@ GetNextSMBIOSStructureInTable(_Inout_ PSMBIOS_TABLE_HEADER* CurrentStructure) STATIC NTSTATUS GetStringAtIndexFromSMBIOSTable(_In_ PSMBIOS_TABLE_HEADER Table, - _In_ INT Index, + _In_ UINT32 Index, _In_ PVOID Buffer, _In_ SIZE_T BufferSize) { PAGED_CODE(); - INT current_string_char_index = 0; - INT string_count = 0; - PCHAR current_string_char = (PCHAR)((UINT64)Table + Table->Length); - PCHAR next_string_char = current_string_char + 1; + UINT32 current_string_char_index = 0; + UINT32 string_count = 0; + PCHAR current_string_char = (PCHAR)((UINT64)Table + Table->Length); + PCHAR next_string_char = current_string_char + 1; for (;;) { if (*current_string_char == NULL_TERMINATOR && @@ -730,16 +619,16 @@ ParseSMBIOSTable(_Out_ PVOID Buffer, { PAGED_CODE(); - NTSTATUS status = STATUS_UNSUCCESSFUL; - PVOID firmware_table_buffer = NULL; - ULONG firmware_table_buffer_size = 0; - ULONG bytes_returned = 0; - PRAW_SMBIOS_DATA smbios_data = NULL; - PSMBIOS_TABLE_HEADER smbios_table_header = NULL; - PRAW_SMBIOS_TABLE_01 smbios_baseboard_information = NULL; + NTSTATUS status = STATUS_UNSUCCESSFUL; + PVOID buffer = NULL; + ULONG buffer_size = 0; + ULONG bytes_copied = 0; + PRAW_SMBIOS_DATA smbios_data = NULL; + PSMBIOS_TABLE_HEADER header = NULL; + PRAW_SMBIOS_TABLE_01 baseboard = NULL; - status = ImpExGetSystemFirmwareTable( - SMBIOS_TABLE, 0, NULL, 0, &firmware_table_buffer_size); + status = + ImpExGetSystemFirmwareTable(SMBIOS_TABLE, 0, NULL, 0, &buffer_size); /* * Because we pass a null buffer here, the NTSTATUS result will be a @@ -747,23 +636,24 @@ ParseSMBIOSTable(_Out_ PVOID Buffer, * the return bytes returned (which indicate required buffer size) is * above 0. */ - if (firmware_table_buffer_size == NULL) { + if (buffer_size == NULL) { DEBUG_ERROR( "ExGetSystemFirmwareTable call 1 failed to get required buffer size."); return STATUS_BUFFER_TOO_SMALL; } - firmware_table_buffer = ImpExAllocatePool2( - POOL_FLAG_NON_PAGED, firmware_table_buffer_size, POOL_TAG_INTEGRITY); + buffer = ImpExAllocatePool2(POOL_FLAG_NON_PAGED, + buffer_size, + POOL_TAG_INTEGRITY); - if (!firmware_table_buffer) + if (!buffer) return STATUS_MEMORY_NOT_ALLOCATED; status = ImpExGetSystemFirmwareTable(SMBIOS_TABLE, NULL, - firmware_table_buffer, - firmware_table_buffer_size, - &bytes_returned); + buffer, + buffer_size, + &bytes_copied); if (!NT_SUCCESS(status)) { DEBUG_ERROR("ExGetSystemFirmwareTable call 2 failed with status %x", @@ -771,8 +661,8 @@ ParseSMBIOSTable(_Out_ PVOID Buffer, goto end; } - smbios_data = GetRawSmbiosData(firmware_table_buffer); - smbios_table_header = GetSmbiosTableHeader(smbios_data); + smbios_data = GetRawSmbiosData(buffer); + header = GetSmbiosTableHeader(smbios_data); /* * The System Information table is equal to Type == 2 and contains the @@ -783,11 +673,13 @@ ParseSMBIOSTable(_Out_ PVOID Buffer, * https://www.dmtf.org/sites/default/files/standards/documents/DSP0134_2.7.1.pdf * line 823 */ - while (smbios_table_header->Type != TableIndex) - GetNextSMBIOSStructureInTable(&smbios_table_header); + while (header->Type != TableIndex) + GetNextSMBIOSStructureInTable(&header); - status = GetStringAtIndexFromSMBIOSTable( - smbios_table_header, TableSubIndex, Buffer, BufferSize); + status = GetStringAtIndexFromSMBIOSTable(header, + TableSubIndex, + Buffer, + BufferSize); if (!NT_SUCCESS(status)) { DEBUG_ERROR("GetStringAtIndexFromSMBIOSTable failed with status %x", @@ -797,8 +689,8 @@ ParseSMBIOSTable(_Out_ PVOID Buffer, end: - if (firmware_table_buffer) - ImpExFreePoolWithTag(firmware_table_buffer, POOL_TAG_INTEGRITY); + if (buffer) + ImpExFreePoolWithTag(buffer, POOL_TAG_INTEGRITY); return status; } @@ -819,25 +711,25 @@ ComputeHashOfSections(_In_ PIMAGE_SECTION_HEADER DiskSection, return STATUS_INVALID_BUFFER_SIZE; } - status = - ComputeHashOfBuffer((UINT64)DiskSection + sizeof(IMAGE_SECTION_HEADER), - DiskSection->SizeOfRawData, - DiskHash, - DiskHashSize); + status = CryptHashBuffer_sha256( + (UINT64)DiskSection + sizeof(IMAGE_SECTION_HEADER), + DiskSection->SizeOfRawData, + DiskHash, + DiskHashSize); if (!NT_SUCCESS(status)) { - DEBUG_ERROR("ComputeHashOfBuffer failed with status %x", status); + DEBUG_ERROR("CryptHashBuffer_sha256 failed with status %x", status); return status; } - status = ComputeHashOfBuffer((UINT64)MemorySection + - sizeof(IMAGE_SECTION_HEADER), - MemorySection->SizeOfRawData, - MemoryHash, - MemoryHashSize); + status = CryptHashBuffer_sha256( + (UINT64)MemorySection + sizeof(IMAGE_SECTION_HEADER), + MemorySection->SizeOfRawData, + MemoryHash, + MemoryHashSize); if (!NT_SUCCESS(status)) { - DEBUG_ERROR("ComputeHashOfBuffer 2 failed with status %x", status); + DEBUG_ERROR("CryptHashBuffer_sha256 2 failed with status %x", status); return status; } @@ -852,23 +744,18 @@ CompareHashes(_In_ PVOID Hash1, _In_ PVOID Hash2, _In_ UINT32 Length) return RtlCompareMemory(Hash1, Hash2, Length) == Length ? TRUE : FALSE; } -typedef struct _VAL_INTEGRITY_HEADER { - INTEGRITY_CHECK_HEADER integrity_check_header; - IMAGE_SECTION_HEADER section_header; - CHAR section_base[]; - -} VAL_INTEGRITY_HEADER, *PVAL_INTEGRITY_HEADER; - STATIC VOID ReportInvalidProcessModule(_In_ PPROCESS_MODULE_INFORMATION Module) { - NTSTATUS status = STATUS_UNSUCCESSFUL; - UINT32 report_size = CryptRequestRequiredBufferLength( + NTSTATUS status = STATUS_UNSUCCESSFUL; + UINT32 len = 0; + PPROCESS_MODULE_VALIDATION_REPORT report = NULL; + + len = CryptRequestRequiredBufferLength( sizeof(PROCESS_MODULE_VALIDATION_REPORT)); - PPROCESS_MODULE_VALIDATION_REPORT report = - ImpExAllocatePool2(POOL_FLAG_NON_PAGED, report_size, REPORT_POOL_TAG); + report = ImpExAllocatePool2(POOL_FLAG_NON_PAGED, len, REPORT_POOL_TAG); if (!report) return; @@ -877,18 +764,20 @@ ReportInvalidProcessModule(_In_ PPROCESS_MODULE_INFORMATION Module) report->image_base = Module->module_base; report->image_size = Module->module_size; - RtlCopyMemory( - report->module_path, Module->module_path, sizeof(report->module_path)); - status = CryptEncryptBuffer(report, report_size); + RtlCopyMemory(report->module_path, + Module->module_path, + sizeof(report->module_path)); + + status = CryptEncryptBuffer(report, len); if (!NT_SUCCESS(status)) { DEBUG_ERROR("CryptEncryptBuffer: %lx", status); - ImpExFreePoolWithTag(report, report_size); + ImpExFreePoolWithTag(report, len); return; } - IrpQueueSchedulePacket(report, report_size); + IrpQueueSchedulePacket(report, len); } /* @@ -966,8 +855,10 @@ ValidateProcessLoadedModule(_Inout_ PIRP Irp) goto end; } - status = MapDiskImageIntoVirtualAddressSpace( - §ion_handle, §ion, &module_path, §ion_size); + status = MapDiskImageIntoVirtualAddressSpace(§ion_handle, + §ion, + &module_path, + §ion_size); if (!NT_SUCCESS(status)) { DEBUG_ERROR("MapDiskImageIntoVirtualAddressSpace failed with status %x", @@ -975,8 +866,11 @@ ValidateProcessLoadedModule(_Inout_ PIRP Irp) goto end; } - status = StoreModuleExecutableRegionsInBuffer( - &disk_buffer, section, section_size, &bytes_written, FALSE); + status = StoreModuleExecutableRegionsInBuffer(&disk_buffer, + section, + section_size, + &bytes_written, + FALSE); if (!NT_SUCCESS(status)) { DEBUG_ERROR( @@ -1043,8 +937,11 @@ HashUserModule(_In_ PPROCESS_MAP_MODULE_ENTRY Entry, */ ImpKeStackAttachProcess(session->process, &apc_state); - status = StoreModuleExecutableRegionsInBuffer( - &memory_buffer, Entry->base, Entry->size, &bytes_written, FALSE); + status = StoreModuleExecutableRegionsInBuffer(&memory_buffer, + Entry->base, + Entry->size, + &bytes_written, + FALSE); ImpKeUnstackDetachProcess(&apc_state); @@ -1055,13 +952,13 @@ HashUserModule(_In_ PPROCESS_MAP_MODULE_ENTRY Entry, goto end; } - status = ComputeHashOfBuffer(memory_buffer->section_base, - memory_buffer->section_header.SizeOfRawData, - &memory_hash, - &memory_hash_size); + status = CryptHashBuffer_sha256(memory_buffer->section_base, + memory_buffer->section_header.SizeOfRawData, + &memory_hash, + &memory_hash_size); if (!NT_SUCCESS(status)) { - DEBUG_ERROR("ComputeHashOfBuffer failed with status %x", status); + DEBUG_ERROR("CryptHashBuffer_sha256 failed with status %x", status); goto end; } @@ -1120,32 +1017,35 @@ GetHardDiskDriveSerialNumber(_Inout_ PVOID ConfigDrive0Serial, { PAGED_CODE(); - NTSTATUS status = STATUS_UNSUCCESSFUL; - HANDLE handle = NULL; - OBJECT_ATTRIBUTES attributes = {0}; - IO_STATUS_BLOCK status_block = {0}; - STORAGE_PROPERTY_QUERY query = {0}; - STORAGE_DESCRIPTOR_HEADER storage_descriptor_header = {0}; - PSTORAGE_DEVICE_DESCRIPTOR device_descriptor = NULL; - UNICODE_STRING physical_drive_path = {0}; - PCHAR serial_number = NULL; - SIZE_T serial_length = 0; + NTSTATUS status = STATUS_UNSUCCESSFUL; + HANDLE handle = NULL; + OBJECT_ATTRIBUTES attributes = {0}; + IO_STATUS_BLOCK status_block = {0}; + STORAGE_PROPERTY_QUERY query = {0}; + STORAGE_DESCRIPTOR_HEADER header = {0}; + PSTORAGE_DEVICE_DESCRIPTOR descriptor = NULL; + UNICODE_STRING path = {0}; + PCHAR serial_number = NULL; + SIZE_T serial_length = 0; - ImpRtlInitUnicodeString(&physical_drive_path, - L"\\DosDevices\\PhysicalDrive0"); + ImpRtlInitUnicodeString(&path, L"\\DosDevices\\PhysicalDrive0"); /* * No need to use the flag OBJ_FORCE_ACCESS_CHECK since we arent passing * a handle given to us from usermode. */ InitializeObjectAttributes(&attributes, - &physical_drive_path, + &path, OBJ_CASE_INSENSITIVE | OBJ_KERNEL_HANDLE, NULL, NULL); - status = ImpZwOpenFile( - &handle, GENERIC_READ, &attributes, &status_block, NULL, NULL); + status = ImpZwOpenFile(&handle, + GENERIC_READ, + &attributes, + &status_block, + NULL, + NULL); if (!NT_SUCCESS(status)) { DEBUG_ERROR("ZwOpenFile on PhysicalDrive0 failed with status %x", @@ -1163,7 +1063,7 @@ GetHardDiskDriveSerialNumber(_Inout_ PVOID ConfigDrive0Serial, IOCTL_STORAGE_QUERY_PROPERTY, &query, sizeof(STORAGE_PROPERTY_QUERY), - &storage_descriptor_header, + &header, sizeof(STORAGE_DESCRIPTOR_HEADER)); if (!NT_SUCCESS(status)) { @@ -1172,11 +1072,11 @@ GetHardDiskDriveSerialNumber(_Inout_ PVOID ConfigDrive0Serial, goto end; } - device_descriptor = ImpExAllocatePool2(POOL_FLAG_NON_PAGED, - storage_descriptor_header.Size, - POOL_TAG_INTEGRITY); + descriptor = ImpExAllocatePool2(POOL_FLAG_NON_PAGED, + header.Size, + POOL_TAG_INTEGRITY); - if (!device_descriptor) { + if (!descriptor) { status = STATUS_MEMORY_NOT_ALLOCATED; goto end; } @@ -1189,8 +1089,8 @@ GetHardDiskDriveSerialNumber(_Inout_ PVOID ConfigDrive0Serial, IOCTL_STORAGE_QUERY_PROPERTY, &query, sizeof(STORAGE_PROPERTY_QUERY), - device_descriptor, - storage_descriptor_header.Size); + descriptor, + header.Size); if (!NT_SUCCESS(status)) { DEBUG_ERROR("ZwDeviceIoControlFile second call failed with status %x", @@ -1198,10 +1098,10 @@ GetHardDiskDriveSerialNumber(_Inout_ PVOID ConfigDrive0Serial, goto end; } - if (!device_descriptor->SerialNumberOffset) + if (!descriptor->SerialNumberOffset) goto end; - serial_number = GetStorageDescriptorSerialNumber(device_descriptor); + serial_number = GetStorageDescriptorSerialNumber(descriptor); serial_length = GetStorageDescriptorSerialLength(serial_number); if (serial_length > ConfigDrive0MaxSize) { @@ -1216,8 +1116,8 @@ end: if (handle) ImpZwClose(handle); - if (device_descriptor) - ImpExFreePoolWithTag(device_descriptor, POOL_TAG_INTEGRITY); + if (descriptor) + ImpExFreePoolWithTag(descriptor, POOL_TAG_INTEGRITY); return status; } @@ -1276,7 +1176,7 @@ MeasureReads(_In_ PVOID Address, _In_ ULONG Count) KeRaiseIrql(HIGH_LEVEL, &irql); _disable(); - for (ULONG iteration = 0; iteration < Count; iteration++) + for (UINT32 iteration = 0; iteration < Count; iteration++) read_average += MeasureInstructionRead(Address); _enable(); @@ -1375,9 +1275,9 @@ InitiateEptFunctionAddressArrays() { PAGED_CODE(); - UNICODE_STRING current_function; + UNICODE_STRING current_function = {0}; - for (INT index = 0; index < EPT_CONTROL_FUNCTIONS_COUNT; index++) { + for (UINT32 index = 0; index < EPT_CONTROL_FUNCTIONS_COUNT; index++) { ImpRtlInitUnicodeString(¤t_function, CONTROL_FUNCTIONS[index]); CONTROL_FUNCTION_ADDRESSES[index] = ImpMmGetSystemRoutineAddress(¤t_function); @@ -1386,7 +1286,7 @@ InitiateEptFunctionAddressArrays() return STATUS_UNSUCCESSFUL; } - for (INT index = 0; index < EPT_PROTECTED_FUNCTIONS_COUNT; index++) { + for (UINT32 index = 0; index < EPT_PROTECTED_FUNCTIONS_COUNT; index++) { ImpRtlInitUnicodeString(¤t_function, PROTECTED_FUNCTIONS[index]); PROTECTED_FUNCTION_ADDRESSES[index] = ImpMmGetSystemRoutineAddress(¤t_function); @@ -1419,7 +1319,7 @@ DetectEptHooksInKeyFunctions() return status; } - for (INT index = 0; index < EPT_CONTROL_FUNCTIONS_COUNT; index++) { + for (UINT32 index = 0; index < EPT_CONTROL_FUNCTIONS_COUNT; index++) { status = GetAverageReadTimeAtRoutine(CONTROL_FUNCTION_ADDRESSES[index], &instruction_time); @@ -1442,9 +1342,10 @@ DetectEptHooksInKeyFunctions() if (control_average == 0) return STATUS_UNSUCCESSFUL; - for (INT index = 0; index < EPT_PROTECTED_FUNCTIONS_COUNT; index++) { - status = GetAverageReadTimeAtRoutine( - PROTECTED_FUNCTION_ADDRESSES[index], &instruction_time); + for (UINT32 index = 0; index < EPT_PROTECTED_FUNCTIONS_COUNT; index++) { + status = + GetAverageReadTimeAtRoutine(PROTECTED_FUNCTION_ADDRESSES[index], + &instruction_time); if (!NT_SUCCESS(status)) { DEBUG_ERROR("DetectEptPresentOnFunction failed with status %x", @@ -1500,8 +1401,11 @@ StoreModuleExecutableRegionsx86(_In_ PRTL_MODULE_EXTENDED_INFO Module, ImpKeStackAttachProcess(process, &apc_state); - status = StoreModuleExecutableRegionsInBuffer( - Buffer, Module->ImageBase, Module->ImageSize, BufferSize, TRUE); + status = StoreModuleExecutableRegionsInBuffer(Buffer, + Module->ImageBase, + Module->ImageSize, + BufferSize, + TRUE); ImpKeUnstackDetachProcess(&apc_state); @@ -1528,46 +1432,45 @@ DeferredModuleHashingCallback(_In_ PDEVICE_OBJECT DeviceObject, UNREFERENCED_PARAMETER(Context); UNREFERENCED_PARAMETER(DeviceObject); - NTSTATUS status = STATUS_UNSUCCESSFUL; - RTL_MODULE_EXTENDED_INFO module = {0}; - PDRIVER_LIST_HEAD driver_list = GetDriverList(); - PLIST_ENTRY deferred_head = &GetDriverList()->deferred_list; - PLIST_ENTRY list_entry = NULL; - PDRIVER_LIST_ENTRY entry = NULL; + NTSTATUS status = STATUS_UNSUCCESSFUL; + RTL_MODULE_EXTENDED_INFO module = {0}; + PDRIVER_LIST_HEAD list = GetDriverList(); + PLIST_ENTRY head = &GetDriverList()->deferred_list; + PLIST_ENTRY entry = NULL; + PDRIVER_LIST_ENTRY driver = NULL; - Enablex86Hashing(driver_list); + Enablex86Hashing(list); - list_entry = RemoveHeadList(deferred_head); + entry = RemoveHeadList(head); - if (list_entry == deferred_head) + if (entry == head) goto end; - while (list_entry != deferred_head) { - entry = - CONTAINING_RECORD(list_entry, DRIVER_LIST_ENTRY, deferred_entry); + while (entry != head) { + driver = CONTAINING_RECORD(entry, DRIVER_LIST_ENTRY, deferred_entry); - DriverListEntryToExtendedModuleInfo(entry, &module); + DriverListEntryToExtendedModuleInfo(driver, &module); DEBUG_VERBOSE("Hashing Deferred Module: %s", module.FullPathName); - status = HashModule(&module, &entry->text_hash); + status = HashModule(&module, &driver->text_hash); if (!NT_SUCCESS(status)) { DEBUG_ERROR("HashModule-x86 failed with status %x", status); - entry->hashed = FALSE; - list_entry = RemoveHeadList(deferred_head); + driver->hashed = FALSE; + entry = RemoveHeadList(head); continue; } - entry->hashed = TRUE; - list_entry = RemoveHeadList(deferred_head); + driver->hashed = TRUE; + entry = RemoveHeadList(head); } end: DEBUG_VERBOSE("All deferred modules hashed."); - ImpIoFreeWorkItem(driver_list->work_item); - driver_list->work_item = NULL; + ImpIoFreeWorkItem(list->work_item); + list->work_item = NULL; } NTSTATUS @@ -1620,8 +1523,9 @@ HashModule(_In_ PRTL_MODULE_EXTENDED_INFO Module, _Out_ PVOID Hash) * Once the WinLogon process has started, we can then hash new * x86 modules. */ - status = StoreModuleExecutableRegionsx86( - Module, (PVOID)&memory_buffer, &memory_buffer_size); + status = StoreModuleExecutableRegionsx86(Module, + (PVOID)&memory_buffer, + &memory_buffer_size); } else { status = StoreModuleExecutableRegionsInBuffer((PVOID)&memory_buffer, @@ -1638,10 +1542,10 @@ HashModule(_In_ PRTL_MODULE_EXTENDED_INFO Module, _Out_ PVOID Hash) goto end; } - status = ComputeHashOfBuffer(memory_buffer->section_base, - memory_buffer->section_header.SizeOfRawData, - &memory_hash, - &memory_hash_size); + status = CryptHashBuffer_sha256(memory_buffer->section_base, + memory_buffer->section_header.SizeOfRawData, + &memory_hash, + &memory_hash_size); if (!NT_SUCCESS(status)) { DEBUG_VERBOSE("ComputeHashOfSections failed with status %x", status); @@ -1673,12 +1577,14 @@ STATIC VOID ReportModifiedSystemImage(_In_ PRTL_MODULE_EXTENDED_INFO Module) { - NTSTATUS status = STATUS_UNSUCCESSFUL; - UINT32 packet_size = CryptRequestRequiredBufferLength( + NTSTATUS status = STATUS_UNSUCCESSFUL; + UINT32 len = 0; + PSYSTEM_MODULE_INTEGRITY_CHECK_REPORT report = NULL; + + len = CryptRequestRequiredBufferLength( sizeof(SYSTEM_MODULE_INTEGRITY_CHECK_REPORT)); - PSYSTEM_MODULE_INTEGRITY_CHECK_REPORT report = - ImpExAllocatePool2(POOL_FLAG_NON_PAGED, packet_size, REPORT_POOL_TAG); + report = ImpExAllocatePool2(POOL_FLAG_NON_PAGED, len, REPORT_POOL_TAG); if (!report) return; @@ -1688,18 +1594,19 @@ ReportModifiedSystemImage(_In_ PRTL_MODULE_EXTENDED_INFO Module) report->image_base = Module->ImageBase; report->image_size = Module->ImageSize; - RtlCopyMemory( - report->path_name, Module->FullPathName, sizeof(report->path_name)); + RtlCopyMemory(report->path_name, + Module->FullPathName, + sizeof(report->path_name)); - status = CryptEncryptBuffer(report, packet_size); + status = CryptEncryptBuffer(report, len); if (!NT_SUCCESS(status)) { DEBUG_ERROR("CryptEncryptBuffer: %lx", status); - ImpExFreePoolWithTag(report, packet_size); + ImpExFreePoolWithTag(report, len); return; } - IrpQueueSchedulePacket(report, packet_size); + IrpQueueSchedulePacket(report, len); } VOID @@ -1709,8 +1616,9 @@ ValidateSystemModule(_In_ PRTL_MODULE_EXTENDED_INFO Module) PDRIVER_LIST_ENTRY entry = NULL; PVOID hash = NULL; - hash = ExAllocatePool2( - POOL_FLAG_NON_PAGED, SHA_256_HASH_LENGTH, POOL_TAG_INTEGRITY); + hash = ExAllocatePool2(POOL_FLAG_NON_PAGED, + SHA_256_HASH_LENGTH, + POOL_TAG_INTEGRITY); if (!hash) return; @@ -1764,33 +1672,36 @@ STATIC VOID ReportModifiedSelfDriverImage(_In_ PRTL_MODULE_EXTENDED_INFO Module) { - NTSTATUS status = STATUS_UNSUCCESSFUL; - UINT32 packet_size = CryptRequestRequiredBufferLength( + NTSTATUS status = STATUS_UNSUCCESSFUL; + UINT32 len = 0; + PDRIVER_SELF_INTEGRITY_CHECK_REPORT packet = NULL; + + len = CryptRequestRequiredBufferLength( sizeof(DRIVER_SELF_INTEGRITY_CHECK_REPORT)); - PDRIVER_SELF_INTEGRITY_CHECK_REPORT report = - ImpExAllocatePool2(POOL_FLAG_NON_PAGED, packet_size, REPORT_POOL_TAG); + packet = ImpExAllocatePool2(POOL_FLAG_NON_PAGED, len, REPORT_POOL_TAG); - if (!report) + if (!packet) return; - INIT_REPORT_PACKET(report, REPORT_SELF_DRIVER_PATCHED, 0); + INIT_REPORT_PACKET(packet, REPORT_SELF_DRIVER_PATCHED, 0); - report->image_base = Module->ImageBase; - report->image_size = Module->ImageSize; + packet->image_base = Module->ImageBase; + packet->image_size = Module->ImageSize; - RtlCopyMemory( - report->path_name, Module->FullPathName, sizeof(report->path_name)); + RtlCopyMemory(packet->path_name, + Module->FullPathName, + sizeof(packet->path_name)); - status = CryptEncryptBuffer(report, packet_size); + status = CryptEncryptBuffer(packet, len); if (!NT_SUCCESS(status)) { DEBUG_ERROR("CryptEncryptBuffer: %lx", status); - ImpExFreePoolWithTag(report, packet_size); + ImpExFreePoolWithTag(packet, len); return; } - IrpQueueSchedulePacket(report, packet_size); + IrpQueueSchedulePacket(packet, len); } NTSTATUS @@ -1819,8 +1730,9 @@ ValidateOurDriverImage() goto end; } - memory_hash = ExAllocatePool2( - POOL_FLAG_NON_PAGED, SHA_256_HASH_LENGTH, POOL_TAG_INTEGRITY); + memory_hash = ExAllocatePool2(POOL_FLAG_NON_PAGED, + SHA_256_HASH_LENGTH, + POOL_TAG_INTEGRITY); if (!memory_hash) goto end; @@ -1926,9 +1838,12 @@ SystemModuleVerificationDispatchFunction(_In_ PDEVICE_OBJECT DeviceObject, { UNREFERENCED_PARAMETER(DeviceObject); + UINT32 count = 0; + UINT32 max = 0; + IncrementActiveThreadCount(Context); - UINT32 count = GetCurrentVerificationIndex(Context); + count = GetCurrentVerificationIndex(Context); /* * theres a race condition here, where if the max is taken after a thread @@ -1937,11 +1852,13 @@ SystemModuleVerificationDispatchFunction(_In_ PDEVICE_OBJECT DeviceObject, * will be off by one. To fix just need to calculate the block max before * threads are dispatched. todo! */ - UINT32 max = GetCurrentVerificationMaxIndex(Context, count); + max = GetCurrentVerificationMaxIndex(Context, count); for (; count < max && count < Context->total_count; count++) { if (!InterlockedCompareExchange( - &Context->dispatcher_info[count].validated, TRUE, FALSE)) { + &Context->dispatcher_info[count].validated, + TRUE, + FALSE)) { ValidateSystemModule(&Context->module_info[count]); } } @@ -1992,6 +1909,7 @@ InitialiseSystemModuleVerificationContext(PSYS_MODULE_VAL_CONTEXT Context) NTSTATUS status = STATUS_UNSUCCESSFUL; SYSTEM_MODULES modules = {0}; PMODULE_DISPATCHER_HEADER dispatcher = NULL; + UINT32 count = 0; status = GetSystemModuleInformation(&modules); @@ -2002,10 +1920,10 @@ InitialiseSystemModuleVerificationContext(PSYS_MODULE_VAL_CONTEXT Context) DEBUG_VERBOSE("driver count: %lx", modules.module_count); - dispatcher = ImpExAllocatePool2(POOL_FLAG_NON_PAGED, - modules.module_count * - sizeof(MODULE_DISPATCHER_HEADER), - POOL_TAG_INTEGRITY); + count = modules.module_count * sizeof(MODULE_DISPATCHER_HEADER); + + dispatcher = + ImpExAllocatePool2(POOL_FLAG_NON_PAGED, count, POOL_TAG_INTEGRITY); if (!dispatcher) { ImpExFreePoolWithTag(modules.address, SYSTEM_MODULES_POOL); @@ -2060,7 +1978,7 @@ STATIC VOID DispatchVerificationWorkerThreads(_In_ PSYS_MODULE_VAL_CONTEXT Context) { - for (INT index = 0; index < VERIFICATION_THREAD_COUNT; index++) { + for (UINT32 index = 0; index < VERIFICATION_THREAD_COUNT; index++) { Context->work_items[index] = ImpIoAllocateWorkItem(GetDriverDeviceObject()); @@ -2084,8 +2002,10 @@ SystemModuleVerificationDispatcher() if (context->complete) { DEBUG_VERBOSE( "System modules integrity check complete. Freeing items."); + context->active = FALSE; context->complete = FALSE; + FreeModuleVerificationItems(context); FreeWorkItems(context); return STATUS_SUCCESS; @@ -2212,14 +2132,16 @@ STATIC NTSTATUS AllocateHeartbeatObjects(_Inout_ PHEARTBEAT_CONFIGURATION Configuration) { - Configuration->dpc = ImpExAllocatePool2( - POOL_FLAG_NON_PAGED, sizeof(KDPC), POOL_TAG_HEARTBEAT); + Configuration->dpc = ImpExAllocatePool2(POOL_FLAG_NON_PAGED, + sizeof(KDPC), + POOL_TAG_HEARTBEAT); if (!Configuration->dpc) return STATUS_INSUFFICIENT_RESOURCES; - Configuration->timer = ImpExAllocatePool2( - POOL_FLAG_NON_PAGED, sizeof(KTIMER), POOL_TAG_HEARTBEAT); + Configuration->timer = ImpExAllocatePool2(POOL_FLAG_NON_PAGED, + sizeof(KTIMER), + POOL_TAG_HEARTBEAT); if (!Configuration->timer) { ImpExFreePoolWithTag(Configuration->dpc, POOL_TAG_HEARTBEAT); @@ -2242,27 +2164,28 @@ STATIC LARGE_INTEGER GenerateHeartbeatDueTime() { - LARGE_INTEGER ticks = {0}; + UINT64 interval = 0; + LARGE_INTEGER ticks = {0}; + LARGE_INTEGER due_time = {0}; + KeQueryTickCount(&ticks); - UINT64 interval = - HEARTBEAT_NANOSECONDS_LOW + - (TICKS_TO_100_NS_INTERVALS(ticks.QuadPart) % - (HEARTBEAT_NANOSECONDS_HIGH - HEARTBEAT_NANOSECONDS_LOW)); + interval = HEARTBEAT_NANOSECONDS_LOW + + (TICKS_TO_100_NS_INTERVALS(ticks.QuadPart) % + (HEARTBEAT_NANOSECONDS_HIGH - HEARTBEAT_NANOSECONDS_LOW)); - LARGE_INTEGER due_time = {.QuadPart = -interval}; + due_time.QuadPart = -interval; return due_time; } FORCEINLINE STATIC VOID -InitialiseHeartbeatObjects(_Inout_ PHEARTBEAT_CONFIGURATION Configuration) +InitialiseHeartbeatObjects(_Inout_ PHEARTBEAT_CONFIGURATION Config) { - KeInitializeDpc(Configuration->dpc, HeartbeatDpcRoutine, Configuration); - KeInitializeTimer(Configuration->timer); - KeSetTimer( - Configuration->timer, GenerateHeartbeatDueTime(), Configuration->dpc); + KeInitializeDpc(Config->dpc, HeartbeatDpcRoutine, Config); + KeInitializeTimer(Config->timer); + KeSetTimer(Config->timer, GenerateHeartbeatDueTime(), Config->dpc); } FORCEINLINE @@ -2302,12 +2225,12 @@ IncrementHeartbeatCounter(_In_ PHEARTBEAT_CONFIGURATION Configuration) FORCEINLINE STATIC PHEARTBEAT_PACKET -BuildHeartbeatPacket(_In_ UINT32 PacketSize) +BuildHeartbeatPacket(_In_ UINT32 Size) { - PIRP_QUEUE_HEAD queue = GetIrpQueueHead(); + PIRP_QUEUE_HEAD queue = GetIrpQueueHead(); + PHEARTBEAT_PACKET packet = NULL; - PHEARTBEAT_PACKET packet = - ImpExAllocatePool2(POOL_FLAG_NON_PAGED, PacketSize, POOL_TAG_HEARTBEAT); + packet = ImpExAllocatePool2(POOL_FLAG_NON_PAGED, Size, POOL_TAG_HEARTBEAT); if (!packet) return NULL; @@ -2324,6 +2247,7 @@ BuildHeartbeatPacket(_In_ UINT32 PacketSize) packet->total_heartbeats_completed = queue->total_heartbeats_completed; packet->total_irps_completed = queue->total_irps_completed; packet->total_reports_completed = queue->total_reports_completed; + KeReleaseGuardedMutex(&queue->lock); return packet; @@ -2395,8 +2319,10 @@ HeartbeatDpcRoutine(_In_ PKDPC Dpc, PHEARTBEAT_CONFIGURATION config = (PHEARTBEAT_CONFIGURATION)DeferredContext; - IoQueueWorkItem( - config->work_item, HeartbeatWorkItem, NormalWorkQueue, config); + IoQueueWorkItem(config->work_item, + HeartbeatWorkItem, + NormalWorkQueue, + config); } /* diff --git a/driver/modules.c b/driver/modules.c index ff4b50b..0b08568 100644 --- a/driver/modules.c +++ b/driver/modules.c @@ -1,15 +1,15 @@ #include "modules.h" +#include "apc.h" #include "callbacks.h" +#include "containers/tree.h" +#include "crypt.h" #include "driver.h" -#include "io.h" #include "ia32.h" #include "imports.h" -#include "apc.h" -#include "thread.h" +#include "io.h" #include "pe.h" -#include "crypt.h" -#include "containers/tree.h" +#include "thread.h" #define WHITELISTED_MODULE_TAG 'whte' @@ -55,18 +55,6 @@ typedef struct _WHITELISTED_REGIONS { } WHITELISTED_REGIONS, *PWHITELISTED_REGIONS; -typedef struct _NMI_POOLS { - PVOID thread_data_pool; - PVOID stack_frames; - PVOID nmi_context; - -} NMI_POOLS, *PNMI_POOLS; - -typedef struct _MODULE_VALIDATION_FAILURE_HEADER { - INT module_count; - -} MODULE_VALIDATION_FAILURE_HEADER, *PMODULE_VALIDATION_FAILURE_HEADER; - typedef struct _NMI_CONTEXT { UINT64 interrupted_rip; UINT64 interrupted_rsp; @@ -76,19 +64,39 @@ typedef struct _NMI_CONTEXT { } NMI_CONTEXT, *PNMI_CONTEXT; +#define DPC_STACKWALK_STACKFRAME_COUNT 10 + +/* the first 3 frames are isr handlers which we dont care about */ +#define DPC_STACKWALK_FRAMES_TO_SKIP 3 + +typedef struct _DPC_CONTEXT { + UINT64 stack_frame[DPC_STACKWALK_STACKFRAME_COUNT]; + UINT16 frames_captured; + volatile BOOLEAN executed; + +} DPC_CONTEXT, *PDPC_CONTEXT; + +// clang-format off + STATIC VOID -PopulateWhitelistedModuleBuffer(_Inout_ PWHITELISTED_REGIONS Whitelist, - _In_ PSYSTEM_MODULES SystemModules); +PopulateWhitelistedModuleBuffer( + _Inout_ PWHITELISTED_REGIONS Whitelist, + _In_ PSYSTEM_MODULES SystemModules +); STATIC NTSTATUS -ValidateDriverObjectsWrapper(_In_ PSYSTEM_MODULES SystemModules); +ValidateDriverObjectsWrapper( + _In_ PSYSTEM_MODULES SystemModules +); STATIC NTSTATUS -AnalyseNmiData(_In_ PNMI_CONTEXT NmiContext, - _In_ PSYSTEM_MODULES SystemModules); +AnalyseNmiData( + _In_ PNMI_CONTEXT NmiContext, + _In_ PSYSTEM_MODULES SystemModules +); STATIC NTSTATUS @@ -96,26 +104,35 @@ LaunchNonMaskableInterrupt(); STATIC VOID -ApcRundownRoutine(_In_ PRKAPC Apc); +ApcRundownRoutine( + _In_ PRKAPC Apc +); STATIC VOID -ApcKernelRoutine(_In_ PRKAPC Apc, - _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); +ApcKernelRoutine( + _In_ PRKAPC Apc, + _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); STATIC VOID -ApcNormalRoutine(_In_opt_ PVOID NormalContext, - _In_opt_ PVOID SystemArgument1, - _In_opt_ PVOID SystemArgument2); +ApcNormalRoutine( + _In_opt_ PVOID NormalContext, + _In_opt_ PVOID SystemArgument1, + _In_opt_ PVOID SystemArgument2 +); STATIC VOID -ValidateThreadViaKernelApcCallback(_In_ PTHREAD_LIST_ENTRY ThreadListEntry, - _Inout_opt_ PVOID Context); +ValidateThreadViaKernelApcCallback( + _In_ PTHREAD_LIST_ENTRY ThreadListEntry, + _Inout_opt_ PVOID Context +); + +// clang-format on #ifdef ALLOC_PRAGMA # pragma alloc_text(PAGE, FindSystemModuleByName) @@ -168,19 +185,21 @@ PopulateWhitelistedModuleBuffer(_Inout_ PWHITELISTED_REGIONS Whitelist, { PAGED_CODE(); - for (INT index = 0; index < WHITELISTED_MODULE_COUNT; index++) { - LPCSTR entry = WHITELISTED_MODULES[index]; + LPCSTR entry = NULL; + PRTL_MODULE_EXTENDED_INFO module = NULL; + PWHITELISTED_REGIONS region = NULL; - PRTL_MODULE_EXTENDED_INFO module = - FindSystemModuleByName(entry, SystemModules); + for (UINT32 index = 0; index < WHITELISTED_MODULE_COUNT; index++) { + entry = WHITELISTED_MODULES[index]; + module = FindSystemModuleByName(entry, SystemModules); /* not everyone will contain all whitelisted modules */ if (!module) continue; - PWHITELISTED_REGIONS region = &Whitelist[index]; - region->base = (UINT64)module->ImageBase; - region->end = (UINT64)module->ImageBase + module->ImageSize; + region = &Whitelist[index]; + region->base = (UINT64)module->ImageBase; + region->end = (UINT64)module->ImageBase + module->ImageSize; } } @@ -199,25 +218,25 @@ DoesDriverHaveInvalidDispatchRoutine(_In_ PDRIVER_OBJECT Driver, { PAGED_CODE(); - UINT64 dispatch_function = 0; - UINT64 module_base = 0; - UINT64 module_end = 0; + UINT64 dispatch_function = 0; + UINT64 module_base = 0; + UINT64 module_end = 0; + PRTL_MODULE_EXTENDED_INFO module = NULL; dispatch_function = GetDriverMajorDispatchFunction(Driver); - if (dispatch_function == NULL) + if (!dispatch_function) return FALSE; - PRTL_MODULE_EXTENDED_INFO module = - (PRTL_MODULE_EXTENDED_INFO)Modules->address; + module = (PRTL_MODULE_EXTENDED_INFO)Modules->address; - for (INT index = 0; index < Modules->module_count; index++) { + for (UINT32 index = 0; index < Modules->module_count; index++) { if (module[index].ImageBase != Driver->DriverStart) continue; /* make sure our driver has a device object which is required * for IOCTL */ - if (Driver->DeviceObject == NULL) + if (!Driver->DeviceObject) return FALSE; module_base = (UINT64)module[index].ImageBase; @@ -265,14 +284,18 @@ DoesDriverObjectHaveBackingModule(_In_ PSYSTEM_MODULES ModuleInformation, { PAGED_CODE(); - PRTL_MODULE_EXTENDED_INFO module = - (PRTL_MODULE_EXTENDED_INFO)ModuleInformation->address; + PRTL_MODULE_EXTENDED_INFO modules = NULL; + PRTL_MODULE_EXTENDED_INFO entry = NULL; - for (INT index = 0; index < ModuleInformation->module_count; index++) { - if (module[index].ImageSize == 0 || module[index].ImageBase == 0) + modules = (PRTL_MODULE_EXTENDED_INFO)ModuleInformation->address; + + for (UINT32 index = 0; index < ModuleInformation->module_count; index++) { + entry = &modules[index]; + + if (entry->ImageSize == 0 || entry->ImageBase == 0) return STATUS_INVALID_MEMBER; - if (module[index].ImageBase == DriverObject->DriverStart) { + if (entry->ImageBase == DriverObject->DriverStart) { return TRUE; } } @@ -300,30 +323,32 @@ GetSystemModuleInformation(_Out_ PSYSTEM_MODULES ModuleInformation) { PAGED_CODE(); + ULONG size = 0; + NTSTATUS status = STATUS_UNSUCCESSFUL; + PRTL_MODULE_EXTENDED_INFO buffer = NULL; + if (!ModuleInformation) return STATUS_INVALID_PARAMETER; - ULONG size = 0; - NTSTATUS status = STATUS_UNSUCCESSFUL; - - status = RtlQueryModuleInformation( - &size, sizeof(RTL_MODULE_EXTENDED_INFO), NULL); + status = RtlQueryModuleInformation(&size, + sizeof(RTL_MODULE_EXTENDED_INFO), + NULL); if (!NT_SUCCESS(status)) { DEBUG_ERROR("RtlQueryModuleInformation failed with status %x", status); return status; } - PRTL_MODULE_EXTENDED_INFO buffer = - ExAllocatePool2(POOL_FLAG_NON_PAGED, size, SYSTEM_MODULES_POOL); + buffer = ExAllocatePool2(POOL_FLAG_NON_PAGED, size, SYSTEM_MODULES_POOL); if (!buffer) { DEBUG_ERROR("Failed to allocate pool LOL"); return STATUS_MEMORY_NOT_ALLOCATED; } - status = RtlQueryModuleInformation( - &size, sizeof(RTL_MODULE_EXTENDED_INFO), buffer); + status = RtlQueryModuleInformation(&size, + sizeof(RTL_MODULE_EXTENDED_INFO), + buffer); if (!NT_SUCCESS(status)) { DEBUG_ERROR("RtlQueryModuleInformation 2 failed with status %x", @@ -332,8 +357,9 @@ GetSystemModuleInformation(_Out_ PSYSTEM_MODULES ModuleInformation) return STATUS_ABANDONED; } - InitSystemModulesStructure( - ModuleInformation, buffer, size / sizeof(RTL_MODULE_EXTENDED_INFO)); + InitSystemModulesStructure(ModuleInformation, + buffer, + ARRAYLEN(size, RTL_MODULE_EXTENDED_INFO)); return status; } @@ -342,12 +368,13 @@ STATIC VOID ReportInvalidDriverObject(_In_ PDRIVER_OBJECT Driver, _In_ UINT32 ReportSubType) { - NTSTATUS status = STATUS_UNSUCCESSFUL; - UINT32 packet_size = - CryptRequestRequiredBufferLength(sizeof(MODULE_VALIDATION_FAILURE)); + UINT32 len = 0; + NTSTATUS status = STATUS_UNSUCCESSFUL; + ANSI_STRING string = {0}; + PMODULE_VALIDATION_FAILURE report = NULL; - PMODULE_VALIDATION_FAILURE report = ImpExAllocatePool2( - POOL_FLAG_NON_PAGED, packet_size, POOL_TAG_INTEGRITY); + len = CryptRequestRequiredBufferLength(sizeof(MODULE_VALIDATION_FAILURE)); + report = ImpExAllocatePool2(POOL_FLAG_NON_PAGED, len, POOL_TAG_INTEGRITY); if (!report) return; @@ -357,7 +384,6 @@ ReportInvalidDriverObject(_In_ PDRIVER_OBJECT Driver, _In_ UINT32 ReportSubType) report->driver_base_address = Driver->DriverStart; report->driver_size = Driver->DriverSize; - ANSI_STRING string = {0}; string.Length = 0; string.MaximumLength = MODULE_REPORT_DRIVER_NAME_BUFFER_SIZE; string.Buffer = &report->driver_name; @@ -365,7 +391,7 @@ ReportInvalidDriverObject(_In_ PDRIVER_OBJECT Driver, _In_ UINT32 ReportSubType) /* Continue regardless of result */ ImpRtlUnicodeStringToAnsiString(&string, &Driver->DriverName, FALSE); - status = CryptEncryptBuffer(report, packet_size); + status = CryptEncryptBuffer(report, len); if (!NT_SUCCESS(status)) { DEBUG_ERROR("CryptEncryptBuffer: %lx", status); @@ -373,7 +399,7 @@ ReportInvalidDriverObject(_In_ PDRIVER_OBJECT Driver, _In_ UINT32 ReportSubType) return; } - IrpQueueSchedulePacket(report, packet_size); + IrpQueueSchedulePacket(report, len); } FORCEINLINE @@ -384,6 +410,14 @@ GetNextObject(_In_ POBJECT_DIRECTORY_ENTRY Entry) return Entry->ChainLink; } +FORCEINLINE +STATIC +PVOID +GetObjectFromDirectory(_In_ POBJECT_DIRECTORY_ENTRY Entry) +{ + return Entry->Object; +} + STATIC VOID ValidateDriverObjects(_In_ PSYSTEM_MODULES Modules, @@ -392,9 +426,10 @@ ValidateDriverObjects(_In_ PSYSTEM_MODULES Modules, { NTSTATUS status = STATUS_UNSUCCESSFUL; POBJECT_DIRECTORY_ENTRY entry = Entry; + PDRIVER_OBJECT driver = NULL; while (entry) { - PDRIVER_OBJECT driver = entry->Object; + driver = GetObjectFromDirectory(entry); if (!DoesDriverObjectHaveBackingModule(Modules, driver)) { ReportInvalidDriverObject(driver, REPORT_SUBTYPE_NO_BACKING_MODULE); @@ -418,29 +453,36 @@ ValidateDriverObjectsWrapper(_In_ PSYSTEM_MODULES SystemModules) { PAGED_CODE(); - HANDLE handle = NULL; - OBJECT_ATTRIBUTES attributes = {0}; - PVOID directory = {0}; - UNICODE_STRING directory_name = {0}; - PWHITELISTED_REGIONS whitelist = NULL; - NTSTATUS status = STATUS_UNSUCCESSFUL; - POBJECT_DIRECTORY directory_object = NULL; + HANDLE handle = NULL; + OBJECT_ATTRIBUTES oa = {0}; + PVOID dir = {0}; + UNICODE_STRING dir_name = {0}; + PWHITELISTED_REGIONS wl = NULL; + NTSTATUS status = STATUS_UNSUCCESSFUL; + POBJECT_DIRECTORY dir_object = NULL; + POBJECT_DIRECTORY_ENTRY bucket = NULL; - ImpRtlInitUnicodeString(&directory_name, L"\\Driver"); + ImpRtlInitUnicodeString(&dir_name, L"\\Driver"); - InitializeObjectAttributes( - &attributes, &directory_name, OBJ_CASE_INSENSITIVE, NULL, NULL); + InitializeObjectAttributes(&oa, + &dir_name, + OBJ_CASE_INSENSITIVE, + NULL, + NULL); - status = - ImpZwOpenDirectoryObject(&handle, DIRECTORY_ALL_ACCESS, &attributes); + status = ImpZwOpenDirectoryObject(&handle, DIRECTORY_ALL_ACCESS, &oa); if (!NT_SUCCESS(status)) { DEBUG_ERROR("ZwOpenDirectoryObject failed with status %x", status); return status; } - status = ImpObReferenceObjectByHandle( - handle, DIRECTORY_ALL_ACCESS, NULL, KernelMode, &directory, NULL); + status = ImpObReferenceObjectByHandle(handle, + DIRECTORY_ALL_ACCESS, + NULL, + KernelMode, + &dir, + NULL); if (!NT_SUCCESS(status)) { DEBUG_ERROR("ObReferenceObjectByHandle failed with status %x", status); @@ -461,19 +503,19 @@ ValidateDriverObjectsWrapper(_In_ PSYSTEM_MODULES SystemModules) * accessed the most can be accessed quickly */ - directory_object = (POBJECT_DIRECTORY)directory; + dir_object = (POBJECT_DIRECTORY)dir; - ImpExAcquirePushLockExclusiveEx(&directory_object->Lock, NULL); + ImpExAcquirePushLockExclusiveEx(&dir_object->Lock, NULL); - whitelist = ImpExAllocatePool2(POOL_FLAG_NON_PAGED, - WHITELISTED_MODULE_COUNT * - sizeof(WHITELISTED_REGIONS), - WHITELISTED_MODULE_TAG); + wl = ImpExAllocatePool2( + POOL_FLAG_NON_PAGED, + WHITELISTED_MODULE_COUNT * sizeof(WHITELISTED_REGIONS), + WHITELISTED_MODULE_TAG); - if (!whitelist) + if (!wl) goto end; - PopulateWhitelistedModuleBuffer(whitelist, SystemModules); + PopulateWhitelistedModuleBuffer(wl, SystemModules); if (!NT_SUCCESS(status)) { DEBUG_ERROR("PopulateWhitelistedModuleBuffer failed with status %x", @@ -481,17 +523,17 @@ ValidateDriverObjectsWrapper(_In_ PSYSTEM_MODULES SystemModules) goto end; } - for (INT index = 0; index < NUMBER_HASH_BUCKETS; index++) { - POBJECT_DIRECTORY_ENTRY entry = directory_object->HashBuckets[index]; - ValidateDriverObjects(SystemModules, entry, whitelist); + for (UINT32 index = 0; index < NUMBER_HASH_BUCKETS; index++) { + bucket = dir_object->HashBuckets[index]; + ValidateDriverObjects(SystemModules, bucket, wl); } end: - if (whitelist) - ImpExFreePoolWithTag(whitelist, WHITELISTED_MODULE_TAG); + if (wl) + ImpExFreePoolWithTag(wl, WHITELISTED_MODULE_TAG); - ImpExReleasePushLockExclusiveEx(&directory_object->Lock, 0); - ImpObDereferenceObject(directory); + ImpExReleasePushLockExclusiveEx(&dir_object->Lock, 0); + ImpObDereferenceObject(dir); ImpZwClose(handle); return STATUS_SUCCESS; @@ -510,21 +552,18 @@ HandleValidateDriversIOCTL() { PAGED_CODE(); - NTSTATUS status = STATUS_UNSUCCESSFUL; - ULONG buffer_size = 0; - SYSTEM_MODULES system_modules = {0}; + NTSTATUS status = STATUS_UNSUCCESSFUL; + ULONG length = 0; + SYSTEM_MODULES modules = {0}; - /* Fix annoying visual studio linting error */ - RtlZeroMemory(&system_modules, sizeof(SYSTEM_MODULES)); - - status = GetSystemModuleInformation(&system_modules); + status = GetSystemModuleInformation(&modules); if (!NT_SUCCESS(status)) { DEBUG_ERROR("GetSystemModuleInformation failed with status %x", status); return status; } - status = ValidateDriverObjectsWrapper(&system_modules); + status = ValidateDriverObjectsWrapper(&modules); if (!NT_SUCCESS(status)) { DEBUG_ERROR("ValidateDriverObjects failed with status %x", status); @@ -533,7 +572,9 @@ HandleValidateDriversIOCTL() end: - ImpExFreePoolWithTag(system_modules.address, SYSTEM_MODULES_POOL); + if (modules.address) + ImpExFreePoolWithTag(modules.address, SYSTEM_MODULES_POOL); + return status; } @@ -542,7 +583,7 @@ end: * boolean and remove the out variable. */ BOOLEAN -IsInstructionPointerInInvalidRegion(_In_ UINT64 RIP, +IsInstructionPointerInInvalidRegion(_In_ UINT64 Rip, _In_ PSYSTEM_MODULES SystemModules) { PAGED_CODE(); @@ -551,11 +592,11 @@ IsInstructionPointerInInvalidRegion(_In_ UINT64 RIP, (PRTL_MODULE_EXTENDED_INFO)SystemModules->address; /* Note that this does not check for HAL or PatchGuard Execution */ - for (INT index = 0; index < SystemModules->module_count; index++) { + for (UINT32 index = 0; index < SystemModules->module_count; index++) { UINT64 base = (UINT64)modules[index].ImageBase; UINT64 end = base + modules[index].ImageSize; - if (RIP >= base && RIP <= end) { + if (Rip >= base && Rip <= end) { return FALSE; } } @@ -580,12 +621,12 @@ STATIC VOID ReportNmiBlocking() { - NTSTATUS status = STATUS_UNSUCCESSFUL; - UINT32 packet_size = - CryptRequestRequiredBufferLength(sizeof(NMI_CALLBACK_FAILURE)); + NTSTATUS status = STATUS_UNSUCCESSFUL; + UINT32 len = 0; + PNMI_CALLBACK_FAILURE report = NULL; - PNMI_CALLBACK_FAILURE report = - ImpExAllocatePool2(POOL_FLAG_NON_PAGED, packet_size, REPORT_POOL_TAG); + len = CryptRequestRequiredBufferLength(sizeof(NMI_CALLBACK_FAILURE)); + report = ImpExAllocatePool2(POOL_FLAG_NON_PAGED, len, REPORT_POOL_TAG); if (!report) return STATUS_INSUFFICIENT_RESOURCES; @@ -596,7 +637,7 @@ ReportNmiBlocking() report->invalid_rip = NULL; report->were_nmis_disabled = TRUE; - status = CryptEncryptBuffer(report, packet_size); + status = CryptEncryptBuffer(report, len); if (!NT_SUCCESS(status)) { DEBUG_ERROR("CryptEncryptBuffer: %lx", status); @@ -604,22 +645,19 @@ ReportNmiBlocking() return; } - IrpQueueSchedulePacket(report, packet_size); + IrpQueueSchedulePacket(report, len); } STATIC VOID ReportMissingCidTableEntry(_In_ PNMI_CONTEXT Context) { - DEBUG_WARNING("Thread: %llx was not found in the pspcid table.", - Context->kthread); + NTSTATUS status = STATUS_UNSUCCESSFUL; + UINT32 len = 0; + PHIDDEN_SYSTEM_THREAD_REPORT report = NULL; - NTSTATUS status = STATUS_UNSUCCESSFUL; - UINT32 packet_size = - CryptRequestRequiredBufferLength(sizeof(HIDDEN_SYSTEM_THREAD_REPORT)); - - PHIDDEN_SYSTEM_THREAD_REPORT report = - ImpExAllocatePool2(POOL_FLAG_NON_PAGED, packet_size, REPORT_POOL_TAG); + len = CryptRequestRequiredBufferLength(sizeof(HIDDEN_SYSTEM_THREAD_REPORT)); + report = ImpExAllocatePool2(POOL_FLAG_NON_PAGED, len, REPORT_POOL_TAG); if (!report) return; @@ -633,7 +671,7 @@ ReportMissingCidTableEntry(_In_ PNMI_CONTEXT Context) RtlCopyMemory(report->thread, Context->kthread, sizeof(report->thread)); - status = CryptEncryptBuffer(report, packet_size); + status = CryptEncryptBuffer(report, len); if (!NT_SUCCESS(status)) { DEBUG_ERROR("CryptEncryptBuffer: %lx", status); @@ -641,7 +679,7 @@ ReportMissingCidTableEntry(_In_ PNMI_CONTEXT Context) return; } - IrpQueueSchedulePacket(report, packet_size); + IrpQueueSchedulePacket(report, len); } STATIC @@ -649,12 +687,12 @@ VOID ReportInvalidRipFoundDuringNmi(_In_ PNMI_CONTEXT Context, _In_ UINT32 ReportSubCode) { - NTSTATUS status = STATUS_UNSUCCESSFUL; - UINT32 packet_size = - CryptRequestRequiredBufferLength(sizeof(HIDDEN_SYSTEM_THREAD_REPORT)); + NTSTATUS status = STATUS_UNSUCCESSFUL; + UINT32 len = 0; + PNMI_CALLBACK_FAILURE report = NULL; - PNMI_CALLBACK_FAILURE report = - ImpExAllocatePool2(POOL_FLAG_NON_PAGED, packet_size, REPORT_POOL_TAG); + len = CryptRequestRequiredBufferLength(sizeof(HIDDEN_SYSTEM_THREAD_REPORT)); + report = ImpExAllocatePool2(POOL_FLAG_NON_PAGED, len, REPORT_POOL_TAG); if (!report) return; @@ -665,7 +703,7 @@ ReportInvalidRipFoundDuringNmi(_In_ PNMI_CONTEXT Context, report->invalid_rip = Context->interrupted_rip; report->were_nmis_disabled = FALSE; - status = CryptEncryptBuffer(report, packet_size); + status = CryptEncryptBuffer(report, len); if (!NT_SUCCESS(status)) { DEBUG_ERROR("CryptEncryptBuffer: %lx", status); @@ -673,7 +711,7 @@ ReportInvalidRipFoundDuringNmi(_In_ PNMI_CONTEXT Context, return; } - IrpQueueSchedulePacket(report, packet_size); + IrpQueueSchedulePacket(report, len); } #define INSTRUCTION_UD2_BYTE_1 0x0F @@ -709,7 +747,8 @@ DoesRetInstructionCauseException(_In_ UINT64 ReturnAddress) if (opcodes[0] == INSTRUCTION_INT3_BYTE_1) return TRUE; - DEBUG_VERBOSE("Ret address instruction doesnt unconditionally throw exception"); + DEBUG_VERBOSE( + "Ret address instruction doesnt unconditionally throw exception"); return FALSE; } @@ -719,19 +758,22 @@ DoesRetInstructionCauseException(_In_ UINT64 ReturnAddress) */ STATIC NTSTATUS -AnalyseNmiData(_In_ PNMI_CONTEXT NmiContext, _In_ PSYSTEM_MODULES SystemModules) +AnalyseNmiData(_In_ PNMI_CONTEXT NmiContext, _In_ PSYSTEM_MODULES Modules) { PAGED_CODE(); - NTSTATUS status = STATUS_UNSUCCESSFUL; - BOOLEAN flag = FALSE; + NTSTATUS status = STATUS_UNSUCCESSFUL; + BOOLEAN flag = FALSE; + PNMI_CONTEXT context = NULL; - if (!NmiContext || !SystemModules) + if (!NmiContext || !Modules) return STATUS_INVALID_PARAMETER; - for (INT core = 0; core < ImpKeQueryActiveProcessorCount(0); core++) { + for (UINT32 core = 0; core < ImpKeQueryActiveProcessorCount(0); core++) { + context = &NmiContext[core]; + /* Make sure our NMIs were run */ - if (!NmiContext[core].callback_count) { + if (!context->callback_count) { ReportNmiBlocking(); return STATUS_SUCCESS; } @@ -739,7 +781,7 @@ AnalyseNmiData(_In_ PNMI_CONTEXT NmiContext, _In_ PSYSTEM_MODULES SystemModules) DEBUG_VERBOSE( "Analysing Nmi Data for: cpu number: %i callback count: %lx", core, - NmiContext[core].callback_count); + context->callback_count); /* * Our NMI callback allows us to interrupt every running thread @@ -764,19 +806,20 @@ AnalyseNmiData(_In_ PNMI_CONTEXT NmiContext, _In_ PSYSTEM_MODULES SystemModules) * PsGetNextProcess ? */ - if (!DoesThreadHaveValidCidEntry(NmiContext[core].kthread)) - ReportMissingCidTableEntry(&NmiContext[core]); + if (!DoesThreadHaveValidCidEntry(context->kthread)) + ReportMissingCidTableEntry(context); - if (NmiContext[core].user_thread) + if (IsInstructionPointerInInvalidRegion(context->interrupted_rip, + Modules)) + ReportInvalidRipFoundDuringNmi(context, 0); + + if (context->user_thread) continue; - if (DoesRetInstructionCauseException(NmiContext[core].interrupted_rip)) + if (DoesRetInstructionCauseException(context->interrupted_rip)) ReportInvalidRipFoundDuringNmi( - &NmiContext[core], REPORT_SUBTYPE_EXCEPTION_THROWING_RET); - - if (IsInstructionPointerInInvalidRegion( - NmiContext[core].interrupted_rip, SystemModules)) - ReportInvalidRipFoundDuringNmi(&NmiContext[core], 0); + context, + REPORT_SUBTYPE_EXCEPTION_THROWING_RET); } return STATUS_SUCCESS; @@ -836,12 +879,6 @@ NmiCallback(_Inout_opt_ PVOID Context, _In_ BOOLEAN Handled) context->kthread = PsGetCurrentThread(); context->callback_count++; - DEBUG_VERBOSE( - "[NMI CALLBACK]: Core Number: %lx, Interrupted RIP: %llx, Interrupted RSP: %llx", - core, - machine_frame->rip, - machine_frame->rsp); - return TRUE; } @@ -853,13 +890,16 @@ LaunchNonMaskableInterrupt() { PAGED_CODE(); - PKAFFINITY_EX affinity = ImpExAllocatePool2( - POOL_FLAG_NON_PAGED, sizeof(KAFFINITY_EX), PROC_AFFINITY_POOL); + PKAFFINITY_EX affinity = NULL; + LARGE_INTEGER delay = {0}; + + affinity = ImpExAllocatePool2(POOL_FLAG_NON_PAGED, + sizeof(KAFFINITY_EX), + PROC_AFFINITY_POOL); if (!affinity) return STATUS_MEMORY_NOT_ALLOCATED; - LARGE_INTEGER delay = {0}; delay.QuadPart -= NMI_DELAY_TIME; for (ULONG core = 0; core < ImpKeQueryActiveProcessorCount(0); core++) { @@ -889,9 +929,11 @@ HandleNmiIOCTL() PVOID handle = NULL; SYSTEM_MODULES modules = {0}; PNMI_CONTEXT context = NULL; + UINT32 size = 0; - UINT32 size = ImpKeQueryActiveProcessorCount(0) * sizeof(NMI_CONTEXT); + size = ImpKeQueryActiveProcessorCount(0) * sizeof(NMI_CONTEXT); + /* Ensure we don't continue if another NMI operation is in progress */ if (IsNmiInProgress()) return STATUS_ALREADY_COMMITTED; @@ -974,12 +1016,12 @@ STATIC VOID ReportApcStackwalkViolation(_In_ UINT64 Rip) { - NTSTATUS status = STATUS_UNSUCCESSFUL; - UINT32 packet_size = - CryptRequestRequiredBufferLength(sizeof(APC_STACKWALK_REPORT)); + NTSTATUS status = STATUS_UNSUCCESSFUL; + UINT32 len = 0; + PAPC_STACKWALK_REPORT report = NULL; - PAPC_STACKWALK_REPORT report = - ImpExAllocatePool2(POOL_FLAG_NON_PAGED, packet_size, REPORT_POOL_TAG); + len = CryptRequestRequiredBufferLength(sizeof(APC_STACKWALK_REPORT)); + report = ImpExAllocatePool2(POOL_FLAG_NON_PAGED, len, REPORT_POOL_TAG); if (!report) return; @@ -990,7 +1032,7 @@ ReportApcStackwalkViolation(_In_ UINT64 Rip) report->invalid_rip = Rip; // report->driver ?? todo! - status = CryptEncryptBuffer(report, packet_size); + status = CryptEncryptBuffer(report, len); if (!NT_SUCCESS(status)) { DEBUG_ERROR("CryptEncryptBuffer: %lx", status); @@ -998,7 +1040,7 @@ ReportApcStackwalkViolation(_In_ UINT64 Rip) return; } - IrpQueueSchedulePacket(report, packet_size); + IrpQueueSchedulePacket(report, len); } /* @@ -1015,45 +1057,47 @@ ApcKernelRoutine(_In_ PRKAPC Apc, { PAGED_CODE(); - NTSTATUS status = STATUS_UNSUCCESSFUL; - PVOID buffer = NULL; - INT frames_captured = 0; - PUINT64 frames = 0; - BOOLEAN flag = FALSE; - PAPC_STACKWALK_CONTEXT context = NULL; - PTHREAD_LIST_ENTRY thread_list_entry = NULL; + NTSTATUS status = STATUS_UNSUCCESSFUL; + PVOID buffer = NULL; + INT frames_captured = 0; + UINT64 frame = 0; + PAPC_STACKWALK_CONTEXT context = NULL; + PTHREAD_LIST_ENTRY entry = NULL; context = (PAPC_STACKWALK_CONTEXT)Apc->NormalContext; - FindThreadListEntryByThreadAddress(KeGetCurrentThread(), - &thread_list_entry); + FindThreadListEntryByThreadAddress(KeGetCurrentThread(), &entry); - if (!thread_list_entry) + if (!entry) return; - buffer = ImpExAllocatePool2( - POOL_FLAG_NON_PAGED, STACK_FRAME_POOL_SIZE, POOL_TAG_APC); + buffer = ImpExAllocatePool2(POOL_FLAG_NON_PAGED, + STACK_FRAME_POOL_SIZE, + POOL_TAG_APC); if (!buffer) goto free; - frames_captured = ImpRtlCaptureStackBackTrace( - NULL, STACK_FRAME_POOL_SIZE / sizeof(UINT64), buffer, NULL); + frames_captured = + ImpRtlCaptureStackBackTrace(NULL, + STACK_FRAME_POOL_SIZE / sizeof(UINT64), + buffer, + NULL); if (!frames_captured) goto free; - for (INT index = 0; index < frames_captured; index++) { - frames = (PUINT64)buffer; + for (UINT32 index = 0; index < frames_captured; index++) { + frame = ((PUINT64)buffer)[index]; /* * Apc->NormalContext holds the address of our context data * structure that we passed into KeInitializeApc as the last * argument. */ - if (IsInstructionPointerInInvalidRegion(frames[index], - context->modules)) - ReportApcStackwalkViolation(frames[index]); + if (IsInstructionPointerInInvalidRegion(frame, context->modules)) { + ReportApcStackwalkViolation(frame); + } } free: @@ -1063,8 +1107,8 @@ free: FreeApcAndDecrementApcCount(Apc, APC_CONTEXT_ID_STACKWALK); - thread_list_entry->apc = NULL; - thread_list_entry->apc_queued = FALSE; + entry->apc = NULL; + entry->apc_queued = FALSE; } /* @@ -1085,25 +1129,25 @@ ApcNormalRoutine(_In_opt_ PVOID NormalContext, STATIC VOID -ValidateThreadViaKernelApcCallback(_In_ PTHREAD_LIST_ENTRY ThreadListEntry, +ValidateThreadViaKernelApcCallback(_In_ PTHREAD_LIST_ENTRY Entry, _Inout_opt_ PVOID Context) { PAGED_CODE(); PKAPC apc = NULL; - BOOLEAN apc_status = FALSE; PLONG flags = NULL; - PCHAR previous_mode = NULL; + PCHAR prev_mode = NULL; PUCHAR state = NULL; BOOLEAN apc_queueable = FALSE; - LPCSTR process_name = NULL; - PAPC_STACKWALK_CONTEXT context = (PAPC_STACKWALK_CONTEXT)Context; + LPCSTR proc_name = NULL; + PAPC_STACKWALK_CONTEXT context = NULL; + + context = (PAPC_STACKWALK_CONTEXT)Context; if (!ARGUMENT_PRESENT(Context)) return; - process_name = - ImpPsGetProcessImageFileName(ThreadListEntry->owning_process); + proc_name = ImpPsGetProcessImageFileName(Entry->owning_process); /* * Its possible to set the KThread->ApcQueueable flag to false ensuring @@ -1112,40 +1156,36 @@ ValidateThreadViaKernelApcCallback(_In_ PTHREAD_LIST_ENTRY ThreadListEntry, * before before queueing ours. Since we filter out any system threads * this should be fine... c: */ - 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); + flags = RVA(PLONG, Entry->thread, KTHREAD_MISC_FLAGS_OFFSET); + prev_mode = RVA(PCHAR, Entry->thread, KTHREAD_PREVIOUS_MODE_OFFSET); + state = RVA(PUCHAR, Entry->thread, KTHREAD_STATE_OFFSET); /* * For now, lets only check for system threads. However, we also want to * check for threads executing in kernel mode, i.e KTHREAD->PreviousMode * == UserMode. */ - if (ThreadListEntry->owning_process != PsInitialSystemProcess) + if (Entry->owning_process != PsInitialSystemProcess) return; - if (ThreadListEntry->thread == KeGetCurrentThread() || - !ThreadListEntry->thread) + if (Entry->thread == KeGetCurrentThread() || !Entry->thread) return; DEBUG_VERBOSE( "Validating thread: %llx, process name: %s via kernel APC stackwalk.", - ThreadListEntry->thread, - process_name); + Entry->thread, + proc_name); SetFlag(*flags, KTHREAD_MISC_FLAGS_ALERTABLE); SetFlag(*flags, KTHREAD_MISC_FLAGS_APC_QUEUEABLE); - apc = (PKAPC)ImpExAllocatePool2( - POOL_FLAG_NON_PAGED, sizeof(KAPC), POOL_TAG_APC); + apc = ImpExAllocatePool2(POOL_FLAG_NON_PAGED, sizeof(KAPC), POOL_TAG_APC); if (!apc) return; ImpKeInitializeApc(apc, - ThreadListEntry->thread, + Entry->thread, OriginalApcEnvironment, ApcKernelRoutine, ApcRundownRoutine, @@ -1153,16 +1193,14 @@ ValidateThreadViaKernelApcCallback(_In_ PTHREAD_LIST_ENTRY ThreadListEntry, KernelMode, Context); - apc_status = ImpKeInsertQueueApc(apc, NULL, NULL, IO_NO_INCREMENT); - - if (!apc_status) { + if (!ImpKeInsertQueueApc(apc, NULL, NULL, IO_NO_INCREMENT)) { DEBUG_ERROR("KeInsertQueueApc failed with no status."); ImpExFreePoolWithTag(apc, POOL_TAG_APC); return; } - ThreadListEntry->apc = apc; - ThreadListEntry->apc_queued = TRUE; + Entry->apc = apc; + Entry->apc_queued = TRUE; IncrementApcCount(APC_CONTEXT_ID_STACKWALK); } @@ -1206,15 +1244,17 @@ ValidateThreadsViaKernelApc() return STATUS_SUCCESS; } - context = ImpExAllocatePool2( - POOL_FLAG_NON_PAGED, sizeof(APC_STACKWALK_CONTEXT), POOL_TAG_APC); + context = ImpExAllocatePool2(POOL_FLAG_NON_PAGED, + sizeof(APC_STACKWALK_CONTEXT), + POOL_TAG_APC); if (!context) return STATUS_MEMORY_NOT_ALLOCATED; context->header.context_id = APC_CONTEXT_ID_STACKWALK; - context->modules = ImpExAllocatePool2( - POOL_FLAG_NON_PAGED, sizeof(SYSTEM_MODULES), POOL_TAG_APC); + context->modules = ImpExAllocatePool2(POOL_FLAG_NON_PAGED, + sizeof(SYSTEM_MODULES), + POOL_TAG_APC); if (!context->modules) { ImpExFreePoolWithTag(context, POOL_TAG_APC); @@ -1233,8 +1273,7 @@ ValidateThreadsViaKernelApc() InsertApcContext(context); SetApcAllocationInProgress(context); - RtlRbTreeEnumerate( - GetThreadTree(), ValidateThreadViaKernelApcCallback, context); + ENUMERATE_THREADS(ValidateThreadViaKernelApcCallback, context); UnsetApcAllocationInProgress(context); return status; } @@ -1248,18 +1287,6 @@ FreeApcStackwalkApcContextInformation(_Inout_ PAPC_STACKWALK_CONTEXT Context) ImpExFreePoolWithTag(Context->modules, POOL_TAG_APC); } -#define DPC_STACKWALK_STACKFRAME_COUNT 10 - -/* the first 3 frames are isr handlers which we dont care about */ -#define DPC_STACKWALK_FRAMES_TO_SKIP 3 - -typedef struct _DPC_CONTEXT { - UINT64 stack_frame[DPC_STACKWALK_STACKFRAME_COUNT]; - UINT16 frames_captured; - volatile BOOLEAN executed; - -} DPC_CONTEXT, *PDPC_CONTEXT; - VOID DpcStackwalkCallbackRoutine(_In_ PKDPC Dpc, _In_opt_ PVOID DeferredContext, @@ -1269,17 +1296,19 @@ DpcStackwalkCallbackRoutine(_In_ PKDPC Dpc, UNREFERENCED_PARAMETER(Dpc); UNREFERENCED_PARAMETER(SystemArgument2); + PDPC_CONTEXT context = NULL; + if (!ARGUMENT_PRESENT(DeferredContext)) return; - PDPC_CONTEXT context = - &((PDPC_CONTEXT)DeferredContext)[KeGetCurrentProcessorNumber()]; + context = &((PDPC_CONTEXT)DeferredContext)[KeGetCurrentProcessorNumber()]; context->frames_captured = ImpRtlCaptureStackBackTrace(DPC_STACKWALK_FRAMES_TO_SKIP, DPC_STACKWALK_STACKFRAME_COUNT, &context->stack_frame, NULL); + InterlockedExchange(&context->executed, TRUE); #pragma warning(push) @@ -1292,31 +1321,18 @@ DpcStackwalkCallbackRoutine(_In_ PKDPC Dpc, context->frames_captured); } -STATIC -BOOLEAN -CheckForDpcCompletion(_In_ PDPC_CONTEXT Context) -{ - for (UINT32 index = 0; index < ImpKeQueryActiveProcessorCount(0); index++) { - if (!InterlockedExchange(&Context[index].executed, - Context[index].executed)) - return FALSE; - } - - return TRUE; -} - STATIC VOID ReportDpcStackwalkViolation(_In_ PDPC_CONTEXT Context, _In_ UINT64 Frame, _In_ UINT32 ReportSubtype) { - NTSTATUS status = STATUS_UNSUCCESSFUL; - UINT32 packet_size = - CryptRequestRequiredBufferLength(sizeof(DPC_STACKWALK_REPORT)); + NTSTATUS status = STATUS_UNSUCCESSFUL; + UINT32 len = 0; + PDPC_STACKWALK_REPORT report = NULL; - PDPC_STACKWALK_REPORT report = - ImpExAllocatePool2(POOL_FLAG_NON_PAGED, packet_size, REPORT_POOL_TAG); + len = CryptRequestRequiredBufferLength(sizeof(DPC_STACKWALK_REPORT)); + report = ImpExAllocatePool2(POOL_FLAG_NON_PAGED, len, REPORT_POOL_TAG); if (!report) return; @@ -1331,7 +1347,7 @@ ReportDpcStackwalkViolation(_In_ PDPC_CONTEXT Context, // - 0x50, // APC_STACKWALK_BUFFER_SIZE); - status = CryptEncryptBuffer(report, packet_size); + status = CryptEncryptBuffer(report, len); if (!NT_SUCCESS(status)) { DEBUG_ERROR("CryptEncryptBuffer: %lx", status); @@ -1339,7 +1355,7 @@ ReportDpcStackwalkViolation(_In_ PDPC_CONTEXT Context, return; } - IrpQueueSchedulePacket(report, packet_size); + IrpQueueSchedulePacket(report, len); } STATIC @@ -1348,6 +1364,7 @@ ValidateDpcStackFrame(_In_ PDPC_CONTEXT Context, _In_ PSYSTEM_MODULES Modules) { NTSTATUS status = STATUS_UNSUCCESSFUL; BOOLEAN flag = FALSE; + UINT64 rip = 0; /* With regards to this, lets only check the interrupted rip */ if (DoesRetInstructionCauseException(Context->stack_frame[0])) @@ -1356,7 +1373,7 @@ ValidateDpcStackFrame(_In_ PDPC_CONTEXT Context, _In_ PSYSTEM_MODULES Modules) REPORT_SUBTYPE_EXCEPTION_THROWING_RET); for (UINT32 frame = 0; frame < Context->frames_captured; frame++) { - UINT64 rip = Context->stack_frame[frame]; + rip = Context->stack_frame[frame]; if (IsInstructionPointerInInvalidRegion(rip, Modules)) ReportDpcStackwalkViolation(Context, rip, 0); @@ -1368,11 +1385,17 @@ VOID ValidateDpcCapturedStack(_In_ PSYSTEM_MODULES Modules, _In_ PDPC_CONTEXT Context) { - BOOLEAN flag = FALSE; - PDPC_STACKWALK_REPORT report = NULL; - UINT32 count = ImpKeQueryActiveProcessorCount(0); + BOOLEAN flag = FALSE; + PDPC_CONTEXT context = NULL; + UINT32 count = ImpKeQueryActiveProcessorCount(0); for (UINT32 core = 0; core < count; core++) { + context = &Context[core]; + + if (!context->executed) + DEBUG_WARNING("DPC Stackwalk routine not executed. Core: %lx", + core); + ValidateDpcStackFrame(&Context[core], Modules); } } @@ -1390,8 +1413,9 @@ DispatchStackwalkToEachCpuViaDpc() NTSTATUS status = STATUS_UNSUCCESSFUL; PDPC_CONTEXT context = NULL; SYSTEM_MODULES modules = {0}; - UINT32 size = ImpKeQueryActiveProcessorCount(0) * sizeof(DPC_CONTEXT); + UINT32 size = 0; + size = ImpKeQueryActiveProcessorCount(0) * sizeof(DPC_CONTEXT); context = ImpExAllocatePool2(POOL_FLAG_NON_PAGED, size, POOL_TAG_DPC); if (!context) @@ -1409,12 +1433,11 @@ DispatchStackwalkToEachCpuViaDpc() * the DPC queue and executed immediately.*/ ImpKeGenericCallDpc(DpcStackwalkCallbackRoutine, context); - while (!CheckForDpcCompletion(context)) - YieldProcessor(); + /* Flush all DPC's in the system to ensure ours have run */ + KeFlushQueuedDpcs(); ValidateDpcCapturedStack(&modules, context); - DEBUG_VERBOSE("Finished validating cores via dpc"); end: if (modules.address) @@ -1615,12 +1638,12 @@ STATIC VOID ReportDataTableInvalidRoutine(_In_ TABLE_ID TableId, _In_ UINT64 Address) { - NTSTATUS status = STATUS_UNSUCCESSFUL; - UINT32 packet_size = - CryptRequestRequiredBufferLength(sizeof(DATA_TABLE_ROUTINE_REPORT)); + NTSTATUS status = STATUS_UNSUCCESSFUL; + UINT32 len = 0; + PDATA_TABLE_ROUTINE_REPORT report = NULL; - PDATA_TABLE_ROUTINE_REPORT report = - ImpExAllocatePool2(POOL_FLAG_NON_PAGED, packet_size, REPORT_POOL_TAG); + len = CryptRequestRequiredBufferLength(sizeof(DATA_TABLE_ROUTINE_REPORT)); + report = ImpExAllocatePool2(POOL_FLAG_NON_PAGED, len, REPORT_POOL_TAG); if (!report) return; @@ -1634,9 +1657,10 @@ ReportDataTableInvalidRoutine(_In_ TABLE_ID TableId, _In_ UINT64 Address) report->address = Address; report->table_id = TableId; report->index = 0; + RtlCopyMemory(report->routine, Address, DATA_TABLE_ROUTINE_BUF_SIZE); - status = CryptEncryptBuffer(report, packet_size); + status = CryptEncryptBuffer(report, len); if (!NT_SUCCESS(status)) { DEBUG_ERROR("CryptEncryptBuffer: %lx", status); @@ -1644,7 +1668,7 @@ ReportDataTableInvalidRoutine(_In_ TABLE_ID TableId, _In_ UINT64 Address) return; } - IrpQueueSchedulePacket(report, packet_size); + IrpQueueSchedulePacket(report, len); } NTSTATUS @@ -1693,19 +1717,25 @@ NTSTATUS GetDriverObjectByDriverName(_In_ PUNICODE_STRING DriverName, _Out_ PDRIVER_OBJECT* DriverObject) { - HANDLE handle = NULL; - OBJECT_ATTRIBUTES attributes = {0}; - PVOID directory = {0}; - UNICODE_STRING directory_name = {0}; - NTSTATUS status = STATUS_UNSUCCESSFUL; - POBJECT_DIRECTORY directory_object = NULL; + HANDLE handle = NULL; + OBJECT_ATTRIBUTES attributes = {0}; + PVOID dir = {0}; + UNICODE_STRING dir_name = {0}; + NTSTATUS status = STATUS_UNSUCCESSFUL; + POBJECT_DIRECTORY dir_object = NULL; + POBJECT_DIRECTORY_ENTRY entry = NULL; + POBJECT_DIRECTORY_ENTRY sub_entry = NULL; + PDRIVER_OBJECT driver = NULL; *DriverObject = NULL; - ImpRtlInitUnicodeString(&directory_name, L"\\Driver"); + ImpRtlInitUnicodeString(&dir_name, L"\\Driver"); - InitializeObjectAttributes( - &attributes, &directory_name, OBJ_CASE_INSENSITIVE, NULL, NULL); + InitializeObjectAttributes(&attributes, + &dir_name, + OBJ_CASE_INSENSITIVE, + NULL, + NULL); status = ImpZwOpenDirectoryObject(&handle, DIRECTORY_ALL_ACCESS, &attributes); @@ -1715,8 +1745,12 @@ GetDriverObjectByDriverName(_In_ PUNICODE_STRING DriverName, return status; } - status = ImpObReferenceObjectByHandle( - handle, DIRECTORY_ALL_ACCESS, NULL, KernelMode, &directory, NULL); + status = ImpObReferenceObjectByHandle(handle, + DIRECTORY_ALL_ACCESS, + NULL, + KernelMode, + &dir, + NULL); if (!NT_SUCCESS(status)) { DEBUG_ERROR("ObReferenceObjectByHandle failed with status %x", status); @@ -1724,34 +1758,35 @@ GetDriverObjectByDriverName(_In_ PUNICODE_STRING DriverName, return status; } - directory_object = (POBJECT_DIRECTORY)directory; + dir_object = (POBJECT_DIRECTORY)dir; - ImpExAcquirePushLockExclusiveEx(&directory_object->Lock, NULL); + ImpExAcquirePushLockExclusiveEx(&dir_object->Lock, NULL); - for (INT index = 0; index < NUMBER_HASH_BUCKETS; index++) { - POBJECT_DIRECTORY_ENTRY entry = directory_object->HashBuckets[index]; + for (UINT32 index = 0; index < NUMBER_HASH_BUCKETS; index++) { + entry = dir_object->HashBuckets[index]; if (!entry) continue; - POBJECT_DIRECTORY_ENTRY sub_entry = entry; + sub_entry = entry; while (sub_entry) { - PDRIVER_OBJECT current_driver = sub_entry->Object; + driver = GetObjectFromDirectory(sub_entry); - if (!RtlCompareUnicodeString( - DriverName, ¤t_driver->DriverName, FALSE)) { - *DriverObject = current_driver; + if (!RtlCompareUnicodeString(DriverName, + &driver->DriverName, + FALSE)) { + *DriverObject = driver; goto end; } - sub_entry = sub_entry->ChainLink; + sub_entry = GetNextObject(sub_entry); } } end: - ImpExReleasePushLockExclusiveEx(&directory_object->Lock, 0); - ImpObDereferenceObject(directory); + ImpExReleasePushLockExclusiveEx(&dir_object->Lock, 0); + ImpObDereferenceObject(dir); ImpZwClose(handle); return STATUS_SUCCESS; } @@ -1759,13 +1794,13 @@ end: PVOID FindDriverBaseNoApi(_In_ PDRIVER_OBJECT DriverObject, _In_ PWCH Name) { - PKLDR_DATA_TABLE_ENTRY first = - (PKLDR_DATA_TABLE_ENTRY)DriverObject->DriverSection; + PKLDR_DATA_TABLE_ENTRY first = NULL; + PKLDR_DATA_TABLE_ENTRY entry = NULL; /* first entry contains invalid data, 2nd entry is the kernel */ - PKLDR_DATA_TABLE_ENTRY entry = - ((PKLDR_DATA_TABLE_ENTRY)DriverObject->DriverSection) - ->InLoadOrderLinks.Flink->Flink; + first = (PKLDR_DATA_TABLE_ENTRY)DriverObject->DriverSection; + entry = ((PKLDR_DATA_TABLE_ENTRY)DriverObject->DriverSection) + ->InLoadOrderLinks.Flink->Flink; while (entry->InLoadOrderLinks.Flink != first) { /* todo: write our own unicode string comparison function, since @@ -1781,7 +1816,6 @@ FindDriverBaseNoApi(_In_ PDRIVER_OBJECT DriverObject, _In_ PWCH Name) return NULL; } -STATIC VOID ValidateDispatchTableRoutines(_In_ PVOID* Table, _In_ UINT32 Entries) { @@ -1952,12 +1986,12 @@ VOID ReportWin32kBase_DxgInterfaceViolation(_In_ UINT32 TableIndex, _In_ UINT64 Address) { - NTSTATUS status = STATUS_UNSUCCESSFUL; - UINT32 packet_size = - CryptRequestRequiredBufferLength(sizeof(DATA_TABLE_ROUTINE_REPORT)); + NTSTATUS status = STATUS_UNSUCCESSFUL; + UINT32 len = 0; + PDATA_TABLE_ROUTINE_REPORT report = NULL; - PDATA_TABLE_ROUTINE_REPORT report = - ImpExAllocatePool2(POOL_FLAG_NON_PAGED, packet_size, REPORT_POOL_TAG); + len = CryptRequestRequiredBufferLength(sizeof(DATA_TABLE_ROUTINE_REPORT)); + report = ImpExAllocatePool2(POOL_FLAG_NON_PAGED, len, REPORT_POOL_TAG); if (!report) return; @@ -1970,7 +2004,7 @@ ReportWin32kBase_DxgInterfaceViolation(_In_ UINT32 TableIndex, // todo! report->routine = ?? // todo: maybe get routine by name from index ? - status = CryptEncryptBuffer(report, packet_size); + status = CryptEncryptBuffer(report, len); if (!NT_SUCCESS(status)) { DEBUG_ERROR("CryptEncryptBuffer: %lx", status); @@ -1978,7 +2012,7 @@ ReportWin32kBase_DxgInterfaceViolation(_In_ UINT32 TableIndex, return; } - IrpQueueSchedulePacket(report, packet_size); + IrpQueueSchedulePacket(report, len); } STATIC @@ -1992,6 +2026,7 @@ ValidateWin32kBase_gDxgInterface() KAPC_STATE apc = {0}; PKPROCESS winlogon = NULL; PVOID* dxg_interface = NULL; + PVOID entry = NULL; status = GetSystemModuleInformation(&modules); @@ -2036,7 +2071,7 @@ ValidateWin32kBase_gDxgInterface() if (!dxg_interface[index]) continue; - PVOID entry = FindChainedPointerEnding(dxg_interface[index]); + entry = FindChainedPointerEnding(dxg_interface[index]); #if DEBUG DEBUG_INFO("chain entry test: %p", entry); diff --git a/driver/modules.h b/driver/modules.h index 337f264..d2e7ab5 100644 --- a/driver/modules.h +++ b/driver/modules.h @@ -55,7 +55,7 @@ VOID FreeApcStackwalkApcContextInformation(_Inout_ PAPC_STACKWALK_CONTEXT Context); BOOLEAN -IsInstructionPointerInInvalidRegion(_In_ UINT64 RIP, +IsInstructionPointerInInvalidRegion(_In_ UINT64 Rip, _In_ PSYSTEM_MODULES SystemModules); PVOID diff --git a/driver/pe.c b/driver/pe.c index d183533..45c698a 100644 --- a/driver/pe.c +++ b/driver/pe.c @@ -11,7 +11,7 @@ PeGetNtHeaderSafe(_In_ PVOID Image) if (dos->e_magic != IMAGE_DOS_SIGNATURE) return NULL; - return CONVERT_RELATIVE_ADDRESS(PNT_HEADER_64, Image, dos->e_lfanew); + return RVA(PNT_HEADER_64, Image, dos->e_lfanew); } PNT_HEADER_64 @@ -22,7 +22,7 @@ PeGetNtHeader(_In_ PVOID Image) if (dos->e_magic != IMAGE_DOS_SIGNATURE) return NULL; - return CONVERT_RELATIVE_ADDRESS(PNT_HEADER_64, Image, dos->e_lfanew); + return RVA(PNT_HEADER_64, Image, dos->e_lfanew); } PIMAGE_DATA_DIRECTORY @@ -59,7 +59,7 @@ PeGetExportDirectory(_In_ PVOID Image, if (!ExportDataDirectory->VirtualAddress || !ExportDataDirectory->Size) return NULL; - return CONVERT_RELATIVE_ADDRESS( + return RVA( PIMAGE_EXPORT_DIRECTORY, Image, ExportDataDirectory->VirtualAddress); } @@ -73,7 +73,7 @@ PeGetExportDirectorySafe(_In_ PVOID Image, if (!ExportDataDirectory->VirtualAddress || !ExportDataDirectory->Size) return NULL; - return CONVERT_RELATIVE_ADDRESS( + return RVA( PIMAGE_EXPORT_DIRECTORY, Image, ExportDataDirectory->VirtualAddress); } @@ -118,16 +118,16 @@ PeFindExportByName(_In_ PVOID Image, _In_ PCHAR Name) return NULL; PUINT32 functions = - CONVERT_RELATIVE_ADDRESS(PUINT32, Image, export->AddressOfFunctions); + RVA(PUINT32, Image, export->AddressOfFunctions); PUINT32 names = - CONVERT_RELATIVE_ADDRESS(PUINT32, Image, export->AddressOfNames); + RVA(PUINT32, Image, export->AddressOfNames); PUINT16 ordinals = - CONVERT_RELATIVE_ADDRESS(PUINT16, Image, export->AddressOfNameOrdinals); + RVA(PUINT16, Image, export->AddressOfNameOrdinals); for (UINT32 index = 0; index < export->NumberOfNames; index++) { - PCHAR export = CONVERT_RELATIVE_ADDRESS(PCHAR, Image, names[index]); + PCHAR export = RVA(PCHAR, Image, names[index]); if (!strcmp(Name, export)) - return CONVERT_RELATIVE_ADDRESS( + return RVA( PVOID, Image, functions[ordinals[index]]); }