te-gen

main.c

/* SPDX-License-Identifier: MIT */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <fcntl.h>
#include <unistd.h>
#include <sys/stat.h>
#include <sys/mman.h>

#include "pe32p.h"
#include "te.h"

#define ALIGN_UP(val, align) ((val + align - 1) & ~(align - 1))
#define ALIGN_DOWN(val, align) (val & ~(align - 1))

int main(int argc, char *argv[])
{
  int fd;
  struct stat sb;
  void *buf;
  IMAGE_DOS_HEADER *doshdr;
  IMAGE_NT_HEADERS *nthdr;
  IMAGE_FILE_HEADER *fhdr;
  IMAGE_OPTIONAL_HEADER *opthdr;
  IMAGE_SECTION_HEADER *sechdr;
  EFI_TE_IMAGE_HEADER *tehdr;
  uint32_t file_align;
  uint32_t sec_align;
  size_t pehdr_size;
  size_t te_size;

  if (argc < 3) {
    fprintf(stderr, "Usage: %s INPUT OUTPUT\n", argv[0]);
    return 1;
  }

  fd = open(argv[1], O_RDONLY);
  if (fd < 0) {
    fprintf(stderr, "Could not open file: %s\n", argv[1]);
    return 1;
  }

  fstat(fd, &sb);
  buf = mmap(NULL, sb.st_size, PROT_READ, MAP_SHARED, fd, 0);

  close(fd);

  tehdr = (EFI_TE_IMAGE_HEADER *)malloc(sizeof(EFI_TE_IMAGE_HEADER));
  if (!tehdr) {
    fprintf(stderr, "Could not allocate memory\n");
    return 1;
  }

  doshdr = (IMAGE_DOS_HEADER *)buf;
  if (doshdr->e_magic != MAGIC_MZ) {
    fprintf(stderr, "Unsupported file format\n");
    return 1;
  }

  nthdr = (IMAGE_NT_HEADERS *)(buf + doshdr->e_lfanew);
  if (nthdr->Signature != MAGIC_PE) {
    fprintf(stderr, "Unsupported file format\n");
    return 1;
  }

  fhdr = &nthdr->FileHeader;
  opthdr = &nthdr->OptionalHeader;

  file_align = opthdr->FileAlignment;
  sec_align = opthdr->SectionAlignment;

  if (file_align > sec_align) {
    fprintf(stderr, "SectionAlignment is smaller than FileAlignment\n");
    return 1;
  }

  pehdr_size = doshdr->e_lfanew + sizeof(IMAGE_NT_HEADERS) \
    + sizeof(IMAGE_SECTION_HEADER) * fhdr->NumberOfSections;
  pehdr_size = ALIGN_UP(pehdr_size, file_align);

  te_size = ALIGN_UP(sizeof(EFI_TE_IMAGE_HEADER), sec_align);
  for (int i = 0; i < fhdr->NumberOfSections; i++) {
    sechdr = (IMAGE_SECTION_HEADER *)
      (buf + doshdr->e_lfanew + sizeof(IMAGE_NT_HEADERS) \
       + sizeof(IMAGE_SECTION_HEADER) * i);
    te_size += ALIGN_UP(sechdr->SizeOfRawData, sec_align);
  }

  void *te_buf = malloc(te_size);
  if (!te_buf) {
    fprintf(stderr, "Could not allocate memory\n");
    return 1;
  }

  tehdr = (EFI_TE_IMAGE_HEADER *)te_buf;

  /* Fill TE image header */
  tehdr->Signature = EFI_TE_IMAGE_HEADER_SIGNATURE;
  tehdr->Machine = fhdr->Machine;
  tehdr->NumberOfSections = fhdr->NumberOfSections;
  tehdr->Subsystem = opthdr->Subsystem;
  tehdr->StrippedSize = pehdr_size - ALIGN_UP(sizeof(EFI_TE_IMAGE_HEADER), file_align);
  tehdr->AddressOfEntryPoint = opthdr->AddressOfEntryPoint;
  tehdr->BaseOfCode = opthdr->BaseOfCode;
  tehdr->ImageBase = opthdr->ImageBase;

  tehdr->DataDirectory[TE_DIRECTORY_BASE_RELOC].VirtualAddress \
    = opthdr->DataDirectory[PE_DIRECTORY_BASE_RELOC].VirtualAddress;
  tehdr->DataDirectory[TE_DIRECTORY_BASE_RELOC].Size \
    = opthdr->DataDirectory[PE_DIRECTORY_BASE_RELOC].Size;

  tehdr->DataDirectory[TE_DIRECTORY_DEBUG].VirtualAddress \
    = opthdr->DataDirectory[PE_DIRECTORY_DEBUG].VirtualAddress;
  tehdr->DataDirectory[TE_DIRECTORY_DEBUG].Size \
    = opthdr->DataDirectory[PE_DIRECTORY_DEBUG].Size;

  void *ptr = te_buf + ALIGN_UP(sizeof(EFI_TE_IMAGE_HEADER), sec_align);
  for (int i = 0; i < fhdr->NumberOfSections; i++) {
    sechdr = (IMAGE_SECTION_HEADER *)
      (buf + doshdr->e_lfanew + sizeof(IMAGE_NT_HEADERS) \
       + sizeof(IMAGE_SECTION_HEADER) * i);
    memcpy(ptr, buf + pehdr_size + sechdr->PointerToRawData, \
            sechdr->SizeOfRawData);
    ptr += ALIGN_UP(sechdr->SizeOfRawData, sec_align);
  }

  FILE *fp = fopen(argv[2], "wb");
  fwrite(te_buf, te_size, 1, fp);
  fclose(fp);

  return 0;
}