/*++

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

Module Name:

   efilib.c

Abstract:

   Provides EFI utilities.

--*/

#include "efilib.h"
#include "mm.h"

EFI_GUID EfiSimpleTextInputExProtocol = EFI_SIMPLE_TEXT_INPUT_EX_PROTOCOL_GUID;

EFI_STATUS
EfiGetEfiStatusCode (
    IN NTSTATUS Status
    )

/*++

Routine Description:

    Converts an NT status code into an EFI status code.

Arguments:

    Status - The NT status code.

Return Value:

    The EFI status code.

--*/

{
    switch (Status) {
    case STATUS_SUCCESS:
        return EFI_SUCCESS;
    case STATUS_DRIVER_UNABLE_TO_LOAD:
        return EFI_LOAD_ERROR;
    case STATUS_INVALID_PARAMETER:
        return EFI_INVALID_PARAMETER;
    case STATUS_NOT_SUPPORTED:
        return EFI_UNSUPPORTED;
    case STATUS_INVALID_BUFFER_SIZE:
        return EFI_BAD_BUFFER_SIZE;
    case STATUS_BUFFER_TOO_SMALL:
        return EFI_BUFFER_TOO_SMALL;
    case STATUS_IO_DEVICE_ERROR:
        return EFI_DEVICE_ERROR;
    case STATUS_MEDIA_WRITE_PROTECTED:
        return EFI_WRITE_PROTECTED;
    case STATUS_INSUFFICIENT_RESOURCES:
    case STATUS_INSUFFICIENT_NVRAM_RESOURCES:
        return EFI_OUT_OF_RESOURCES;
    case STATUS_DISK_CORRUPT_ERROR:
        return EFI_VOLUME_CORRUPTED;
    case STATUS_DISK_FULL:
        return EFI_VOLUME_FULL;
    case STATUS_NO_MEDIA:
        return EFI_NO_MEDIA;
    case STATUS_MEDIA_CHANGED:
        return EFI_MEDIA_CHANGED;
    case STATUS_NOT_FOUND:
        return EFI_NOT_FOUND;
    case STATUS_ACCESS_DENIED:
        return EFI_ACCESS_DENIED;
    case STATUS_NO_MATCH:
        return EFI_NO_MAPPING;
    case STATUS_TIMEOUT:
        return EFI_TIMEOUT;
    case STATUS_DEVICE_NOT_READY:
        return EFI_NOT_STARTED;
    case STATUS_DEVICE_ALREADY_ATTACHED:
        return EFI_ALREADY_STARTED;
    case STATUS_REQUEST_ABORTED:
        return EFI_ABORTED;
    default:
        return EFI_NO_MAPPING;
    }
}

NTSTATUS
EfiGetNtStatusCode (
    IN EFI_STATUS Status
    )

/*++

Routine Description:

    Converts an EFI status code into an NT status code.

Arguments:

    Status - The EFI status code.

Return Value:

    The NT status code.

--*/

{
    switch (Status) {
    case EFI_SUCCESS:
        return STATUS_SUCCESS;
    case EFI_LOAD_ERROR:
        return STATUS_DRIVER_UNABLE_TO_LOAD;
    case EFI_INVALID_PARAMETER:
        return STATUS_INVALID_PARAMETER;
    case EFI_UNSUPPORTED:
        return STATUS_NOT_SUPPORTED;
    case EFI_BAD_BUFFER_SIZE:
        return STATUS_INVALID_BUFFER_SIZE;
    case EFI_BUFFER_TOO_SMALL:
        return STATUS_BUFFER_TOO_SMALL;
    case EFI_DEVICE_ERROR:
        return STATUS_IO_DEVICE_ERROR;
    case EFI_WRITE_PROTECTED:
        return STATUS_MEDIA_WRITE_PROTECTED;
    case EFI_OUT_OF_RESOURCES:
        return STATUS_INSUFFICIENT_NVRAM_RESOURCES;
    case EFI_VOLUME_CORRUPTED:
        return STATUS_DISK_CORRUPT_ERROR;
    case EFI_VOLUME_FULL:
        return STATUS_DISK_FULL;
    case EFI_NO_MEDIA:
        return STATUS_NO_MEDIA;
    case EFI_MEDIA_CHANGED:
        return STATUS_MEDIA_CHANGED;
    case EFI_NOT_FOUND:
    case EFI_NOT_READY:
        return STATUS_NOT_FOUND;
    case EFI_ACCESS_DENIED:
    case EFI_SECURITY_VIOLATION:
        return STATUS_ACCESS_DENIED;
    case EFI_NO_MAPPING:
        return STATUS_NO_MATCH;
    case EFI_TIMEOUT:
    case EFI_NO_RESPONSE:
        return STATUS_TIMEOUT;
    case EFI_NOT_STARTED:
        return STATUS_DEVICE_NOT_READY;
    case EFI_ALREADY_STARTED:
        return STATUS_DEVICE_ALREADY_ATTACHED;
    case EFI_ABORTED:
        return STATUS_REQUEST_ABORTED;
    default:
        return STATUS_UNSUCCESSFUL;
    }
}

