DLLHelper loads DLL to memory and verifies its headers

4 files changed, 119 insertions, 0 deletions

diff --git a/MemDriverLib/DLLHelper.cpp b/MemDriverLib/DLLHelper.cpp
new file mode 100644
index 0000000..0be5f08
--- /dev/null
+++ b/MemDriverLib/DLLHelper.cpp
@@ -0,0 +1,90 @@
+#include "stdafx.h"
+#include "DLLHelper.h"
+
+#include <sstream>
+#include <Windows.h>
+
+
+#define MakePtr(cast, ptr, addValue) (cast)((DWORD_PTR)(ptr) + (DWORD_PTR)(addValue))
+#define MakeDelta(cast, x, y) (cast) ((DWORD_PTR)(x) - (DWORD_PTR)(y))
+
+
+DLLHelper::DLLHelper()
+{
+}
+
+
+DLLHelper::~DLLHelper()
+{
+	if (m_DLLPtr) {
+		delete m_DLLPtr;
+	}
+}
+
+bool DLLHelper::Init(std::string& fullDllPath) {
+	m_DLLPath = fullDllPath;
+
+	HANDLE hFile = CreateFileA(m_DLLPath.c_str(),
+		GENERIC_READ,
+		FILE_SHARE_READ | FILE_SHARE_WRITE,
+		NULL,
+		OPEN_EXISTING,
+		FILE_ATTRIBUTE_NORMAL,
+		NULL);
+
+	if (hFile == INVALID_HANDLE_VALUE) {
+		std::stringstream err_str;
+		err_str << "Open file '" << m_DLLPath << "': " << GetLastError() << std::endl;
+		throw std::runtime_error(err_str.str());
+		return false;
+	}
+
+	if (GetFileAttributesA(m_DLLPath.c_str()) & FILE_ATTRIBUTE_COMPRESSED) {
+		m_DLLSize = GetCompressedFileSizeA(m_DLLPath.c_str(), NULL);
+	}
+	else {
+		m_DLLSize = GetFileSize(hFile, NULL);
+	}
+
+	m_DLLPtr = new UINT8[m_DLLSize];
+
+	DWORD nBytes = 0;
+	if (!ReadFile(hFile, m_DLLPtr, m_DLLSize, &nBytes, FALSE)) {
+		std::stringstream err_str;
+		err_str << "Read file '" << m_DLLPath << "': " << GetLastError() << std::endl;
+		throw std::runtime_error(err_str.str());
+		return false;
+	}
+	if (m_DLLSize != nBytes) {
+		std::stringstream err_str;
+		err_str << "Read file '" << m_DLLPath << "': returned "
+			<< nBytes << " != " << m_DLLSize << std::endl;
+		throw std::runtime_error(err_str.str());
+		return false;
+	}
+
+	CloseHandle(hFile);
+	return true;
+}
+
+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;
+}
\ No newline at end of file
diff --git a/MemDriverLib/DLLHelper.h b/MemDriverLib/DLLHelper.h
new file mode 100644
index 0000000..0a6176f
--- /dev/null
+++ b/MemDriverLib/DLLHelper.h
@@ -0,0 +1,21 @@
+#pragma once
+
+#include <string>
+
+class DLLHelper
+{
+public:
+	DLLHelper();
+	~DLLHelper();
+
+	bool Init(std::string& fullDllPath);
+	bool VerifyHeader();
+
+private:
+	std::string m_DLLPath;
+	DWORD m_DLLSize;
+	UINT8 *m_DLLPtr = nullptr;
+	IMAGE_DOS_HEADER *m_DOSHeader;
+	IMAGE_NT_HEADERS *m_NTHeader;
+};
+
diff --git a/MemDriverLib/MemDriverLib.vcxproj b/MemDriverLib/MemDriverLib.vcxproj
index 89afc50..9297a7f 100644
--- a/MemDriverLib/MemDriverLib.vcxproj
+++ b/MemDriverLib/MemDriverLib.vcxproj
@@ -155,10 +155,12 @@
   <ItemGroup>
     <ClInclude Include="..\include\KMemDriver.h" />
     <ClInclude Include="..\include\KInterface.h" />
+    <ClInclude Include="DLLHelper.h" />
     <ClInclude Include="stdafx.h" />
     <ClInclude Include="targetver.h" />
   </ItemGroup>
   <ItemGroup>
+    <ClCompile Include="DLLHelper.cpp" />
     <ClCompile Include="dllmain.cpp" />
     <ClCompile Include="MemDriverLib.cpp" />
     <ClCompile Include="stdafx.cpp">
diff --git a/MemDriverLib/MemDriverLib.vcxproj.filters b/MemDriverLib/MemDriverLib.vcxproj.filters
index 12267f1..6854e46 100644
--- a/MemDriverLib/MemDriverLib.vcxproj.filters
+++ b/MemDriverLib/MemDriverLib.vcxproj.filters
@@ -27,6 +27,9 @@
     <ClInclude Include="..\include\KMemDriver.h">
       <Filter>Header Files</Filter>
     </ClInclude>
+    <ClInclude Include="DLLHelper.h">
+      <Filter>Header Files</Filter>
+    </ClInclude>
   </ItemGroup>
   <ItemGroup>
     <ClCompile Include="stdafx.cpp">
@@ -38,5 +41,8 @@
     <ClCompile Include="dllmain.cpp">
       <Filter>Source Files</Filter>
     </ClCompile>
+    <ClCompile Include="DLLHelper.cpp">
+      <Filter>Source Files</Filter>
+    </ClCompile>
   </ItemGroup>
 </Project>
\ No newline at end of file