#include <efi.h>
#include <efigpt.h>
#include <efilib.h>

#include "pe_parser.h"
#include "ukiboot-efi.h"

#define IMAGE_DOS_SIGNATURE 23117
#define IMAGE_NT_SIGNATURE  17744

typedef struct _IMAGE_DOS_HEADER {
    UINT16 e_magic;
    UINT16 e_cblp;
    UINT16 e_cp;
    UINT16 e_crlc;
    UINT16 e_cparhdr;
    UINT16 e_minalloc;
    UINT16 e_maxalloc;
    UINT16 e_ss;
    UINT16 e_sp;
    UINT16 e_csum;
    UINT16 e_ip;
    UINT16 e_cs;
    UINT16 e_lfarlc;
    UINT16 e_ovno;
    UINT16 e_res[4];
    UINT16 e_oemid;
    UINT16 e_oeminfo;
    UINT16 e_res2[10];
    INT32 e_lfanew;
} IMAGE_DOS_HEADER;

typedef struct _IMAGE_FILE_HEADER {
    UINT16 Machine;
    UINT16 NumberOfSections;
    UINT32 TimeDateStamp;
    UINT32 PointerToSymbolTable;
    UINT32 NumberOfSymbols;
    UINT16 SizeOfOptionalHeader;
    UINT16 Characteristics;
} IMAGE_FILE_HEADER;

typedef struct _IMAGE_NT_HEADERS {
    UINT32 Signature;
    IMAGE_FILE_HEADER FileHeader;
} IMAGE_NT_HEADERS;

#define IMAGE_SIZEOF_SHORT_NAME 8

typedef struct _IMAGE_SECTION_HEADER {
    UINT8 Name[IMAGE_SIZEOF_SHORT_NAME];
    union {
	UINT32 PhysicalAddress;
	UINT32 VirtualSize;
    } Misc;
    UINT32 VirtualAddress;
    UINT32 SizeOfRawData;
    UINT32 PointerToRawData;
    UINT32 PointerToRelocations;
    UINT32 PointerToLinenumbers;
    UINT16 NumberOfRelocations;
    UINT16 NumberOfLinenumbers;
    UINT32 Characteristics;
} IMAGE_SECTION_HEADER;

#define IMAGE_DIRECTORY_ENTRY_SECURITY 4
typedef struct _IMAGE_DATA_DIRECTORY {
    UINT32 VirtualAddress;
    UINT32 Size;
} IMAGE_DATA_DIRECTORY;

#define IMAGE_NUMBEROF_DIRECTORY_ENTRIES 16
#define IMAGE_NT_OPTIONAL_HDR64_MAGIC	 0x20b
typedef struct _IMAGE_OPTIONAL_HEADER64 {
    UINT16 Magic;
    UINT8 MajorLinkerVersion;
    UINT8 MinorLinkerVersion;
    UINT32 SizeOfCode;
    UINT32 SizeOfInitializedData;
    UINT32 SizeOfUninitializedData;
    UINT32 AddressOfEntryPoint;
    UINT32 BaseOfCode;
    UINT64 ImageBase;
    UINT32 SectionAlignment;
    UINT32 FileAlignment;
    UINT16 MajorOperatingSystemVersion;
    UINT16 MinorOperatingSystemVersion;
    UINT16 MajorImageVersion;
    UINT16 MinorImageVersion;
    UINT16 MajorSubsystemVersion;
    UINT16 MinorSubsystemVersion;
    UINT32 Win32VersionValue;
    UINT32 SizeOfImage;
    UINT32 SizeOfHeaders;
    UINT32 CheckSum;
    UINT16 Subsystem;
    UINT16 DllCharacteristics;
    UINT64 SizeOfStackReserve;
    UINT64 SizeOfStackCommit;
    UINT64 SizeOfHeapReserve;
    UINT64 SizeOfHeapCommit;
    UINT32 LoaderFlags;
    UINT32 NumberOfRvaAndSizes;
    IMAGE_DATA_DIRECTORY DataDirectory[IMAGE_NUMBEROF_DIRECTORY_ENTRIES];
} IMAGE_OPTIONAL_HEADER64;

