Conform BlExitBootServices() to UEFI Specification

sdk/xtdk/bltypes.h

@@ -85,7 +85,7 @@ typedef EFI_STATUS (*PBL_EXECIMAGE_GET_MACHINE_TYPE)(IN PVOID ImagePointer, OUT
 typedef EFI_STATUS (*PBL_EXECIMAGE_GET_SUBSYSTEM)(IN PVOID ImagePointer, OUT PUSHORT SubSystem);
 typedef EFI_STATUS (*PBL_EXECIMAGE_LOAD_IMAGE)(IN PEFI_FILE_HANDLE FileHandle, IN LOADER_MEMORY_TYPE MemoryType, IN PVOID VirtualAddress, OUT PVOID *ImagePointer);
 typedef EFI_STATUS (*PBL_EXECIMAGE_RELOCATE_IMAGE)(IN PVOID ImagePointer, IN EFI_VIRTUAL_ADDRESS Address);
-typedef EFI_STATUS (*PBL_EXIT_BOOT_SERVICES)(IN UINT_PTR MapKey);
+typedef EFI_STATUS (*PBL_EXIT_BOOT_SERVICES)();
 typedef EFI_STATUS (*PBL_FIND_BOOT_PROTOCOL)(IN PWCHAR SystemType, OUT PEFI_GUID BootProtocolGuid);
 typedef EFI_STATUS (*PBL_FREE_PAGES)(IN UINT64 Size, IN EFI_PHYSICAL_ADDRESS Memory);
 typedef EFI_STATUS (*PBL_FREE_POOL)(IN PVOID Memory);

xtldr/efiutils.c

@@ -12,32 +12,54 @@
 /**
  * Exits EFI boot services.
  *
- * @param MapKey
- *        Identifies the current memory map of the system.
- *
  * @return This routine returns status code.
  *
  * @since XT 1.0
  */
 XTCDECL
 EFI_STATUS
-BlExitBootServices(IN UINT_PTR MapKey)
+BlExitBootServices()
 {
+    PEFI_MEMORY_MAP MemoryMap;
     EFI_STATUS Status;
+    ULONG Counter;
 
-    /* Attempt to exit boot services */
+    /* Boot Services might be partially shutdown, so mark them as unavailable */
-    Status = EfiSystemTable->BootServices->ExitBootServices(EfiImageHandle, MapKey);
+    BlpStatus.BootServices = FALSE;
+
+    /* Allocate buffer for EFI memory map */
+    Status = BlMemoryAllocatePool(sizeof(EFI_MEMORY_MAP), (PVOID*)&MemoryMap);
     if(Status != STATUS_EFI_SUCCESS)
     {
-        /* Retry as UEFI spec says to do it twice */
+        /* Memory allocation failure */
-        Status = EfiSystemTable->BootServices->ExitBootServices(EfiImageHandle, MapKey);
+        BlDebugPrint(L"ERROR: Memory allocation failure (Status Code: 0x%lx)\n", Status);
+        return Status;
     }
 
-    /* Make sure boot services were successfully exited */
+    /* Zero fill the buffer and initialize counter */
-    if(Status == STATUS_EFI_SUCCESS)
+    RtlZeroMemory(MemoryMap, sizeof(EFI_MEMORY_MAP));
+    Counter = 0xFF;
+
+    /* Attempt to exit boot services */
+    while(Counter > 0)
     {
-        /* Mark EFI Boot Services as no longer available */
+        /* Get memory map each time as it can change between two calls */
-        BlpStatus.BootServices = FALSE;
+        Status = BlGetMemoryMap(MemoryMap);
+        if(Status != STATUS_EFI_SUCCESS)
+        {
+            /* Failed to get new memory map */
+            return Status;
+        }
+
+        /* Exit boot services */
+        Status = EfiSystemTable->BootServices->ExitBootServices(EfiImageHandle, MemoryMap->MapKey);
+        if(Status == STATUS_EFI_SUCCESS)
+        {
+            break;
+        }
+
+        /* Decrement counter */
+        Counter--;
     }
 
     /* Return EFI status code */

xtldr/includes/xtldr.h

@@ -85,7 +85,7 @@ BlEnumerateBlockDevices();
 
 XTCDECL
 EFI_STATUS
-BlExitBootServices(IN UINT_PTR MapKey);
+BlExitBootServices();
 
 XTCDECL
 EFI_STATUS

xtldr/modules/xtos_o/amd64/memory.c

@@ -37,7 +37,6 @@ XtEnablePaging(IN PLIST_ENTRY MemoryMappings,
 {
     PLOADER_MEMORY_MAPPING Mapping;
     EFI_PHYSICAL_ADDRESS Address;
-    PEFI_MEMORY_MAP MemoryMap;
     PLIST_ENTRY ListEntry, ModulesList, ModulesListEntry;
     PXTBL_MODULE_INFO ModuleInfo;
     EFI_STATUS Status;
@@ -120,39 +119,12 @@ XtEnablePaging(IN PLIST_ENTRY MemoryMappings,
     /* Map zero page as well */
     XtMapVirtualMemory(MemoryMappings, 0, 0, 1, PtePointer);
 
-    /* Allocate and zero-fill buffer for EFI memory map */
-    XtLdrProtocol->Memory.AllocatePool(sizeof(EFI_MEMORY_MAP), (PVOID*)&MemoryMap);
-    RtlZeroMemory(MemoryMap, sizeof(EFI_MEMORY_MAP));
-
-    /* Get EFI memory map and prepare for exiting boot services */
-    XtLdrProtocol->Debug.Print(L"Exiting EFI boot services\n");
-    Status = XtLdrProtocol->Memory.GetMemoryMap(MemoryMap);
-    if(Status != STATUS_EFI_SUCCESS)
-    {
-        /* Unable to get memory map */
-        return Status;
-    }
-
     /* Exit EFI Boot Services */
-    Status = XtLdrProtocol->Util.ExitBootServices(MemoryMap->MapKey);
+    XtLdrProtocol->Debug.Print(L"Exiting EFI boot services\n");
-
+    Status = XtLdrProtocol->Util.ExitBootServices();
-    /* 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)
-    {
         XtLdrProtocol->Console.Print(L"Failed to exit boot services (Status code: %lx)\n", Status);
         return STATUS_EFI_ABORTED;
     }

xtldr/modules/xtos_o/i686/memory.c

@@ -218,36 +218,9 @@ XtEnablePaging(IN PLIST_ENTRY MemoryMappings,
     /* Map zero page as well */
     XtMapVirtualMemory(MemoryMappings, 0, 0, 1, PtePointer);
 
-    /* Zero-fill buffer for EFI memory map */
-    RtlZeroMemory(MemoryMap, sizeof(EFI_MEMORY_MAP));
-
-    /* Get EFI memory map and prepare for exiting boot services */
-    XtLdrProtocol->Debug.Print(L"Exiting EFI boot services\n");
-    Status = XtLdrProtocol->Memory.GetMemoryMap(MemoryMap);
-    if(Status != STATUS_EFI_SUCCESS)
-    {
-        /* Unable to get memory map */
-        return Status;
-    }
-
     /* Exit EFI Boot Services */
-    Status = XtLdrProtocol->Util.ExitBootServices(MemoryMap->MapKey);
+    XtLdrProtocol->Debug.Print(L"Exiting EFI boot services\n");
-
+    Status = XtLdrProtocol->Util.ExitBootServices();
-    /* 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)
     {
         /* Failed to exit boot services */