[BOOTMGR]: Fix bug in test.
[reactos.git] / reactos / boot / environ / app / bootmgr / bootmgr.c
index ae31bb0..8b1ef5d 100644 (file)
@@ -1,7 +1,7 @@
 /*
  * COPYRIGHT:       See COPYING.ARM in the top level directory
  * PROJECT:         ReactOS UEFI Boot Manager
- * FILE:            boot/environ/app/bootmgr/bootmgr.c
+ * FILE:            boot/environ/app/bootmgr/bootmgr.cla
  * PURPOSE:         Boot Manager Entrypoint
  * PROGRAMMER:      Alex Ionescu (alex.ionescu@reactos.org)
  */
@@ -274,7 +274,7 @@ BmpFwGetApplicationDirectoryPath (
                                   AppPathLength - sizeof(UNICODE_NULL));
                     PathCopy[AppPathLength] = UNICODE_NULL;
 
-                    /* Finally, initialize the outoing string */
+                    /* Finally, initialize the outgoing string */
                     RtlInitUnicodeString(ApplicationDirectoryPath, PathCopy);
                 }
                 else
@@ -1083,36 +1083,12 @@ BlXmiInitialize (
     return STATUS_SUCCESS;
 }
 
-VOID
-BlImgQueryCodeIntegrityBootOptions (
-    _In_ PBL_LOADED_APPLICATION_ENTRY ApplicationEntry,
-    _Out_ PBOOLEAN IntegrityChecksDisabled, 
-    _Out_ PBOOLEAN TestSigningEnabled
-    )
-{
-    
-    NTSTATUS Status;
-    BOOLEAN Value;
-
-    /* Check if /DISABLEINTEGRITYCHECKS is on */
-    Status = BlGetBootOptionBoolean(ApplicationEntry->BcdData,
-                                    BcdLibraryBoolean_DisableIntegrityChecks,
-                                    &Value);
-    *IntegrityChecksDisabled = NT_SUCCESS(Status) && (Value);
-
-    /* Check if /TESTSIGNING is on */
-    Status = BlGetBootOptionBoolean(ApplicationEntry->BcdData,
-                                    BcdLibraryBoolean_AllowPrereleaseSignatures,
-                                    &Value);
-    *TestSigningEnabled = NT_SUCCESS(Status) && (Value);
-}
-
 NTSTATUS
 BmFwVerifySelfIntegrity (
     VOID
     )
 {
-    /* Check if we're booted by UEFI off the DVD directlry */
+    /* Check if we're booted by UEFI off the DVD directly */
     if ((BlpBootDevice->DeviceType == LocalDevice) &&
         (BlpBootDevice->Local.Type == CdRomDevice) &&
         (BlpApplicationFlags & BL_APPLICATION_FLAG_CONVERTED_FROM_EFI))
@@ -1427,7 +1403,7 @@ BmpPopulateBootEntryList (
 
         /* Check if this is a real mode startup.com */
         if ((ObjectType.Application.ObjectCode == BCD_OBJECT_TYPE_APPLICATION) &&
-            (ObjectType.Application.ImageCode = BCD_IMAGE_TYPE_REAL_MODE) &&
+            (ObjectType.Application.ImageCode == BCD_IMAGE_TYPE_REAL_MODE) &&
             (ObjectType.Application.ApplicationCode == BCD_APPLICATION_TYPE_STARTUPCOM))
         {
             /* Check if PXE soft reboot will occur */
@@ -1494,7 +1470,7 @@ BmGetBootSequence (
     )
 {
     PBL_LOADED_APPLICATION_ENTRY* Sequence;
-    ULONG Count;
+    ULONG Count = SequenceListCount;
     NTSTATUS Status;
 
     /* Allocate the sequence list */
@@ -1753,7 +1729,7 @@ BmGetBootMenuPolicy (
             return MenuPolicyLegacy;
         }
 
-        /* Use the correct opetion ID */
+        /* Use the correct option ID */
         OptionId = BcdResumeInteger_BootMenuPolicy;
     }
 
@@ -1807,7 +1783,7 @@ BmpProcessBootEntry (
     /* If the legacy menu must be shown, or if we have a boot entry */
     if ((BmGetBootMenuPolicy(BootEntry) != MenuPolicyStandard) || (BootEntry))
     {
-        /* Check if any key has been presseed */
+        /* Check if any key has been pressed */
         BmDisplayGetBootMenuStatus(&MenuStatus);
         if (MenuStatus.AnyKey)
         {
@@ -2029,17 +2005,709 @@ Quickie:
     return Status;
 }
 
+NTSTATUS 
+BmLaunchRecoverySequence (
+    _In_ PBL_LOADED_APPLICATION_ENTRY BootEntry,
+    _In_ ULONG LaunchCode
+    )
+{
+    NTSTATUS Status;
+    PBL_LOADED_APPLICATION_ENTRY RecoveryEntry;
+    HANDLE BcdHandle;
+    PGUID RecoverySequence;
+    ULONG Count, i, RecoveryIndex, SequenceCount;
+    PBL_LOADED_APPLICATION_ENTRY* Sequence;
+
+    /* Initialize locals */
+    RecoveryIndex = 0;
+    Sequence = NULL;
+    RecoverySequence = NULL;
+    Count = 0;
+    BcdHandle = NULL;
+
+    /* Open the BCD*/
+    Status = BmOpenDataStore(&BcdHandle);
+    if (!NT_SUCCESS(Status))
+    {
+        goto Quickie;
+    }
+
+    /* Get the recovery sequence list */
+    Status = BlGetBootOptionGuidList(BootEntry->BcdData,
+                                     BcdLibraryObjectList_RecoverySequence,
+                                     &RecoverySequence,
+                                     &SequenceCount);
+    if (!NT_SUCCESS(Status))
+    {
+        goto Quickie;
+    }
+
+    /* Get the sequence of boot entries out of it */
+    Status = BmGetBootSequence(BcdHandle,
+                               RecoverySequence,
+                               SequenceCount,
+                               BL_APPLICATION_ENTRY_RECOVERY,
+                               &Sequence,
+                               &Count);
+    if (!NT_SUCCESS(Status))
+    {
+        goto Quickie;
+    }
+
+    /* Was the BCD open? */
+    if (BcdHandle)
+    {
+        /* Close it */
+        BmCloseDataStore(BcdHandle);
+    }
+
+    /* Now go over every entry in the sequence */
+    for (i = 0; i < Count; ++i)
+    {
+        /* Check the code for this recovery launch */
+        if (LaunchCode == 2 || LaunchCode == 5)
+        {
+            /* Remove the override if there is one, and set it to 4 */
+            BlRemoveBootOption(Sequence[i]->BcdData, BcdLibraryInteger_DisplayMessageOverride);
+            BlAppendBootOptionInteger(Sequence[i],
+                                      BcdLibraryInteger_DisplayMessageOverride,
+                                      4);
+        }
+        else if (LaunchCode == 3)
+        {
+            /* Remove the override if there is one, and set it to 10 */
+            BlRemoveBootOption(Sequence[i]->BcdData, BcdLibraryInteger_DisplayMessageOverride);
+            BlAppendBootOptionInteger(Sequence[i],
+                                        BcdLibraryInteger_DisplayMessageOverride,
+                                        10);
+        }
+
+        /* Launch the boot entry for this part of the recovery sequence */
+        Status = BmpLaunchBootEntry(Sequence[i], NULL, LaunchCode, FALSE);
+        if (!NT_SUCCESS(Status))
+        {
+            break;
+        }
+    }
+
+Quickie:
+    /* Did we have a sequence of entries? */
+    if (Sequence)
+    {   
+        /* Loop through each one */
+        for (RecoveryIndex = 0; RecoveryIndex < Count; RecoveryIndex++)
+        {
+            /* Does this index have an allocated boot entry? */
+            RecoveryEntry = Sequence[RecoveryIndex];
+            if (RecoveryEntry)
+            {
+                /* Destroy it */
+                BlDestroyBootEntry(RecoveryEntry);
+            }
+        }
+
+        /* Free the sequence itself */
+        BlMmFreeHeap(Sequence);
+    }
+
+    /* Was there a sequence list? */
+    if (RecoverySequence)
+    {
+        /* Free it */
+        BlMmFreeHeap(RecoverySequence);
+    }
+  
+    /* Return back to caller */
+    return Status;
+}
+
+ULONG
+BmDisplayDumpError (
+    _In_ PBL_LOADED_APPLICATION_ENTRY BootEntry,
+    _In_ ULONG LaunchCode
+    )
+{
+    ULONG BootError;
+    NTSTATUS Status;
+    BOOLEAN Restart, NoError;
+
+    /* Assume we'll just reboot */
+    BootError = Reboot;
+
+    /* Should we reboot? */
+    Status = BlGetBootOptionBoolean(BlpApplicationEntry.BcdData,
+                                    BcdLibraryBoolean_RestartOnFailure,
+                                    &Restart);
+    if ((NT_SUCCESS(Status)) && (Restart))
+    {
+        return BootError;
+    }
+
+    /* Should we not show errors, and thus, reboot? */
+    Status = BlGetBootOptionBoolean(BlpApplicationEntry.BcdData,
+                                    BcdBootMgrBoolean_NoErrorDisplay,
+                                    &NoError);
+    if ((NT_SUCCESS(Status)) && (NoError))
+    {
+        return BootError;
+    }
+
+    /* Is there an internal boot error? */
+    if (BmpInternalBootError)
+    {
+        /* Return it -- but it's a pointer? */
+        return (ULONG)BmpInternalBootError; // ???
+    }
+
+    /* Otherwise, show the menu to see what to do */
+    EfiPrintf(L"Error menu not yet implemented\r\n");
+    return BootError;
+}
+
+NTSTATUS
+BmpCreateDevices (
+    _In_ PBL_LOADED_APPLICATION_ENTRY BootEntry
+    )
+{
+    ULONG NextOffset, DataOffset, ListOffset;
+    PBL_BCD_OPTION Option, ListOption;
+    BcdElementType ElementType;
+    PBCD_DEVICE_OPTION BcdDevice;
+
+    /* Starting at offset 0, loop every BCD option */
+    NextOffset = 0;
+    do
+    {
+        /* Get the current option, and its offset */
+        Option = (PBL_BCD_OPTION)((ULONG_PTR)BootEntry->BcdData + NextOffset);
+        NextOffset = Option->NextEntryOffset;
+
+        /* If it's empty, ignore it */
+        if (Option->Empty)
+        {
+            continue;
+        }
+
+        /* If it's not a device option, ignore it */
+        ElementType.PackedValue = Option->Type;
+        if (ElementType.Format != BCD_TYPE_DEVICE)
+        {
+            continue;
+        }
+
+        /* Get the data offset */
+        DataOffset = Option->DataOffset;
+
+        /* Extract the device out of it */
+        BcdDevice = (PBCD_DEVICE_OPTION)((ULONG_PTR)BootEntry->BcdData + DataOffset);
+
+        /* If the device is already fully specified, no need to build it */
+        if (!(BcdDevice->DeviceDescriptor.Flags & 1))
+        {
+            continue;
+        }
+
+        /* Otherwise, check if there's any list options as well */
+        ListOption = NULL;
+        ListOffset = Option->ListOffset;
+        if (Option->ListOffset)
+        {
+            ListOption = (PBL_BCD_OPTION)((ULONG_PTR)BootEntry->BcdData + ListOffset);
+        }
+
+        /* And now call BlCreateDevice to build the full device descriptor */
+        EfiPrintf(L"Unspecified devices not yet supported: %p\r\n", ListOption);
+        return STATUS_NOT_SUPPORTED;
+    } while (NextOffset != 0);
+
+    /* Devices created successfully */
+    return STATUS_SUCCESS;
+}
+
+NTSTATUS
+BmpTransferExecution (
+    _In_ PBL_LOADED_APPLICATION_ENTRY BootEntry,
+    _Out_ PULONG LaunchCode,
+    _Out_ PBOOLEAN Recover
+    )
+{
+    PWCHAR AppPath;
+    NTSTATUS Status;
+    PBL_DEVICE_DESCRIPTOR AppDevice;
+    BL_RETURN_ARGUMENTS ReturnArgs;
+    BOOLEAN AdvancedOptions;
+    ULONG AppHandle;
+
+    /* Get the application path */
+    Status = BlGetBootOptionString(BootEntry->BcdData,
+                                   BcdLibraryString_ApplicationPath,
+                                   &AppPath);
+    if (!NT_SUCCESS(Status))
+    {
+        /* If we couldn't find one, set this to NULL */
+        AppPath = NULL;
+    }
+
+    /* Check if this is a PXE startup.com */
+    if (BootEntry->Flags & BL_APPLICATION_ENTRY_STARTUP)
+    {
+#if BL_NET_SUPPORT
+        /* Do soft reboot to launch it */
+        Status = BlNetSoftReboot(BootEntry);
+#else
+        EfiPrintf(L"Net boot not supported\r\n");
+        Status = STATUS_NOT_SUPPORTED;
+#endif
+        /* Nothing else for us to do */
+        goto Quickie;
+    }
+
+    /* Loop as long as boot was not cancelled */
+    do
+    {
+        /* Load the boot application */
+        Status = BlImgLoadBootApplication(BootEntry, &AppHandle);
+
+        /* Did we not find it? */
+        if (Status == STATUS_NOT_FOUND)
+        {
+            /* Get the device for the boot application */
+            Status = BlGetBootOptionDevice(BootEntry->BcdData,
+                                           BcdLibraryDevice_ApplicationDevice,
+                                           &AppDevice,
+                                           NULL);
+            if (!NT_SUCCESS(Status))
+            {
+                /* Force re-enumeration */
+                Status = BlFwEnumerateDevice(AppDevice);
+            }
+
+            /* Did re-enumeration work? */
+            if (!NT_SUCCESS(Status))
+            {
+                /* Nope, raise a fatal error */
+                BmFatalErrorEx(BL_FATAL_ERROR_APP_LOAD,
+                               (ULONG_PTR)AppPath,
+                               Status,
+                               0,
+                               0);
+                goto Quickie;
+            }
+
+            /* Yes, try booting it again */
+            Status = BlImgLoadBootApplication(BootEntry, &AppHandle);
+        }
+
+        /* Was boot cancelled?*/
+        if (Status == STATUS_CANCELLED)
+        {
+            /* Should we display the menu, or is there no launch sequence? */
+            if ((BmGetBootMenuPolicy(BootEntry) != MenuPolicyStandard) ||
+                !(MiscGetBootOption(BootEntry->BcdData,
+                                    BcdLibraryObjectList_RecoverySequence)))
+            {
+                /* Bail out, the menu will take care of it */
+                goto Quickie;
+            }
+
+            /* No menu and there's a sequence, launch it */
+            *LaunchCode = 4;
+            *Recover = TRUE;
+            goto Quickie;
+        }
+
+        /* STATUS_FVE_LOCKED_VOLUME -- bitlocker volume is locked */
+        if (Status == 0xC0210000)
+        {
+            /* Launch recovery mode */
+            *LaunchCode = 4;
+            *Recover = TRUE;
+            goto Quickie;
+        }
+
+        /* Was there some other error launching the boot application? */
+        if (!NT_SUCCESS(Status))
+        {
+            /* Raise a fatal error */
+            BmFatalErrorEx(BL_FATAL_ERROR_APP_LOAD,
+                           (ULONG_PTR)AppPath,
+                           Status,
+                           0,
+                           0);
+            goto Quickie;
+        }
+
+        /* Zero out the return arguments */
+        RtlZeroMemory(&ReturnArgs, sizeof(ReturnArgs));
+
+        /* Log to ETW this launch */
+        //BmpLogApplicationLaunchEvent(&BootEntry->Guid, AppPath);
+
+        /* Launch the boot application*/
+        Status = BlImgStartBootApplication(AppHandle, &ReturnArgs);
+
+#if BL_BITLOCKER_SUPPORT
+        /* Bitlocker stuff */
+        BlFveSecureBootCheckpointAppReturn(BootEntry, &ReturnArgs);
+#endif
+
+        /* Log in the boot status log the launch */
+        //BlBsdLogEntry(1, 0x12, &BootEntry->Guid, 0x14);
+
+        /* Unloac the boot application if we've returned */
+        BlImgUnloadBootApplication(AppHandle);
+
+        /* Keep going unless STATUS_RESTART_BOOT_APPLICATION */
+    } while (Status != 0xC0000453);
+
+    /* We've come back. Assume we need to launch the recovery sequence */
+    *Recover = TRUE;
+
+    /* Why did we get back? */
+    if (ReturnArgs.Flags & 1)
+    {
+        /* Flag 1 -- should we display advanced options? */
+        Status = BlGetBootOptionBoolean(BootEntry->BcdData,
+                                        BcdLibraryBoolean_DisplayAdvancedOptions,
+                                        &AdvancedOptions);
+        if ((NT_SUCCESS(Status)) && (AdvancedOptions))
+        {
+            /* Yes, so return with code 2 */
+            *LaunchCode = 2;
+        }
+        else
+        {
+            /* No, return with code 1 */
+            *LaunchCode = 1;
+        }
+    }
+    else if (ReturnArgs.Flags & 4)
+    {
+        /* Flag 4 -- unknown */
+        *LaunchCode = 1;
+    }
+    else if (ReturnArgs.Flags & 8)
+    {
+        /* Flag 5 -- unknown */
+        *LaunchCode = 5;
+    }
+    else if (ReturnArgs.Flags & 0x10)
+    {
+        /* Flag 6 -- unknown */
+        *LaunchCode = 6;
+    }
+    else if (ReturnArgs.Flags & 0x20)
+    {
+        /* Flag 7 -- unknown */
+        *LaunchCode = 7;
+    }
+    else if (ReturnArgs.Flags & BL_RETURN_ARGUMENTS_NO_PAE_FLAG)
+    {
+        /* PAE is not supported -- refuse to boot */
+        *Recover = FALSE;
+        BmFatalErrorEx(BL_FATAL_ERROR_NO_PAE, Status, 0, 0, 0);
+    }
+
+Quickie:
+    /* All done, did we have an application path? */
+    if (AppPath)
+    {
+        /* Free it */
+        BlMmFreeHeap(AppPath);
+    }
+
+    /* Back to the caller now */
+    return Status;
+}
+
 NTSTATUS
 BmpLaunchBootEntry (
     _In_ PBL_LOADED_APPLICATION_ENTRY BootEntry,
     _Out_ PULONG EntryIndex,
-    _In_ ULONG Unknown,
+    _In_ ULONG LaunchCode,
     _In_ BOOLEAN LaunchWinRe
     )
 {
-    EfiPrintf(L"Boot launch not yet implemented\r\n");
-    EfiStall(1000000000);
-    return STATUS_NOT_IMPLEMENTED;
+    HANDLE BcdHandle;
+    NTSTATUS Status;
+    GUID ObjectId;
+    BOOLEAN DoRecovery, AutoRecovery, DoSequence, RestartOnFailure;
+    ULONG ErrorCode;
+    BOOLEAN AdvancedOneTime, EditOneTime;
+
+    /* Check if this is the OS loader */
+    if (BootEntry->Flags & BL_APPLICATION_ENTRY_WINLOAD)
+    {
+        /* Check if one-time advanced options should be shown */
+        if (MiscGetBootOption(BootEntry->BcdData,
+                              BcdOSLoaderBoolean_AdvancedOptionsOneTime))
+        {
+            /* Open the BCD */
+            BcdHandle = NULL;
+            Status = BmOpenDataStore(BcdHandle);
+            if (NT_SUCCESS(Status))
+            {
+                /* Delete the option from the BCD, so it doesn't happen again */
+                ObjectId = BootEntry->Guid;
+                BmPurgeOption(BcdHandle,
+                              &ObjectId,
+                              BcdOSLoaderBoolean_AdvancedOptionsOneTime);
+                BmCloseDataStore(BcdHandle);
+            }
+        }
+
+        /* Check if one-time options editor should be shown */
+        if (MiscGetBootOption(BootEntry->BcdData,
+                              BcdOSLoaderBoolean_OptionsEditOneTime))
+        {
+            /* Open the BCD */
+            BcdHandle = NULL;
+            Status = BmOpenDataStore(BcdHandle);
+            if (NT_SUCCESS(Status))
+            {
+                /* Delete the option from the BCD, so it doesn't happen again */
+                ObjectId = BootEntry->Guid;
+                BmPurgeOption(BcdHandle,
+                              &ObjectId,
+                              BcdOSLoaderBoolean_OptionsEditOneTime);
+                BmCloseDataStore(BcdHandle);
+            }
+        }
+    }
+
+TryAgain:
+    /* Disable recovery mode */
+    DoRecovery = FALSE;
+
+    /* Store globally which entry we are trying to boot */
+    BmpSelectedBootEntry = BootEntry;
+
+    /* Create any devices that aren't yet fully defined for this boot entry */
+    Status = BmpCreateDevices(BootEntry);
+    if (!NT_SUCCESS(Status))
+    {
+        /* That failed -- can we launch the recovery environment? */
+        if (!LaunchWinRe)
+        {
+            return Status;
+        }
+
+        /* Yes, so return with the WinRe launch code */
+        LaunchCode = 2;
+        goto Quickie;
+    }
+
+    /* Is this an OS loader/ */
+    if (BootEntry->Flags & BL_APPLICATION_ENTRY_WINLOAD)
+    {
+        /* Is the one-time advanced options menu option present? */
+        Status = BlGetBootOptionBoolean(BootEntry->BcdData,
+                                        BcdOSLoaderBoolean_AdvancedOptionsOneTime,
+                                        &AdvancedOneTime);
+        if (NT_SUCCESS(Status))
+        {
+            /* Is it turned on? */
+            if (AdvancedOneTime)
+            {
+                /* Set the option this once */
+                BlAppendBootOptionBoolean(BootEntry,
+                                          BcdLibraryBoolean_DisplayAdvancedOptions);
+            }
+            else
+            {
+                /* It's not, so disable the option if active */
+                BlRemoveBootOption(BootEntry->BcdData,
+                                   BcdLibraryBoolean_DisplayAdvancedOptions);
+            }
+
+            /* Remove the one-time option. We've already purged it earlier */
+            BlRemoveBootOption(BootEntry->BcdData,
+                               BcdOSLoaderBoolean_AdvancedOptionsOneTime);
+        }
+
+        /* Is the one-time options editor menu option present? */
+        Status = BlGetBootOptionBoolean(BootEntry->BcdData,
+                                        BcdOSLoaderBoolean_OptionsEditOneTime,
+                                        &EditOneTime);
+        if (NT_SUCCESS(Status))
+        {
+            /* Is it turned on? */
+            if (EditOneTime)
+            {
+                /* Set the option this once */
+                BlAppendBootOptionBoolean(BootEntry,
+                                          BcdLibraryBoolean_DisplayOptionsEdit);
+            }
+            else
+            {
+                /* It's not, so disable the option if active */
+                BlRemoveBootOption(BootEntry->BcdData,
+                                   BcdLibraryBoolean_DisplayOptionsEdit);
+            }
+
+            /* Remove the one-time option. We've already purged it earlier */
+            BlRemoveBootOption(BootEntry->BcdData,
+                               BcdOSLoaderBoolean_OptionsEditOneTime);
+        }
+    }
+
+    /* BCD handling done, transfer execution to this entry */
+    Status = BmpTransferExecution(BootEntry, &LaunchCode, &DoRecovery);
+    if (!LaunchWinRe)
+    {
+        return Status;
+    }
+
+    /* Check if boot was successful, or  cancelled and we're not doing WinRE */
+    if (((NT_SUCCESS(Status)) || (Status == STATUS_CANCELLED)) && !(DoRecovery))
+    {
+        return Status;
+    }
+
+    /* Boot failed -- are we doing recovery? */
+    if (!DoRecovery)
+    {
+        /* Nope, bail out */
+        LaunchCode = 2;
+        goto Quickie;
+    }
+
+Quickie:
+    /* Get the recovery sequence */
+    if (MiscGetBootOption(BootEntry->BcdData, BcdLibraryObjectList_RecoverySequence))
+    {
+        /* Check if the launch depends on auto-recovery being enabled or not */
+        if ((LaunchCode == 3) || (LaunchCode == 5) || (LaunchCode == 6))
+        {
+            Status = BlGetBootOptionBoolean(BootEntry->BcdData,
+                                            BcdLibraryBoolean_AutoRecoveryEnabled,
+                                            &AutoRecovery);
+            if (NT_SUCCESS(Status))
+            {
+                /* Override the setting */
+                DoRecovery = AutoRecovery;
+            }
+        }
+    }
+    else
+    {
+        /* There's no recovery setting */
+        DoRecovery = FALSE;
+    }
+
+    /* Check if we should restart on failure */
+    RestartOnFailure = FALSE;
+    BlGetBootOptionBoolean(BlpApplicationEntry.BcdData,
+                           BcdLibraryBoolean_RestartOnFailure,
+                           &RestartOnFailure);
+
+    /* Do the sequence if recovery is on, unless we should restart instead */
+    DoSequence = RestartOnFailure ? FALSE : DoRecovery;
+    while (1)
+    {
+        /* Are we doing the recovery sequence? */
+        if (DoSequence)
+        {
+            /* Because of automatic recovery? */
+            if (AutoRecovery)
+            {
+#if BL_BITLOCKER_SUPPORT
+                /* Do bitlocker stuff */
+                BlFveRegisterBootEntryForTrustedWimBoot(BootEntry, TRUE);
+#endif
+            }
+
+            /* Launch the recovery sequence*/
+            Status = BmLaunchRecoverySequence(BootEntry, LaunchCode);
+
+            /* Was it launched automatically? */
+            if (AutoRecovery)
+            {
+#if BL_BITLOCKER_SUPPORT
+                /* Do bitlocker stuff */
+                BlFveRegisterBootEntryForTrustedWimBoot(BootEntry, FALSE);
+#endif
+
+                /* No need to do this again */
+                AutoRecovery = FALSE;
+            }
+
+            /* Did the recovery sequence work? */
+            if (NT_SUCCESS(Status))
+            {
+                /* All good */
+                return STATUS_SUCCESS;
+            }
+
+            /* Remove the sequence, don't do it again */
+            BlRemoveBootOption(BootEntry->BcdData, BcdLibraryObjectList_RecoverySequence);
+        }
+
+        /* Recovery sequence also failed, show fatal error */
+        if (!BmpInternalBootError)
+        {
+            BmFatalErrorEx(BL_FATAL_ERROR_GENERIC, Status, 0, 0, 0);
+        }
+
+        /* Display the error menu */
+        ErrorCode = BmDisplayDumpError(BootEntry, LaunchCode);
+        BmErrorPurge();
+
+        /* See what the user wants to do */
+        switch (ErrorCode)
+        {
+            case TryAgain:
+                /* Try again */
+                goto TryAgain;
+
+            case NextOs:
+                /* Boot the next entry*/
+                break;
+
+            case OsSelection:
+                /* Cancel the boot*/
+                return STATUS_CANCELLED;
+
+            case RecoverOem:
+                /* Custom OEM recovery -- open the BCD */
+                Status = BmOpenDataStore(BcdHandle);
+                if (NT_SUCCESS(Status))
+                {
+                    /* See what the custom sequence is */
+                    Status = BmProcessCustomAction(BcdHandle, NULL);
+                }
+
+                /* All done, close the BCD */
+                if (BcdHandle)
+                {
+                    BmCloseDataStore(BcdHandle);
+                }
+                return Status;
+
+            case AdvancedOptions:
+                /* Show the advanced options next iteration */
+                BlAppendBootOptionBoolean(BootEntry, BcdOSLoaderBoolean_AdvancedOptionsOneTime);
+                goto TryAgain;
+
+            case BootOptions:
+                /* Show the options editor next iteration */
+                BlAppendBootOptionBoolean(BootEntry, BcdOSLoaderBoolean_OptionsEditOneTime);
+                goto TryAgain;
+
+            case Recover:
+                /* Try the recovery sequence next time*/
+                DoSequence = TRUE;
+                LaunchCode = 1;
+                goto TryAgain;
+
+            default:
+                /* Something unknown */
+                return STATUS_CANCELLED;
+        }
+    }
+
+    /* We are booting the next OS, so return success as to not kill the boot */
+    return STATUS_SUCCESS;
 }
 
 /*++
@@ -2056,6 +2724,7 @@ BmpLaunchBootEntry (
  *
  *--*/
 NTSTATUS
+NTAPI
 BmMain (
     _In_ PBOOT_APPLICATION_PARAMETER_BLOCK BootParameters
     )
@@ -2226,10 +2895,32 @@ BmMain (
         }
     }
 
+
+    /* TEST MODE */
+    EfiPrintf(L"Performing memory allocator tests...\r\n");
+    {
+        NTSTATUS Status;
+        PHYSICAL_ADDRESS PhysicalAddress;
+        PBL_MEMORY_DESCRIPTOR Found;
+
+        /* Allocate 1 physical page */
+        PhysicalAddress.QuadPart = 0;
+        Status = BlMmAllocatePhysicalPages(&PhysicalAddress, BlLoaderData, 1, 0, 1);
+        EfiPrintf(L"Allocation status: %lx at address: %llx\r\n", Status, PhysicalAddress.QuadPart);
+        EfiStall(10000);
+
+        Found = MmMdFindDescriptor(BL_MM_INCLUDE_UNMAPPED_ALLOCATED, 0, PhysicalAddress.QuadPart >> PAGE_SHIFT);
+        EfiPrintf(L"Found descriptor: %p %llx\r\n", Found, Found->BasePage);
+
+        Status = BlMmFreePhysicalPages(PhysicalAddress);
+        EfiPrintf(L"Memory free status: %lx\r\n", Status);
+    }
+
+
     /* Write out the first XML tag */
     BlXmiWrite(L"<bootmgr/>");
 
-    /* Check for factory resset */
+    /* Check for factory reset */
     BlSecureBootCheckForFactoryReset();
 
     /* Load the revocation list */
@@ -2278,12 +2969,14 @@ BmMain (
                                     &CustomActions);
     if ((NT_SUCCESS(Status)) && (CustomActions))
     {
-        /* We don't suppport this yet */
+        /* We don't support this yet */
         EfiPrintf(L"Not implemented\r\n");
         Status = STATUS_NOT_IMPLEMENTED;
         goto Failure;
     }
 
+    //BlResourceFindMessage(BM_MSG_TEST);
+
     /* At last, enter the boot selection stage */
     SequenceId = 0;
     GetEntry = FALSE;