1 files changed, 44 insertions, 49 deletions

diff --git a/drivers/firmware/efi/libstub/mem.c b/drivers/firmware/efi/libstub/mem.c
index feef8d4be113..45841ef55a9f 100644
--- a/drivers/firmware/efi/libstub/mem.c
+++ b/drivers/firmware/efi/libstub/mem.c
@@ -5,71 +5,66 @@
 
 #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
+ * @install_cfg_tbl:	whether or not to install the boot memory map as a
+ *			configuration table
  *
  * 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,
+				bool install_cfg_tbl)
 {
-	efi_memory_desc_t *m = NULL;
+	int memtype = install_cfg_tbl ? EFI_ACPI_RECLAIM_MEMORY
+				      : EFI_LOADER_DATA;
+	efi_guid_t tbl_guid = LINUX_EFI_BOOT_MEMMAP_GUID;
+	struct efi_boot_memmap *m, tmp;
 	efi_status_t status;
-	unsigned long key;
-	u32 desc_version;
-
-	*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);
+	unsigned long size;
+
+	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, memtype, sizeof(*m) + size,
+			     (void **)&m);
 	if (status != EFI_SUCCESS)
-		goto fail;
-
-	*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);
+		return status;
+
+	if (install_cfg_tbl) {
 		/*
-		 * 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()
+		 * Installing a configuration table might allocate memory, and
+		 * this may modify the memory map. This means we should install
+		 * the configuration table first, and re-install or delete it
+		 * as needed.
 		 */
-		*map->map_size += *map->desc_size * EFI_MMAP_NR_SLACK_SLOTS;
-		*map->buff_size = *map->map_size;
-		goto again;
+		status = efi_bs_call(install_configuration_table, &tbl_guid, m);
+		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);
-	}
+	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 uninstall_table;
+
+	*map = m;
+	return EFI_SUCCESS;
 
-fail:
-	*map->map = m;
+uninstall_table:
+	if (install_cfg_tbl)
+		efi_bs_call(install_configuration_table, &tbl_guid, NULL);
+free_map:
+	efi_bs_call(free_pool, m);
 	return status;
 }