feat(Util): Added checks for invalid DLL

Closes #4

Author: Yimura

src/Util.hpp

@@ -1,6 +1,16 @@
 #pragma once
 #include "Common.hpp"
 
+enum class DllValidityError
+{
+	VALID,
+	ACCESS_FAILURE,
+	TOO_SMALL,
+	ALLOCATION_FAILURE,
+	NOT_A_DLL,
+	INVALID_PLATFORM
+};
+
 namespace Util
 {
 	inline const wchar_t* CharToWchar(const char* a)
@@ -11,4 +21,68 @@ namespace Util
 
 		return wchar;
 	}
+
+	inline DllValidityError CheckIfFileIsValidDll(const std::filesystem::path& dllFile)
+	{
+		std::ifstream fileStream(dllFile, std::ios::binary | std::ios::ate);
+
+		if (fileStream.fail())
+		{
+			fileStream.close();
+
+			return DllValidityError::ACCESS_FAILURE;
+		}
+
+		const auto fileSize = fileStream.tellg();
+		if (fileSize < 0x1000)
+		{
+			fileStream.close();
+
+			return DllValidityError::TOO_SMALL;
+		}
+
+		auto* pSrcData = new uint8_t[static_cast<uintptr_t>(fileSize)];
+		if (!pSrcData)
+		{
+			fileStream.close();
+
+			return DllValidityError::ALLOCATION_FAILURE;
+		}
+
+		fileStream.seekg(0, std::ios::beg);
+		fileStream.read(reinterpret_cast<char*>(pSrcData), fileSize);
+		fileStream.close();
+
+		if (reinterpret_cast<IMAGE_DOS_HEADER*>(pSrcData)->e_magic != 0x5A4D)
+		{
+			delete[] pSrcData;
+
+			return DllValidityError::NOT_A_DLL;
+		}
+
+		const auto* pOldNtHeader	= reinterpret_cast<IMAGE_NT_HEADERS*>(pSrcData + reinterpret_cast<IMAGE_DOS_HEADER*>(pSrcData)->e_lfanew);
+		const auto* pOldFileHeader	= &pOldNtHeader->FileHeader;
+
+#ifdef _WIN64
+		if (pOldFileHeader->Machine != IMAGE_FILE_MACHINE_AMD64)
+		{
+			delete[] pSrcData;
+			delete pOldFileHeader;
+			delete pOldNtHeader;
+
+			return DllValidityError::INVALID_PLATFORM;
+		}
+#else
+		if (pOldFileHeader->Machine != IMAGE_FILE_MACHINE_I386)
+		{
+			delete[] pSrcData;
+			delete pOldFileHeader;
+			delete pOldNtHeader;
+
+			return DllValidityError::INVALID_PLATFORM;
+		}
+#endif
+
+		return DllValidityError::VALID;
+	}
 }

src/main.cpp

@@ -59,6 +59,34 @@ int main(int argc, const char** argv)
 	if (!dllFile.is_absolute())
 		dllFile = std::filesystem::absolute(dllFile);
 	LOG(INFO) << "Starting injection for " << dllFile.filename().string();
+
+	switch (Util::CheckIfFileIsValidDll(dllFile))
+	{
+	case DllValidityError::ACCESS_FAILURE:
+		LOG(WARNING) << "Failed to access DLL on disk.";
+
+		return 1;
+	case DllValidityError::TOO_SMALL:
+		LOG(WARNING) << "DLL file seems inconceivably small, request to inject ignored.";
+
+		return 1;
+	case DllValidityError::ALLOCATION_FAILURE:
+		LOG(WARNING) << "Failed to allocate memory when checking DLL file.";
+
+		return 1;
+	case DllValidityError::NOT_A_DLL:
+		LOG(WARNING) << "The file given does not appear to be a valid DLL.";
+
+		return 1;
+	case DllValidityError::INVALID_PLATFORM:
+		LOG(WARNING) << "The DLL given did not match the target platform the injector.";
+
+		return 1;
+	case DllValidityError::VALID:
+		LOG(INFO) << "DLL seems valid, proceeding with injection.";
+
+		break;
+	}
 
 	const int processId = Injector::GetProcessId(targetApplication.data());
 	if (processId == -1)