added an extensible dynamic symbol resolver including support for injected DLLs

2 files changed, 128 insertions, 45 deletions

diff --git a/MemDriverLib/DLLHelper.cpp b/MemDriverLib/DLLHelper.cpp
index f028b83..18d60e3 100644
--- a/MemDriverLib/DLLHelper.cpp
+++ b/MemDriverLib/DLLHelper.cpp
@@ -9,18 +9,113 @@
 #define MakePtr(cast, ptr, addValue) (cast)((DWORD_PTR)(ptr) + (DWORD_PTR)(addValue))
 #define MakeDelta(cast, x, y) (cast) ((DWORD_PTR)(x) - (DWORD_PTR)(y))
 
+static HMODULE sym_res_loadlib(IN const char * const module_name,
+	IN PVOID const symbol_resolver_user_data);
+static FARPROC sym_res_getproc(IN HMODULE const module_base,
+	IN const char * const proc_name, IN PVOID const symbol_resolver_user_data);
+static BOOL sym_res_freelib(IN HMODULE const module_base,
+	IN PVOID const symbol_resolver_user_data);
+
+const struct symbol_resolver_data sym_loadlib = symbol_resolver_data(sym_res_loadlib,
+	sym_res_getproc, sym_res_freelib);
+
+static HMODULE sym_res_loadlib(IN const char * const module_name,
+	IN PVOID const symbol_resolver_user_data)
+{
+	UNREFERENCED_PARAMETER(symbol_resolver_user_data);
+
+	return LoadLibraryA(module_name);
+}
+
+static FARPROC sym_res_getproc(IN HMODULE const module_base,
+	IN const char * const proc_name, IN PVOID const symbol_resolver_user_data)
+{
+	UNREFERENCED_PARAMETER(symbol_resolver_user_data);
+
+	return GetProcAddress(module_base, proc_name);
+}
+
+static BOOL sym_res_freelib(IN HMODULE const module_base, IN PVOID const symbol_resolver_user_data)
+{
+	UNREFERENCED_PARAMETER(symbol_resolver_user_data);
+
+	return FreeLibrary(module_base);
+}
+
+SymbolResolver::SymbolResolver(struct symbol_resolver_data const * const srd, PVOID symbol_resolver_user_Data)
+	: srd(srd), symbol_resolver_user_data(symbol_resolver_user_data)
+{
+}
+
+SymbolResolver::~SymbolResolver()
+{
+}
+
+HMODULE SymbolResolver::LoadLibrary(IN const char * const module_name)
+{
+	return srd->loadlib(module_name, symbol_resolver_user_data);
+}
+
+FARPROC SymbolResolver::GetProcAddress(IN HMODULE const module_base,
+	IN const char * const proc_name)
+{
+	return srd->getproc(module_base, proc_name, symbol_resolver_user_data);
+}
+
+BOOL SymbolResolver::FreeLibrary(IN HMODULE const module_base)
+{
+	return srd->freelib(module_base, symbol_resolver_user_data);
+}
+
+template<SIZE_T s>
+bool SymbolResolver::ResolveAllFunctionSymbols(ResolvedDllArray<s>& rda)
+{
+	bool result = true;
+
+	for (auto& unresolved : rda) {
+		unresolved.moduleBase = this->LoadLibrary(unresolved.baseDllName);
+		if (!unresolved.moduleBase) {
+			result = false;
+			continue;
+		}
+
+		unresolved.resolvedProc = this->GetProcAddress(unresolved.moduleBase,
+			unresolved.functionName);
+		if (!unresolved.resolvedProc) {
+			result = false;
+		}
+	}
+
+	return result;
+}
+
+template<SIZE_T s>
+bool SymbolResolver::CleanupAllFunctionSymbols(ResolvedDllArray<s>& rda)
+{
+	bool result = true;
+
+	for (auto& unresolved : rda) {
+		result = this->FreeLibrary(unresolved.moduleBase);
+		unresolved.moduleBase = NULL;
+		unresolved.resolvedProc = NULL;
+	}
+
+	return result;
+}
 
