michael@0: /* -*- Mode: C++; tab-width: 2; indent-tabs-mode: nil; c-basic-offset: 2 -*- */
michael@0: 
michael@0: /* This Source Code Form is subject to the terms of the Mozilla Public
michael@0:  * License, v. 2.0. If a copy of the MPL was not distributed with this file,
michael@0:  * You can obtain one at http://mozilla.org/MPL/2.0/. */
michael@0: 
michael@0: // System headers (alphabetical)
michael@0: #include <fcntl.h>
michael@0: #include <io.h>
michael@0: #include <share.h>
michael@0: #include <stdio.h>
michael@0: #include <stdlib.h>
michael@0: #include <sys/stat.h>
michael@0: #include <windows.h>
michael@0: 
michael@0: // Mozilla headers (alphabetical)
michael@0: #include "mozilla/FileUtils.h"  // ScopedClose
michael@0: #include "nsAutoPtr.h"          // nsAutoArrayPtr
michael@0: 
michael@0: /*
michael@0: Icon files are made up of:
michael@0: 
michael@0: IconHeader
michael@0: IconDirEntry1
michael@0: IconDirEntry2
michael@0: ...
michael@0: IconDirEntryN
michael@0: IconData1
michael@0: IconData2
michael@0: ...
michael@0: IconDataN
michael@0: 
michael@0: Each IconData must be added as a new RT_ICON resource to the exe. Then
michael@0: an RT_GROUP_ICON resource must be added that contains an equivalent
michael@0: header:
michael@0: 
michael@0: IconHeader
michael@0: IconResEntry1
michael@0: IconResEntry2
michael@0: ...
michael@0: IconResEntryN
michael@0: */
michael@0: 
michael@0: #pragma pack(push, 2)
michael@0: typedef struct
michael@0: {
michael@0:   WORD Reserved;
michael@0:   WORD ResourceType;
michael@0:   WORD ImageCount;
michael@0: } IconHeader;
michael@0: 
michael@0: typedef struct
michael@0: {
michael@0:   BYTE Width;
michael@0:   BYTE Height;
michael@0:   BYTE Colors;
michael@0:   BYTE Reserved;
michael@0:   WORD Planes;
michael@0:   WORD BitsPerPixel;
michael@0:   DWORD ImageSize;
michael@0:   DWORD ImageOffset;
michael@0: } IconDirEntry;
michael@0: 
michael@0: typedef struct
michael@0: {
michael@0:   BYTE Width;
michael@0:   BYTE Height;
michael@0:   BYTE Colors;
michael@0:   BYTE Reserved;
michael@0:   WORD Planes;
michael@0:   WORD BitsPerPixel;
michael@0:   DWORD ImageSize;
michael@0:   WORD ResourceID;    // This field is the one difference to above
michael@0: } IconResEntry;
michael@0: #pragma pack(pop)
michael@0: 
michael@0: namespace {
michael@0:   /**
michael@0:    * ScopedResourceUpdate is a RAII wrapper for Windows resource updating
michael@0:    *
michael@0:    * Instances |EndUpdateResourceW()| their handles when they go out of scope.
michael@0:    * They pass |TRUE| as the second argument to |EndUpdateResourceW()|, which
michael@0:    * causes the resource update to be aborted (changes are discarded).
michael@0:    */
michael@0:   struct ScopedResourceUpdateTraits
michael@0:   {
michael@0:     typedef HANDLE type;
michael@0:     static type empty() { return nullptr; }
michael@0:     static void release(type handle) {
michael@0:       if(nullptr != handle) {
michael@0:         EndUpdateResourceW(handle, TRUE); // Discard changes
michael@0:       }
michael@0:     }
michael@0:   };
michael@0: 
michael@0:   typedef mozilla::Scoped<ScopedResourceUpdateTraits> ScopedResourceUpdate;
michael@0: };
michael@0: 
michael@0: #ifdef __MINGW32__
michael@0: extern "C"
michael@0: #endif
michael@0: int
michael@0: wmain(int argc, wchar_t** argv)
michael@0: {
michael@0:   if (argc != 3) {
michael@0:     printf("Usage: redit <exe file> <icon file>\n");
michael@0:     return 1;
michael@0:   }
michael@0: 
michael@0:   mozilla::ScopedClose file;
michael@0:   if (0 != _wsopen_s(&file.rwget(),
michael@0:                      argv[2],
michael@0:                      _O_BINARY | _O_RDONLY,
michael@0:                      _SH_DENYWR,
michael@0:                      _S_IREAD)
michael@0:   || (-1 == file)) {
michael@0:     fprintf(stderr, "Unable to open icon file.\n");
michael@0:     return 1;
michael@0:   }
michael@0: 
michael@0:   // Load all the data from the icon file
michael@0:   long filesize = _filelength(file);
michael@0:   nsAutoArrayPtr<BYTE> data(new BYTE[filesize]);
michael@0:   if(!data) {
michael@0:     fprintf(stderr, "Failed to allocate memory for icon file.\n");
michael@0:     return 1;
michael@0:   }
michael@0:   _read(file, data, filesize);
michael@0: 
michael@0:   IconHeader* header = reinterpret_cast<IconHeader*>(data.get());
michael@0: 
michael@0:   // Open the target library for updating
michael@0:   ScopedResourceUpdate updateRes(BeginUpdateResourceW(argv[1], FALSE));
michael@0:   if (nullptr == updateRes) {
michael@0:     fprintf(stderr, "Unable to open library for modification.\n");
michael@0:     return 1;
michael@0:   }
michael@0: 
michael@0:   // Allocate the group resource entry
michael@0:   long groupSize = sizeof(IconHeader)
michael@0:                  + header->ImageCount * sizeof(IconResEntry);
michael@0:   nsAutoArrayPtr<BYTE> group(new BYTE[groupSize]);
michael@0:   if(!group) {
michael@0:     fprintf(stderr, "Failed to allocate memory for new images.\n");
michael@0:     return 1;
michael@0:   }
michael@0:   memcpy(group, data, sizeof(IconHeader));
michael@0: 
michael@0:   IconDirEntry* sourceIcon =
michael@0:                     reinterpret_cast<IconDirEntry*>(data
michael@0:                                                   + sizeof(IconHeader));
michael@0:   IconResEntry* targetIcon =
michael@0:                     reinterpret_cast<IconResEntry*>(group
michael@0:                                                   + sizeof(IconHeader));
michael@0: 
michael@0:   for (int id = 1; id <= header->ImageCount; id++) {
michael@0:     // Add the individual icon
michael@0:     if (!UpdateResourceW(updateRes, RT_ICON, MAKEINTRESOURCE(id),
michael@0:                          MAKELANGID(LANG_NEUTRAL, SUBLANG_NEUTRAL),
michael@0:                          data + sourceIcon->ImageOffset,
michael@0:                          sourceIcon->ImageSize)) {
michael@0:       fprintf(stderr, "Unable to update resource (RT_ICON).\n");
michael@0:       return 1;
michael@0:     }
michael@0:     // Copy the data for this icon
michael@0:     // (note that the structs have different sizes)
michael@0:     memcpy(targetIcon, sourceIcon, sizeof(IconResEntry));
michael@0:     targetIcon->ResourceID = id;
michael@0:     sourceIcon++;
michael@0:     targetIcon++;
michael@0:   }
michael@0: 
michael@0:   if (!UpdateResourceW(updateRes, RT_GROUP_ICON, L"MAINICON",
michael@0:                        MAKELANGID(LANG_NEUTRAL, SUBLANG_NEUTRAL),
michael@0:                        group, groupSize)) {
michael@0:     fprintf(stderr, "Unable to update resource (RT_GROUP_ICON).\n");
michael@0:     return 1;
michael@0:   }
michael@0: 
michael@0:   // Save the modifications
michael@0:   if(!EndUpdateResourceW(updateRes.forget(), FALSE)) {
michael@0:     fprintf(stderr, "Unable to write changes to library.\n");
michael@0:     return 1;
michael@0:   }
michael@0: 
michael@0:   return 0;
michael@0: }