fix: update the PE checksum after committing changes

rcedit.gyp

@@ -10,6 +10,9 @@
         'src/rcedit.rc',
         'src/version.h',
       ],
+      'libraries': [
+        'imagehlp.lib',
+      ],
       'msvs_settings': {
         'VCLinkerTool': {
           'SubSystem': 1, # console executable

rcedit.vcxproj

@@ -45,7 +45,7 @@
       <RuntimeLibrary>MultiThreaded</RuntimeLibrary>
     </ClCompile>
     <Link>
-      <AdditionalDependencies></AdditionalDependencies>
+      <AdditionalDependencies>imagehlp.lib</AdditionalDependencies>
       <OutputFile>$(OutDir)$(ProjectName)$(TargetExt)</OutputFile>
       <SubSystem>Console</SubSystem>
     </Link>

src/rescle.cc

@@ -9,6 +9,7 @@
 
 #include <assert.h>
 #include <atlstr.h>
+#include <imagehlp.h>
 #include <sstream> // wstringstream
 #include <iomanip> // setw, setfill
 #include <fstream>
@@ -93,8 +94,8 @@ std::wstring ReadFileToString(const wchar_t* filename) {
 
 class ScopedFile {
  public:
-  ScopedFile(const WCHAR* path)
-    : file_(CreateFileW(path, GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL)) {}
+  ScopedFile(const WCHAR* path, DWORD access = GENERIC_READ)
+    : file_(CreateFileW(path, access, FILE_SHARE_READ, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL)) {}
   ~ScopedFile() { CloseHandle(file_); }
 
   operator HANDLE() { return file_; }
@@ -568,8 +569,8 @@ bool ResourceUpdater::ChangeRcData(UINT id, const WCHAR* pathToResource) {
 
   wchar_t abspath[MAX_PATH] = { 0 };
   const auto filePath = _wfullpath(abspath, pathToResource, MAX_PATH) ? abspath : pathToResource;
-  ScopedFile newRcDataFile(filePath);
-  if (newRcDataFile == INVALID_HANDLE_VALUE) {
+  ScopedFile newRcDataFile(filePath);
+  if (newRcDataFile == INVALID_HANDLE_VALUE) {
     fprintf(stderr, "Cannot open new data file '%ws'\n", filePath);
     return false;
   }
@@ -581,13 +582,13 @@ bool ResourceUpdater::ChangeRcData(UINT id, const WCHAR* pathToResource) {
   }
 
   auto& rcData = rcDataLngPairIt->second[id];
-  rcData.clear();
+  rcData.clear();
   rcData.resize(dwFileSize);
-
-  DWORD dwBytesRead{ 0 };
+
+  DWORD dwBytesRead{ 0 };
   if (!ReadFile(newRcDataFile, rcData.data(), dwFileSize, &dwBytesRead, NULL)) {
     fprintf(stderr, "Cannot read file '%ws'\n", filePath);
-    return false;
+    return false;
   }
 
   return true;
@@ -701,7 +702,7 @@ bool ResourceUpdater::Commit() {
   FreeLibrary(module_);
   module_ = NULL;
 
-  ScopedResourceUpdater ru(filename_.c_str(), false);
+  ScopedResourceUpdater ru(filename_, false);
   if (ru.Get() == NULL) {
     return false;
   }
@@ -979,8 +980,8 @@ BOOL CALLBACK ResourceUpdater::OnEnumResourceManifest(HMODULE hModule, LPCTSTR l
   return TRUE;   // Keep going
 }
 
-ScopedResourceUpdater::ScopedResourceUpdater(const WCHAR* filename, bool deleteOld)
-    : handle_(BeginUpdateResourceW(filename, deleteOld)) {
+ScopedResourceUpdater::ScopedResourceUpdater(std::wstring filename, bool deleteOld)
+    : filename_(filename), handle_(BeginUpdateResourceW(filename.c_str(), deleteOld)) {
 }
 
 ScopedResourceUpdater::~ScopedResourceUpdater() {
@@ -995,7 +996,7 @@ HANDLE ScopedResourceUpdater::Get() const {
 
 bool ScopedResourceUpdater::Commit() {
   commited_ = true;
-  return EndUpdate(true);
+  return EndUpdate(true) && UpdateChecksum();
 }
 
 bool ScopedResourceUpdater::EndUpdate(bool doesCommit) {
@@ -1005,4 +1006,45 @@ bool ScopedResourceUpdater::EndUpdate(bool doesCommit) {
   return bResult ? true : false;
 }
 
+bool ScopedResourceUpdater::UpdateChecksum() {
+  const WCHAR* path = filename_.c_str();
+  ScopedFile binFile(path, GENERIC_READ|GENERIC_WRITE);
+  if (binFile == INVALID_HANDLE_VALUE) {
+    fprintf(stderr, "Cannot open '%ws'\n", path);
+    return false;
+  }
+  DWORD fileSize = GetFileSize(binFile, NULL);
+  if (fileSize == INVALID_FILE_SIZE) {
+    fprintf(stderr, "Cannot get file size for '%ws'\n", path);
+    return false;
+  }
+
+  HANDLE hFilemap = CreateFileMapping(binFile, NULL, PAGE_READWRITE, 0, 0, NULL);
+  if (!hFilemap) {
+    fprintf(stderr, "CreateFileMapping(%ws) failed\n", path);
+    return false;
+  }
+
+  void* view = MapViewOfFile(hFilemap, FILE_MAP_ALL_ACCESS, 0, 0, 0);
+  if (!view) {
+    fprintf(stderr, "Cannot map '%ws' into memory\n", path);
+    CloseHandle(hFilemap);
+    return false;
+  }
+
+  bool ret = false;
+  DWORD originalSum, checkSum;
+  IMAGE_NT_HEADERS *ntHeader = CheckSumMappedFile(view, fileSize, &originalSum, &checkSum);
+  if (!ntHeader) {
+    fprintf(stderr, "CheckSumMappedFile failed\n");
+  } else {
+    ntHeader->OptionalHeader.CheckSum = checkSum;
+    ret = true;
+  }
+
+  UnmapViewOfFile(view);
+  CloseHandle(hFilemap);
+  return ret;
+}
+
 }  // namespace rescle

src/rescle.h

@@ -167,15 +167,17 @@ class ResourceUpdater {
 
 class ScopedResourceUpdater {
  public:
-  ScopedResourceUpdater(const WCHAR* filename, bool deleteOld);
+  ScopedResourceUpdater(std::wstring filename, bool deleteOld);
   ~ScopedResourceUpdater();
 
   HANDLE Get() const;
   bool Commit();
 
  private:
   bool EndUpdate(bool doesCommit);
+  bool UpdateChecksum();
 
+  std::wstring filename_;
   HANDLE handle_;
   bool commited_ = false;
 };