/*++

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;

EFI_DEVICE_PATH *
EfiInitpGetDeviceNode (
    IN EFI_DEVICE_PATH *DevicePath
    )

/*++

Routine Description:

    Searches an EFI device path for the last device path node
    before a file path node.

Arguments:

    DevicePath - EFI device path to search.

Return Value:

    Pointer to the last device path node.

--*/

{
    EFI_DEVICE_PATH *Node;

    //
    // Check if the current node is the end of the path.
    //
    if (IsDevicePathEndType(DevicePath)) {
        return DevicePath;
    }

    //
    // Find the last non-filepath node.
    //
    Node = NextDevicePathNode(DevicePath);
    while (!IsDevicePathEndType(Node)) {
        if (DevicePathType(Node) == MEDIA_DEVICE_PATH && DevicePathSubType(Node) == MEDIA_FILEPATH_DP) {
            break;
        }

        DevicePath = Node;
        Node = NextDevicePathNode(Node);
    }

    return DevicePath;
}

NTSTATUS
EfiInitTranslateDevicePath (
    IN EFI_DEVICE_PATH *EfiDevicePath,
    IN OUT PBOOT_DEVICE BootDevice,
    IN ULONG BufferSize
    )

/*++

Routine Description:

    Translates an EFI device path into boot device format.

Arguments:

    EfiDevicePath - The EFI device path to be translated.

    BootDevice - Pointer to the destination device structure.

    BufferSize - The amount of available space in the buffer.

Return Value:

    STATUS_SUCCESS if successful,
    STATUS_INVALID_PARAMETER if the buffer is too small.
    STATUS_UNSUCCESSFUL if the path could not be translated.

--*/

{
    EFI_DEVICE_PATH *DeviceNode;
    MEMMAP_DEVICE_PATH *MemmapNode;
    HARDDRIVE_DEVICE_PATH *HarddriveNode;
    PBOOT_BLOCK_IDENTIFIER BlockDevice;

    //
    // Check for available buffer space.
    //
    if (BufferSize < sizeof(BOOT_DEVICE)) {
        return STATUS_INVALID_PARAMETER;
    }
    BootDevice->Size = sizeof(BOOT_DEVICE);

    //
    // Memory mapped device paths are treated as ramdisks.
    //
    if (DevicePathType(EfiDevicePath) == HARDWARE_DEVICE_PATH && DevicePathSubType(EfiDevicePath) == HW_MEMMAP_DP) {
        MemmapNode = (MEMMAP_DEVICE_PATH *)EfiDevicePath;
        BlockDevice = &BootDevice->Block;
        BootDevice->Type = BOOT_DEVICE_TYPE_BLOCK;
        BlockDevice->Type = BOOT_BLOCK_DEVICE_TYPE_RAMDISK;
        BlockDevice->Ramdisk.ImageBase.QuadPart = MemmapNode->StartingAddress;
        BlockDevice->Ramdisk.ImageSize = MemmapNode->EndingAddress - MemmapNode->StartingAddress;
        BlockDevice->Ramdisk.ImageOffset = 0;
        return STATUS_SUCCESS;
    }

    //
    // Get the device node, the device the application was loaded from.
    // TODO: Only media devices and ramdisks are currently supported.
    //
    DeviceNode = EfiInitpGetDeviceNode(EfiDevicePath);
    if (DevicePathType(DeviceNode) != MEDIA_DEVICE_PATH) {
        return STATUS_UNSUCCESSFUL;
    }

    //
    // Check device node subtype.
    //
    switch (DevicePathSubType(DeviceNode)) {
    case MEDIA_HARDDRIVE_DP:
        HarddriveNode = (HARDDRIVE_DEVICE_PATH *)DeviceNode;

        //
        // Use correct block device and partition format.
        //
        if (HarddriveNode->SignatureType != SIGNATURE_TYPE_MBR) {
            BlockDevice = &BootDevice->PartitionEx.Parent;
            BootDevice->Type = BOOT_DEVICE_TYPE_PARTITION_EX;
        } else {
            BlockDevice = &BootDevice->Partition.Parent;
            BootDevice->Type = BOOT_DEVICE_TYPE_PARTITION;
        }
        BlockDevice->Type = BOOT_BLOCK_DEVICE_TYPE_HARDDRIVE;

        //
        // Initialize partition based on the drive's partitioning system.
        //
        switch (HarddriveNode->SignatureType) {
        case SIGNATURE_TYPE_MBR:
            BlockDevice->Harddrive.PartitionType = BOOT_HARDDRIVE_PARTITION_TYPE_MBR;
            BlockDevice->Harddrive.Mbr.Signature = *((ULONG *)HarddriveNode->Signature);
            BootDevice->Partition.Mbr.PartitionNumber = HarddriveNode->PartitionNumber;
            break;
        case SIGNATURE_TYPE_GUID:
            BootDevice->Attributes |= BOOT_DEVICE_ATTRIBUTE_NO_PARENT_SIGNATURE;
            BlockDevice->Harddrive.PartitionType = BOOT_HARDDRIVE_PARTITION_TYPE_GPT;
            RtlCopyMemory(&BootDevice->PartitionEx.Gpt.PartitionIdentifier, &HarddriveNode->Signature, sizeof(HarddriveNode->Signature));
            break;
        default:
            BlockDevice->Harddrive.PartitionType = BOOT_HARDDRIVE_PARTITION_TYPE_RAW;
            BlockDevice->Harddrive.Raw.DriveNumber = 0;
        }

        break;
    case MEDIA_CDROM_DP:
        BootDevice->Type = BOOT_DEVICE_TYPE_BLOCK;
        BootDevice->Block.Type = BOOT_BLOCK_DEVICE_TYPE_CDROM;
        BootDevice->Block.Cdrom.DriveNumber = 0;
        break;
    default:
        return STATUS_UNSUCCESSFUL;
    }

    return STATUS_SUCCESS;
}

