1 files changed, 11 insertions, 0 deletions
diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/SmmCpuMemoryManagement.c b/UefiCpuPkg/PiSmmCpuDxeSmm/SmmCpuMemoryManagement.c
index 8888c26ffd..8ed4940076 100644
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/SmmCpuMemoryManagement.c
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/SmmCpuMemoryManagement.c
@@ -1033,6 +1033,8 @@ SetMemMapAttributes (
   IA32_MAP_ENTRY                        *Map;
   UINTN                                 Count;
   UINT64                                MemoryAttribute;
+  BOOLEAN                               WpEnabled;
+  BOOLEAN                               CetEnabled;
 
   SmmGetSystemConfigurationTable (&gEdkiiPiSmmMemoryAttributesTableGuid, (VOID **)&MemoryAttributesTable);
   if (MemoryAttributesTable == NULL) {
@@ -1077,6 +1079,8 @@ SetMemMapAttributes (
 
   ASSERT_RETURN_ERROR (Status);
 
+  DisableReadOnlyPageWriteProtect (&WpEnabled, &CetEnabled);
+
   MemoryMap = MemoryMapStart;
   for (Index = 0; Index < MemoryMapEntryCount; Index++) {
     DEBUG ((DEBUG_VERBOSE, "SetAttribute: Memory Entry - 0x%lx, 0x%x\n", MemoryMap->PhysicalStart, MemoryMap->NumberOfPages));
@@ -1105,6 +1109,7 @@ SetMemMapAttributes (
     MemoryMap = NEXT_MEMORY_DESCRIPTOR (MemoryMap, DescriptorSize);
   }
 
+  EnableReadOnlyPageWriteProtect (WpEnabled, CetEnabled);
   FreePool (Map);
 
   PatchSmmSaveStateMap ();
@@ -1411,11 +1416,15 @@ SetUefiMemMapAttributes (
   UINTN                  MemoryMapEntryCount;
   UINTN                  Index;
   EFI_MEMORY_DESCRIPTOR  *Entry;
+  BOOLEAN                WpEnabled;
+  BOOLEAN                CetEnabled;
 
   PERF_FUNCTION_BEGIN ();
 
   DEBUG ((DEBUG_INFO, "SetUefiMemMapAttributes\n"));
 
+  DisableReadOnlyPageWriteProtect (&WpEnabled, &CetEnabled);
+
   if (mUefiMemoryMap != NULL) {
     MemoryMapEntryCount = mUefiMemoryMapSize/mUefiDescriptorSize;
     MemoryMap           = mUefiMemoryMap;
@@ -1494,6 +1503,8 @@ SetUefiMemMapAttributes (
     }
   }
 
+  EnableReadOnlyPageWriteProtect (WpEnabled, CetEnabled);
+
   //
   // Do not free mUefiMemoryAttributesTable, it will be checked in IsSmmCommBufferForbiddenAddress().
   //