[BOOT:LIB] More initialization and cleanup

Started BlpMmDestroy(), MmMdDestroy(), MmPaDestroy(),
EfiSetWatchdogTimer(), EfiOpenProtocol(), EfiConInExSetState(), and
BlDestroyLibrary().
Completed BlpFwInitialize().
Improved InitializeLibrary().

BOOT/ENVIRON/INC/bootlib.h

@@ -426,11 +426,13 @@ extern BOOT_APPLICATION_ENTRY BlpApplicationEntry;
 
 #if defined(_EFI)
 extern PBOOT_FIRMWARE_DATA EfiFirmwareParameters;
+extern EFI_HANDLE EfiImageHandle;
 extern EFI_SYSTEM_TABLE *EfiST;
 extern EFI_BOOT_SERVICES *EfiBS;
 extern EFI_RUNTIME_SERVICES *EfiRT;
 extern SIMPLE_TEXT_OUTPUT_INTERFACE *EfiConOut;
 extern SIMPLE_INPUT_INTERFACE *EfiConIn;
+extern EFI_SIMPLE_TEXT_INPUT_EX_PROTOCOL *EfiConInEx;
 #endif
 
 VOID
@@ -469,6 +471,11 @@ BlpFwInitialize (
     IN PBOOT_FIRMWARE_DATA FirmwareData
     );
 
+NTSTATUS
+BlpMmDestroy (
+    IN ULONG Stage
+    );
+
 NTSTATUS
 BlpMmInitializeConstraints (
     VOID

BOOT/ENVIRON/INC/efiapi.h

@@ -16,48 +16,13 @@ Abstract:
 #ifndef _EFIAPI_H
 #define _EFIAPI_H
 
+#define EFI_MAKE_REVISION(Major, Minor) (((Major) << 16) | (Minor))
 #define EFI_SPECIFICATION_MAJOR_REVISION 1
 #define EFI_SPECIFICATION_MINOR_REVISION 02
-#define EFI_SPECIFICATION_VERSION ((EFI_SPECIFICATION_MAJOR_REVISION << 16) || EFI_SPECIFICATION_MINOR_REVISION)
+#define EFI_SPECIFICATION_VERSION EFI_MAKE_REVISION(EFI_SPECIFICATION_MAJOR_REVISION, EFI_SPECIFICATION_MINOR_REVISION)
 
 INTERFACE_DECL(_EFI_SYSTEM_TABLE);
 
-/*
- * Loaded image protocol definitions.
- */
-
-typedef
-EFI_STATUS
-(EFIAPI *EFI_IMAGE_UNLOAD) (
-    IN EFI_HANDLE ImageHandle
-    );
-
-#define EFI_LOADED_IMAGE_PROTOCOL_GUID \
-    { 0x5b1b31a1, 0x9562, 0x11d2, { 0x8e, 0x3f, 0x00, 0xa0, 0xc9, 0x69, 0x72, 0x3b } }
-#define LOADED_IMAGE_PROTOCOL EFI_LOADED_IMAGE_PROTOCOL_GUID
-
-#define EFI_IMAGE_INFORMATION_REVISION 0x1000
-
-typedef struct {
-    UINT32 Revision;
-    EFI_HANDLE ParentHandle;
-    struct _EFI_SYSTEM_TABLE *SystemTable;
-
-    EFI_HANDLE DeviceHandle;
-    EFI_DEVICE_PATH *FilePath;
-    VOID *Reserved;
-
-    UINT32 LoadOptionsSize;
-    VOID *LoadOptions;
-
-    VOID *ImageBase;
-    UINT64 ImageSize;
-    EFI_MEMORY_TYPE ImageCodeType;
-    EFI_MEMORY_TYPE ImageDataType;
-
-    EFI_IMAGE_UNLOAD Unload;
-} EFI_LOADED_IMAGE;
-
 /*
  * EFI table header.
  */
@@ -141,6 +106,42 @@ EFI_STATUS
     OUT VOID **Interface
     );
 
