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