-bool LoadAndTestLibraryEntry(const char * const fullDllPath)
+bool SymbolResolver::LoadAndTestLibraryEntry(const char * const fullDllPath)
 {
-	HMODULE TestDLLModule = LoadLibraryA(fullDllPath);
-	LibEntry_FN LibEntryProc = (LibEntry_FN)GetProcAddress(TestDLLModule, "LibEntry");
+	HMODULE TestDLLModule = this->LoadLibrary(fullDllPath);
+	LibEntry_FN LibEntryProc = (LibEntry_FN)this->GetProcAddress(TestDLLModule,
+		"LibEntry");
 	if (LibEntryProc) {
-		LibEntryProc();
+		LibEntryProc(NULL);
 		return true;
 	}
 	else {
 		return false;
 	}
+	this->FreeLibrary(TestDLLModule);
 }
 
 bool VerifyPeHeader(UINT8 const * const buf, SIZE_T siz, IMAGE_NT_HEADERS ** const return_NTHeader)
@@ -49,14 +144,15 @@ bool VerifyPeHeader(UINT8 const * const buf, SIZE_T siz, IMAGE_NT_HEADERS ** con
 	return true;
 }
 
-static FARPROC GetRemoteProcAddress(HMODULE localMod, HMODULE remoteMod, char *func_name)
+static FARPROC GetRemoteProcAddress(SymbolResolver& symres,
+	HMODULE localMod, HMODULE remoteMod, char *func_name)
 {
 	/*
 	 * Account for potential differences in base address
 	 * of modules in different processes.
 	 */
 	ULONGLONG delta = MakeDelta(ULONGLONG, remoteMod, localMod);
-	return MakePtr(FARPROC, GetProcAddress(localMod, func_name), delta);
+	return MakePtr(FARPROC, symres.GetProcAddress(localMod, func_name), delta);
 }
 
 static HMODULE GetRemoteModuleHandle(char *module_name,
@@ -111,7 +207,8 @@ static LPVOID GetPtrFromRVA(DWORD rva, IMAGE_NT_HEADERS *pNTHeader, PBYTE imageB
 	return (PVOID)(imageBase + rva - delta);
 }
 
-DLLHelper::DLLHelper()
+DLLHelper::DLLHelper(SymbolResolver& symres)
+	: m_symbolResolver(symres)
 {
 }
 
@@ -233,7 +330,7 @@ bool DLLHelper::FixImports()
 	while ((module_name = (char *)GetPtrFromRVA((DWORD)(impDesc->Name), m_NTHeader,
 		(PBYTE)m_DLLPtr)))
 	{
-		HMODULE localMod = LoadLibraryA(module_name);
+		HMODULE localMod = m_symbolResolver.LoadLibrary(module_name);
 		HMODULE remoteMod = GetRemoteModuleHandle(module_name, modules);
 
 		if (!remoteMod) {
@@ -254,8 +351,8 @@ bool DLLHelper::FixImports()
 			iibn = (IMAGE_IMPORT_BY_NAME *)GetPtrFromRVA((DWORD)(itd->u1.AddressOfData),
 				m_NTHeader, (PBYTE)m_DLLPtr);
 
-			itd->u1.Function = MakePtr(ULONGLONG, GetRemoteProcAddress(localMod,
-				remoteMod, (char *)iibn->Name), 0);
+			itd->u1.Function = MakePtr(ULONGLONG, GetRemoteProcAddress(m_symbolResolver,
+				localMod, remoteMod, (char *)iibn->Name), 0);
 
 			itd++;
 		}
diff --git a/MemDriverLib/PatternScanner.cpp b/MemDriverLib/PatternScanner.cpp
index 3e1b8ee..a274b61 100644
--- a/MemDriverLib/PatternScanner.cpp
+++ b/MemDriverLib/PatternScanner.cpp
@@ -11,19 +11,14 @@
 #include <Windows.h>
 
 
-static bool map_file_loadlib(MODULE_DATA& module, PVOID * const buffer,
+static bool map_file_loadlib(SymbolResolver& symres, MODULE_DATA& module, PVOID * const buffer,
 	SIZE_T * const size, PVOID const user_ptr);
-static bool map_file_loadlib_cleanup(MODULE_DATA& module, PVOID buffer,
-	PVOID const user_ptr);
-static bool map_file_kmem(MODULE_DATA& module, PVOID * const buffer,
-	SIZE_T * const size, PVOID const user_ptr);
-static bool map_file_kmem_cleanup(MODULE_DATA& module, PVOID buffer,
+static bool map_file_loadlib_cleanup(SymbolResolver& symres, MODULE_DATA& module, PVOID buffer,
 	PVOID const user_ptr);
 
 const struct map_file_data map_loadlib = map_file_data(map_file_loadlib, map_file_loadlib_cleanup, true);
-const struct map_file_data map_kmem = map_file_data(map_file_kmem, map_file_kmem_cleanup, false);
 
-bool map_file_loadlib(MODULE_DATA& module, PVOID * const buffer,
+bool map_file_loadlib(SymbolResolver& symres, MODULE_DATA& module, PVOID * const buffer,
 	SIZE_T * const size, PVOID const user_ptr)
 {
 	HMODULE hMod;
@@ -45,7 +40,7 @@ bool map_file_loadlib(MODULE_DATA& module, PVOID * const buffer,
 		}
 	}
 
-	hMod = LoadLibraryA(module.FullDllPath);
+	hMod = symres.LoadLibrary(module.FullDllPath);
 
 	if (user_data) {
 		if (dir_cookies.size() > 1) {
@@ -69,24 +64,14 @@ bool map_file_loadlib(MODULE_DATA& module, PVOID * const buffer,
 	}
 }
 
-bool map_file_loadlib_cleanup(MODULE_DATA& module, PVOID buffer, PVOID const user_ptr)
-{
-	return FreeLibrary((HMODULE)buffer);
-}
-
-bool map_file_kmem(MODULE_DATA& module, PVOID * const buffer,
-	SIZE_T * const size, PVOID const user_ptr)
+bool map_file_loadlib_cleanup(SymbolResolver& symres, MODULE_DATA& module, PVOID buffer, PVOID const user_ptr)
 {
-	return false;
-}
-
-bool map_file_kmem_cleanup(MODULE_DATA& module, PVOID buffer, PVOID const user_ptr)
-{
-	return false;
+	return symres.FreeLibrary((HMODULE)buffer);
 }
 
-PatternScanner::PatternScanner(struct map_file_data const * const mfd, PVOID map_file_user_data)
-	: mfd(mfd), map_file_user_data(map_file_user_data)
+PatternScanner::PatternScanner(SymbolResolver& symres,
+	struct map_file_data const * const mfd, PVOID map_file_user_data)
+	: m_symbolResolver(symres), mfd(mfd), map_file_user_data(map_file_user_data)
 {
 	if (!mfd) {
 		throw std::runtime_error("MapFileData was NULL");
@@ -137,14 +122,19 @@ bool PatternScanner::checkPattern(MODULE_DATA& module, const char * const patter
 	return true;
 }
 
-bool PatternScanner::doScan(UINT8 *buf, SIZE_T size, std::vector<UINT64>& foundOffsets)
+#include <iostream>
+bool PatternScanner::doScan(std::string& pattern, UINT8 *buf, SIZE_T size, std::vector<UINT64>& foundOffsets)
 {
+	//std::wcout << "BLAAAAAAAAAAAAA" << std::endl;
+	//std::wstring bla(str_pattern.begin(), str_pattern.end());
+	//std::wcout << bla << std::endl;
+	std::cout << pattern << std::endl;
 	return false;
 }
 
-#include <iostream>
 bool PatternScanner::Scan(MODULE_DATA& module, const char * const pattern)
 {
+	bool result;
 	std::string validPattern;
 	IMAGE_NT_HEADERS *ntHeader;
 	IMAGE_SECTION_HEADER *secHeader;
@@ -156,7 +146,7 @@ bool PatternScanner::Scan(MODULE_DATA& module, const char * const pattern)
 		return false;
 	}
 
-	if (!mfd->mapfile(module, (PVOID *)&mappedBuffer, &mappedSize, map_file_user_data))
+	if (!mfd->mapfile(m_symbolResolver, module, (PVOID *)&mappedBuffer, &mappedSize, map_file_user_data))
 	{
 		return false;
 	}
@@ -180,20 +170,16 @@ bool PatternScanner::Scan(MODULE_DATA& module, const char * const pattern)
 			virtualSize = secHeader->VirtualAddress - virtualSize;
 			nBytes += virtualSize;
 		}
+		result = false;
 	}
 	else {
-		doScan(mappedBuffer, mappedSize, foundOffsets);
+		result = doScan(validPattern, mappedBuffer, mappedSize, foundOffsets);
 	}
 
-	if (!mfd->mapcleanup(module, mappedBuffer, map_file_user_data))
+	if (!mfd->mapcleanup(m_symbolResolver, module, mappedBuffer, map_file_user_data))
 	{
 		return false;
 	}
 
-	//std::wcout << "BLAAAAAAAAAAAAA" << std::endl;
-	//std::wstring bla(str_pattern.begin(), str_pattern.end());
-	//std::wcout << bla << std::endl;
-	std::cout << validPattern << std::endl;
-
-	return true;
+	return result;
 }
\ No newline at end of file