/*++

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 "bootlib.h"
#include "bootmgr.h"
#include "efi.h"

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

NTSTATUS
EfiInitpAppendPathString (
    IN PWCHAR Destination,
    IN ULONG BufferSize,
    IN PWCHAR Source,
    IN ULONG SourceSize,
    IN OUT PULONG BufferUsed
    )

/*++

Routine Description:

    Appends a soure path to a destination path.

Arguments:

    Destination - Path to append Source to.

    BufferSize - Maximum number of bytes to append to Destination.

    Source - Source path to append to Destination.

    SourceSize - Size of Source, in bytes.

    BufferUsed - Pointer to a ULONG recieving the number of bytes appended in.

Return Value:

    STATUS_SUCCESS if successful,
    STATUS_INVALID_PARAMETER if Destination is not valid,
    STATUS_BUFFER_TOO_SMALL if BufferSize is too small.

--*/

{
    ULONG Position;

    //
    // Verify that Source uses wide characters.
    //
    if (SourceSize % sizeof(WCHAR) != 0) {
        return STATUS_INVALID_PARAMETER;
    }

    //
    // Remove NULL terminator.
    //
    if (SourceSize >= sizeof(WCHAR)) {
        Position = (SourceSize / sizeof(WCHAR)) - 1;
        if (Source[Position] == UNICODE_NULL) {
            SourceSize -= sizeof(UNICODE_NULL);
        }
    }

    //
    // Remove leading separator.
    //
    if (SourceSize >= sizeof(WCHAR)) {
        if (Source[0] == L'\\') {
            Source++;
            SourceSize -= sizeof(WCHAR);
        }
    }

    //
    // Remove trailing separator.
    //
    if (SourceSize >= sizeof(WCHAR)) {
        Position = (SourceSize / sizeof(WCHAR)) - 1;
        if (Source[Position] == L'\\') {
            SourceSize -= sizeof(WCHAR);
        }
    }

    //
    // Check if Source is empty.
    //
    if (SourceSize == 0) {
        *BufferUsed = 0;
        return STATUS_SUCCESS;
    }

    //
    // Make sure the buffer is large enough.
    //
    if (BufferSize < SourceSize + sizeof(WCHAR)) {
        return STATUS_BUFFER_TOO_SMALL;
    }

    //
    // Append separator and Source to Destination.
    //
    Destination[0] = L'\\';
    RtlCopyMemory(Destination + 1, Source, SourceSize);

    *BufferUsed = SourceSize + sizeof(WCHAR);
    return STATUS_SUCCESS;
}

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 *.

Return Value:

    EFI_DEVICE_PATH *.

--*/

{
    EFI_DEVICE_PATH *Node;

    if (IsDevicePathEndType(DevicePath)) {
        return DevicePath;
    }

    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 a BOOT_DEVICE.

Arguments:

    EfiDevicePath - Path to be translated.

    BootDevice - Pointer to a buffer that recieves the device.

    BufferSize - Amount of available bytes 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;

    if (BufferSize < sizeof(BOOT_DEVICE)) {
        return STATUS_INVALID_PARAMETER;
    }
    BootDevice->Size = sizeof(BOOT_DEVICE);

    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;
    }

    //
    // TODO: Only media devices and ramdisks are currently supported.
    //
    DeviceNode = EfiInitpGetDeviceNode(EfiDevicePath);
    if (DevicePathType(DeviceNode) != MEDIA_DEVICE_PATH) {
        return STATUS_UNSUCCESSFUL;
    }

    switch (DevicePathSubType(DeviceNode)) {
    case MEDIA_HARDDRIVE_DP:
        HarddriveNode = (HARDDRIVE_DEVICE_PATH *)DeviceNode;

        //
        // MBR disks still use the old partition struct.
        //
        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;

        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_OPTION Option,
    IN ULONG BufferSize
    )

/*++

Routine Description:

    Converts an EFI device path into option format.

Arguments:

    EfiDevicePath - Path to be converted.

    OptionType - The data type to be assigned to Option->Type.

    Option - Pointer to a buffer that recieves the option.

    BufferSize - The amount of available bytes in the buffer.

Return Value:

    STATUS_SUCCESS if successful.
    STATUS_INVALID_PARAMETER if the buffer is too small.
    Any status code returned by EfiInitTranslateDevicePath().

--*/

{
    NTSTATUS Status;
    PBCDE_DEVICE DeviceElement;

    if (BufferSize < sizeof(BOOT_APPLICATION_OPTION) + FIELD_OFFSET(BCDE_DEVICE, Device)) {
        return STATUS_INVALID_PARAMETER;
    }

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

    Option->Type = OptionType;
    Option->DataOffset = sizeof(BOOT_APPLICATION_OPTION);
    Option->DataSize = FIELD_OFFSET(BCDE_DEVICE, Device) + DeviceElement->Device.Size;

    return STATUS_SUCCESS;
}

