[PATCH v2 2/8] efi: libstub: avoid efi_get_memory_map() for allocating the virt map

[Ard Biesheuvel <ardb@kernel.org>]

The virt map is a set of efi_memory_desc_t descriptors that are passed
to SetVirtualAddressMap() to inform the firmware about the desired
virtual mapping of the regions marked as EFI_MEMORY_RUNTIME. The only
reason we currently call the efi_get_memory_map() helper is that it
gives us an allocation that is guaranteed to be of sufficient size.
However, efi_get_memory_map() has grown some additional complexity over
the years, and today, we're actually better off calling the EFI boot
service directly with a zero size, which tells us how much memory should
be enough for the virt map.

While at it, avoid fiddling with the VA map allocation if we will not be
using it anyway, i.e., if efi_novamap is true.

Signed-off-by: Ard Biesheuvel <ardb@kernel.org>
---
 drivers/firmware/efi/libstub/efi-stub.c | 31 ++++++++++++++++
 drivers/firmware/efi/libstub/efistub.h  |  2 +
 drivers/firmware/efi/libstub/fdt.c      | 39 ++++++++------------
 3 files changed, 49 insertions(+), 23 deletions(-)

diff --git a/drivers/firmware/efi/libstub/efi-stub.c b/drivers/firmware/efi/libstub/efi-stub.c
index 4bf751484e8b..90d44834e33e 100644
--- a/drivers/firmware/efi/libstub/efi-stub.c
+++ b/drivers/firmware/efi/libstub/efi-stub.c
@@ -321,6 +321,35 @@ efi_status_t __efiapi efi_pe_entry(efi_handle_t handle,
 	return status;
 }
 
+/*
+ * efi_allocate_virtmap() - create a pool allocation for the virtmap
+ *
+ * Create an allocation that is of sufficient size to hold all the memory
+ * descriptors that will be passed to SetVirtualAddressMap() to inform the
+ * firmware about the virtual mapping that will be used under the OS to call
+ * into the firmware.
+ */
+efi_status_t efi_alloc_virtmap(efi_memory_desc_t **virtmap,
+			       unsigned long *desc_size, u32 *desc_ver)
+{
+	unsigned long size, mmap_key;
+	efi_status_t status;
+
+	/*
+	 * Use the size of the current memory map as an upper bound for the
+	 * size of the buffer we need to pass to SetVirtualAddressMap() to
+	 * cover all EFI_MEMORY_RUNTIME regions.
+	 */
+	size = 0;
+	status = efi_bs_call(get_memory_map, &size, NULL, &mmap_key, desc_size,
+			     desc_ver);
+	if (status != EFI_BUFFER_TOO_SMALL)
+		return EFI_LOAD_ERROR;
+
+	return efi_bs_call(allocate_pool, EFI_LOADER_DATA, size,
+			   (void **)virtmap);
+}
+
 /*
  * efi_get_virtmap() - create a virtual mapping for the EFI memory map
  *
@@ -336,6 +365,8 @@ void efi_get_virtmap(efi_memory_desc_t *memory_map, unsigned long map_size,
 	efi_memory_desc_t *in, *out = runtime_map;
 	int l;
 
+	*count = 0;
+
 	for (l = 0; l < map_size; l += desc_size) {
 		u64 paddr, size;
 
diff --git a/drivers/firmware/efi/libstub/efistub.h b/drivers/firmware/efi/libstub/efistub.h
index b0ae0a454404..e9d466822b67 100644
--- a/drivers/firmware/efi/libstub/efistub.h
+++ b/drivers/firmware/efi/libstub/efistub.h
@@ -864,6 +864,8 @@ efi_status_t allocate_new_fdt_and_exit_boot(void *handle,
 
 void *get_fdt(unsigned long *fdt_size);
 
+efi_status_t efi_alloc_virtmap(efi_memory_desc_t **virtmap,
+			       unsigned long *desc_size, u32 *desc_ver);
 void efi_get_virtmap(efi_memory_desc_t *memory_map, unsigned long map_size,
 		     unsigned long desc_size, efi_memory_desc_t *runtime_map,
 		     int *count);
diff --git a/drivers/firmware/efi/libstub/fdt.c b/drivers/firmware/efi/libstub/fdt.c
index 804f542be3f2..bb141be6dfce 100644
--- a/drivers/firmware/efi/libstub/fdt.c
+++ b/drivers/firmware/efi/libstub/fdt.c
@@ -199,7 +199,7 @@ static efi_status_t update_fdt_memmap(void *fdt, struct efi_boot_memmap *map)
 
 struct exit_boot_struct {
 	efi_memory_desc_t	*runtime_map;
-	int			*runtime_entry_count;
+	int			runtime_entry_count;
 	void			*new_fdt_addr;
 };
 
@@ -212,8 +212,9 @@ static efi_status_t exit_boot_func(struct efi_boot_memmap *map,
 	 * populate @runtime_map with copies of just the EFI_MEMORY_RUNTIME
 	 * entries so that we can pass it straight to SetVirtualAddressMap()
 	 */
