/*++

Copyright (c) 2024, Quinn Stephens.
Provided under the BSD 3-Clause license.

Module Name:

   efiinit.c

Abstract:

   Provides EFI initialization utilities.

--*/

#include <ntrtl.h>
#include <string.h>
#include <wchar.h>
#include "bootmgr.h"
#include "efi.h"

UCHAR EfiInitScratch[2048];
const EFI_GUID EfiLoadedImageProtocol = LOADED_IMAGE_PROTOCOL;
const EFI_GUID EfiDevicePathProtocol = DEVICE_PATH_PROTOCOL;

VOID
EfiInitpCreateApplicationEntry (
    IN EFI_SYSTEM_TABLE *SystemTable,
    IN OUT PBOOT_APPLICATION_ENTRY Entry,
    IN ULONG BufferSize,
    IN EFI_DEVICE_PATH *DevicePath,
    IN EFI_DEVICE_PATH *FilePath,
    IN PWCHAR LoadOptions,
    IN ULONG LoadOptionsSize,
    OUT PULONG BufferUsed,
    OUT PBOOT_DEVICE *BootDevice
    )

/*++

Routine Description:

    Creates an application entry structure for the boot application.

Arguments:

    SystemTable - Pointer to the EFI system table.

    Entry - A buffer to put the entry in.

    BufferSize - The amount of available space in the buffer.

    DevicePath - The device path for the application.

    FilePath - The file path for the application.

    LoadOptions - Firmware load options string.

    LoadOptionsSize - Length of the string pointed to by LoadOptions.

    BufferUsed - Returns the amount of buffer space used by the routine.

    BootDevice - Returns a pointer to the device the application was loaded from.

Return Value:

    None.

--*/

{
    PWCHAR BcdOptionString;
    BOOLEAN BcdIdentifierSet;
    UNICODE_STRING UnicodeString;

    *BufferUsed = 0;
    *BootDevice = NULL;
    BcdIdentifierSet = FALSE;

    //
    // Require enough space for the application entry.
    //
    if (BufferSize < sizeof(BOOT_APPLICATION_ENTRY)) {
        return;
    }

    //
    // Terminate load options string.
    //
    LoadOptionsSize /= sizeof(WCHAR);
    if (LoadOptionsSize != 0 && wcsnlen(LoadOptions, LoadOptionsSize) == LoadOptionsSize) {
        LoadOptions[LoadOptionsSize - 1] = L'\0';
    }

    //
    // Set up application entry structure.
    //
    RtlZeroMemory(Entry, sizeof(BOOT_APPLICATION_ENTRY));
    Entry->Signature = BOOT_APPLICATION_ENTRY_SIGNATURE;
    *BufferUsed = sizeof(BOOT_APPLICATION_ENTRY);

    //
    // Parse BCD GUID if present.
    //
    if (LoadOptions != NULL && (BcdOptionString = wcsstr(LoadOptions, L"BCDOBJECT=")) != NULL) {
        RtlInitUnicodeString(&UnicodeString, (PWCHAR)((PUCHAR)BcdOptionString + sizeof(L"BCDOBJECT=") - sizeof(UNICODE_NULL)));
        if (NT_SUCCESS(RtlGUIDFromString(&UnicodeString, &Entry->BcdIdentifier))) {
            BcdIdentifierSet = TRUE;
        }
    }

    if (!BcdIdentifierSet) {
        Entry->Attributes |= BOOT_APPLICATION_ENTRY_BCD_IDENTIFIER_NOT_SET;
    }

    //
    // TODO: This routine is not fully implemented.
    //
    (VOID)SystemTable;
    (VOID)DevicePath;
    (VOID)FilePath;
}

PBOOT_INPUT_PARAMETERS
EfiInitCreateInputParameters (
    IN EFI_HANDLE ImageHandle,
    IN EFI_SYSTEM_TABLE *SystemTable
    )

/*++

Routine Description:

    Creates the input parameter structure for the boot application.

Arguments:

    ImageHandle - Handle for the boot manager image.

    SystemTable - Pointer to the EFI system table.

Return Value:

    Pointer to parameter structure on success or NULL on failure.

--*/

