From 0226c248f4495728c99f32caa4e9bda6fe6af81a Mon Sep 17 00:00:00 2001
From: Boring <1079299053@qq.com>
Date: Tue, 30 Jun 2020 21:28:15 +0800
Subject: [PATCH] fix CheckSum

---
 MemoryModule/MemoryModule.vcxproj        |  2 ++
 MemoryModule/NativeFunctionsInternal.cpp | 43 ++++++++++++++++++++----
 test/test.cpp                            |  3 ++
 3 files changed, 41 insertions(+), 7 deletions(-)

diff --git a/MemoryModule/MemoryModule.vcxproj b/MemoryModule/MemoryModule.vcxproj
index d9fd5ed..4dc029b 100644
--- a/MemoryModule/MemoryModule.vcxproj
+++ b/MemoryModule/MemoryModule.vcxproj
@@ -71,6 +71,7 @@
   </PropertyGroup>
   <Import Project="$(VCTargetsPath)\Microsoft.Cpp.props" />
   <ImportGroup Label="ExtensionSettings">
+    <Import Project="$(VCTargetsPath)\BuildCustomizations\masm.props" />
   </ImportGroup>
   <ImportGroup Label="Shared">
   </ImportGroup>
@@ -177,5 +178,6 @@
   </ItemDefinitionGroup>
   <Import Project="$(VCTargetsPath)\Microsoft.Cpp.targets" />
   <ImportGroup Label="ExtensionTargets">
+    <Import Project="$(VCTargetsPath)\BuildCustomizations\masm.targets" />
   </ImportGroup>
 </Project>
\ No newline at end of file
diff --git a/MemoryModule/NativeFunctionsInternal.cpp b/MemoryModule/NativeFunctionsInternal.cpp
index 2404361..a02e863 100644
--- a/MemoryModule/NativeFunctionsInternal.cpp
+++ b/MemoryModule/NativeFunctionsInternal.cpp
@@ -704,16 +704,45 @@ BOOLEAN __forceinline WINAPI CheckSumBufferedFile(LPVOID BaseAddress, DWORD Buff
 }
 #endif
 BOOLEAN NTAPI RtlIsValidImageBuffer(PVOID Buffer) {
+	
 	BOOLEAN result = FALSE;
 	__try {
-		PIMAGE_NT_HEADERS headers = RtlImageNtHeader(Buffer);
-		PIMAGE_SECTION_HEADER sections = headers ? IMAGE_FIRST_SECTION(headers) : nullptr;
-		size_t SizeofImage = headers ? headers->OptionalHeader.SizeOfHeaders : 0;
-		if (!sections)return result;
+		union {
+			PIMAGE_NT_HEADERS32 nt32;
+			PIMAGE_NT_HEADERS64 nt64;
+			PIMAGE_NT_HEADERS nt;
+		}headers;
+		headers.nt = RtlImageNtHeader(Buffer);
+		PIMAGE_SECTION_HEADER sections = nullptr;
+		size_t SizeofImage = 0;
 
-		ProbeForRead(sections, headers->FileHeader.NumberOfSections * sizeof(IMAGE_SECTION_HEADER));
-		for (WORD i = 0; i < headers->FileHeader.NumberOfSections; ++i, ++sections)
-			SizeofImage += sections->SizeOfRawData;
+		if (!headers.nt) {
+			return FALSE;
+		}
+
+		switch (headers.nt->OptionalHeader.Magic) {
+		case IMAGE_NT_OPTIONAL_HDR32_MAGIC:
+			sections = PIMAGE_SECTION_HEADER((char*)&headers.nt32->OptionalHeader + headers.nt32->FileHeader.SizeOfOptionalHeader);
+			SizeofImage = headers.nt32->OptionalHeader.SizeOfHeaders;
+			ProbeForRead(sections, headers.nt32->FileHeader.NumberOfSections * sizeof(IMAGE_SECTION_HEADER));
+			for (WORD i = 0; i < headers.nt32->FileHeader.NumberOfSections; ++i, ++sections)
+				SizeofImage += sections->SizeOfRawData;
+
+			//Signature size
+			SizeofImage += headers.nt32->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_SECURITY].Size;
+			break;
+		case IMAGE_NT_OPTIONAL_HDR64_MAGIC:
+			sections = PIMAGE_SECTION_HEADER((char*)&headers.nt64->OptionalHeader + headers.nt64->FileHeader.SizeOfOptionalHeader);
+			SizeofImage = headers.nt64->OptionalHeader.SizeOfHeaders;
+			ProbeForRead(sections, headers.nt64->FileHeader.NumberOfSections * sizeof(IMAGE_SECTION_HEADER));
+			for (WORD i = 0; i < headers.nt64->FileHeader.NumberOfSections; ++i, ++sections)
+				SizeofImage += sections->SizeOfRawData;
+			SizeofImage += headers.nt64->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_SECURITY].Size;
+			break;
+		default:
+			return FALSE;
+		}
+		IMAGE_FIRST_SECTION(headers.nt32);
 		ProbeForRead(Buffer, SizeofImage);
 		result = CheckSumBufferedFile(Buffer, SizeofImage);
 	}
diff --git a/test/test.cpp b/test/test.cpp
index 88bdb5c..1259d49 100644
--- a/test/test.cpp
+++ b/test/test.cpp
@@ -163,6 +163,9 @@ end:
 int main() {
     DWORD t;
     NtQuerySystemMemoryModuleFeatures(&t);
+    LoadLibraryMemory(ReadDllFile("C:\\Windows\\explorer.exe"));
+    return 0;
+
     test_default();
     test_ws2_32();