nt-lite/code/boot/uefi/boot_pe.c

// Portable fixed-image PE32+ loader. This file has no UEFI dependencies so firmware glue stays out of image parsing.

#include "boot_pe.h"

#include <stdbool.h>
#include <stdint.h>
#include <string.h>

#define BOOT_PE_DOS_MAGIC 0x5A4Du
#define BOOT_PE_NT_MAGIC 0x00004550u
#define BOOT_PE_MACHINE_AMD64 0x8664u
#define BOOT_PE_OPTIONAL_MAGIC_PE32_PLUS 0x020Bu
#define BOOT_PE_ORDINAL_FLAG64 0x8000000000000000ull
#define BOOT_PE_DIRECTORY_EXPORT 0u
#define BOOT_PE_DIRECTORY_IMPORT 1u
#define BOOT_PE_DIRECTORY_RELOCATION 5u
#define BOOT_PE_OPTIONAL_MIN_BYTES 112u
#define BOOT_PE_SECTION_BYTES 40u
#define BOOT_PE_IMPORT_DESCRIPTOR_BYTES 20u
#define BOOT_PE_PAGE_SIZE 4096u

typedef struct BOOT_PE_VIEW {
    const uint8_t* File;
    size_t FileSize;
    size_t NtOffset;
    size_t OptionalOffset;
    size_t SectionOffset;
    uint16_t SectionCount;
    uint16_t OptionalSize;
    uint32_t DirectoryCount;
    uint64_t ImageBase;
    uint32_t SizeOfImage;
    uint32_t SizeOfHeaders;
    uint32_t EntryRva;
    uint32_t Checksum;
} BOOT_PE_VIEW;

// Returns whether adding two sizes would overflow size_t.
static bool BootPeSizeAdd(size_t a, size_t b, size_t* out) {
    if (a > SIZE_MAX - b)
        return false;
    *out = a + b;
    return true;
}

// Returns whether multiplying two sizes would overflow size_t.
static bool BootPeSizeMultiply(size_t a, size_t b, size_t* out) {
    if (a != 0 && b > SIZE_MAX / a)
        return false;
    *out = a * b;
    return true;
}

// Returns whether one byte range is fully contained in a buffer.
static bool BootPeRangeValid(size_t offset, size_t length, size_t bufferSize) {
    return offset <= bufferSize && length <= bufferSize - offset;
}

// Reads one little-endian 16-bit field without assuming alignment.
static uint16_t BootPeReadU16(const void* address) {
    uint16_t value;
    memcpy(&value, address, sizeof(value));
    return value;
}

// Reads one little-endian 32-bit field without assuming alignment.
static uint32_t BootPeReadU32(const void* address) {
    uint32_t value;
    memcpy(&value, address, sizeof(value));
    return value;
}

// Reads one little-endian 64-bit field without assuming alignment.
static uint64_t BootPeReadU64(const void* address) {
    uint64_t value;
    memcpy(&value, address, sizeof(value));
    return value;
}

// Writes one little-endian 64-bit field without assuming alignment.
static void BootPeWriteU64(void* address, uint64_t value) {
    memcpy(address, &value, sizeof(value));
}

