1 files changed, 11 insertions, 7 deletions

diff --git a/drivers/firmware/efi/libstub/efi-stub-helper.c b/drivers/firmware/efi/libstub/efi-stub-helper.c
index a920fec8fe88..d073e3946383 100644
--- a/drivers/firmware/efi/libstub/efi-stub-helper.c
+++ b/drivers/firmware/efi/libstub/efi-stub-helper.c
@@ -66,25 +66,29 @@ efi_status_t efi_get_memory_map(efi_system_table_t *sys_table_arg,
         unsigned long key;
         u32 desc_version;
 
-        *map_size = sizeof(*m) * 32;
-again:
+        *map_size = 0;
+        *desc_size = 0;
+        key = 0;
+        status = efi_call_early(get_memory_map, map_size, NULL,
+                                &key, desc_size, &desc_version);
+        if (status != EFI_BUFFER_TOO_SMALL)
+                return EFI_LOAD_ERROR;
+
         /*
          * Add an additional efi_memory_desc_t because we're doing an
          * allocation which may be in a new descriptor region.
          */
-        *map_size += sizeof(*m);
+        *map_size += *desc_size;
         status = efi_call_early(allocate_pool, EFI_LOADER_DATA,
                                 *map_size, (void **)&m);
         if (status != EFI_SUCCESS)
                 goto fail;
 
-        *desc_size = 0;
-        key = 0;
         status = efi_call_early(get_memory_map, map_size, m,
                                 &key, desc_size, &desc_version);
         if (status == EFI_BUFFER_TOO_SMALL) {
                 efi_call_early(free_pool, m);
-                goto again;
+                return EFI_LOAD_ERROR;
         }
 
         if (status != EFI_SUCCESS)
@@ -101,7 +105,7 @@ fail:
 }
 
 
-unsigned long __init get_dram_base(efi_system_table_t *sys_table_arg)
+unsigned long get_dram_base(efi_system_table_t *sys_table_arg)
 {
         efi_status_t status;
         unsigned long map_size;