Skip to content

Commit 446bb91

Browse files
bb107bb107
committed
Resolve IAT after insert LdrDataTableEntry
1 parent a74d29e commit 446bb91

5 files changed

Lines changed: 230 additions & 239 deletions

File tree

‎MemoryModule/MemoryModule.cpp‎

Lines changed: 124 additions & 115 deletions
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,129 @@ static SYSTEM_INFO sysInfo = []()->SYSTEM_INFO {
6262
return tmp;
6363
}();
6464

65+
NTSTATUS MemoryResolveImportTable(
66+
_In_ LPBYTE base,
67+
_In_ PIMAGE_NT_HEADERS lpNtHeaders,
68+
_In_ PMEMORYMODULE hMemoryModule) {
69+
NTSTATUS status = STATUS_SUCCESS;
70+
PIMAGE_IMPORT_DESCRIPTOR importDesc = nullptr;
71+
DWORD count = 0;
72+
73+
do {
74+
__try {
75+
PIMAGE_DATA_DIRECTORY dir = GET_HEADER_DICTIONARY(lpNtHeaders, IMAGE_DIRECTORY_ENTRY_IMPORT);
76+
PIMAGE_IMPORT_DESCRIPTOR iat = nullptr;
77+
78+
if (dir && dir->Size) {
79+
iat = importDesc = PIMAGE_IMPORT_DESCRIPTOR(lpNtHeaders->OptionalHeader.ImageBase + dir->VirtualAddress);
80+
}
81+
82+
if (iat) {
83+
while (iat->Name) {
84+
++count;
85+
++iat;
86+
}
87+
}
88+
89+
if (importDesc && count) {
90+
if (!(hMemoryModule->hModulesList = new HMODULE[count])) {
91+
status = STATUS_NO_MEMORY;
92+
break;
93+
}
94+
95+
RtlZeroMemory(
96+
hMemoryModule->hModulesList,
97+
sizeof(HMODULE) * count
98+
);
99+
100+
for (DWORD i = 0; i < count; ++i, ++importDesc) {
101+
uintptr_t* thunkRef;
102+
FARPROC* funcRef;
103+
HMODULE handle = LoadLibraryA((LPCSTR)(base + importDesc->Name));
104+
105+
if (!handle) {
106+
status = STATUS_DLL_NOT_FOUND;
107+
break;
108+
}
109+
110+
hMemoryModule->hModulesList[hMemoryModule->dwModulesCount++] = handle;
111+
thunkRef = (uintptr_t*)(base + (importDesc->OriginalFirstThunk ? importDesc->OriginalFirstThunk : importDesc->FirstThunk));
112+
funcRef = (FARPROC*)(base + importDesc->FirstThunk);
113+
while (*thunkRef) {
114+
*funcRef = GetProcAddress(
115+
handle,
116+
IMAGE_SNAP_BY_ORDINAL(*thunkRef) ? (LPCSTR)IMAGE_ORDINAL(*thunkRef) : (LPCSTR)PIMAGE_IMPORT_BY_NAME(base + (*thunkRef))->Name
117+
);
118+
if (!*funcRef) {
119+
status = STATUS_ENTRYPOINT_NOT_FOUND;
120+
break;
121+
}
122+
++thunkRef;
123+
++funcRef;
124+
}
125+
126+
if (!NT_SUCCESS(status))break;
127+
}
128+
129+
}
130+
}
131+
__except (EXCEPTION_EXECUTE_HANDLER) {
132+
status = GetExceptionCode();
133+
}
134+
} while (false);
135+
136+
if (!NT_SUCCESS(status)) {
137+
for (DWORD i = 0; i < hMemoryModule->dwModulesCount; ++i)
138+
FreeLibrary(hMemoryModule->hModulesList[i]);
139+
140+
delete[]hMemoryModule->hModulesList;
141+
hMemoryModule->hModulesList = nullptr;
142+
hMemoryModule->dwModulesCount = 0;
143+
}
144+
145+
return status;
146+
}
147+
148+
NTSTATUS MemorySetSectionProtection(
149+
_In_ LPBYTE base,
150+
_In_ PIMAGE_NT_HEADERS lpNtHeaders) {
151+
NTSTATUS status = STATUS_SUCCESS;
152+
PIMAGE_SECTION_HEADER section = IMAGE_FIRST_SECTION(lpNtHeaders);
153+
154+
//
155+
// Determine whether it is a .NET assembly
156+
//
157+
auto& com = lpNtHeaders->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_COM_DESCRIPTOR];
158+
bool CorImage = com.Size && com.VirtualAddress;
159+
160+
for (DWORD i = 0; i < lpNtHeaders->FileHeader.NumberOfSections; ++i, ++section) {
161+
LPVOID address = LPBYTE(base) + section->VirtualAddress;
162+
SIZE_T size = AlignValueUp(section->Misc.VirtualSize, lpNtHeaders->OptionalHeader.SectionAlignment);
163+
164+
if (section->Characteristics & IMAGE_SCN_MEM_DISCARDABLE && !CorImage) {
165+
//
166+
// If it is a .NET assembly, we cannot release this memory block
167+
//
168+
#pragma warning(disable:6250)
169+
VirtualFree(address, size, MEM_DECOMMIT);
170+
#pragma warning(default:6250)
171+
}
172+
else {
173+
BOOL executable = (section->Characteristics & IMAGE_SCN_MEM_EXECUTE) != 0,
174+
readable = (section->Characteristics & IMAGE_SCN_MEM_READ) != 0,
175+
writeable = (section->Characteristics & IMAGE_SCN_MEM_WRITE) != 0;
176+
DWORD protect = ProtectionFlags[executable][readable][writeable], oldProtect;
177+
178+
if (section->Characteristics & IMAGE_SCN_MEM_NOT_CACHED) protect |= PAGE_NOCACHE;
179+
180+
status = NtProtectVirtualMemory(NtCurrentProcess(), &address, &size, protect, &oldProtect);
181+
if (!NT_SUCCESS(status))break;
182+
}
183+
}
184+
185+
return status;
186+
}
187+
65188
NTSTATUS MemoryLoadLibrary(
66189
_Out_ HMEMORYMODULE* MemoryModuleHandle,
67190
_In_ LPCVOID data,
@@ -98,12 +221,6 @@ NTSTATUS MemoryLoadLibrary(
98221
__leave;
99222
}
100223

101-
//
102-
// Determine whether it is a .NET assembly
103-
//
104-
auto& com = old_header->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_COM_DESCRIPTOR];
105-
CorImage = com.Size && com.VirtualAddress;
106-
107224
//
108225
// Match machine type
109226
//
@@ -183,6 +300,7 @@ NTSTATUS MemoryLoadLibrary(
183300
hMemoryModule->Signature = MEMORY_MODULE_SIGNATURE;
184301
hMemoryModule->SizeofHeaders = old_header->OptionalHeader.SizeOfHeaders;
185302
hMemoryModule->lpReserved = (LPVOID)data;
303+
hMemoryModule->dwReferenceCount = 1;
186304

187305
do {
188306
//
@@ -267,115 +385,6 @@ NTSTATUS MemoryLoadLibrary(
267385
}
268386
if (!NT_SUCCESS(status))break;
269387

270-
//
271-
// Build import table
272-
//
273-
PIMAGE_IMPORT_DESCRIPTOR importDesc = nullptr;
274-
DWORD count = 0;
275-
276-
__try {
277-
PIMAGE_DATA_DIRECTORY dir = GET_HEADER_DICTIONARY(new_header, IMAGE_DIRECTORY_ENTRY_IMPORT);
278-
PIMAGE_IMPORT_DESCRIPTOR iat = nullptr;
279-
280-
status = STATUS_SUCCESS;
281-
282-
if (dir && dir->Size) {
283-
iat = importDesc = PIMAGE_IMPORT_DESCRIPTOR(new_header->OptionalHeader.ImageBase + dir->VirtualAddress);
284-
}
285-
286-
if (iat) {
287-
while (iat->Name) {
288-
++count;
289-
++iat;
290-
}
291-
}
292-
293-
if (importDesc && count) {
294-
if (!(hMemoryModule->hModulesList = new HMODULE[count])) {
295-
status = STATUS_NO_MEMORY;
296-
break;
297-
}
298-
299-
RtlZeroMemory(
300-
hMemoryModule->hModulesList,
301-
sizeof(HMODULE) * count
302-
);
303-
304-
for (DWORD i = 0; i < count; ++i, ++importDesc) {
305-
uintptr_t* thunkRef;
306-
FARPROC* funcRef;
307-
HMODULE handle = LoadLibraryA((LPCSTR)(base + importDesc->Name));
308-
309-
if (!handle) {
310-
status = STATUS_DLL_NOT_FOUND;
311-
break;
312-
}
313-
314-
hMemoryModule->hModulesList[hMemoryModule->dwModulesCount++] = handle;
315-
thunkRef = (uintptr_t*)(base + (importDesc->OriginalFirstThunk ? importDesc->OriginalFirstThunk : importDesc->FirstThunk));
316-
funcRef = (FARPROC*)(base + importDesc->FirstThunk);
317-
while (*thunkRef) {
318-
*funcRef = GetProcAddress(
319-
handle,
320-
IMAGE_SNAP_BY_ORDINAL(*thunkRef) ? (LPCSTR)IMAGE_ORDINAL(*thunkRef) : (LPCSTR)PIMAGE_IMPORT_BY_NAME(base + (*thunkRef))->Name
321-
);
322-
if (!*funcRef) {
323-
status = STATUS_ENTRYPOINT_NOT_FOUND;
324-
break;
325-
}
326-
++thunkRef;
327-
++funcRef;
328-
}
329-
330-
if (!NT_SUCCESS(status))break;
331-
}
332-
333-
}
334-
}
335-
__except (EXCEPTION_EXECUTE_HANDLER) {
336-
status = GetExceptionCode();
337-
}
338-
if (!NT_SUCCESS(status)) {
339-
for (DWORD i = 0; i < hMemoryModule->dwModulesCount; ++i)
340-
FreeLibrary(hMemoryModule->hModulesList[i]);
341-
342-
delete[]hMemoryModule->hModulesList;
343-
hMemoryModule->hModulesList = nullptr;
344-
hMemoryModule->dwModulesCount = 0;
345-
346-
break;
347-
}
348-
349-
//
350-
// Set section memory protect
351-
//
352-
section = IMAGE_FIRST_SECTION(new_header);
353-
for (DWORD i = 0; i < new_header->FileHeader.NumberOfSections; ++i, ++section) {
354-
LPVOID address = LPBYTE(base) + section->VirtualAddress;
355-
SIZE_T size = AlignValueUp(section->Misc.VirtualSize, new_header->OptionalHeader.SectionAlignment);
356-
357-
if (section->Characteristics & IMAGE_SCN_MEM_DISCARDABLE && !CorImage) {
358-
//
359-
// If it is a .NET assembly, we cannot release this memory block
360-
//
361-
#pragma warning(disable:6250)
362-
VirtualFree(address, size, MEM_DECOMMIT);
363-
#pragma warning(default:6250)
364-
}
365-
else {
366-
BOOL executable = (section->Characteristics & IMAGE_SCN_MEM_EXECUTE) != 0,
367-
readable = (section->Characteristics & IMAGE_SCN_MEM_READ) != 0,
368-
writeable = (section->Characteristics & IMAGE_SCN_MEM_WRITE) != 0;
369-
DWORD protect = ProtectionFlags[executable][readable][writeable], oldProtect;
370-
371-
if (section->Characteristics & IMAGE_SCN_MEM_NOT_CACHED) protect |= PAGE_NOCACHE;
372-
373-
status = NtProtectVirtualMemory(NtCurrentProcess(), &address, &size, protect, &oldProtect);
374-
if (!NT_SUCCESS(status))break;
375-
}
376-
}
377-
if (!NT_SUCCESS(status))break;
378-
379388
__try {
380389
*MemoryModuleHandle = (HMEMORYMODULE)base;
381390
}

‎MemoryModule/MemoryModule.h‎

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -50,7 +50,7 @@ typedef struct _MEMORYMODULE {
5050

5151
HMODULE* hModulesList; //Import module handles
5252
DWORD dwModulesCount; //number of module handles
53-
DWORD dwReserved;
53+
DWORD dwReferenceCount;
5454

5555
DWORD dwImageFileSize;
5656
DWORD headers_align; //headers_align == OptionalHeaders.BaseOfCode;
@@ -70,6 +70,17 @@ extern "C" {
7070
_In_ DWORD size
7171
);
7272

73+
NTSTATUS MemoryResolveImportTable(
74+
_In_ LPBYTE base,
75+
_In_ PIMAGE_NT_HEADERS lpNtHeaders,
76+
_In_ PMEMORYMODULE hMemoryModule
77+
);
78+
79+
NTSTATUS MemorySetSectionProtection(
80+
_In_ LPBYTE base,
81+
_In_ PIMAGE_NT_HEADERS lpNtHeaders
82+
);
83+
7384
bool MemoryFreeLibrary(HMEMORYMODULE);
7485

7586
bool WINAPI IsValidMemoryModuleHandle(HMEMORYMODULE hModule);

‎MemoryModule/MmpTls.cpp‎

Lines changed: 0 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -297,8 +297,6 @@ NTSTATUS NTAPI HookNtCreateThread(
297297
Context.Rdx = ULONG64(_Context);
298298
#endif
299299

300-
EnterCriticalSection(&MmpTlspLock);
301-
302300
status = OriginNtCreateThread(
303301
ThreadHandle,
304302
DesiredAccess,
@@ -313,8 +311,6 @@ NTSTATUS NTAPI HookNtCreateThread(
313311
RtlFreeHeap(RtlProcessHeap(), 0, _Context);
314312
}
315313

316-
LeaveCriticalSection(&MmpTlspLock);
317-
318314
return status;
319315
}
320316

@@ -338,8 +334,6 @@ NTSTATUS NTAPI HookNtCreateThreadEx(
338334
Context->ThreadStartRoutine = PTHREAD_START_ROUTINE(StartRoutine);
339335
Context->ThreadParameter = Argument;
340336

341-
EnterCriticalSection(&MmpTlspLock);
342-
343337
NTSTATUS status = OriginNtCreateThreadEx(
344338
ThreadHandle,
345339
DesiredAccess,
@@ -357,8 +351,6 @@ NTSTATUS NTAPI HookNtCreateThreadEx(
357351
RtlFreeHeap(RtlProcessHeap(), 0, Context);
358352
}
359353

360-
LeaveCriticalSection(&MmpTlspLock);
361-
362354
return status;
363355
}
364356

0 commit comments

Comments
 (0)