Add Memory Protection Nonstop Protocol

MdePkg/Include/Protocol/MemoryProtectionNonstopMode.h

@@ -0,0 +1,63 @@
+/** @file
+  Provides an interface for recovering from page faults without rebooting.
+
+  NOTE: This API should not be installed if the platform does not want to allow
+        page faults to be cleared. For most cases, this API should only be installed
+        on platforms in manufacturing mode to test the function of memory protections.
+
+  Copyright (c) Microsoft Corporation.
+  SPDX-License-Identifier: BSD-2-Clause-Patent
+**/
+
+#ifndef MEMORY_PROTECTION_NONSTOP_MODE_H_
+#define MEMORY_PROTECTION_NONSTOP_MODE_H_
+
+#define MEMORY_PROTECTION_NONSTOP_MODE_PROTOCOL_GUID \
+  { \
+    {0x896306BA, 0x1722, 0x414D, {0x93, 0xCF, 0x14, 0x24, 0x4F, 0x5E, 0x31, 0x5D } \
+  }
+
+typedef struct _MEMORY_PROTECTION_NONSTOP_MODE_PROTOCOL MEMORY_PROTECTION_NONSTOP_MODE_PROTOCOL;
+
+/**
+  Marks the faulting pages as present, R/W, and executable so execution can resume.
+
+  @param[in] ExceptionType  Exception type.
+  @param[in] SystemContext  Pointer to EFI_SYSTEM_CONTEXT.
+
+  @retval EFI_SUCCESS              Page attributes were cleared.
+  @retval RETURN_ACCESS_DENIED     The attributes for the page could not be modified.
+  @retval RETURN_OUT_OF_RESOURCES  There are not enough system resources to modify the attributes of
+                                   the memory resource range.
+**/
+typedef
+EFI_STATUS
+(EFIAPI *CLEAR_PAGE_FAULT)(
+  IN EFI_EXCEPTION_TYPE   InterruptType,
+  IN EFI_SYSTEM_CONTEXT   SystemContext
+  );
+
+/**
+  Restores the page table attributes of a page whose attributes were cleared via a call to
+  CLEAR_PAGE_FAULT.
+
+**/
+typedef
+EFI_STATUS
+(EFIAPI *RESET_PAGE_ATTRIBUTES)(
+  VOID
+  );
+
+///
+/// EFI Nonstop Protocol provides services to clear a page fault
+/// by reseting the faulting page's attributes and to restore
+/// the previously faulting page to its original attributes.
+///
+struct _MEMORY_PROTECTION_NONSTOP_MODE_PROTOCOL {
+  CLEAR_PAGE_FAULT         ClearPageFault;
+  RESET_PAGE_ATTRIBUTES    ResetPageAttributes;
+};
+
+extern EFI_GUID  gMemoryProtectionNonstopModeProtocolGuid;
+
+#endif

MdePkg/MdePkg.ci.yaml

@@ -95,7 +95,9 @@
             "UnitTestFrameworkPkg/UnitTestFrameworkPkg.dec"
         ],
         # For UEFI shell based apps
-        "AcceptableDependencies-UEFI_APPLICATION":[],
+        "AcceptableDependencies-UEFI_APPLICATION":[
+            "UnitTestFrameworkPkg/UnitTestFrameworkPkg.dec" # TCBZ3519 MU_CHANGE
+        ],
         "IgnoreInf": []
     },
 

MdePkg/MdePkg.dec

@@ -1991,6 +1991,11 @@
 #   0x80000002 | Reserved bits must be set to zero.
 #   0x80000003 | Incorrect progress code provided.
 #
+  # MU_CHANGE START
+  ## Include/Protocol/MemoryProtectionNonstopMode.h
+  # 896306BA-1722-414D-93CF-14244F5E315D
+  gMemoryProtectionNonstopModeProtocolGuid = {0x896306BA, 0x1722, 0x414D, {0x93, 0xCF, 0x14, 0x24, 0x4F, 0x5E, 0x31, 0x5D }}
+  # MU_CHANGE END
 
 [PcdsFeatureFlag]
   ## Indicates if the component name protocol will be installed.<BR><BR>

UefiCpuPkg/CpuDxe/CpuDxe.c

@@ -9,7 +9,7 @@
 #include "CpuDxe.h"
 #include "CpuMp.h"
 #include "CpuPageTable.h"
-
+#include <Library/DeviceStateLib.h> // MU_CHANGE
 #define CPU_INTERRUPT_NUM  256
 
 //
@@ -1283,6 +1283,12 @@ InitializeCpu (
   ASSERT_EFI_ERROR (Status);
 
   InitializeMpSupport ();
+  // MU_CHANGE START
+  if ((GetDeviceState () & DEVICE_STATE_UNIT_TEST_MODE) != 0) {
+    InstallMemoryProtectionNonstopModeProtocol ();
+  }
+
+  // MU_CHANGE END
 
   return Status;
 }

