/*++

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

Module Name:

   efimm.c

Abstract:

   Provides EFI memory manager routines.

--*/

#include "bootmgr.h"
#include "efi.h"
#include "efilib.h"
#include "mm.h"

#define _1MiB 1048576

#define EFI_PAGE(NtPage) (((NtPage) << PAGE_SHIFT) >> EFI_PAGE_SHIFT)
#define NT_PAGE(EfiPage) (((EfiPage) << EFI_PAGE_SHIFT) >> PAGE_SHIFT)

extern EFI_BOOT_SERVICES *EfiBS;
extern PBOOT_DEVICE BlpBootDevice;

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,
    Other NTSTATUS value if an error occurs.

--*/

{
    return EfiGetNtStatusCode(
        EfiBS->GetMemoryMap(
            MemoryMapSize,
            MemoryMap,
            MapKey,
            DescriptorSize,
            DescriptorVersion
        )
    );
}

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,
    Other NTSTATUS value if an error occurs.

--*/

{
    return EfiGetNtStatusCode(
        EfiBS->AllocatePages(
            Type,
            MemoryType,
            Pages,
            Memory
        )
    );
}

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,
    Other NTSTATUS value if an error occurs.

--*/

{
    return EfiGetNtStatusCode(
        EfiBS->FreePages(
            Memory,
            Pages
        )
    );
}

MEMORY_TYPE
BlMmTranslateEfiMemoryType (
    IN EFI_MEMORY_TYPE EfiMemoryType
    )

/*++

Routine Description:

    Translates an EFI memory type to an NT memory type.

Arguments:

    EfiMemoryType - the EFI memory type.

Return Value:

    The NT memory type.

--*/

{
    switch (EfiMemoryType) {
    case EfiConventionalMemory:
        return MEMORY_TYPE_FREE;
    case EfiLoaderCode:
    case EfiLoaderData:
        return MEMORY_TYPE_BOOT_APPLICATION;
    case EfiBootServicesCode:
    case EfiBootServicesData:
        return MEMORY_TYPE_BOOT_SERVICES;
    case EfiRuntimeServicesCode:
        return MEMORY_TYPE_RUNTIME_SERVICES_CODE;
    case EfiRuntimeServicesData:
        return MEMORY_TYPE_RUNTIME_SERVICES_DATA;
    case EfiUnusableMemory:
        return MEMORY_TYPE_UNUSABLE;
    case EfiACPIReclaimMemory:
        return MEMORY_TYPE_ACPI_RECLAIM;
    case EfiACPIMemoryNVS:
        return MEMORY_TYPE_ACPI_NVS;
    case EfiMemoryMappedIO:
        return MEMORY_TYPE_MMIO;
    case EfiMemoryMappedIOPortSpace:
        return MEMORY_TYPE_MMIO_PORT_SPACE;
    case EfiPalCode:
        return MEMORY_TYPE_PAL_CODE;
    case EfiPersistentMemory:
        return MEMORY_TYPE_PERSISTENT;
    case EfiReservedMemoryType:
    default:
        if ((ULONG)EfiMemoryType < MAXLONG) {
            return MEMORY_TYPE_RESERVED;
        } else {
            return (MEMORY_TYPE)EfiMemoryType;
        }
    }
}

ULONG
MmFwpGetOsAttributeType (
    IN UINT64 EfiAttributes
    )

/*++

Routine Description:

    Translates EFI memory descriptor attributes to NT memory descriptor attributes.

Arguments:

    EfiAttributes - the EFI attributes.

Return Value:

    The NT attributes.

--*/

