efi: Move exit_boot_services into a function

At present this code is inline in the app and stub. But they do the same
thing. The difference is that the stub does it immediately and the app
doesn't want to do it until the end (when it boots a kernel) or not at
all, if returning to UEFI.

Move it into a function so it can be called as needed.

Add a comment showing how to store the memory map so that it can be
accessed within the app if needed, for debugging purposes only. The map
can change without notice.

Signed-off-by: Simon Glass <sjg@chromium.org>

4 files changed, 124 insertions, 55 deletions

diff --git a/include/efi.h b/include/efi.h
index 877a2e5a8d..dc9907fa16 100644
--- a/include/efi.h
+++ b/include/efi.h
@@ -407,6 +407,12 @@ static inline struct efi_mem_desc *efi_get_next_mem_desc(
  * @sys_table: Pointer to system table
  * @boot: Pointer to boot-services table
  * @run: Pointer to runtime-services table
+ * @memmap_key: Key returned from get_memory_map()
+ * @memmap_desc: List of memory-map records
+ * @memmap_alloc: Amount of memory allocated for memory map list
+ * @memmap_size Size of memory-map list in bytes
+ * @memmap_desc_size: Size of an individual memory-map record, in bytes
+ * @memmap_version: Memory-map version
  *
  * @use_pool_for_malloc: true if all allocation should go through the EFI 'pool'
  *	methods allocate_pool() and free_pool(); false to use 'pages' methods
@@ -424,6 +430,12 @@ struct efi_priv {
 	struct efi_system_table *sys_table;
 	struct efi_boot_services *boot;
 	struct efi_runtime_services *run;
+	efi_uintn_t memmap_key;
+	struct efi_mem_desc *memmap_desc;
+	efi_uintn_t memmap_alloc;
+	efi_uintn_t memmap_size;
+	efi_uintn_t memmap_desc_size;
+	u32 memmap_version;
 
 	/* app: */
 	bool use_pool_for_malloc;
@@ -578,4 +590,24 @@ void efi_putc(struct efi_priv *priv, const char ch);
  */
 int efi_info_get(enum efi_entry_t type, void **datap, int *sizep);
 
+/**
+ * efi_store_memory_map() - Collect the memory-map info from EFI
+ *
+ * Collect the memory info and store it for later use, e.g. in calling
+ * exit_boot_services()
+ *
+ * @priv:	Pointer to private EFI structure
+ * @return 0 if OK, non-zero on error
+ */
+int efi_store_memory_map(struct efi_priv *priv);
+
+/**
+ * efi_call_exit_boot_services() - Handle the exit-boot-service procedure
+ *
+ * Tell EFI we don't want their boot services anymore
+ *
+ * Return: 0 if OK, non-zero on error
+ */
+int efi_call_exit_boot_services(void);
+
 #endif /* _LINUX_EFI_H */
diff --git a/lib/efi/efi.c b/lib/efi/efi.c
index cd6bf47b18..aa42f1842f 100644
--- a/lib/efi/efi.c
+++ b/lib/efi/efi.c
@@ -135,3 +135,75 @@ void efi_free(struct efi_priv *priv, void *ptr)
 
 	boot->free_pool(ptr);
 }
+
+int efi_store_memory_map(struct efi_priv *priv)
+{
+	struct efi_boot_services *boot = priv->sys_table->boottime;
+	efi_uintn_t size, desc_size;
+	efi_status_t ret;
+
+	/* Get the memory map so we can switch off EFI */
+	size = 0;
+	ret = boot->get_memory_map(&size, NULL, &priv->memmap_key,
+				   &priv->memmap_desc_size,
+				   &priv->memmap_version);
+	if (ret != EFI_BUFFER_TOO_SMALL) {
+		/*
+		 * Note this function avoids using printf() since it is not
+		 * available in the stub
+		 */
+		printhex2(EFI_BITS_PER_LONG);
+		putc(' ');
+		printhex2(ret);
+		puts(" No memory map\n");
+		return ret;
+	}
+	/*
+	 * Since doing a malloc() may change the memory map and also we want to
+	 * be able to read the memory map in efi_call_exit_boot_services()
+	 * below, after more changes have happened
+	 */
+	priv->memmap_alloc = size + 1024;
+	priv->memmap_size = priv->memmap_alloc;
+	priv->memmap_desc = efi_malloc(priv, size, &ret);
+	if (!priv->memmap_desc) {
+		printhex2(ret);
+		puts(" No memory for memory descriptor\n");
+		return ret;
+	}
+
+	ret = boot->get_memory_map(&priv->memmap_size, priv->memmap_desc,
+				   &priv->memmap_key, &desc_size,
+				   &priv->memmap_version);
+	if (ret) {
+		printhex2(ret);
+		puts(" Can't get memory map\n");
+		return ret;
+	}
+
+	return 0;
+}
+
+int efi_call_exit_boot_services(void)
+{
+	struct efi_priv *priv = efi_get_priv();
+	const struct efi_boot_services *boot = priv->boot;
+	efi_uintn_t size;
+	u32 version;
+	efi_status_t ret;
+
+	size = priv->memmap_alloc;
+	ret = boot->get_memory_map(&size, priv->memmap_desc,
+				   &priv->memmap_key,
+				   &priv->memmap_desc_size, &version);
+	if (ret) {
+		printhex2(ret);
+		puts(" Can't get memory map\n");
+		return ret;
+	}
+	ret = boot->exit_boot_services(priv->parent_image, priv->memmap_key);
+	if (ret)
+		return ret;
+
+	return 0;
+}
diff --git a/lib/efi/efi_app.c b/lib/efi/efi_app.c
index d60f2f6c28..5c2593bc4d 100644
--- a/lib/efi/efi_app.c
+++ b/lib/efi/efi_app.c
@@ -321,6 +321,15 @@ efi_status_t EFIAPI efi_main(efi_handle_t image,
 		return ret;
 	}
 
+	/*
+	 * We could store the EFI memory map here, but it changes all the time,
+	 * so this is only useful for debugging.
+	 *
+	 * ret = efi_store_memory_map(priv);
+	 * if (ret)
+	 *	return ret;
+	 */
+
 	printf("starting\n");
 
 	board_init_f(GD_FLG_SKIP_RELOC);
diff --git a/lib/efi/efi_stub.c b/lib/efi/efi_stub.c
index c89ae7c907..646cde3214 100644
--- a/lib/efi/efi_stub.c
+++ b/lib/efi/efi_stub.c
@@ -304,15 +304,12 @@ efi_status_t EFIAPI efi_main(efi_handle_t image,
 {
 	struct efi_priv local_priv, *priv = &local_priv;
 	struct efi_boot_services *boot = sys_table->boottime;
-	struct efi_mem_desc *desc;
 	struct efi_entry_memmap map;
 	struct efi_gop *gop;
 	struct efi_entry_gopmode mode;
 	struct efi_entry_systable table;
 	efi_guid_t efi_gop_guid = EFI_GRAPHICS_OUTPUT_PROTOCOL_GUID;
-	efi_uintn_t key, desc_size, size;
 	efi_status_t ret;
-	u32 version;
 	int cs32;
 
 	ret = efi_init(priv, "Payload", image, sys_table);
@@ -327,24 +324,11 @@ efi_status_t EFIAPI efi_main(efi_handle_t image,
 	if (cs32 < 0)
 		return EFI_UNSUPPORTED;
 
-	/* Get the memory map so we can switch off EFI */
-	size = 0;
-	ret = boot->get_memory_map(&size, NULL, &key, &desc_size, &version);
-	if (ret != EFI_BUFFER_TOO_SMALL) {
-		printhex2(EFI_BITS_PER_LONG);
-		putc(' ');
-		printhex2(ret);
-		puts(" No memory map\n");
-		return ret;
-	}
-	size += 1024;	/* Since doing a malloc() may change the memory map! */
-	desc = efi_malloc(priv, size, &ret);
-	if (!desc) {
-		printhex2(ret);
-		puts(" No memory for memory descriptor\n");
+	ret = efi_store_memory_map(priv);
+	if (ret)
 		return ret;
-	}
-	ret = setup_info_table(priv, size + 128);
+
+	ret = setup_info_table(priv, priv->memmap_size + 128);
 	if (ret)
 		return ret;
 
@@ -360,48 +344,20 @@ efi_status_t EFIAPI efi_main(efi_handle_t image,
 			       sizeof(struct efi_gop_mode_info));
 	}
 
-	ret = boot->get_memory_map(&size, desc, &key, &desc_size, &version);
-	if (ret) {
-		printhex2(ret);
-		puts(" Can't get memory map\n");
-		return ret;
-	}
-
 	table.sys_table = (ulong)sys_table;
 	add_entry_addr(priv, EFIET_SYS_TABLE, &table, sizeof(table), NULL, 0);
 
-	ret = boot->exit_boot_services(image, key);
-	if (ret) {
-		/*
-		 * Unfortunately it happens that we cannot exit boot services
-		 * the first time. But the second time it work. I don't know
-		 * why but this seems to be a repeatable problem. To get
-		 * around it, just try again.
-		 */
-		printhex2(ret);
-		puts(" Can't exit boot services\n");
-		size = sizeof(desc);
-		ret = boot->get_memory_map(&size, desc, &key, &desc_size,
-					   &version);
-		if (ret) {
-			printhex2(ret);
-			puts(" Can't get memory map\n");
-			return ret;
-		}
-		ret = boot->exit_boot_services(image, key);
-		if (ret) {
-			printhex2(ret);
-			puts(" Can't exit boot services 2\n");
-			return ret;
-		}
-	}
+	ret = efi_call_exit_boot_services();
+	if (ret)
+		return ret;
 
 	/* The EFI UART won't work now, switch to a debug one */
 	use_uart = true;
 
-	map.version = version;
-	map.desc_size = desc_size;
-	add_entry_addr(priv, EFIET_MEMORY_MAP, &map, sizeof(map), desc, size);
+	map.version = priv->memmap_version;
+	map.desc_size = priv->memmap_desc_size;
+	add_entry_addr(priv, EFIET_MEMORY_MAP, &map, sizeof(map),
+		       priv->memmap_desc, priv->memmap_size);
 	add_entry_addr(priv, EFIET_END, NULL, 0, 0, 0);
 
 	memcpy((void *)CONFIG_SYS_TEXT_BASE, _binary_u_boot_bin_start,