pe.c 16.4 KB raw

#include <stdint.h>
#include <stddef.h>
#include <lib/misc.h>
#include <lib/libc.h>
#include <lib/pe.h>
#include <lib/print.h>
#include <lib/rand.h>
#include <mm/pmm.h>
​
#define FIXED_HIGHER_HALF_OFFSET_64 ((uint64_t)0xffffffff80000000)
​
#define IMAGE_DOS_SIGNATURE 0x5a4d
​
typedef struct _IMAGE_DOS_HEADER {
    uint16_t e_magic;
    uint16_t e_cblp;
    uint16_t e_cp;
    uint16_t e_crlc;
    uint16_t e_cparhdr;
    uint16_t e_minalloc;
    uint16_t e_maxalloc;
    uint16_t e_ss;
    uint16_t e_sp;
    uint16_t e_csum;
    uint16_t e_ip;
    uint16_t e_cs;
    uint16_t e_lfarlc;
    uint16_t e_ovno;
    uint16_t e_res[4];
    uint16_t e_oemid;
    uint16_t e_oeminfo;
    uint16_t e_res2[10];
    uint32_t e_lfanew;
} IMAGE_DOS_HEADER;
​
#define IMAGE_FILE_MACHINE_I386 0x14c
#define IMAGE_FILE_MACHINE_AMD64 0x8664
#define IMAGE_FILE_MACHINE_ARM64 0xaa64
#define IMAGE_FILE_MACHINE_RISCV64 0x5064
#define IMAGE_FILE_MACHINE_LOONGARCH64 0x6264
​
#define IMAGE_FILE_RELOCS_STRIPPED 1
#define IMAGE_FILE_EXECUTABLE_IMAGE 2
​
typedef struct {
    uint16_t Machine;
    uint16_t NumberOfSections;
    uint32_t TimeDateStamp;
    uint32_t PointerToSymbolTable;
    uint32_t NumberOfSymbols;
    uint16_t SizeOfOptionalHeader;
    uint16_t Characteristics;
} IMAGE_FILE_HEADER;
​
typedef struct {
    uint32_t VirtualAddress;
    uint32_t Size;
} IMAGE_DATA_DIRECTORY;
​
#define IMAGE_NT_OPTIONAL_HDR32_MAGIC 0x10b
#define IMAGE_NT_OPTIONAL_HDR64_MAGIC 0x20b
​
#define IMAGE_DIRECTORY_ENTRY_EXPORT 0
#define IMAGE_DIRECTORY_ENTRY_IMPORT 1
#define IMAGE_DIRECTORY_ENTRY_RESOURCE 2
#define IMAGE_DIRECTORY_ENTRY_EXCEPTION 3
#define IMAGE_DIRECTORY_ENTRY_SECURITY 4
#define IMAGE_DIRECTORY_ENTRY_BASERELOC 5
#define IMAGE_DIRECTORY_ENTRY_DEBUG 6
#define IMAGE_DIRECTORY_ENTRY_ARCHITECTURE 7
#define IMAGE_DIRECTORY_ENTRY_GLOBALPTR 8
#define IMAGE_DIRECTORY_ENTRY_TLS 9
#define IMAGE_DIRECTORY_ENTRY_LOAD_CONFIG 10
#define IMAGE_DIRECTORY_ENTRY_BOUND_IMPORT 11
#define IMAGE_DIRECTORY_ENTRY_IAT 12
#define IMAGE_DIRECTORY_ENTRY_DELAY_IMPORT 13
#define IMAGE_DIRECTORY_ENTRY_COM_DESCRIPTOR 14
​
typedef struct {
    uint16_t Magic;
    uint8_t MajorLinkerVersion;
    uint8_t MinorLinkerVersion;
    uint32_t SizeOfCode;
    uint32_t SizeOfInitializedData;
    uint32_t SizeOfUninitializedData;
    uint32_t AddressOfEntryPoint;
    uint32_t BaseOfCode;
    uint64_t ImageBase;
    uint32_t SectionAlignment;
    uint32_t FileAlignment;
    uint16_t MajorOperatingSystemVersion;
    uint16_t MinorOperatingSystemVersion;
    uint16_t MajorImageVersion;
    uint16_t MinorImageVersion;
    uint16_t MajorSubsystemVersion;
    uint16_t MinorSubsystemVersion;
    uint32_t Win32VersionValue;
    uint32_t SizeOfImage;
    uint32_t SizeOfHeaders;
    uint32_t CheckSum;
    uint16_t Subsystem;
    uint16_t DllCharacteristics;
    uint64_t SizeOfStackReserve;
    uint64_t SizeOfStackCommit;
    uint64_t SizeOfHeapReserve;
    uint64_t SizeOfHeapCommit;
    uint32_t LoaderFlags;
    uint32_t NumberOfRvaAndSizes;
    IMAGE_DATA_DIRECTORY DataDirectory[16];
} IMAGE_OPTIONAL_HEADER64;
​
#define IMAGE_NT_SIGNATURE 0x4550
​
typedef struct {
    uint32_t Signature;
    IMAGE_FILE_HEADER FileHeader;
    IMAGE_OPTIONAL_HEADER64 OptionalHeader;
} IMAGE_NT_HEADERS64;
​
#define IMAGE_SCN_MEM_EXECUTE 0x20000000
#define IMAGE_SCN_MEM_READ 0x40000000
#define IMAGE_SCN_MEM_WRITE 0x80000000
​
typedef struct {
    char Name[8];
    uint32_t VirtualSize;
    uint32_t VirtualAddress;
    uint32_t SizeOfRawData;
    uint32_t PointerToRawData;
    uint32_t PointerToRelocations;
    uint32_t PointerToLinenumbers;
    uint16_t NumberOfRelocations;
    uint16_t NumberOfLinenumbers;
    uint32_t Characteristics;
} IMAGE_SECTION_HEADER;
​
typedef struct {
    union {
        uint32_t Characteristics;
        uint32_t OriginalFirstThunk;
    };
    uint32_t TimeDateStamp;
    uint32_t ForwarderChain;
    uint32_t Name;
    uint32_t FirstThunk;
} IMAGE_IMPORT_DESCRIPTOR;
​
#define IMAGE_REL_BASED_ABSOLUTE 0
#define IMAGE_REL_BASED_HIGHLOW 3
#define IMAGE_REL_BASED_DIR64 10
​
typedef struct {
    uint32_t VirtualAddress;
    uint32_t SizeOfBlock;
} IMAGE_BASE_RELOCATION_BLOCK;
​
static void pe64_validate(uint8_t *image, size_t file_size) {
    IMAGE_DOS_HEADER *dos_hdr = (IMAGE_DOS_HEADER *)image;
​
    if (file_size < sizeof(IMAGE_DOS_HEADER)) {
        panic(true, "pe: File too small for DOS header");
    }
​
    if (dos_hdr->e_magic != IMAGE_DOS_SIGNATURE) {
        panic(true, "pe: Not a valid PE file");
    }
​
    if (file_size < sizeof(IMAGE_NT_HEADERS64)) {
        panic(true, "pe: File too small for NT headers");
    }
​
    if (dos_hdr->e_lfanew > file_size - sizeof(IMAGE_NT_HEADERS64)) {
        panic(true, "pe: e_lfanew offset out of bounds");
    }
​
    IMAGE_NT_HEADERS64 *nt_hdrs = (IMAGE_NT_HEADERS64 *)(image + dos_hdr->e_lfanew);
​
    if (nt_hdrs->Signature != IMAGE_NT_SIGNATURE) {
        panic(true, "pe: Not a valid PE file");
    }
​
    if (nt_hdrs->OptionalHeader.Magic != IMAGE_NT_OPTIONAL_HDR64_MAGIC) {
        panic(true, "pe: Not a valid PE32+ file");
    }
​
    if (nt_hdrs->FileHeader.SizeOfOptionalHeader < sizeof(IMAGE_OPTIONAL_HEADER64)) {
        panic(true, "pe: SizeOfOptionalHeader too small");
    }
​
#if defined(__x86_64__) || defined(__i386__)
    if (nt_hdrs->FileHeader.Machine != IMAGE_FILE_MACHINE_AMD64) {
        panic(true, "pe: Not an x86-64 PE file");
    }
#elif defined(__aarch64__)
    if (nt_hdrs->FileHeader.Machine != IMAGE_FILE_MACHINE_ARM64) {
        panic(true, "pe: Not an ARM64 PE file");
    }
#elif defined (__riscv) && (__riscv_xlen == 64)
    if (nt_hdrs->FileHeader.Machine != IMAGE_FILE_MACHINE_RISCV64) {
        panic(true, "pe: Not a RISC-V PE file");
    }
#elif defined (__loongarch__) && (__loongarch_grlen == 64)
    if (nt_hdrs->FileHeader.Machine != IMAGE_FILE_MACHINE_LOONGARCH64) {
        panic(true, "pe: Not a loongarch64 PE file");
    }
#else
#error Unknown architecture
#endif
}
​
int pe_bits(uint8_t *image, size_t image_size) {
    if (image_size < sizeof(IMAGE_DOS_HEADER)) {
        return -1;
    }
​
    IMAGE_DOS_HEADER *dos_hdr = (IMAGE_DOS_HEADER *)image;
​
    if (dos_hdr->e_magic != IMAGE_DOS_SIGNATURE) {
        return -1;
    }
​
    if (image_size < sizeof(IMAGE_NT_HEADERS64)) {
        return -1;
    }
​
    if ((size_t)dos_hdr->e_lfanew > image_size - sizeof(IMAGE_NT_HEADERS64)) {
        return -1;
    }
​
    IMAGE_NT_HEADERS64 *nt_hdrs = (IMAGE_NT_HEADERS64 *)(image + dos_hdr->e_lfanew);
​
    if (nt_hdrs->Signature != IMAGE_NT_SIGNATURE) {
        return -1;
    }
​
    switch (nt_hdrs->FileHeader.Machine) {
        case IMAGE_FILE_MACHINE_I386:
            return 32;
        case IMAGE_FILE_MACHINE_AMD64:
        case IMAGE_FILE_MACHINE_ARM64:
        case IMAGE_FILE_MACHINE_RISCV64:
        case IMAGE_FILE_MACHINE_LOONGARCH64:
            return 64;
    }
​
    return -1;
}
​
bool pe64_load(uint8_t *image, size_t file_size, uint64_t *entry_point, uint64_t *_slide, uint32_t alloc_type, bool kaslr, struct mem_range **_ranges, uint64_t *_ranges_count, uint64_t *physical_base, uint64_t *virtual_base, uint64_t *_image_size, uint64_t *image_size_before_bss, bool *_is_reloc) {
    pe64_validate(image, file_size);
​
    IMAGE_DOS_HEADER *dos_hdr = (IMAGE_DOS_HEADER *)image;
    IMAGE_NT_HEADERS64 *nt_hdrs = (IMAGE_NT_HEADERS64 *)(image + dos_hdr->e_lfanew);
​
    // Validate SizeOfOptionalHeader doesn't cause sections pointer to go out of bounds
    uint64_t sections_offset = (uint64_t)dos_hdr->e_lfanew + sizeof(uint32_t) + sizeof(IMAGE_FILE_HEADER) + nt_hdrs->FileHeader.SizeOfOptionalHeader;
    uint64_t sections_end = sections_offset + (uint64_t)nt_hdrs->FileHeader.NumberOfSections * sizeof(IMAGE_SECTION_HEADER);
    if (sections_end > file_size) {
        panic(true, "pe: Section headers extend beyond file bounds");
    }
​
    IMAGE_SECTION_HEADER *sections = (IMAGE_SECTION_HEADER *)((uintptr_t)&nt_hdrs->OptionalHeader + nt_hdrs->FileHeader.SizeOfOptionalHeader);
​
    bool is_reloc = true;
​
    if (nt_hdrs->FileHeader.Characteristics & IMAGE_FILE_RELOCS_STRIPPED) {
        is_reloc = false;
    }
​
    if (_is_reloc) {
        *_is_reloc = is_reloc;
    }
​
    uint64_t image_base = nt_hdrs->OptionalHeader.ImageBase;
    uint64_t image_size = nt_hdrs->OptionalHeader.SizeOfImage;
    uint64_t alignment = nt_hdrs->OptionalHeader.SectionAlignment;
​
    if (alignment == 0) {
        panic(true, "pe: SectionAlignment is zero");
    }
​
    if (alignment > 1 && (alignment & (alignment - 1)) != 0) {
        panic(true, "pe: SectionAlignment is not a power of 2");
    }
​
    bool lower_to_higher = false;
​
    if (image_base < FIXED_HIGHER_HALF_OFFSET_64) {
        if (!is_reloc) {
            panic(true, "pe: Lower half images are not allowed");
        }
​
        lower_to_higher = true;
    }
​
    uint64_t slide = 0;
    size_t try_count = 0;
    size_t max_simulated_tries = 0x10000;
​
    if (lower_to_higher) {
        slide = FIXED_HIGHER_HALF_OFFSET_64 - image_base;
    }
​
    *physical_base = (uintptr_t)ext_mem_alloc_type_aligned(image_size, alloc_type, alignment);
    *virtual_base = image_base;
​
    // Validate SizeOfHeaders doesn't exceed file size or image size
    if (nt_hdrs->OptionalHeader.SizeOfHeaders > file_size
     || nt_hdrs->OptionalHeader.SizeOfHeaders > image_size) {
        panic(true, "pe: SizeOfHeaders exceeds file or image size");
    }
​
    memcpy((void *)(uintptr_t)*physical_base, image, nt_hdrs->OptionalHeader.SizeOfHeaders);
​
    if (_image_size) {
        *_image_size = image_size;
    }
​
    if (is_reloc && kaslr) {
again:
        slide = (rand32() & ~(alignment - 1)) + (lower_to_higher ? FIXED_HIGHER_HALF_OFFSET_64 - image_base : 0);
​
        if (*virtual_base + slide + image_size < 0xffffffff80000000 /* this comparison relies on overflow */) {
            if (++try_count == max_simulated_tries) {
                panic(true, "pe: Image wants to load too high");
            }
            goto again;
        }
    }
​
    for (size_t i = 0; i < nt_hdrs->FileHeader.NumberOfSections; i++) {
        IMAGE_SECTION_HEADER *section = &sections[i];
​
        uintptr_t section_base = *physical_base + section->VirtualAddress;
        uint32_t section_raw_size = section->VirtualSize < section->SizeOfRawData ? section->VirtualSize : section->SizeOfRawData;
​
        // Validate section doesn't write past the image buffer
        if ((uint64_t)section->VirtualAddress + section_raw_size > image_size) {
            panic(true, "pe: Section %U exceeds image bounds", (uint64_t)i);
        }
​
        // Validate section data doesn't exceed file bounds
        if ((uint64_t)section->PointerToRawData + section_raw_size > file_size) {
            panic(true, "pe: Section %U data extends beyond file bounds", (uint64_t)i);
        }
​
        memcpy((void *)section_base, image + section->PointerToRawData, section_raw_size);
    }
​
    if (nt_hdrs->OptionalHeader.NumberOfRvaAndSizes < IMAGE_DIRECTORY_ENTRY_BASERELOC + 1) {
        panic(true, "pe: NumberOfRvaAndSizes too small for import/reloc directories");
    }
​
    IMAGE_DATA_DIRECTORY *import_dir = &nt_hdrs->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT];
    IMAGE_DATA_DIRECTORY *reloc_dir = &nt_hdrs->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_BASERELOC];
