[syslinux] [PATCH] efi/main: add retry to exit_boot()

[Tom Huibregtse]

Sometimes the UEFI memory map changes between GetMemoryMap() and
ExitBootServices(), making the MapKey value incorrect. Per the UEFI
spec, the memory map should then be fetched again and exiting retried.

Signed-off-by: Tom Huibregtse <thuibregtse at xes-inc.com>
---
 efi/main.c | 81 ++++++++++++++++++++++++++++++++++++++++++++++--------
 1 file changed, 70 insertions(+), 11 deletions(-)

diff --git a/efi/main.c b/efi/main.c
index 6a748412..d7fb637e 100644
--- a/efi/main.c
+++ b/efi/main.c
@@ -1008,16 +1008,81 @@ static int exit_boot(struct boot_params *bp)
 	EFI_MEMORY_DESCRIPTOR *map;
 	EFI_STATUS status;
 	uint32_t e820_type;
-	UINTN i, nr_entries, key, desc_sz;
+	UINTN i, nr_entries, key, desc_sz, buf_sz, map_sz;
 	UINT32 desc_ver;
 
+	/*
+	 * Call once with empty buffer size to
+	 * see how large the buffer should be.
+	 */
+	buf_sz = 0;
+	status = uefi_call_wrapper(BS->GetMemoryMap,
+				   5,
+				   &buf_sz,
+				   map,
+				   &key,
+				   &desc_sz,
+				   &desc_ver);
+	if (!buf_sz)
+		return -1;
+
+	/*
+	 * Allocate a pool double the size of the map in
+	 * case if the first call to ExitBootServices()
+	 * fails and the memory map grows before
+	 * GetMemoryMap() is called again.
+	 */
+	buf_sz *= 2;
+	map = AllocatePool(buf_sz);
+
 	/* Build efi memory map */
-	map = get_memory_map(&nr_entries, &key, &desc_sz, &desc_ver);
-	if (!map)
+	map_sz = buf_sz;
+	status = uefi_call_wrapper(BS->GetMemoryMap,
+				   5,
+				   &map_sz,
+				   map,
+				   &key,
+				   &desc_sz,
+				   &desc_ver);
+	if (!map) {
+		FreePool(map);
+		return -1;
+	}
+
+	status = uefi_call_wrapper(BS->ExitBootServices, 2, image_handle, key);
+
+	if (status == EFI_INVALID_PARAMETER) {
+		/*
+		 * Retry GetMemoryMap() and ExitBootServices() if the
+		 * first try fails  per UEFI Spec v2.6, Section 6.4:
+		 * EFI_BOOT_SERVICES.ExitBootServices.
+		 */
+		map_sz = buf_sz;
+		status = uefi_call_wrapper(BS->GetMemoryMap,
+					   5,
+					   &map_sz,
+					   map,
+					   &key,
+					   &desc_sz,
+					   &desc_ver);
+
+		/* FreePool() cannot be called after ExitBootServices() */
+		if (status != EFI_SUCCESS)
+			return -1;
+
+		status = uefi_call_wrapper(BS->ExitBootServices,
+					   2,
+					   image_handle,
+					   key);
+	}
+
+	if (status != EFI_SUCCESS) {
+		printf("Failed to exit boot services: 0x%016lx\n", status);
 		return -1;
+	}
 
 	bp->efi.memmap = (uint32_t)(unsigned long)map;
-	bp->efi.memmap_size = nr_entries * desc_sz;
+	bp->efi.memmap_size = map_sz;
 	bp->efi.systab = (uint32_t)(unsigned long)ST;
 	bp->efi.desc_size = desc_sz;
 	bp->efi.desc_version = desc_ver;
@@ -1033,6 +1098,7 @@ static int exit_boot(struct boot_params *bp)
 	 * e820 map because it will most likely have changed in the
 	 * interim.
 	 */
+	nr_entries = map_sz / desc_sz;
 	e = e820buf = bp->e820_map;
 	for (i = 0; i < nr_entries && i < E820MAX; i++) {
 		struct e820_entry *prev = NULL;
@@ -1089,13 +1155,6 @@ static int exit_boot(struct boot_params *bp)
 
 	bp->e820_entries = e - e820buf;
 
-	status = uefi_call_wrapper(BS->ExitBootServices, 2, image_handle, key);
-	if (status != EFI_SUCCESS) {
-		printf("Failed to exit boot services: 0x%016lx\n", status);
-		FreePool(map);
-		return -1;
-	}
-
 	return 0;
 }
 
-- 
2.27.0