Files
bb107-MemoryModulePP/MemoryModule/BaseAddressIndex.cpp
T
2022-10-11 07:47:31 +08:00

43 lines
1.8 KiB
C++

#include "stdafx.h"
NTSTATUS NTAPI RtlInsertModuleBaseAddressIndexNode(
_In_ PLDR_DATA_TABLE_ENTRY DataTableEntry,
_In_ PVOID BaseAddress) {
auto LdrpModuleBaseAddressIndex = MmpGlobalDataPtr->MmpBaseAddressIndex->LdrpModuleBaseAddressIndex;
if (!LdrpModuleBaseAddressIndex)return STATUS_UNSUCCESSFUL;
PLDR_DATA_TABLE_ENTRY_WIN8 LdrNode = decltype(LdrNode)((size_t)LdrpModuleBaseAddressIndex - offsetof(LDR_DATA_TABLE_ENTRY_WIN8, BaseAddressIndexNode));
bool bRight = false;
const auto i = offsetof(LDR_DATA_TABLE_ENTRY_WIN8, BaseAddressIndexNode);
while (true) {
if (BaseAddress < LdrNode->DllBase) {
if (!LdrNode->BaseAddressIndexNode.Left)break;
LdrNode = decltype(LdrNode)((size_t)LdrNode->BaseAddressIndexNode.Left - offsetof(LDR_DATA_TABLE_ENTRY_WIN8, BaseAddressIndexNode));
}
else if (BaseAddress > LdrNode->DllBase) {
if (!LdrNode->BaseAddressIndexNode.Right) {
bRight = true;
break;
}
LdrNode = decltype(LdrNode)((size_t)LdrNode->BaseAddressIndexNode.Right - offsetof(LDR_DATA_TABLE_ENTRY_WIN8, BaseAddressIndexNode));
}
else {
LdrNode->DdagNode->LoadCount++;
if (RtlIsWindowsVersionOrGreater(10, 0, 0)) {
PLDR_DATA_TABLE_ENTRY_WIN10(LdrNode)->ReferenceCount++;
}
return STATUS_SUCCESS;
}
}
RtlRbInsertNodeEx(LdrpModuleBaseAddressIndex, &LdrNode->BaseAddressIndexNode, bRight, &PLDR_DATA_TABLE_ENTRY_WIN8(DataTableEntry)->BaseAddressIndexNode);
return STATUS_SUCCESS;
}
NTSTATUS NTAPI RtlRemoveModuleBaseAddressIndexNode(_In_ PLDR_DATA_TABLE_ENTRY DataTableEntry) {
static auto tree{ MmpGlobalDataPtr->MmpBaseAddressIndex->LdrpModuleBaseAddressIndex };
if (!tree->Root)return STATUS_UNSUCCESSFUL;
RtlRbRemoveNode(tree, &PLDR_DATA_TABLE_ENTRY_WIN8(DataTableEntry)->BaseAddressIndexNode);
return STATUS_SUCCESS;
}