NTSTATUS
EfiInitpConvertEfiFilePath (
    IN EFI_DEVICE_PATH *EfiFilePath,
    IN BCDE_DATA_TYPE OptionType,
    IN OUT PBOOT_APPLICATION_OPTION Option,
    IN ULONG BufferSize
    )

/*++

Routine Description:

    Converts an EFI file path into option format.

Arguments:

    EfiFilePath - Path to be converted.

    OptionType - The data type to be assigned to Option->Type.

    Option - Pointer to a buffer that recieves the option.

    BufferSize - The amount of available bytes in the buffer.

Return Value:

    STATUS_SUCCESS if successful.
    STATUS_INVALID_PARAMETER if the buffer is too small.

--*/

{
    NTSTATUS Status;
    EFI_DEVICE_PATH *Node;
    PWCHAR PathStart, Position;
    ULONG BufferRemaining, Length, Appended;

    if (BufferSize < sizeof(BOOT_APPLICATION_OPTION)) {
        return STATUS_INVALID_PARAMETER;
    }

    RtlZeroMemory(Option, sizeof(BOOT_APPLICATION_OPTION));
    Option->Type = OptionType;
    Option->DataOffset = sizeof(BOOT_APPLICATION_OPTION);

    //
    // Loop through nodes and add one at a time.
    //
    Option->DataSize = 0;
    BufferRemaining = BufferSize - sizeof(BOOT_APPLICATION_OPTION);
    Node = EfiFilePath;
    PathStart = (PWCHAR)((PUCHAR)Option + Option->DataOffset);
    Position = PathStart;
    while (!IsDevicePathEndType(Node)) {
        //
        // Ignore non-filepath nodes.
        //
        if (DevicePathType(Node) != MEDIA_DEVICE_PATH || DevicePathSubType(Node) != MEDIA_FILEPATH_DP) {
            Node = NextDevicePathNode(Node);
            continue;
        }

        //
        // Find the length of this path.
        //
        Status = RtlULongSub(DevicePathNodeLength(Node), FIELD_OFFSET(FILEPATH_DEVICE_PATH, PathName), &Length);
        if (!NT_SUCCESS(Status)) {
            return Status;
        }

        //
        // Append this path to the path string.
        //
        Status = EfiInitpAppendPathString(Position, BufferRemaining, &((FILEPATH_DEVICE_PATH *)Node)->PathName[0], Length, &Appended);
        if (!NT_SUCCESS(Status)) {
            return Status;
        }

        //
        // Update counters & position.
        //
        Option->DataSize += Appended;
        BufferRemaining -= Appended;
        Position = (PWCHAR)((PUCHAR)Position + Appended);
        Node = NextDevicePathNode(Node);
    }

    //
    // NULL-terminate path string.
    //
    if (BufferRemaining < sizeof(UNICODE_NULL)) {
        return STATUS_INVALID_PARAMETER;
    }
    *Position = L'\0';
    Option->DataSize += sizeof(UNICODE_NULL);

    //
    // The option is invalid if the path is empty.
    //
    if (Position == PathStart) {
        Option->IsInvalid = TRUE;
    }

    return STATUS_SUCCESS;
}

