Replaced passing allocator to functions with library global memory management support types

Domito.sln.DotSettings

@@ -4,6 +4,7 @@
 	<s:Boolean x:Key="/Default/UserDictionary/Words/=CALG/@EntryIndexedValue">True</s:Boolean>
 	<s:Boolean x:Key="/Default/UserDictionary/Words/=COFF/@EntryIndexedValue">True</s:Boolean>
 	<s:Boolean x:Key="/Default/UserDictionary/Words/=Domito/@EntryIndexedValue">True</s:Boolean>
+	<s:Boolean x:Key="/Default/UserDictionary/Words/=Domito_0027s/@EntryIndexedValue">True</s:Boolean>
 	<s:Boolean x:Key="/Default/UserDictionary/Words/=LPWIN/@EntryIndexedValue">True</s:Boolean>
 	<s:Boolean x:Key="/Default/UserDictionary/Words/=PDOMITO/@EntryIndexedValue">True</s:Boolean>
 	<s:Boolean x:Key="/Default/UserDictionary/Words/=PKCS/@EntryIndexedValue">True</s:Boolean>

README.md

@@ -14,3 +14,6 @@ for Windows NT3.51/NT4/2000/XP/2003/Vista/7/8 kernel mode,
 both 32 and 64 bit platforms](http://alter.org.ua/en/docs/nt_kernel/procaddr/)
 - [Use ci.dll API for validating Authenticode signature of files](https://github.com/Ido-Moshe-Github/CiDllDemo)
 - [Helper functions for calculating the authenticode digest for a portable executable file](https://github.com/mihaly044/pedigest)
+- Custom memory allocator exposure used by SDL
+  - [SDL_stdinc.h](https://github.com/libsdl-org/SDL/blob/main/include/SDL3/SDL_stdinc.h)
+  - [SDL_malloc.c](https://github.com/libsdl-org/SDL/blob/main/src/stdlib/SDL_malloc.c)

include/Domito.h

@@ -4,16 +4,16 @@
 
 
 /********************************************************************************
- * Memory management, misc.                                                     *
+ * Memory management                                                            *
  ********************************************************************************/
 
- //
+//
- // Custom allocator for function that allocate pool memory.
+// Allocator function the library uses.
- // 
+// 
 typedef
 _IRQL_requires_same_
 _Function_class_(EVT_DOMITO_ALLOCATE_ROUTINE)
-__drv_allocatesMem(Mem)
+__drv_allocatesMem(mem)
 PVOID
 NTAPI
 EVT_DOMITO_ALLOCATE_ROUTINE(
@@ -21,6 +21,48 @@ EVT_DOMITO_ALLOCATE_ROUTINE(
 );
 typedef EVT_DOMITO_ALLOCATE_ROUTINE* PFN_DOMITO_ALLOCATE_ROUTINE;
 
+//
+// Freeing function the library uses.
+// 
+typedef
+_IRQL_requires_same_
+_Function_class_(EVT_DOMITO_FREE_ROUTINE)
+void
+NTAPI
+EVT_DOMITO_FREE_ROUTINE(
+    _In_ __drv_freesMem(mem) PVOID Memory
+);
+typedef EVT_DOMITO_FREE_ROUTINE* PFN_DOMITO_FREE_ROUTINE;
+
+//
+// Get the original set of Domito memory functions.
+// 
+EXTERN_C
+void
+DomitoGetOriginalMemoryFunctions(
+    _Out_opt_ PFN_DOMITO_ALLOCATE_ROUTINE* Allocator,
+    _Out_opt_ PFN_DOMITO_FREE_ROUTINE* Free
+);
+
+//
+// Get the current set of Domito memory functions.
+//  
+EXTERN_C
+void
+DomitoGetMemoryFunctions(
+    _Out_opt_ PFN_DOMITO_ALLOCATE_ROUTINE* Allocator,
+    _Out_opt_ PFN_DOMITO_FREE_ROUTINE* Free
+);
+
+//
+//  Replace Domito's memory allocation functions with a custom set
+EXTERN_C
+void
+DomitoSetMemoryFunctions(
+    _In_opt_ PFN_DOMITO_ALLOCATE_ROUTINE Allocator,
+    _In_opt_ PFN_DOMITO_FREE_ROUTINE Free
+);
+
 
 /********************************************************************************
  * Cryptography                                                                 *
@@ -104,7 +146,6 @@ _IRQL_requires_max_(PASSIVE_LEVEL)
 EXTERN_C
 NTSTATUS
 DomitoFindModuleBaseAddress(
-    _In_ PFN_DOMITO_ALLOCATE_ROUTINE Allocator,
     _In_ STRING ModuleName,
     _Inout_opt_ PVOID* ModuleBase
 );
@@ -175,7 +216,6 @@ _IRQL_requires_max_(PASSIVE_LEVEL)
 EXTERN_C
 NTSTATUS
 DomitoCalculatePortableExecutableDigest(
-    _In_ PFN_DOMITO_ALLOCATE_ROUTINE Allocator,
     _In_ PUCHAR pPeBytes,
     _In_ ULONG PeSize,
     _Out_ PUINT32 pDigestCalgOut,
@@ -194,7 +234,6 @@ _IRQL_requires_max_(PASSIVE_LEVEL)
 EXTERN_C
 NTSTATUS
 DomitoGetProcessImageName(
-    _In_ PFN_DOMITO_ALLOCATE_ROUTINE Allocator,
     _In_ ULONG ProcessId,
     _Inout_ PUNICODE_STRING* ProcessImageName
 );
@@ -205,7 +244,6 @@ _IRQL_requires_max_(PASSIVE_LEVEL)
 EXTERN_C
 NTSTATUS
 DomitoValidateFileLegacyMode(
-    _In_ PFN_DOMITO_ALLOCATE_ROUTINE Allocator,
     _In_ HANDLE FileHandle,
     _In_ PVOID Hash,
     _In_ UINT32 HashSize,

src/Domito.cpp

@@ -7,6 +7,10 @@
 #include "ci.h"
 
 
+/********************************************************************************
+ * NtDll and other internal types                                               *
+ ********************************************************************************/
+
 // Structure representing a loaded module
 typedef struct _SYSTEM_MODULE_INFORMATION_ENTRY
 {
@@ -79,12 +83,95 @@ typedef NTSTATUS(*QUERY_INFO_PROCESS) (
 static QUERY_INFO_PROCESS ZwQueryInformationProcess;
 
 
+/********************************************************************************
+ * Memory management                                                            *
+ ********************************************************************************/
+
+#define DOMITO_POOL_TAG     'imoD'
+
+static PVOID NTAPI DomitoDefaultMalloc(size_t s)
+{
+#pragma warning(disable:4996)
+    return ExAllocatePoolWithTag(NonPagedPool, s, DOMITO_POOL_TAG);
+#pragma warninf(default:4996)
+}
+
+static void NTAPI DomitoDefaultFree(PVOID p)
+{
+    ExFreePoolWithTag(p, DOMITO_POOL_TAG);
+}
+
+static struct
+{
+    PFN_DOMITO_ALLOCATE_ROUTINE Allocate;
+
+    PFN_DOMITO_FREE_ROUTINE Free;
+} G_Memory = {
+    DomitoDefaultMalloc,
+    DomitoDefaultFree
+};
+
+void
+DomitoGetOriginalMemoryFunctions(
+    _Out_opt_ PFN_DOMITO_ALLOCATE_ROUTINE* Allocator,
+    _Out_opt_ PFN_DOMITO_FREE_ROUTINE* Free
+)
+{
+    if (Allocator)
+    {
+        *Allocator = DomitoDefaultMalloc;
+    }
+
+    if (Free)
+    {
+        *Free = DomitoDefaultFree;
+    }
+}
+
+void
+DomitoGetMemoryFunctions(
+    _Out_opt_ PFN_DOMITO_ALLOCATE_ROUTINE* Allocator,
+    _Out_opt_ PFN_DOMITO_FREE_ROUTINE* Free
+)
+{
+    if (Allocator)
+    {
+        *Allocator = G_Memory.Allocate;
+    }
+
+    if (Free)
+    {
+        *Free = G_Memory.Free;
+    }
+}
+
+void
+DomitoSetMemoryFunctions(
+    _In_opt_ PFN_DOMITO_ALLOCATE_ROUTINE Allocator,
+    _In_opt_ PFN_DOMITO_FREE_ROUTINE Free
+)
+{
+    if (Allocator)
+    {
+        G_Memory.Allocate = Allocator;
+    }
+
+    if (Free)
+    {
+        G_Memory.Free = Free;
+    }
+}
+
+
+/********************************************************************************
+ * Library functions                                                            *
+ ********************************************************************************/
+
 _Success_(return == STATUS_SUCCESS)
 _Must_inspect_result_
 _IRQL_requires_max_(PASSIVE_LEVEL)
 NTSTATUS
 DomitoFindModuleBaseAddress(
-    _In_ PFN_DOMITO_ALLOCATE_ROUTINE Allocator,
     _In_ STRING ModuleName,
     _Inout_opt_ PVOID * ModuleBase
 )
@@ -108,7 +195,7 @@ DomitoFindModuleBaseAddress(
     }
 
     // Allocate memory for the module information
-    moduleInfo = (PSYSTEM_MODULE_INFORMATION)Allocator(
+    moduleInfo = (PSYSTEM_MODULE_INFORMATION)G_Memory.Allocate(
         bufferSize
     );
 
@@ -127,7 +214,7 @@ DomitoFindModuleBaseAddress(
 
     if (!NT_SUCCESS(status))
     {
-        ExFreePool(moduleInfo);
+        G_Memory.Free(moduleInfo);
         return status;
     }
 
@@ -152,7 +239,7 @@ DomitoFindModuleBaseAddress(
         }
     }
 
-    ExFreePool(moduleInfo);
+    G_Memory.Free(moduleInfo);
 
     return status;
 }
@@ -371,7 +458,6 @@ _Must_inspect_result_
 _IRQL_requires_max_(PASSIVE_LEVEL)
 NTSTATUS
 DomitoCalculatePortableExecutableDigest(
-    _In_ PFN_DOMITO_ALLOCATE_ROUTINE Allocator,
     _In_ PUCHAR pPeBytes,
     _In_ ULONG PeSize,
     _Out_ PUINT32 pDigestCalgOut,
@@ -424,7 +510,7 @@ DomitoCalculatePortableExecutableDigest(
     // TODO: Not sure if  16 * 512 * 512 is right. Do something better!
     //
     ULONG copySize = phDos->e_lfanew + sizeof(IMAGE_FILE_HEADER) + 4 + 0x40;
-    const PUCHAR pBuf = (PUCHAR)Allocator(16 * 512 * 512);
+    const PUCHAR pBuf = (PUCHAR)G_Memory.Allocate(16 * 512 * 512);
     if (!pBuf)
     {
         return STATUS_INSUFFICIENT_RESOURCES;
@@ -480,7 +566,7 @@ DomitoCalculatePortableExecutableDigest(
     //
     // Allocate a buffer to store the resulting hash
     //
-    pHash = (PUCHAR)Allocator(hashLength);
+    pHash = (PUCHAR)G_Memory.Allocate(hashLength);
     if (!pHash)
     {
         status = STATUS_INSUFFICIENT_RESOURCES;
@@ -631,7 +717,7 @@ DomitoCalculatePortableExecutableDigest(
 cleanup:
     if (pBuf)
     {
-        ExFreePool(pBuf);
+        G_Memory.Free(pBuf);
     }
 
     if (hbHash)
@@ -653,7 +739,6 @@ _IRQL_requires_max_(PASSIVE_LEVEL)
 #pragma code_seg("PAGED")
 NTSTATUS
 DomitoGetProcessImageName(
-    _In_ PFN_DOMITO_ALLOCATE_ROUTINE Allocator,
     _In_ ULONG ProcessId,
     _Inout_ PUNICODE_STRING * ProcessImageName
 )
@@ -664,19 +749,14 @@ DomitoGetProcessImageName(
 
     PAGED_CODE(); // this eliminates the possibility of the IDLE Thread/Process
     
-    if (Allocator == NULL)
+    if (ProcessId == 0 || ProcessId == 4)
     {
         return STATUS_INVALID_PARAMETER_1;
     }
 
-    if (ProcessId == 0 || ProcessId == 4)
-    {
-        return STATUS_INVALID_PARAMETER_2;
-    }
-
     if (ProcessImageName == NULL)
     {
-        return STATUS_INVALID_PARAMETER_3;
+        return STATUS_INVALID_PARAMETER_2;
     }
 
     CLIENT_ID cid;
@@ -730,7 +810,7 @@ DomitoGetProcessImageName(
         goto cleanUp;
     }
 
-    *ProcessImageName = (PUNICODE_STRING)Allocator(returnedLength);
+    *ProcessImageName = (PUNICODE_STRING)G_Memory.Allocate(returnedLength);
 
     if (*ProcessImageName == NULL)
     {
@@ -747,7 +827,7 @@ DomitoGetProcessImageName(
         &returnedLength
     )))
     {
-        ExFreePool(*ProcessImageName);
+        G_Memory.Free(*ProcessImageName);
     }
 
 cleanUp:
@@ -767,7 +847,6 @@ _IRQL_requires_max_(PASSIVE_LEVEL)
 #pragma code_seg("PAGED")
 NTSTATUS
 DomitoValidateFileLegacyMode(
-    _In_ PFN_DOMITO_ALLOCATE_ROUTINE Allocator,
     _In_ HANDLE FileHandle,
     _In_ PVOID Hash,
     _In_ UINT32 HashSize,
@@ -800,7 +879,7 @@ DomitoValidateFileLegacyMode(
         if (SecurityDirectory->Size != 0u &&
             SecurityDirectory->VirtualAddress != 0u)
         {
-            certDirectory = Allocator(SecurityDirectory->Size);
+            certDirectory = G_Memory.Allocate(SecurityDirectory->Size);
             if (certDirectory == NULL)
             {
                 status = STATUS_INSUFFICIENT_RESOURCES;
@@ -904,7 +983,7 @@ DomitoValidateFileLegacyMode(
 
     if (certDirectory)
     {
-        ExFreePool(certDirectory);
+        G_Memory.Free(certDirectory);
     }
 
     return status;