- Fix various warnings in inflib (thanks to Ged for help with what appears to be...
[reactos.git] / reactos / tools / pefixup.c
1 /*
2 * PE Fixup Utility
3 * Copyright (C) 2005 Filip Navara
4 *
5 * The purpose of this utility is fix PE binaries generated by binutils and
6 * to manipulate flags that can't be set by binutils.
7 *
8 * Currently two features are implemented:
9 *
10 * - Setting flags on PE sections for use by drivers. The sections
11 * .text, .data, .idata, .bss are marked as non-pageable and
12 * non-discarable, section PAGE is marked as pageable and section
13 * INIT is marked as discaradable.
14 *
15 * - Sorting of export name table in executables. DLLTOOL has bug
16 * in sorting algorithm when the --kill-at flag is used. The exports
17 * are sorted in the decorated form and so the fastcall symbols are
18 * incorrectly put at the beginning of export table. This option
19 * allow to correct sort the table, so binary search can be used
20 * to process them.
21 */
22
23 #include <stdio.h>
24 #include <stdlib.h>
25 #include <string.h>
26 #include <fcntl.h>
27 #ifndef _WIN32
28 #include <unistd.h>
29 #endif
30
31 #ifndef O_BINARY
32 #define O_BINARY 0
33 #endif
34
35 /* The following definitions are ripped from MinGW W32API headers. We don't
36 use these headers directly in order to allow compilation on Linux hosts. */
37
38 typedef unsigned char BYTE, *PBYTE;
39 typedef unsigned short WORD;
40 typedef unsigned int DWORD;
41 typedef int LONG;
42 typedef long LONG_PTR;
43
44 #define IMAGE_NUMBEROF_DIRECTORY_ENTRIES 16
45 #define IMAGE_SIZEOF_SHORT_NAME 8
46 #define IMAGE_DOS_SIGNATURE 0x5A4D
47 #define IMAGE_NT_SIGNATURE 0x00004550
48 #define IMAGE_SCN_MEM_DISCARDABLE 0x2000000
49 #define IMAGE_SCN_MEM_NOT_PAGED 0x8000000
50 #define FIELD_OFFSET(t,f) ((LONG)(LONG_PTR)&(((t*)0)->f))
51 #define IMAGE_FIRST_SECTION(h) ((PIMAGE_SECTION_HEADER) ((unsigned long)h+FIELD_OFFSET(IMAGE_NT_HEADERS,OptionalHeader)+((PIMAGE_NT_HEADERS)(h))->FileHeader.SizeOfOptionalHeader))
52 #define IMAGE_DIRECTORY_ENTRY_EXPORT 0
53
54 #pragma pack(2)
55 typedef struct _IMAGE_DOS_HEADER {
56 WORD e_magic;
57 WORD e_cblp;
58 WORD e_cp;
59 WORD e_crlc;
60 WORD e_cparhdr;
61 WORD e_minalloc;
62 WORD e_maxalloc;
63 WORD e_ss;
64 WORD e_sp;
65 WORD e_csum;
66 WORD e_ip;
67 WORD e_cs;
68 WORD e_lfarlc;
69 WORD e_ovno;
70 WORD e_res[4];
71 WORD e_oemid;
72 WORD e_oeminfo;
73 WORD e_res2[10];
74 LONG e_lfanew;
75 } IMAGE_DOS_HEADER,*PIMAGE_DOS_HEADER;
76 #pragma pack(4)
77 #pragma pack(4)
78 typedef struct _IMAGE_EXPORT_DIRECTORY {
79 DWORD Characteristics;
80 DWORD TimeDateStamp;
81 WORD MajorVersion;
82 WORD MinorVersion;
83 DWORD Name;
84 DWORD Base;
85 DWORD NumberOfFunctions;
86 DWORD NumberOfNames;
87 DWORD AddressOfFunctions;
88 DWORD AddressOfNames;
89 DWORD AddressOfNameOrdinals;
90 } IMAGE_EXPORT_DIRECTORY,*PIMAGE_EXPORT_DIRECTORY;
91 typedef struct _IMAGE_FILE_HEADER {
92 WORD Machine;
93 WORD NumberOfSections;
94 DWORD TimeDateStamp;
95 DWORD PointerToSymbolTable;
96 DWORD NumberOfSymbols;
97 WORD SizeOfOptionalHeader;
98 WORD Characteristics;
99 } IMAGE_FILE_HEADER, *PIMAGE_FILE_HEADER;
100 typedef struct _IMAGE_DATA_DIRECTORY {
101 DWORD VirtualAddress;
102 DWORD Size;
103 } IMAGE_DATA_DIRECTORY,*PIMAGE_DATA_DIRECTORY;
104 typedef struct _IMAGE_OPTIONAL_HEADER {
105 WORD Magic;
106 BYTE MajorLinkerVersion;
107 BYTE MinorLinkerVersion;
108 DWORD SizeOfCode;
109 DWORD SizeOfInitializedData;
110 DWORD SizeOfUninitializedData;
111 DWORD AddressOfEntryPoint;
112 DWORD BaseOfCode;
113 DWORD BaseOfData;
114 DWORD ImageBase;
115 DWORD SectionAlignment;
116 DWORD FileAlignment;
117 WORD MajorOperatingSystemVersion;
118 WORD MinorOperatingSystemVersion;
119 WORD MajorImageVersion;
120 WORD MinorImageVersion;
121 WORD MajorSubsystemVersion;
122 WORD MinorSubsystemVersion;
123 DWORD Reserved1;
124 DWORD SizeOfImage;
125 DWORD SizeOfHeaders;
126 DWORD CheckSum;
127 WORD Subsystem;
128 WORD DllCharacteristics;
129 DWORD SizeOfStackReserve;
130 DWORD SizeOfStackCommit;
131 DWORD SizeOfHeapReserve;
132 DWORD SizeOfHeapCommit;
133 DWORD LoaderFlags;
134 DWORD NumberOfRvaAndSizes;
135 IMAGE_DATA_DIRECTORY DataDirectory[IMAGE_NUMBEROF_DIRECTORY_ENTRIES];
136 } IMAGE_OPTIONAL_HEADER,*PIMAGE_OPTIONAL_HEADER;
137 typedef struct _IMAGE_NT_HEADERS {
138 DWORD Signature;
139 IMAGE_FILE_HEADER FileHeader;
140 IMAGE_OPTIONAL_HEADER OptionalHeader;
141 } IMAGE_NT_HEADERS,*PIMAGE_NT_HEADERS;
142 typedef struct _IMAGE_SECTION_HEADER {
143 BYTE Name[IMAGE_SIZEOF_SHORT_NAME];
144 union {
145 DWORD PhysicalAddress;
146 DWORD VirtualSize;
147 } Misc;
148 DWORD VirtualAddress;
149 DWORD SizeOfRawData;
150 DWORD PointerToRawData;
151 DWORD PointerToRelocations;
152 DWORD PointerToLinenumbers;
153 WORD NumberOfRelocations;
154 WORD NumberOfLinenumbers;
155 DWORD Characteristics;
156 } IMAGE_SECTION_HEADER,*PIMAGE_SECTION_HEADER;
157 #pragma pack(4)
158
159 /* End of ripped definitions */
160
161 typedef struct _export_t {
162 DWORD name;
163 WORD ordinal;
164 } export_t;
165
166 unsigned char *buffer;
167 PIMAGE_DOS_HEADER dos_header;
168 PIMAGE_NT_HEADERS nt_header;
169
170 static inline WORD dtohs(WORD in)
171 {
172 PBYTE in_ptr = (PBYTE)&in;
173 return in_ptr[0] | (in_ptr[1] << 8);
174 }
175
176 static inline WORD htods(WORD in)
177 {
178 WORD out;
179 PBYTE out_ptr = (PBYTE)&out;
180 out_ptr[0] = in; out_ptr[1] = in >> 8;
181 return out;
182 }
183
184 static inline DWORD dtohl(DWORD in)
185 {
186 PBYTE in_ptr = (PBYTE)&in;
187 return in_ptr[0] | (in_ptr[1] << 8) | (in_ptr[2] << 16) | (in_ptr[3] << 24);
188 }
189
190 static inline DWORD htodl(DWORD in)
191 {
192 DWORD out;
193 PBYTE out_ptr = (PBYTE)&out;
194 out_ptr[0] = in ; out_ptr[1] = in >> 8;
195 out_ptr[2] = in >> 16; out_ptr[3] = in >> 24;
196 return out;
197 }
198
199 void *rva_to_ptr(DWORD rva)
200 {
201 PIMAGE_SECTION_HEADER section_header;
202 unsigned int i;
203
204 for (i = 0, section_header = IMAGE_FIRST_SECTION(nt_header);
205 i < dtohl(nt_header->FileHeader.NumberOfSections);
206 i++, section_header++)
207 {
208 if (rva >= dtohl(section_header->VirtualAddress) &&
209 rva < dtohl(section_header->VirtualAddress) +
210 dtohl(section_header->Misc.VirtualSize))
211 {
212 return buffer + rva - dtohl(section_header->VirtualAddress) +
213 dtohl(section_header->PointerToRawData);
214 }
215 }
216
217 return NULL;
218 }
219
220 int export_compare_func(const void *a, const void *b)
221 {
222 const export_t *ap = a;
223 const export_t *bp = b;
224 char *an = rva_to_ptr(ap->name);
225 char *bn = rva_to_ptr(bp->name);
226 return strcmp(an, bn);
227 }
228
229 int main(int argc, char **argv)
230 {
231 int fd_in, fd_out;
232 long len;
233 char hdrbuf[4] = { }, elfhdr[4] = { '\177', 'E', 'L', 'F' };
234 PIMAGE_SECTION_HEADER section_header;
235 PIMAGE_DATA_DIRECTORY data_dir;
236 unsigned int i;
237 unsigned long checksum;
238 int fixup_exports = 0;
239 int fixup_sections = 0;
240
241 /*
242 * Process parameters.
243 */
244
245 if (argc < 2)
246 {
247 printf("Usage: %s <filename> <options>\n"
248 "Options:\n"
249 " -sections Sets section flags for PE image.\n"
250 " -exports Sort the names in export table.\n",
251 argv[0]);
252 return 1;
253 }
254
255 for (i = 2; i < argc; i++)
256 {
257 if (!strcmp(argv[i], "-sections"))
258 fixup_sections = 1;
259 else if (!strcmp(argv[i], "-exports"))
260 fixup_exports = 1;
261 else
262 { printf("Invalid option: %s\n", argv[i]); return 1; }
263 }
264
265 if (fixup_sections == 0 && fixup_exports == 0)
266 {
267 printf("Nothing to do.\n");
268 return 0;
269 }
270
271 /*
272 * Read the whole file to memory.
273 */
274
275 fd_in = open(argv[1], O_RDONLY | O_BINARY);
276 if (fd_in == 0)
277 {
278 printf("Can't open input file.\n");
279 return 1;
280 }
281
282 /*
283 * PowerPC ReactOS uses elf, so doesn't need pefixup
284 */
285 len = read(fd_in, hdrbuf, sizeof(elfhdr));
286 if (!memcmp(hdrbuf, elfhdr, sizeof(elfhdr)))
287 return 0;
288
289 len = lseek(fd_in, 0, SEEK_END);
290 if (len < sizeof(IMAGE_DOS_HEADER))
291 {
292 close(fd_in);
293 printf("'%s' isn't a PE image (too short)\n", argv[1]);
294 return 1;
295 }
296
297 /* Lower down we overwrite the byte at len, so here, we need at least
298 * one more byte than len. We'll be guaranteed one or two now. */
299 buffer = malloc((len + 2) & ~1);
300 if (buffer == NULL)
301 {
302 close(fd_in);
303 printf("Not enough memory available.\n");
304 return 1;
305 }
306
307 /* Read the whole input file into a buffer */
308 lseek(fd_in, 0, SEEK_SET);
309 read(fd_in, buffer, len);
310 /* Here is where the block end overwrite was */
311 if (len & 1)
312 buffer[len] = 0;
313
314 close(fd_in);
315
316 /*
317 * Check the headers and save pointers to them.
318 */
319
320 dos_header = (PIMAGE_DOS_HEADER)buffer;
321 nt_header = (PIMAGE_NT_HEADERS)(buffer + dtohl(dos_header->e_lfanew));
322
323 if (dtohs(dos_header->e_magic) != IMAGE_DOS_SIGNATURE ||
324 dtohl(nt_header->Signature) != IMAGE_NT_SIGNATURE)
325 {
326 printf("'%s' isn't a PE image (bad headers)\n", argv[1]);
327 free(buffer);
328 return 1;
329 }
330
331 if (fixup_exports)
332 {
333 /* Sort export directory */
334 data_dir = &nt_header->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT];
335 if (dtohl(data_dir->Size) != 0)
336 {
337 PIMAGE_EXPORT_DIRECTORY export_directory;
338 DWORD *name_ptr;
339 WORD *ordinal_ptr;
340 export_t *exports;
341
342 export_directory = (PIMAGE_EXPORT_DIRECTORY)rva_to_ptr(dtohl(data_dir->VirtualAddress));
343 if (export_directory != NULL)
344 {
345 exports = malloc(sizeof(export_t) * dtohl(export_directory->NumberOfNames));
346 if (exports == NULL)
347 {
348 printf("Not enough memory.\n");
349 free(buffer);
350 return 1;
351 }
352
353 name_ptr = (DWORD *)rva_to_ptr(dtohl(export_directory->AddressOfNames));
354 ordinal_ptr = (WORD *)rva_to_ptr(dtohl(export_directory->AddressOfNameOrdinals));
355
356 for (i = 0; i < dtohl(export_directory->NumberOfNames); i++)
357 {
358 exports[i].name = dtohl(name_ptr[i]);
359 exports[i].ordinal = dtohl(ordinal_ptr[i]);
360 }
361
362 qsort(exports, dtohl(export_directory->NumberOfNames), sizeof(export_t),
363 export_compare_func);
364
365 for (i = 0; i < dtohl(export_directory->NumberOfNames); i++)
366 {
367 name_ptr[i] = htodl(exports[i].name);
368 ordinal_ptr[i] = htodl(exports[i].ordinal);
369 }
370
371 free(exports);
372 }
373 }
374 }
375
376 if (fixup_sections)
377 {
378 /* Update section flags */
379 for (i = 0, section_header = IMAGE_FIRST_SECTION(nt_header);
380 i < dtohl(nt_header->OptionalHeader.NumberOfRvaAndSizes);
381 i++, section_header++)
382 {
383 if (!strcmp((char*)section_header->Name, ".text") ||
384 !strcmp((char*)section_header->Name, ".data") ||
385 !strcmp((char*)section_header->Name, ".idata") ||
386 !strcmp((char*)section_header->Name, ".bss"))
387 {
388 section_header->Characteristics |= htodl(IMAGE_SCN_MEM_NOT_PAGED);
389 section_header->Characteristics &= htodl(~IMAGE_SCN_MEM_DISCARDABLE);
390 }
391 else if (!strcmp((char*)section_header->Name, "INIT"))
392 {
393 section_header->Characteristics |= htodl(IMAGE_SCN_MEM_DISCARDABLE);
394 }
395 else if (!strcmp((char*)section_header->Name, "PAGE"))
396 {
397 section_header->Characteristics |= htodl(IMAGE_SCN_MEM_NOT_PAGED);
398 }
399 }
400 }
401
402 /* Recalculate checksum */
403 nt_header->OptionalHeader.CheckSum = 0;
404 checksum = 0;
405 for (i = 0; i < len; i += 2)
406 {
407 checksum += *(unsigned short *)(buffer + i);
408 checksum = (checksum + (checksum >> 16)) & 0xffff;
409 }
410 checksum += len;
411 nt_header->OptionalHeader.CheckSum = htods(checksum);
412
413 /* Write the output file */
414 fd_out = open(argv[1], O_WRONLY | O_BINARY);
415 write(fd_out, buffer, len);
416 close(fd_out);
417
418 return 0;
419 }