From fda70c9beb7b10e01de89e24e338629bf2984f0a Mon Sep 17 00:00:00 2001 From: Boring <1079299053@qq.com> Date: Wed, 30 Nov 2022 08:20:11 +0800 Subject: [PATCH] Refactor the RtlFindMemoryBlockFromModuleSection function --- MemoryModule/Initialize.cpp | 14 ++--- MemoryModule/Utils.cpp | 112 ++++++++++++++++++++++-------------- MemoryModule/Utils.h | 26 +++------ 3 files changed, 85 insertions(+), 67 deletions(-) diff --git a/MemoryModule/Initialize.cpp b/MemoryModule/Initialize.cpp index edd6fe9..a4f845d 100644 --- a/MemoryModule/Initialize.cpp +++ b/MemoryModule/Initialize.cpp @@ -15,11 +15,11 @@ PRTL_RB_TREE FindLdrpModuleBaseAddressIndex() { BYTE count = 0; PRTL_RB_TREE tmp = nullptr; SEARCH_CONTEXT SearchContext{}; - SearchContext.MemoryBuffer = &node; - SearchContext.BufferLength = sizeof(size_t); + SearchContext.SearchPattern = (LPBYTE)&node; + SearchContext.PatternSize = sizeof(size_t); while (NT_SUCCESS(RtlFindMemoryBlockFromModuleSection((HMODULE)nt10->DllBase, ".data", &SearchContext))) { if (count++)return nullptr; - tmp = (decltype(tmp))SearchContext.MemoryBlockInSection; + tmp = (decltype(tmp))SearchContext.Result; } if (count && tmp && tmp->Root && tmp->Min) { LdrpModuleBaseAddressIndex = tmp; @@ -56,7 +56,7 @@ PVOID FindLdrpInvertedFunctionTable32() { PIMAGE_NT_HEADERS NtdllHeaders = RtlImageNtHeader(hNtdll), ModuleHeaders = nullptr; _RTL_INVERTED_FUNCTION_TABLE_ENTRY_WIN7_32 entry{}; LPCSTR lpSectionName = ".data"; - SEARCH_CONTEXT SearchContext{ SearchContext.MemoryBuffer = &entry,SearchContext.BufferLength = sizeof(entry) }; + SEARCH_CONTEXT SearchContext{ SearchContext.SearchPattern = (LPBYTE)&entry,SearchContext.PatternSize = sizeof(entry) }; PLIST_ENTRY ListHead = &NtCurrentPeb()->Ldr->InMemoryOrderModuleList, ListEntry = ListHead->Flink; PLDR_DATA_TABLE_ENTRY CurEntry = nullptr; @@ -81,7 +81,7 @@ PVOID FindLdrpInvertedFunctionTable32() { entry = { RtlEncodeSystemPointer((PVOID)SEHTable),(DWORD)hModule,ModuleHeaders->OptionalHeader.SizeOfImage,(PVOID)SEHCount }; while (NT_SUCCESS(RtlFindMemoryBlockFromModuleSection(hNtdll, lpSectionName, &SearchContext))) { - PRTL_INVERTED_FUNCTION_TABLE_WIN7_32 tab = decltype(tab)(SearchContext.OutBufferPtr - Offset); + PRTL_INVERTED_FUNCTION_TABLE_WIN7_32 tab = decltype(tab)(SearchContext.Result - Offset); //Note: Same memory layout for RTL_INVERTED_FUNCTION_TABLE_ENTRY in Windows 10 x86 and x64. if (RtlIsWindowsVersionOrGreater(6, 2, 0) && tab->MaxCount == 0x200 && !tab->NextEntrySEHandlerTableEncoded) return tab; @@ -111,7 +111,7 @@ PVOID FindLdrpInvertedFunctionTable64() { _RTL_INVERTED_FUNCTION_TABLE_ENTRY_64 entry{}; LPCSTR lpSectionName = ".data"; PIMAGE_DATA_DIRECTORY dir = nullptr; - SEARCH_CONTEXT SearchContext{ SearchContext.MemoryBuffer = &entry,SearchContext.BufferLength = sizeof(entry) }; + SEARCH_CONTEXT SearchContext{ SearchContext.SearchPattern = (LPBYTE)&entry,SearchContext.PatternSize = sizeof(entry) }; //Windows 8 if (RtlVerifyVersion(6, 2, 0, RTL_VERIFY_FLAGS_MAJOR_VERSION | RTL_VERIFY_FLAGS_MINOR_VERSION)) { @@ -147,7 +147,7 @@ PVOID FindLdrpInvertedFunctionTable64() { }; while (NT_SUCCESS(RtlFindMemoryBlockFromModuleSection(hNtdll, lpSectionName, &SearchContext))) { - PRTL_INVERTED_FUNCTION_TABLE_64 tab = decltype(tab)(SearchContext.OutBufferPtr - 0x10); + PRTL_INVERTED_FUNCTION_TABLE_64 tab = decltype(tab)(SearchContext.Result - 0x10); if (RtlIsWindowsVersionOrGreater(6, 2, 0) && tab->MaxCount == 0x200 && !tab->Overflow) return tab; else if (tab->MaxCount == 0x200 && !tab->Epoch) return tab; } diff --git a/MemoryModule/Utils.cpp b/MemoryModule/Utils.cpp index 040e815..74bf5ff 100644 --- a/MemoryModule/Utils.cpp +++ b/MemoryModule/Utils.cpp @@ -159,68 +159,96 @@ SIZE_T NTAPI _RtlCompareMemory( const VOID* Source1, const VOID* Source2, SIZE_T Length) { - return decltype(&_RtlCompareMemory)(GetProcAddress((HMODULE)MmpGlobalDataPtr->MmpBaseAddressIndex->NtdllLdrEntry->DllBase,"RtlCompareMemory"))(Source1, Source2, Length); + return decltype(&_RtlCompareMemory)(GetProcAddress((HMODULE)MmpGlobalDataPtr->MmpBaseAddressIndex->NtdllLdrEntry->DllBase, "RtlCompareMemory"))(Source1, Source2, Length); } #define RtlCompareMemory _RtlCompareMemory #endif NTSTATUS NTAPI RtlFindMemoryBlockFromModuleSection( - _In_ HMODULE hModule, - _In_ LPCSTR lpSectionName, + _In_ HMODULE ModuleHandle, + _In_ LPCSTR SectionName, _Inout_ PSEARCH_CONTEXT SearchContext) { NTSTATUS status = STATUS_SUCCESS; - size_t begin = 0, buffer = 0; - DWORD Length = 0, bufferLength = 0; __try { - begin = SearchContext->OutBufferPtr; - Length = SearchContext->RemainingLength; - buffer = SearchContext->InBufferPtr; - bufferLength = SearchContext->BufferLength; - if (!buffer || !bufferLength) { - SearchContext->OutBufferPtr = 0; - SearchContext->RemainingLength = 0; - return STATUS_INVALID_PARAMETER; + + // + // checks if no search pattern and length are provided + // + + if (!SearchContext->SearchPattern || !SearchContext->PatternSize) { + SearchContext->Result = nullptr; + SearchContext->MemoryBlockSize = 0; + + status = STATUS_INVALID_PARAMETER; + __leave; } - if (!begin) { - PIMAGE_NT_HEADERS headers = RtlImageNtHeader(hModule); - PIMAGE_SECTION_HEADER section = nullptr; - if (!headers)return STATUS_INVALID_PARAMETER_1; - section = IMAGE_FIRST_SECTION(headers); - for (WORD i = 0; i < headers->FileHeader.NumberOfSections; ++i) { - if (!_strnicmp(lpSectionName, (LPCSTR)section->Name, 8)) { - begin = SearchContext->OutBufferPtr = (size_t)hModule + section->VirtualAddress; - Length = SearchContext->RemainingLength = section->Misc.VirtualSize; - break; - } - ++section; - } - if (!begin || !Length || Length < bufferLength) { - SearchContext->OutBufferPtr = 0; - SearchContext->RemainingLength = 0; - return STATUS_NOT_FOUND; - } + + if (SearchContext->Result) { + ++SearchContext->Result; + --SearchContext->MemoryBlockSize; } else { - begin++; - Length--; - } - status = STATUS_NOT_FOUND; - for (DWORD i = 0; i < Length - bufferLength; ++begin, ++i) { - if (RtlCompareMemory((PVOID)begin, (PVOID)buffer, bufferLength) == bufferLength) { - SearchContext->OutBufferPtr = begin; - --SearchContext->RemainingLength; - return STATUS_SUCCESS; + + // + // if it is the first search, find the length and start address of the specified section + // + + PIMAGE_NT_HEADERS headers = RtlImageNtHeader(ModuleHandle); + PIMAGE_SECTION_HEADER section = nullptr; + + if (headers) { + section = IMAGE_FIRST_SECTION(headers); + for (WORD i = 0; i < headers->FileHeader.NumberOfSections; ++i) { + if (!_strnicmp(SectionName, (LPCSTR)section->Name, 8)) { + SearchContext->Result = (LPBYTE)ModuleHandle + section->VirtualAddress; + SearchContext->MemoryBlockSize = section->Misc.VirtualSize; + break; + } + + ++section; + } + + if (!SearchContext->Result || !SearchContext->MemoryBlockSize || SearchContext->MemoryBlockSize < SearchContext->PatternSize) { + SearchContext->Result = nullptr; + SearchContext->MemoryBlockSize = 0; + status = STATUS_NOT_FOUND; + __leave; + } + } + else { + status = STATUS_INVALID_PARAMETER_1; + __leave; } } + + // + // perform a linear search on the pattern + // + + LPBYTE end = SearchContext->Result + SearchContext->MemoryBlockSize - SearchContext->PatternSize; + while (SearchContext->Result <= end) { + if (RtlCompareMemory(SearchContext->SearchPattern, SearchContext->Result, SearchContext->PatternSize) == SearchContext->PatternSize) { + __leave; + } + + ++SearchContext->Result; + --SearchContext->MemoryBlockSize; + } + + // + // if the search fails, clear the output parameters + // + + SearchContext->Result = nullptr; + SearchContext->MemoryBlockSize = 0; + status = STATUS_NOT_FOUND; } __except (EXCEPTION_EXECUTE_HANDLER) { status = GetExceptionCode(); } - SearchContext->OutBufferPtr = 0; - SearchContext->RemainingLength = 0; return status; } diff --git a/MemoryModule/Utils.h b/MemoryModule/Utils.h index 21f4b0d..bd1a241 100644 --- a/MemoryModule/Utils.h +++ b/MemoryModule/Utils.h @@ -1,28 +1,18 @@ #pragma once typedef struct _SEARCH_CONTEXT { - union { - IN PVOID MemoryBuffer; - size_t InBufferPtr; - }; - union { - IN DWORD BufferLength; - size_t reserved0; - }; - union { - OUT PVOID MemoryBlockInSection; - size_t OutBufferPtr; - }; - union { - DWORD RemainingLength; - size_t reserved1; - }; + IN LPBYTE SearchPattern; + IN SIZE_T PatternSize; + + OUT LPBYTE Result; + SIZE_T MemoryBlockSize; + }SEARCH_CONTEXT, * PSEARCH_CONTEXT; NTSTATUS NTAPI RtlFindMemoryBlockFromModuleSection( - _In_ HMODULE hModule, - _In_ LPCSTR lpSectionName, + _In_ HMODULE ModuleHandle, + _In_ LPCSTR SectionName, _Inout_ PSEARCH_CONTEXT SearchContext );