uefi: More sane handling of exiting boot services

stage23/lib/blib.c

@@ -78,26 +78,64 @@ bool efi_exit_boot_services(void) {
     EFI_STATUS status;
 
     EFI_MEMORY_DESCRIPTOR tmp_mmap[1];
-    UINTN mmap_size = sizeof(tmp_mmap);
-    UINTN mmap_key = 0, desc_size = 0, desc_ver = 0;
+    efi_mmap_size = sizeof(tmp_mmap);
+    UINTN mmap_key = 0;
 
     uefi_call_wrapper(gBS->GetMemoryMap, 5,
-        &mmap_size, tmp_mmap, &mmap_key, &desc_size, &desc_ver);
+        &efi_mmap_size, tmp_mmap, &mmap_key, &efi_desc_size, &efi_desc_ver);
 
-    status = uefi_call_wrapper(gBS->ExitBootServices, 2, efi_image_handle, mmap_key);
+    efi_mmap_size += 4096;
 
+    status = uefi_call_wrapper(gBS->AllocatePool, 3,
+        EfiLoaderData, efi_mmap_size, &efi_mmap);
     if (status)
-        panic("efi: Failed to exit boot services");
+        goto fail;
+
+    status = uefi_call_wrapper(gBS->GetMemoryMap, 5,
+        &efi_mmap_size, efi_mmap, &mmap_key, &efi_desc_size, &efi_desc_ver);
+    if (status)
+        goto fail;
+
+    // Be gone, UEFI!
+    status = uefi_call_wrapper(gBS->ExitBootServices, 2, efi_image_handle, mmap_key);
 
     asm volatile ("cli" ::: "memory");
 
+    if (status)
+        goto fail;
+
     pmm_reclaim_uefi_mem();
 
+    // Go through new EFI memmap and free up bootloader entries
+    size_t entry_count = efi_mmap_size / efi_desc_size;
+
+    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;
+            }
+        }
+    }
+
     efi_boot_services_exited = true;
 
     printv("efi: Exited boot services.\n");
 
     return true;
+
+fail:
+    panic("efi: Failed to exit boot services");
 }
 
 #endif

stage23/mm/pmm.h

@@ -19,6 +19,11 @@
 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);

stage23/mm/pmm.s2.c

@@ -51,6 +51,11 @@ 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:
@@ -310,6 +315,9 @@ 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

stage23/protos/linux.c

@@ -544,8 +544,6 @@ void linux_load(char *config, char *cmdline) {
     boot_params->efi_info.efi_memdesc_size    = efi_desc_size;
     boot_params->efi_info.efi_memdesc_version = efi_desc_ver;
 
-    pmm_release_uefi_mem();
-
     efi_exit_boot_services();
 #endif