ImportTable.cpp/h - Finish support for resolver callbacks

// flag to prepend or append the resolver list
MmRegisterImportTableResolver

// new function to clear resolver list
MmClearImportTableResolvers

// expose function to remove a resolver
MmRemoveImportTableResolver

Author: Dan Dees

MemoryModule/ImportTable.cpp

@@ -132,9 +132,10 @@ NTSTATUS MemoryResolveImportTable(
 	return status;
 }
 
-HANDLE WINAPI MmRegisterImportTableResolver(
+extern "C" HANDLE WINAPI MmRegisterImportTableResolver(
 	_In_ MM_IAT_RESOLVER_ENTRY LoadLibraryProv,
-	_In_ MM_IAT_FREE_ENTRY FreeLibraryProv) {
+	_In_ MM_IAT_FREE_ENTRY FreeLibraryProv,
+	_In_ int do_append) {
 
 	HANDLE heap = RtlProcessHeap();
 	PMM_IAT_RESOLVER resolver = static_cast<PMM_IAT_RESOLVER>(RtlAllocateHeap(heap, 0, sizeof(MM_IAT_RESOLVER)));
@@ -145,15 +146,20 @@ HANDLE WINAPI MmRegisterImportTableResolver(
 		resolver->ReferenceCount = 1;
 		resolver->LoadLibraryProv = LoadLibraryProv;
 		resolver->FreeLibraryProv = FreeLibraryProv;
-		InsertTailList(&MmpGlobalDataPtr->MmpIat->MmpIatResolverList, &resolver->InMmpIatResolverList);
+        if (do_append) {
+            InsertTailList(&MmpGlobalDataPtr->MmpIat->MmpIatResolverList, &resolver->InMmpIatResolverList);
+        }
+        else {
+            InsertHeadList(&MmpGlobalDataPtr->MmpIat->MmpIatResolverList, &resolver->InMmpIatResolverList);
+        }
 
 		LeaveCriticalSection(&MmpGlobalDataPtr->MmpIat->MmpIatResolverListLock);
 	}
 
 	return resolver;
 }
 
-_Success_(return)
+extern "C" _Success_(return)
 BOOL WINAPI MmRemoveImportTableResolver(_In_ HANDLE hMmIatResolver) {
 
 	HANDLE heap = RtlProcessHeap();
@@ -176,3 +182,16 @@ BOOL WINAPI MmRemoveImportTableResolver(_In_ HANDLE hMmIatResolver) {
 
 	return RtlFreeHeap(heap, 0, hMmIatResolver);
 }
+
+extern "C" void MmClearImportTableResolvers() {
+    EnterCriticalSection(&MmpGlobalDataPtr->MmpIat->MmpIatResolverListLock);
+    while (!IsListEmpty(&MmpGlobalDataPtr->MmpIat->MmpIatResolverList)) {
+        const auto entry = RemoveTailList(&MmpGlobalDataPtr->MmpIat->MmpIatResolverList);
+       if (!IsListEmpty(&MmpGlobalDataPtr->MmpIat->MmpIatResolverList)) {
+        // Don't delete head
+        const auto resolver = CONTAINING_RECORD(entry, MM_IAT_RESOLVER, InMmpIatResolverList);
+        RtlFreeHeap(RtlProcessHeap(), 0, resolver);
+       }
+    } // while
+    LeaveCriticalSection(&MmpGlobalDataPtr->MmpIat->MmpIatResolverListLock);
+}

MemoryModule/ImportTable.h

@@ -22,10 +22,17 @@ NTSTATUS MemoryResolveImportTable(
 	_In_ PMEMORYMODULE hMemoryModule
 );
 
+extern "C" {
+
 HANDLE WINAPI MmRegisterImportTableResolver(
 	_In_ MM_IAT_RESOLVER_ENTRY LoadLibraryProv,
-	_In_ MM_IAT_FREE_ENTRY FreeLibraryProv
+	_In_ MM_IAT_FREE_ENTRY FreeLibraryProv,
+	_In_ int do_append
 );
 
 _Success_(return)
 BOOL WINAPI MmRemoveImportTableResolver(_In_ HANDLE hMmIatResolver);
+
+void MmClearImportTableResolvers() ;
+
+}

MemoryModule/MemoryModulePP.def

@@ -16,3 +16,7 @@ LdrQuerySystemMemoryModuleFeatures
 
 MmpGlobalDataPtr
 GetMmpGlobalDataPtr
+
+MmRegisterImportTableResolver
+MmRemoveImportTableResolver
+MmClearImportTableResolvers