Skip to content

Commit fda70c9

Browse files
bb107bb107
committed
Refactor the RtlFindMemoryBlockFromModuleSection function
1 parent dcf5f3c commit fda70c9

3 files changed

Lines changed: 83 additions & 65 deletions

File tree

‎MemoryModule/Initialize.cpp‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -15,11 +15,11 @@ PRTL_RB_TREE FindLdrpModuleBaseAddressIndex() {
1515
BYTE count = 0;
1616
PRTL_RB_TREE tmp = nullptr;
1717
SEARCH_CONTEXT SearchContext{};
18-
SearchContext.MemoryBuffer = &node;
19-
SearchContext.BufferLength = sizeof(size_t);
18+
SearchContext.SearchPattern = (LPBYTE)&node;
19+
SearchContext.PatternSize = sizeof(size_t);
2020
while (NT_SUCCESS(RtlFindMemoryBlockFromModuleSection((HMODULE)nt10->DllBase, ".data", &SearchContext))) {
2121
if (count++)return nullptr;
22-
tmp = (decltype(tmp))SearchContext.MemoryBlockInSection;
22+
tmp = (decltype(tmp))SearchContext.Result;
2323
}
2424
if (count && tmp && tmp->Root && tmp->Min) {
2525
LdrpModuleBaseAddressIndex = tmp;
@@ -56,7 +56,7 @@ PVOID FindLdrpInvertedFunctionTable32() {
5656
PIMAGE_NT_HEADERS NtdllHeaders = RtlImageNtHeader(hNtdll), ModuleHeaders = nullptr;
5757
_RTL_INVERTED_FUNCTION_TABLE_ENTRY_WIN7_32 entry{};
5858
LPCSTR lpSectionName = ".data";
59-
SEARCH_CONTEXT SearchContext{ SearchContext.MemoryBuffer = &entry,SearchContext.BufferLength = sizeof(entry) };
59+
SEARCH_CONTEXT SearchContext{ SearchContext.SearchPattern = (LPBYTE)&entry,SearchContext.PatternSize = sizeof(entry) };
6060
PLIST_ENTRY ListHead = &NtCurrentPeb()->Ldr->InMemoryOrderModuleList,
6161
ListEntry = ListHead->Flink;
6262
PLDR_DATA_TABLE_ENTRY CurEntry = nullptr;
@@ -81,7 +81,7 @@ PVOID FindLdrpInvertedFunctionTable32() {
8181
entry = { RtlEncodeSystemPointer((PVOID)SEHTable),(DWORD)hModule,ModuleHeaders->OptionalHeader.SizeOfImage,(PVOID)SEHCount };
8282

8383
while (NT_SUCCESS(RtlFindMemoryBlockFromModuleSection(hNtdll, lpSectionName, &SearchContext))) {
84-
PRTL_INVERTED_FUNCTION_TABLE_WIN7_32 tab = decltype(tab)(SearchContext.OutBufferPtr - Offset);
84+
PRTL_INVERTED_FUNCTION_TABLE_WIN7_32 tab = decltype(tab)(SearchContext.Result - Offset);
8585

8686
//Note: Same memory layout for RTL_INVERTED_FUNCTION_TABLE_ENTRY in Windows 10 x86 and x64.
8787
if (RtlIsWindowsVersionOrGreater(6, 2, 0) && tab->MaxCount == 0x200 && !tab->NextEntrySEHandlerTableEncoded) return tab;
@@ -111,7 +111,7 @@ PVOID FindLdrpInvertedFunctionTable64() {
111111
_RTL_INVERTED_FUNCTION_TABLE_ENTRY_64 entry{};
112112
LPCSTR lpSectionName = ".data";
113113
PIMAGE_DATA_DIRECTORY dir = nullptr;
114-
SEARCH_CONTEXT SearchContext{ SearchContext.MemoryBuffer = &entry,SearchContext.BufferLength = sizeof(entry) };
114+
SEARCH_CONTEXT SearchContext{ SearchContext.SearchPattern = (LPBYTE)&entry,SearchContext.PatternSize = sizeof(entry) };
115115

116116
//Windows 8
117117
if (RtlVerifyVersion(6, 2, 0, RTL_VERIFY_FLAGS_MAJOR_VERSION | RTL_VERIFY_FLAGS_MINOR_VERSION)) {
@@ -147,7 +147,7 @@ PVOID FindLdrpInvertedFunctionTable64() {
147147
};
148148

149149
while (NT_SUCCESS(RtlFindMemoryBlockFromModuleSection(hNtdll, lpSectionName, &SearchContext))) {
150-
PRTL_INVERTED_FUNCTION_TABLE_64 tab = decltype(tab)(SearchContext.OutBufferPtr - 0x10);
150+
PRTL_INVERTED_FUNCTION_TABLE_64 tab = decltype(tab)(SearchContext.Result - 0x10);
151151
if (RtlIsWindowsVersionOrGreater(6, 2, 0) && tab->MaxCount == 0x200 && !tab->Overflow) return tab;
152152
else if (tab->MaxCount == 0x200 && !tab->Epoch) return tab;
153153
}

‎MemoryModule/Utils.cpp‎

Lines changed: 67 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -159,68 +159,96 @@ SIZE_T NTAPI _RtlCompareMemory(
159159
const VOID* Source1,
160160
const VOID* Source2,
161161
SIZE_T Length) {
162-
return decltype(&_RtlCompareMemory)(GetProcAddress((HMODULE)MmpGlobalDataPtr->MmpBaseAddressIndex->NtdllLdrEntry->DllBase,"RtlCompareMemory"))(Source1, Source2, Length);
162+
return decltype(&_RtlCompareMemory)(GetProcAddress((HMODULE)MmpGlobalDataPtr->MmpBaseAddressIndex->NtdllLdrEntry->DllBase, "RtlCompareMemory"))(Source1, Source2, Length);
163163
}
164164
#define RtlCompareMemory _RtlCompareMemory
165165
#endif
166166

167167
NTSTATUS NTAPI RtlFindMemoryBlockFromModuleSection(
168-
_In_ HMODULE hModule,
169-
_In_ LPCSTR lpSectionName,
168+
_In_ HMODULE ModuleHandle,
169+
_In_ LPCSTR SectionName,
170170
_Inout_ PSEARCH_CONTEXT SearchContext) {
171171

172172
NTSTATUS status = STATUS_SUCCESS;
173-
size_t begin = 0, buffer = 0;
174-
DWORD Length = 0, bufferLength = 0;
175173

176174
__try {
177-
begin = SearchContext->OutBufferPtr;
178-
Length = SearchContext->RemainingLength;
179-
buffer = SearchContext->InBufferPtr;
180-
bufferLength = SearchContext->BufferLength;
181-
if (!buffer || !bufferLength) {
182-
SearchContext->OutBufferPtr = 0;
183-
SearchContext->RemainingLength = 0;
184-
return STATUS_INVALID_PARAMETER;
175+
176+
//
177+
// checks if no search pattern and length are provided
178+
//
179+
180+
if (!SearchContext->SearchPattern || !SearchContext->PatternSize) {
181+
SearchContext->Result = nullptr;
182+
SearchContext->MemoryBlockSize = 0;
183+
184+
status = STATUS_INVALID_PARAMETER;
185+
__leave;
185186
}
186-
if (!begin) {
187-
PIMAGE_NT_HEADERS headers = RtlImageNtHeader(hModule);
187+
188+
if (SearchContext->Result) {
189+
++SearchContext->Result;
190+
--SearchContext->MemoryBlockSize;
191+
}
192+
else {
193+
194+
//
195+
// if it is the first search, find the length and start address of the specified section
196+
//
197+
198+
PIMAGE_NT_HEADERS headers = RtlImageNtHeader(ModuleHandle);
188199
PIMAGE_SECTION_HEADER section = nullptr;
189-
if (!headers)return STATUS_INVALID_PARAMETER_1;
190-
section = IMAGE_FIRST_SECTION(headers);
191-
for (WORD i = 0; i < headers->FileHeader.NumberOfSections; ++i) {
192-
if (!_strnicmp(lpSectionName, (LPCSTR)section->Name, 8)) {
193-
begin = SearchContext->OutBufferPtr = (size_t)hModule + section->VirtualAddress;
194-
Length = SearchContext->RemainingLength = section->Misc.VirtualSize;
195-
break;
200+
201+
if (headers) {
202+
section = IMAGE_FIRST_SECTION(headers);
203+
for (WORD i = 0; i < headers->FileHeader.NumberOfSections; ++i) {
204+
if (!_strnicmp(SectionName, (LPCSTR)section->Name, 8)) {
205+
SearchContext->Result = (LPBYTE)ModuleHandle + section->VirtualAddress;
206+
SearchContext->MemoryBlockSize = section->Misc.VirtualSize;
207+
break;
208+
}
209+
210+
++section;
211+
}
212+
213+
if (!SearchContext->Result || !SearchContext->MemoryBlockSize || SearchContext->MemoryBlockSize < SearchContext->PatternSize) {
214+
SearchContext->Result = nullptr;
215+
SearchContext->MemoryBlockSize = 0;
216+
status = STATUS_NOT_FOUND;
217+
__leave;
196218
}
197-
++section;
198219
}
199-
if (!begin || !Length || Length < bufferLength) {
200-
SearchContext->OutBufferPtr = 0;
201-
SearchContext->RemainingLength = 0;
202-
return STATUS_NOT_FOUND;
220+
else {
221+
status = STATUS_INVALID_PARAMETER_1;
222+
__leave;
203223
}
204224
}
205-
else {
206-
begin++;
207-
Length--;
208-
}
209-
status = STATUS_NOT_FOUND;
210-
for (DWORD i = 0; i < Length - bufferLength; ++begin, ++i) {
211-
if (RtlCompareMemory((PVOID)begin, (PVOID)buffer, bufferLength) == bufferLength) {
212-
SearchContext->OutBufferPtr = begin;
213-
--SearchContext->RemainingLength;
214-
return STATUS_SUCCESS;
225+
226+
//
227+
// perform a linear search on the pattern
228+
//
229+
230+
LPBYTE end = SearchContext->Result + SearchContext->MemoryBlockSize - SearchContext->PatternSize;
231+
while (SearchContext->Result <= end) {
232+
if (RtlCompareMemory(SearchContext->SearchPattern, SearchContext->Result, SearchContext->PatternSize) == SearchContext->PatternSize) {
233+
__leave;
215234
}
235+
236+
++SearchContext->Result;
237+
--SearchContext->MemoryBlockSize;
216238
}
239+
240+
//
241+
// if the search fails, clear the output parameters
242+
//
243+
244+
SearchContext->Result = nullptr;
245+
SearchContext->MemoryBlockSize = 0;
246+
status = STATUS_NOT_FOUND;
217247
}
218248
__except (EXCEPTION_EXECUTE_HANDLER) {
219249
status = GetExceptionCode();
220250
}
221251

222-
SearchContext->OutBufferPtr = 0;
223-
SearchContext->RemainingLength = 0;
224252
return status;
225253
}
226254

‎MemoryModule/Utils.h‎

Lines changed: 9 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -1,28 +1,18 @@
11
#pragma once
22

33
typedef struct _SEARCH_CONTEXT {
4-
union {
5-
IN PVOID MemoryBuffer;
6-
size_t InBufferPtr;
7-
};
8-
union {
9-
IN DWORD BufferLength;
10-
size_t reserved0;
11-
};
12-
13-
union {
14-
OUT PVOID MemoryBlockInSection;
15-
size_t OutBufferPtr;
16-
};
17-
union {
18-
DWORD RemainingLength;
19-
size_t reserved1;
20-
};
4+
5+
IN LPBYTE SearchPattern;
6+
IN SIZE_T PatternSize;
7+
8+
OUT LPBYTE Result;
9+
SIZE_T MemoryBlockSize;
10+
2111
}SEARCH_CONTEXT, * PSEARCH_CONTEXT;
2212

2313
NTSTATUS NTAPI RtlFindMemoryBlockFromModuleSection(
24-
_In_ HMODULE hModule,
25-
_In_ LPCSTR lpSectionName,
14+
_In_ HMODULE ModuleHandle,
15+
_In_ LPCSTR SectionName,
2616
_Inout_ PSEARCH_CONTEXT SearchContext
2717
);
2818

0 commit comments

Comments
 (0)