VOID
EfiInitpCreateApplicationEntry (
    IN EFI_SYSTEM_TABLE *SystemTable,
    IN OUT PBOOT_INPUT_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 for the boot application.

Arguments:

    SystemTable - Pointer to the EFI system table.

    Entry - Pointer to a buffer that recieves the entry.

    BufferSize - The amount of available bytes in the buffer.

    EfiDevicePath - The application's device path.

    EfiFilePath - The application's file path.

    LoadOptions - Firmware load options string.

    LoadOptionsSize - Size in bytes of the string pointed to by LoadOptions.

    Flags - Unused.

    BufferUsed - Pointer to a ULONG that recieves the buffer space used by this routine.

    BootDevice - Pointer to a PBOOT_DEVICE that recieves the device the application was loaded from.

Return Value:

    None.

--*/

{
    NTSTATUS Status;
    ULONG BufferRemaining, OptionsSize, Size;
    PWCHAR BcdOptionString;
    BOOLEAN BcdIdentifierSet;
    UNICODE_STRING UnicodeString;
    PBOOT_APPLICATION_OPTION Option, PrevOption;
    PBCDE_DEVICE BootDeviceElement;

    (VOID)SystemTable;
    (VOID)Flags;

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

    BufferRemaining = BufferSize;
    if (BufferRemaining < sizeof(BOOT_INPUT_APPLICATION_ENTRY)) {
        return;
    }

    RtlZeroMemory(Entry, sizeof(BOOT_INPUT_APPLICATION_ENTRY));
    Entry->Signature = BOOT_INPUT_APPLICATION_ENTRY_SIGNATURE;
    BufferRemaining -= FIELD_OFFSET(BOOT_INPUT_APPLICATION_ENTRY, Options);

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

    //
    // Parse BCD GUID option 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_INPUT_APPLICATION_ENTRY_NO_BCD_IDENTIFIER;
    }

    //
    // Convert the EFI device path into a boot device option.
    //
    Option = &Entry->Options;
    Status = EfiInitpConvertEfiDevicePath(EfiDevicePath, BCDE_DATA_TYPE_APPLICATION_DEVICE, Option, BufferRemaining);
    if (!NT_SUCCESS(Status)) {
        Option->IsInvalid = TRUE;
        goto exit;
    }
    BootDeviceElement = (PBCDE_DEVICE)((PUCHAR)Option + Option->DataOffset);
    *BootDevice = &BootDeviceElement->Device;
    Size = BlGetBootOptionSize(Option);
    OptionsSize += Size;
    BufferRemaining -= Size;

    //
    // Convert the EFI file path into a boot file path option.
    // TODO: UDP/PXE are not supported.
    //
    PrevOption = Option;
    Option = (PBOOT_APPLICATION_OPTION)((PUCHAR)&Entry->Options + OptionsSize);
    Status = EfiInitpConvertEfiFilePath(EfiFilePath, BCDE_DATA_TYPE_APPLICATION_PATH, Option, BufferRemaining);
    if (!NT_SUCCESS(Status)) {
        goto exit;
    }
    PrevOption->NextOptionOffset = (PUCHAR)Option - (PUCHAR)&Entry->Options;
    Size = BlGetBootOptionSize(Option);
    OptionsSize += Size;
    BufferRemaining -= Size;

    //
    // TODO: Additional options in LoadOptions are not parsed.
    //
    PrevOption = Option;
    Option = (PBOOT_APPLICATION_OPTION)((PUCHAR)&Entry->Options + OptionsSize);
    // Status = Unknown(LoadOptions, &Entry->Options, RemainingSize, &OptionsSize, &PrevOption, &Size);
    if (!NT_SUCCESS(Status)) {
        goto exit;
    }

exit:
    *BufferUsed = BufferSize - BufferRemaining;
}

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 - EFI handle for the boot application image.

    SystemTable - Pointer to the EFI system table.

Return Value:

    Pointer to parameter structure if successful.
    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;
    PMEMORY_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);

    Status = SystemTable->BootServices->HandleProtocol(
        ImageHandle,
        (EFI_GUID*)&EfiLoadedImageProtocol,
        (VOID**)&LoadedImage
    );
    if (Status != EFI_SUCCESS) {
        return NULL;
    }

    Status = SystemTable->BootServices->HandleProtocol(
        LoadedImage->DeviceHandle,
        (EFI_GUID*)&EfiDevicePathProtocol,
        (VOID**)&DevicePath
    );
    if (Status != EFI_SUCCESS) {
        return NULL;
    }

    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_NONE;
    InputParameters->ImageBase = LoadedImage->ImageBase;
    InputParameters->ImageSize = LoadedImage->ImageSize;

    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(MEMORY_DESCRIPTOR);
    MemoryInfo->BasePageOffset = FIELD_OFFSET(MEMORY_DESCRIPTOR, BasePage);

    MemoryDescriptor = (PMEMORY_DESCRIPTOR)(&EfiInitScratch[ScratchUsed]);
    ScratchUsed += sizeof(MEMORY_DESCRIPTOR);
    MemoryDescriptor->BasePage = (UINTN)InputParameters->ImageBase >> PAGE_SHIFT;
    MemoryDescriptor->Pages = ALIGN_UP(InputParameters->ImageSize, PAGE_SIZE) >> PAGE_SHIFT;
    MemoryDescriptor->Attributes = MEMORY_ATTRIBUTE_WB;
    MemoryDescriptor->Type = MEMORY_TYPE_BOOT_APPLICATION;

    InputParameters->ApplicationEntryOffset = ScratchUsed;
    EfiInitpCreateApplicationEntry(
        SystemTable,
        (PBOOT_INPUT_APPLICATION_ENTRY)(&EfiInitScratch[ScratchUsed]),
        sizeof(EfiInitScratch) - ScratchUsed,
        DevicePath,
        LoadedImage->FilePath,
        LoadedImage->LoadOptions,
        LoadedImage->LoadOptionsSize,
        0,
        &ApplicationEntrySize,
        &BootDevice
    );
    ScratchUsed += ApplicationEntrySize;

    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);
    }

    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;

    InputParameters->ReturnDataOffset = ScratchUsed;
    ReturnData = (PBOOT_RETURN_DATA)(&EfiInitScratch[ScratchUsed]);
    ScratchUsed += sizeof(BOOT_RETURN_DATA);
    ReturnData->Version = BOOT_RETURN_DATA_VERSION;

    InputParameters->Size = ScratchUsed;
    if (InputParameters->Size > sizeof(EfiInitScratch)) {
        return NULL;
    }

    return InputParameters;
}