some code cleaning

.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:

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

driver/containers/tree.h

@@ -3,8 +3,8 @@
 
 #include "../common.h"
 
-#define RB_TREE_EQUAL 0
+#define RB_TREE_EQUAL        0
-#define RB_TREE_LESS_THAN 1
+#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);
 

driver/crypt.c

@@ -58,14 +58,16 @@ CryptEncryptImportsArray(_In_ PUINT64 Array, _In_ UINT32 Entries)
         __m256i load_block    = {0};
         __m256i xored_block   = {0};
 
-        RtlCopyMemory(
+        RtlCopyMemory(&current_block,
-            &current_block, &Array[block_index * block_size], sizeof(__m256i));
+                      &Array[block_index * block_size],
+                      sizeof(__m256i));
 
         load_block  = _mm256_loadu_si256(&current_block);
         xored_block = _mm256_xor_si256(load_block, *imports_key);
 
-        RtlCopyMemory(
+        RtlCopyMemory(&Array[block_index * block_size],
-            &Array[block_index * block_size], &xored_block, sizeof(__m256i));
+                      &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(
+    RtlCopyMemory(&load_block,
-        &load_block, &Array[BlockIndex * block_size], sizeof(__m256i));
+                  &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(
+    CryptFindContainingBlockForArrayIndex(EntryIndex,
-        EntryIndex, block_size, &containing_block_index, &block_sub_index);
+                                          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(
+    session->key_object = ExAllocatePool2(POOL_FLAG_NON_PAGED,
-        POOL_FLAG_NON_PAGED, session->key_object_length, POOL_TAG_CRYPT);
+                                          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(
+    status = BCryptOpenAlgorithmProvider(handle,
-        handle, BCRYPT_AES_ALGORITHM, NULL, BCRYPT_PROV_DISPATCH);
+                                         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;
 }

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

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

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]]);
     }