path: root/boot/common/efi_loader.cpp

// ============================================================================
// efi_loader.cpp - Core UEFI loader logic (architecture-independent)
//
// Sequence:
//   1. Load kernel ELF from ESP
//   2. Parse ELF and load segments into memory
//   3. Get framebuffer via GOP
//   4. Find ACPI/FDT tables
//   5. Get final memory map + ExitBootServices
//   6. Populate BootInfo and return to arch entry for the final jump
// ============================================================================

#include <boot/efi.h>
#include <elf/elf.h>
#include <boot/boot_info.h>

// TODO: move it to include
// Stringify helper for kernel file path (wide-string literal from config macro).
#define _WIDE(x) L##x
#define WIDE(x)  _WIDE(x)

// Parse file and allocate memory for kernel
extern ElfLoadResult elf_load(const uint8_t* data, uint64_t size, EFI_BOOT_SERVICES* bs);


namespace efi {

static EFI_SYSTEM_TABLE*  gST = nullptr; /* global System Table, has pointers to console, boot services, runtime services and other*/
static EFI_BOOT_SERVICES* gBS = nullptr; /* global Boot Services, function table, that has memory allocation, protocol discovery, GetMemoryMap, ExitBootServices and other */

// Max memory-map entries
//TODO: fix problem with that
static constexpr uint64_t MAX_MEMORY_REGIONS = CONFIG_MAX_MEMORY_REGIONS;

// Low-level halt - safe to call even after ExitBootServices
[[noreturn]] static void panic_halt() noexcept {
    for (;;) {
#if defined(EFI_ARCH_X86_64)
        asm volatile("hlt");
#elif defined(EFI_ARCH_AARCH64)
        asm volatile("wfi");
#endif
    }
}

// Console helpers
void print(const CHAR16* msg) {
    gST->ConOut->OutputString(gST->ConOut, const_cast<CHAR16*>(msg));
}

void print_hex(uint64_t val) {
    CHAR16 buf[19];
    buf[0] = u'0';
    buf[1] = u'x';
    for (int i = 15; i >= 0; i--) {
        uint8_t nibble = (val >> (i * 4)) & 0xFu;
        buf[17 - i] = nibble < 10 ? (u'0' + nibble) : (u'A' + nibble - 10);
    }
    buf[18] = 0;
    print(buf);
}

[[noreturn]] void panic(const CHAR16* msg) {
    print(L"[PANIC] ");
    print(msg);
    print(L"\r\n");
    panic_halt();
}

// Type-safe protocol lookup helpers
//
// UEFI's HandleProtocol/LocateProtocol take void** which requires an explicit
// reinterpret_cast at every call site.  These templates absorb the cast.
// TODO: move it to efi.hpp
template<typename T>
[[nodiscard]] static EFI_STATUS handle_protocol(EFI_HANDLE h, EFI_GUID& guid, T** out) {
    return gBS->HandleProtocol(h, &guid, reinterpret_cast<void**>(out));
}

template<typename T>
[[nodiscard]] static EFI_STATUS locate_protocol(EFI_GUID& guid, T** out) {
    return gBS->LocateProtocol(&guid, nullptr, reinterpret_cast<void**>(out));
}

// RAII wrapper for EFI_FILE_PROTOCOL
//
// Closes the underlying file handle when it goes out of scope.
// Non-copyable; move is not needed in the loader's simple linear flow.
// TODO: maybe move it efi.h
class ScopedFile {
public:
    ScopedFile() = default;
    ScopedFile(const ScopedFile&) = delete;
    ScopedFile& operator=(const ScopedFile&) = delete;

    ~ScopedFile() { if (h) h->Close(h); }

    EFI_FILE_PROTOCOL* operator->() const noexcept { return h; }
    explicit operator bool()         const noexcept { return h != nullptr; }

