Owen Rudge : imagehlp: Check for 64-bit PE files in CheckSumMappedFile.

[Alexandre Julliard]

Module: wine
Branch: master
Commit: ba5a86e349bc4f5943413b2d89cd1f1165b93ad7
URL:    http://source.winehq.org/git/wine.git/?a=commit;h=ba5a86e349bc4f5943413b2d89cd1f1165b93ad7

Author: Owen Rudge <orudge at codeweavers.com>
Date:   Fri Dec  4 10:12:00 2009 -0600

imagehlp: Check for 64-bit PE files in CheckSumMappedFile.

---

 dlls/imagehlp/modify.c |   32 ++++++++++++++++++++++++++------
 1 files changed, 26 insertions(+), 6 deletions(-)

diff --git a/dlls/imagehlp/modify.c b/dlls/imagehlp/modify.c
index 02636bd..debccc0 100644
--- a/dlls/imagehlp/modify.c
+++ b/dlls/imagehlp/modify.c
@@ -89,11 +89,14 @@ PIMAGE_NT_HEADERS WINAPI CheckSumMappedFile(
   LPVOID BaseAddress, DWORD FileLength,
   LPDWORD HeaderSum, LPDWORD CheckSum)
 {
-  PIMAGE_NT_HEADERS Header;
+  IMAGE_DOS_HEADER *dos = (IMAGE_DOS_HEADER *) BaseAddress;
+  PIMAGE_NT_HEADERS32 Header32;
+  PIMAGE_NT_HEADERS64 Header64;
+  DWORD *ChecksumFile;
   DWORD CalcSum;
   DWORD HdrSum;
 
-  FIXME("(%p, %d, %p, %p): stub\n",
+  TRACE("(%p, %d, %p, %p)\n",
     BaseAddress, FileLength, HeaderSum, CheckSum
   );
 
@@ -101,8 +104,25 @@ PIMAGE_NT_HEADERS WINAPI CheckSumMappedFile(
 				BaseAddress,
 				(FileLength + 1) / sizeof(WORD));
 
-  Header = RtlImageNtHeader(BaseAddress);
-  HdrSum = Header->OptionalHeader.CheckSum;
+  if (dos->e_magic != IMAGE_DOS_SIGNATURE)
+    return NULL;
+
+  Header32 = (IMAGE_NT_HEADERS32 *)((char *)dos + dos->e_lfanew);
+
+  if (Header32->Signature != IMAGE_NT_SIGNATURE)
+    return NULL;
+
+  if (Header32->OptionalHeader.Magic == IMAGE_NT_OPTIONAL_HDR32_MAGIC)
+    ChecksumFile = &Header32->OptionalHeader.CheckSum;
+  else if (Header32->OptionalHeader.Magic == IMAGE_NT_OPTIONAL_HDR64_MAGIC)
+  {
+    Header64 = (IMAGE_NT_HEADERS64 *)Header32;
+    ChecksumFile = &Header64->OptionalHeader.CheckSum;
+  }
+  else
+    return NULL;
+
+  HdrSum = *ChecksumFile;
 
   /* Subtract image checksum from calculated checksum. */
   /* fix low word of checksum */
@@ -129,9 +149,9 @@ PIMAGE_NT_HEADERS WINAPI CheckSumMappedFile(
   CalcSum += FileLength;
 
   *CheckSum = CalcSum;
-  *HeaderSum = Header->OptionalHeader.CheckSum;
+  *HeaderSum = *ChecksumFile;
 
-  return Header;
+  return (PIMAGE_NT_HEADERS) Header32;
 }
 
 /***********************************************************************