​
    if (import_dir->Size != 0) {
        if (import_dir->VirtualAddress >= image_size ||
            sizeof(IMAGE_IMPORT_DESCRIPTOR) > image_size - import_dir->VirtualAddress) {
            panic(true, "pe: Import directory VirtualAddress out of bounds");
        }
        IMAGE_IMPORT_DESCRIPTOR *import_desc = (IMAGE_IMPORT_DESCRIPTOR *)((uintptr_t)*physical_base + import_dir->VirtualAddress);
​
        if (import_desc->Name != 0) {
            panic(true, "pe: Kernel must not have any imports");
        }
    }
​
    if (reloc_dir->VirtualAddress != 0) {
        if (reloc_dir->VirtualAddress >= image_size ||
            reloc_dir->Size > image_size - reloc_dir->VirtualAddress) {
            panic(true, "pe: Relocation directory VirtualAddress out of bounds");
        }
        size_t reloc_block_offset = 0;
​
        while (reloc_dir->Size - reloc_block_offset >= sizeof(IMAGE_BASE_RELOCATION_BLOCK)) {
            IMAGE_BASE_RELOCATION_BLOCK *block = (IMAGE_BASE_RELOCATION_BLOCK *)((uintptr_t)*physical_base + reloc_dir->VirtualAddress + reloc_block_offset);
​
            // Validate SizeOfBlock to prevent infinite loop (if 0) and underflow (if too small)
            if (block->SizeOfBlock < sizeof(IMAGE_BASE_RELOCATION_BLOCK)) {
                panic(true, "pe: Invalid relocation block size");
            }
​
            if (block->SizeOfBlock > reloc_dir->Size - reloc_block_offset) {
                panic(true, "pe: Relocation block size exceeds directory");
            }
​
            if (block->VirtualAddress >= image_size) {
                panic(true, "pe: Relocation block VirtualAddress out of bounds");
            }
​
            uintptr_t block_base = *physical_base + block->VirtualAddress;
            size_t entries = (block->SizeOfBlock - sizeof(IMAGE_BASE_RELOCATION_BLOCK)) / sizeof(uint16_t);
            uint16_t *relocs = (uint16_t *)(block + 1);
​
            for (size_t i = 0; i < entries; i++) {
                uint16_t type = relocs[i] >> 12;
                uint16_t offset = relocs[i] & 0xfff;
​
                if (type == IMAGE_REL_BASED_ABSOLUTE) {
                    continue;
                }
​
                size_t write_size;
                switch (type) {
                    case IMAGE_REL_BASED_HIGHLOW:
                        if (lower_to_higher) {
                            panic(true, "pe: 32-bit relocations are incompatible with higher-half loading");
                        }
                        write_size = 4;
                        break;
                    case IMAGE_REL_BASED_DIR64: write_size = 8; break;
                    default:
                        panic(true, "pe: Unsupported relocation type %u", type);
                        __builtin_unreachable();
                }
​
                if ((uint64_t)block->VirtualAddress + offset + write_size > image_size) {
                    panic(true, "pe: Relocation offset out of bounds");
                }
​
                switch (type) {
                    case IMAGE_REL_BASED_HIGHLOW:
                        *(uint32_t *)(block_base + offset) += slide;
                        break;
                    case IMAGE_REL_BASED_DIR64:
                        *(uint64_t *)(block_base + offset) += slide;
                        break;
                }
            }
​
            reloc_block_offset += block->SizeOfBlock;
        }
    }