UefiCpuPkg/CpuDxe/CpuDxe.inf

@@ -42,6 +42,7 @@
   TimerLib
   PeCoffGetEntryPointLib
   DxeMemoryProtectionHobLib ## MU_CHANGE
+  DeviceStateLib            ## MU_CHANGE
 
 [Sources]
   CpuDxe.c
@@ -66,6 +67,7 @@
   gEfiMpServiceProtocolGuid                     ## PRODUCES
   gEfiMemoryAttributeProtocolGuid               ## TCBZ3519 MU_CHANGE PRODUCES
   gEfiSmmBase2ProtocolGuid                      ## SOMETIMES_CONSUMES
+  gMemoryProtectionNonstopModeProtocolGuid      ## MU_CHANGE: PRODUCES
 
 [Guids]
   gIdleLoopEventGuid                            ## CONSUMES           ## Event

UefiCpuPkg/CpuDxe/CpuPageTable.c

@@ -98,6 +98,12 @@ EFI_SMM_BASE2_PROTOCOL         *mSmmBase2 = NULL;
 UINTN  *mPFEntryCount;
 UINT64                    *(*mLastPFEntryPointer)[MAX_PF_ENTRY_COUNT];
 
+// MU_CHANGE START
+UINTN   mPFCount = 0;
+UINT64  mPageFaultAddresses[MAX_PF_ENTRY_COUNT];
+UINT64  mPageFaultAttributes[MAX_PF_ENTRY_COUNT];
+// MU_CHANGE END
+
 /**
  Check if current execution environment is in SMM mode or not, via
  EFI_SMM_BASE2_PROTOCOL.
@@ -1714,6 +1720,134 @@ EFI_MEMORY_ATTRIBUTE_PROTOCOL  mMemoryAttributeProtocol = {
   EfiClearMemoryAttributes,
 };
 
+// MU_CHANGE START
+
+/**
+  Marks the faulting pages as present, R/W, and executable so execution can resume.
+
+  @param[in] ExceptionType  Exception type.
+  @param[in] SystemContext  Pointer to EFI_SYSTEM_CONTEXT.
+
+  @retval EFI_SUCCESS              Page attributes were cleared.
+  @retval RETURN_ACCESS_DENIED     The attributes for the page could not be modified.
+  @retval RETURN_OUT_OF_RESOURCES  There are not enough system resources to modify the attributes of
+                                   the memory resource range.
+**/
+EFI_STATUS
+EFIAPI
+ClearPageFault (
+  IN EFI_EXCEPTION_TYPE  ExceptionType,
+  IN EFI_SYSTEM_CONTEXT  SystemContext
+  )
+{
+  EFI_STATUS                     Status;
+  UINT64                         PFAddress;
+  PAGE_TABLE_LIB_PAGING_CONTEXT  PagingContext;
+  PAGE_ATTRIBUTE                 PageAttribute;
+  UINT64                         OldAttributes;
+  UINT64                         NewAttributes;
+  UINT64                         *PageEntry;
+  UINTN                          CpuIndex;
+  UINTN                          PageNumber;
+
+  PFAddress = AsmReadCr2 () & ~EFI_PAGE_MASK;
+  MpInitLibWhoAmI (&CpuIndex);
+  GetCurrentPagingContext (&PagingContext);
+  //
+  // Memory operation cross page boundary, like "rep mov" instruction, will
+  // cause infinite loop. We have to make sure that current page and the page
+  // followed are both in PRESENT state.
+  //
+  PageNumber = 2;
+  while (PageNumber > 0) {
+    PageEntry = GetPageTableEntry (&PagingContext, PFAddress, &PageAttribute);
+    ASSERT (PageEntry != NULL);
+    if (PageEntry != NULL) {
+      OldAttributes = GetAttributesFromPageEntry (PageEntry);
+      if (((OldAttributes & EFI_MEMORY_RP) != 0) || ((OldAttributes & EFI_MEMORY_XP) != 0)) {
+        NewAttributes  = OldAttributes;
+        NewAttributes &= ~(EFI_MEMORY_RP | EFI_MEMORY_XP);
+        DEBUG ((DEBUG_INFO, "%a - Clearing page fault at address: 0x%x\n", __FUNCTION__, PFAddress));
+        Status = AssignMemoryPageAttributes (
+                   &PagingContext,
+                   PFAddress,
+                   EFI_PAGE_SIZE,
+                   NewAttributes,
+                   NULL
+                   );
+        if (!EFI_ERROR (Status)) {
+          if (mPFCount < ARRAY_SIZE (mPageFaultAddresses)) {
+            mPageFaultAttributes[mPFCount]  = OldAttributes;
+            mPageFaultAddresses[mPFCount++] = PFAddress;
+          }
+        } else {
+          DEBUG ((DEBUG_INFO, "%a - Failed to clear page fault at address: 0x%x\n", __FUNCTION__, PFAddress));
+        }
+      }
+    }
+
+    PFAddress += EFI_PAGE_SIZE;
+    --PageNumber;
+  }
+
+  return EFI_SUCCESS;
+}
+
+/**
+  Restores the page table attributes of a page whose attributes were cleared via a call to
+  CLEAR_PAGE_FAULT.
+
+  @retval EFI_SUCCESS              Page attributes were restored.
+  @retval EFI_DEVICE_ERROR         The attributes could not be modified for one or more of the pages being restored.
+**/
+EFI_STATUS
+EFIAPI
+ResetPageAttributes (
+  VOID
+  )
+{
+  EFI_STATUS                     Status = EFI_SUCCESS;
+  UINT64                         PFAddress;
+  PAGE_TABLE_LIB_PAGING_CONTEXT  PagingContext;
+  PAGE_ATTRIBUTE                 PageAttribute;
+  UINT64                         Attributes;
+  UINT64                         *PageEntry;
+
+  while (mPFCount > 0) {
+    PFAddress = mPageFaultAddresses[--mPFCount];
+    GetCurrentPagingContext (&PagingContext);
+
+    PageEntry = GetPageTableEntry (&PagingContext, PFAddress, &PageAttribute);
+
+    if (PageEntry != NULL) {
+      Attributes = mPageFaultAttributes[mPFCount];
+      DEBUG ((DEBUG_INFO, "%a - Restoring page attributes at address: 0x%x\n", __FUNCTION__, PFAddress));
+      if (EFI_ERROR (
+            AssignMemoryPageAttributes (
+              &PagingContext,
+              PFAddress,
+              EFI_PAGE_SIZE,
+              Attributes,
+              NULL
+              )
+            ))
+      {
+        DEBUG ((DEBUG_ERROR, "%a - Unable to set memory attributes at address: 0x%x\n", __FUNCTION__, PFAddress));
+        Status = EFI_DEVICE_ERROR;
+      }
+
+      mPageFaultAddresses[mPFCount] = 0;
+    }
+  }
+
+  return Status;
+}
+
+MEMORY_PROTECTION_NONSTOP_MODE_PROTOCOL  mMemoryNonstopModeProtocol = {
+  ClearPageFault,
+  ResetPageAttributes
+};
+
 /**
   Install Memory Attribute Protocol.
 **/