NTSTATUS
EfiInitpConvertEfiDevicePath (
    IN EFI_DEVICE_PATH *EfiDevicePath,
    IN BCDE_DATA_TYPE OptionType,
    IN OUT PBOOT_APPLICATION_ENTRY_OPTION Option,
    IN ULONG BufferSize
    )

/*++

Routine Description:

    Converts an EFI device path into BCD format.

Arguments:

    EfiDevicePath - The EFI device path to be converted.

    OptionType - The data type to be assigned to option.

    Option - Pointer to the destination option structure.

    BufferSize - The amount of available space in the buffer.

Return Value:

    STATUS_SUCCESS if successful.
    other NTSTATUS value if failure occurs.

--*/

{
    NTSTATUS Status;
    PBCDE_DEVICE DeviceElement;

    //
    // Check for available buffer space.
    //
    if (BufferSize < sizeof(BOOT_APPLICATION_ENTRY_OPTION) + FIELD_OFFSET(BCDE_DEVICE, Device)) {
        return STATUS_INVALID_PARAMETER;
    }

    //
    // Translate device path.
    //
    RtlZeroMemory(Option, sizeof(BOOT_APPLICATION_ENTRY_OPTION));
    DeviceElement = (PBCDE_DEVICE)((PUCHAR)Option + sizeof(BOOT_APPLICATION_ENTRY_OPTION));
    Status = EfiInitTranslateDevicePath(
        EfiDevicePath,
        &DeviceElement->Device,
        BufferSize - (sizeof(BOOT_APPLICATION_ENTRY_OPTION) + FIELD_OFFSET(BCDE_DEVICE, Device))
    );
    if (!NT_SUCCESS(Status)) {
        return Status;
    }

    //
    // Set up option structure.
    //
    Option->Type = OptionType;
    Option->DataOffset = sizeof(BOOT_APPLICATION_ENTRY_OPTION);
    Option->DataSize = FIELD_OFFSET(BCDE_DEVICE, Device) + DeviceElement->Device.Size;

    return STATUS_SUCCESS;
}

VOID
EfiInitpCreateApplicationEntry (
    IN EFI_SYSTEM_TABLE *SystemTable,
    IN OUT PBOOT_APPLICATION_ENTRY Entry,
    IN ULONG BufferSize,
    IN EFI_DEVICE_PATH *EfiDevicePath,
    IN EFI_DEVICE_PATH *EfiFilePath,
    IN PWCHAR LoadOptions,
    IN ULONG LoadOptionsSize,
    IN ULONG Flags,
    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.

    EfiDevicePath - The device path for the application.

    EfiFilePath - The file path for the application.

    LoadOptions - Firmware load options string.

    LoadOptionsSize - Length of the string pointed to by LoadOptions.

    Flags - Unused.

    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.

--*/

{
    NTSTATUS Status;
    ULONG BufferRemaining;
    PWCHAR BcdOptionString;
    BOOLEAN BcdIdentifierSet;
    UNICODE_STRING UnicodeString;
    PBOOT_APPLICATION_ENTRY_OPTION Option;
    PBCDE_DEVICE BootDeviceElement;

    (VOID)SystemTable;
    (VOID)EfiDevicePath;
    (VOID)EfiFilePath;
    (VOID)Flags;

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

    //
    // Require enough space for the application entry.
    //
    BufferRemaining = BufferSize;
    if (BufferRemaining < 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_NO_BCD_IDENTIFIER;
    }

    //
    // Convert the EFI device path into a boot device option.
    //
    Option = &Entry->Options;
    BufferRemaining -= FIELD_OFFSET(BOOT_APPLICATION_ENTRY, Options);
    Status = EfiInitpConvertEfiDevicePath(EfiDevicePath, BCDE_DATA_TYPE_APPLICATION_DEVICE, Option, BufferRemaining);
    if (!NT_SUCCESS(Status)) {
        Option->IsInvalid = TRUE;
        *BufferUsed = sizeof(BOOT_APPLICATION_ENTRY_OPTION);
        return;
    }

    BootDeviceElement = (PBCDE_DEVICE)((PUCHAR)Option + Option->DataOffset);
    *BootDevice = &BootDeviceElement->Device;

    //
    // TODO: This routine is not fully implemented.
    //
}

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,
        0,
        &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;
}