// Parses the fixed PE header fields shared by file and materialized-image validation.
static BOOT_PE_STATUS BootPeParseHeaders(const void* buffer, size_t bufferSize,
                                         bool requireRawSections, BOOT_PE_VIEW* viewOut) {
    if (!buffer || !viewOut || bufferSize < 64u)
        return BootPeInvalidImage;

    const uint8_t* bytes = (const uint8_t*)buffer;
    if (BootPeReadU16(bytes) != BOOT_PE_DOS_MAGIC)
        return BootPeInvalidImage;

    size_t ntOffset = (size_t)BootPeReadU32(bytes + 0x3Cu);
    if (!BootPeRangeValid(ntOffset, 24u, bufferSize))
        return BootPeInvalidImage;
    if (BootPeReadU32(bytes + ntOffset) != BOOT_PE_NT_MAGIC)
        return BootPeInvalidImage;
    if (BootPeReadU16(bytes + ntOffset + 4u) != BOOT_PE_MACHINE_AMD64)
        return BootPeWrongMachine;

    uint16_t sectionCount = BootPeReadU16(bytes + ntOffset + 6u);
    uint16_t optionalSize = BootPeReadU16(bytes + ntOffset + 20u);
    size_t optionalOffset = ntOffset + 24u;
    if (optionalSize < BOOT_PE_OPTIONAL_MIN_BYTES ||
        !BootPeRangeValid(optionalOffset, optionalSize, bufferSize)) {
        return BootPeInvalidImage;
    }
    if (BootPeReadU16(bytes + optionalOffset) != BOOT_PE_OPTIONAL_MAGIC_PE32_PLUS)
        return BootPeInvalidImage;

    uint32_t entryRva = BootPeReadU32(bytes + optionalOffset + 16u);
    uint64_t imageBase = BootPeReadU64(bytes + optionalOffset + 24u);
    uint32_t sizeOfImage = BootPeReadU32(bytes + optionalOffset + 56u);
    uint32_t sizeOfHeaders = BootPeReadU32(bytes + optionalOffset + 60u);
    uint32_t checksum = BootPeReadU32(bytes + optionalOffset + 64u);
    uint32_t directoryCount = BootPeReadU32(bytes + optionalOffset + 108u);

    if (imageBase == 0 || (imageBase & (BOOT_PE_PAGE_SIZE - 1u)) != 0 ||
        sizeOfImage == 0 || sizeOfHeaders == 0 || sizeOfHeaders > sizeOfImage ||
        entryRva >= sizeOfImage) {
        return BootPeInvalidImage;
    }

    size_t directoryBytes = 0;
    uint32_t availableDirectories = (uint32_t)((optionalSize - BOOT_PE_OPTIONAL_MIN_BYTES) / 8u);
    if (directoryCount > availableDirectories)
        return BootPeInvalidImage;
    if (!BootPeSizeMultiply((size_t)directoryCount, 8u, &directoryBytes) ||
        !BootPeRangeValid(optionalOffset + BOOT_PE_OPTIONAL_MIN_BYTES, directoryBytes, bufferSize)) {
        return BootPeInvalidImage;
    }

    size_t sectionOffset = 0;
    size_t sectionBytes = 0;
    if (!BootPeSizeAdd(optionalOffset, optionalSize, &sectionOffset) ||
        !BootPeSizeMultiply(sectionCount, BOOT_PE_SECTION_BYTES, &sectionBytes) ||
        !BootPeRangeValid(sectionOffset, sectionBytes, bufferSize)) {
        return BootPeInvalidImage;
    }

    if (requireRawSections) {
        if (sizeOfHeaders > bufferSize)
            return BootPeInvalidImage;

        for (uint16_t index = 0; index < sectionCount; ++index) {
            const uint8_t* section = bytes + sectionOffset + (size_t)index * BOOT_PE_SECTION_BYTES;
            uint32_t virtualSize = BootPeReadU32(section + 8u);
            uint32_t virtualAddress = BootPeReadU32(section + 12u);
            uint32_t rawSize = BootPeReadU32(section + 16u);
            uint32_t rawOffset = BootPeReadU32(section + 20u);
            uint32_t mappedSize = virtualSize > rawSize ? virtualSize : rawSize;

            if (virtualAddress > sizeOfImage || mappedSize > sizeOfImage - virtualAddress)
                return BootPeInvalidImage;
            if (rawSize != 0 && !BootPeRangeValid(rawOffset, rawSize, bufferSize))
                return BootPeInvalidImage;
        }
    }

    BOOT_PE_VIEW view = {0};
    view.File = bytes;
    view.FileSize = bufferSize;
    view.NtOffset = ntOffset;
    view.OptionalOffset = optionalOffset;
    view.SectionOffset = sectionOffset;
    view.SectionCount = sectionCount;
    view.OptionalSize = optionalSize;
    view.DirectoryCount = directoryCount;
    view.ImageBase = imageBase;
    view.SizeOfImage = sizeOfImage;
    view.SizeOfHeaders = sizeOfHeaders;
    view.EntryRva = entryRva;
    view.Checksum = checksum;
    *viewOut = view;
    return BootPeSuccess;
}

// Reads one PE data-directory entry when the optional header declares it.
static bool BootPeGetDirectory(const BOOT_PE_VIEW* view, uint32_t index,
                               uint32_t* rvaOut, uint32_t* sizeOut) {
    if (!view || index >= view->DirectoryCount)
        return false;
    const uint8_t* directory = view->File + view->OptionalOffset + BOOT_PE_OPTIONAL_MIN_BYTES + (size_t)index * 8u;
    if (rvaOut) *rvaOut = BootPeReadU32(directory);
    if (sizeOut) *sizeOut = BootPeReadU32(directory + 4u);
    return true;
}

