From ec63fe3fa6c1c09dab1ece358ed0ed28a7df57c7 Mon Sep 17 00:00:00 2001
From: Rafal Kupiec <belliash@codingworkshop.eu.org>
Date: Sat, 2 Dec 2023 16:02:24 +0100
Subject: [PATCH] Retry exiting boot services if first try fails

---
 xtldr/amd64/memory.c | 8 ++++++++
 xtldr/i686/memory.c  | 8 ++++++++
 2 files changed, 16 insertions(+)

diff --git a/xtldr/amd64/memory.c b/xtldr/amd64/memory.c
index 64acabf1..ee297236 100644
--- a/xtldr/amd64/memory.c
+++ b/xtldr/amd64/memory.c
@@ -119,6 +119,14 @@ BlEnablePaging(IN PLIST_ENTRY MemoryMappings,
     Status = EfiSystemTable->BootServices->ExitBootServices(EfiImageHandle, MemoryMap->MapKey);
     if(Status != STATUS_EFI_SUCCESS)
     {
+        /* Retry as UEFI spec says to do it twice */
+        Status = EfiSystemTable->BootServices->ExitBootServices(EfiImageHandle, MemoryMap->MapKey);
+    }
+
+    /* Check if exitted boot services successfully */
+    if(Status != STATUS_EFI_SUCCESS)
+    {
+        /* Failed to exit boot services */
         BlDbgPrint(L"Failed to exit boot services (Status code: %lx)\n", Status);
         return STATUS_EFI_ABORTED;
     }
diff --git a/xtldr/i686/memory.c b/xtldr/i686/memory.c
index 2ba4c3bd..ae99f4cb 100644
--- a/xtldr/i686/memory.c
+++ b/xtldr/i686/memory.c
@@ -216,6 +216,14 @@ BlEnablePaging(IN PLIST_ENTRY MemoryMappings,
     Status = EfiSystemTable->BootServices->ExitBootServices(EfiImageHandle, MemoryMap->MapKey);
     if(Status != STATUS_EFI_SUCCESS)
     {
+        /* Retry as UEFI spec says to do it twice */
+        Status = EfiSystemTable->BootServices->ExitBootServices(EfiImageHandle, MemoryMap->MapKey);
+    }
+
+    /* Check if exitted boot services successfully */
+    if(Status != STATUS_EFI_SUCCESS)
+    {
+        /* Failed to exit boot services */
         BlDbgPrint(L"Failed to exit boot services (Status code: %lx)\n", Status);
         return STATUS_EFI_ABORTED;
     }