EFI_STATUS
get_pe_size(UINT8 *buffer, UINTN size, UINTN *size_out)
{
    if (size < sizeof(IMAGE_DOS_HEADER)) {
	return EFI_INVALID_PARAMETER;
    }

    IMAGE_DOS_HEADER *dos_header = (IMAGE_DOS_HEADER *)buffer;

    if (dos_header->e_magic != IMAGE_DOS_SIGNATURE) {
	log(L"Invalid DOS header");
	return EFI_INVALID_PARAMETER;
    }

    if (dos_header->e_lfanew < (INT32)sizeof(IMAGE_DOS_HEADER)) {
	log(L"Invalid e_lfanew");
	return EFI_INVALID_PARAMETER;
    }

    UINTN nt_header_start = dos_header->e_lfanew;
    UINTN nt_header_size = sizeof(IMAGE_NT_HEADERS);

    if (nt_header_start >= size || nt_header_size > size - nt_header_start) {
	log(L"Buffer too small for NT header");
	return EFI_INVALID_PARAMETER;
    }

    IMAGE_NT_HEADERS *nt_header = (IMAGE_NT_HEADERS *)(buffer + nt_header_start);
    if (nt_header->Signature != IMAGE_NT_SIGNATURE) {
	log(L"Invalid NT header");
	return EFI_INVALID_PARAMETER;
    }

    UINTN optional_header_size = nt_header->FileHeader.SizeOfOptionalHeader;
    UINTN num_sections = nt_header->FileHeader.NumberOfSections;

    if (num_sections > 96) {
	log(L"Too many PE sections");
	return EFI_INVALID_PARAMETER;
    }

    UINTN headers_end = nt_header_start + nt_header_size;
    if (optional_header_size > size - headers_end) {
	log(L"Buffer too small for optional header");
	return EFI_INVALID_PARAMETER;
    }

    IMAGE_OPTIONAL_HEADER64 *optional_header = (IMAGE_OPTIONAL_HEADER64 *)(buffer + headers_end);
    if (optional_header->Magic != IMAGE_NT_OPTIONAL_HDR64_MAGIC) {
	log(L"Wrong optional nt header");
	return EFI_INVALID_PARAMETER;
    }

    UINTN sections_start = headers_end + optional_header_size;
    UINTN section_entry_size = sizeof(IMAGE_SECTION_HEADER);
    UINTN sections_total = num_sections * section_entry_size;

    if (sections_total > size - sections_start) {
	log(L"Buffer too small for section table");
	return EFI_INVALID_PARAMETER;
    }

    UINTN end_of_file = sections_start + sections_total;

    for (UINTN i = 0; i < num_sections; i++) {
	IMAGE_SECTION_HEADER *section =
	    (IMAGE_SECTION_HEADER *)(buffer + sections_start + i * section_entry_size);
	UINTN raw_offset = section->PointerToRawData;
	UINTN raw_size = section->SizeOfRawData;
	if (raw_size > 0 && raw_offset > (UINTN)-1 - raw_size) {
	    log(L"PE section offset + size overflow");
	    return EFI_INVALID_PARAMETER;
	}
	UINTN section_end = raw_offset + raw_size;
	if (section_end > end_of_file) {
	    end_of_file = section_end;
	}
    }

    /* The Attribute Certificate is referenced by IMAGE_DIRECTORY_ENTRY_SECURITY
     * in the optional header. Other entries in the data directory are refering
     * to addresses inside the sections, but the certificates are actually
     * file offsets outside of the regular sections. So we need to extend the file
     * size to fit the certificates as well.
     */
    IMAGE_DATA_DIRECTORY *security_dir =
	&optional_header->DataDirectory[IMAGE_DIRECTORY_ENTRY_SECURITY];
    if (security_dir->VirtualAddress != 0) {
	UINTN sec_start = security_dir->VirtualAddress;
	UINTN sec_size = security_dir->Size;
	if (sec_size > 0 && sec_start > (UINTN)-1 - sec_size) {
	    log(L"Security directory offset + size overflow");
	    return EFI_INVALID_PARAMETER;
	}
	UINTN sec_end = sec_start + sec_size;
	if (sec_end > end_of_file) {
	    end_of_file = sec_end;
	}
    }

    *size_out = end_of_file;
    return EFI_SUCCESS;
}
