1 files changed, 24 insertions, 50 deletions

diff --git a/drivers/firmware/efi/libstub/mem.c b/drivers/firmware/efi/libstub/mem.c
index feef8d4be113..c92b7dbc6dfe 100644
--- a/drivers/firmware/efi/libstub/mem.c
+++ b/drivers/firmware/efi/libstub/mem.c
@@ -5,71 +5,45 @@
 
 #include "efistub.h"
 
-static inline bool mmap_has_headroom(unsigned long buff_size,
-				     unsigned long map_size,
-				     unsigned long desc_size)
-{
-	unsigned long slack = buff_size - map_size;
-
-	return slack / desc_size >= EFI_MMAP_NR_SLACK_SLOTS;
-}
-
 /**
  * efi_get_memory_map() - get memory map
- * @map:	on return pointer to memory map
+ * @map:		pointer to memory map pointer to which to assign the
+ *			newly allocated memory map
  *
  * Retrieve the UEFI memory map. The allocated memory leaves room for
  * up to EFI_MMAP_NR_SLACK_SLOTS additional memory map entries.
  *
  * Return:	status code
  */
-efi_status_t efi_get_memory_map(struct efi_boot_memmap *map)
+efi_status_t efi_get_memory_map(struct efi_boot_memmap **map)
 {
-	efi_memory_desc_t *m = NULL;
+	struct efi_boot_memmap *m, tmp;
 	efi_status_t status;
-	unsigned long key;
-	u32 desc_version;
+	unsigned long size;
 
-	*map->desc_size =	sizeof(*m);
-	*map->map_size =	*map->desc_size * 32;
-	*map->buff_size =	*map->map_size;
-again:
-	status = efi_bs_call(allocate_pool, EFI_LOADER_DATA,
-			     *map->map_size, (void **)&m);
+	tmp.map_size = 0;
+	status = efi_bs_call(get_memory_map, &tmp.map_size, NULL, &tmp.map_key,
+			     &tmp.desc_size, &tmp.desc_ver);
+	if (status != EFI_BUFFER_TOO_SMALL)
+		return EFI_LOAD_ERROR;
+
+	size = tmp.map_size + tmp.desc_size * EFI_MMAP_NR_SLACK_SLOTS;
+	status = efi_bs_call(allocate_pool, EFI_LOADER_DATA, sizeof(*m) + size,
+			     (void **)&m);
 	if (status != EFI_SUCCESS)
-		goto fail;
+		return status;
 
-	*map->desc_size = 0;
-	key = 0;
-	status = efi_bs_call(get_memory_map, map->map_size, m,
-			     &key, map->desc_size, &desc_version);
-	if (status == EFI_BUFFER_TOO_SMALL ||
-	    !mmap_has_headroom(*map->buff_size, *map->map_size,
-			       *map->desc_size)) {
-		efi_bs_call(free_pool, m);
-		/*
-		 * Make sure there is some entries of headroom so that the
-		 * buffer can be reused for a new map after allocations are
-		 * no longer permitted.  Its unlikely that the map will grow to
-		 * exceed this headroom once we are ready to trigger
-		 * ExitBootServices()
-		 */
-		*map->map_size += *map->desc_size * EFI_MMAP_NR_SLACK_SLOTS;
-		*map->buff_size = *map->map_size;
-		goto again;
-	}
+	m->buff_size = m->map_size = size;
+	status = efi_bs_call(get_memory_map, &m->map_size, m->map, &m->map_key,
+			     &m->desc_size, &m->desc_ver);
+	if (status != EFI_SUCCESS)
+		goto free_map;
 
-	if (status == EFI_SUCCESS) {
-		if (map->key_ptr)
-			*map->key_ptr = key;
-		if (map->desc_ver)
-			*map->desc_ver = desc_version;
-	} else {
-		efi_bs_call(free_pool, m);
-	}
+	*map = m;
+	return EFI_SUCCESS;
 
-fail:
-	*map->map = m;
+free_map:
+	efi_bs_call(free_pool, m);
 	return status;
 }