From e6cc0e4b93d8c7f5fc4dd4f8726e34b65a8354b3 Mon Sep 17 00:00:00 2001
From: mintsuki <mintsuki@protonmail.com>
Date: Tue, 6 Jul 2021 08:21:20 +0200
Subject: [PATCH] uefi: Even saner handling of exiting boot services

---
 stage23/lib/blib.c  | 19 ++++++-------------
 stage23/mm/pmm.h    |  5 -----
 stage23/mm/pmm.s2.c | 10 +---------
 3 files changed, 7 insertions(+), 27 deletions(-)

diff --git a/stage23/lib/blib.c b/stage23/lib/blib.c
index a827d259..1f708b23 100644
--- a/stage23/lib/blib.c
+++ b/stage23/lib/blib.c
@@ -86,6 +86,10 @@ bool efi_exit_boot_services(void) {
 
     efi_mmap_size += 4096;
 
+    status = uefi_call_wrapper(gBS->FreePool, 1, efi_mmap);
+    if (status)
+        goto fail;
+
     status = uefi_call_wrapper(gBS->AllocatePool, 3,
         EfiLoaderData, efi_mmap_size, &efi_mmap);
     if (status)
@@ -112,19 +116,8 @@ bool efi_exit_boot_services(void) {
     for (size_t i = 0; i < entry_count; i++) {
         EFI_MEMORY_DESCRIPTOR *entry = (void *)efi_mmap + i * efi_desc_size;
 
-        uint64_t base = entry->PhysicalStart;
-        uint64_t length = entry->NumberOfPages * 4096;
-
-        // Find for a match in the untouched memory map
-        for (size_t j = 0; j < untouched_memmap_entries; j++) {
-            if (untouched_memmap[j].type != MEMMAP_USABLE)
-                continue;
-
-            if (untouched_memmap[j].base == base && untouched_memmap[j].length == length) {
-                // It's a match!
-                entry->Type = EfiConventionalMemory;
-                break;
-            }
+        if (entry->Type == 0x80000000) {
+            entry->Type = EfiConventionalMemory;
         }
     }
 
diff --git a/stage23/mm/pmm.h b/stage23/mm/pmm.h
index 7483478d..f0b8a30d 100644
--- a/stage23/mm/pmm.h
+++ b/stage23/mm/pmm.h
@@ -19,11 +19,6 @@
 extern struct e820_entry_t memmap[];
 extern size_t memmap_entries;
 
-#if defined (uefi)
-extern struct e820_entry_t untouched_memmap[];
-extern size_t untouched_memmap_entries;
-#endif
-
 void init_memmap(void);
 struct e820_entry_t *get_memmap(size_t *entries);
 struct e820_entry_t *get_raw_memmap(size_t *entry_count);
diff --git a/stage23/mm/pmm.s2.c b/stage23/mm/pmm.s2.c
index c0717559..921db84d 100644
--- a/stage23/mm/pmm.s2.c
+++ b/stage23/mm/pmm.s2.c
@@ -51,11 +51,6 @@ void *conv_mem_alloc(size_t count) {
 struct e820_entry_t memmap[MEMMAP_MAX_ENTRIES];
 size_t memmap_entries = 0;
 
-#if defined (uefi)
-struct e820_entry_t untouched_memmap[MEMMAP_MAX_ENTRIES];
-size_t untouched_memmap_entries = 0;
-#endif
-
 static const char *memmap_type(uint32_t type) {
     switch (type) {
         case MEMMAP_USABLE:
@@ -315,9 +310,6 @@ void init_memmap(void) {
 
     sanitise_entries(memmap, &memmap_entries, false);
 
-    memcpy(untouched_memmap, memmap, memmap_entries * sizeof(struct e820_entry_t));
-    untouched_memmap_entries = memmap_entries;
-
     allocations_disallowed = false;
 
     // Let's leave 64MiB to the firmware
@@ -331,7 +323,7 @@ void init_memmap(void) {
         EFI_PHYSICAL_ADDRESS base = memmap[i].base;
 
         status = uefi_call_wrapper(gBS->AllocatePages, 4,
-          AllocateAddress, EfiLoaderData, memmap[i].length / 4096, &base);
+          AllocateAddress, 0x80000000, memmap[i].length / 4096, &base);
 
         if (status)
             panic("pmm: AllocatePages failure (%x)", status);