NTSTATUS
EfiGetMemoryMap (
    IN OUT UINTN *MemoryMapSize,
    IN OUT EFI_MEMORY_DESCRIPTOR *MemoryMap,
    IN OUT UINTN *MapKey,
    IN OUT UINTN *DescriptorSize,
    IN OUT UINT32 *DescriptorVersion
    )

/*++

Routine Description:

    Wrapper for EFI_BOOT_SERVICES.GetMemoryMap(). Gets the firmware memory map
    and places it into a buffer.

Arguments:

    MemoryMapSize - pointer to the size of the buffer.

    MemoryMap - pointer to the buffer to store the memory map in.

    MapKey - ponter to the memory map key.

    DescriptorSize - pointer to the size of each memory map descriptor.

    DescriptorVersion - pointer to the version of memory map descriptors.

Return Value:

    STATUS_SUCCESS if successful.
    STATUS_INVALID_PARAMETER if MemoryMapSize and/or MemoryMap are invalid.
    STATUS_BUFFER_TOO_SMALL if MemoryMapSize is too small.

--*/

{
    EXECUTION_CONTEXT_TYPE ContextType;
    EFI_STATUS EfiStatus;

    ContextType = CurrentExecutionContext->Type;
    if (ContextType != ExecutionContextFirmware) {
        //
        // TODO: Translate addresses here.
        // Need MmArchTranslateVirtualAddress().
        //

        BlpArchSwitchContext(ExecutionContextFirmware);
    }

    EfiStatus = EfiBS->GetMemoryMap(
        MemoryMapSize,
        MemoryMap,
        MapKey,
        DescriptorSize,
        DescriptorVersion
    );

    if (ContextType != ExecutionContextFirmware) {
        BlpArchSwitchContext(ContextType);
    }

    return EfiGetNtStatusCode(EfiStatus);
}

NTSTATUS
EfiAllocatePages (
    IN EFI_ALLOCATE_TYPE Type,
    IN EFI_MEMORY_TYPE MemoryType,
    IN UINTN Pages,
    IN OUT EFI_PHYSICAL_ADDRESS *Memory
    )

/*++

Routine Description:

    Wrapper for EFI_BOOT_SERVICES.AllocatePages(). Allocates contiguous pages
    of physical memory.

Arguments:

    Type - the type of allocation.

    MemoryType - the type of memory to allocate.

    Pages - the number of pages to allocate.

    Memory - pointer to a physical address of the allocation.

Return Value:

    STATUS_SUCCESS if successful.
    STATUS_INVALID_PARAMETER if Type, MemoryType, and/or Memory are invalid.
    STATUS_INSUFFICIENT_NVRAM_RESOURCES if the pages could not be allocated.
    STATUS_NOT_FOUND if the pages could not be found.

--*/

{
    EXECUTION_CONTEXT_TYPE ContextType;
    EFI_STATUS EfiStatus;

    ContextType = CurrentExecutionContext->Type;
    if (ContextType != ExecutionContextFirmware) {
        BlpArchSwitchContext(ExecutionContextFirmware);
    }

    EfiStatus = EfiBS->AllocatePages(
        Type,
        MemoryType,
        Pages,
        Memory
    );

    if (ContextType != ExecutionContextFirmware) {
        BlpArchSwitchContext(ContextType);
    }

    return EfiGetNtStatusCode(EfiStatus);
}

NTSTATUS
EfiFreePages (
    IN EFI_PHYSICAL_ADDRESS Memory,
    IN UINTN Pages
    )

/*++

Routine Description:

    Wrapper for EFI_BOOT_SERVICES.FreePages(). Frees contiguous pages
    of physical memory.

Arguments:

    Memory - physical address of the pages to be freed.

    Pages - the number of pages to free.

Return Value:

    STATUS_SUCCESS if successful.
    STATUS_NOT_FOUND if the allocation was not found.
    STATUS_INVALID_PARAMETER Memory and/or Pages are invalid.

--*/

{
    EXECUTION_CONTEXT_TYPE ContextType;
    EFI_STATUS EfiStatus;

    ContextType = CurrentExecutionContext->Type;
    if (ContextType != ExecutionContextFirmware) {
        BlpArchSwitchContext(ExecutionContextFirmware);
    }

    EfiStatus = EfiBS->FreePages(
        Memory,
        Pages
    );

    if (ContextType != ExecutionContextFirmware) {
        BlpArchSwitchContext(ContextType);
    }

    return EfiGetNtStatusCode(EfiStatus);
}

NTSTATUS
EfiSetWatchdogTimer (
    IN UINTN Timeout,
    IN UINT64 WatchdogCode,
    IN UINTN DataSize,
    IN CHAR16 *WatchdogData
    )

