sdk/tools/pefixup.c

   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 }