Merge branch 'ImportTableResolver' into MmpTlsFiber

This commit is contained in:
Boring
2023-12-14 19:44:48 +08:00
11 changed files with 338 additions and 190 deletions
+1 -89
View File
@@ -85,86 +85,6 @@ NTSTATUS MmpInitializeStructure(DWORD ImageFileSize, LPCVOID ImageFileBuffer, PI
return STATUS_SUCCESS;
}
NTSTATUS MemoryResolveImportTable(
_In_ LPBYTE base,
_In_ PIMAGE_NT_HEADERS lpNtHeaders,
_In_ PMEMORYMODULE hMemoryModule) {
NTSTATUS status = STATUS_SUCCESS;
PIMAGE_IMPORT_DESCRIPTOR importDesc = nullptr;
DWORD count = 0;
do {
__try {
PIMAGE_DATA_DIRECTORY dir = GET_HEADER_DICTIONARY(lpNtHeaders, IMAGE_DIRECTORY_ENTRY_IMPORT);
PIMAGE_IMPORT_DESCRIPTOR iat = nullptr;
if (dir && dir->Size) {
iat = importDesc = PIMAGE_IMPORT_DESCRIPTOR(lpNtHeaders->OptionalHeader.ImageBase + dir->VirtualAddress);
}
if (iat) {
while (iat->Name) {
++count;
++iat;
}
}
if (importDesc && count) {
hMemoryModule->hModulesList = (HMODULE*)RtlAllocateHeap(NtCurrentPeb()->ProcessHeap, HEAP_ZERO_MEMORY, sizeof(HMODULE) * count);
if (!hMemoryModule->hModulesList) {
status = STATUS_NO_MEMORY;
break;
}
for (DWORD i = 0; i < count; ++i, ++importDesc) {
uintptr_t* thunkRef;
FARPROC* funcRef;
HMODULE handle = LoadLibraryA((LPCSTR)(base + importDesc->Name));
if (!handle) {
status = STATUS_DLL_NOT_FOUND;
break;
}
hMemoryModule->hModulesList[hMemoryModule->dwModulesCount++] = handle;
thunkRef = (uintptr_t*)(base + (importDesc->OriginalFirstThunk ? importDesc->OriginalFirstThunk : importDesc->FirstThunk));
funcRef = (FARPROC*)(base + importDesc->FirstThunk);
while (*thunkRef) {
*funcRef = GetProcAddress(
handle,
IMAGE_SNAP_BY_ORDINAL(*thunkRef) ? (LPCSTR)IMAGE_ORDINAL(*thunkRef) : (LPCSTR)PIMAGE_IMPORT_BY_NAME(base + (*thunkRef))->Name
);
if (!*funcRef) {
status = STATUS_ENTRYPOINT_NOT_FOUND;
break;
}
++thunkRef;
++funcRef;
}
if (!NT_SUCCESS(status))break;
}
}
}
__except (EXCEPTION_EXECUTE_HANDLER) {
status = GetExceptionCode();
}
} while (false);
if (!NT_SUCCESS(status)) {
for (DWORD i = 0; i < hMemoryModule->dwModulesCount; ++i)
FreeLibrary(hMemoryModule->hModulesList[i]);
RtlFreeHeap(NtCurrentPeb()->ProcessHeap, 0, hMemoryModule->hModulesList);
hMemoryModule->hModulesList = nullptr;
hMemoryModule->dwModulesCount = 0;
}
return status;
}
NTSTATUS MemorySetSectionProtection(
_In_ LPBYTE base,
_In_ PIMAGE_NT_HEADERS lpNtHeaders) {
@@ -404,15 +324,7 @@ BOOL MemoryFreeLibrary(HMEMORYMODULE mod) {
if (!module) return FALSE;
if (module->loadFromLdrLoadDllMemory && !module->underUnload)return FALSE;
if (module->hModulesList) {
for (DWORD i = 0; i < module->dwModulesCount; ++i) {
if (module->hModulesList[i]) {
FreeLibrary(module->hModulesList[i]);
}
}
RtlFreeHeap(NtCurrentPeb()->ProcessHeap, 0, module->hModulesList);
}
if (module->hModulesList)MemoryFreeImportTable(module);
if (module->codeBase) VirtualFree(mod, 0, MEM_RELEASE);
return TRUE;