{
    ULONG NtAttributes;

    NtAttributes = 0;

    if (EfiAttributes & EFI_MEMORY_UC) {
        NtAttributes |= MEMORY_ATTRIBUTE_UC;
    }

    if (EfiAttributes & EFI_MEMORY_WC) {
        NtAttributes |= MEMORY_ATTRIBUTE_WC;
    }

    if (EfiAttributes & EFI_MEMORY_WT) {
        NtAttributes |= MEMORY_ATTRIBUTE_WT;
    }

    if (EfiAttributes & EFI_MEMORY_WB) {
        NtAttributes |= MEMORY_ATTRIBUTE_WB;
    }

    if (EfiAttributes & EFI_MEMORY_UCE) {
        NtAttributes |= MEMORY_ATTRIBUTE_UCE;
    }

    if (EfiAttributes & EFI_MEMORY_WP) {
        NtAttributes |= MEMORY_ATTRIBUTE_WP;
    }

    if (EfiAttributes & EFI_MEMORY_RP) {
        NtAttributes |= MEMORY_ATTRIBUTE_RP;
    }

    if (EfiAttributes & EFI_MEMORY_XP) {
        NtAttributes |= MEMORY_ATTRIBUTE_XP;
    }

    if (EfiAttributes & EFI_MEMORY_RUNTIME) {
        NtAttributes |= MEMORY_ATTRIBUTE_RUNTIME;
    }

    return NtAttributes;
}

NTSTATUS
MmFwGetMemoryMap (
   IN OUT PMEMORY_DESCRIPTOR_LIST Mdl,
   IN ULONG Flags
   )

/*++

Routine Description:

    Gets the memory map from EFI and converts it into an MDL.

Arguments:

    Mdl - Pointer to the MDL to store new descriptors in.

    Flags - Unused.

Return Value:

    STATUS_SUCCESS if successful,
    STATUS_INVALID_PARAMETER if Mdl is invalid.

--*/