    EFI_FILE_PROTOCOL*& get(){
        return h;
    }
private:
    EFI_FILE_PROTOCOL* h = nullptr;
};

// Load kernel file from ESP
//
// Returns a pointer to the raw ELF bytes allocated with AllocatePages.
// The caller (efi_loader_main) is responsible for keeping the pages alive
// until ExitBootServices is called (they are freed automatically thereafter
// because they are typed EfiLoaderData).
// image_handle: opaque pointer, that UEFI firmware passes to efi_main - it identifies loaded EFI application (our loader)
// out_size: size of the file
// return: pointer to file in allocated memory

uint8_t* load_kernel_file(EFI_HANDLE image_handle, uint64_t* out_size) {
    EFI_STATUS status;

    // Find the boot device via the loaded-image protocol.
    EFI_GUID lip_guid = EFI_LOADED_IMAGE_PROTOCOL_GUID;
    EFI_LOADED_IMAGE_PROTOCOL* loaded_image = nullptr;
    status = handle_protocol(image_handle, lip_guid, &loaded_image);
    if (EFI_ERROR(status))
        panic(L"Failed to get LoadedImageProtocol");

    // Open the filesystem on that device.
    EFI_GUID sfsp_guid = EFI_SIMPLE_FILE_SYSTEM_PROTOCOL_GUID;
    EFI_SIMPLE_FILE_SYSTEM_PROTOCOL* fs = nullptr;
    status = handle_protocol(loaded_image->DeviceHandle, sfsp_guid, &fs);
    if (EFI_ERROR(status))
        panic(L"Failed to get SimpleFileSystemProtocol");

    // Open the ESP root directory.
    ScopedFile root;
    status = fs->OpenVolume(fs, &root.get());
    if (EFI_ERROR(status))
        panic(L"Failed to open ESP volume");

    // Open kernel.elf.
    ScopedFile kernel_file;
    status = root->Open(root.get(), &kernel_file.get(),
                        const_cast<CHAR16*>(WIDE("\\" CONFIG_KERNEL_FILE_NAME)),
                        EFI_FILE_MODE_READ, 0);
    if (EFI_ERROR(status))
        panic(L"Failed to open " WIDE("\\" CONFIG_KERNEL_FILE_NAME));

    // Query file size via GetInfo.
    EFI_GUID fi_guid    = EFI_FILE_INFO_ID;
    UINTN    info_size  = sizeof(EFI_FILE_INFO) + 256; // 256 bytes extra for filename
    uint8_t  info_buf[sizeof(EFI_FILE_INFO) + 256];
    status = kernel_file->GetInfo(kernel_file.get(), &fi_guid, &info_size, info_buf);
    if (EFI_ERROR(status))
        panic(L"Failed to get kernel file info");

    const auto* file_info = reinterpret_cast<const EFI_FILE_INFO*>(info_buf);
    const uint64_t file_size = file_info->FileSize;

    // Allocate pages to hold the file.
    // Round-up division to convert a byte size into a number of 4 KiB pages
    const UINTN pages = (file_size + ELF_PAGE_MASK) / ELF_PAGE_SIZE;
    EFI_PHYSICAL_ADDRESS addr = 0;
    status = gBS->AllocatePages(AllocateAnyPages, EfiLoaderData, pages, &addr);
    if (EFI_ERROR(status))
        panic(L"Failed to allocate memory for kernel");

    // Read the file into the allocated buffer.
    auto* file_data = reinterpret_cast<uint8_t*>(addr);
    UINTN read_size = file_size;
    status = kernel_file->Read(kernel_file.get(), &read_size, file_data);
    if (EFI_ERROR(status))
        panic(L"Failed to read kernel file");

    *out_size = file_size;
    return file_data;
}

// Locate GOP framebuffer
//
// Graphics Output Protocol: https://uefi.org/specs/UEFI/2.10/12_Protocols_Console_Support.html
// Returns false if GOP is unavailable (non-fatal — the kernel can run headless).

bool get_framebuffer(FramebufferInfo* fb) {
    EFI_GUID gop_guid = EFI_GRAPHICS_OUTPUT_PROTOCOL_GUID;
    EFI_GRAPHICS_OUTPUT_PROTOCOL* gop = nullptr;
    EFI_STATUS status = locate_protocol(gop_guid, &gop);
    if (EFI_ERROR(status) || !gop || !gop->Mode || !gop->Mode->Info)
        return false;

    const auto* mode = gop->Mode;
    const auto* info = mode->Info;

    fb->base   = mode->FrameBufferBase;
    fb->width  = info->HorizontalResolution;
    fb->height = info->VerticalResolution;
    fb->pitch  = info->PixelsPerScanLine * 4;

    switch (info->PixelFormat) {
        case PixelRedGreenBlueReserved8BitPerColor: fb->format = PixelFormat::RGB;  break;
        case PixelBlueGreenRedReserved8BitPerColor: fb->format = PixelFormat::BGR;  break;
        default:                                    fb->format = PixelFormat::Mask; break;
    }
    return true;
}

// Search the UEFI Configuration Table
//
// The System Table has an array of {GUID, pointer} pairs published by various
// firmware subsystems.  Returns the VendorTable pointer, or 0 if not found.

uint64_t find_config_table(const EFI_GUID& target_guid) {
    for (UINTN i = 0; i < gST->NumberOfTableEntries; i++) {
        if (guid_equal(gST->ConfigurationTable[i].VendorGuid, target_guid))
            return reinterpret_cast<uint64_t>(gST->ConfigurationTable[i].VendorTable);
    }
    return 0;
}

// Convert UEFI memory descriptor type to MemoryRegionType

MemoryRegionType convert_memory_type(uint32_t efi_type) {
    switch (efi_type) {
        case EfiConventionalMemory:  return MemoryRegionType::Usable; /* Free RAM - the kernel can allocate this */
        case EfiACPIReclaimMemory:   return MemoryRegionType::AcpiReclaimable; /* ACPI tables - usable after the kernel is done parsing them */
        case EfiACPIMemoryNVS:       return MemoryRegionType::AcpiNvs; /* ACPI Non_volatile Storage - must be preserved */
        case EfiBootServicesCode:
        case EfiBootServicesData:
        case EfiLoaderCode:
        case EfiLoaderData:          return MemoryRegionType::BootloaderReclaimable; /* Memory used by our EFI loader - the kernel can reclaim it later */
        default:                     return MemoryRegionType::Reserved; /* Hardware-reserved, MMIO registers, firmware and other */
    }
}

// Pointer arithmetic helper for the packed memory descriptor array
//
// UEFI memory descriptors are desc_size bytes each (not sizeof(EFI_MEMORY_DESCRIPTOR)
//  firmware may use a larger struct with extra fields at the end).

[[nodiscard]] static const EFI_MEMORY_DESCRIPTOR* mem_desc_at(const void* map, UINTN index, UINTN desc_size) noexcept {
    return reinterpret_cast<const EFI_MEMORY_DESCRIPTOR*>(
        reinterpret_cast<uintptr_t>(map) + index * desc_size
    );
}

} // namespace efi

