added an extensible dynamic symbol resolver including support for injected DLLs

5 files changed, 197 insertions, 58 deletions

diff --git a/Hunted/Hunted.cpp b/Hunted/Hunted.cpp
index 1ab0a47..458a383 100644
--- a/Hunted/Hunted.cpp
+++ b/Hunted/Hunted.cpp
@@ -304,7 +304,8 @@ class Vec3_tpl<float>   size(12):
 									std::wcout << L"VUnlink failed" << std::endl;
 								}
 #endif						
-								DLLHelper dll;
+								SymbolResolver sresolv;
+								DLLHelper dll(sresolv);
 								if (!dll.Init(targetPID, "./TestDLL.dll")) {
 									std::wcout << L"DLL Init failed" << std::endl;
 								}
@@ -351,7 +352,7 @@ class Vec3_tpl<float>   size(12):
 										<< std::wstring(dir.begin(), dir.end()) << std::endl;
 								}
 
-								PatternScanner pscan(&map_loadlib, &llua);
+								PatternScanner pscan(sresolv, &map_loadlib, &llua);
 								pscan.Scan(md, "01 23 45 67 89 ?? ab cd ef ?? AB CD EF FF");
 
 								// TODO: get gEnv with 0F B7 00 48 83 C4 28 C3
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
diff --git a/include/DLLHelper.h b/include/DLLHelper.h
index 70058f2..21bb9aa 100644
--- a/include/DLLHelper.h
+++ b/include/DLLHelper.h
@@ -4,15 +4,63 @@
 #include <Windows.h>
 
 
-typedef void(*LibEntry_FN)(void);
+struct ResolvedDllEntry {
+	const char * const baseDllName;
+	const char * const functionName;
+
+	HMODULE moduleBase;
+	FARPROC resolvedProc;
+};
+template<SIZE_T s>
+using ResolvedDllArray = std::array<struct ResolvedDllEntry, s>;
+
+typedef HMODULE(*load_library_cb)(IN const char * const module_name,
+	IN PVOID const symbol_resolver_user_data);
+typedef FARPROC(*get_proc_address_cb)(IN HMODULE const module_base,
+	IN const char * const proc_name, IN PVOID const symbol_resolver_user_data);
+typedef BOOL(*free_library_cb)(IN HMODULE const module_base,
+	IN PVOID const symbol_resolver_user_data);
+
+struct symbol_resolver_data {
+	explicit symbol_resolver_data(load_library_cb _loadlib, get_proc_address_cb _getproc, free_library_cb _freelib)
+		: loadlib(_loadlib), getproc(_getproc), freelib(_freelib) {}
+	load_library_cb loadlib;
+	get_proc_address_cb getproc;
+	free_library_cb freelib;
+};
+
+extern const struct symbol_resolver_data sym_loadlib;
+
+typedef void(*LibEntry_FN)(PVOID user_ptr);
 
-static inline bool LoadAndTestLibraryEntry(const char * const fullDllPath);
 bool VerifyPeHeader(UINT8 const * const buf, SIZE_T siz, IMAGE_NT_HEADERS ** const return_NTHeader);
 
+class SymbolResolver
+{
+public:
+	explicit SymbolResolver(struct symbol_resolver_data const * const srd = &sym_loadlib,
+		PVOID const symbol_resolver_user_data = NULL);
+	~SymbolResolver();
+
+	HMODULE LoadLibrary(IN const char * const module_name);
+	FARPROC GetProcAddress(IN HMODULE const module_base,
+		IN const char * const proc_name);
+	BOOL FreeLibrary(IN HMODULE const module_base);
+
+	template<SIZE_T s>
+	bool ResolveAllFunctionSymbols(ResolvedDllArray<s>& rda);
+	template<SIZE_T s>
+	bool CleanupAllFunctionSymbols(ResolvedDllArray<s>& rda);
+	bool LoadAndTestLibraryEntry(const char * const fullDllPath);
+private:
+	struct symbol_resolver_data const * const srd;
+	PVOID symbol_resolver_user_data;
+};
+
 class DLLHelper
 {
 public:
-	DLLHelper();
+	DLLHelper(SymbolResolver& symres);
 	~DLLHelper();
 
 	bool Init(HANDLE targetPID, const char * const fullDllPath);
@@ -39,14 +87,16 @@ public:
 		return (UINT64)m_TargetBaseAddress;
 	}
 	UINT64 GetDllProcAddress(const char * const proc_name) {
-		HMODULE hBase = LoadLibraryA(m_DLLPath.c_str());
-		FARPROC hEntry = GetProcAddress(hBase, proc_name);
+		HMODULE hBase = m_symbolResolver.LoadLibrary(m_DLLPath.c_str());
+		FARPROC hEntry = m_symbolResolver.GetProcAddress(hBase, proc_name);
 		UINT64 result = ((UINT64)hEntry - (UINT64)hBase) + (UINT64)m_TargetBaseAddress;
-		FreeLibrary(hBase);
+		m_symbolResolver.FreeLibrary(hBase);
 		return result;
 	}
 
 private:
+	SymbolResolver& m_symbolResolver;
+
 	HANDLE m_TargetPID = 0;
 	std::string m_DLLPath;
 	DWORD m_DLLSize = 0;
diff --git a/include/PatternScanner.h b/include/PatternScanner.h
index 1c53349..aa011f8 100644
--- a/include/PatternScanner.h
+++ b/include/PatternScanner.h
@@ -6,9 +6,9 @@
 #include <vector>
 
 
-typedef bool(*map_file_cb)(IN MODULE_DATA&, OUT PVOID * const,
+typedef bool(*map_file_cb)(SymbolResolver& symres, IN MODULE_DATA&, OUT PVOID * const,
 	OUT SIZE_T * const, IN PVOID const);
-typedef bool(*map_file_cleanup_cb)(IN MODULE_DATA&,
+typedef bool(*map_file_cleanup_cb)(SymbolResolver& symres, IN MODULE_DATA&,
 	IN PVOID, IN PVOID const);
 
 struct map_file_data {
@@ -29,7 +29,9 @@ extern const struct map_file_data map_kmem;
 class PatternScanner
 {
 public:
-	explicit PatternScanner(struct map_file_data const * const mfd = &map_loadlib, PVOID map_file_user_data = NULL);
+	explicit PatternScanner(SymbolResolver& symres,
+		struct map_file_data const * const mfd = &map_loadlib,
+		PVOID const map_file_user_data = NULL);
 	~PatternScanner();
 	void SetScanLowAddress(UINT64 startAddress) {
 		m_LowAddress = startAddress;
@@ -40,11 +42,14 @@ public:
 	bool Scan(MODULE_DATA& module, const char * const pattern);
 private:
 	bool checkPattern(MODULE_DATA& module, const char * const pattern, std::string& result);
-	bool doScan(UINT8 *buf, SIZE_T size, std::vector<UINT64>& foundOffsets);
+	bool doScan(std::string& pattern, UINT8 *buf, SIZE_T size, std::vector<UINT64>& foundOffsets);
+
+	SymbolResolver& m_symbolResolver;
 
 	struct map_file_data const * const mfd;
+	PVOID const map_file_user_data;
+
 	UINT64 m_LowAddress = 0x0;
 	UINT64 m_HighAddress = ((UINT64)-1);
-	PVOID map_file_user_data;
 };