{
    ULONG ScratchUsed = 0;
    ULONG ApplicationEntrySize = 0;
    EFI_STATUS Status;
    EFI_PHYSICAL_ADDRESS BadPageAddress;
    EFI_LOADED_IMAGE *LoadedImage;
    EFI_DEVICE_PATH *DevicePath;
    PBOOT_INPUT_PARAMETERS InputParameters;
    PBOOT_MEMORY_INFO MemoryInfo;
    PBOOT_MEMORY_DESCRIPTOR MemoryDescriptor;
    PBOOT_DEVICE BootDevice;
    PBOOT_FIRMWARE_DATA FirmwareData;
    PBOOT_RETURN_DATA ReturnData;

    //
    // Page 0x102 may be broken on some machines.
    // It is mapped here so that it does not get used.
    //
    BadPageAddress = 0x102 << PAGE_SHIFT;
    SystemTable->BootServices->AllocatePages(AllocateAddress, EfiLoaderData, 1, &BadPageAddress);

    //
    // Get boot manager image information.
    //
    Status = SystemTable->BootServices->HandleProtocol(
        ImageHandle,
        (EFI_GUID*)&EfiLoadedImageProtocol,
        (VOID**)&LoadedImage
    );
    if (Status != EFI_SUCCESS) {
        return NULL;
    }

    //
    // Get boot manager image device path.
    //
    Status = SystemTable->BootServices->HandleProtocol(
        LoadedImage->DeviceHandle,
        (EFI_GUID*)&EfiDevicePathProtocol,
        (VOID**)&DevicePath
    );
    if (Status != EFI_SUCCESS) {
        return NULL;
    }

    //
    // Create input parameters structure.
    //
    InputParameters = (PBOOT_INPUT_PARAMETERS)(&EfiInitScratch[ScratchUsed]);
    ScratchUsed += sizeof(BOOT_INPUT_PARAMETERS);
    InputParameters->Signature = BOOT_INPUT_PARAMETERS_SIGNATURE;
    InputParameters->Version = BOOT_INPUT_PARAMETERS_VERSION;
    InputParameters->MachineType = BOOT_MACHINE_TYPE;
    InputParameters->TranslationType = BOOT_TRANSLATION_TYPE;
    InputParameters->ImageBase = LoadedImage->ImageBase;
    InputParameters->ImageSize = LoadedImage->ImageSize;

    //
    // Create memory info structure.
    //
    InputParameters->MemoryInfoOffset = ScratchUsed;
    MemoryInfo = (PBOOT_MEMORY_INFO)(&EfiInitScratch[ScratchUsed]);
    ScratchUsed += sizeof(BOOT_MEMORY_INFO);
    MemoryInfo->Version = BOOT_MEMORY_INFO_VERSION;
    MemoryInfo->MdlOffset = sizeof(BOOT_MEMORY_INFO);
    MemoryInfo->DescriptorCount = 1;
    MemoryInfo->DescriptorSize = sizeof(BOOT_MEMORY_DESCRIPTOR);
    MemoryInfo->BasePageOffset = FIELD_OFFSET(BOOT_MEMORY_DESCRIPTOR, BasePage);

    //
    // Create a memory descriptor for the boot manager image.
    //
    MemoryDescriptor = (PBOOT_MEMORY_DESCRIPTOR)(&EfiInitScratch[ScratchUsed]);
    ScratchUsed += sizeof(BOOT_MEMORY_DESCRIPTOR);
    MemoryDescriptor->BasePage = (UINTN)InputParameters->ImageBase >> PAGE_SHIFT;
    MemoryDescriptor->Pages = ALIGN_UP(InputParameters->ImageSize, PAGE_SIZE) >> PAGE_SHIFT;
    MemoryDescriptor->Attributes = MEMORY_ATTRIBUTE_CACHE_WB;
    MemoryDescriptor->Type = MEMORY_TYPE_BOOT_APPLICATION;

    //
    // Create an application entry for the boot application.
    //
    InputParameters->ApplicationEntryOffset = ScratchUsed;
    EfiInitpCreateApplicationEntry(
        SystemTable,
        (PBOOT_APPLICATION_ENTRY)(&EfiInitScratch[ScratchUsed]),
        sizeof(EfiInitScratch) - ScratchUsed,
        DevicePath,
        LoadedImage->FilePath,
        LoadedImage->LoadOptions,
        LoadedImage->LoadOptionsSize,
        &ApplicationEntrySize,
        &BootDevice
    );
    ScratchUsed += ApplicationEntrySize;

    //
    // Copy application device to scratch area.
    //
    InputParameters->BootDeviceOffset = ScratchUsed;
    if (BootDevice != NULL) {
        RtlCopyMemory(&EfiInitScratch[ScratchUsed], BootDevice, BootDevice->Size);
        ScratchUsed += BootDevice->Size;
    } else {
        RtlZeroMemory(&EfiInitScratch[ScratchUsed], sizeof(BOOT_DEVICE));
        ScratchUsed += sizeof(BOOT_DEVICE);
    }

    //
    // Create firmware data structure.
    //
    InputParameters->FirmwareDataOffset = ScratchUsed;
    FirmwareData = (PBOOT_FIRMWARE_DATA)(&EfiInitScratch[ScratchUsed]);
    ScratchUsed += sizeof(BOOT_FIRMWARE_DATA);
    FirmwareData->Version = BOOT_FIRMWARE_DATA_VERSION;
    FirmwareData->Reserved = 0;
    FirmwareData->ImageHandle = ImageHandle;
    FirmwareData->SystemTable = SystemTable;

    //
    // Create return data structure.
    //
    InputParameters->ReturnDataOffset = ScratchUsed;
    ReturnData = (PBOOT_RETURN_DATA)(&EfiInitScratch[ScratchUsed]);
    ScratchUsed += sizeof(BOOT_RETURN_DATA);
    ReturnData->Version = BOOT_RETURN_DATA_VERSION;

    //
    // Set and validate total size.
    //
    InputParameters->Size = ScratchUsed;
    if (InputParameters->Size > sizeof(EfiInitScratch)) {
        return NULL;
    }

    return InputParameters;
}