​
    if (image_size_before_bss) {
        *image_size_before_bss = image_size;
    }
​
    *virtual_base += slide;
    *entry_point = *virtual_base + nt_hdrs->OptionalHeader.AddressOfEntryPoint;
​
    if (_slide) {
        *_slide = slide;
    }
​
    if (_ranges && _ranges_count) {
        size_t range_count = 0;
​
        bool headers_within_section = false;
​
        for (size_t i = 0; i < nt_hdrs->FileHeader.NumberOfSections; i++) {
            IMAGE_SECTION_HEADER *section = &sections[i];
​
            if (section->VirtualAddress == 0) {
                headers_within_section = true;
            }
​
            range_count++;
        }
​
        if (!headers_within_section) {
            range_count++;
        }
​
        struct mem_range *ranges = ext_mem_alloc_counted(range_count, sizeof(struct mem_range));
​
        *_ranges = ranges;
        *_ranges_count = range_count;
​
        size_t range_index = 0;
​
        if (!headers_within_section) {
            struct mem_range *range = &ranges[range_index++];
            range->base = *virtual_base;
            range->length = ALIGN_UP(nt_hdrs->OptionalHeader.SizeOfHeaders, 0x1000, panic(true, "pe: Alignment overflow"));
            range->permissions = MEM_RANGE_R;
        }
​
        for (size_t i = 0; i < nt_hdrs->FileHeader.NumberOfSections; i++) {
            IMAGE_SECTION_HEADER *section = &sections[i];
​
            uintptr_t misalign = section->VirtualAddress % alignment;
​
            struct mem_range *range = &ranges[range_index++];
            range->base = *virtual_base + ALIGN_DOWN(section->VirtualAddress, alignment);
            range->length = ALIGN_UP(section->VirtualSize + misalign, alignment, panic(true, "pe: Alignment overflow"));
​
            if (section->Characteristics & IMAGE_SCN_MEM_EXECUTE) {
                range->permissions |= MEM_RANGE_X;
            }
​
            if (section->Characteristics & IMAGE_SCN_MEM_WRITE) {
                range->permissions |= MEM_RANGE_W;
            }
​
            if (section->Characteristics & IMAGE_SCN_MEM_READ) {
                range->permissions |= MEM_RANGE_R;
            }
        }
    }
