diff options
Diffstat (limited to 'post-process-pe.c')
| -rw-r--r-- | post-process-pe.c | 501 |
1 files changed, 501 insertions, 0 deletions
diff --git a/post-process-pe.c b/post-process-pe.c new file mode 100644 index 00000000..8414a5fa --- /dev/null +++ b/post-process-pe.c @@ -0,0 +1,501 @@ +// SPDX-License-Identifier: BSD-2-Clause-Patent +/* + * post-process-pe.c - fix up timestamps and checksums in broken PE files + * Copyright Peter Jones <pjones@redhat.com> + */ + +#define _GNU_SOURCE 1 + +#include <err.h> +#include <fcntl.h> +#include <getopt.h> +#include <inttypes.h> +#include <limits.h> +#include <stdint.h> +#include <stdio.h> +#include <stdlib.h> +#include <string.h> +#include <sys/mman.h> +#include <sys/param.h> +#include <sys/stat.h> +#include <sys/types.h> +#include <unistd.h> + +#ifndef PAGE_SIZE +#define PAGE_SIZE 4096 +#endif + +static int verbosity; +#define ERROR 0 +#define WARNING 1 +#define INFO 2 +#define NOISE 3 +#define MIN_VERBOSITY ERROR +#define MAX_VERBOSITY NOISE +#define debug(level, ...) \ + ({ \ + if (verbosity >= (level)) { \ + printf("%s():%d: ", __func__, __LINE__); \ + printf(__VA_ARGS__); \ + } \ + 0; \ + }) + +typedef uint8_t UINT8; +typedef uint16_t UINT16; +typedef uint32_t UINT32; +typedef uint64_t UINT64; + +typedef uint16_t CHAR16; + +typedef unsigned long UINTN; + +typedef struct { + UINT32 Data1; + UINT16 Data2; + UINT16 Data3; + UINT8 Data4[8]; +} EFI_GUID; + +#include "include/peimage.h" + +#if defined(__GNUC__) && defined(__GNUC_MINOR__) +#define GNUC_PREREQ(maj, min) \ + ((__GNUC__ << 16) + __GNUC_MINOR__ >= ((maj) << 16) + (min)) +#else +#define GNUC_PREREQ(maj, min) 0 +#endif + +#if defined(__clang__) && defined(__clang_major__) && defined(__clang_minor__) +#define CLANG_PREREQ(maj, min) \ + ((__clang_major__ > (maj)) || \ + (__clang_major__ == (maj) && __clang_minor__ >= (min))) +#else +#define CLANG_PREREQ(maj, min) 0 +#endif + +#if GNUC_PREREQ(5, 1) || CLANG_PREREQ(3, 8) +#define add(a0, a1, s) __builtin_add_overflow(a0, a1, s) +#define sub(s0, s1, d) __builtin_sub_overflow(s0, s1, d) +#define mul(f0, f1, p) __builtin_mul_overflow(f0, f1, p) +#else +#define add(a0, a1, s) \ + ({ \ + (*s) = ((a0) + (a1)); \ + 0; \ + }) +#define sub(s0, s1, d) \ + ({ \ + (*d) = ((s0) - (s1)); \ + 0; \ + }) +#define mul(f0, f1, p) \ + ({ \ + (*p) = ((f0) * (f1)); \ + 0; \ + }) +#endif +#define div(d0, d1, q) \ + ({ \ + unsigned int ret_ = ((d1) == 0); \ + if (ret_ == 0) \ + (*q) = ((d0) / (d1)); \ + ret_; \ + }) + +static int +image_is_64_bit(EFI_IMAGE_OPTIONAL_HEADER_UNION *PEHdr) +{ + /* .Magic is the same offset in all cases */ + if (PEHdr->Pe32Plus.OptionalHeader.Magic == + EFI_IMAGE_NT_OPTIONAL_HDR64_MAGIC) + return 1; + return 0; +} + +static void +load_pe(const char *const file, void *const data, const size_t datasize, + PE_COFF_LOADER_IMAGE_CONTEXT *ctx) +{ + EFI_IMAGE_DOS_HEADER *DOSHdr = data; + EFI_IMAGE_OPTIONAL_HEADER_UNION *PEHdr = data; + size_t HeaderWithoutDataDir, SectionHeaderOffset, OptHeaderSize; + size_t FileAlignment = 0; + size_t sz0 = 0, sz1 = 0; + uintptr_t loc = 0; + + debug(NOISE, "datasize:%zu sizeof(PEHdr->Pe32):%zu\n", datasize, + sizeof(PEHdr->Pe32)); + if (datasize < sizeof(PEHdr->Pe32)) + errx(1, "%s: Invalid image size %zu (%zu < %zu)", file, + datasize, datasize, sizeof(PEHdr->Pe32)); + + debug(NOISE, + "DOSHdr->e_magic:0x%02hx EFI_IMAGE_DOS_SIGNATURE:0x%02hx\n", + DOSHdr->e_magic, EFI_IMAGE_DOS_SIGNATURE); + if (DOSHdr->e_magic != EFI_IMAGE_DOS_SIGNATURE) + errx(1, + "%s: Invalid DOS header signature 0x%04hx (expected 0x%04hx)", + file, DOSHdr->e_magic, EFI_IMAGE_DOS_SIGNATURE); + + debug(NOISE, "DOSHdr->e_lfanew:%u datasize:%zu\n", DOSHdr->e_lfanew, + datasize); + if (DOSHdr->e_lfanew >= datasize || + add((uintptr_t)data, DOSHdr->e_lfanew, &loc)) + errx(1, "%s: invalid pe header location", file); + + ctx->PEHdr = PEHdr = (EFI_IMAGE_OPTIONAL_HEADER_UNION *)loc; + debug(NOISE, "PE signature:0x%04x EFI_IMAGE_NT_SIGNATURE:0x%04x\n", + PEHdr->Pe32.Signature, EFI_IMAGE_NT_SIGNATURE); + if (PEHdr->Pe32.Signature != EFI_IMAGE_NT_SIGNATURE) + errx(1, "%s: Unsupported image type", file); + + if (image_is_64_bit(PEHdr)) { + debug(NOISE, "image is 64bit\n"); + ctx->NumberOfRvaAndSizes = + PEHdr->Pe32Plus.OptionalHeader.NumberOfRvaAndSizes; + ctx->SizeOfHeaders = + PEHdr->Pe32Plus.OptionalHeader.SizeOfHeaders; + ctx->ImageSize = PEHdr->Pe32Plus.OptionalHeader.SizeOfImage; + ctx->SectionAlignment = + PEHdr->Pe32Plus.OptionalHeader.SectionAlignment; + FileAlignment = PEHdr->Pe32Plus.OptionalHeader.FileAlignment; + OptHeaderSize = sizeof(EFI_IMAGE_OPTIONAL_HEADER64); + } else { + debug(NOISE, "image is 32bit\n"); + ctx->NumberOfRvaAndSizes = + PEHdr->Pe32.OptionalHeader.NumberOfRvaAndSizes; + ctx->SizeOfHeaders = PEHdr->Pe32.OptionalHeader.SizeOfHeaders; + ctx->ImageSize = (UINT64)PEHdr->Pe32.OptionalHeader.SizeOfImage; + ctx->SectionAlignment = + PEHdr->Pe32.OptionalHeader.SectionAlignment; + FileAlignment = PEHdr->Pe32.OptionalHeader.FileAlignment; + OptHeaderSize = sizeof(EFI_IMAGE_OPTIONAL_HEADER32); + } + + if (FileAlignment % 2 != 0) + errx(1, "%s: Invalid file alignment %ld", file, FileAlignment); + + if (FileAlignment == 0) + FileAlignment = 0x200; + if (ctx->SectionAlignment == 0) + ctx->SectionAlignment = PAGE_SIZE; + if (ctx->SectionAlignment < FileAlignment) + ctx->SectionAlignment = FileAlignment; + + ctx->NumberOfSections = PEHdr->Pe32.FileHeader.NumberOfSections; + + debug(NOISE, + "Number of RVAs:%"PRIu64" EFI_IMAGE_NUMBER_OF_DIRECTORY_ENTRIES:%d\n", + ctx->NumberOfRvaAndSizes, EFI_IMAGE_NUMBER_OF_DIRECTORY_ENTRIES); + if (EFI_IMAGE_NUMBER_OF_DIRECTORY_ENTRIES < ctx->NumberOfRvaAndSizes) + errx(1, "%s: invalid number of RVAs (%lu entries, max is %d)", + file, ctx->NumberOfRvaAndSizes, + EFI_IMAGE_NUMBER_OF_DIRECTORY_ENTRIES); + + if (mul(sizeof(EFI_IMAGE_DATA_DIRECTORY), + EFI_IMAGE_NUMBER_OF_DIRECTORY_ENTRIES, &sz0) || + sub(OptHeaderSize, sz0, &HeaderWithoutDataDir) || + sub(PEHdr->Pe32.FileHeader.SizeOfOptionalHeader, + HeaderWithoutDataDir, &sz0) || + mul(ctx->NumberOfRvaAndSizes, sizeof(EFI_IMAGE_DATA_DIRECTORY), + &sz1) || + (sz0 != sz1)) { + if (mul(sizeof(EFI_IMAGE_DATA_DIRECTORY), + EFI_IMAGE_NUMBER_OF_DIRECTORY_ENTRIES, &sz0)) + debug(ERROR, + "sizeof(EFI_IMAGE_DATA_DIRECTORY) * EFI_IMAGE_NUMBER_OF_DIRECTORY_ENTRIES overflows\n"); + else + debug(ERROR, + "sizeof(EFI_IMAGE_DATA_DIRECTORY) * EFI_IMAGE_NUMBER_OF_DIRECTORY_ENTRIES = %zu\n", + sz0); + if (sub(OptHeaderSize, sz0, &HeaderWithoutDataDir)) + debug(ERROR, + "OptHeaderSize (%zu) - HeaderWithoutDataDir (%zu) overflows\n", + OptHeaderSize, HeaderWithoutDataDir); + else + debug(ERROR, + "OptHeaderSize (%zu) - HeaderWithoutDataDir (%zu) = %zu\n", + OptHeaderSize, sz0, HeaderWithoutDataDir); + + if (sub(PEHdr->Pe32.FileHeader.SizeOfOptionalHeader, + HeaderWithoutDataDir, &sz0)) { + debug(ERROR, + "PEHdr->Pe32.FileHeader.SizeOfOptionalHeader (%d) - %zu overflows\n", + PEHdr->Pe32.FileHeader.SizeOfOptionalHeader, + HeaderWithoutDataDir); + } else { + debug(ERROR, + "PEHdr->Pe32.FileHeader.SizeOfOptionalHeader (%d)- %zu = %zu\n", + PEHdr->Pe32.FileHeader.SizeOfOptionalHeader, + HeaderWithoutDataDir, sz0); + } + if (mul(ctx->NumberOfRvaAndSizes, + sizeof(EFI_IMAGE_DATA_DIRECTORY), &sz1)) + debug(ERROR, + "ctx->NumberOfRvaAndSizes (%zu) * sizeof(EFI_IMAGE_DATA_DIRECTORY) overflows\n", + ctx->NumberOfRvaAndSizes); + else + debug(ERROR, + "ctx->NumberOfRvaAndSizes (%zu) * sizeof(EFI_IMAGE_DATA_DIRECTORY) = %zu\n", + ctx->NumberOfRvaAndSizes, sz1); + debug(ERROR, + "space after image header:%zu data directory size:%zu\n", + sz0, sz1); + + errx(1, "%s: image header overflows data directory", file); + } + + if (add(DOSHdr->e_lfanew, sizeof(UINT32), &SectionHeaderOffset) || + add(SectionHeaderOffset, sizeof(EFI_IMAGE_FILE_HEADER), + &SectionHeaderOffset) || + add(SectionHeaderOffset, + PEHdr->Pe32.FileHeader.SizeOfOptionalHeader, + &SectionHeaderOffset)) { + debug(ERROR, "SectionHeaderOffset:%" PRIu32 " + %zu + %zu + %d", + DOSHdr->e_lfanew, sizeof(UINT32), + sizeof(EFI_IMAGE_FILE_HEADER), + PEHdr->Pe32.FileHeader.SizeOfOptionalHeader); + errx(1, "%s: SectionHeaderOffset would overflow", file); + } + + if (sub(ctx->ImageSize, SectionHeaderOffset, &sz0) || + div(sz0, EFI_IMAGE_SIZEOF_SECTION_HEADER, &sz0) || + (sz0 <= ctx->NumberOfSections)) { + debug(ERROR, "(%" PRIu64 " - %zu) / %d > %d\n", ctx->ImageSize, + SectionHeaderOffset, EFI_IMAGE_SIZEOF_SECTION_HEADER, + ctx->NumberOfSections); + errx(1, "%s: image sections overflow image size", file); + } + + if (sub(ctx->SizeOfHeaders, SectionHeaderOffset, &sz0) || + div(sz0, EFI_IMAGE_SIZEOF_SECTION_HEADER, &sz0) || + (sz0 < ctx->NumberOfSections)) { + debug(ERROR, "(%zu - %zu) / %d >= %d\n", ctx->SizeOfHeaders, + SectionHeaderOffset, EFI_IMAGE_SIZEOF_SECTION_HEADER, + ctx->NumberOfSections); + errx(1, "%s: image sections overflow section headers", file); + } + + if (sub((uintptr_t)PEHdr, (uintptr_t)data, &sz0) || + add(sz0, sizeof(EFI_IMAGE_OPTIONAL_HEADER_UNION), &sz0) || + (sz0 > datasize)) { + errx(1, "%s: PE Image size %zu > %zu", file, sz0, datasize); + } + + if (PEHdr->Pe32.FileHeader.Characteristics & EFI_IMAGE_FILE_RELOCS_STRIPPED) + errx(1, "%s: Unsupported image - Relocations have been stripped", file); + + if (image_is_64_bit(PEHdr)) { + ctx->ImageAddress = PEHdr->Pe32Plus.OptionalHeader.ImageBase; + ctx->EntryPoint = + PEHdr->Pe32Plus.OptionalHeader.AddressOfEntryPoint; + ctx->RelocDir = &PEHdr->Pe32Plus.OptionalHeader.DataDirectory + [EFI_IMAGE_DIRECTORY_ENTRY_BASERELOC]; + ctx->SecDir = &PEHdr->Pe32Plus.OptionalHeader.DataDirectory + [EFI_IMAGE_DIRECTORY_ENTRY_SECURITY]; + } else { + ctx->ImageAddress = PEHdr->Pe32.OptionalHeader.ImageBase; + ctx->EntryPoint = + PEHdr->Pe32.OptionalHeader.AddressOfEntryPoint; + ctx->RelocDir = &PEHdr->Pe32.OptionalHeader.DataDirectory + [EFI_IMAGE_DIRECTORY_ENTRY_BASERELOC]; + ctx->SecDir = &PEHdr->Pe32.OptionalHeader.DataDirectory + [EFI_IMAGE_DIRECTORY_ENTRY_SECURITY]; + } + + if (add((uintptr_t)PEHdr, PEHdr->Pe32.FileHeader.SizeOfOptionalHeader, + &loc) || + add(loc, sizeof(UINT32), &loc) || + add(loc, sizeof(EFI_IMAGE_FILE_HEADER), &loc)) + errx(1, "%s: invalid location for first section", file); + + ctx->FirstSection = (EFI_IMAGE_SECTION_HEADER *)loc; + + if (ctx->ImageSize < ctx->SizeOfHeaders) + errx(1, + "%s: Image size %"PRIu64" is smaller than header size %lu", + file, ctx->ImageSize, ctx->SizeOfHeaders); + + if (sub((uintptr_t)ctx->SecDir, (uintptr_t)data, &sz0) || + sub(datasize, sizeof(EFI_IMAGE_DATA_DIRECTORY), &sz1) || + sz0 > sz1) + errx(1, + "%s: security direcory offset %zu past data directory at %zu", + file, sz0, sz1); + + if (ctx->SecDir->VirtualAddress > datasize || + (ctx->SecDir->VirtualAddress == datasize && + ctx->SecDir->Size > 0)) + errx(1, "%s: Security directory extends past end", file); +} + +static void +fix_timestamp(PE_COFF_LOADER_IMAGE_CONTEXT *ctx) +{ + uint32_t ts; + + if (image_is_64_bit(ctx->PEHdr)) { + ts = ctx->PEHdr->Pe32Plus.FileHeader.TimeDateStamp; + } else { + ts = ctx->PEHdr->Pe32.FileHeader.TimeDateStamp; + } + + if (ts != 0) { + debug(INFO, "Updating timestamp from 0x%08x to 0\n", ts); + if (image_is_64_bit(ctx->PEHdr)) { + ctx->PEHdr->Pe32Plus.FileHeader.TimeDateStamp = 0; + } else { + ctx->PEHdr->Pe32.FileHeader.TimeDateStamp = 0; + } + } +} + +static void +fix_checksum(PE_COFF_LOADER_IMAGE_CONTEXT *ctx, void *map, size_t mapsize) +{ + uint32_t old; + uint32_t checksum = 0; + uint16_t word; + uint8_t *data = map; + + if (image_is_64_bit(ctx->PEHdr)) { + old = ctx->PEHdr->Pe32Plus.OptionalHeader.CheckSum; + ctx->PEHdr->Pe32Plus.OptionalHeader.CheckSum = 0; + } else { + old = ctx->PEHdr->Pe32.OptionalHeader.CheckSum; + ctx->PEHdr->Pe32.OptionalHeader.CheckSum = 0; + } + debug(NOISE, "old checksum was 0x%08x\n", old); + + for (size_t i = 0; i < mapsize - 1; i += 2) { + word = (data[i + 1] << 8ul) | data[i]; + checksum += word; + checksum = 0xffff & (checksum + (checksum >> 0x10)); + } + debug(NOISE, "checksum = 0x%08x + 0x%08zx = 0x%08zx\n", checksum, + mapsize, checksum + mapsize); + + checksum += mapsize; + + if (checksum != old) + debug(INFO, "Updating checksum from 0x%08x to 0x%08x\n", + old, checksum); + + if (image_is_64_bit(ctx->PEHdr)) { + ctx->PEHdr->Pe32Plus.OptionalHeader.CheckSum = checksum; + } else { + ctx->PEHdr->Pe32.OptionalHeader.CheckSum = checksum; + } +} + +static void +handle_one(char *f) +{ + int fd; + int rc; + struct stat statbuf; + size_t sz; + void *map; + int failed = 0; + + PE_COFF_LOADER_IMAGE_CONTEXT ctx = { 0, 0 }; + + fd = open(f, O_RDWR | O_EXCL); + if (fd < 0) + err(1, "Could not open \"%s\"", f); + + rc = fstat(fd, &statbuf); + if (rc < 0) + err(1, "Could not stat \"%s\"", f); + + sz = statbuf.st_size; + + map = mmap(NULL, sz, PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0); + if (map == MAP_FAILED) + err(1, "Could not map \"%s\"", f); + + load_pe(f, map, sz, &ctx); + + fix_timestamp(&ctx); + + fix_checksum(&ctx, map, sz); + + rc = msync(map, sz, MS_SYNC); + if (rc < 0) { + warn("msync(%p, %zu, MS_SYNC) failed", map, sz); + failed = 1; + } + munmap(map, sz); + if (rc < 0) { + warn("munmap(%p, %zu) failed", map, sz); + failed = 1; + } + rc = close(fd); + if (rc < 0) { + warn("close(%d) failed", fd); + failed = 1; + } + if (failed) + exit(1); +} + +static void __attribute__((__noreturn__)) usage(int status) +{ + FILE *out = status ? stderr : stdout; + + fprintf(out, + "Usage: post-process-pe [OPTIONS] file0 [file1 [.. fileN]]\n"); + fprintf(out, "Options:\n"); + fprintf(out, " -q Be more quiet\n"); + fprintf(out, " -v Be more verbose\n"); + fprintf(out, " -h Print this help text and exit\n"); + + exit(status); +} + +int main(int argc, char **argv) +{ + int i; + struct option options[] = { + {.name = "help", + .val = '?', + }, + {.name = "usage", + .val = '?', + }, + {.name = "quiet", + .val = 'q', + }, + {.name = "verbose", + .val = 'v', + }, + {.name = ""} + }; + int longindex = -1; + + while ((i = getopt_long(argc, argv, "hqsv", options, &longindex)) != -1) { + switch (i) { + case 'h': + case '?': + usage(longindex == -1 ? 1 : 0); + break; + case 'q': + verbosity = MAX(verbosity - 1, MIN_VERBOSITY); + break; + case 'v': + verbosity = MIN(verbosity + 1, MAX_VERBOSITY); + break; + } + } + + if (optind == argc) + usage(1); + + for (i = optind; i < argc; i++) + handle_one(argv[i]); + + return 0; +} + +// vim:fenc=utf-8:tw=75:noet |