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

This commit is contained in:
2023-07-01 16:33:13 +02:00
parent 0df66bdea2
commit 5541c05d1d
4 changed files with 154 additions and 33 deletions

View File

@ -7,6 +7,10 @@
#include "ci.h"
/********************************************************************************
* NtDll and other internal types *
********************************************************************************/
// Structure representing a loaded module
typedef struct _SYSTEM_MODULE_INFORMATION_ENTRY
{
@ -79,13 +83,96 @@ 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,
_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
)
@ -663,20 +748,15 @@ DomitoGetProcessImageName(
HANDLE hProcess = NULL;
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,9 +983,9 @@ DomitoValidateFileLegacyMode(
if (certDirectory)
{
ExFreePool(certDirectory);
G_Memory.Free(certDirectory);
}
return status;
}
#pragma code_seg()