refactoring

This commit is contained in:
Boring
2022-10-03 17:25:02 +08:00
parent 1af4e81913
commit ff8a663aca
21 changed files with 581 additions and 400 deletions
+70 -74
View File
@@ -1,7 +1,7 @@
#include "stdafx.h"
#include <cstddef>
static NTSTATUS NTAPI RtlFreeDependencies(IN PLDR_DATA_TABLE_ENTRY_WIN10 LdrEntry) {
static NTSTATUS RtlFreeDependencies(_In_ PLDR_DATA_TABLE_ENTRY_WIN10 LdrEntry) {
_LDR_DDAG_NODE* DependentDdgeNode = nullptr;
PLDR_DATA_TABLE_ENTRY_WIN10 ModuleEntry = nullptr;
_LDRP_CSLIST* head = (decltype(head))LdrEntry->DdagNode->Dependencies, * entry = head;
@@ -81,7 +81,7 @@ static NTSTATUS NTAPI RtlFreeDependencies(IN PLDR_DATA_TABLE_ENTRY_WIN10 LdrEntr
return STATUS_SUCCESS;
}
PLDR_DATA_TABLE_ENTRY NTAPI RtlAllocateDataTableEntry(IN PVOID BaseAddress) {
PLDR_DATA_TABLE_ENTRY NTAPI RtlAllocateDataTableEntry(_In_ PVOID BaseAddress) {
PLDR_DATA_TABLE_ENTRY LdrEntry = nullptr;
PIMAGE_NT_HEADERS NtHeader;
HANDLE heap = NtCurrentPeb()->ProcessHeap;
@@ -89,22 +89,22 @@ PLDR_DATA_TABLE_ENTRY NTAPI RtlAllocateDataTableEntry(IN PVOID BaseAddress) {
/* Make sure the header is valid */
if (NtHeader = RtlImageNtHeader(BaseAddress)) {
/* Allocate an entry */
LdrEntry = (PLDR_DATA_TABLE_ENTRY)RtlAllocateHeap(heap, HEAP_ZERO_MEMORY, LdrpDataTableEntrySize());
LdrEntry = (PLDR_DATA_TABLE_ENTRY)RtlAllocateHeap(heap, HEAP_ZERO_MEMORY, MmpGlobalDataPtr->LdrDataTableEntrySize);
}
/* Return the entry */
return LdrEntry;
}
bool NTAPI RtlInitializeLdrDataTableEntry(
OUT PLDR_DATA_TABLE_ENTRY LdrEntry,
IN DWORD dwFlags,
IN PVOID BaseAddress,
IN UNICODE_STRING& DllBaseName,
IN UNICODE_STRING& DllFullName) {
RtlZeroMemory(LdrEntry, LdrpDataTableEntrySize());
BOOL NTAPI RtlInitializeLdrDataTableEntry(
_Out_ PLDR_DATA_TABLE_ENTRY LdrEntry,
_In_ DWORD dwFlags,
_In_ PVOID BaseAddress,
_In_ UNICODE_STRING& DllBaseName,
_In_ UNICODE_STRING& DllFullName) {
RtlZeroMemory(LdrEntry, MmpGlobalDataPtr->LdrDataTableEntrySize);
PIMAGE_NT_HEADERS headers = RtlImageNtHeader(BaseAddress);
if (!headers)return false;
if (!headers)return FALSE;
HANDLE heap = NtCurrentPeb()->ProcessHeap;
bool FlagsProcessed = false;
@@ -119,24 +119,24 @@ bool NTAPI RtlInitializeLdrDataTableEntry(
}
}
switch (NtWindowsVersion()) {
case win10:
case win10_1:
case win10_2: {
switch (MmpGlobalDataPtr->WindowsVersion) {
case WINDOWS_VERSION::win10:
case WINDOWS_VERSION::win10_1:
case WINDOWS_VERSION::win10_2: {
auto entry = (PLDR_DATA_TABLE_ENTRY_WIN10)LdrEntry;
entry->ReferenceCount = 1;
}
case win8:
case win8_1: {
case WINDOWS_VERSION::win8:
case WINDOWS_VERSION::win8_1: {
auto entry = (PLDR_DATA_TABLE_ENTRY_WIN8)LdrEntry;
const static bool IsWin8 = RtlIsWindowsVersionInScope(6, 2, 0, 6, 3, -1);
NtQuerySystemTime(&entry->LoadTime);
entry->OriginalBase = headers->OptionalHeader.ImageBase;
entry->BaseNameHashValue = LdrHashEntry(DllBaseName, false);
entry->LoadReason = LoadReasonDynamicLoad;
if (!NT_SUCCESS(RtlInsertModuleBaseAddressIndexNode(LdrEntry, BaseAddress)))return false;
if (!NT_SUCCESS(RtlInsertModuleBaseAddressIndexNode(LdrEntry, BaseAddress)))return FALSE;
if (!(entry->DdagNode = (decltype(entry->DdagNode))
RtlAllocateHeap(heap, HEAP_ZERO_MEMORY, IsWin8 ? sizeof(_LDR_DDAG_NODE_WIN8) : sizeof(_LDR_DDAG_NODE))))return false;
RtlAllocateHeap(heap, HEAP_ZERO_MEMORY, IsWin8 ? sizeof(_LDR_DDAG_NODE_WIN8) : sizeof(_LDR_DDAG_NODE))))return FALSE;
//RtlInitializeListEntry(&entry->NodeModuleLink);
//RtlInitializeListEntry(&entry->DdagNode->Modules);
//RtlInitializeSingleEntry(&entry->DdagNode->CondenseLink);
@@ -156,23 +156,23 @@ bool NTAPI RtlInitializeLdrDataTableEntry(
FlagsProcessed = true;
}
case win7: {
if (LdrpDataTableEntrySize() == sizeof(LDR_DATA_TABLE_ENTRY_WIN7)) {
case WINDOWS_VERSION::win7: {
if (MmpGlobalDataPtr->LdrDataTableEntrySize == sizeof(LDR_DATA_TABLE_ENTRY_WIN7)) {
auto entry = (PLDR_DATA_TABLE_ENTRY_WIN7)LdrEntry;
entry->OriginalBase = headers->OptionalHeader.ImageBase;
NtQuerySystemTime(&entry->LoadTime);
}
}
case vista: {
if (LdrpDataTableEntrySize() == sizeof(LDR_DATA_TABLE_ENTRY_VISTA) ||
LdrpDataTableEntrySize() == sizeof(LDR_DATA_TABLE_ENTRY_WIN7)) {
case WINDOWS_VERSION::vista: {
if (MmpGlobalDataPtr->LdrDataTableEntrySize == sizeof(LDR_DATA_TABLE_ENTRY_VISTA) ||
MmpGlobalDataPtr->LdrDataTableEntrySize == sizeof(LDR_DATA_TABLE_ENTRY_WIN7)) {
auto entry = (PLDR_DATA_TABLE_ENTRY_VISTA)LdrEntry;
RtlInitializeListEntry(&entry->ForwarderLinks);
RtlInitializeListEntry(&entry->StaticLinks);
RtlInitializeListEntry(&entry->ServiceTagLinks);
InitializeListHead(&entry->ForwarderLinks);
InitializeListHead(&entry->StaticLinks);
InitializeListHead(&entry->ServiceTagLinks);
}
}
case xp: {
case WINDOWS_VERSION::xp: {
LdrEntry->DllBase = BaseAddress;
LdrEntry->SizeOfImage = headers->OptionalHeader.SizeOfImage;
LdrEntry->TimeDateStamp = headers->FileHeader.TimeDateStamp;
@@ -184,30 +184,30 @@ bool NTAPI RtlInitializeLdrDataTableEntry(
LdrEntry->Flags = LDRP_IMAGE_DLL | LDRP_ENTRY_INSERTED | LDRP_ENTRY_PROCESSED | LDRP_PROCESS_ATTACH_CALLED;
if (CorImage)LdrEntry->Flags |= LDRP_COR_IMAGE;
}
RtlInitializeListEntry(&LdrEntry->HashLinks);
return true;
InitializeListHead(&LdrEntry->HashLinks);
return TRUE;
}
default:return false;
default:return FALSE;
}
}
bool NTAPI RtlFreeLdrDataTableEntry(IN PLDR_DATA_TABLE_ENTRY LdrEntry) {
BOOL NTAPI RtlFreeLdrDataTableEntry(_In_ PLDR_DATA_TABLE_ENTRY LdrEntry) {
HANDLE heap = NtCurrentPeb()->ProcessHeap;
switch (NtWindowsVersion()) {
case win10:
case win10_1:
case win10_2:
case win8:
case win8_1: {
switch (MmpGlobalDataPtr->WindowsVersion) {
case WINDOWS_VERSION::win10:
case WINDOWS_VERSION::win10_1:
case WINDOWS_VERSION::win10_2:
case WINDOWS_VERSION::win8:
case WINDOWS_VERSION::win8_1: {
auto entry = (PLDR_DATA_TABLE_ENTRY_WIN10)LdrEntry;
RtlFreeDependencies(entry);
RtlFreeHeap(heap, 0, entry->DdagNode);
RtlRemoveModuleBaseAddressIndexNode(LdrEntry);
}
case win7:
case vista: {
if (LdrpDataTableEntrySize() == sizeof(LDR_DATA_TABLE_ENTRY_VISTA) ||
LdrpDataTableEntrySize() == sizeof(LDR_DATA_TABLE_ENTRY_WIN7)) {
case WINDOWS_VERSION::win7:
case WINDOWS_VERSION::vista: {
if (MmpGlobalDataPtr->LdrDataTableEntrySize == sizeof(LDR_DATA_TABLE_ENTRY_VISTA) ||
MmpGlobalDataPtr->LdrDataTableEntrySize == sizeof(LDR_DATA_TABLE_ENTRY_WIN7)) {
PLDR_DATA_TABLE_ENTRY_VISTA entry = (decltype(entry))LdrEntry;
PLIST_ENTRY head = &entry->ForwarderLinks, next = head->Flink;
while (head != next) {
@@ -218,7 +218,7 @@ bool NTAPI RtlFreeLdrDataTableEntry(IN PLDR_DATA_TABLE_ENTRY LdrEntry) {
}
}
}
case xp: {
case WINDOWS_VERSION::xp: {
RtlFreeHeap(heap, 0, LdrEntry->BaseDllName.Buffer);
RtlFreeHeap(heap, 0, LdrEntry->FullDllName.Buffer);
RemoveEntryList(&LdrEntry->InLoadOrderLinks);
@@ -226,13 +226,15 @@ bool NTAPI RtlFreeLdrDataTableEntry(IN PLDR_DATA_TABLE_ENTRY LdrEntry) {
RemoveEntryList(&LdrEntry->InInitializationOrderLinks);
RemoveEntryList(&LdrEntry->HashLinks);
RtlFreeHeap(heap, 0, LdrEntry);
return true;
return TRUE;
}
default:return false;
default:return FALSE;
}
}
NTSTATUS NTAPI RtlUpdateReferenceCount(IN OUT PMEMORYMODULE pModule, IN DWORD Flags) {
NTSTATUS NTAPI RtlUpdateReferenceCount(
_Inout_ PMEMORYMODULE pModule,
_In_ DWORD Flags) {
if (Flags != FLAG_REFERENCE && Flags != FLAG_DEREFERENCE)return STATUS_INVALID_PARAMETER_2;
if (Flags == FLAG_REFERENCE && pModule->dwReferenceCount != 0xffffffff)
@@ -243,14 +245,16 @@ NTSTATUS NTAPI RtlUpdateReferenceCount(IN OUT PMEMORYMODULE pModule, IN DWORD Fl
return STATUS_SUCCESS;
}
NTSTATUS NTAPI RtlGetReferenceCount(IN PMEMORYMODULE pModule, OUT PULONG Count) {
NTSTATUS NTAPI RtlGetReferenceCount(
_In_ PMEMORYMODULE pModule,
_Out_ PULONG Count) {
*Count = pModule->dwReferenceCount;
return STATUS_SUCCESS;
}
VOID NTAPI RtlInsertMemoryTableEntry(IN PLDR_DATA_TABLE_ENTRY LdrEntry) {
VOID NTAPI RtlInsertMemoryTableEntry(_In_ PLDR_DATA_TABLE_ENTRY LdrEntry) {
PPEB_LDR_DATA PebData = NtCurrentPeb()->Ldr;
PLIST_ENTRY LdrpHashTable = MmpGlobalDataPtr->MmpLdrEntry.LdrpHashTable;
ULONG i;
@@ -265,19 +269,28 @@ VOID NTAPI RtlInsertMemoryTableEntry(IN PLDR_DATA_TABLE_ENTRY LdrEntry) {
InsertTailList(&PebData->InInitializationOrderModuleList, &LdrEntry->InInitializationOrderLinks);
}
VOID NTAPI RtlRbInsertNodeEx(IN PRTL_RB_TREE Tree, IN PRTL_BALANCED_NODE Parent, IN BOOLEAN Right, OUT PRTL_BALANCED_NODE Node) {
decltype(&RtlRbInsertNodeEx)_RtlRbInsertNodeEx = decltype(_RtlRbInsertNodeEx)(RtlGetNtProcAddress("RtlRbInsertNodeEx"));
VOID NTAPI RtlRbInsertNodeEx(
_In_ PRTL_RB_TREE Tree,
_In_ PRTL_BALANCED_NODE Parent,
_In_ BOOLEAN Right,
_Out_ PRTL_BALANCED_NODE Node) {
static decltype(&RtlRbInsertNodeEx)_RtlRbInsertNodeEx = decltype(_RtlRbInsertNodeEx)(RtlGetNtProcAddress("RtlRbInsertNodeEx"));
RtlZeroMemory(Node, sizeof(*Node));
if (!_RtlRbInsertNodeEx)return;
return _RtlRbInsertNodeEx(Tree, Parent, Right, Node);
}
VOID NTAPI RtlRbRemoveNode(IN PRTL_RB_TREE Tree, IN PRTL_BALANCED_NODE Node) {
decltype(&RtlRbRemoveNode)_RtlRbRemoveNode = decltype(_RtlRbRemoveNode)(RtlGetNtProcAddress("RtlRbRemoveNode"));
VOID NTAPI RtlRbRemoveNode(
_In_ PRTL_RB_TREE Tree,
_In_ PRTL_BALANCED_NODE Node) {
static decltype(&RtlRbRemoveNode)_RtlRbRemoveNode = decltype(_RtlRbRemoveNode)(RtlGetNtProcAddress("RtlRbRemoveNode"));
if (!_RtlRbRemoveNode)return;
return _RtlRbRemoveNode(Tree, Node);
}
PLDR_DATA_TABLE_ENTRY NTAPI RtlFindLdrTableEntryByHandle(PVOID BaseAddress) {
PLDR_DATA_TABLE_ENTRY NTAPI RtlFindLdrTableEntryByHandle(_In_ PVOID BaseAddress) {
PLIST_ENTRY ListHead = &NtCurrentPeb()->Ldr->InLoadOrderModuleList, ListEntry = ListHead->Flink;
PLDR_DATA_TABLE_ENTRY CurEntry;
while (ListEntry != ListHead) {
@@ -290,21 +303,21 @@ PLDR_DATA_TABLE_ENTRY NTAPI RtlFindLdrTableEntryByHandle(PVOID BaseAddress) {
return nullptr;
}
PLDR_DATA_TABLE_ENTRY NTAPI RtlFindLdrTableEntryByBaseName(PCWSTR BaseName) {
PLDR_DATA_TABLE_ENTRY NTAPI RtlFindLdrTableEntryByBaseName(_In_z_ PCWSTR BaseName) {
PLIST_ENTRY ListHead = &NtCurrentPeb()->Ldr->InLoadOrderModuleList, ListEntry = ListHead->Flink;
PLDR_DATA_TABLE_ENTRY CurEntry;
while (ListEntry != ListHead) {
CurEntry = CONTAINING_RECORD(ListEntry, LDR_DATA_TABLE_ENTRY, InLoadOrderLinks);
ListEntry = ListEntry->Flink;
if (!wcsnicmp(BaseName, CurEntry->BaseDllName.Buffer, (CurEntry->BaseDllName.Length / sizeof(wchar_t)) - 4) ||
!wcsnicmp(BaseName, CurEntry->BaseDllName.Buffer, CurEntry->BaseDllName.Length / sizeof(wchar_t))) {
if (!_wcsnicmp(BaseName, CurEntry->BaseDllName.Buffer, (CurEntry->BaseDllName.Length / sizeof(wchar_t)) - 4) ||
!_wcsnicmp(BaseName, CurEntry->BaseDllName.Buffer, CurEntry->BaseDllName.Length / sizeof(wchar_t))) {
return CurEntry;
}
}
return nullptr;
}
ULONG NTAPI LdrHashEntry(IN UNICODE_STRING& str, IN bool _xor) {
ULONG NTAPI LdrHashEntry(_In_ UNICODE_STRING& str, _In_ BOOL _xor) {
ULONG result = 0;
if (RtlIsWindowsVersionOrGreater(6, 2, 0)) {
RtlHashUnicodeString(&str, TRUE, HASH_STRING_ALGORITHM_DEFAULT, &result);
@@ -316,20 +329,3 @@ ULONG NTAPI LdrHashEntry(IN UNICODE_STRING& str, IN bool _xor) {
if (_xor)result &= (LDR_HASH_TABLE_ENTRIES - 1);
return result;
}
size_t NTAPI LdrpDataTableEntrySize() {
static size_t size = 0;
if (size)return size;
switch (NtWindowsVersion()) {
case xp:return size = sizeof(LDR_DATA_TABLE_ENTRY_XP);
case vista:return size = sizeof(LDR_DATA_TABLE_ENTRY_VISTA);
case win7:return size = sizeof(LDR_DATA_TABLE_ENTRY_WIN7);
case win8:return size = sizeof(LDR_DATA_TABLE_ENTRY_WIN8);
case win8_1:return size = sizeof(LDR_DATA_TABLE_ENTRY_WIN8_1);
case win10:return size = sizeof(LDR_DATA_TABLE_ENTRY_WIN10);
case win10_1:return size = sizeof(LDR_DATA_TABLE_ENTRY_WIN10_1);
case win10_2:return size = sizeof(LDR_DATA_TABLE_ENTRY_WIN10_2);
default:return size = sizeof(LDR_DATA_TABLE_ENTRY_WIN10_2);
}
}