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