​
    return true;
}

tab: 2 4 8 wrap: off on

1	#include <stdint.h>
2	#include <stddef.h>
3	#include <lib/misc.h>
4	#include <lib/libc.h>
5	#include <lib/pe.h>
6	#include <lib/print.h>
7	#include <lib/rand.h>
8	#include <mm/pmm.h>
9
10	#define FIXED_HIGHER_HALF_OFFSET_64 ((uint64_t)0xffffffff80000000)
11
12	#define IMAGE_DOS_SIGNATURE 0x5a4d
13
14	typedef struct _IMAGE_DOS_HEADER {
15	uint16_t e_magic;
16	uint16_t e_cblp;
17	uint16_t e_cp;
18	uint16_t e_crlc;
19	uint16_t e_cparhdr;
20	uint16_t e_minalloc;
21	uint16_t e_maxalloc;
22	uint16_t e_ss;
23	uint16_t e_sp;
24	uint16_t e_csum;
25	uint16_t e_ip;
26	uint16_t e_cs;
27	uint16_t e_lfarlc;
28	uint16_t e_ovno;
29	uint16_t e_res[4];
30	uint16_t e_oemid;
31	uint16_t e_oeminfo;
32	uint16_t e_res2[10];
33	uint32_t e_lfanew;
34	} IMAGE_DOS_HEADER;
35
36	#define IMAGE_FILE_MACHINE_I386 0x14c
37	#define IMAGE_FILE_MACHINE_AMD64 0x8664
38	#define IMAGE_FILE_MACHINE_ARM64 0xaa64
39	#define IMAGE_FILE_MACHINE_RISCV64 0x5064
40	#define IMAGE_FILE_MACHINE_LOONGARCH64 0x6264
41
42	#define IMAGE_FILE_RELOCS_STRIPPED 1
43	#define IMAGE_FILE_EXECUTABLE_IMAGE 2
44
45	typedef struct {
46	uint16_t Machine;
47	uint16_t NumberOfSections;
48	uint32_t TimeDateStamp;
49	uint32_t PointerToSymbolTable;
50	uint32_t NumberOfSymbols;
51	uint16_t SizeOfOptionalHeader;
52	uint16_t Characteristics;
53	} IMAGE_FILE_HEADER;
54
55	typedef struct {
56	uint32_t VirtualAddress;
57	uint32_t Size;
58	} IMAGE_DATA_DIRECTORY;
59
60	#define IMAGE_NT_OPTIONAL_HDR32_MAGIC 0x10b
61	#define IMAGE_NT_OPTIONAL_HDR64_MAGIC 0x20b
62
63	#define IMAGE_DIRECTORY_ENTRY_EXPORT 0
64	#define IMAGE_DIRECTORY_ENTRY_IMPORT 1
65	#define IMAGE_DIRECTORY_ENTRY_RESOURCE 2
66	#define IMAGE_DIRECTORY_ENTRY_EXCEPTION 3
67	#define IMAGE_DIRECTORY_ENTRY_SECURITY 4
68	#define IMAGE_DIRECTORY_ENTRY_BASERELOC 5
69	#define IMAGE_DIRECTORY_ENTRY_DEBUG 6
70	#define IMAGE_DIRECTORY_ENTRY_ARCHITECTURE 7
71	#define IMAGE_DIRECTORY_ENTRY_GLOBALPTR 8
72	#define IMAGE_DIRECTORY_ENTRY_TLS 9
73	#define IMAGE_DIRECTORY_ENTRY_LOAD_CONFIG 10
74	#define IMAGE_DIRECTORY_ENTRY_BOUND_IMPORT 11
75	#define IMAGE_DIRECTORY_ENTRY_IAT 12
76	#define IMAGE_DIRECTORY_ENTRY_DELAY_IMPORT 13
77	#define IMAGE_DIRECTORY_ENTRY_COM_DESCRIPTOR 14
78
79	typedef struct {
80	uint16_t Magic;
81	uint8_t MajorLinkerVersion;
82	uint8_t MinorLinkerVersion;
83	uint32_t SizeOfCode;
84	uint32_t SizeOfInitializedData;
85	uint32_t SizeOfUninitializedData;
86	uint32_t AddressOfEntryPoint;
87	uint32_t BaseOfCode;
88	uint64_t ImageBase;
89	uint32_t SectionAlignment;
90	uint32_t FileAlignment;
91	uint16_t MajorOperatingSystemVersion;
92	uint16_t MinorOperatingSystemVersion;
93	uint16_t MajorImageVersion;
94	uint16_t MinorImageVersion;
95	uint16_t MajorSubsystemVersion;
96	uint16_t MinorSubsystemVersion;
97	uint32_t Win32VersionValue;
98	uint32_t SizeOfImage;
99	uint32_t SizeOfHeaders;
100	uint32_t CheckSum;
101	uint16_t Subsystem;
102	uint16_t DllCharacteristics;
103	uint64_t SizeOfStackReserve;
104	uint64_t SizeOfStackCommit;
105	uint64_t SizeOfHeapReserve;
106	uint64_t SizeOfHeapCommit;
107	uint32_t LoaderFlags;
108	uint32_t NumberOfRvaAndSizes;
109	IMAGE_DATA_DIRECTORY DataDirectory[16];
110	} IMAGE_OPTIONAL_HEADER64;
111
112	#define IMAGE_NT_SIGNATURE 0x4550
113
114	typedef struct {
115	uint32_t Signature;
116	IMAGE_FILE_HEADER FileHeader;
117	IMAGE_OPTIONAL_HEADER64 OptionalHeader;
118	} IMAGE_NT_HEADERS64;
119
120	#define IMAGE_SCN_MEM_EXECUTE 0x20000000
121	#define IMAGE_SCN_MEM_READ 0x40000000
122	#define IMAGE_SCN_MEM_WRITE 0x80000000
123
124	typedef struct {
125	char Name[8];
126	uint32_t VirtualSize;
127	uint32_t VirtualAddress;
128	uint32_t SizeOfRawData;
129	uint32_t PointerToRawData;
130	uint32_t PointerToRelocations;
131	uint32_t PointerToLinenumbers;
132	uint16_t NumberOfRelocations;
133	uint16_t NumberOfLinenumbers;
134	uint32_t Characteristics;
135	} IMAGE_SECTION_HEADER;
136
137	typedef struct {
138	union {
139	uint32_t Characteristics;
140	uint32_t OriginalFirstThunk;
141	};
142	uint32_t TimeDateStamp;
143	uint32_t ForwarderChain;
144	uint32_t Name;
145	uint32_t FirstThunk;
146	} IMAGE_IMPORT_DESCRIPTOR;
147
148	#define IMAGE_REL_BASED_ABSOLUTE 0
149	#define IMAGE_REL_BASED_HIGHLOW 3
150	#define IMAGE_REL_BASED_DIR64 10
151
152	typedef struct {
153	uint32_t VirtualAddress;
154	uint32_t SizeOfBlock;
155	} IMAGE_BASE_RELOCATION_BLOCK;
156
157	static void pe64_validate(uint8_t *image, size_t file_size) {
158	IMAGE_DOS_HEADER dos_hdr = (IMAGE_DOS_HEADER )image;
159
160	if (file_size < sizeof(IMAGE_DOS_HEADER)) {
161	panic(true, "pe: File too small for DOS header");
162	}
163
164	if (dos_hdr->e_magic != IMAGE_DOS_SIGNATURE) {
165	panic(true, "pe: Not a valid PE file");
166	}
167
168	if (file_size < sizeof(IMAGE_NT_HEADERS64)) {
169	panic(true, "pe: File too small for NT headers");
170	}
171
172	if (dos_hdr->e_lfanew > file_size - sizeof(IMAGE_NT_HEADERS64)) {
173	panic(true, "pe: e_lfanew offset out of bounds");
174	}
175
176	IMAGE_NT_HEADERS64 nt_hdrs = (IMAGE_NT_HEADERS64 )(image + dos_hdr->e_lfanew);
177
178	if (nt_hdrs->Signature != IMAGE_NT_SIGNATURE) {
179	panic(true, "pe: Not a valid PE file");
180	}
181
182	if (nt_hdrs->OptionalHeader.Magic != IMAGE_NT_OPTIONAL_HDR64_MAGIC) {
183	panic(true, "pe: Not a valid PE32+ file");
184	}
185
186	if (nt_hdrs->FileHeader.SizeOfOptionalHeader < sizeof(IMAGE_OPTIONAL_HEADER64)) {
187	panic(true, "pe: SizeOfOptionalHeader too small");
188	}
189
190	#if defined(__x86_64__) \|\| defined(__i386__)
191	if (nt_hdrs->FileHeader.Machine != IMAGE_FILE_MACHINE_AMD64) {
192	panic(true, "pe: Not an x86-64 PE file");
193	}
194	#elif defined(__aarch64__)
195	if (nt_hdrs->FileHeader.Machine != IMAGE_FILE_MACHINE_ARM64) {
196	panic(true, "pe: Not an ARM64 PE file");
197	}
198	#elif defined (__riscv) && (__riscv_xlen == 64)
199	if (nt_hdrs->FileHeader.Machine != IMAGE_FILE_MACHINE_RISCV64) {
200	panic(true, "pe: Not a RISC-V PE file");
201	}
202	#elif defined (__loongarch__) && (__loongarch_grlen == 64)
203	if (nt_hdrs->FileHeader.Machine != IMAGE_FILE_MACHINE_LOONGARCH64) {
204	panic(true, "pe: Not a loongarch64 PE file");
205	}
206	#else
207	#error Unknown architecture
208	#endif
209	}
210
211	int pe_bits(uint8_t *image, size_t image_size) {
212	if (image_size < sizeof(IMAGE_DOS_HEADER)) {
213	return -1;
214	}
215
216	IMAGE_DOS_HEADER dos_hdr = (IMAGE_DOS_HEADER )image;
217
218	if (dos_hdr->e_magic != IMAGE_DOS_SIGNATURE) {
219	return -1;
220	}
221
222	if (image_size < sizeof(IMAGE_NT_HEADERS64)) {
223	return -1;
224	}
225
226	if ((size_t)dos_hdr->e_lfanew > image_size - sizeof(IMAGE_NT_HEADERS64)) {
227	return -1;
228	}
229
230	IMAGE_NT_HEADERS64 nt_hdrs = (IMAGE_NT_HEADERS64 )(image + dos_hdr->e_lfanew);
231
232	if (nt_hdrs->Signature != IMAGE_NT_SIGNATURE) {
233	return -1;
234	}
235
236	switch (nt_hdrs->FileHeader.Machine) {
237	case IMAGE_FILE_MACHINE_I386:
238	return 32;
239	case IMAGE_FILE_MACHINE_AMD64:
240	case IMAGE_FILE_MACHINE_ARM64:
241	case IMAGE_FILE_MACHINE_RISCV64:
242	case IMAGE_FILE_MACHINE_LOONGARCH64:
243	return 64;
244	}
245
246	return -1;
247	}
248
249	bool pe64_load(uint8_t image, size_t file_size, uint64_t entry_point, uint64_t _slide, uint32_t alloc_type, bool kaslr, struct mem_range _ranges, uint64_t _ranges_count, uint64_t physical_base, uint64_t virtual_base, uint64_t _image_size, uint64_t image_size_before_bss, bool *_is_reloc) {
250	pe64_validate(image, file_size);
251
252	IMAGE_DOS_HEADER dos_hdr = (IMAGE_DOS_HEADER )image;
253	IMAGE_NT_HEADERS64 nt_hdrs = (IMAGE_NT_HEADERS64 )(image + dos_hdr->e_lfanew);
254
255	// Validate SizeOfOptionalHeader doesn't cause sections pointer to go out of bounds
256	uint64_t sections_offset = (uint64_t)dos_hdr->e_lfanew + sizeof(uint32_t) + sizeof(IMAGE_FILE_HEADER) + nt_hdrs->FileHeader.SizeOfOptionalHeader;
257	uint64_t sections_end = sections_offset + (uint64_t)nt_hdrs->FileHeader.NumberOfSections * sizeof(IMAGE_SECTION_HEADER);
258	if (sections_end > file_size) {
259	panic(true, "pe: Section headers extend beyond file bounds");
260	}
261
262	IMAGE_SECTION_HEADER sections = (IMAGE_SECTION_HEADER )((uintptr_t)&nt_hdrs->OptionalHeader + nt_hdrs->FileHeader.SizeOfOptionalHeader);
263
264	bool is_reloc = true;
265
266	if (nt_hdrs->FileHeader.Characteristics & IMAGE_FILE_RELOCS_STRIPPED) {
267	is_reloc = false;
268	}
269
270	if (_is_reloc) {
271	*_is_reloc = is_reloc;
272	}
273
274	uint64_t image_base = nt_hdrs->OptionalHeader.ImageBase;
275	uint64_t image_size = nt_hdrs->OptionalHeader.SizeOfImage;
276	uint64_t alignment = nt_hdrs->OptionalHeader.SectionAlignment;
277
278	if (alignment == 0) {
279	panic(true, "pe: SectionAlignment is zero");
280	}
281
282	if (alignment > 1 && (alignment & (alignment - 1)) != 0) {
283	panic(true, "pe: SectionAlignment is not a power of 2");
284	}
285
286	bool lower_to_higher = false;
287
288	if (image_base < FIXED_HIGHER_HALF_OFFSET_64) {
289	if (!is_reloc) {
290	panic(true, "pe: Lower half images are not allowed");
291	}
292
293	lower_to_higher = true;
294	}
295
296	uint64_t slide = 0;
297	size_t try_count = 0;
298	size_t max_simulated_tries = 0x10000;
299
300	if (lower_to_higher) {
301	slide = FIXED_HIGHER_HALF_OFFSET_64 - image_base;
302	}
303
304	*physical_base = (uintptr_t)ext_mem_alloc_type_aligned(image_size, alloc_type, alignment);
305	*virtual_base = image_base;
306
307	// Validate SizeOfHeaders doesn't exceed file size or image size
308	if (nt_hdrs->OptionalHeader.SizeOfHeaders > file_size
309	\|\| nt_hdrs->OptionalHeader.SizeOfHeaders > image_size) {
310	panic(true, "pe: SizeOfHeaders exceeds file or image size");
311	}
312
313	memcpy((void )(uintptr_t)physical_base, image, nt_hdrs->OptionalHeader.SizeOfHeaders);
314
315	if (_image_size) {
316	*_image_size = image_size;
317	}
318
319	if (is_reloc && kaslr) {
320	again:
321	slide = (rand32() & ~(alignment - 1)) + (lower_to_higher ? FIXED_HIGHER_HALF_OFFSET_64 - image_base : 0);
322
323	if (virtual_base + slide + image_size < 0xffffffff80000000 / this comparison relies on overflow */) {
324	if (++try_count == max_simulated_tries) {
325	panic(true, "pe: Image wants to load too high");
326	}
327	goto again;
328	}
329	}
330
331	for (size_t i = 0; i < nt_hdrs->FileHeader.NumberOfSections; i++) {
332	IMAGE_SECTION_HEADER *section = &sections[i];
333
334	uintptr_t section_base = *physical_base + section->VirtualAddress;
335	uint32_t section_raw_size = section->VirtualSize < section->SizeOfRawData ? section->VirtualSize : section->SizeOfRawData;
336
337	// Validate section doesn't write past the image buffer
338	if ((uint64_t)section->VirtualAddress + section_raw_size > image_size) {
339	panic(true, "pe: Section %U exceeds image bounds", (uint64_t)i);
340	}
341
342	// Validate section data doesn't exceed file bounds
343	if ((uint64_t)section->PointerToRawData + section_raw_size > file_size) {
344	panic(true, "pe: Section %U data extends beyond file bounds", (uint64_t)i);
345	}
346
347	memcpy((void *)section_base, image + section->PointerToRawData, section_raw_size);
348	}
349
350	if (nt_hdrs->OptionalHeader.NumberOfRvaAndSizes < IMAGE_DIRECTORY_ENTRY_BASERELOC + 1) {
351	panic(true, "pe: NumberOfRvaAndSizes too small for import/reloc directories");
352	}
353
354	IMAGE_DATA_DIRECTORY *import_dir = &nt_hdrs->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT];
355	IMAGE_DATA_DIRECTORY *reloc_dir = &nt_hdrs->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_BASERELOC];
356
357	if (import_dir->Size != 0) {
358	if (import_dir->VirtualAddress >= image_size \|\|
359	sizeof(IMAGE_IMPORT_DESCRIPTOR) > image_size - import_dir->VirtualAddress) {
360	panic(true, "pe: Import directory VirtualAddress out of bounds");
361	}
362	IMAGE_IMPORT_DESCRIPTOR import_desc = (IMAGE_IMPORT_DESCRIPTOR )((uintptr_t)*physical_base + import_dir->VirtualAddress);
363
364	if (import_desc->Name != 0) {
365	panic(true, "pe: Kernel must not have any imports");
366	}
367	}
368
369	if (reloc_dir->VirtualAddress != 0) {
370	if (reloc_dir->VirtualAddress >= image_size \|\|
371	reloc_dir->Size > image_size - reloc_dir->VirtualAddress) {
372	panic(true, "pe: Relocation directory VirtualAddress out of bounds");
373	}
374	size_t reloc_block_offset = 0;
375
376	while (reloc_dir->Size - reloc_block_offset >= sizeof(IMAGE_BASE_RELOCATION_BLOCK)) {
377	IMAGE_BASE_RELOCATION_BLOCK block = (IMAGE_BASE_RELOCATION_BLOCK )((uintptr_t)*physical_base + reloc_dir->VirtualAddress + reloc_block_offset);
378
379	// Validate SizeOfBlock to prevent infinite loop (if 0) and underflow (if too small)
380	if (block->SizeOfBlock < sizeof(IMAGE_BASE_RELOCATION_BLOCK)) {
381	panic(true, "pe: Invalid relocation block size");
382	}
383
384	if (block->SizeOfBlock > reloc_dir->Size - reloc_block_offset) {
385	panic(true, "pe: Relocation block size exceeds directory");
386	}
387
388	if (block->VirtualAddress >= image_size) {
389	panic(true, "pe: Relocation block VirtualAddress out of bounds");
390	}
391
392	uintptr_t block_base = *physical_base + block->VirtualAddress;
393	size_t entries = (block->SizeOfBlock - sizeof(IMAGE_BASE_RELOCATION_BLOCK)) / sizeof(uint16_t);
394	uint16_t relocs = (uint16_t )(block + 1);
395
396	for (size_t i = 0; i < entries; i++) {
397	uint16_t type = relocs[i] >> 12;
398	uint16_t offset = relocs[i] & 0xfff;
399
400	if (type == IMAGE_REL_BASED_ABSOLUTE) {
401	continue;
402	}
403
404	size_t write_size;
405	switch (type) {
406	case IMAGE_REL_BASED_HIGHLOW:
407	if (lower_to_higher) {
408	panic(true, "pe: 32-bit relocations are incompatible with higher-half loading");
409	}
410	write_size = 4;
411	break;
412	case IMAGE_REL_BASED_DIR64: write_size = 8; break;
413	default:
414	panic(true, "pe: Unsupported relocation type %u", type);
415	__builtin_unreachable();
416	}
417
418	if ((uint64_t)block->VirtualAddress + offset + write_size > image_size) {
419	panic(true, "pe: Relocation offset out of bounds");
420	}
421
422	switch (type) {
423	case IMAGE_REL_BASED_HIGHLOW:
424	(uint32_t )(block_base + offset) += slide;
425	break;
426	case IMAGE_REL_BASED_DIR64:
427	(uint64_t )(block_base + offset) += slide;
428	break;
429	}
430	}
431
432	reloc_block_offset += block->SizeOfBlock;
433	}
434	}
435
436	if (image_size_before_bss) {
437	*image_size_before_bss = image_size;
438	}
439
440	*virtual_base += slide;
441	entry_point = virtual_base + nt_hdrs->OptionalHeader.AddressOfEntryPoint;
442
443	if (_slide) {
444	*_slide = slide;
445	}
446
447	if (_ranges && _ranges_count) {
448	size_t range_count = 0;
449
450	bool headers_within_section = false;
451
452	for (size_t i = 0; i < nt_hdrs->FileHeader.NumberOfSections; i++) {
453	IMAGE_SECTION_HEADER *section = &sections[i];
454
455	if (section->VirtualAddress == 0) {
456	headers_within_section = true;
457	}
458
459	range_count++;
460	}
461
462	if (!headers_within_section) {
463	range_count++;
464	}
465
466	struct mem_range *ranges = ext_mem_alloc_counted(range_count, sizeof(struct mem_range));
467
468	*_ranges = ranges;
469	*_ranges_count = range_count;
470
471	size_t range_index = 0;
472
473	if (!headers_within_section) {
474	struct mem_range *range = &ranges[range_index++];
475	range->base = *virtual_base;
476	range->length = ALIGN_UP(nt_hdrs->OptionalHeader.SizeOfHeaders, 0x1000, panic(true, "pe: Alignment overflow"));
477	range->permissions = MEM_RANGE_R;
478	}
479
480	for (size_t i = 0; i < nt_hdrs->FileHeader.NumberOfSections; i++) {
481	IMAGE_SECTION_HEADER *section = &sections[i];
482
483	uintptr_t misalign = section->VirtualAddress % alignment;
484
485	struct mem_range *range = &ranges[range_index++];
486	range->base = *virtual_base + ALIGN_DOWN(section->VirtualAddress, alignment);
487	range->length = ALIGN_UP(section->VirtualSize + misalign, alignment, panic(true, "pe: Alignment overflow"));
488
489	if (section->Characteristics & IMAGE_SCN_MEM_EXECUTE) {
490	range->permissions \|= MEM_RANGE_X;
491	}
492
493	if (section->Characteristics & IMAGE_SCN_MEM_WRITE) {
494	range->permissions \|= MEM_RANGE_W;
495	}
496
497	if (section->Characteristics & IMAGE_SCN_MEM_READ) {
498	range->permissions \|= MEM_RANGE_R;
499	}
500	}
501	}
502
503	return true;
504	}