// Computes the standard PE checksum while treating the stored CheckSum field as zero.
static uint32_t BootPeComputeChecksum(const uint8_t* file, size_t fileSize, size_t checksumOffset) {
    uint64_t sum = 0;
    size_t offset = 0;

    while (offset + 1u < fileSize) {
        if (offset == checksumOffset || offset == checksumOffset + 2u) {
            offset += 2u;
            continue;
        }

        sum += BootPeReadU16(file + offset);
        sum = (sum & 0xFFFFu) + (sum >> 16);
        offset += 2u;
    }

    if (offset < fileSize)
        sum += file[offset];

    sum = (sum & 0xFFFFu) + (sum >> 16);
    sum += sum >> 16;
    sum = (sum & 0xFFFFu) + fileSize;
    return (uint32_t)sum;
}

// Returns whether an image RVA range is fully mapped inside SizeOfImage.
static bool BootPeImageRangeValid(uint32_t rva, size_t length, size_t imageSize) {
    return (size_t)rva <= imageSize && length <= imageSize - (size_t)rva;
}

// Returns whether an image-relative string has a NUL terminator before the image ends.
static bool BootPeImageStringValid(const uint8_t* image, size_t imageSize, uint32_t rva) {
    if ((size_t)rva >= imageSize)
        return false;
    for (size_t offset = (size_t)rva; offset < imageSize; ++offset) {
        if (image[offset] == 0)
            return true;
    }
    return false;
}

// Compares two ASCII module names without case sensitivity.
static bool BootPeAsciiEqualInsensitive(const char* a, const char* b) {
    if (!a || !b)
        return false;
    while (*a && *b) {
        char ca = *a++;
        char cb = *b++;
        if (ca >= 'A' && ca <= 'Z') ca = (char)(ca + ('a' - 'A'));
        if (cb >= 'A' && cb <= 'Z') cb = (char)(cb + ('a' - 'A'));
        if (ca != cb)
            return false;
    }
    return *a == *b;
}

// Finds one named export RVA and rejects forwarded exports.
static BOOT_PE_STATUS BootPeFindExport(const uint8_t* provider, size_t providerSize,
                                       const BOOT_PE_VIEW* providerView, const char* functionName,
                                       uint32_t* functionRvaOut) {
    uint32_t exportRva = 0;
    uint32_t exportSize = 0;
    if (!BootPeGetDirectory(providerView, BOOT_PE_DIRECTORY_EXPORT, &exportRva, &exportSize) ||
        exportRva == 0 || exportSize < 40u || !BootPeImageRangeValid(exportRva, exportSize, providerSize)) {
        return BootPeImportNotFound;
    }

    const uint8_t* exports = provider + exportRva;
    uint32_t functionCount = BootPeReadU32(exports + 20u);
    uint32_t nameCount = BootPeReadU32(exports + 24u);
    uint32_t functionsRva = BootPeReadU32(exports + 28u);
    uint32_t namesRva = BootPeReadU32(exports + 32u);
    uint32_t ordinalsRva = BootPeReadU32(exports + 36u);

    if (!BootPeImageRangeValid(functionsRva, (size_t)functionCount * 4u, providerSize) ||
        !BootPeImageRangeValid(namesRva, (size_t)nameCount * 4u, providerSize) ||
        !BootPeImageRangeValid(ordinalsRva, (size_t)nameCount * 2u, providerSize)) {
        return BootPeInvalidImage;
    }

    for (uint32_t index = 0; index < nameCount; ++index) {
        uint32_t nameRva = BootPeReadU32(provider + namesRva + (size_t)index * 4u);
        if (!BootPeImageStringValid(provider, providerSize, nameRva))
            return BootPeInvalidImage;
        if (strcmp((const char*)(provider + nameRva), functionName) != 0)
            continue;

        uint16_t ordinal = BootPeReadU16(provider + ordinalsRva + (size_t)index * 2u);
        if (ordinal >= functionCount)
            return BootPeInvalidImage;

        uint32_t functionRva = BootPeReadU32(provider + functionsRva + (size_t)ordinal * 4u);
        if (!BootPeImageRangeValid(functionRva, 1u, providerSize))
            return BootPeInvalidImage;
        if (functionRva >= exportRva && (uint64_t)functionRva < (uint64_t)exportRva + exportSize)
            return BootPeUnsupportedImport;

        *functionRvaOut = functionRva;
        return BootPeSuccess;
    }

    return BootPeImportNotFound;
}