// ============================================================================
// Main loader — called by arch-specific entry.cpp
// ============================================================================

bool efi_loader_main(EFI_HANDLE image_handle, EFI_SYSTEM_TABLE* system_table, BootInfo* boot_info) {
    efi::gST = system_table;
    efi::gBS = system_table->BootServices;

    efi::gST->ConOut->ClearScreen(efi::gST->ConOut);
    efi::print(L"Bastion UEFI Loader\r\n");
    efi::print(L"=====================\r\n\r\n");

    // 1. Load kernel ELF from ESP.
    efi::print(L"Loading " WIDE("\\" CONFIG_KERNEL_FILE_NAME) " from ESP...\r\n");
    uint64_t kernel_file_size = 0;
    uint8_t* kernel_file = efi::load_kernel_file(image_handle, &kernel_file_size);
    efi::print(L"  Size: "); efi::print_hex(kernel_file_size); efi::print(L"\r\n");

    // 2. Parse and load ELF segments into memory.
    efi::print(L"Parsing ELF...\r\n");
    ElfLoadResult elf = elf_load(kernel_file, kernel_file_size, efi::gBS);
    if (!elf.success) {
        switch (elf.error) {
            case ELF_ERR_FILE_TOO_SMALL:  efi::panic(L"ELF: file too small");
            case ELF_ERR_INVALID_MAGIC:   efi::panic(L"ELF: bad magic");
            case ELF_ERR_INVALID_CLASS:   efi::panic(L"ELF: bad class (need ELF64)");
            case ELF_ERR_INVALID_IDENT:   efi::panic(L"ELF: bad endianness (need LE)");
            case ELF_ERR_INVALID_ARCH:    efi::panic(L"ELF: wrong machine type");
            case ELF_ERR_INVALID_TYPE:    efi::panic(L"ELF: not an executable");
            case ELF_ERR_NO_LOAD_SEGS:    efi::panic(L"ELF: no loadable segments");
            case ELF_ERR_ALLOC_FAILED:
                efi::print(L"ELF: AllocatePages failed at "); efi::print_hex(elf.virt_base);
                efi::print(L"  EFI status: "); efi::print_hex(elf.efi_alloc_status);
                efi::panic(L"\r\n");
            default:
                efi::panic(L"ELF: unknown error");
        }
    }
    efi::print(L"  Entry: "); efi::print_hex(elf.entry_point); efi::print(L"\r\n");
    efi::print(L"  Phys:  "); efi::print_hex(elf.phys_base);   efi::print(L"\r\n");

    // 3. Acquire framebuffer (non-fatal if absent).
    efi::print(L"Locating GOP...\r\n");
    if (!efi::get_framebuffer(&boot_info->framebuffer)) {
        efi::print(L"  WARNING: No GOP framebuffer found\r\n");
        #pragma message ("maybe panic too?")
    } else {
        efi::print(L"  FB at: "); efi::print_hex(boot_info->framebuffer.base); efi::print(L"\r\n");
    }

    // 4. Find firmware tables (ACPI RSDP or FDT depending on arch).
#ifdef CONFIG_ARCH_X86
    EFI_GUID acpi_guid = EFI_ACPI_20_TABLE_GUID;
    boot_info->rsdp_address = efi::find_config_table(acpi_guid);
#endif
#ifdef CONFIG_ARCH_AARCH64
    EFI_GUID fdt_guid = EFI_DTB_TABLE_GUID;
    boot_info->fdt_address = efi::find_config_table(fdt_guid);
#endif

    // 5. GetMemoryMap → ExitBootServices
    /*
        Memory map is array of entries
        each entry is
            - Base addr - where the region starts
            - Length - size of region
            - Type - look MemoryRegionType
        For example, 256MiB Qemu may have around 120 entreis.
        For now kernel prints all of them at boot, then sums up the Usable regions to report total free RAM
     
     map_size - buffer size on input; We pass in how big our buffer is. UEFI fills in how many bytes it actuall used.
     map_key - opaque snapshot token - must match when calling ExitBootServices.
                  (If an allocation happens between GetMemoryMap and ExitBootServices,
                  the key goes stale and ExitBootServices returns EFI_INVALID_PARAMETER)
     desc_size - actual size of each EFI_MEMORY_DESCRIPTOR entry in the returned map
                    (may be larger than sizeof(EFI_MEMORY_DESCRIPTOR) in newer
                     firmware — never assume the struct size).
    */
    efi::print(L"Exiting boot services...\r\n");

    UINTN map_size = 0, map_key = 0, desc_size = 0;
    [[maybe_unused]] uint32_t desc_version = 0; //unused
    EFI_MEMORY_DESCRIPTOR* efi_map = nullptr;

    // Get required size
    // First call: probe the required buffer size.
    efi::gBS->GetMemoryMap(&map_size, nullptr, &map_key, &desc_size, &desc_version);
    // We add desc_size*4 because next allocation will change memory map, and 4 descriptors mush be enough
    map_size += desc_size * 4;
    efi::gBS->AllocatePool(EfiLoaderData, map_size, reinterpret_cast<void**>(&efi_map));

    // Second call: fill the buffer.
    EFI_STATUS status = efi::gBS->GetMemoryMap(&map_size, efi_map, &map_key, &desc_size, &desc_version);
    if (EFI_ERROR(status))
        efi::panic(L"GetMemoryMap failed");

    // Point of No Return
    efi::print(L"EFI Point of No Return\r\n");
    status = efi::gBS->ExitBootServices(image_handle, map_key);
    if (EFI_ERROR(status)) {
        // Map changed between GetMemoryMap and ExitBootServices — retry once.
        // Do NOT allocate between this GetMemoryMap and ExitBootServices.
        efi::gBS->GetMemoryMap(&map_size, efi_map, &map_key, &desc_size, &desc_version);
        status = efi::gBS->ExitBootServices(image_handle, map_key);
        if (EFI_ERROR(status)) {
            // EFI services are in an unknown state — use the low-level halt;
            // calling efi::panic() (which uses ConOut) is unsafe here.
            efi::panic_halt();
        }
    }

    // ═══ NO MORE UEFI CALLS FROM THIS POINT ═══

    // 6. Convert the UEFI memory map to our compact MemoryRegion format.
    //
    // Convert big UEFI-specific struct (with fields like catching flags, vaddreses, attr and other)
    //      to our small map (maybe in future i will use more fierds from UEFI map)
    // Note: MAX_MEMORY_REGIONS (256) covers typical hardware.  Systems with
    // many RAM sticks, large MMIO holes, or complex firmware may exceed this.
    // TODO: compute the required count from map_size/desc_size before capping.
    // TODO: read how Linux solves it

    static MemoryRegion new_memory_map[efi::MAX_MEMORY_REGIONS];
    uint64_t new_memory_map_count = 0;

    const UINTN entry_count = map_size / desc_size;
    for (UINTN i = 0; i < entry_count && new_memory_map_count < efi::MAX_MEMORY_REGIONS; i++, new_memory_map_count++) {
        const auto* desc = efi::mem_desc_at(efi_map, i, desc_size);
        new_memory_map[new_memory_map_count].base   = desc->PhysicalStart;
        new_memory_map[new_memory_map_count].length = desc->NumberOfPages * ELF_PAGE_SIZE;
        new_memory_map[new_memory_map_count].type   = efi::convert_memory_type(desc->Type);
    }

    // Mark the kernel's physical pages as KernelAndModules so the
    // PMM (Phase 3) does not hand them out as free RAM.
    // The kernel was allocated with EfiLoaderData, so convert_memory_type()
    // classified it as BootLoaderReclaimable - the PMM would free it in Phase 3.
    {
        const uint64_t k_base = elf.phys_base;
        const uint64_t k_end  = k_base + elf.total_size;
        for (uint64_t i = 0; i < new_memory_map_count; ++i) {
            const uint64_t region_end = new_memory_map[i].base + new_memory_map[i].length;
            if (new_memory_map[i].base < k_end && region_end > k_base)
                new_memory_map[i].type = MemoryRegionType::KernelAndModules;
        }
    }

    // Mark the framebuffer region so the PMM never hands it out as free RAM.
    // GOP framebuffer memory isn't guaranteed a distinct EFI type (some firmware
    // reports it as EfiConventionalMemory), so we fix it up manually here now
    // that we know the framebuffer address and size.
    if (boot_info->framebuffer.base != 0) {
        const uint64_t fb_base = boot_info->framebuffer.base;
        const uint64_t fb_size = static_cast<uint64_t>(boot_info->framebuffer.height)
                               * boot_info->framebuffer.pitch;
        const uint64_t fb_end  = fb_base + fb_size;

        for (uint64_t i = 0; i < new_memory_map_count; i++) {
            const uint64_t region_end = new_memory_map[i].base + new_memory_map[i].length;
            if (new_memory_map[i].base < fb_end && region_end > fb_base)
                new_memory_map[i].type = MemoryRegionType::Framebuffer;
        }
    }

    // 7. Populate BootInfo for the kernel.
    boot_info->magic              = BOOT_INFO_MAGIC;
    boot_info->memory_map         = new_memory_map;
    boot_info->memory_map_count   = new_memory_map_count;
    boot_info->kernel_phys_base   = elf.phys_base;
    boot_info->kernel_virt_base   = elf.virt_base;
    boot_info->kernel_size        = elf.total_size;
    boot_info->kernel_entry_point = elf.entry_point;
    boot_info->hhdm_base          = 0;

    return true;
}