[BOOT] Execution contexts and more refactoring

BOOT/ENVIRON/LIB/EFI/eficon.c

@@ -17,6 +17,7 @@ Abstract:
 #include <wchar.h>
 #include "bootlib.h"
 
+extern PEXECUTION_CONTEXT CurrentExecutionContext;
 extern SIMPLE_TEXT_OUTPUT_INTERFACE *EfiConOut;
 
 VOID
@@ -41,7 +42,18 @@ Return Value:
 --*/
 
 {
+    EXECUTION_CONTEXT_TYPE ContextType;
+
+    ContextType = CurrentExecutionContext->Type;
+    if (ContextType != ExecutionContextFirmware) {
+        BlpArchSwitchContext(ExecutionContextFirmware);
+    }
+
     EfiConOut->OutputString(EfiConOut, String);
+
+    if (ContextType != ExecutionContextFirmware) {
+        BlpArchSwitchContext(ContextType);
+    }
 }
 
 VOID

BOOT/ENVIRON/LIB/EFI/efifw.c

@@ -13,9 +13,12 @@ Abstract:
 
 --*/
 
+#include <ntrtl.h>
 #include "bootlib.h"
 #include "efi.h"
 
+BOOT_FIRMWARE_DATA EfiFirmwareData;
+PBOOT_FIRMWARE_DATA EfiFirmwareParameters;
 EFI_SYSTEM_TABLE *EfiST;
 EFI_BOOT_SERVICES *EfiBS;
 EFI_RUNTIME_SERVICES *EfiRT;
@@ -36,7 +39,10 @@ Routine Description:
 
 Arguments:
 
-    Stage - 0 or 1.
+    Stage - Which stage of initialization to perform.
+
+        Stage 0: Initialize global firmware-related data and pointers.
+            Once this stage is complete, ConsolePrint() and ConsolePrintf() can be used.
 
     FirmwareData - Pointer to BOOT_FIRMWARE_DATA.
 
@@ -53,6 +59,9 @@ Return Value:
     }
 
     if (Stage == 0) {
+        RtlCopyMemory(&EfiFirmwareData, FirmwareData, sizeof(BOOT_FIRMWARE_DATA));
+        EfiFirmwareParameters = &EfiFirmwareData;
+
         EfiST = FirmwareData->SystemTable;
         EfiBS = EfiST->BootServices;
         EfiRT = EfiST->RuntimeServices;
@@ -60,5 +69,9 @@ Return Value:
         EfiConIn = EfiST->ConIn;
     }
 
+    //
+    // TODO: Implement stage 1 initialization.
+    //
+
     return STATUS_SUCCESS;
 }

BOOT/ENVIRON/LIB/EFI/efiinit.c

@@ -615,7 +615,6 @@ Return Value:
     PMEMORY_DESCRIPTOR MemoryDescriptor;
     PBOOT_DEVICE BootDevice;
     PBOOT_FIRMWARE_DATA FirmwareData;
-    PBOOT_PLATFORM_DATA PlatformData;
     PBOOT_RETURN_DATA ReturnData;
 
     ScratchUsed = 0;
@@ -697,19 +696,14 @@ Return Value:
     InputParameters->FirmwareDataOffset = ScratchUsed;
     FirmwareData = (PBOOT_FIRMWARE_DATA)(&EfiInitScratch[ScratchUsed]);
     ScratchUsed += sizeof(BOOT_FIRMWARE_DATA);
+    RtlZeroMemory(FirmwareData, sizeof(BOOT_FIRMWARE_DATA));
     FirmwareData->Version = BOOT_FIRMWARE_DATA_VERSION;
-    FirmwareData->Reserved = 0;
     FirmwareData->ImageHandle = ImageHandle;
     FirmwareData->SystemTable = SystemTable;
-
-    InputParameters->PlatformDataOffset = ScratchUsed;
-    PlatformData = (PBOOT_PLATFORM_DATA)(&EfiInitScratch[ScratchUsed]);
-    ScratchUsed += sizeof(BOOT_PLATFORM_DATA);
-    PlatformData->Reserved = 0;
 #if defined(__i386__) || defined(__x86_64__)
