efi: Free malloc regions on exit

[grub.git] / grub-core / kern / efi / mm.c

diff --git a/grub-core/kern/efi/mm.c b/grub-core/kern/efi/mm.c
index ac2a4c5..c48e9b5 100644 (file)
--- a/grub-core/kern/efi/mm.c
+++ b/grub-core/kern/efi/mm.c
@@ -49,6 +49,70 @@ static grub_efi_uintn_t finish_desc_size;
 static grub_efi_uint32_t finish_desc_version;
 int grub_efi_is_finished = 0;
 
+/*
+ * We need to roll back EFI allocations on exit. Remember allocations that
+ * we'll free on exit.
+ */
+struct efi_allocation;
+struct efi_allocation {
+       grub_efi_physical_address_t address;
+       grub_efi_uint64_t pages;
+       struct efi_allocation *next;
+};
+static struct efi_allocation *efi_allocated_memory;
+
+static void
+grub_efi_store_alloc (grub_efi_physical_address_t address,
+                         grub_efi_uintn_t pages)
+{
+  grub_efi_boot_services_t *b;
+  struct efi_allocation *alloc;
+  grub_efi_status_t status;
+
+  b = grub_efi_system_table->boot_services;
+  status = efi_call_3 (b->allocate_pool, GRUB_EFI_LOADER_DATA,
+                           sizeof(*alloc), (void**)&alloc);
+
+  if (status == GRUB_EFI_SUCCESS)
+    {
+      alloc->next = efi_allocated_memory;
+      alloc->address = address;
+      alloc->pages = pages;
+      efi_allocated_memory = alloc;
+    }
+  else
+      grub_printf ("Could not malloc memory to remember EFI allocation. "
+                   "Exiting GRUB won't free all memory.\n");
+}
+
+static void
+grub_efi_drop_alloc (grub_efi_physical_address_t address,
+                           grub_efi_uintn_t pages)
+{
+  struct efi_allocation *ea, *eap;
+  grub_efi_boot_services_t *b;
+
+  b = grub_efi_system_table->boot_services;
+
+  for (eap = NULL, ea = efi_allocated_memory; ea; eap = ea, ea = ea->next)
+    {
+      if (ea->address != address || ea->pages != pages)
+         continue;
+
+      /* Remove the current entry from the list. */
+      if (eap)
+        eap->next = ea->next;
+      else
+        efi_allocated_memory = ea->next;
+
+      /* Then free the memory backing it. */
+      efi_call_1 (b->free_pool, ea);
+
+      /* And leave, we're done. */
+      break;
+    }
+}
+
 /* Allocate pages. Return the pointer to the first of allocated pages.  */
 void *
 grub_efi_allocate_pages_real (grub_efi_physical_address_t address,
@@ -79,6 +143,8 @@ grub_efi_allocate_pages_real (grub_efi_physical_address_t address,
        return 0;
     }
 
+  grub_efi_store_alloc (address, pages);
+
   return (void *) ((grub_addr_t) address);
 }
 
@@ -108,6 +174,8 @@ grub_efi_free_pages (grub_efi_physical_address_t address,
 
   b = grub_efi_system_table->boot_services;
   efi_call_2 (b->free_pages, address, pages);
+
+  grub_efi_drop_alloc (address, pages);
 }
 
 #if defined (__i386__) || defined (__x86_64__)
@@ -422,6 +490,20 @@ add_memory_regions (grub_efi_memory_descriptor_t *memory_map,
     grub_fatal ("too little memory");
 }
 
+void
+grub_efi_memory_fini (void)
+{
+  /*
+   * Free all stale allocations. grub_efi_free_pages() will remove
+   * the found entry from the list and it will always find the first
+   * list entry (efi_allocated_memory is the list start). Hence we
+   * remove all entries from the list until none is left altogether.
+   */
+  while (efi_allocated_memory)
+      grub_efi_free_pages (efi_allocated_memory->address,
+                           efi_allocated_memory->pages);
+}
+
 #if 0
 /* Print the memory map.  */
 static void