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