From 60ce943cab2f56a09a6cbaacbf755891fb7f2c72 Mon Sep 17 00:00:00 2001
From: Rafal Kupiec <belliash@codingworkshop.eu.org>
Date: Sat, 6 Jan 2024 23:23:52 +0100
Subject: [PATCH] Fix memory mappings and exiting boot services for i686

---
 xtldr/modules/xtos_o/i686/memory.c | 21 +++++++++++++++------
 1 file changed, 15 insertions(+), 6 deletions(-)

diff --git a/xtldr/modules/xtos_o/i686/memory.c b/xtldr/modules/xtos_o/i686/memory.c
index 141fe42..3ccf644 100644
--- a/xtldr/modules/xtos_o/i686/memory.c
+++ b/xtldr/modules/xtos_o/i686/memory.c
@@ -175,13 +175,8 @@ XtEnablePaging(IN PLIST_ENTRY MemoryMappings,
     }
 
     /* Map XTLDR code */
-    Status = XtAddVirtualMemoryMapping(MemoryMappings, ImageProtocol->ImageBase, ImageProtocol->ImageBase,
+    XtAddVirtualMemoryMapping(MemoryMappings, ImageProtocol->ImageBase, ImageProtocol->ImageBase,
                                        EFI_SIZE_TO_PAGES(ImageProtocol->ImageSize), LoaderFirmwareTemporary);
-    if(Status != STATUS_EFI_SUCCESS)
-    {
-        /* Mapping the boot loader code failed */
-        return Status;
-    }
 
     /* Add page mapping itself to memory mapping */
     Status = XtAddVirtualMemoryMapping(MemoryMappings, NULL, *PtePointer, 1, LoaderMemoryData);
@@ -238,6 +233,20 @@ XtEnablePaging(IN PLIST_ENTRY MemoryMappings,
     /* Exit EFI Boot Services */
     Status = XtLdrProtocol->Util.ExitBootServices(MemoryMap->MapKey);
 
+    /* Check if exitted boot services successfully */
+    if(Status != STATUS_EFI_SUCCESS)
+    {
+        /* Failed to exit boot services */
+        Status = XtLdrProtocol->Memory.GetMemoryMap(MemoryMap);
+        if(Status != STATUS_EFI_SUCCESS)
+        {
+            /* Unable to get memory map */
+            return Status;
+        }
+
+        Status = XtLdrProtocol->Util.ExitBootServices(MemoryMap->MapKey);
+    }
+
     /* Check if exitted boot services successfully */
     if(Status != STATUS_EFI_SUCCESS)
     {