Merge pull request #15 from bb107/initialize

Initialize

Author: bb107

MemoryModule/BaseAddressIndex.cpp

@@ -1,35 +1,9 @@
 #include "stdafx.h"
 
-PRTL_RB_TREE NTAPI RtlFindLdrpModuleBaseAddressIndex() {
-	static PRTL_RB_TREE LdrpModuleBaseAddressIndex = nullptr;
-	if (LdrpModuleBaseAddressIndex)return LdrpModuleBaseAddressIndex;
-
-	PLDR_DATA_TABLE_ENTRY_WIN10 nt10 = decltype(nt10)(RtlFindNtdllLdrEntry());
-	PRTL_BALANCED_NODE node = nullptr;
-	if (!nt10 || !RtlIsWindowsVersionOrGreater(6, 2, 0))return nullptr;
-	node = &nt10->BaseAddressIndexNode;
-	while (node->ParentValue & (~7)) node = decltype(node)(node->ParentValue & (~7));
-
-	if (!node->Red) {
-		BYTE count = 0;
-		PRTL_RB_TREE tmp = nullptr;
-		SEARCH_CONTEXT SearchContext{};
-		SearchContext.MemoryBuffer = &node;
-		SearchContext.BufferLength = sizeof(size_t);
-		while (NT_SUCCESS(RtlFindMemoryBlockFromModuleSection((HMODULE)nt10->DllBase, ".data", &SearchContext))) {
-			if (count++)return nullptr;
-			tmp = (decltype(tmp))SearchContext.MemoryBlockInSection;
-		}
-		if (count && tmp && tmp->Root && tmp->Min) {
-			LdrpModuleBaseAddressIndex = tmp;
-		}
-	}
-
-	return LdrpModuleBaseAddressIndex;
-}
-
-NTSTATUS NTAPI RtlInsertModuleBaseAddressIndexNode(IN PLDR_DATA_TABLE_ENTRY DataTableEntry, IN PVOID BaseAddress) {
-	static auto LdrpModuleBaseAddressIndex = RtlFindLdrpModuleBaseAddressIndex();
+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));
@@ -60,8 +34,8 @@ NTSTATUS NTAPI RtlInsertModuleBaseAddressIndexNode(IN PLDR_DATA_TABLE_ENTRY Data
 	return STATUS_SUCCESS;
 }
 
-NTSTATUS NTAPI RtlRemoveModuleBaseAddressIndexNode(IN PLDR_DATA_TABLE_ENTRY DataTableEntry) {
-	static auto tree{ RtlFindLdrpModuleBaseAddressIndex() };
+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;

MemoryModule/BaseAddressIndex.h

@@ -1,7 +1,8 @@
 #pragma once
 
-PRTL_RB_TREE NTAPI RtlFindLdrpModuleBaseAddressIndex();
+NTSTATUS NTAPI RtlInsertModuleBaseAddressIndexNode(
+	_In_ PLDR_DATA_TABLE_ENTRY DataTableEntry,
+	_In_ PVOID BaseAddress
+);
 
-NTSTATUS NTAPI RtlInsertModuleBaseAddressIndexNode(IN PLDR_DATA_TABLE_ENTRY DataTableEntry, IN PVOID BaseAddress);
-
-NTSTATUS NTAPI RtlRemoveModuleBaseAddressIndexNode(IN PLDR_DATA_TABLE_ENTRY DataTableEntry);
+NTSTATUS NTAPI RtlRemoveModuleBaseAddressIndexNode(_In_ PLDR_DATA_TABLE_ENTRY DataTableEntry);