{
    NTSTATUS Status;
    UINTN EfiMapSize;
    UINTN EfiMapKey;
    UINTN EfiDescriptorSize;
    UINT32 EfiDescriptorVersion;
    UINTN PageCount, EfiPageCount;
    EFI_PHYSICAL_ADDRESS EfiBuffer;
    EFI_MEMORY_DESCRIPTOR *EfiMap;
    BOOLEAN IsRamdisk;
    MEMORY_TYPE MemoryType;
    UINT64 EfiStartPage, EfiEndPage;
    UINT64 EfiRamdiskStartPage, EfiRamdiskEndPage;
    PMEMORY_DESCRIPTOR NtDescriptor;

    (VOID)Flags;

    EfiBuffer = (EFI_PHYSICAL_ADDRESS)0;

    if (Mdl == NULL) {
        return STATUS_INVALID_PARAMETER;
    }

    MmMdFreeList(Mdl);

    //
    // Get the memory map from firmware.
    // This is a "fake" call to actually just get
    // the required buffer size and other info.
    //
    EfiMapSize = 0;
    Status = EfiGetMemoryMap(&EfiMapSize, NULL, &EfiMapKey, &EfiDescriptorSize, &EfiDescriptorVersion);
    if (Status != STATUS_BUFFER_TOO_SMALL) {
        DebugPrint(L"MmFwGetMemoryMap(): EfiGetMemoryMap() failed\r\n");

        //
        // Make sure status is not successful, just in case
        // EfiGetMemoryMap() somehow succeeded.
        //
        if (NT_SUCCESS(Status)) {
            Status = STATUS_UNSUCCESSFUL;
        }

        goto exit;
    }

    EfiMapSize += 4 * EfiDescriptorSize;
    PageCount = (ALIGN_UP(EfiMapSize, PAGE_SIZE) >> PAGE_SHIFT) + 1;
    EfiPageCount = EFI_PAGE(PageCount);
    Status = EfiAllocatePages(AllocateAnyPages, EfiLoaderData, EfiPageCount, &EfiBuffer);
    if (!NT_SUCCESS(Status)) {
        DebugPrint(L"MmFwGetMemoryMap(): EfiAllocatePages() failed\r\n");
        goto exit;
    }

    EfiMap = (EFI_MEMORY_DESCRIPTOR *)EfiBuffer;
    Status = EfiGetMemoryMap(&EfiMapSize, EfiMap, &EfiMapKey, &EfiDescriptorSize, &EfiDescriptorVersion);
    if (!NT_SUCCESS(Status)) {
        DebugPrint(L"MmFwGetMemoryMap(): EfiGetMemoryMap() failed\r\n");
        goto exit;
    }

    if (EfiDescriptorSize < sizeof(EFI_MEMORY_DESCRIPTOR) || EfiMapSize % EfiDescriptorSize) {
        DebugPrint(L"MmFwGetMemoryMap(): Invalid EFI descriptor/map sizes\r\n");
        Status = STATUS_UNSUCCESSFUL;
        goto exit;
    }

    if (BlpBootDevice->Type == BOOT_DEVICE_TYPE_BLOCK && BlpBootDevice->Block.Type == BOOT_BLOCK_DEVICE_TYPE_RAMDISK) {
        IsRamdisk = TRUE;
        EfiRamdiskStartPage = BlpBootDevice->Block.Ramdisk.ImageBase.QuadPart >> EFI_PAGE_SHIFT;
        EfiRamdiskEndPage = EfiRamdiskStartPage + ((BlpBootDevice->Block.Ramdisk.ImageBase.QuadPart + BlpBootDevice->Block.Ramdisk.ImageSize) >> EFI_PAGE_SHIFT);
    } else {
        IsRamdisk = FALSE;
    }

    for (
        ;
        EfiMapSize > 0;
        EfiMapSize -= EfiDescriptorSize,
        EfiMap = (EFI_MEMORY_DESCRIPTOR *)((PUCHAR)EfiMap + EfiDescriptorSize)
    ) {
        //
        // Skip desciptors with no pages.
        //
        if (EfiMap->NumberOfPages == 0) {
            continue;
        }

        MemoryType = BlMmTranslateEfiMemoryType(EfiMap->Type);
        if (MemoryType == MEMORY_TYPE_FREE) {
            EfiStartPage = ALIGN_UP(EfiMap->PhysicalStart, EFI_PAGE_SIZE) >> EFI_PAGE_SHIFT;

        } else {
            EfiStartPage = EfiMap->PhysicalStart >> EFI_PAGE_SHIFT;
        }
        EfiEndPage = ALIGN_DOWN(EfiStartPage + EfiMap->NumberOfPages, EFI_PAGE(1));

        //
        // Regions under 1MiB are mapped differently.
        //
        if (EfiStartPage < (_1MiB >> EFI_PAGE_SHIFT)) {
            //
            // Reserve region at page 0.
            //
            if (EfiStartPage == 0) {
                NtDescriptor = MmMdInitDescriptor(
                    NT_PAGE(EfiStartPage),
                    0,
                    1,
                    MmFwpGetOsAttributeType(EfiMap->Attribute),
                    MEMORY_TYPE_RESERVED
                );
                if (NtDescriptor == NULL) {
                    Status = STATUS_INSUFFICIENT_RESOURCES;
                    goto exit;
                }

                Status = MmMdAddDescriptorToList(Mdl, NtDescriptor, MDL_OPERATION_FLAGS_TRUNCATE);
                if (!NT_SUCCESS(Status)) {
                    goto exit;
                }

                //
                // Continue if this descriptor was only one page.
                //
                if (EfiEndPage == 1) {
                    continue;
                }
            }

            //
            // For regions crossing over the 1MiB boundary,
            // create two descriptors. One for <1MiB and one
            // for over >=1MiB.
            //
            if (EfiEndPage > (_1MiB >> EFI_PAGE_SHIFT)) {
                NtDescriptor = MmMdInitDescriptor(
                    NT_PAGE(EfiStartPage),
                    0,
                    (_1MiB >> PAGE_SHIFT) - NT_PAGE(EfiStartPage),
                    MmFwpGetOsAttributeType(EfiMap->Attribute),
                    MemoryType
                );
                if (NtDescriptor == NULL) {
                    Status = STATUS_INSUFFICIENT_RESOURCES;
                    goto exit;
                }

                if (NtDescriptor->Type == MEMORY_TYPE_FREE) {
                    NtDescriptor->Attributes |= MEMORY_ATTRIBUTE_BELOW_1MIB;
                }

                Status = MmMdAddDescriptorToList(Mdl, NtDescriptor, MDL_OPERATION_FLAGS_TRUNCATE);
                if (!NT_SUCCESS(Status)) {
                    goto exit;
                }

                //
                // Continue to creating the >=1MiB mapping.
                //
                EfiStartPage = _1MiB >> EFI_PAGE_SHIFT;
            }
        }

        if (IsRamdisk) {
            if (
                EfiStartPage <= EfiRamdiskStartPage
                && EfiEndPage >= EfiRamdiskEndPage
            ) {
                if (NT_PAGE(EfiStartPage) < NT_PAGE(EfiRamdiskStartPage)) {
                    NtDescriptor = MmMdInitDescriptor(
                        NT_PAGE(EfiStartPage),
                        0,
                        NT_PAGE(EfiRamdiskStartPage) - NT_PAGE(EfiStartPage),
                        MmFwpGetOsAttributeType(EfiMap->Attribute),
                        MemoryType
                    );
                    if (NtDescriptor == NULL) {
                        Status = STATUS_INSUFFICIENT_RESOURCES;
                        goto exit;
                    }

                    Status = MmMdAddDescriptorToList(Mdl, NtDescriptor, MDL_OPERATION_FLAGS_TRUNCATE);
                    if (!NT_SUCCESS(Status)) {
                        goto exit;
                    }
                }

                //
                // Create a memory descriptor for the ramdisk.
                //
                NtDescriptor = MmMdInitDescriptor(
                    NT_PAGE(EfiRamdiskStartPage),
                    0,
                    NT_PAGE(EfiRamdiskEndPage) - NT_PAGE(EfiRamdiskStartPage),
                    MmFwpGetOsAttributeType(EfiMap->Attribute),
                    MemoryType
                );
                if (NtDescriptor == NULL) {
                    Status = STATUS_INSUFFICIENT_RESOURCES;
                    goto exit;
                }

                Status = MmMdAddDescriptorToList(Mdl, NtDescriptor, MDL_OPERATION_FLAGS_TRUNCATE);
                if (!NT_SUCCESS(Status)) {
                    goto exit;
                }

                //
                // Continue if there is no more memory to map inside the ramdisk.
                //
                if (NT_PAGE(EfiEndPage) <= NT_PAGE(EfiRamdiskEndPage)) {
                    continue;
                }

                EfiStartPage = EFI_PAGE(NT_PAGE(EfiRamdiskEndPage - 1) + 1);
            } else if (
                NT_PAGE(EfiStartPage) < NT_PAGE(EfiRamdiskStartPage)
                && NT_PAGE(EfiEndPage) > NT_PAGE(EfiRamdiskStartPage)
            ) {
                //
                // Remove the region inside the start of the ramdisk.
                //
                EfiEndPage = EfiRamdiskStartPage;
            } else if (
                NT_PAGE(EfiStartPage) < NT_PAGE(EfiRamdiskEndPage)
                && NT_PAGE(EfiEndPage) > NT_PAGE(EfiRamdiskEndPage)
            ) {
                //
                // Remove the region inside the end of the ramdisk.
                //
                EfiStartPage = EfiRamdiskEndPage;
            }

            //
            // Continue if the region is now empty.
            //
            if (NT_PAGE(EfiStartPage) == NT_PAGE(EfiEndPage)) {
                continue;
            }
        }

        //
        // Create a memory descriptor for the region.
        //
        NtDescriptor = MmMdInitDescriptor(
            NT_PAGE(EfiStartPage),
            0,
            NT_PAGE(EfiEndPage) - NT_PAGE(EfiStartPage),
            MmFwpGetOsAttributeType(EfiMap->Attribute),
            MemoryType
        );
        if (NtDescriptor == NULL) {
            Status = STATUS_INSUFFICIENT_RESOURCES;
            goto exit;
        }

        //
        // Set attribute if below 1MiB.
        //
        if (NtDescriptor->Type == MEMORY_TYPE_FREE && EfiEndPage <= (_1MiB >> EFI_PAGE_SHIFT)) {
            NtDescriptor->Attributes |= MEMORY_ATTRIBUTE_BELOW_1MIB;
        }

        Status = MmMdAddDescriptorToList(Mdl, NtDescriptor, MDL_OPERATION_FLAGS_TRUNCATE);
        if (!NT_SUCCESS(Status)) {
            goto exit;
        }
    }

    //
    // TODO: Account for possible MDL changes due to EfiFreePages(EfiBuffer).
    //

exit:
    if (EfiBuffer) {
        EfiFreePages(EfiBuffer, EfiPageCount);
    }

    if (!NT_SUCCESS(Status)) {
        MmMdFreeList(Mdl);
    }

    return Status;
}