From a90cf727c524def4a78d903a59421e7c760edcd8 Mon Sep 17 00:00:00 2001
From: Rafal Kupiec <belliash@codingworkshop.eu.org>
Date: Tue, 2 Jan 2024 23:19:51 +0100
Subject: [PATCH] Implement BlGetMemoryMap() routine

---
 sdk/xtdk/bltypes.h      |  2 ++
 xtldr2/includes/xtldr.h |  4 +++
 xtldr2/memory.c         | 63 +++++++++++++++++++++++++++++++++++++++++
 xtldr2/protocol.c       |  1 +
 4 files changed, 70 insertions(+)

diff --git a/sdk/xtdk/bltypes.h b/sdk/xtdk/bltypes.h
index 062fb78..0d84500 100644
--- a/sdk/xtdk/bltypes.h
+++ b/sdk/xtdk/bltypes.h
@@ -73,6 +73,7 @@ typedef EFI_STATUS (*PBL_EXIT_BOOT_SERVICES)(IN UINT_PTR MapKey);
 typedef EFI_STATUS (*PBL_FIND_BOOT_PROTOCOL)(IN PWCHAR SystemType, OUT PEFI_GUID BootProtocolGuid);
 typedef EFI_STATUS (*PBL_FREE_PAGES)(IN UINT64 Size, IN EFI_PHYSICAL_ADDRESS Memory);
 typedef EFI_STATUS (*PBL_FREE_POOL)(IN PVOID Memory);
+typedef EFI_STATUS (*PBL_GET_MEMORY_MAP)(OUT PEFI_MEMORY_MAP MemoryMap);
 typedef INT_PTR (*PBL_GET_SECURE_BOOT_STATUS)();
 typedef EFI_STATUS (*PBL_INVOKE_BOOT_PROTOCOL)(IN PLIST_ENTRY OptionsList);
 typedef EFI_STATUS (*PBL_OPEN_VOLUME)(IN PEFI_DEVICE_PATH_PROTOCOL DevicePath, OUT PEFI_HANDLE DiskHandle, OUT PEFI_FILE_HANDLE *FsHandle);
@@ -235,6 +236,7 @@ typedef struct _XTBL_LOADER_PROTOCOL
         PBL_ALLOCATE_POOL AllocatePool;
         PBL_FREE_PAGES FreePages;
         PBL_FREE_POOL FreePool;
+        PBL_GET_MEMORY_MAP GetMemoryMap;
     } Memory;
     struct
     {
diff --git a/xtldr2/includes/xtldr.h b/xtldr2/includes/xtldr.h
index 3480f68..bebc8ca 100644
--- a/xtldr2/includes/xtldr.h
+++ b/xtldr2/includes/xtldr.h
@@ -102,6 +102,10 @@ XTCDECL
 PWCHAR
 BlGetConfigValue(IN CONST PWCHAR ConfigName);
 
+XTCDECL
+EFI_STATUS
+BlGetMemoryMap(OUT PEFI_MEMORY_MAP MemoryMap);
+
 XTCDECL
 INT_PTR
 BlGetSecureBootStatus();
diff --git a/xtldr2/memory.c b/xtldr2/memory.c
index 19e9397..8df0448 100644
--- a/xtldr2/memory.c
+++ b/xtldr2/memory.c
@@ -9,6 +9,69 @@
 #include <xtldr.h>
 
 
+/**
+ * Returns the memory descriptors which define a memory map of all the physical memory ranges reserved by the UEFI.
+ *
+ * @param MemoryMap
+ *        Supplies a pointer to the buffer where memory map will be written.
+ *
+ * @return This routine returns a status code.
+ *
+ * @since XT 1.0
+ */
+XTCDECL
+EFI_STATUS
+BlGetMemoryMap(OUT PEFI_MEMORY_MAP MemoryMap)
+{
+    EFI_STATUS Status;
+
+    if(MemoryMap == NULL)
+    {
+        return STATUS_EFI_INVALID_PARAMETER;
+    }
+
+    MemoryMap->Map = NULL;
+    MemoryMap->MapSize = 0;
+
+    /* Get memory map */
+    do
+    {
+        /* Attempt do get EFI memory map */
+        Status = EfiSystemTable->BootServices->GetMemoryMap(&MemoryMap->MapSize, MemoryMap->Map, &MemoryMap->MapKey,
+                                                            &MemoryMap->DescriptorSize, &MemoryMap->DescriptorVersion);
+        if(Status == STATUS_EFI_SUCCESS)
+        {
+            /* Go further if succeeded */
+            break;
+        }
+        else if(Status != STATUS_EFI_BUFFER_TOO_SMALL)
+        {
+            /* Some error occurred */
+            if(MemoryMap->Map)
+            {
+                /* Free allocated memory */
+                BlMemoryFreePool(MemoryMap->Map);
+            }
+            return Status;
+        }
+
+        /* Allocate the desired amount of memory */
+        MemoryMap->MapSize += 2 * MemoryMap->DescriptorSize;
+        BlMemoryAllocatePool(MemoryMap->MapSize, (PVOID *)&MemoryMap->Map);
+    }
+    while(Status == STATUS_EFI_BUFFER_TOO_SMALL);
+
+    /* Make sure memory map is set */
+    if(MemoryMap->Map == NULL)
+    {
+        /* Something went wrong */
+        return STATUS_EFI_NO_MAPPING;
+    }
+
+    /* Return success */
+    return STATUS_EFI_SUCCESS;
+}
+
 /**
  * This routine allocates one or more 4KB pages.
  *
diff --git a/xtldr2/protocol.c b/xtldr2/protocol.c
index 23f131d..3fdb77b 100644
--- a/xtldr2/protocol.c
+++ b/xtldr2/protocol.c
@@ -592,6 +592,7 @@ BlpRegisterXtLoaderProtocol()
     BlpLdrProtocol.Memory.AllocatePool = BlMemoryAllocatePool;
     BlpLdrProtocol.Memory.FreePages = BlMemoryFreePages;
     BlpLdrProtocol.Memory.FreePool = BlMemoryFreePool;
+    BlpLdrProtocol.Memory.GetMemoryMap = BlGetMemoryMap;
     BlpLdrProtocol.Protocol.Open = BlOpenProtocol;
     BlpLdrProtocol.Protocol.Register = BlRegisterProtocol;
     BlpLdrProtocol.Tui.DisplayErrorDialog = BlDisplayErrorDialog;