:: commit a81f094d1dfde767c588f9960dcc4dc76d255f43

diff --git a/stage23/entry.s3.c b/stage23/entry.s3.c
index c210e2e4..46a07769 100644
--- a/stage23/entry.s3.c
+++ b/stage23/entry.s3.c
@@ -28,6 +28,7 @@ EFI_STATUS EFIAPI efi_main(EFI_HANDLE ImageHandle, EFI_SYSTEM_TABLE *SystemTable
     gST = SystemTable;
     gBS = SystemTable->BootServices;
     gRT = SystemTable->RuntimeServices;
+    efi_image_handle = ImageHandle;
 
     print("Limine " LIMINE_VERSION "\n\n", print);
 
@@ -87,6 +88,10 @@ void stage3_common(struct volume *boot_volume) {
     char *cmdline;
     char *config = menu(&cmdline);
 
+#if defined (uefi)
+    efi_exit_boot_services();
+#endif
+
     char *proto = config_get_value(config, 0, "PROTOCOL");
     if (proto == NULL) {
         panic("PROTOCOL not specified");
diff --git a/stage23/lib/blib.c b/stage23/lib/blib.c
index 6eee5e33..0d6019e7 100644
--- a/stage23/lib/blib.c
+++ b/stage23/lib/blib.c
@@ -7,11 +7,13 @@
 #include <lib/trace.h>
 #include <lib/real.h>
 #include <fs/file.h>
+#include <mm/pmm.h>
 
 #if defined (uefi)
 EFI_SYSTEM_TABLE *gST;
 EFI_BOOT_SERVICES *gBS;
 EFI_RUNTIME_SERVICES *gRT;
+EFI_HANDLE efi_image_handle;
 #endif
 
 bool parse_resolution(int *width, int *height, int *bpp, const char *buf) {
@@ -63,3 +65,33 @@ uint64_t sqrt(uint64_t a_nInput) {
 
     return res;
 }
+
+#if defined (uefi)
+
+bool efi_boot_services_exited = false;
+
+bool efi_exit_boot_services(void) {
+    EFI_MEMORY_DESCRIPTOR tmp_mmap[1];
+    UINTN mmap_size = sizeof(tmp_mmap);
+    UINTN mmap_key = 0, desc_size = 0, desc_ver = 0;
+
+    uefi_call_wrapper(gBS->GetMemoryMap, 5,
+        &mmap_size, tmp_mmap, &mmap_key, &desc_size, &desc_ver);
+
+    EFI_MEMORY_DESCRIPTOR *efi_mmap = ext_mem_alloc(mmap_size);
+
+    uefi_call_wrapper(gBS->GetMemoryMap, 5,
+        &mmap_size, efi_mmap, &mmap_key, &desc_size, &desc_ver);
+
+    uefi_call_wrapper(gBS->ExitBootServices, 2, efi_image_handle, mmap_key);
+
+    pmm_mmap_efi2ours(efi_mmap, desc_size, mmap_size / desc_size);
+
+    efi_boot_services_exited = true;
+
+    print("efi: Exited boot services.\n");
+
+    return true;
+}
+
+#endif
diff --git a/stage23/lib/blib.h b/stage23/lib/blib.h
index ad55a5d1..2fdd4e48 100644
--- a/stage23/lib/blib.h
+++ b/stage23/lib/blib.h
@@ -14,6 +14,10 @@
 extern EFI_SYSTEM_TABLE *gST;
 extern EFI_BOOT_SERVICES *gBS;
 extern EFI_RUNTIME_SERVICES *gRT;
+extern EFI_HANDLE efi_image_handle;
+
+extern bool efi_boot_services_exited;
+bool efi_exit_boot_services(void);
 #endif
 
 extern int boot_drive;
diff --git a/stage23/mm/pmm.h b/stage23/mm/pmm.h
index b5173f63..5f0aff1f 100644
--- a/stage23/mm/pmm.h
+++ b/stage23/mm/pmm.h
@@ -17,6 +17,13 @@
 extern struct e820_entry_t memmap[];
 extern size_t memmap_entries;
 
+#if defined (uefi)
+#include <efi.h>
+
+bool pmm_mmap_efi2ours(EFI_MEMORY_DESCRIPTOR *efi_mmap,
+                       size_t desc_size, size_t entry_count);
+#endif
+
 void init_memmap(void);
 struct e820_entry_t *get_memmap(size_t *entries);
 void print_memmap(struct e820_entry_t *mm, size_t size);
diff --git a/stage23/mm/pmm.s2.c b/stage23/mm/pmm.s2.c
index 31e2cc74..06addbbb 100644
--- a/stage23/mm/pmm.s2.c
+++ b/stage23/mm/pmm.s2.c
@@ -214,9 +214,8 @@ void *ext_mem_alloc_type(size_t count, uint32_t type) {
     return ext_mem_alloc_aligned_type(count, 4, type);
 }
 
-#if defined (bios)
 // Allocate memory top down, hopefully without bumping into kernel or modules
-void *ext_mem_alloc_aligned_type(size_t count, size_t alignment, uint32_t type) {
+static void *_ext_mem_alloc_aligned_type(size_t count, size_t alignment, uint32_t type) {
     if (allocations_disallowed)
         panic("Extended memory allocations disallowed");
 
@@ -257,11 +256,63 @@ void *ext_mem_alloc_aligned_type(size_t count, size_t alignment, uint32_t type)
 
     panic("High memory allocator: Out of memory");
 }
+
+#if defined (bios)
+void *ext_mem_alloc_aligned_type(size_t count, size_t alignment, uint32_t type) {
+    return _ext_mem_alloc_aligned_type(count, alignment, type);
+}
 #endif
 
 #if defined (uefi)
+bool pmm_mmap_efi2ours(EFI_MEMORY_DESCRIPTOR *efi_mmap,
+                       size_t desc_size, size_t entry_count) {
+    for (size_t i = 0; i < entry_count; i++) {
+        EFI_MEMORY_DESCRIPTOR *entry = (void *)efi_mmap + i * desc_size;
+
+        uint32_t our_type;
+        switch (entry->Type) {
+            case EfiReservedMemoryType:
+            case EfiRuntimeServicesCode:
+            case EfiRuntimeServicesData:
+            case EfiUnusableMemory:
+            case EfiMemoryMappedIO:
+            case EfiMemoryMappedIOPortSpace:
+            case EfiPalCode:
+            default:
+                our_type = MEMMAP_RESERVED; break;
+            case EfiBootServicesCode:
+            case EfiBootServicesData:
+            case EfiLoaderCode:
+            case EfiLoaderData:
+                our_type = MEMMAP_BOOTLOADER_RECLAIMABLE; break;
+            case EfiACPIReclaimMemory:
+                our_type = MEMMAP_ACPI_RECLAIMABLE; break;
+            case EfiACPIMemoryNVS:
+                our_type = MEMMAP_ACPI_NVS; break;
+            case EfiConventionalMemory:
+                our_type = MEMMAP_USABLE; break;
+        }
+
+        memmap[memmap_entries].type = our_type;
+        memmap[memmap_entries].base = entry->PhysicalStart;
+        memmap[memmap_entries].length = entry->NumberOfPages * 4096;
+
+        memmap_entries++;
+    }
+
+    sanitise_entries(false);
+
+    print_memmap(memmap, memmap_entries);
+
+    allocations_disallowed = false;
+
+    return true;
+}
+
 void *ext_mem_alloc_aligned_type(size_t count, size_t alignment, uint32_t type) {
-    (void)type;
+    if (efi_boot_services_exited) {
+        return _ext_mem_alloc_aligned_type(count, alignment, type);
+    }
 
     EFI_STATUS status;
 
@@ -275,7 +326,11 @@ void *ext_mem_alloc_aligned_type(size_t count, size_t alignment, uint32_t type)
     }
 
     if (status) {
-        panic("Memory allocation error %x\n", status);
+        panic("Memory allocation error %x", status);
+    }
+
+    if ((uintptr_t)ret % alignment) {
+        panic("Memory alloc align bad");
     }
 
     memset(ret, 0, count);