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