-	efi_get_virtmap(*map->map, *map->map_size, *map->desc_size,
-			p->runtime_map, p->runtime_entry_count);
+	if (!efi_novamap)
+		efi_get_virtmap(*map->map, *map->map_size, *map->desc_size,
+				p->runtime_map, &p->runtime_entry_count);
 
 	return update_fdt_memmap(p->new_fdt_addr, map);
 }
@@ -246,29 +247,24 @@ efi_status_t allocate_new_fdt_and_exit_boot(void *handle,
 	unsigned long map_size, desc_size, buff_size;
 	u32 desc_ver;
 	unsigned long mmap_key;
-	efi_memory_desc_t *memory_map, *runtime_map;
+	efi_memory_desc_t *memory_map;
 	efi_status_t status;
-	int runtime_entry_count;
 	struct efi_boot_memmap map;
 	struct exit_boot_struct priv;
 
-	map.map		= &runtime_map;
 	map.map_size	= &map_size;
 	map.desc_size	= &desc_size;
 	map.desc_ver	= &desc_ver;
 	map.key_ptr	= &mmap_key;
 	map.buff_size	= &buff_size;
 
-	/*
-	 * Get a copy of the current memory map that we will use to prepare
-	 * the input for SetVirtualAddressMap(). We don't have to worry about
-	 * subsequent allocations adding entries, since they could not affect
-	 * the number of EFI_MEMORY_RUNTIME regions.
-	 */
-	status = efi_get_memory_map(&map);
-	if (status != EFI_SUCCESS) {
-		efi_err("Unable to retrieve UEFI memory map.\n");
-		return status;
+	if (!efi_novamap) {
+		status = efi_alloc_virtmap(&priv.runtime_map, &desc_size,
+					   &desc_ver);
+		if (status != EFI_SUCCESS) {
+			efi_err("Unable to retrieve UEFI memory map.\n");
+			return status;
+		}
 	}
 
 	efi_info("Exiting boot services...\n");
@@ -289,10 +285,7 @@ efi_status_t allocate_new_fdt_and_exit_boot(void *handle,
 		goto fail_free_new_fdt;
 	}
 
-	runtime_entry_count		= 0;
-	priv.runtime_map		= runtime_map;
-	priv.runtime_entry_count	= &runtime_entry_count;
-	priv.new_fdt_addr		= (void *)*new_fdt_addr;
+	priv.new_fdt_addr = (void *)*new_fdt_addr;
 
 	status = efi_exit_boot_services(handle, &map, &priv, exit_boot_func);
 
@@ -304,8 +297,8 @@ efi_status_t allocate_new_fdt_and_exit_boot(void *handle,
 
 		/* Install the new virtual address map */
 		svam = efi_system_table->runtime->set_virtual_address_map;
-		status = svam(runtime_entry_count * desc_size, desc_size,
-			      desc_ver, runtime_map);
+		status = svam(priv.runtime_entry_count * desc_size, desc_size,
+			      desc_ver, priv.runtime_map);
 
 		/*
 		 * We are beyond the point of no return here, so if the call to
@@ -337,7 +330,7 @@ efi_status_t allocate_new_fdt_and_exit_boot(void *handle,
 	efi_free(MAX_FDT_SIZE, *new_fdt_addr);
 
 fail:
-	efi_system_table->boottime->free_pool(runtime_map);
+	efi_bs_call(free_pool, priv.runtime_map);
 
 	return EFI_LOAD_ERROR;
 }
-- 
2.35.1