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