// Validates one fixed AMD64 PE32+ file and returns the image metadata required by the bootloader.
BOOT_PE_STATUS BootPeInspect(const void* file, size_t fileSize, BOOT_PE_INFO* infoOut) {
    BOOT_PE_VIEW view;
    BOOT_PE_STATUS status = BootPeParseHeaders(file, fileSize, true, &view);
    if (status != BootPeSuccess)
        return status;

    uint32_t relocationRva = 0;
    uint32_t relocationSize = 0;
    if (BootPeGetDirectory(&view, BOOT_PE_DIRECTORY_RELOCATION, &relocationRva, &relocationSize) &&
        (relocationRva != 0 || relocationSize != 0)) {
        return BootPeRelocationsPresent;
    }

    if (view.Checksum == 0)
        return BootPeChecksumMissing;

    size_t checksumOffset = view.OptionalOffset + 64u;
    uint32_t computed = BootPeComputeChecksum(view.File, view.FileSize, checksumOffset);
    if (computed != view.Checksum)
        return BootPeChecksumMismatch;

    if (infoOut) {
        infoOut->ImageBase = view.ImageBase;
        infoOut->SizeOfImage = view.SizeOfImage;
        infoOut->SizeOfHeaders = view.SizeOfHeaders;
        infoOut->EntryRva = view.EntryRva;
        infoOut->Checksum = view.Checksum;
    }
    return BootPeSuccess;
}

// Copies a validated fixed image into a zeroed image-sized destination.
BOOT_PE_STATUS BootPeMaterialize(const void* file, size_t fileSize, void* image,
                                 size_t imageCapacity, BOOT_PE_INFO* infoOut) {
    BOOT_PE_INFO info;
    BOOT_PE_STATUS status = BootPeInspect(file, fileSize, &info);
    if (status != BootPeSuccess)
        return status;
    if (!image || imageCapacity < info.SizeOfImage)
        return BootPeBufferTooSmall;

    BOOT_PE_VIEW view;
    status = BootPeParseHeaders(file, fileSize, true, &view);
    if (status != BootPeSuccess)
        return status;

    memset(image, 0, info.SizeOfImage);
    memcpy(image, file, info.SizeOfHeaders);

    for (uint16_t index = 0; index < view.SectionCount; ++index) {
        const uint8_t* section = view.File + view.SectionOffset + (size_t)index * BOOT_PE_SECTION_BYTES;
        uint32_t virtualAddress = BootPeReadU32(section + 12u);
        uint32_t rawSize = BootPeReadU32(section + 16u);
        uint32_t rawOffset = BootPeReadU32(section + 20u);
        if (rawSize != 0)
            memcpy((uint8_t*)image + virtualAddress, view.File + rawOffset, rawSize);
    }

    if (infoOut)
        *infoOut = info;
    return BootPeSuccess;
}

// Returns whether a materialized image declares any import descriptors.
BOOT_PE_STATUS BootPeRequireNoImports(const void* image, size_t imageSize) {
    BOOT_PE_VIEW view;
    BOOT_PE_STATUS status = BootPeParseHeaders(image, imageSize, false, &view);
    if (status != BootPeSuccess)
        return status;

    uint32_t importRva = 0;
    uint32_t importSize = 0;
    if (BootPeGetDirectory(&view, BOOT_PE_DIRECTORY_IMPORT, &importRva, &importSize) &&
        (importRva != 0 || importSize != 0)) {
        return BootPeImportsPresent;
    }
    return BootPeSuccess;
}