@@ -1734,3 +1868,26 @@ InstallEfiMemoryAttributeProtocol (
 }
 
 // TCBZ3519 MU_CHANGE END
+
+// MU_CHANGE START
+
+/**
+  Install Memory Protection Nonstop Protocol.
+**/
+VOID
+InstallMemoryProtectionNonstopModeProtocol (
+  VOID
+  )
+{
+  EFI_STATUS  Status;
+
+  Status = gBS->InstallMultipleProtocolInterfaces (
+                  &mCpuHandle,
+                  &gMemoryProtectionNonstopModeProtocolGuid,
+                  &mMemoryNonstopModeProtocol,
+                  NULL
+                  );
+  ASSERT_EFI_ERROR (Status);
+}
+
+// MU_CHANGE END

UefiCpuPkg/CpuDxe/CpuPageTable.h

@@ -11,6 +11,7 @@
 
 #include <IndustryStandard/PeImage.h>
 #include <Protocol/MemoryAttribute.h>             // TCBZ3519 MU_CHANGE
+#include <Protocol/MemoryProtectionNonstopMode.h> // MU_CHANGE
 
 #define PAGE_TABLE_LIB_PAGING_CONTEXT_IA32_X64_ATTRIBUTES_PSE              BIT0
 #define PAGE_TABLE_LIB_PAGING_CONTEXT_IA32_X64_ATTRIBUTES_PAE              BIT1
@@ -166,4 +167,16 @@ InstallEfiMemoryAttributeProtocol (
 
 // TCBZ3519 MU_CHANGE END
 
+// MU_CHANGE START
+
+/**
+ Install the Memory Protection Nonstop Protocol
+**/
+VOID
+InstallMemoryProtectionNonstopModeProtocol (
+  VOID
+  );
+
+// MU_CHANGE END
+
 #endif

UefiCpuPkg/UefiCpuPkg.dsc

@@ -83,6 +83,7 @@
 ##MSCHANGE Begin
   DxeMemoryProtectionHobLib|MdeModulePkg/Library/MemoryProtectionHobLibNull/DxeMemoryProtectionHobLibNull.inf
   MmMemoryProtectionHobLib|MdeModulePkg/Library/MemoryProtectionHobLibNull/MmMemoryProtectionHobLibNull.inf
+  DeviceStateLib|MdeModulePkg/Library/DeviceStateLib/DeviceStateLib.inf
   BaseBinSecurityLib|MdePkg/Library/BaseBinSecurityLibNull/BaseBinSecurityLibNull.inf
 [LibraryClasses.X64, LibraryClasses.IA32]
   HwResetSystemLib|MdeModulePkg/Library/BaseResetSystemLibNull/BaseResetSystemLibNull.inf