/*++

Routine Description:

    Wrapper for EFI_BOOT_SERVICES.SetWatchdogTimer().
    Sets the watchdog timer.

Arguments:

    Timeout - The number of seconds to set the timer to.
        Setting this to 0 disables the timer.

    WatchdogCode - The code to set when an event occurs.

    DataSize - The size in bytes of WatchdogData.

    WatchdogData - Optional pointer to a string containing
        a description of the timer, possibly accompanied
        by binary data.

Return Value:

    STATUS_SUCCESS if successful.
    STATUS_INVALID_PARAMETER if WatchdogCode is invalid.
    STATUS_NOT_SUPPORTED if there is no Watchdog timer.
    STATUS_IO_DEVICE_ERROR if the timer could not be set.

--*/

{
    EXECUTION_CONTEXT_TYPE ContextType;
    EFI_STATUS EfiStatus;

    if (Timeout != 0 && WatchdogCode <= 0xffff) {
        return STATUS_INVALID_PARAMETER;
    }

    ContextType = CurrentExecutionContext->Type;
    if (ContextType != ExecutionContextFirmware) {
        //
        // TODO: Translate addresses here.
        // Need MmArchTranslateVirtualAddress().
        //

        BlpArchSwitchContext(ExecutionContextFirmware);
    }

    EfiStatus = EfiBS->SetWatchdogTimer(
        Timeout,
        WatchdogCode,
        DataSize,
        WatchdogData
    );

    if (ContextType != ExecutionContextFirmware) {
        BlpArchSwitchContext(ContextType);
    }

    return EfiGetNtStatusCode(EfiStatus);
}

NTSTATUS
EfiOpenProtocol (
    IN EFI_HANDLE Handle,
    IN EFI_GUID *Protocol,
    IN OUT VOID **Interface
    )

/*++

Routine Description:

    Wrapper that uses either EFI_BOOT_SERVICES.OpenProtocol() or HandleProtocol().
    Opens a handle to a protocol and finds its interface.

Arguments:

    Handle - The handle to the protocol to open.

    Protocol - The GUID of the protocol.

    Intercce - Pointer that recieves the address of the interface.

Return Value:

    STATUS_SUCCESS if successful.
    Error code on failure.

--*/

{
    EXECUTION_CONTEXT_TYPE ContextType;
    EFI_STATUS EfiStatus;
    NTSTATUS Status;

    if (MmTranslationType != TRANSLATION_TYPE_NONE) {
        //
        // TODO: Translate addresses.
        // Need EfiVmOpenProtocol().
        //
        DebugPrint(L"EfiOpenProtocol(): Virtual address translation not supported\r\n");
        return STATUS_NOT_IMPLEMENTED;
    }

    ContextType = CurrentExecutionContext->Type;
    if (ContextType != ExecutionContextFirmware) {
        BlpArchSwitchContext(ExecutionContextFirmware);
    }

    //
    // If supported, use OpenProtocol() (EFI 1.10+).
    // It helps ensure the protocol is not uninstalled
    // when still in use.
    //
    if (EfiST->Hdr.Revision >= EFI_MAKE_REVISION(1, 10)) {
        EfiStatus = EfiBS->OpenProtocol(
            Handle,
            Protocol,
            Interface,
            EfiImageHandle,
            NULL,
            EFI_OPEN_PROTOCOL_GET_PROTOCOL
        );
    } else {
        EfiStatus = EfiBS->HandleProtocol(
            Handle,
            Protocol,
            Interface
        );
    }

    if (ContextType != ExecutionContextFirmware) {
        BlpArchSwitchContext(ContextType);
    }

    //
    // Convert EFI status to NT status.
    //
    Status = EfiGetNtStatusCode(EfiStatus);
    if (!NT_SUCCESS(Status)) {
        *Interface = NULL;
    }

    return Status;
}

NTSTATUS
EfiConInExSetState (
    IN EFI_SIMPLE_TEXT_INPUT_EX_PROTOCOL *Protocol,
    IN EFI_KEY_TOGGLE_STATE *KeyToggleState
    )

/*++

Routine Description:

    Wrapper around EFI_SIMPLE_TEXT_INPUT_EX_PROTOCOL.SetState().
    Sets an input protocol's state.

Arguments:

    Protocol - The protocol to set the state of.

    KeyToggle State - The state to set.

Return Value:

    STATUS_SUCCESS if successful.
    STATUS_IO_DEVICE_ERROR if the state could not be set.
    STATUS_UNSUPPORTED if State is not supported by the device.

--*/

{
    EXECUTION_CONTEXT_TYPE ContextType;
    EFI_STATUS EfiStatus;

    ContextType = CurrentExecutionContext->Type;
    if (ContextType != ExecutionContextFirmware) {
        //
        // TODO: Translate addresses here.
        // Need MmArchTranslateVirtualAddress().
        //

        BlpArchSwitchContext(ExecutionContextFirmware);
    }

    EfiStatus = Protocol->SetState(
        Protocol,
        KeyToggleState
    );

    if (ContextType != ExecutionContextFirmware) {
        BlpArchSwitchContext(ContextType);
    }

    return EfiGetNtStatusCode(EfiStatus);
}