1//************************************ bs::framework - Copyright 2018 Marko Pintera **************************************//
2//*********** Licensed under the MIT license. See LICENSE.md for full terms. This notice is not to be removed. ***********//
3#include "Utility/BsIconUtility.h"
4#include "Image/BsPixelData.h"
5#include "Image/BsColor.h"
6#include "Error/BsException.h"
7
8#define MSDOS_SIGNATURE 0x5A4D
9#define PE_SIGNATURE 0x00004550
10#define PE_32BIT_SIGNATURE 0x10B
11#define PE_64BIT_SIGNATURE 0x20B
12#define PE_NUM_DIRECTORY_ENTRIES 16
13#define PE_SECTION_UNINITIALIZED_DATA 0x00000080
14#define PE_IMAGE_DIRECTORY_ENTRY_RESOURCE 2
15#define PE_IMAGE_RT_ICON 3
16
17namespace bs
18{
19 /** MS-DOS header found at the beggining in a PE format file. */
20 struct MSDOSHeader
21 {
22 UINT16 signature;
23 UINT16 lastSize;
24 UINT16 numBlocks;
25 UINT16 numReloc;
26 UINT16 hdrSize;
27 UINT16 minAlloc;
28 UINT16 maxAlloc;
29 UINT16 ss;
30 UINT16 sp;
31 UINT16 checksum;
32 UINT16 ip;
33 UINT16 cs;
34 UINT16 relocPos;
35 UINT16 numOverlay;
36 UINT16 reserved1[4];
37 UINT16 oemId;
38 UINT16 oemInfo;
39 UINT16 reserved2[10];
40 UINT32 lfanew;
41 };
42
43 /** COFF header found in a PE format file. */
44 struct COFFHeader
45 {
46 UINT16 machine;
47 UINT16 numSections;
48 UINT32 timeDateStamp;
49 UINT32 ptrSymbolTable;
50 UINT32 numSymbols;
51 UINT16 sizeOptHeader;
52 UINT16 characteristics;
53 };
54
55 /** Contains address and size of data areas in a PE image. */
56 struct PEDataDirectory
57 {
58 UINT32 virtualAddress;
59 UINT32 size;
60 };
61
62 /** Optional header in a 32-bit PE format file. */
63 struct PEOptionalHeader32
64 {
65 UINT16 signature;
66 UINT8 majorLinkerVersion;
67 UINT8 minorLinkerVersion;
68 UINT32 sizeCode;
69 UINT32 sizeInitializedData;
70 UINT32 sizeUninitializedData;
71 UINT32 addressEntryPoint;
72 UINT32 baseCode;
73 UINT32 baseData;
74 UINT32 baseImage;
75 UINT32 alignmentSection;
76 UINT32 alignmentFile;
77 UINT16 majorOSVersion;
78 UINT16 minorOSVersion;
79 UINT16 majorImageVersion;
80 UINT16 minorImageVersion;
81 UINT16 majorSubsystemVersion;
82 UINT16 minorSubsystemVersion;
83 UINT32 reserved;
84 UINT32 sizeImage;
85 UINT32 sizeHeaders;
86 UINT32 checksum;
87 UINT16 subsystem;
88 UINT16 characteristics;
89 UINT32 sizeStackReserve;
90 UINT32 sizeStackCommit;
91 UINT32 sizeHeapReserve;
92 UINT32 sizeHeapCommit;
93 UINT32 loaderFlags;
94 UINT32 NumRvaAndSizes;
95 PEDataDirectory dataDirectory[16];
96 };
97
98 /** Optional header in a 64-bit PE format file. */
99 struct PEOptionalHeader64
100 {
101 UINT16 signature;
102 UINT8 majorLinkerVersion;
103 UINT8 minorLinkerVersion;
104 UINT32 sizeCode;
105 UINT32 sizeInitializedData;
106 UINT32 sizeUninitializedData;
107 UINT32 addressEntryPoint;
108 UINT32 baseCode;
109 UINT64 baseImage;
110 UINT32 alignmentSection;
111 UINT32 alignmentFile;
112 UINT16 majorOSVersion;
113 UINT16 minorOSVersion;
114 UINT16 majorImageVersion;
115 UINT16 minorImageVersion;
116 UINT16 majorSubsystemVersion;
117 UINT16 minorSubsystemVersion;
118 UINT32 reserved;
119 UINT32 sizeImage;
120 UINT32 sizeHeaders;
121 UINT32 checksum;
122 UINT16 subsystem;
123 UINT16 characteristics;
124 UINT64 sizeStackReserve;
125 UINT64 sizeStackCommit;
126 UINT64 sizeHeapReserve;
127 UINT64 sizeHeapCommit;
128 UINT32 loaderFlags;
129 UINT32 NumRvaAndSizes;
130 PEDataDirectory dataDirectory[16];
131 };
132
133 /** A section header in a PE format file. */
134 struct PESectionHeader
135 {
136 char name[8];
137 UINT32 virtualSize;
138 UINT32 relativeVirtualAddress;
139 UINT32 physicalSize;
140 UINT32 physicalAddress;
141 UINT8 deprecated[12];
142 UINT32 flags;
143 };
144
145 /** A resource table header within a .rsrc section in a PE format file. */
146 struct PEImageResourceDirectory
147 {
148 UINT32 flags;
149 UINT32 timeDateStamp;
150 UINT16 majorVersion;
151 UINT16 minorVersion;
152 UINT16 numNamedEntries;
153 UINT16 numIdEntries;
154 };
155
156 /** A single entry in a resource table within a .rsrc section in a PE format file. */
157 struct PEImageResourceEntry
158 {
159 UINT32 type;
160 UINT32 offsetDirectory : 31;
161 UINT32 isDirectory : 1;
162 };
163
164 /** An entry in a resource table referencing resource data. Found within a .rsrc section in a PE format file. */
165 struct PEImageResourceEntryData
166 {
167 UINT32 offsetData;
168 UINT32 size;
169 UINT32 codePage;
170 UINT32 resourceHandle;
171 };
172
173 /** Header used in icon file format. */
174 struct IconHeader
175 {
176 UINT32 size;
177 INT32 width;
178 INT32 height;
179 UINT16 planes;
180 UINT16 bitCount;
181 UINT32 compression;
182 UINT32 sizeImage;
183 INT32 xPelsPerMeter;
184 INT32 yPelsPerMeter;
185 UINT32 clrUsed;
186 UINT32 clrImportant;
187 };
188
189 void IconUtility::updateIconExe(const Path& path, const Map<UINT32, SPtr<PixelData>>& pixelsPerSize)
190 {
191 // A PE file is structured as such:
192 // - MSDOS Header
193 // - PE Signature
194 // - COFF Header
195 // - PE Optional Header
196 // - One or multiple sections
197 // - .code
198 // - .data
199 // - ...
200 // - .rsrc
201 // - icon/cursor/etc data
202
203 std::fstream stream;
204 stream.open(path.toPlatformString().c_str(), std::ios::in | std::ios::out | std::ios::binary);
205
206 // First check magic number to ensure file is even an executable
207 UINT16 magicNum;
208 stream.read((char*)&magicNum, sizeof(magicNum));
209 if (magicNum != MSDOS_SIGNATURE)
210 BS_EXCEPT(InvalidStateException, "Provided file is not a valid executable.");
211
212 // Read the MSDOS header and skip over it
213 stream.seekg(0);
214
215 MSDOSHeader msdosHeader;
216 stream.read((char*)&msdosHeader, sizeof(MSDOSHeader));
217
218 // Read PE signature
219 stream.seekg(msdosHeader.lfanew);
220
221 UINT32 peSignature;
222 stream.read((char*)&peSignature, sizeof(peSignature));
223
224 if (peSignature != PE_SIGNATURE)
225 BS_EXCEPT(InvalidStateException, "Provided file is not in PE format.");
226
227 // Read COFF header
228 COFFHeader coffHeader;
229 stream.read((char*)&coffHeader, sizeof(COFFHeader));
230
231 if (coffHeader.sizeOptHeader == 0) // .exe files always have an optional header
232 BS_EXCEPT(InvalidStateException, "Provided file is not a valid executable.");
233
234 UINT32 numSectionHeaders = coffHeader.numSections;
235
236 // Read optional header
237 auto optionalHeaderPos = stream.tellg();
238
239 UINT16 optionalHeaderSignature;
240 stream.read((char*)&optionalHeaderSignature, sizeof(optionalHeaderSignature));
241
242 PEDataDirectory* dataDirectory = nullptr;
243 stream.seekg(optionalHeaderPos);
244 if (optionalHeaderSignature == PE_32BIT_SIGNATURE)
245 {
246 PEOptionalHeader32 optionalHeader;
247 stream.read((char*)&optionalHeader, sizeof(optionalHeader));
248
249 dataDirectory = optionalHeader.dataDirectory + PE_IMAGE_DIRECTORY_ENTRY_RESOURCE;
250 }
251 else if (optionalHeaderSignature == PE_64BIT_SIGNATURE)
252 {
253 PEOptionalHeader64 optionalHeader;
254 stream.read((char*)&optionalHeader, sizeof(optionalHeader));
255
256 dataDirectory = optionalHeader.dataDirectory + PE_IMAGE_DIRECTORY_ENTRY_RESOURCE;
257 }
258 else
259 BS_EXCEPT(InvalidStateException, "Unrecognized PE format.");
260
261 // Read section headers
262 auto sectionHeaderPos = optionalHeaderPos + (std::ifstream::pos_type)coffHeader.sizeOptHeader;
263 stream.seekg(sectionHeaderPos);
264
265 PESectionHeader* sectionHeaders = bs_stack_alloc<PESectionHeader>(numSectionHeaders);
266 stream.read((char*)sectionHeaders, sizeof(PESectionHeader) * numSectionHeaders);
267
268 // Look for .rsrc section header
269 std::function<void(PEImageResourceDirectory*, PEImageResourceDirectory*, UINT8*, UINT32)> setIconData =
270 [&](PEImageResourceDirectory* base, PEImageResourceDirectory* current, UINT8* imageData, UINT32 sectionAddress)
271 {
272 UINT32 numEntries = current->numIdEntries; // Not supporting name entries
273 PEImageResourceEntry* entries = (PEImageResourceEntry*)(current + 1);
274
275 for (UINT32 i = 0; i < numEntries; i++)
276 {
277 // Only at root does the type identify resource type
278 if (base == current && entries[i].type != PE_IMAGE_RT_ICON)
279 continue;
280
281 if (entries[i].isDirectory)
282 {
283 PEImageResourceDirectory* child = (PEImageResourceDirectory*)(((UINT8*)base) + entries[i].offsetDirectory);
284 setIconData(base, child, imageData, sectionAddress);
285 }
286 else
287 {
288 PEImageResourceEntryData* data = (PEImageResourceEntryData*)(((UINT8*)base) + entries[i].offsetDirectory);
289
290 UINT8* iconData = imageData + (data->offsetData - sectionAddress);
291 updateIconData(iconData, pixelsPerSize);
292 }
293 }
294 };
295
296 for (UINT32 i = 0; i < numSectionHeaders; i++)
297 {
298 if (sectionHeaders[i].flags & PE_SECTION_UNINITIALIZED_DATA)
299 continue;
300
301 if (strcmp(sectionHeaders[i].name, ".rsrc") == 0)
302 {
303 UINT32 imageSize = sectionHeaders[i].physicalSize;
304 UINT8* imageData = (UINT8*)bs_stack_alloc(imageSize);
305
306 stream.seekg(sectionHeaders[i].physicalAddress);
307 stream.read((char*)imageData, imageSize);
308
309 UINT32 resourceDirOffset = dataDirectory->virtualAddress - sectionHeaders[i].relativeVirtualAddress;
310 PEImageResourceDirectory* resourceDirectory = (PEImageResourceDirectory*)&imageData[resourceDirOffset];
311
312 setIconData(resourceDirectory, resourceDirectory, imageData, sectionHeaders[i].relativeVirtualAddress);
313 stream.seekp(sectionHeaders[i].physicalAddress);
314 stream.write((char*)imageData, imageSize);
315
316 bs_stack_free(imageData);
317 }
318 }
319
320 bs_stack_free(sectionHeaders);
321 stream.close();
322 }
323
324 void IconUtility::updateIconData(UINT8* iconData, const Map<UINT32, SPtr<PixelData>>& pixelsPerSize)
325 {
326 IconHeader* iconHeader = (IconHeader*)iconData;
327
328 if (iconHeader->size != sizeof(IconHeader) || iconHeader->compression != 0
329 || iconHeader->planes != 1 || iconHeader->bitCount != 32)
330 {
331 // Unsupported format
332 return;
333 }
334
335 UINT8* iconPixels = iconData + sizeof(IconHeader);
336 UINT32 width = iconHeader->width;
337 UINT32 height = iconHeader->height / 2;
338
339 auto iterFind = pixelsPerSize.find(width);
340 if (iterFind == pixelsPerSize.end() || iterFind->second->getWidth() != width
341 || iterFind->second->getHeight() != height)
342 {
343 // No icon of this size provided
344 return;
345 }
346
347 // Write colors
348 SPtr<PixelData> srcPixels = iterFind->second;
349 UINT32* colorData = (UINT32*)iconPixels;
350
351 UINT32 idx = 0;
352 for (INT32 y = (INT32)height - 1; y >= 0; y--)
353 {
354 for (UINT32 x = 0; x < width; x++)
355 colorData[idx++] = srcPixels->getColorAt(x, y).getAsBGRA();
356 }
357
358 // Write AND mask
359 UINT32 colorDataSize = width * height * sizeof(UINT32);
360 UINT8* maskData = iconPixels + colorDataSize;
361
362 UINT32 numPackedPixels = width / 8; // One per bit in byte
363
364 for (INT32 y = (INT32)height - 1; y >= 0; y--)
365 {
366 UINT8 mask = 0;
367 for (UINT32 packedX = 0; packedX < numPackedPixels; packedX++)
368 {
369 for (UINT32 pixelIdx = 0; pixelIdx < 8; pixelIdx++)
370 {
371 UINT32 x = packedX * 8 + pixelIdx;
372 Color color = srcPixels->getColorAt(x, y);
373 if (color.a < 0.25f)
374 mask |= 1 << (7 - pixelIdx);
375 }
376
377 *maskData = mask;
378 maskData++;
379 }
380 }
381 }
382}
383