xref: /reactos/sdk/tools/pefixup.c (revision 3adf4508)
1 /*
2  * PE Fixup Utility
3  * Copyright (C) 2005 Filip Navara
4  * Copyright (C) 2020 Mark Jansen
5  *
6  * The purpose of this utility is fix PE binaries generated by binutils and
7  * to manipulate flags that can't be set by binutils.
8  *
9  * Currently one features is implemented:
10  *
11  * - Updating the PE header to use a LOAD_CONFIG,
12  *   when the struct is exported with the name '_load_config_used'
13  */
14 
15 #include <stdio.h>
16 #include <stdlib.h>
17 #include <string.h>
18 
19 // host_includes
20 #include <typedefs.h>
21 #include <pecoff.h>
22 #include "../../dll/win32/dbghelp/compat.h"
23 
24 static const char* g_ApplicationName;
25 static const char* g_Target;
26 
27 enum fixup_mode
28 {
29     MODE_NONE = 0,
30     MODE_LOADCONFIG,
31     MODE_KERNELDRIVER,
32     MODE_WDMDRIVER,
33     MODE_KERNELDLL,
34     MODE_KERNEL
35 };
36 
37 void *rva_to_ptr(unsigned char *buffer, PIMAGE_NT_HEADERS nt_header, DWORD rva)
38 {
39     unsigned int i;
40     PIMAGE_SECTION_HEADER section_header = IMAGE_FIRST_SECTION(nt_header);
41 
42     for (i = 0; i < nt_header->FileHeader.NumberOfSections; i++, section_header++)
43     {
44         if (rva >= section_header->VirtualAddress &&
45             rva < section_header->VirtualAddress + section_header->Misc.VirtualSize)
46         {
47             return buffer + rva - section_header->VirtualAddress + section_header->PointerToRawData;
48         }
49     }
50 
51     return NULL;
52 }
53 
54 static void error(const char* message, ...)
55 {
56     va_list args;
57 
58     fprintf(stderr, "%s ERROR: '%s': ", g_ApplicationName, g_Target);
59 
60     va_start(args, message);
61     vfprintf(stderr, message, args);
62     va_end(args);
63 }
64 
65 static void fix_checksum(unsigned char *buffer, size_t len, PIMAGE_NT_HEADERS nt_header)
66 {
67     unsigned int checksum = 0;
68     size_t n;
69 
70     nt_header->OptionalHeader.CheckSum = 0;
71 
72     for (n = 0; n < len; n += 2)
73     {
74         checksum += *(unsigned short *)(buffer + n);
75         checksum = (checksum + (checksum >> 16)) & 0xffff;
76     }
77 
78     checksum += (unsigned int)len;
79     nt_header->OptionalHeader.CheckSum = checksum;
80 }
81 
82 static int add_loadconfig(unsigned char *buffer, PIMAGE_NT_HEADERS nt_header)
83 {
84     PIMAGE_DATA_DIRECTORY export_dir;
85     PIMAGE_EXPORT_DIRECTORY export_directory;
86     PDWORD name_ptr, function_ptr;
87     PWORD ordinal_ptr;
88     DWORD n;
89 
90     export_dir = &nt_header->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_EXPORT];
91     if (export_dir->Size == 0)
92     {
93         error("No export directory\n");
94         return 1;
95     }
96 
97     export_directory = rva_to_ptr(buffer, nt_header, export_dir->VirtualAddress);
98     if (export_directory == NULL)
99     {
100         error("Invalid rva for export directory\n");
101         return 1;
102     }
103 
104     name_ptr = rva_to_ptr(buffer, nt_header, export_directory->AddressOfNames);
105     ordinal_ptr = rva_to_ptr(buffer, nt_header, export_directory->AddressOfNameOrdinals);
106     function_ptr = rva_to_ptr(buffer, nt_header, export_directory->AddressOfFunctions);
107 
108     for (n = 0; n < export_directory->NumberOfNames; n++)
109     {
110         const char* name = rva_to_ptr(buffer, nt_header, name_ptr[n]);
111         if (!strcmp(name, "_load_config_used"))
112         {
113             PIMAGE_DATA_DIRECTORY load_config_dir;
114             DWORD load_config_rva = function_ptr[ordinal_ptr[n]];
115             PDWORD load_config_ptr = rva_to_ptr(buffer, nt_header, load_config_rva);
116 
117             /* Update the DataDirectory pointer / size
118                The first entry of the LOAD_CONFIG struct is the size, use that as DataDirectory.Size */
119             load_config_dir = &nt_header->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_LOAD_CONFIG];
120             load_config_dir->VirtualAddress = load_config_rva;
121             load_config_dir->Size = *load_config_ptr;
122 
123             return 0;
124         }
125     }
126 
127     error("Export '_load_config_used' not found\n");
128     return 1;
129 }
130 
131 static int driver_fixup(enum fixup_mode mode, unsigned char *buffer, PIMAGE_NT_HEADERS nt_header)
132 {
133     /* GNU LD just doesn't know what a driver is, and has notably no idea of paged vs non-paged sections */
134     for (unsigned int i = 0; i < nt_header->FileHeader.NumberOfSections; i++)
135     {
136         PIMAGE_SECTION_HEADER Section = IMAGE_FIRST_SECTION(nt_header) + i;
137 
138         /* LD puts alignment crap that nobody asked for */
139         Section->Characteristics &= ~IMAGE_SCN_ALIGN_MASK;
140 
141         /* LD overdoes it and puts the initialized flag everywhere */
142         if (Section->Characteristics & IMAGE_SCN_CNT_CODE)
143             Section->Characteristics &= ~IMAGE_SCN_CNT_INITIALIZED_DATA;
144 
145         if (strncmp((char*)Section->Name, ".rsrc", 5) == 0)
146         {
147             /* .rsrc is discardable for driver images, WDM drivers and Kernel-Mode DLLs */
148             if (mode == MODE_KERNELDRIVER || mode == MODE_WDMDRIVER || mode == MODE_KERNELDLL)
149             {
150                 Section->Characteristics |= IMAGE_SCN_MEM_DISCARDABLE;
151             }
152 
153             /* For some reason, .rsrc is made writable by windres */
154             Section->Characteristics &= ~IMAGE_SCN_MEM_WRITE;
155             continue;
156         }
157 
158         /* Known sections which can be discarded */
159         if (strncmp((char*)Section->Name, "INIT", 4) == 0)
160         {
161             Section->Characteristics |= IMAGE_SCN_MEM_DISCARDABLE;
162             continue;
163         }
164 
165         /* Known sections which can be paged */
166         if ((strncmp((char*)Section->Name, "PAGE", 4) == 0)
167             || (strncmp((char*)Section->Name, ".rsrc", 5) == 0)
168             || (strncmp((char*)Section->Name, ".edata", 6) == 0)
169             || (strncmp((char*)Section->Name, ".reloc", 6) == 0))
170         {
171             continue;
172         }
173 
174         /* If it's discardable, don't set the flag */
175         if (Section->Characteristics & IMAGE_SCN_MEM_DISCARDABLE)
176             continue;
177 
178         Section->Characteristics |= IMAGE_SCN_MEM_NOT_PAGED;
179     }
180 
181     return 0;
182 }
183 
184 /* NOTE: This function tokenizes its parameter in place.
185  * Its format is: name[=newname][,[[!]{CDEIKOMPRSUW}][A{1248PTSX}]] */
186 static int change_section_attribs(char* section_attribs, unsigned char* buffer, PIMAGE_NT_HEADERS nt_header)
187 {
188     char *sectionName, *newSectionName;
189     char *attribsSpec, *alignSpec;
190     char *ptr;
191     DWORD dwAttribs[2]; /* Attributes to [0]: set; [1]: filter */
192     unsigned int i;
193     PIMAGE_SECTION_HEADER SectionTable, Section;
194 
195     if (!section_attribs || !*section_attribs)
196     {
197         error("Section attributes specification is empty.\n");
198         return 1;
199     }
200 
201     sectionName = section_attribs;
202 
203     /* Find the optional new section name and attributes specifications */
204     newSectionName = strchr(section_attribs, '=');
205     attribsSpec = strchr(section_attribs, ',');
206     if (newSectionName && attribsSpec)
207     {
208         /* The attributes specification must be after the new name */
209         if (!(newSectionName < attribsSpec))
210         {
211             error("Invalid section attributes specification.\n");
212             return 1;
213         }
214     }
215     if (newSectionName)
216         *newSectionName++ = 0;
217     if (attribsSpec)
218         *attribsSpec++ = 0;
219 
220     /* An original section name must be specified */
221     if (!*sectionName)
222     {
223         error("Invalid section attributes specification.\n");
224         return 1;
225     }
226     /* If a new section name begins, it must be not empty */
227     if (newSectionName && !*newSectionName)
228     {
229         error("Invalid section attributes specification.\n");
230         return 1;
231     }
232 
233     /* The alignment specification is inside the attributes specification */
234     alignSpec = NULL;
235     if (attribsSpec)
236     {
237         /* Check for the first occurrence of the 'A' separator, case-insensitive */
238         for (ptr = attribsSpec; *ptr; ++ptr)
239         {
240             if (toupper(*ptr) == 'A')
241             {
242                 alignSpec = ptr;
243                 break;
244             }
245         }
246     }
247     if (alignSpec)
248         *alignSpec++ = 0;
249 
250     /* But it's not supported at the moment! */
251     if (alignSpec && *alignSpec)
252     {
253         fprintf(stdout, "%s WARNING: '%s': %s", g_ApplicationName, g_Target,
254                 "Section alignment specification not currently supported! Ignoring.\n");
255     }
256 
257     /* Parse the attributes specification */
258     dwAttribs[0] = dwAttribs[1] = 0;
259     if (attribsSpec && *attribsSpec)
260     {
261         for (i = 0, ptr = attribsSpec; *ptr; ++ptr)
262         {
263             if (*ptr == '!')
264             {
265                 /* The next attribute needs to be removed.
266                  * Any successive '!' gets collapsed. */
267                 i = 1;
268                 continue;
269             }
270 
271             switch (toupper(*ptr))
272             {
273             case 'C':
274                 dwAttribs[i%2] |= IMAGE_SCN_CNT_CODE;
275                 break;
276             case 'D':
277                 dwAttribs[i%2] |= IMAGE_SCN_MEM_DISCARDABLE;
278                 break;
279             case 'E':
280                 dwAttribs[i%2] |= IMAGE_SCN_MEM_EXECUTE;
281                 break;
282             case 'I':
283                 dwAttribs[i%2] |= IMAGE_SCN_CNT_INITIALIZED_DATA;
284                 break;
285             case 'K': /* Remove the not-cached attribute */
286                 dwAttribs[(i+1)%2] |= IMAGE_SCN_MEM_NOT_CACHED;
287                 break;
288             case 'M':
289                 dwAttribs[i%2] |= IMAGE_SCN_LNK_REMOVE;
290                 break;
291             case 'O':
292                 dwAttribs[i%2] |= IMAGE_SCN_LNK_INFO;
293                 break;
294             case 'P': /* Remove the not-paged attribute */
295                 dwAttribs[(i+1)%2] |= IMAGE_SCN_MEM_NOT_PAGED;
296                 break;
297             case 'R':
298                 dwAttribs[i%2] |= IMAGE_SCN_MEM_READ;
299                 break;
300             case 'S':
301                 dwAttribs[i%2] |= IMAGE_SCN_MEM_SHARED;
302                 break;
303             case 'U':
304                 dwAttribs[i%2] |= IMAGE_SCN_CNT_UNINITIALIZED_DATA;
305                 break;
306             case 'W':
307                 dwAttribs[i%2] |= IMAGE_SCN_MEM_WRITE;
308                 break;
309 
310             default:
311                 error("Invalid section attributes specification.\n");
312                 return 1;
313             }
314 
315             /* Got an attribute; reset the state */
316             i = 0;
317         }
318         /* If the state was not reset, the attributes specification is invalid */
319         if (i != 0)
320         {
321             error("Invalid section attributes specification.\n");
322             return 1;
323         }
324     }
325 
326     /* Find all sections with the given name, rename them and change their attributes */
327     Section = NULL;
328     SectionTable = IMAGE_FIRST_SECTION(nt_header);
329     for (i = 0; i < nt_header->FileHeader.NumberOfSections; ++i)
330     {
331         if (strncmp((char*)SectionTable[i].Name, sectionName, ARRAY_SIZE(SectionTable[i].Name)) != 0)
332             continue;
333 
334         Section = &SectionTable[i];
335 
336         if (newSectionName && *newSectionName)
337         {
338             memset(Section->Name, 0, sizeof(Section->Name));
339             strncpy((char*)Section->Name, newSectionName, ARRAY_SIZE(Section->Name));
340         }
341 
342         /* First filter attributes out, then add the new ones.
343          * The new attributes override any removed ones. */
344         Section->Characteristics &= ~dwAttribs[1];
345         Section->Characteristics |=  dwAttribs[0];
346     }
347 
348     /* If no section was found, return an error */
349     if (!Section)
350     {
351         error("Section '%s' does not exist.\n", sectionName);
352         return 1;
353     }
354 
355     return 0;
356 }
357 
358 static
359 void
360 print_usage(void)
361 {
362     printf("Usage: %s <options> <filename>\n\n", g_ApplicationName);
363     printf("<options> can be one of the following options:\n"
364            "  --loadconfig          Fix the LOAD_CONFIG directory entry;\n"
365            "  --kernelmodedriver    Fix code, data and resource sections for driver images;\n"
366            "  --wdmdriver           Fix code, data and resource sections for WDM drivers;\n"
367            "  --kerneldll           Fix code, data and resource sections for Kernel-Mode DLLs;\n"
368            "  --kernel              Fix code, data and resource sections for kernels;\n"
369            "\n"
370            "and/or a combination of the following ones:\n"
371            "  --section:name[=newname][,[[!]{CDEIKOMPRSUW}][A{1248PTSX}]]\n"
372            "                        Overrides the attributes of a section, optionally\n"
373            "                        changing its name and its alignment.\n");
374 }
375 
376 int main(int argc, char **argv)
377 {
378     int result = 1;
379     enum fixup_mode mode = MODE_NONE;
380     int i;
381     FILE* file;
382     size_t len;
383     unsigned char *buffer;
384     PIMAGE_DOS_HEADER dos_header;
385     PIMAGE_NT_HEADERS nt_header;
386 
387     g_ApplicationName = argv[0];
388 
389     /* Check for options */
390     for (i = 1; i < argc; ++i)
391     {
392         if (!(argv[i][0] == '-' && argv[i][1] == '-'))
393         {
394             /* We are out of options (they come first before
395              * anything else, and cannot come after). */
396             break;
397         }
398 
399         if (strcmp(&argv[i][2], "loadconfig") == 0)
400         {
401             if (mode != MODE_NONE)
402                 goto mode_error;
403             mode = MODE_LOADCONFIG;
404         }
405         else if (strcmp(&argv[i][2], "kernelmodedriver") == 0)
406         {
407             if (mode != MODE_NONE)
408                 goto mode_error;
409             mode = MODE_KERNELDRIVER;
410         }
411         else if (strcmp(&argv[i][2], "wdmdriver") == 0)
412         {
413             if (mode != MODE_NONE)
414                 goto mode_error;
415             mode = MODE_WDMDRIVER;
416         }
417         else if (strcmp(&argv[i][2], "kerneldll") == 0)
418         {
419             if (mode != MODE_NONE)
420                 goto mode_error;
421             mode = MODE_KERNELDLL;
422         }
423         else if (strcmp(&argv[i][2], "kernel") == 0)
424         {
425             if (mode != MODE_NONE)
426                 goto mode_error;
427             mode = MODE_KERNEL;
428         }
429         else if (strncmp(&argv[i][2], "section:", 8) == 0)
430         {
431             /* Section attributes override, will be handled later */
432         }
433         else
434         {
435             fprintf(stderr, "%s ERROR: Unknown option: '%s'.\n", g_ApplicationName, argv[i]);
436             goto failure;
437     mode_error:
438             fprintf(stderr, "%s ERROR: Specific mode already set.\n", g_ApplicationName);
439     failure:
440             print_usage();
441             return 1;
442         }
443     }
444     /* Stop now if we don't have any option or file */
445     if ((i <= 1) || (i >= argc))
446     {
447         print_usage();
448         return 1;
449     }
450 
451     g_Target = argv[i];
452 
453     /* Read the whole file to memory */
454     file = fopen(g_Target, "r+b");
455     if (!file)
456     {
457         fprintf(stderr, "%s ERROR: Can't open '%s'.\n", g_ApplicationName, g_Target);
458         return 1;
459     }
460 
461     fseek(file, 0, SEEK_END);
462     len = ftell(file);
463     if (len < sizeof(IMAGE_DOS_HEADER))
464     {
465         fclose(file);
466         error("Image size too small to be a PE image\n");
467         return 1;
468     }
469 
470     /* Add one byte extra for the case where the input file size is odd.
471        We rely on this in our checksum calculation. */
472     buffer = calloc(len + 1, 1);
473     if (buffer == NULL)
474     {
475         fclose(file);
476         error("Not enough memory available (Needed %lu bytes).\n", len + 1);
477         return 1;
478     }
479 
480     /* Read the whole input file into a buffer */
481     fseek(file, 0, SEEK_SET);
482     fread(buffer, 1, len, file);
483 
484     /* Check the headers and save pointers to them */
485     dos_header = (PIMAGE_DOS_HEADER)buffer;
486     if (dos_header->e_magic != IMAGE_DOS_SIGNATURE)
487     {
488         error("Invalid DOS signature: %x\n", dos_header->e_magic);
489         goto Quit;
490     }
491 
492     nt_header = (PIMAGE_NT_HEADERS)(buffer + dos_header->e_lfanew);
493     if (nt_header->Signature != IMAGE_NT_SIGNATURE)
494     {
495         error("Invalid PE signature: %x\n", nt_header->Signature);
496         goto Quit;
497     }
498 
499     result = 0;
500 
501     /* Apply mode fixups */
502     if (mode != MODE_NONE)
503     {
504         if (mode == MODE_LOADCONFIG)
505             result = add_loadconfig(buffer, nt_header);
506         else
507             result = driver_fixup(mode, buffer, nt_header);
508     }
509 
510     /* Apply any section attributes override */
511     for (i = 1; (i < argc) && (result == 0); ++i)
512     {
513         /* Ignore anything but the section specifications */
514         if (!(argv[i][0] == '-' && argv[i][1] == '-'))
515             break;
516         if (strncmp(&argv[i][2], "section:", 8) != 0)
517             continue;
518 
519         result = change_section_attribs(&argv[i][10], buffer, nt_header);
520     }
521 
522     if (!result)
523     {
524         /* Success. Recalculate the checksum only if this is not a reproducible build file */
525         if (nt_header->OptionalHeader.CheckSum != 0)
526             fix_checksum(buffer, len, nt_header);
527 
528         /* We could optimize by only writing the changed parts, but keep it simple for now */
529         fseek(file, 0, SEEK_SET);
530         fwrite(buffer, 1, len, file);
531     }
532 
533 Quit:
534     free(buffer);
535     fclose(file);
536 
537     return result;
538 }
539