[PATCH v2] efi: Free malloc regions on exit

[Alexander Graf]

When we exit grub, we don't free all the memory that we allocated earlier
for our heap region. This can cause problems with setups where you try
to descend the boot order using "exit" entries, such as PXE -> HD boot
scenarios.

Signed-off-by: Alexander Graf <address@hidden>

---

v1 -> v2:

  - add comment explaining the number of regions
  - move nr of regions into a define
  - add warning if we exceed the number of freeable regions
  - reset region counter to 0 on fini
---
 grub-core/kern/efi/init.c |  1 +
 grub-core/kern/efi/mm.c   | 36 ++++++++++++++++++++++++++++++++++++
 include/grub/efi/efi.h    |  1 +
 3 files changed, 38 insertions(+)

diff --git a/grub-core/kern/efi/init.c b/grub-core/kern/efi/init.c
index e9c85de..b848014 100644
--- a/grub-core/kern/efi/init.c
+++ b/grub-core/kern/efi/init.c
@@ -77,4 +77,5 @@ grub_efi_fini (void)
 {
   grub_efidisk_fini ();
   grub_console_fini ();
+  grub_efi_memory_fini ();
 }
diff --git a/grub-core/kern/efi/mm.c b/grub-core/kern/efi/mm.c
index 20a47aa..6f1cb1e 100644
--- a/grub-core/kern/efi/mm.c
+++ b/grub-core/kern/efi/mm.c
@@ -49,6 +49,18 @@ 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 up to a randomly
+ * chosen 16 allocations that we'll free on exit. In most cases only one
+ * will be used.
+ */
+#define MAX_REVERSIBLE_EFI_ALLOCATIONS 16
+static struct efi_allocation {
+       grub_efi_physical_address_t start_addr;
+       grub_efi_uint64_t pages;
+} efi_allocated_memory[MAX_REVERSIBLE_EFI_ALLOCATIONS];
+unsigned int efi_allocated_memory_idx = 0;
+
 /* Allocate pages. Return the pointer to the first of allocated pages.  */
 void *
 grub_efi_allocate_pages (grub_efi_physical_address_t address,
@@ -408,6 +420,17 @@ add_memory_regions (grub_efi_memory_descriptor_t 
*memory_map,
                    (void *) ((grub_addr_t) start),
                    (unsigned) pages);
 
+      if (efi_allocated_memory_idx < ARRAY_SIZE(efi_allocated_memory)) {
+        efi_allocated_memory[efi_allocated_memory_idx].start_addr = start;
+        efi_allocated_memory[efi_allocated_memory_idx].pages = pages;
+        efi_allocated_memory_idx++;
+
+        if (efi_allocated_memory_idx == ARRAY_SIZE(efi_allocated_memory)) {
+          grub_printf ("EFI memory allocations exceeded limit. Exiting grub2 "
+                       "won't free all memory.\n");
+        }
+      }
+
       grub_mm_init_region (addr, PAGES_TO_BYTES (pages));
 
       required_pages -= pages;
@@ -419,6 +442,19 @@ add_memory_regions (grub_efi_memory_descriptor_t 
*memory_map,
     grub_fatal ("too little memory");
 }
 
+void
+grub_efi_memory_fini (void)
+{
+  unsigned int i;
+
+  for (i = 0; i < efi_allocated_memory_idx; i++) {
+    grub_efi_free_pages (efi_allocated_memory[i].start_addr,
+                         efi_allocated_memory[i].pages);
+  }
+
+  efi_allocated_memory_idx = 0;
+}
+
 #if 0
 /* Print the memory map.  */
 static void
diff --git a/include/grub/efi/efi.h b/include/grub/efi/efi.h
index 0e6fd86..545e7ce 100644
--- a/include/grub/efi/efi.h
+++ b/include/grub/efi/efi.h
@@ -48,6 +48,7 @@ EXPORT_FUNC(grub_efi_get_memory_map) (grub_efi_uintn_t 
*memory_map_size,
                                      grub_efi_uintn_t *map_key,
                                      grub_efi_uintn_t *descriptor_size,
                                      grub_efi_uint32_t *descriptor_version);
+void grub_efi_memory_fini (void);
 grub_efi_loaded_image_t *EXPORT_FUNC(grub_efi_get_loaded_image) 
(grub_efi_handle_t image_handle);
 void EXPORT_FUNC(grub_efi_print_device_path) (grub_efi_device_path_t *dp);
 char *EXPORT_FUNC(grub_efi_get_filename) (grub_efi_device_path_t *dp);
-- 
1.8.5.6