ImportTableResolver

MemoryModule/ImportTable.cpp

@@ -0,0 +1,178 @@
+#include "stdafx.h"
+
+typedef struct _MMP_IAT_HANDLE {
+
+	HMODULE hModule;
+	PMM_IAT_RESOLVER lpResolver;
+
+}MMP_IAT_HANDLE, * PMMP_IAT_HANDLE;
+
+HMODULE MmpLoadLibraryA(
+	_In_ LPCSTR lpModuleName,
+	_Out_ PMM_IAT_RESOLVER* lpModuleResolver) {
+
+	HMODULE hModule = nullptr;
+	PMM_IAT_RESOLVER resolver = nullptr;
+
+	EnterCriticalSection(&MmpGlobalDataPtr->MmpIat->MmpIatResolverListLock);
+
+	PLIST_ENTRY lpResolver = MmpGlobalDataPtr->MmpIat->MmpIatResolverList.Flink;
+	while (lpResolver != &MmpGlobalDataPtr->MmpIat->MmpIatResolverList) {
+		PMM_IAT_RESOLVER entry = CONTAINING_RECORD(lpResolver, MM_IAT_RESOLVER, MM_IAT_RESOLVER::InMmpIatResolverList);
+
+		hModule = entry->LoadLibraryProv(lpModuleName);
+		if (hModule) {
+			resolver = entry;
+			++entry->ReferenceCount;
+			break;
+		}
+
+		lpResolver = lpResolver->Flink;
+	}
+
+	LeaveCriticalSection(&MmpGlobalDataPtr->MmpIat->MmpIatResolverListLock);
+
+	*lpModuleResolver = resolver;
+	return hModule;
+}
+
+VOID MemoryFreeImportTable(_In_ PMEMORYMODULE hMemoryModule) {
+
+	EnterCriticalSection(&MmpGlobalDataPtr->MmpIat->MmpIatResolverListLock);
+
+	PMMP_IAT_HANDLE list = (PMMP_IAT_HANDLE)hMemoryModule->hModulesList;
+	for (DWORD i = 0; i < hMemoryModule->dwModulesCount; ++i) {
+		auto entry = list[i];
+		entry.lpResolver->FreeLibraryProv(entry.hModule);
+		--entry.lpResolver->ReferenceCount;
+	}
+
+	LeaveCriticalSection(&MmpGlobalDataPtr->MmpIat->MmpIatResolverListLock);
+
+
+	RtlFreeHeap(NtCurrentPeb()->ProcessHeap, 0, hMemoryModule->hModulesList);
+	hMemoryModule->hModulesList = nullptr;
+	hMemoryModule->dwModulesCount = 0;
+}
+
+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 = &lpNtHeaders->OptionalHeader.DataDirectory[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) {
+				PMMP_IAT_HANDLE handles = (PMMP_IAT_HANDLE)RtlAllocateHeap(NtCurrentPeb()->ProcessHeap, HEAP_ZERO_MEMORY, sizeof(MMP_IAT_HANDLE) * count);
+				hMemoryModule->hModulesList = handles;
+				if (!hMemoryModule->hModulesList) {
+					status = STATUS_NO_MEMORY;
+					break;
+				}
+
+				for (DWORD i = 0; i < count; ++i, ++importDesc) {
+					uintptr_t* thunkRef;
+					FARPROC* funcRef;
+					PMM_IAT_RESOLVER resolver;
+					HMODULE handle = MmpLoadLibraryA((LPCSTR)(base + importDesc->Name), &resolver);
+
+					if (!handle) {
+						status = STATUS_DLL_NOT_FOUND;
+						break;
+					}
+
+					handles[hMemoryModule->dwModulesCount].hModule = handle;
+					handles[hMemoryModule->dwModulesCount++].lpResolver = resolver;
+					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)) {
+		MemoryFreeImportTable(hMemoryModule);
+	}
+
+	return status;
+}
+
+HANDLE WINAPI MmRegisterImportTableResolver(
+	_In_ MM_IAT_RESOLVER_ENTRY LoadLibraryProv,
+	_In_ MM_IAT_FREE_ENTRY FreeLibraryProv) {
+
+	HANDLE heap = RtlProcessHeap();
+	PMM_IAT_RESOLVER resolver = (PMM_IAT_RESOLVER)RtlAllocateHeap(heap, 0, sizeof(MM_IAT_RESOLVER));
+
+	if (resolver) {
+		EnterCriticalSection(&MmpGlobalDataPtr->MmpIat->MmpIatResolverListLock);
+
+		resolver->ReferenceCount = 1;
+		resolver->LoadLibraryProv = LoadLibraryProv;
+		resolver->FreeLibraryProv = FreeLibraryProv;
+		InsertTailList(&MmpGlobalDataPtr->MmpIat->MmpIatResolverList, &resolver->InMmpIatResolverList);
+
+		LeaveCriticalSection(&MmpGlobalDataPtr->MmpIat->MmpIatResolverListLock);
+	}
+
+	return resolver;
+}
+
+_Success_(return)
+BOOL WINAPI MmRemoveImportTableResolver(_In_ HANDLE hMmIatResolver) {
+
+	HANDLE heap = RtlProcessHeap();
+
+	if (hMmIatResolver == &MmpGlobalDataPtr->MmpIat->MmpIatResolverHead) {
+		return FALSE;
+	}
+
+	PMM_IAT_RESOLVER resolver = CONTAINING_RECORD(hMmIatResolver, MM_IAT_RESOLVER, MM_IAT_RESOLVER::InMmpIatResolverList);
+
+	EnterCriticalSection(&MmpGlobalDataPtr->MmpIat->MmpIatResolverListLock);
+
+	if (resolver->ReferenceCount > 1) {
+		LeaveCriticalSection(&MmpGlobalDataPtr->MmpIat->MmpIatResolverListLock);
+		return FALSE;
+	}
+
+	RemoveHeadList(&resolver->InMmpIatResolverList);
+	LeaveCriticalSection(&MmpGlobalDataPtr->MmpIat->MmpIatResolverListLock);
+
+	return RtlFreeHeap(heap, 0, hMmIatResolver);
+}