added function to verify PE header values and removed inlined functions from DLLHelper

2 files changed, 44 insertions, 29 deletions

diff --git a/MemDriverLib/DLLHelper.cpp b/MemDriverLib/DLLHelper.cpp
index 723d63d..f028b83 100644
--- a/MemDriverLib/DLLHelper.cpp
+++ b/MemDriverLib/DLLHelper.cpp
@@ -10,6 +10,45 @@
 #define MakeDelta(cast, x, y) (cast) ((DWORD_PTR)(x) - (DWORD_PTR)(y))
 
 
+bool LoadAndTestLibraryEntry(const char * const fullDllPath)
+{
+	HMODULE TestDLLModule = LoadLibraryA(fullDllPath);
+	LibEntry_FN LibEntryProc = (LibEntry_FN)GetProcAddress(TestDLLModule, "LibEntry");
+	if (LibEntryProc) {
+		LibEntryProc();
+		return true;
+	}
+	else {
+		return false;
+	}
+}
+
+bool VerifyPeHeader(UINT8 const * const buf, SIZE_T siz, IMAGE_NT_HEADERS ** const return_NTHeader)
+{
+	IMAGE_DOS_HEADER *m_DOSHeader;
+
+	if (!return_NTHeader || !buf || siz < sizeof(IMAGE_DOS_HEADER) + sizeof(IMAGE_FILE_HEADER) +
+		sizeof(IMAGE_OPTIONAL_HEADER64))
+	{
+		return false;
+	}
+	*return_NTHeader = NULL;
+	m_DOSHeader = MakePtr(IMAGE_DOS_HEADER *, buf, 0);
+
+	if (m_DOSHeader->e_magic != IMAGE_DOS_SIGNATURE)
+	{
+		return false;
+	}
+
+	*return_NTHeader = MakePtr(IMAGE_NT_HEADERS *, buf, m_DOSHeader->e_lfanew);
+	if ((*return_NTHeader)->Signature != IMAGE_NT_SIGNATURE)
+	{
+		return false;
+	}
+
+	return true;
+}
+
 static FARPROC GetRemoteProcAddress(HMODULE localMod, HMODULE remoteMod, char *func_name)
 {
 	/*
@@ -84,7 +123,7 @@ DLLHelper::~DLLHelper()
 	}
 }
 
-bool DLLHelper::Init(HANDLE targetPID, const char * fullDllPath) {
+bool DLLHelper::Init(HANDLE targetPID, const char * const fullDllPath) {
 	if (!targetPID) {
 		return false;
 	}
@@ -139,21 +178,8 @@ bool DLLHelper::VerifyHeader()
 	if (!m_DLLPtr) {
 		return false;
 	}
-	m_DOSHeader = MakePtr(IMAGE_DOS_HEADER *, m_DLLPtr, 0);
-
-	if (m_DOSHeader->e_magic != IMAGE_DOS_SIGNATURE)
-	{
-		delete m_DLLPtr;
-		return false;
-	}
-	m_NTHeader = MakePtr(IMAGE_NT_HEADERS *, m_DLLPtr, m_DOSHeader->e_lfanew);
-	if (m_NTHeader->Signature != IMAGE_NT_SIGNATURE)
-	{
-		delete m_DLLPtr;
-		return false;
-	}
 
-	return true;
+	return VerifyPeHeader(m_DLLPtr, m_DLLSize, &m_NTHeader);
 }
 
 bool DLLHelper::InitTargetMemory()
diff --git a/include/DLLHelper.h b/include/DLLHelper.h
index fc17711..7fdff37 100644
--- a/include/DLLHelper.h
+++ b/include/DLLHelper.h
@@ -6,18 +6,8 @@
 
 typedef void(*LibEntry_FN)(void);
 
-static inline bool LoadAndTestLibraryEntry(const char * const fullDllPath)
-{
-	HMODULE TestDLLModule = LoadLibraryA(fullDllPath);
-	LibEntry_FN LibEntryProc = (LibEntry_FN)GetProcAddress(TestDLLModule, "LibEntry");
-	if (LibEntryProc) {
-		LibEntryProc();
-		return true;
-	}
-	else {
-		return false;
-	}
-}
+static inline bool LoadAndTestLibraryEntry(const char * const fullDllPath);
+bool VerifyPeHeader(UINT8 const * const buf, SIZE_T siz, IMAGE_NT_HEADERS ** const return_NTHeader);
 
 class DLLHelper
 {
@@ -25,7 +15,7 @@ public:
 	DLLHelper();
 	~DLLHelper();
 
-	bool Init(HANDLE targetPID, const char * fullDllPath);
+	bool Init(HANDLE targetPID, const char * const fullDllPath);
 	bool VerifyHeader();
 	bool InitTargetMemory();
 	bool HasImports() {
@@ -54,7 +44,6 @@ private:
 	std::string m_DLLPath;
 	DWORD m_DLLSize = 0;
 	UINT8 *m_DLLPtr = nullptr;
-	IMAGE_DOS_HEADER *m_DOSHeader = nullptr;
 	IMAGE_NT_HEADERS *m_NTHeader = nullptr;
 	PVOID m_TargetBaseAddress = nullptr;
 };