-    asm volatile("mov %%cr3, %0" :"=r"(PlatformData->CR3));
+    asm volatile("mov %%cr3, %0" :"=r"(FirmwareData->Cr3));
 #endif
-    BlpArchGetDescriptorTableContext(&PlatformData->DescriptorTableContext);
+    BlpArchGetDescriptorTableContext(&FirmwareData->DescriptorTableContext);
 
     InputParameters->ReturnDataOffset = ScratchUsed;
     ReturnData = (PBOOT_RETURN_DATA)(&EfiInitScratch[ScratchUsed]);

BOOT/ENVIRON/LIB/EFI/efimm.c

@@ -62,15 +62,32 @@ Return Value:
 --*/
 
 {
-    return EfiGetNtStatusCode(
-        EfiBS->GetMemoryMap(
-            MemoryMapSize,
-            MemoryMap,
-            MapKey,
-            DescriptorSize,
-            DescriptorVersion
-        )
+    EXECUTION_CONTEXT_TYPE ContextType;
+    EFI_STATUS EfiStatus;
+
+    ContextType = CurrentExecutionContext->Type;
+    if (ContextType != ExecutionContextFirmware) {
+        //
+        // TODO: Translate addresses here.
+        // Need MmArchTranslateVirtualAddress().
+        //
+
+        BlpArchSwitchContext(ExecutionContextFirmware);
+    }
+
+    EfiStatus = EfiBS->GetMemoryMap(
+        MemoryMapSize,
+        MemoryMap,
+        MapKey,
+        DescriptorSize,
+        DescriptorVersion
     );
+
+    if (ContextType != ExecutionContextFirmware) {
+        BlpArchSwitchContext(ContextType);
+    }
+
+    return EfiGetNtStatusCode(EfiStatus);
 }
 
 NTSTATUS
@@ -106,14 +123,26 @@ Return Value:
 --*/
 
 {
-    return EfiGetNtStatusCode(
-        EfiBS->AllocatePages(
-            Type,
-            MemoryType,
-            Pages,
-            Memory
-        )
+    EXECUTION_CONTEXT_TYPE ContextType;
+    EFI_STATUS EfiStatus;
+
+    ContextType = CurrentExecutionContext->Type;
+    if (ContextType != ExecutionContextFirmware) {
+        BlpArchSwitchContext(ExecutionContextFirmware);
+    }
+
+    EfiStatus = EfiBS->AllocatePages(
+        Type,
+        MemoryType,
+        Pages,
+        Memory
     );
+
+    if (ContextType != ExecutionContextFirmware) {
+        BlpArchSwitchContext(ContextType);
+    }
+
+    return EfiGetNtStatusCode(EfiStatus);
 }
 
 NTSTATUS
@@ -143,12 +172,24 @@ Return Value:
 --*/
 
 {
-    return EfiGetNtStatusCode(
-        EfiBS->FreePages(
-            Memory,
-            Pages
-        )
+    EXECUTION_CONTEXT_TYPE ContextType;
+    EFI_STATUS EfiStatus;
+
+    ContextType = CurrentExecutionContext->Type;
+    if (ContextType != ExecutionContextFirmware) {
+        BlpArchSwitchContext(ExecutionContextFirmware);
+    }
+
+    EfiStatus = EfiBS->FreePages(
+        Memory,
+        Pages
     );
+
+    if (ContextType != ExecutionContextFirmware) {
+        BlpArchSwitchContext(ContextType);
+    }
+
+    return EfiGetNtStatusCode(EfiStatus);
 }
 
 MEMORY_TYPE
@@ -616,7 +657,7 @@ Return Value:
     //
     // Remove the current descriptor.
     //
-    Status = MmMdRemoveRegionFromMdlEx(Mdl, NtStartPage, NtPageCount, MDL_OPERATION_FLAGS_PHYSICAL, NULL);
+    Status = MmMdRemoveRegionFromMdl(Mdl, NtStartPage, NtPageCount, MDL_OPERATION_FLAGS_PHYSICAL);
     if (!NT_SUCCESS(Status)) {
         MmMdFreeDescriptor(NtDescriptor);
         goto exit;