added an extensible dynamic symbol resolver including support for injected DLLs

2 files changed, 66 insertions, 11 deletions

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;
 };