// Resolves one consumer's named imports against a single fixed provider image.
BOOT_PE_STATUS BootPeResolveImports(void* consumerImage, size_t consumerSize,
                                    const void* providerImage, size_t providerSize,
                                    const char* providerName, uint64_t providerVirtualBase) {
    if (!consumerImage || !providerImage || !providerName)
        return BootPeInvalidImage;

    BOOT_PE_VIEW consumerView;
    BOOT_PE_VIEW providerView;
    BOOT_PE_STATUS status = BootPeParseHeaders(consumerImage, consumerSize, false, &consumerView);
    if (status != BootPeSuccess)
        return status;
    status = BootPeParseHeaders(providerImage, providerSize, false, &providerView);
    if (status != BootPeSuccess)
        return status;

    uint32_t importRva = 0;
    uint32_t importSize = 0;
    if (!BootPeGetDirectory(&consumerView, BOOT_PE_DIRECTORY_IMPORT, &importRva, &importSize) ||
        importRva == 0 || importSize == 0) {
        return BootPeSuccess;
    }
    if (importSize < BOOT_PE_IMPORT_DESCRIPTOR_BYTES ||
        !BootPeImageRangeValid(importRva, importSize, consumerSize)) {
        return BootPeInvalidImage;
    }

    uint32_t descriptorCount = importSize / BOOT_PE_IMPORT_DESCRIPTOR_BYTES;
    bool terminated = false;
    uint8_t* consumer = (uint8_t*)consumerImage;
    const uint8_t* provider = (const uint8_t*)providerImage;

    for (uint32_t descriptorIndex = 0; descriptorIndex < descriptorCount; ++descriptorIndex) {
        uint8_t* descriptor = consumer + importRva + (size_t)descriptorIndex * BOOT_PE_IMPORT_DESCRIPTOR_BYTES;
        uint32_t originalThunkRva = BootPeReadU32(descriptor);
        uint32_t nameRva = BootPeReadU32(descriptor + 12u);
        uint32_t firstThunkRva = BootPeReadU32(descriptor + 16u);

        if (originalThunkRva == 0 && nameRva == 0 && firstThunkRva == 0) {
            terminated = true;
            break;
        }
        if (!BootPeImageStringValid(consumer, consumerSize, nameRva) ||
            !BootPeAsciiEqualInsensitive((const char*)(consumer + nameRva), providerName)) {
            return BootPeUnsupportedImport;
        }

        uint32_t lookupRva = originalThunkRva ? originalThunkRva : firstThunkRva;
        if (!BootPeImageRangeValid(lookupRva, 8u, consumerSize) ||
            !BootPeImageRangeValid(firstThunkRva, 8u, consumerSize)) {
            return BootPeInvalidImage;
        }

        size_t lookupCapacity = (consumerSize - lookupRva) / 8u;
        size_t writeCapacity = (consumerSize - firstThunkRva) / 8u;
        size_t thunkCapacity = lookupCapacity < writeCapacity ? lookupCapacity : writeCapacity;
        bool thunkTerminated = false;

        for (size_t thunkIndex = 0; thunkIndex < thunkCapacity; ++thunkIndex) {
            uint64_t lookup = BootPeReadU64(consumer + lookupRva + thunkIndex * 8u);
            if (lookup == 0) {
                thunkTerminated = true;
                break;
            }
            if ((lookup & BOOT_PE_ORDINAL_FLAG64) != 0 || lookup > UINT32_MAX)
                return BootPeUnsupportedImport;

            uint32_t importByNameRva = (uint32_t)lookup;
            if (!BootPeImageRangeValid(importByNameRva, 3u, consumerSize) ||
                !BootPeImageStringValid(consumer, consumerSize, importByNameRva + 2u)) {
                return BootPeInvalidImage;
            }

            const char* functionName = (const char*)(consumer + importByNameRva + 2u);
            uint32_t functionRva = 0;
            status = BootPeFindExport(provider, providerSize, &providerView, functionName, &functionRva);
            if (status != BootPeSuccess)
                return status;

            if (providerVirtualBase > UINT64_MAX - functionRva)
                return BootPeInvalidImage;
            BootPeWriteU64(consumer + firstThunkRva + thunkIndex * 8u,
                           providerVirtualBase + functionRva);
        }

        if (!thunkTerminated)
            return BootPeInvalidImage;
    }

    return terminated ? BootPeSuccess : BootPeInvalidImage;
}

// Returns a concise diagnostic name for one portable loader status.
const char* BootPeStatusName(BOOT_PE_STATUS status) {
    switch (status) {
        case BootPeSuccess: return "success";
        case BootPeInvalidImage: return "invalid image";
        case BootPeWrongMachine: return "wrong machine";
        case BootPeChecksumMissing: return "checksum missing";
        case BootPeChecksumMismatch: return "checksum mismatch";
        case BootPeRelocationsPresent: return "relocations present";
        case BootPeBufferTooSmall: return "buffer too small";
        case BootPeImportsPresent: return "imports present";
        case BootPeUnsupportedImport: return "unsupported import";
        case BootPeImportNotFound: return "import not found";
        default: return "unknown PE status";
    }
}