+typedef
+EFI_STATUS
+(EFIAPI *EFI_OPEN_PROTOCOL) (
+    IN EFI_HANDLE Handle,
+    IN EFI_GUID *Protocol,
+    OUT VOID **Interface,
+    IN EFI_HANDLE AgentHandle,
+    IN EFI_HANDLE ControllerHandle,
+    IN UINT32 Attributes
+    );
+
+#define EFI_OPEN_PROTOCOL_BY_HANDLE_PROTOCOL  0x00000001
+#define EFI_OPEN_PROTOCOL_GET_PROTOCOL        0x00000002
+#define EFI_OPEN_PROTOCOL_TEST_PROTOCOL       0x00000004
+#define EFI_OPEN_PROTOCOL_BY_CHILD_CONTROLLER 0x00000008
+#define EFI_OPEN_PROTOCOL_BY_DRIVER           0x00000010
+#define EFI_OPEN_PROTOCOL_EXCLUSIVE           0x00000020
+
+typedef
+EFI_STATUS
+(EFIAPI *EFI_CLOSE_PROTOCOL) (
+    IN EFI_HANDLE Handle,
+    IN EFI_GUID *Protocol,
+    IN EFI_HANDLE AgentHandle,
+    IN EFI_HANDLE ControllerHandle
+    );
+
+typedef
+EFI_STATUS
+(EFIAPI *EFI_SET_WATCHDOG_TIMER) (
+    IN UINTN Timeout,
+    IN UINT64 WatchdogCode,
+    IN UINTN DataSize,
+    IN CHAR16 *WatchdogData OPTIONAL
+    );
+
 typedef struct _EFI_BOOT_SERVICES {
     EFI_TABLE_HEADER Hdr;
 
@@ -178,13 +179,13 @@ typedef struct _EFI_BOOT_SERVICES {
 
     EFI_HANDLE GetNextHighMonotonicCount;
     EFI_HANDLE Stall;
-    EFI_HANDLE SetWatchdogTimer;
+    EFI_SET_WATCHDOG_TIMER SetWatchdogTimer;
 
     EFI_HANDLE ConnectController;
     EFI_HANDLE DisconnectController;
 
-    EFI_HANDLE OpenProtocol;
+    EFI_OPEN_PROTOCOL OpenProtocol;
-    EFI_HANDLE CloseProtocol;
+    EFI_CLOSE_PROTOCOL CloseProtocol;
     EFI_HANDLE OpenProtocolInformation;
 
     EFI_HANDLE ProtocolsPerHandle;

BOOT/ENVIRON/INC/eficon.h

@@ -300,9 +300,85 @@ EFI_STATUS
     );
 
 typedef struct _SIMPLE_INPUT_INTERFACE {
-    EFI_INPUT_RESET    Reset;
+    EFI_INPUT_RESET Reset;
     EFI_INPUT_READ_KEY ReadKeyStroke;
-    EFI_EVENT          WaitForKey;
+    EFI_EVENT WaitForKey;
 } SIMPLE_INPUT_INTERFACE, EFI_SIMPLE_TEXT_IN_PROTOCOL;
 
+INTERFACE_DECL(_EFI_SIMPLE_TEXT_INPUT_EX_PROTOCOL);
+
+#define EFI_SIMPLE_TEXT_INPUT_EX_PROTOCOL_GUID \
+    { 0xdd9e7534, 0x7762, 0x4698, { 0x8c, 0x14, 0xf5, 0x85, 0x17, 0xa6, 0x25, 0xaa } }
+#define SIMPLE_TEXT_INPUT_EX_PROTOCOL EFI_SIMPLE_TEXT_INPUT_EX_PROTOCOL_GUID
+
+typedef UINT8 EFI_KEY_TOGGLE_STATE;
+
+#define EFI_SCROLL_LOCK_ACTIVE 0x01
+#define EFI_NUM_LOCK_ACTIVE    0x02
+#define EFI_CAPS_LOCK_ACTIVE   0x04
+#define EFI_KEY_STATE_EXPOSED  0x40
+#define EFI_TOGGLE_STATE_VALID 0x80
+
+typedef struct {
+    UINT32 KeyShiftState;
+    EFI_KEY_TOGGLE_STATE KeyToggleState;
+} EFI_KEY_STATE;
+
+typedef struct {
+    EFI_INPUT_KEY Key;
+    EFI_KEY_STATE KeyState;
+} EFI_KEY_DATA;
+
+typedef
+EFI_STATUS
+(EFIAPI *EFI_KEY_NOTIFY_FUNCTION) (
+    IN EFI_KEY_DATA *KeyData
+    );
+
+typedef
+EFI_STATUS
+(EFIAPI *EFI_INPUT_RESET_EX) (
+    IN struct _EFI_SIMPLE_TEXT_INPUT_EX_PROTOCOL *This,
+    IN BOOLEAN ExtendedVerification
+    );
+
+typedef
+EFI_STATUS
+(EFIAPI *EFI_INPUT_READ_KEY_EX) (
+    IN struct _EFI_SIMPLE_TEXT_INPUT_EX_PROTOCOL *This,
+    OUT EFI_KEY_DATA *KeyData
+    );
+
+typedef
+EFI_STATUS
+(EFIAPI *EFI_SET_STATE) (
+    IN struct _EFI_SIMPLE_TEXT_INPUT_EX_PROTOCOL *This,
+    IN EFI_KEY_TOGGLE_STATE *KeyToggleState
+    );
+
+typedef
+EFI_STATUS
+(EFIAPI *EFI_REGISTER_KEYSTROKE_NOTIFY) (
+    IN struct _EFI_SIMPLE_TEXT_INPUT_EX_PROTOCOL *This,
+    IN EFI_KEY_DATA *KeyData,
+    IN EFI_KEY_NOTIFY_FUNCTION KeyNotificationFunction,
+    OUT VOID **NotifyHandle
+    );
+
+typedef
+EFI_STATUS
+(EFIAPI *EFI_UNREGISTER_KEYSTROKE_NOTIFY) (
+    IN struct _EFI_SIMPLE_TEXT_INPUT_EX_PROTOCOL *This,
+    IN VOID *NotificationHandle
+    );
+
+typedef struct _EFI_SIMPLE_TEXT_INPUT_EX_PROTOCOL {
+    EFI_INPUT_RESET_EX Reset;
+    EFI_INPUT_READ_KEY_EX ReadKeyStrokeEx;
+    EFI_EVENT WaitForKeyEx;
+    EFI_SET_STATE SetState;
+    EFI_REGISTER_KEYSTROKE_NOTIFY RegisterKeyNotify;
+    EFI_UNREGISTER_KEYSTROKE_NOTIFY UnregisterKeyNotify;
+} EFI_SIMPLE_TEXT_INPUT_EX_PROTOCOL;
+
 #endif

BOOT/ENVIRON/INC/efilib.h

@@ -19,6 +19,8 @@ Abstract:
 #include "bootlib.h"
 #include "efi.h"
 
+extern EFI_GUID EfiSimpleTextInputExProtocol;
+
 PBOOT_APPLICATION_PARAMETERS
 EfiInitCreateInputParametersEx (
     IN EFI_HANDLE ImageHandle,
@@ -42,4 +44,48 @@ EfiGetNtStatusCode (
     IN EFI_STATUS Status
     );
 
+NTSTATUS
+EfiGetMemoryMap (
+    IN OUT UINTN *MemoryMapSize,
+    IN OUT EFI_MEMORY_DESCRIPTOR *MemoryMap,
+    IN OUT UINTN *MapKey,
+    IN OUT UINTN *DescriptorSize,
+    IN OUT UINT32 *DescriptorVersion
+    );
+
+NTSTATUS
+EfiAllocatePages (
+    IN EFI_ALLOCATE_TYPE Type,
+    IN EFI_MEMORY_TYPE MemoryType,
+    IN UINTN Pages,
+    IN OUT EFI_PHYSICAL_ADDRESS *Memory
+    );
+
+NTSTATUS
+EfiFreePages (
+    IN EFI_PHYSICAL_ADDRESS Memory,
+    IN UINTN Pages
+    );
+
+NTSTATUS
+EfiSetWatchdogTimer (
+    IN UINTN Timeout,
+    IN UINT64 WatchdogCode,
+    IN UINTN DataSize,
+    IN CHAR16 *WatchdogData
+    );
+
+NTSTATUS
+EfiOpenProtocol (
+    IN EFI_HANDLE Handle,
+    IN EFI_GUID *Protocol,
+    IN OUT VOID **Interface
+    );
+
+NTSTATUS
+EfiConInExSetState (
+    IN EFI_SIMPLE_TEXT_INPUT_EX_PROTOCOL *Protocol,
+    IN EFI_KEY_TOGGLE_STATE *KeyToggleState
+    );
+
 #endif

BOOT/ENVIRON/INC/efiprot.h

@@ -17,9 +17,41 @@ Abstract:
 #define _EFIPROT_H
 
 /*
- * Device path protocol definitions.
+ * Loaded image protocol definitions.
  */
 
+typedef
+EFI_STATUS
+(EFIAPI *EFI_IMAGE_UNLOAD) (
+    IN EFI_HANDLE ImageHandle
+    );
+
+#define EFI_LOADED_IMAGE_PROTOCOL_GUID \
+    { 0x5b1b31a1, 0x9562, 0x11d2, { 0x8e, 0x3f, 0x00, 0xa0, 0xc9, 0x69, 0x72, 0x3b } }
+#define LOADED_IMAGE_PROTOCOL EFI_LOADED_IMAGE_PROTOCOL_GUID
+
+#define EFI_IMAGE_INFORMATION_REVISION 0x1000
+
+typedef struct {
+    UINT32 Revision;
+    EFI_HANDLE ParentHandle;
+    struct _EFI_SYSTEM_TABLE *SystemTable;
+
+    EFI_HANDLE DeviceHandle;
+    EFI_DEVICE_PATH *FilePath;
+    VOID *Reserved;
+
+    UINT32 LoadOptionsSize;
+    VOID *LoadOptions;
+
+    VOID *ImageBase;
+    UINT64 ImageSize;
+    EFI_MEMORY_TYPE ImageCodeType;
+    EFI_MEMORY_TYPE ImageDataType;
+
+    EFI_IMAGE_UNLOAD Unload;
+} EFI_LOADED_IMAGE;
+
 #define EFI_DEVICE_PATH_PROTOCOL_GUID \
     { 0x9576e91, 0x6d3f, 0x11d2, { 0x8e, 0x39, 0x0, 0xa0, 0xc9, 0x69, 0x72, 0x3b } }
 #define DEVICE_PATH_PROTOCOL EFI_DEVICE_PATH_PROTOCOL_GUID

BOOT/ENVIRON/INC/mm.h

@@ -22,6 +22,8 @@ Abstract:
 #define MDL_OPERATION_FLAGS_PHYSICAL 0x40000000
 #define MDL_OPERATION_FLAGS_VIRTUAL  0x80000000
 
+extern ULONG MmTranslationType;
+
 NTSTATUS
 MmFwGetMemoryMap (
    IN OUT PMEMORY_DESCRIPTOR_LIST Mdl,
@@ -97,12 +99,21 @@ MmMdInitDescriptor (
     IN MEMORY_TYPE Type
     );
 
+NTSTATUS
+MmMdDestroy (
+    );
+
 VOID
 MmMdInitialize (
     IN ULONG Unused,
     IN PBOOT_LIBRARY_PARAMETERS LibraryParameters
     );
 
+NTSTATUS
+MmPaDestroy (
+    IN ULONG Stage
+    );
+
 NTSTATUS
 MmPaInitialize (
     IN PBOOT_MEMORY_INFO MemoryInfo,

BOOT/ENVIRON/LIB/EFI/efifw.c

@@ -16,14 +16,17 @@ Abstract:
 #include <ntrtl.h>
 #include "bootlib.h"
 #include "efi.h"
+#include "efilib.h"
 
 BOOT_FIRMWARE_DATA EfiFirmwareData;
 PBOOT_FIRMWARE_DATA EfiFirmwareParameters;
+EFI_HANDLE EfiImageHandle;
 EFI_SYSTEM_TABLE *EfiST;
 EFI_BOOT_SERVICES *EfiBS;
 EFI_RUNTIME_SERVICES *EfiRT;
 SIMPLE_TEXT_OUTPUT_INTERFACE *EfiConOut;
 SIMPLE_INPUT_INTERFACE *EfiConIn;
+EFI_SIMPLE_TEXT_INPUT_EX_PROTOCOL *EfiConInEx;
 
 NTSTATUS
 BlpFwInitialize (
@@ -54,6 +57,9 @@ Return Value:
 --*/
 
 {
+    NTSTATUS Status;
+    EFI_KEY_TOGGLE_STATE KeyToggleState;
+
     if (FirmwareData == NULL || FirmwareData->Version == 0) {
         return STATUS_INVALID_PARAMETER;
     }
@@ -62,16 +68,33 @@ Return Value:
         RtlCopyMemory(&EfiFirmwareData, FirmwareData, sizeof(BOOT_FIRMWARE_DATA));
         EfiFirmwareParameters = &EfiFirmwareData;
 
+        EfiImageHandle = FirmwareData->ImageHandle;
         EfiST = FirmwareData->SystemTable;
         EfiBS = EfiST->BootServices;
         EfiRT = EfiST->RuntimeServices;
         EfiConOut = EfiST->ConOut;
         EfiConIn = EfiST->ConIn;
-    }
+        EfiConInEx = NULL;
+    } else if (Stage == 1) {
+        //
+        // Open the extended console input protocol.
+        // If successful, tell it to capture partial key events.
+        //
+        Status = EfiOpenProtocol(
+            EfiST->ConsoleInHandle,
+            &EfiSimpleTextInputExProtocol,
+            (VOID**)&EfiConInEx
+        );
+        if (NT_SUCCESS(Status)) {
+            KeyToggleState = EFI_KEY_STATE_EXPOSED | EFI_TOGGLE_STATE_VALID;
+            EfiConInExSetState(EfiConInEx, &KeyToggleState);
+        }
 
-    //
+        //
-    // TODO: Implement stage 1 initialization.
+        // Disable the watchdog timer.
-    //
+        //
+        EfiSetWatchdogTimer(0, 0, 0, NULL);
+    }
 
     return STATUS_SUCCESS;
 }

BOOT/ENVIRON/LIB/EFI/efilib.c

@@ -14,6 +14,9 @@ Abstract:
 --*/
 
 #include "efilib.h"
+#include "mm.h"
+
+EFI_GUID EfiSimpleTextInputExProtocol = EFI_SIMPLE_TEXT_INPUT_EX_PROTOCOL_GUID;
 
 EFI_STATUS
 EfiGetEfiStatusCode (
@@ -84,7 +87,6 @@ Return Value:
     }
 }
 
-
 NTSTATUS
 EfiGetNtStatusCode (
     IN EFI_STATUS Status
@@ -154,4 +156,381 @@ Return Value:
     default:
         return STATUS_UNSUCCESSFUL;
     }
-}
+}
+
+NTSTATUS
+EfiGetMemoryMap (
+    IN OUT UINTN *MemoryMapSize,
+    IN OUT EFI_MEMORY_DESCRIPTOR *MemoryMap,
+    IN OUT UINTN *MapKey,
+    IN OUT UINTN *DescriptorSize,
+    IN OUT UINT32 *DescriptorVersion
+    )
+
+/*++
+
+Routine Description:
+
+    Wrapper for EFI_BOOT_SERVICES.GetMemoryMap(). Gets the firmware memory map
+    and places it into a buffer.
+
+Arguments:
+
+    MemoryMapSize - pointer to the size of the buffer.
+
+    MemoryMap - pointer to the buffer to store the memory map in.
+
+    MapKey - ponter to the memory map key.
+
+    DescriptorSize - pointer to the size of each memory map descriptor.
+
+    DescriptorVersion - pointer to the version of memory map descriptors.
+
+Return Value:
+
+    STATUS_SUCCESS if successful.
+    STATUS_INVALID_PARAMETER if MemoryMapSize and/or MemoryMap are invalid.
+    STATUS_BUFFER_TOO_SMALL if MemoryMapSize is too small.
+
+--*/
+
+{
+    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
+EfiAllocatePages (
+    IN EFI_ALLOCATE_TYPE Type,
+    IN EFI_MEMORY_TYPE MemoryType,
+    IN UINTN Pages,
+    IN OUT EFI_PHYSICAL_ADDRESS *Memory
+    )
+
+/*++
+
+Routine Description:
+
+    Wrapper for EFI_BOOT_SERVICES.AllocatePages(). Allocates contiguous pages
+    of physical memory.
+
+Arguments:
+
+    Type - the type of allocation.
+
+    MemoryType - the type of memory to allocate.
+
+    Pages - the number of pages to allocate.
+
+    Memory - pointer to a physical address of the allocation.
+
+Return Value:
+
+    STATUS_SUCCESS if successful.
+    STATUS_INVALID_PARAMETER if Type, MemoryType, and/or Memory are invalid.
+    STATUS_INSUFFICIENT_NVRAM_RESOURCES if the pages could not be allocated.
+    STATUS_NOT_FOUND if the pages could not be found.
+
+--*/
+
+{
+    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
+EfiFreePages (
+    IN EFI_PHYSICAL_ADDRESS Memory,
+    IN UINTN Pages
+    )
+
+/*++
+
+Routine Description:
+
+    Wrapper for EFI_BOOT_SERVICES.FreePages(). Frees contiguous pages
+    of physical memory.
+
+Arguments:
+
+    Memory - physical address of the pages to be freed.
+
+    Pages - the number of pages to free.
+
+Return Value:
+
+    STATUS_SUCCESS if successful.
+    STATUS_NOT_FOUND if the allocation was not found.
+    STATUS_INVALID_PARAMETER Memory and/or Pages are invalid.
+
+--*/
+
+{
+    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);
+}
+
+NTSTATUS
+EfiSetWatchdogTimer (
+    IN UINTN Timeout,
+    IN UINT64 WatchdogCode,
+    IN UINTN DataSize,
+    IN CHAR16 *WatchdogData
+    )
+
+/*++
+
+Routine Description:
+
+    Wrapper for EFI_BOOT_SERVICES.SetWatchdogTimer().
+    Sets the watchdog timer.
+
+Arguments:
+
+    Timeout - The number of seconds to set the timer to.
+        Setting this to 0 disables the timer.
+
+    WatchdogCode - The code to set when an event occurs.
+
+    DataSize - The size in bytes of WatchdogData.
+
+    WatchdogData - Optional pointer to a string containing
+        a description of the timer, possibly accompanied
+        by binary data.
+
+Return Value:
+
+    STATUS_SUCCESS if successful.
+    STATUS_INVALID_PARAMETER if WatchdogCode is invalid.
+    STATUS_NOT_SUPPORTED if there is no Watchdog timer.
+    STATUS_IO_DEVICE_ERROR if the timer could not be set.
+
+--*/
+
+{
+    EXECUTION_CONTEXT_TYPE ContextType;
+    EFI_STATUS EfiStatus;
+
+    if (Timeout != 0 && WatchdogCode <= 0xffff) {
+        return STATUS_INVALID_PARAMETER;
+    }
+
+    ContextType = CurrentExecutionContext->Type;
+    if (ContextType != ExecutionContextFirmware) {
+        //
+        // TODO: Translate addresses here.
+        // Need MmArchTranslateVirtualAddress().
+        //
+
+        BlpArchSwitchContext(ExecutionContextFirmware);
+    }
+
+    EfiStatus = EfiBS->SetWatchdogTimer(
+        Timeout,
+        WatchdogCode,
+        DataSize,
+        WatchdogData
+    );
+
+    if (ContextType != ExecutionContextFirmware) {
+        BlpArchSwitchContext(ContextType);
+    }
+
+    return EfiGetNtStatusCode(EfiStatus);
+}
+
+NTSTATUS
+EfiOpenProtocol (
+    IN EFI_HANDLE Handle,
+    IN EFI_GUID *Protocol,
+    IN OUT VOID **Interface
+    )
+
+/*++
+
+Routine Description:
+
+    Wrapper that uses either EFI_BOOT_SERVICES.OpenProtocol() or HandleProtocol().
+    Opens a handle to a protocol and finds its interface.
+
+Arguments:
+
+    Handle - The handle to the protocol to open.
+
+    Protocol - The GUID of the protocol.
+
+    Intercce - Pointer that recieves the address of the interface.
+
+Return Value:
+
+    STATUS_SUCCESS if successful.
+    Error code on failure.
+
+--*/
+
+{
+    EXECUTION_CONTEXT_TYPE ContextType;
+    EFI_STATUS EfiStatus;
+    NTSTATUS Status;
+
+    if (MmTranslationType != TRANSLATION_TYPE_NONE) {
+        //
+        // TODO: Translate addresses.
+        // Need EfiVmOpenProtocol().
+        //
+        DebugPrint(L"EfiOpenProtocol(): Virtual address translation not supported\r\n");
+        return STATUS_NOT_IMPLEMENTED;
+    }
+
+    ContextType = CurrentExecutionContext->Type;
+    if (ContextType != ExecutionContextFirmware) {
+        BlpArchSwitchContext(ExecutionContextFirmware);
+    }
+
+    //
+    // If supported, use OpenProtocol() (EFI 1.10+).
+    // It helps ensure the protocol is not uninstalled
+    // when still in use.
+    //
+    if (EfiST->Hdr.Revision >= EFI_MAKE_REVISION(1, 10)) {
+        EfiStatus = EfiBS->OpenProtocol(
+            Handle,
+            Protocol,
+            Interface,
+            EfiImageHandle,
+            NULL,
+            EFI_OPEN_PROTOCOL_GET_PROTOCOL
+        );
+    } else {
+        EfiStatus = EfiBS->HandleProtocol(
+            Handle,
+            Protocol,
+            Interface
+        );
+    }
+
+    if (ContextType != ExecutionContextFirmware) {
+        BlpArchSwitchContext(ContextType);
+    }
+
+    //
+    // Convert EFI status to NT status.
+    //
+    Status = EfiGetNtStatusCode(EfiStatus);
+    if (!NT_SUCCESS(Status)) {
+        *Interface = NULL;
+    }
+
+    return Status;
+}
+
+NTSTATUS
+EfiConInExSetState (
+    IN EFI_SIMPLE_TEXT_INPUT_EX_PROTOCOL *Protocol,
+    IN EFI_KEY_TOGGLE_STATE *KeyToggleState
+    )
+
+/*++
+
+Routine Description:
+
+    Wrapper around EFI_SIMPLE_TEXT_INPUT_EX_PROTOCOL.SetState().
+    Sets an input protocol's state.
+
+Arguments:
+
+    Protocol - The protocol to set the state of.
+
+    KeyToggle State - The state to set.
+
+Return Value:
+
+    STATUS_SUCCESS if successful.
+    STATUS_IO_DEVICE_ERROR if the state could not be set.
+    STATUS_UNSUPPORTED if State is not supported by the device.
+
+--*/
+
+{
+    EXECUTION_CONTEXT_TYPE ContextType;
+    EFI_STATUS EfiStatus;
+
+    ContextType = CurrentExecutionContext->Type;
+    if (ContextType != ExecutionContextFirmware) {
+        //
+        // TODO: Translate addresses here.
+        // Need MmArchTranslateVirtualAddress().
+        //
+
+        BlpArchSwitchContext(ExecutionContextFirmware);
+    }
+
+    EfiStatus = Protocol->SetState(
+        Protocol,
+        KeyToggleState
+    );
+
+    if (ContextType != ExecutionContextFirmware) {
+        BlpArchSwitchContext(ContextType);
+    }
+
+    return EfiGetNtStatusCode(EfiStatus);
+}

BOOT/ENVIRON/LIB/EFI/efimm.c

@@ -23,172 +23,6 @@ Abstract:
 #define EFI_PAGE(NtPage) (((NtPage) << PAGE_SHIFT) >> EFI_PAGE_SHIFT)
 #define NT_PAGE(EfiPage) (((EfiPage) << EFI_PAGE_SHIFT) >> PAGE_SHIFT)
 
-NTSTATUS
-EfiGetMemoryMap (
-    IN OUT UINTN *MemoryMapSize,
-    IN OUT EFI_MEMORY_DESCRIPTOR *MemoryMap,
-    IN OUT UINTN *MapKey,
-    IN OUT UINTN *DescriptorSize,
-    IN OUT UINT32 *DescriptorVersion
-    )
-
-/*++
-
-Routine Description:
-
-    Wrapper for EFI_BOOT_SERVICES.GetMemoryMap(). Gets the firmware memory map
-    and places it into a buffer.
-
-Arguments:
-
-    MemoryMapSize - pointer to the size of the buffer.
-
-    MemoryMap - pointer to the buffer to store the memory map in.
-
-    MapKey - ponter to the memory map key.
-
-    DescriptorSize - pointer to the size of each memory map descriptor.
-
-    DescriptorVersion - pointer to the version of memory map descriptors.
-
-Return Value:
-
-    STATUS_SUCCESS if successful,
-    Other NTSTATUS value if an error occurs.
-
---*/
-
-{
-    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
-EfiAllocatePages (
-    IN EFI_ALLOCATE_TYPE Type,
-    IN EFI_MEMORY_TYPE MemoryType,
-    IN UINTN Pages,
-    IN OUT EFI_PHYSICAL_ADDRESS *Memory
-    )
-
-/*++
-
-Routine Description:
-
-    Wrapper for EFI_BOOT_SERVICES.AllocatePages(). Allocates contiguous pages
-    of physical memory.
-
-Arguments:
-
-    Type - the type of allocation.
-
-    MemoryType - the type of memory to allocate.
-
-    Pages - the number of pages to allocate.
-
-    Memory - pointer to a physical address of the allocation.
-
-Return Value:
-
-    STATUS_SUCCESS if successful,
-    Other NTSTATUS value if an error occurs.
-
---*/
-
-{
-    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
-EfiFreePages (
-    IN EFI_PHYSICAL_ADDRESS Memory,
-    IN UINTN Pages
-    )
-
-/*++
-
-Routine Description:
-
-    Wrapper for EFI_BOOT_SERVICES.FreePages(). Frees contiguous pages
-    of physical memory.
-
-Arguments:
-
-    Memory - physical address of the pages to be freed.
-
-    Pages - the number of pages to free.
-
-Return Value:
-
-    STATUS_SUCCESS if successful,
-    Other NTSTATUS value if an error occurs.
-
---*/
-
-{
-    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
 BlMmTranslateEfiMemoryType (
     IN EFI_MEMORY_TYPE EfiMemoryType

BOOT/ENVIRON/LIB/MM/mm.c

@@ -17,6 +17,56 @@ Abstract:
 #include "bootlib.h"
 #include "mm.h"
 
+ULONG MmTranslationType;
+
+NTSTATUS
+BlpMmDestroy (
+    IN ULONG Stage
+    )
+
+/*++
+
+Routine Description:
+
+    Cleans up after any actions performed by the memory manager.
+    After calling this, the memory manager can no longer be used.
+
+Arguments:
+
+    Stage - Which stage of cleanup to perform.
+        Stage 0: Unknown.
+        Stage 1: Destroy all MM modules.
+
+Return Value:
+
+    STATUS_SUCCESS if successful.
+
+--*/
+
+{
+    NTSTATUS Status, ExitStatus;
+
+    ExitStatus = STATUS_SUCCESS;
+    if (Stage == 1) {
+        Status = MmMdDestroy();
+        if (!NT_SUCCESS(Status)) {
+            ExitStatus = Status;
+        }
+
+        Status = MmPaDestroy(0);
+        if (!NT_SUCCESS(Status)) {
+            ExitStatus = Status;
+        }
+
+        Status = MmPaDestroy(1);
+        if (!NT_SUCCESS(Status)) {
+            ExitStatus = Status;
+        }
+    }
+
+    return ExitStatus;
+}
+
 NTSTATUS
 BlpMmInitializeConstraints (
     VOID
@@ -34,8 +84,7 @@ Arguments:
 
 Return Value:
 
-    STATUS_SUCCESS if successful,
+    STATUS_SUCCESS if successful.
-
 
 --*/
 
@@ -83,8 +132,8 @@ Return Value:
     // Check TranslationType.
     //
     if (
-        TranslationType > TRANSLATION_TYPE_MAX ||
+        TranslationType > TRANSLATION_TYPE_MAX
-        LibraryParameters->TranslationType > TRANSLATION_TYPE_MAX
+        || LibraryParameters->TranslationType > TRANSLATION_TYPE_MAX
     ) {
         DebugPrint(L"BlpMmInitialize(): TranslationType is invalid\r\n");
         return STATUS_INVALID_PARAMETER;
@@ -103,6 +152,8 @@ Return Value:
         return Status;
     }
 
+    MmTranslationType = LibraryParameters->TranslationType;
+
     //
     // TODO: Finish this routine.
     //

BOOT/ENVIRON/LIB/MM/mmmd.c

@@ -885,6 +885,35 @@ Return Value:
     return Descriptor;
 }
 
+NTSTATUS
+MmMdDestroy (
+    )
+
+/*++
+
+Routine Description:
+
+    Cleans up after any actions performed by the memory descriptor manager.
+    After calling this, the memory descriptor manager can no longer be used.
+
+Arguments:
+
+    None.
+
+Return Value:
+
+    STATUS_SUCCESS.
+
+--*/
+
+{
+    //
+    // TODO: Implement this routine.
+    //
+
+    return STATUS_SUCCESS;
+}
+
 VOID
 MmMdInitialize (
     IN ULONG Stage,

BOOT/ENVIRON/LIB/MM/mmpa.c

@@ -58,6 +58,40 @@ Return Value:
     Mdl->Type = MDL_TYPE_PHYSICAL;
 }
 
+NTSTATUS
+MmPaDestroy (
+    IN ULONG Stage
+    )
+
+/*++
+
+Routine Description:
+
+    Cleans up after any actions performed by the page allocator.
+    After calling this, the page allocator can no longer be used.
+
+Arguments:
+
+    Stage - Which stage of cleanup to perform.
+        Stage 0: Unknown.
+        Stage 1: Unknown.
+
+Return Value:
+
+    STATUS_SUCCESS.
+
+--*/
+
+{
+    (VOID)Stage;
+
+    //
+    // TODO: Implement this routine.
+    //
+
+    return STATUS_SUCCESS;
+}
+
 NTSTATUS
 MmPaInitialize (
     IN PBOOT_MEMORY_INFO MemoryInfo,

BOOT/ENVIRON/LIB/bootlib.c

@@ -78,7 +78,7 @@ Return Value:
     //
     Status = BlpFwInitialize(0, FirmwareData);
     if (!NT_SUCCESS(Status)) {
-        return Status;
+        goto Stage0Failed;
     }
 
     ConsolePrint(L"> Alcyone EFI Boot Manager\r\n");
@@ -86,7 +86,8 @@ Return Value:
 
     if (ApplicationEntry->Signature != BOOT_INIT_APPLICATION_ENTRY_SIGNATURE) {
         DebugPrint(L"InitializeLibrary(): ApplicationEntry Signature is invalid\r\n");
-        return STATUS_INVALID_PARAMETER_9;
+        Status = STATUS_INVALID_PARAMETER_9;
+        goto Stage0Failed;
     }
 
     //
@@ -104,14 +105,27 @@ Return Value:
 
     Status = BlpArchInitialize(0);
     if (!NT_SUCCESS(Status)) {
-        return Status;
+        goto Stage0Failed;
     }
 
     Status = BlpMmInitialize(MemoryInfo, ApplicationParameters->TranslationType, LibraryParameters);
     if (!NT_SUCCESS(Status)) {
-        return Status;
+        goto Stage0Failed;
     }
 
+    Status = BlpFwInitialize(1, FirmwareData);
+    if (!NT_SUCCESS(Status)) {
+        goto Stage1Failed;
+    }
+
+    Status = BlpArchInitialize(1);
+    if (!NT_SUCCESS(Status)) {
+        goto Stage1Failed;
+    }
+
+Stage1Failed:
+    BlpMmDestroy(1);
+Stage0Failed:
     return STATUS_SUCCESS;
 }
 
@@ -160,10 +174,25 @@ Arguments:
 
 Return Value:
 
-    STATUS_SUCCESS.
+    STATUS_SUCCESS if successful.
+    Error code on failure.
 
 --*/
 
 {
-    return STATUS_SUCCESS;
+    NTSTATUS Status, ExitStatus;
+
+    ExitStatus = STATUS_SUCCESS;
+
+    Status = BlpMmDestroy(0);
+    if (!NT_SUCCESS(Status)) {
+        ExitStatus = Status;
+    }
+
+    Status = BlpMmDestroy(1);
+    if (!NT_SUCCESS(Status)) {
+        ExitStatus = Status;
+    }
+
+    return ExitStatus;
 }