]> git.proxmox.com Git - mirror_edk2.git/blobdiff - UefiCpuPkg/PiSmmCpuDxeSmm/SmmCpuMemoryManagement.c
UefiCpuPkg: Move AsmRelocateApLoopStart from Mpfuncs.nasm to AmdSev.nasm
[mirror_edk2.git] / UefiCpuPkg / PiSmmCpuDxeSmm / SmmCpuMemoryManagement.c
index 3ad5256f1e03d64b4a93bfe7eee74de80c5173c1..834a756061da80b26ce2086740c8689fd6946d2f 100644 (file)
 /** @file\r
 \r
-Copyright (c) 2016 - 2017, Intel Corporation. All rights reserved.<BR>\r
-This program and the accompanying materials\r
-are licensed and made available under the terms and conditions of the BSD License\r
-which accompanies this distribution.  The full text of the license may be found at\r
-http://opensource.org/licenses/bsd-license.php\r
-\r
-THE PROGRAM IS DISTRIBUTED UNDER THE BSD LICENSE ON AN "AS IS" BASIS,\r
-WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.\r
+Copyright (c) 2016 - 2019, Intel Corporation. All rights reserved.<BR>\r
+SPDX-License-Identifier: BSD-2-Clause-Patent\r
 \r
 **/\r
 \r
 #include "PiSmmCpuDxeSmm.h"\r
 \r
-#define NEXT_MEMORY_DESCRIPTOR(MemoryDescriptor, Size) \\r
-  ((EFI_MEMORY_DESCRIPTOR *)((UINT8 *)(MemoryDescriptor) + (Size)))\r
+//\r
+// attributes for reserved memory before it is promoted to system memory\r
+//\r
+#define EFI_MEMORY_PRESENT      0x0100000000000000ULL\r
+#define EFI_MEMORY_INITIALIZED  0x0200000000000000ULL\r
+#define EFI_MEMORY_TESTED       0x0400000000000000ULL\r
 \r
 #define PREVIOUS_MEMORY_DESCRIPTOR(MemoryDescriptor, Size) \\r
   ((EFI_MEMORY_DESCRIPTOR *)((UINT8 *)(MemoryDescriptor) - (Size)))\r
 \r
-EFI_MEMORY_DESCRIPTOR *mUefiMemoryMap;\r
-UINTN                 mUefiMemoryMapSize;\r
-UINTN                 mUefiDescriptorSize;\r
+EFI_MEMORY_DESCRIPTOR  *mUefiMemoryMap;\r
+UINTN                  mUefiMemoryMapSize;\r
+UINTN                  mUefiDescriptorSize;\r
+\r
+EFI_GCD_MEMORY_SPACE_DESCRIPTOR  *mGcdMemSpace       = NULL;\r
+UINTN                            mGcdMemNumberOfDesc = 0;\r
+\r
+EFI_MEMORY_ATTRIBUTES_TABLE  *mUefiMemoryAttributesTable = NULL;\r
 \r
-PAGE_ATTRIBUTE_TABLE mPageAttributeTable[] = {\r
-  {Page4K,  SIZE_4KB, PAGING_4K_ADDRESS_MASK_64},\r
-  {Page2M,  SIZE_2MB, PAGING_2M_ADDRESS_MASK_64},\r
-  {Page1G,  SIZE_1GB, PAGING_1G_ADDRESS_MASK_64},\r
+PAGE_ATTRIBUTE_TABLE  mPageAttributeTable[] = {\r
+  { Page4K, SIZE_4KB, PAGING_4K_ADDRESS_MASK_64 },\r
+  { Page2M, SIZE_2MB, PAGING_2M_ADDRESS_MASK_64 },\r
+  { Page1G, SIZE_1GB, PAGING_1G_ADDRESS_MASK_64 },\r
 };\r
 \r
+BOOLEAN  mIsShadowStack      = FALSE;\r
+BOOLEAN  m5LevelPagingNeeded = FALSE;\r
+\r
+//\r
+// Global variable to keep track current available memory used as page table.\r
+//\r
+PAGE_TABLE_POOL  *mPageTablePool = NULL;\r
+\r
+//\r
+// If memory used by SMM page table has been mareked as ReadOnly.\r
+//\r
+BOOLEAN  mIsReadOnlyPageTable = FALSE;\r
+\r
 /**\r
-  Return page table base.\r
+  Initialize a buffer pool for page table use only.\r
 \r
-  @return page table base.\r
+  To reduce the potential split operation on page table, the pages reserved for\r
+  page table should be allocated in the times of PAGE_TABLE_POOL_UNIT_PAGES and\r
+  at the boundary of PAGE_TABLE_POOL_ALIGNMENT. So the page pool is always\r
+  initialized with number of pages greater than or equal to the given PoolPages.\r
+\r
+  Once the pages in the pool are used up, this method should be called again to\r
+  reserve at least another PAGE_TABLE_POOL_UNIT_PAGES. But usually this won't\r
+  happen in practice.\r
+\r
+  @param PoolPages  The least page number of the pool to be created.\r
+\r
+  @retval TRUE    The pool is initialized successfully.\r
+  @retval FALSE   The memory is out of resource.\r
 **/\r
-UINTN\r
-GetPageTableBase (\r
-  VOID\r
+BOOLEAN\r
+InitializePageTablePool (\r
+  IN UINTN  PoolPages\r
   )\r
 {\r
-  return (AsmReadCr3 () & PAGING_4K_ADDRESS_MASK_64);\r
+  VOID      *Buffer;\r
+  BOOLEAN   CetEnabled;\r
+  BOOLEAN   WpEnabled;\r
+  IA32_CR0  Cr0;\r
+\r
+  //\r
+  // Always reserve at least PAGE_TABLE_POOL_UNIT_PAGES, including one page for\r
+  // header.\r
+  //\r
+  PoolPages += 1;   // Add one page for header.\r
+  PoolPages  = ((PoolPages - 1) / PAGE_TABLE_POOL_UNIT_PAGES + 1) *\r
+               PAGE_TABLE_POOL_UNIT_PAGES;\r
+  Buffer = AllocateAlignedPages (PoolPages, PAGE_TABLE_POOL_ALIGNMENT);\r
+  if (Buffer == NULL) {\r
+    DEBUG ((DEBUG_ERROR, "ERROR: Out of aligned pages\r\n"));\r
+    return FALSE;\r
+  }\r
+\r
+  //\r
+  // Link all pools into a list for easier track later.\r
+  //\r
+  if (mPageTablePool == NULL) {\r
+    mPageTablePool           = Buffer;\r
+    mPageTablePool->NextPool = mPageTablePool;\r
+  } else {\r
+    ((PAGE_TABLE_POOL *)Buffer)->NextPool = mPageTablePool->NextPool;\r
+    mPageTablePool->NextPool              = Buffer;\r
+    mPageTablePool                        = Buffer;\r
+  }\r
+\r
+  //\r
+  // Reserve one page for pool header.\r
+  //\r
+  mPageTablePool->FreePages = PoolPages - 1;\r
+  mPageTablePool->Offset    = EFI_PAGES_TO_SIZE (1);\r
+\r
+  //\r
+  // If page table memory has been marked as RO, mark the new pool pages as read-only.\r
+  //\r
+  if (mIsReadOnlyPageTable) {\r
+    CetEnabled = ((AsmReadCr4 () & CR4_CET_ENABLE) != 0) ? TRUE : FALSE;\r
+    Cr0.UintN  = AsmReadCr0 ();\r
+    WpEnabled  = (Cr0.Bits.WP != 0) ? TRUE : FALSE;\r
+    if (WpEnabled) {\r
+      if (CetEnabled) {\r
+        //\r
+        // CET must be disabled if WP is disabled. Disable CET before clearing CR0.WP.\r
+        //\r
+        DisableCet ();\r
+      }\r
+\r
+      Cr0.Bits.WP = 0;\r
+      AsmWriteCr0 (Cr0.UintN);\r
+    }\r
+\r
+    SmmSetMemoryAttributes ((EFI_PHYSICAL_ADDRESS)(UINTN)Buffer, EFI_PAGES_TO_SIZE (PoolPages), EFI_MEMORY_RO);\r
+    if (WpEnabled) {\r
+      Cr0.UintN   = AsmReadCr0 ();\r
+      Cr0.Bits.WP = 1;\r
+      AsmWriteCr0 (Cr0.UintN);\r
+\r
+      if (CetEnabled) {\r
+        //\r
+        // re-enable CET.\r
+        //\r
+        EnableCet ();\r
+      }\r
+    }\r
+  }\r
+\r
+  return TRUE;\r
+}\r
+\r
+/**\r
+  This API provides a way to allocate memory for page table.\r
+\r
+  This API can be called more once to allocate memory for page tables.\r
+\r
+  Allocates the number of 4KB pages of type EfiRuntimeServicesData and returns a pointer to the\r
+  allocated buffer.  The buffer returned is aligned on a 4KB boundary.  If Pages is 0, then NULL\r
+  is returned.  If there is not enough memory remaining to satisfy the request, then NULL is\r
+  returned.\r
+\r
+  @param  Pages                 The number of 4 KB pages to allocate.\r
+\r
+  @return A pointer to the allocated buffer or NULL if allocation fails.\r
+\r
+**/\r
+VOID *\r
+AllocatePageTableMemory (\r
+  IN UINTN  Pages\r
+  )\r
+{\r
+  VOID  *Buffer;\r
+\r
+  if (Pages == 0) {\r
+    return NULL;\r
+  }\r
+\r
+  //\r
+  // Renew the pool if necessary.\r
+  //\r
+  if ((mPageTablePool == NULL) ||\r
+      (Pages > mPageTablePool->FreePages))\r
+  {\r
+    if (!InitializePageTablePool (Pages)) {\r
+      return NULL;\r
+    }\r
+  }\r
+\r
+  Buffer = (UINT8 *)mPageTablePool + mPageTablePool->Offset;\r
+\r
+  mPageTablePool->Offset    += EFI_PAGES_TO_SIZE (Pages);\r
+  mPageTablePool->FreePages -= Pages;\r
+\r
+  return Buffer;\r
 }\r
 \r
 /**\r
@@ -55,11 +198,13 @@ PageAttributeToLength (
   )\r
 {\r
   UINTN  Index;\r
-  for (Index = 0; Index < sizeof(mPageAttributeTable)/sizeof(mPageAttributeTable[0]); Index++) {\r
+\r
+  for (Index = 0; Index < sizeof (mPageAttributeTable)/sizeof (mPageAttributeTable[0]); Index++) {\r
     if (PageAttribute == mPageAttributeTable[Index].Attribute) {\r
       return (UINTN)mPageAttributeTable[Index].Length;\r
     }\r
   }\r
+\r
   return 0;\r
 }\r
 \r
@@ -76,44 +221,64 @@ PageAttributeToMask (
   )\r
 {\r
   UINTN  Index;\r
-  for (Index = 0; Index < sizeof(mPageAttributeTable)/sizeof(mPageAttributeTable[0]); Index++) {\r
+\r
+  for (Index = 0; Index < sizeof (mPageAttributeTable)/sizeof (mPageAttributeTable[0]); Index++) {\r
     if (PageAttribute == mPageAttributeTable[Index].Attribute) {\r
       return (UINTN)mPageAttributeTable[Index].AddressMask;\r
     }\r
   }\r
+\r
   return 0;\r
 }\r
 \r
 /**\r
   Return page table entry to match the address.\r
 \r
-  @param[in]   Address          The address to be checked.\r
-  @param[out]  PageAttributes   The page attribute of the page entry.\r
+  @param[in]   PageTableBase      The page table base.\r
+  @param[in]   Enable5LevelPaging If PML5 paging is enabled.\r
+  @param[in]   Address            The address to be checked.\r
+  @param[out]  PageAttributes     The page attribute of the page entry.\r
 \r
   @return The page entry.\r
 **/\r
 VOID *\r
 GetPageTableEntry (\r
-  IN  PHYSICAL_ADDRESS                  Address,\r
-  OUT PAGE_ATTRIBUTE                    *PageAttribute\r
+  IN  UINTN             PageTableBase,\r
+  IN  BOOLEAN           Enable5LevelPaging,\r
+  IN  PHYSICAL_ADDRESS  Address,\r
+  OUT PAGE_ATTRIBUTE    *PageAttribute\r
   )\r
 {\r
-  UINTN                 Index1;\r
-  UINTN                 Index2;\r
-  UINTN                 Index3;\r
-  UINTN                 Index4;\r
-  UINT64                *L1PageTable;\r
-  UINT64                *L2PageTable;\r
-  UINT64                *L3PageTable;\r
-  UINT64                *L4PageTable;\r
-\r
+  UINTN   Index1;\r
+  UINTN   Index2;\r
+  UINTN   Index3;\r
+  UINTN   Index4;\r
+  UINTN   Index5;\r
+  UINT64  *L1PageTable;\r
+  UINT64  *L2PageTable;\r
+  UINT64  *L3PageTable;\r
+  UINT64  *L4PageTable;\r
+  UINT64  *L5PageTable;\r
+\r
+  Index5 = ((UINTN)RShiftU64 (Address, 48)) & PAGING_PAE_INDEX_MASK;\r
   Index4 = ((UINTN)RShiftU64 (Address, 39)) & PAGING_PAE_INDEX_MASK;\r
   Index3 = ((UINTN)Address >> 30) & PAGING_PAE_INDEX_MASK;\r
   Index2 = ((UINTN)Address >> 21) & PAGING_PAE_INDEX_MASK;\r
   Index1 = ((UINTN)Address >> 12) & PAGING_PAE_INDEX_MASK;\r
 \r
-  if (sizeof(UINTN) == sizeof(UINT64)) {\r
-    L4PageTable = (UINT64 *)GetPageTableBase ();\r
+  if (sizeof (UINTN) == sizeof (UINT64)) {\r
+    if (Enable5LevelPaging) {\r
+      L5PageTable = (UINT64 *)PageTableBase;\r
+      if (L5PageTable[Index5] == 0) {\r
+        *PageAttribute = PageNone;\r
+        return NULL;\r
+      }\r
+\r
+      L4PageTable = (UINT64 *)(UINTN)(L5PageTable[Index5] & ~mAddressEncMask & PAGING_4K_ADDRESS_MASK_64);\r
+    } else {\r
+      L4PageTable = (UINT64 *)PageTableBase;\r
+    }\r
+\r
     if (L4PageTable[Index4] == 0) {\r
       *PageAttribute = PageNone;\r
       return NULL;\r
@@ -121,12 +286,14 @@ GetPageTableEntry (
 \r
     L3PageTable = (UINT64 *)(UINTN)(L4PageTable[Index4] & ~mAddressEncMask & PAGING_4K_ADDRESS_MASK_64);\r
   } else {\r
-    L3PageTable = (UINT64 *)GetPageTableBase ();\r
+    L3PageTable = (UINT64 *)PageTableBase;\r
   }\r
+\r
   if (L3PageTable[Index3] == 0) {\r
     *PageAttribute = PageNone;\r
     return NULL;\r
   }\r
+\r
   if ((L3PageTable[Index3] & IA32_PG_PS) != 0) {\r
     // 1G\r
     *PageAttribute = Page1G;\r
@@ -138,6 +305,7 @@ GetPageTableEntry (
     *PageAttribute = PageNone;\r
     return NULL;\r
   }\r
+\r
   if ((L2PageTable[Index2] & IA32_PG_PS) != 0) {\r
     // 2M\r
     *PageAttribute = Page2M;\r
@@ -150,6 +318,7 @@ GetPageTableEntry (
     *PageAttribute = PageNone;\r
     return NULL;\r
   }\r
+\r
   *PageAttribute = Page4K;\r
   return &L1PageTable[Index1];\r
 }\r
@@ -163,20 +332,24 @@ GetPageTableEntry (
 **/\r
 UINT64\r
 GetAttributesFromPageEntry (\r
-  IN  UINT64                            *PageEntry\r
+  IN  UINT64  *PageEntry\r
   )\r
 {\r
   UINT64  Attributes;\r
+\r
   Attributes = 0;\r
   if ((*PageEntry & IA32_PG_P) == 0) {\r
     Attributes |= EFI_MEMORY_RP;\r
   }\r
+\r
   if ((*PageEntry & IA32_PG_RW) == 0) {\r
     Attributes |= EFI_MEMORY_RO;\r
   }\r
+\r
   if ((*PageEntry & IA32_PG_NX) != 0) {\r
     Attributes |= EFI_MEMORY_XP;\r
   }\r
+\r
   return Attributes;\r
 }\r
 \r
@@ -190,17 +363,17 @@ GetAttributesFromPageEntry (
 **/\r
 VOID\r
 ConvertPageEntryAttribute (\r
-  IN  UINT64                            *PageEntry,\r
-  IN  UINT64                            Attributes,\r
-  IN  BOOLEAN                           IsSet,\r
-  OUT BOOLEAN                           *IsModified\r
+  IN  UINT64   *PageEntry,\r
+  IN  UINT64   Attributes,\r
+  IN  BOOLEAN  IsSet,\r
+  OUT BOOLEAN  *IsModified\r
   )\r
 {\r
   UINT64  CurrentPageEntry;\r
   UINT64  NewPageEntry;\r
 \r
   CurrentPageEntry = *PageEntry;\r
-  NewPageEntry = CurrentPageEntry;\r
+  NewPageEntry     = CurrentPageEntry;\r
   if ((Attributes & EFI_MEMORY_RP) != 0) {\r
     if (IsSet) {\r
       NewPageEntry &= ~(UINT64)IA32_PG_P;\r
@@ -208,13 +381,26 @@ ConvertPageEntryAttribute (
       NewPageEntry |= IA32_PG_P;\r
     }\r
   }\r
+\r
   if ((Attributes & EFI_MEMORY_RO) != 0) {\r
     if (IsSet) {\r
       NewPageEntry &= ~(UINT64)IA32_PG_RW;\r
+      if (mIsShadowStack) {\r
+        // Environment setup\r
+        // ReadOnly page need set Dirty bit for shadow stack\r
+        NewPageEntry |= IA32_PG_D;\r
+        // Clear user bit for supervisor shadow stack\r
+        NewPageEntry &= ~(UINT64)IA32_PG_U;\r
+      } else {\r
+        // Runtime update\r
+        // Clear dirty bit for non shadow stack, to protect RO page.\r
+        NewPageEntry &= ~(UINT64)IA32_PG_D;\r
+      }\r
     } else {\r
       NewPageEntry |= IA32_PG_RW;\r
     }\r
   }\r
+\r
   if ((Attributes & EFI_MEMORY_XP) != 0) {\r
     if (mXdSupported) {\r
       if (IsSet) {\r
@@ -224,6 +410,7 @@ ConvertPageEntryAttribute (
       }\r
     }\r
   }\r
+\r
   *PageEntry = NewPageEntry;\r
   if (CurrentPageEntry != NewPageEntry) {\r
     *IsModified = TRUE;\r
@@ -246,13 +433,13 @@ ConvertPageEntryAttribute (
 **/\r
 PAGE_ATTRIBUTE\r
 NeedSplitPage (\r
-  IN  PHYSICAL_ADDRESS                  BaseAddress,\r
-  IN  UINT64                            Length,\r
-  IN  UINT64                            *PageEntry,\r
-  IN  PAGE_ATTRIBUTE                    PageAttribute\r
+  IN  PHYSICAL_ADDRESS  BaseAddress,\r
+  IN  UINT64            Length,\r
+  IN  UINT64            *PageEntry,\r
+  IN  PAGE_ATTRIBUTE    PageAttribute\r
   )\r
 {\r
-  UINT64                PageEntryLength;\r
+  UINT64  PageEntryLength;\r
 \r
   PageEntryLength = PageAttributeToLength (PageAttribute);\r
 \r
@@ -280,14 +467,14 @@ NeedSplitPage (
 **/\r
 RETURN_STATUS\r
 SplitPage (\r
-  IN  UINT64                            *PageEntry,\r
-  IN  PAGE_ATTRIBUTE                    PageAttribute,\r
-  IN  PAGE_ATTRIBUTE                    SplitAttribute\r
+  IN  UINT64          *PageEntry,\r
+  IN  PAGE_ATTRIBUTE  PageAttribute,\r
+  IN  PAGE_ATTRIBUTE  SplitAttribute\r
   )\r
 {\r
-  UINT64   BaseAddress;\r
-  UINT64   *NewPageEntry;\r
-  UINTN    Index;\r
+  UINT64  BaseAddress;\r
+  UINT64  *NewPageEntry;\r
+  UINTN   Index;\r
 \r
   ASSERT (PageAttribute == Page2M || PageAttribute == Page1G);\r
 \r
@@ -302,10 +489,12 @@ SplitPage (
       if (NewPageEntry == NULL) {\r
         return RETURN_OUT_OF_RESOURCES;\r
       }\r
+\r
       BaseAddress = *PageEntry & PAGING_2M_ADDRESS_MASK_64;\r
-      for (Index = 0; Index < SIZE_4KB / sizeof(UINT64); Index++) {\r
+      for (Index = 0; Index < SIZE_4KB / sizeof (UINT64); Index++) {\r
         NewPageEntry[Index] = (BaseAddress + SIZE_4KB * Index) | mAddressEncMask | ((*PageEntry) & PAGE_PROGATE_BITS);\r
       }\r
+\r
       (*PageEntry) = (UINT64)(UINTN)NewPageEntry | mAddressEncMask | PAGE_ATTRIBUTE_BITS;\r
       return RETURN_SUCCESS;\r
     } else {\r
@@ -317,16 +506,18 @@ SplitPage (
     // No need support 1G->4K directly, we should use 1G->2M, then 2M->4K to get more compact page table.\r
     //\r
     ASSERT (SplitAttribute == Page2M || SplitAttribute == Page4K);\r
-    if ((SplitAttribute == Page2M || SplitAttribute == Page4K)) {\r
+    if (((SplitAttribute == Page2M) || (SplitAttribute == Page4K))) {\r
       NewPageEntry = AllocatePageTableMemory (1);\r
       DEBUG ((DEBUG_VERBOSE, "Split - 0x%x\n", NewPageEntry));\r
       if (NewPageEntry == NULL) {\r
         return RETURN_OUT_OF_RESOURCES;\r
       }\r
+\r
       BaseAddress = *PageEntry & PAGING_1G_ADDRESS_MASK_64;\r
-      for (Index = 0; Index < SIZE_4KB / sizeof(UINT64); Index++) {\r
+      for (Index = 0; Index < SIZE_4KB / sizeof (UINT64); Index++) {\r
         NewPageEntry[Index] = (BaseAddress + SIZE_2MB * Index) | mAddressEncMask | IA32_PG_PS | ((*PageEntry) & PAGE_PROGATE_BITS);\r
       }\r
+\r
       (*PageEntry) = (UINT64)(UINTN)NewPageEntry | mAddressEncMask | PAGE_ATTRIBUTE_BITS;\r
       return RETURN_SUCCESS;\r
     } else {\r
@@ -343,6 +534,8 @@ SplitPage (
 \r
   Caller should make sure BaseAddress and Length is at page boundary.\r
 \r
+  @param[in]   PageTableBase    The page table base.\r
+  @param[in]   EnablePML5Paging If PML5 paging is enabled.\r
   @param[in]   BaseAddress      The physical address that is the start address of a memory region.\r
   @param[in]   Length           The size in bytes of the memory region.\r
   @param[in]   Attributes       The bit mask of attributes to modify for the memory region.\r
@@ -364,26 +557,27 @@ SplitPage (
                                    range specified by BaseAddress and Length.\r
 **/\r
 RETURN_STATUS\r
-EFIAPI\r
 ConvertMemoryPageAttributes (\r
-  IN  PHYSICAL_ADDRESS                  BaseAddress,\r
-  IN  UINT64                            Length,\r
-  IN  UINT64                            Attributes,\r
-  IN  BOOLEAN                           IsSet,\r
-  OUT BOOLEAN                           *IsSplitted,  OPTIONAL\r
-  OUT BOOLEAN                           *IsModified   OPTIONAL\r
+  IN  UINTN             PageTableBase,\r
+  IN  BOOLEAN           EnablePML5Paging,\r
+  IN  PHYSICAL_ADDRESS  BaseAddress,\r
+  IN  UINT64            Length,\r
+  IN  UINT64            Attributes,\r
+  IN  BOOLEAN           IsSet,\r
+  OUT BOOLEAN           *IsSplitted   OPTIONAL,\r
+  OUT BOOLEAN           *IsModified   OPTIONAL\r
   )\r
 {\r
-  UINT64                            *PageEntry;\r
-  PAGE_ATTRIBUTE                    PageAttribute;\r
-  UINTN                             PageEntryLength;\r
-  PAGE_ATTRIBUTE                    SplitAttribute;\r
-  RETURN_STATUS                     Status;\r
-  BOOLEAN                           IsEntryModified;\r
-  EFI_PHYSICAL_ADDRESS              MaximumSupportMemAddress;\r
+  UINT64                *PageEntry;\r
+  PAGE_ATTRIBUTE        PageAttribute;\r
+  UINTN                 PageEntryLength;\r
+  PAGE_ATTRIBUTE        SplitAttribute;\r
+  RETURN_STATUS         Status;\r
+  BOOLEAN               IsEntryModified;\r
+  EFI_PHYSICAL_ADDRESS  MaximumSupportMemAddress;\r
 \r
   ASSERT (Attributes != 0);\r
-  ASSERT ((Attributes & ~(EFI_MEMORY_RP | EFI_MEMORY_RO | EFI_MEMORY_XP)) == 0);\r
+  ASSERT ((Attributes & ~EFI_MEMORY_ATTRIBUTE_MASK) == 0);\r
 \r
   ASSERT ((BaseAddress & (SIZE_4KB - 1)) == 0);\r
   ASSERT ((Length & (SIZE_4KB - 1)) == 0);\r
@@ -396,32 +590,36 @@ ConvertMemoryPageAttributes (
   if (BaseAddress > MaximumSupportMemAddress) {\r
     return RETURN_UNSUPPORTED;\r
   }\r
+\r
   if (Length > MaximumSupportMemAddress) {\r
     return RETURN_UNSUPPORTED;\r
   }\r
+\r
   if ((Length != 0) && (BaseAddress > MaximumSupportMemAddress - (Length - 1))) {\r
     return RETURN_UNSUPPORTED;\r
   }\r
 \r
-//  DEBUG ((DEBUG_ERROR, "ConvertMemoryPageAttributes(%x) - %016lx, %016lx, %02lx\n", IsSet, BaseAddress, Length, Attributes));\r
+  //  DEBUG ((DEBUG_ERROR, "ConvertMemoryPageAttributes(%x) - %016lx, %016lx, %02lx\n", IsSet, BaseAddress, Length, Attributes));\r
 \r
   if (IsSplitted != NULL) {\r
     *IsSplitted = FALSE;\r
   }\r
+\r
   if (IsModified != NULL) {\r
     *IsModified = FALSE;\r
   }\r
 \r
   //\r
-  // Below logic is to check 2M/4K page to make sure we donot waist memory.\r
+  // Below logic is to check 2M/4K page to make sure we do not waste memory.\r
   //\r
   while (Length != 0) {\r
-    PageEntry = GetPageTableEntry (BaseAddress, &PageAttribute);\r
+    PageEntry = GetPageTableEntry (PageTableBase, EnablePML5Paging, BaseAddress, &PageAttribute);\r
     if (PageEntry == NULL) {\r
       return RETURN_UNSUPPORTED;\r
     }\r
+\r
     PageEntryLength = PageAttributeToLength (PageAttribute);\r
-    SplitAttribute = NeedSplitPage (BaseAddress, Length, PageEntry, PageAttribute);\r
+    SplitAttribute  = NeedSplitPage (BaseAddress, Length, PageEntry, PageAttribute);\r
     if (SplitAttribute == PageNone) {\r
       ConvertPageEntryAttribute (PageEntry, Attributes, IsSet, &IsEntryModified);\r
       if (IsEntryModified) {\r
@@ -429,22 +627,26 @@ ConvertMemoryPageAttributes (
           *IsModified = TRUE;\r
         }\r
       }\r
+\r
       //\r
       // Convert success, move to next\r
       //\r
       BaseAddress += PageEntryLength;\r
-      Length -= PageEntryLength;\r
+      Length      -= PageEntryLength;\r
     } else {\r
       Status = SplitPage (PageEntry, PageAttribute, SplitAttribute);\r
       if (RETURN_ERROR (Status)) {\r
         return RETURN_UNSUPPORTED;\r
       }\r
+\r
       if (IsSplitted != NULL) {\r
         *IsSplitted = TRUE;\r
       }\r
+\r
       if (IsModified != NULL) {\r
         *IsModified = TRUE;\r
       }\r
+\r
       //\r
       // Just split current page\r
       // Convert success in next around\r
@@ -477,7 +679,7 @@ FlushTlbForAll (
   VOID\r
   )\r
 {\r
-  UINTN       Index;\r
+  UINTN  Index;\r
 \r
   FlushTlbOnCurrentProcessor (NULL);\r
 \r
@@ -494,6 +696,8 @@ FlushTlbForAll (
   This function sets the attributes for the memory region specified by BaseAddress and\r
   Length from their current attributes to the attributes specified by Attributes.\r
 \r
+  @param[in]   PageTableBase    The page table base.\r
+  @param[in]   EnablePML5Paging If PML5 paging is enabled.\r
   @param[in]   BaseAddress      The physical address that is the start address of a memory region.\r
   @param[in]   Length           The size in bytes of the memory region.\r
   @param[in]   Attributes       The bit mask of attributes to set for the memory region.\r
@@ -514,24 +718,25 @@ FlushTlbForAll (
 \r
 **/\r
 EFI_STATUS\r
-EFIAPI\r
 SmmSetMemoryAttributesEx (\r
-  IN  EFI_PHYSICAL_ADDRESS                       BaseAddress,\r
-  IN  UINT64                                     Length,\r
-  IN  UINT64                                     Attributes,\r
-  OUT BOOLEAN                                    *IsSplitted  OPTIONAL\r
+  IN  UINTN                 PageTableBase,\r
+  IN  BOOLEAN               EnablePML5Paging,\r
+  IN  EFI_PHYSICAL_ADDRESS  BaseAddress,\r
+  IN  UINT64                Length,\r
+  IN  UINT64                Attributes,\r
+  OUT BOOLEAN               *IsSplitted  OPTIONAL\r
   )\r
 {\r
   EFI_STATUS  Status;\r
   BOOLEAN     IsModified;\r
 \r
-  Status = ConvertMemoryPageAttributes (BaseAddress, Length, Attributes, TRUE, IsSplitted, &IsModified);\r
-  if (!EFI_ERROR(Status)) {\r
+  Status = ConvertMemoryPageAttributes (PageTableBase, EnablePML5Paging, BaseAddress, Length, Attributes, TRUE, IsSplitted, &IsModified);\r
+  if (!EFI_ERROR (Status)) {\r
     if (IsModified) {\r
       //\r
       // Flush TLB as last step\r
       //\r
-      FlushTlbForAll();\r
+      FlushTlbForAll ();\r
     }\r
   }\r
 \r
@@ -542,6 +747,8 @@ SmmSetMemoryAttributesEx (
   This function clears the attributes for the memory region specified by BaseAddress and\r
   Length from their current attributes to the attributes specified by Attributes.\r
 \r
+  @param[in]   PageTableBase    The page table base.\r
+  @param[in]   EnablePML5Paging If PML5 paging is enabled.\r
   @param[in]   BaseAddress      The physical address that is the start address of a memory region.\r
   @param[in]   Length           The size in bytes of the memory region.\r
   @param[in]   Attributes       The bit mask of attributes to clear for the memory region.\r
@@ -552,34 +759,35 @@ SmmSetMemoryAttributesEx (
                                 BaseAddress and Length cannot be modified.\r
   @retval EFI_INVALID_PARAMETER Length is zero.\r
                                 Attributes specified an illegal combination of attributes that\r
-                                cannot be set together.\r
+                                cannot be cleared together.\r
   @retval EFI_OUT_OF_RESOURCES  There are not enough system resources to modify the attributes of\r
                                 the memory resource range.\r
   @retval EFI_UNSUPPORTED       The processor does not support one or more bytes of the memory\r
                                 resource range specified by BaseAddress and Length.\r
-                                The bit mask of attributes is not support for the memory resource\r
+                                The bit mask of attributes is not supported for the memory resource\r
                                 range specified by BaseAddress and Length.\r
 \r
 **/\r
 EFI_STATUS\r
-EFIAPI\r
 SmmClearMemoryAttributesEx (\r
-  IN  EFI_PHYSICAL_ADDRESS                       BaseAddress,\r
-  IN  UINT64                                     Length,\r
-  IN  UINT64                                     Attributes,\r
-  OUT BOOLEAN                                    *IsSplitted  OPTIONAL\r
+  IN  UINTN                 PageTableBase,\r
+  IN  BOOLEAN               EnablePML5Paging,\r
+  IN  EFI_PHYSICAL_ADDRESS  BaseAddress,\r
+  IN  UINT64                Length,\r
+  IN  UINT64                Attributes,\r
+  OUT BOOLEAN               *IsSplitted  OPTIONAL\r
   )\r
 {\r
   EFI_STATUS  Status;\r
   BOOLEAN     IsModified;\r
 \r
-  Status = ConvertMemoryPageAttributes (BaseAddress, Length, Attributes, FALSE, IsSplitted, &IsModified);\r
-  if (!EFI_ERROR(Status)) {\r
+  Status = ConvertMemoryPageAttributes (PageTableBase, EnablePML5Paging, BaseAddress, Length, Attributes, FALSE, IsSplitted, &IsModified);\r
+  if (!EFI_ERROR (Status)) {\r
     if (IsModified) {\r
       //\r
       // Flush TLB as last step\r
       //\r
-      FlushTlbForAll();\r
+      FlushTlbForAll ();\r
     }\r
   }\r
 \r
@@ -604,19 +812,25 @@ SmmClearMemoryAttributesEx (
                                 the memory resource range.\r
   @retval EFI_UNSUPPORTED       The processor does not support one or more bytes of the memory\r
                                 resource range specified by BaseAddress and Length.\r
-                                The bit mask of attributes is not support for the memory resource\r
+                                The bit mask of attributes is not supported for the memory resource\r
                                 range specified by BaseAddress and Length.\r
 \r
 **/\r
 EFI_STATUS\r
-EFIAPI\r
 SmmSetMemoryAttributes (\r
-  IN  EFI_PHYSICAL_ADDRESS                       BaseAddress,\r
-  IN  UINT64                                     Length,\r
-  IN  UINT64                                     Attributes\r
+  IN  EFI_PHYSICAL_ADDRESS  BaseAddress,\r
+  IN  UINT64                Length,\r
+  IN  UINT64                Attributes\r
   )\r
 {\r
-  return SmmSetMemoryAttributesEx (BaseAddress, Length, Attributes, NULL);\r
+  IA32_CR4  Cr4;\r
+  UINTN     PageTableBase;\r
+  BOOLEAN   Enable5LevelPaging;\r
+\r
+  PageTableBase      = AsmReadCr3 () & PAGING_4K_ADDRESS_MASK_64;\r
+  Cr4.UintN          = AsmReadCr4 ();\r
+  Enable5LevelPaging = (BOOLEAN)(Cr4.Bits.LA57 == 1);\r
+  return SmmSetMemoryAttributesEx (PageTableBase, Enable5LevelPaging, BaseAddress, Length, Attributes, NULL);\r
 }\r
 \r
 /**\r
@@ -632,27 +846,78 @@ SmmSetMemoryAttributes (
                                 BaseAddress and Length cannot be modified.\r
   @retval EFI_INVALID_PARAMETER Length is zero.\r
                                 Attributes specified an illegal combination of attributes that\r
-                                cannot be set together.\r
+                                cannot be cleared together.\r
   @retval EFI_OUT_OF_RESOURCES  There are not enough system resources to modify the attributes of\r
                                 the memory resource range.\r
   @retval EFI_UNSUPPORTED       The processor does not support one or more bytes of the memory\r
                                 resource range specified by BaseAddress and Length.\r
-                                The bit mask of attributes is not support for the memory resource\r
+                                The bit mask of attributes is not supported for the memory resource\r
                                 range specified by BaseAddress and Length.\r
 \r
 **/\r
 EFI_STATUS\r
-EFIAPI\r
 SmmClearMemoryAttributes (\r
-  IN  EFI_PHYSICAL_ADDRESS                       BaseAddress,\r
-  IN  UINT64                                     Length,\r
-  IN  UINT64                                     Attributes\r
+  IN  EFI_PHYSICAL_ADDRESS  BaseAddress,\r
+  IN  UINT64                Length,\r
+  IN  UINT64                Attributes\r
+  )\r
+{\r
+  IA32_CR4  Cr4;\r
+  UINTN     PageTableBase;\r
+  BOOLEAN   Enable5LevelPaging;\r
+\r
+  PageTableBase      = AsmReadCr3 () & PAGING_4K_ADDRESS_MASK_64;\r
+  Cr4.UintN          = AsmReadCr4 ();\r
+  Enable5LevelPaging = (BOOLEAN)(Cr4.Bits.LA57 == 1);\r
+  return SmmClearMemoryAttributesEx (PageTableBase, Enable5LevelPaging, BaseAddress, Length, Attributes, NULL);\r
+}\r
+\r
+/**\r
+  Set ShadowStack memory.\r
+\r
+  @param[in]  Cr3              The page table base address.\r
+  @param[in]  BaseAddress      The physical address that is the start address of a memory region.\r
+  @param[in]  Length           The size in bytes of the memory region.\r
+\r
+  @retval EFI_SUCCESS           The shadow stack memory is set.\r
+**/\r
+EFI_STATUS\r
+SetShadowStack (\r
+  IN  UINTN                 Cr3,\r
+  IN  EFI_PHYSICAL_ADDRESS  BaseAddress,\r
+  IN  UINT64                Length\r
   )\r
 {\r
-  return SmmClearMemoryAttributesEx (BaseAddress, Length, Attributes, NULL);\r
+  EFI_STATUS  Status;\r
+\r
+  mIsShadowStack = TRUE;\r
+  Status         = SmmSetMemoryAttributesEx (Cr3, m5LevelPagingNeeded, BaseAddress, Length, EFI_MEMORY_RO, NULL);\r
+  mIsShadowStack = FALSE;\r
+\r
+  return Status;\r
 }\r
 \r
+/**\r
+  Set not present memory.\r
 \r
+  @param[in]  Cr3              The page table base address.\r
+  @param[in]  BaseAddress      The physical address that is the start address of a memory region.\r
+  @param[in]  Length           The size in bytes of the memory region.\r
+\r
+  @retval EFI_SUCCESS           The not present memory is set.\r
+**/\r
+EFI_STATUS\r
+SetNotPresentPage (\r
+  IN  UINTN                 Cr3,\r
+  IN  EFI_PHYSICAL_ADDRESS  BaseAddress,\r
+  IN  UINT64                Length\r
+  )\r
+{\r
+  EFI_STATUS  Status;\r
+\r
+  Status = SmmSetMemoryAttributesEx (Cr3, m5LevelPagingNeeded, BaseAddress, Length, EFI_MEMORY_RP, NULL);\r
+  return Status;\r
+}\r
 \r
 /**\r
   Retrieves a pointer to the system configuration table from the SMM System Table\r
@@ -672,7 +937,7 @@ SmmGetSystemConfigurationTable (
   OUT VOID      **Table\r
   )\r
 {\r
-  UINTN             Index;\r
+  UINTN  Index;\r
 \r
   ASSERT (TableGuid != NULL);\r
   ASSERT (Table != NULL);\r
@@ -702,11 +967,11 @@ PatchSmmSaveStateMap (
   UINTN  TileSize;\r
 \r
   TileCodeSize = GetSmiHandlerSize ();\r
-  TileCodeSize = ALIGN_VALUE(TileCodeSize, SIZE_4KB);\r
+  TileCodeSize = ALIGN_VALUE (TileCodeSize, SIZE_4KB);\r
   TileDataSize = (SMRAM_SAVE_STATE_MAP_OFFSET - SMM_PSD_OFFSET) + sizeof (SMRAM_SAVE_STATE_MAP);\r
-  TileDataSize = ALIGN_VALUE(TileDataSize, SIZE_4KB);\r
-  TileSize = TileDataSize + TileCodeSize - 1;\r
-  TileSize = 2 * GetPowerOfTwo32 ((UINT32)TileSize);\r
+  TileDataSize = ALIGN_VALUE (TileDataSize, SIZE_4KB);\r
+  TileSize     = TileDataSize + TileCodeSize - 1;\r
+  TileSize     = 2 * GetPowerOfTwo32 ((UINT32)TileSize);\r
 \r
   DEBUG ((DEBUG_INFO, "PatchSmmSaveStateMap:\n"));\r
   for (Index = 0; Index < mMaxNumberOfCpus - 1; Index++) {\r
@@ -768,6 +1033,52 @@ PatchSmmSaveStateMap (
     );\r
 }\r
 \r
+/**\r
+  This function sets GDT/IDT buffer to be RO and XP.\r
+**/\r
+VOID\r
+PatchGdtIdtMap (\r
+  VOID\r
+  )\r
+{\r
+  EFI_PHYSICAL_ADDRESS  BaseAddress;\r
+  UINTN                 Size;\r
+\r
+  //\r
+  // GDT\r
+  //\r
+  DEBUG ((DEBUG_INFO, "PatchGdtIdtMap - GDT:\n"));\r
+\r
+  BaseAddress = mGdtBuffer;\r
+  Size        = ALIGN_VALUE (mGdtBufferSize, SIZE_4KB);\r
+  //\r
+  // The range should have been set to RO\r
+  // if it is allocated with EfiRuntimeServicesCode.\r
+  //\r
+  SmmSetMemoryAttributes (\r
+    BaseAddress,\r
+    Size,\r
+    EFI_MEMORY_XP\r
+    );\r
+\r
+  //\r
+  // IDT\r
+  //\r
+  DEBUG ((DEBUG_INFO, "PatchGdtIdtMap - IDT:\n"));\r
+\r
+  BaseAddress = gcSmiIdtr.Base;\r
+  Size        = ALIGN_VALUE (gcSmiIdtr.Limit + 1, SIZE_4KB);\r
+  //\r
+  // The range should have been set to RO\r
+  // if it is allocated with EfiRuntimeServicesCode.\r
+  //\r
+  SmmSetMemoryAttributes (\r
+    BaseAddress,\r
+    Size,\r
+    EFI_MEMORY_XP\r
+    );\r
+}\r
+\r
 /**\r
   This function sets memory attribute according to MemoryAttributesTable.\r
 **/\r
@@ -776,17 +1087,17 @@ SetMemMapAttributes (
   VOID\r
   )\r
 {\r
-  EFI_MEMORY_DESCRIPTOR                     *MemoryMap;\r
-  EFI_MEMORY_DESCRIPTOR                     *MemoryMapStart;\r
-  UINTN                                     MemoryMapEntryCount;\r
-  UINTN                                     DescriptorSize;\r
-  UINTN                                     Index;\r
-  EDKII_PI_SMM_MEMORY_ATTRIBUTES_TABLE      *MemoryAttributesTable;\r
+  EFI_MEMORY_DESCRIPTOR                 *MemoryMap;\r
+  EFI_MEMORY_DESCRIPTOR                 *MemoryMapStart;\r
+  UINTN                                 MemoryMapEntryCount;\r
+  UINTN                                 DescriptorSize;\r
+  UINTN                                 Index;\r
+  EDKII_PI_SMM_MEMORY_ATTRIBUTES_TABLE  *MemoryAttributesTable;\r
 \r
   SmmGetSystemConfigurationTable (&gEdkiiPiSmmMemoryAttributesTableGuid, (VOID **)&MemoryAttributesTable);\r
   if (MemoryAttributesTable == NULL) {\r
     DEBUG ((DEBUG_INFO, "MemoryAttributesTable - NULL\n"));\r
-    return ;\r
+    return;\r
   }\r
 \r
   DEBUG ((DEBUG_INFO, "MemoryAttributesTable:\n"));\r
@@ -795,9 +1106,9 @@ SetMemMapAttributes (
   DEBUG ((DEBUG_INFO, "  DescriptorSize            - 0x%08x\n", MemoryAttributesTable->DescriptorSize));\r
 \r
   MemoryMapEntryCount = MemoryAttributesTable->NumberOfEntries;\r
-  DescriptorSize = MemoryAttributesTable->DescriptorSize;\r
-  MemoryMapStart = (EFI_MEMORY_DESCRIPTOR *)(MemoryAttributesTable + 1);\r
-  MemoryMap = MemoryMapStart;\r
+  DescriptorSize      = MemoryAttributesTable->DescriptorSize;\r
+  MemoryMapStart      = (EFI_MEMORY_DESCRIPTOR *)(MemoryAttributesTable + 1);\r
+  MemoryMap           = MemoryMapStart;\r
   for (Index = 0; Index < MemoryMapEntryCount; Index++) {\r
     DEBUG ((DEBUG_INFO, "Entry (0x%x)\n", MemoryMap));\r
     DEBUG ((DEBUG_INFO, "  Type              - 0x%x\n", MemoryMap->Type));\r
@@ -805,42 +1116,43 @@ SetMemMapAttributes (
     DEBUG ((DEBUG_INFO, "  VirtualStart      - 0x%016lx\n", MemoryMap->VirtualStart));\r
     DEBUG ((DEBUG_INFO, "  NumberOfPages     - 0x%016lx\n", MemoryMap->NumberOfPages));\r
     DEBUG ((DEBUG_INFO, "  Attribute         - 0x%016lx\n", MemoryMap->Attribute));\r
-    MemoryMap = NEXT_MEMORY_DESCRIPTOR(MemoryMap, DescriptorSize);\r
+    MemoryMap = NEXT_MEMORY_DESCRIPTOR (MemoryMap, DescriptorSize);\r
   }\r
 \r
   MemoryMap = MemoryMapStart;\r
   for (Index = 0; Index < MemoryMapEntryCount; Index++) {\r
     DEBUG ((DEBUG_VERBOSE, "SetAttribute: Memory Entry - 0x%lx, 0x%x\n", MemoryMap->PhysicalStart, MemoryMap->NumberOfPages));\r
     switch (MemoryMap->Type) {\r
-    case EfiRuntimeServicesCode:\r
-      SmmSetMemoryAttributes (\r
-        MemoryMap->PhysicalStart,\r
-        EFI_PAGES_TO_SIZE((UINTN)MemoryMap->NumberOfPages),\r
-        EFI_MEMORY_RO\r
-        );\r
-      break;\r
-    case EfiRuntimeServicesData:\r
-      SmmSetMemoryAttributes (\r
-        MemoryMap->PhysicalStart,\r
-        EFI_PAGES_TO_SIZE((UINTN)MemoryMap->NumberOfPages),\r
-        EFI_MEMORY_XP\r
-        );\r
-      break;\r
-    default:\r
-      SmmSetMemoryAttributes (\r
-        MemoryMap->PhysicalStart,\r
-        EFI_PAGES_TO_SIZE((UINTN)MemoryMap->NumberOfPages),\r
-        EFI_MEMORY_XP\r
-        );\r
-      break;\r
+      case EfiRuntimeServicesCode:\r
+        SmmSetMemoryAttributes (\r
+          MemoryMap->PhysicalStart,\r
+          EFI_PAGES_TO_SIZE ((UINTN)MemoryMap->NumberOfPages),\r
+          EFI_MEMORY_RO\r
+          );\r
+        break;\r
+      case EfiRuntimeServicesData:\r
+        SmmSetMemoryAttributes (\r
+          MemoryMap->PhysicalStart,\r
+          EFI_PAGES_TO_SIZE ((UINTN)MemoryMap->NumberOfPages),\r
+          EFI_MEMORY_XP\r
+          );\r
+        break;\r
+      default:\r
+        SmmSetMemoryAttributes (\r
+          MemoryMap->PhysicalStart,\r
+          EFI_PAGES_TO_SIZE ((UINTN)MemoryMap->NumberOfPages),\r
+          EFI_MEMORY_XP\r
+          );\r
+        break;\r
     }\r
-    MemoryMap = NEXT_MEMORY_DESCRIPTOR(MemoryMap, DescriptorSize);\r
+\r
+    MemoryMap = NEXT_MEMORY_DESCRIPTOR (MemoryMap, DescriptorSize);\r
   }\r
 \r
   PatchSmmSaveStateMap ();\r
   PatchGdtIdtMap ();\r
 \r
-  return ;\r
+  return;\r
 }\r
 \r
 /**\r
@@ -859,27 +1171,27 @@ SortMemoryMap (
   IN UINTN                      DescriptorSize\r
   )\r
 {\r
-  EFI_MEMORY_DESCRIPTOR       *MemoryMapEntry;\r
-  EFI_MEMORY_DESCRIPTOR       *NextMemoryMapEntry;\r
-  EFI_MEMORY_DESCRIPTOR       *MemoryMapEnd;\r
-  EFI_MEMORY_DESCRIPTOR       TempMemoryMap;\r
+  EFI_MEMORY_DESCRIPTOR  *MemoryMapEntry;\r
+  EFI_MEMORY_DESCRIPTOR  *NextMemoryMapEntry;\r
+  EFI_MEMORY_DESCRIPTOR  *MemoryMapEnd;\r
+  EFI_MEMORY_DESCRIPTOR  TempMemoryMap;\r
 \r
-  MemoryMapEntry = MemoryMap;\r
+  MemoryMapEntry     = MemoryMap;\r
   NextMemoryMapEntry = NEXT_MEMORY_DESCRIPTOR (MemoryMapEntry, DescriptorSize);\r
-  MemoryMapEnd = (EFI_MEMORY_DESCRIPTOR *) ((UINT8 *) MemoryMap + MemoryMapSize);\r
+  MemoryMapEnd       = (EFI_MEMORY_DESCRIPTOR *)((UINT8 *)MemoryMap + MemoryMapSize);\r
   while (MemoryMapEntry < MemoryMapEnd) {\r
     while (NextMemoryMapEntry < MemoryMapEnd) {\r
       if (MemoryMapEntry->PhysicalStart > NextMemoryMapEntry->PhysicalStart) {\r
-        CopyMem (&TempMemoryMap, MemoryMapEntry, sizeof(EFI_MEMORY_DESCRIPTOR));\r
-        CopyMem (MemoryMapEntry, NextMemoryMapEntry, sizeof(EFI_MEMORY_DESCRIPTOR));\r
-        CopyMem (NextMemoryMapEntry, &TempMemoryMap, sizeof(EFI_MEMORY_DESCRIPTOR));\r
+        CopyMem (&TempMemoryMap, MemoryMapEntry, sizeof (EFI_MEMORY_DESCRIPTOR));\r
+        CopyMem (MemoryMapEntry, NextMemoryMapEntry, sizeof (EFI_MEMORY_DESCRIPTOR));\r
+        CopyMem (NextMemoryMapEntry, &TempMemoryMap, sizeof (EFI_MEMORY_DESCRIPTOR));\r
       }\r
 \r
       NextMemoryMapEntry = NEXT_MEMORY_DESCRIPTOR (NextMemoryMapEntry, DescriptorSize);\r
     }\r
 \r
-    MemoryMapEntry      = NEXT_MEMORY_DESCRIPTOR (MemoryMapEntry, DescriptorSize);\r
-    NextMemoryMapEntry  = NEXT_MEMORY_DESCRIPTOR (MemoryMapEntry, DescriptorSize);\r
+    MemoryMapEntry     = NEXT_MEMORY_DESCRIPTOR (MemoryMapEntry, DescriptorSize);\r
+    NextMemoryMapEntry = NEXT_MEMORY_DESCRIPTOR (MemoryMapEntry, DescriptorSize);\r
   }\r
 }\r
 \r
@@ -901,21 +1213,21 @@ IsUefiPageNotPresent (
   )\r
 {\r
   switch (MemoryMap->Type) {\r
-  case EfiLoaderCode:\r
-  case EfiLoaderData:\r
-  case EfiBootServicesCode:\r
-  case EfiBootServicesData:\r
-  case EfiConventionalMemory:\r
-  case EfiUnusableMemory:\r
-  case EfiACPIReclaimMemory:\r
-    return TRUE;\r
-  default:\r
-    return FALSE;\r
+    case EfiLoaderCode:\r
+    case EfiLoaderData:\r
+    case EfiBootServicesCode:\r
+    case EfiBootServicesData:\r
+    case EfiConventionalMemory:\r
+    case EfiUnusableMemory:\r
+    case EfiACPIReclaimMemory:\r
+      return TRUE;\r
+    default:\r
+      return FALSE;\r
   }\r
 }\r
 \r
 /**\r
-  Merge continous memory map entries whose type is\r
+  Merge continuous memory map entries whose type is\r
   EfiLoaderCode/Data, EfiBootServicesCode/Data, EfiConventionalMemory,\r
   EfiUnusableMemory, EfiACPIReclaimMemory, because the memory described by\r
   these entries will be set as NOT present in SMM page table.\r
@@ -936,24 +1248,25 @@ MergeMemoryMapForNotPresentEntry (
   IN UINTN                      DescriptorSize\r
   )\r
 {\r
-  EFI_MEMORY_DESCRIPTOR       *MemoryMapEntry;\r
-  EFI_MEMORY_DESCRIPTOR       *MemoryMapEnd;\r
-  UINT64                      MemoryBlockLength;\r
-  EFI_MEMORY_DESCRIPTOR       *NewMemoryMapEntry;\r
-  EFI_MEMORY_DESCRIPTOR       *NextMemoryMapEntry;\r
+  EFI_MEMORY_DESCRIPTOR  *MemoryMapEntry;\r
+  EFI_MEMORY_DESCRIPTOR  *MemoryMapEnd;\r
+  UINT64                 MemoryBlockLength;\r
+  EFI_MEMORY_DESCRIPTOR  *NewMemoryMapEntry;\r
+  EFI_MEMORY_DESCRIPTOR  *NextMemoryMapEntry;\r
 \r
-  MemoryMapEntry = MemoryMap;\r
+  MemoryMapEntry    = MemoryMap;\r
   NewMemoryMapEntry = MemoryMap;\r
-  MemoryMapEnd = (EFI_MEMORY_DESCRIPTOR *) ((UINT8 *) MemoryMap + *MemoryMapSize);\r
+  MemoryMapEnd      = (EFI_MEMORY_DESCRIPTOR *)((UINT8 *)MemoryMap + *MemoryMapSize);\r
   while ((UINTN)MemoryMapEntry < (UINTN)MemoryMapEnd) {\r
-    CopyMem (NewMemoryMapEntry, MemoryMapEntry, sizeof(EFI_MEMORY_DESCRIPTOR));\r
+    CopyMem (NewMemoryMapEntry, MemoryMapEntry, sizeof (EFI_MEMORY_DESCRIPTOR));\r
     NextMemoryMapEntry = NEXT_MEMORY_DESCRIPTOR (MemoryMapEntry, DescriptorSize);\r
 \r
     do {\r
-      MemoryBlockLength = (UINT64) (EFI_PAGES_TO_SIZE((UINTN)MemoryMapEntry->NumberOfPages));\r
+      MemoryBlockLength = (UINT64)(EFI_PAGES_TO_SIZE ((UINTN)MemoryMapEntry->NumberOfPages));\r
       if (((UINTN)NextMemoryMapEntry < (UINTN)MemoryMapEnd) &&\r
-          IsUefiPageNotPresent(MemoryMapEntry) && IsUefiPageNotPresent(NextMemoryMapEntry) &&\r
-          ((MemoryMapEntry->PhysicalStart + MemoryBlockLength) == NextMemoryMapEntry->PhysicalStart)) {\r
+          IsUefiPageNotPresent (MemoryMapEntry) && IsUefiPageNotPresent (NextMemoryMapEntry) &&\r
+          ((MemoryMapEntry->PhysicalStart + MemoryBlockLength) == NextMemoryMapEntry->PhysicalStart))\r
+      {\r
         MemoryMapEntry->NumberOfPages += NextMemoryMapEntry->NumberOfPages;\r
         if (NewMemoryMapEntry != MemoryMapEntry) {\r
           NewMemoryMapEntry->NumberOfPages += NextMemoryMapEntry->NumberOfPages;\r
@@ -967,13 +1280,89 @@ MergeMemoryMapForNotPresentEntry (
       }\r
     } while (TRUE);\r
 \r
-    MemoryMapEntry = NEXT_MEMORY_DESCRIPTOR (MemoryMapEntry, DescriptorSize);\r
+    MemoryMapEntry    = NEXT_MEMORY_DESCRIPTOR (MemoryMapEntry, DescriptorSize);\r
     NewMemoryMapEntry = NEXT_MEMORY_DESCRIPTOR (NewMemoryMapEntry, DescriptorSize);\r
   }\r
 \r
   *MemoryMapSize = (UINTN)NewMemoryMapEntry - (UINTN)MemoryMap;\r
 \r
-  return ;\r
+  return;\r
+}\r
+\r
+/**\r
+  This function caches the GCD memory map information.\r
+**/\r
+VOID\r
+GetGcdMemoryMap (\r
+  VOID\r
+  )\r
+{\r
+  UINTN                            NumberOfDescriptors;\r
+  EFI_GCD_MEMORY_SPACE_DESCRIPTOR  *MemSpaceMap;\r
+  EFI_STATUS                       Status;\r
+  UINTN                            Index;\r
+\r
+  Status = gDS->GetMemorySpaceMap (&NumberOfDescriptors, &MemSpaceMap);\r
+  if (EFI_ERROR (Status)) {\r
+    return;\r
+  }\r
+\r
+  mGcdMemNumberOfDesc = 0;\r
+  for (Index = 0; Index < NumberOfDescriptors; Index++) {\r
+    if ((MemSpaceMap[Index].GcdMemoryType == EfiGcdMemoryTypeReserved) &&\r
+        ((MemSpaceMap[Index].Capabilities & (EFI_MEMORY_PRESENT | EFI_MEMORY_INITIALIZED | EFI_MEMORY_TESTED)) ==\r
+         (EFI_MEMORY_PRESENT | EFI_MEMORY_INITIALIZED))\r
+        )\r
+    {\r
+      mGcdMemNumberOfDesc++;\r
+    }\r
+  }\r
+\r
+  mGcdMemSpace = AllocateZeroPool (mGcdMemNumberOfDesc * sizeof (EFI_GCD_MEMORY_SPACE_DESCRIPTOR));\r
+  ASSERT (mGcdMemSpace != NULL);\r
+  if (mGcdMemSpace == NULL) {\r
+    mGcdMemNumberOfDesc = 0;\r
+    gBS->FreePool (MemSpaceMap);\r
+    return;\r
+  }\r
+\r
+  mGcdMemNumberOfDesc = 0;\r
+  for (Index = 0; Index < NumberOfDescriptors; Index++) {\r
+    if ((MemSpaceMap[Index].GcdMemoryType == EfiGcdMemoryTypeReserved) &&\r
+        ((MemSpaceMap[Index].Capabilities & (EFI_MEMORY_PRESENT | EFI_MEMORY_INITIALIZED | EFI_MEMORY_TESTED)) ==\r
+         (EFI_MEMORY_PRESENT | EFI_MEMORY_INITIALIZED))\r
+        )\r
+    {\r
+      CopyMem (\r
+        &mGcdMemSpace[mGcdMemNumberOfDesc],\r
+        &MemSpaceMap[Index],\r
+        sizeof (EFI_GCD_MEMORY_SPACE_DESCRIPTOR)\r
+        );\r
+      mGcdMemNumberOfDesc++;\r
+    }\r
+  }\r
+\r
+  gBS->FreePool (MemSpaceMap);\r
+}\r
+\r
+/**\r
+  Get UEFI MemoryAttributesTable.\r
+**/\r
+VOID\r
+GetUefiMemoryAttributesTable (\r
+  VOID\r
+  )\r
+{\r
+  EFI_STATUS                   Status;\r
+  EFI_MEMORY_ATTRIBUTES_TABLE  *MemoryAttributesTable;\r
+  UINTN                        MemoryAttributesTableSize;\r
+\r
+  Status = EfiGetSystemConfigurationTable (&gEfiMemoryAttributesTableGuid, (VOID **)&MemoryAttributesTable);\r
+  if (!EFI_ERROR (Status) && (MemoryAttributesTable != NULL)) {\r
+    MemoryAttributesTableSize  = sizeof (EFI_MEMORY_ATTRIBUTES_TABLE) + MemoryAttributesTable->DescriptorSize * MemoryAttributesTable->NumberOfEntries;\r
+    mUefiMemoryAttributesTable = AllocateCopyPool (MemoryAttributesTableSize, MemoryAttributesTable);\r
+    ASSERT (mUefiMemoryAttributesTable != NULL);\r
+  }\r
 }\r
 \r
 /**\r
@@ -984,30 +1373,30 @@ GetUefiMemoryMap (
   VOID\r
   )\r
 {\r
-  EFI_STATUS            Status;\r
-  UINTN                 MapKey;\r
-  UINT32                DescriptorVersion;\r
-  EFI_MEMORY_DESCRIPTOR *MemoryMap;\r
-  UINTN                 UefiMemoryMapSize;\r
+  EFI_STATUS             Status;\r
+  UINTN                  MapKey;\r
+  UINT32                 DescriptorVersion;\r
+  EFI_MEMORY_DESCRIPTOR  *MemoryMap;\r
+  UINTN                  UefiMemoryMapSize;\r
 \r
   DEBUG ((DEBUG_INFO, "GetUefiMemoryMap\n"));\r
 \r
   UefiMemoryMapSize = 0;\r
-  MemoryMap = NULL;\r
-  Status = gBS->GetMemoryMap (\r
-                  &UefiMemoryMapSize,\r
-                  MemoryMap,\r
-                  &MapKey,\r
-                  &mUefiDescriptorSize,\r
-                  &DescriptorVersion\r
-                  );\r
+  MemoryMap         = NULL;\r
+  Status            = gBS->GetMemoryMap (\r
+                             &UefiMemoryMapSize,\r
+                             MemoryMap,\r
+                             &MapKey,\r
+                             &mUefiDescriptorSize,\r
+                             &DescriptorVersion\r
+                             );\r
   ASSERT (Status == EFI_BUFFER_TOO_SMALL);\r
 \r
   do {\r
     Status = gBS->AllocatePool (EfiBootServicesData, UefiMemoryMapSize, (VOID **)&MemoryMap);\r
     ASSERT (MemoryMap != NULL);\r
     if (MemoryMap == NULL) {\r
-      return ;\r
+      return;\r
     }\r
 \r
     Status = gBS->GetMemoryMap (\r
@@ -1024,17 +1413,27 @@ GetUefiMemoryMap (
   } while (Status == EFI_BUFFER_TOO_SMALL);\r
 \r
   if (MemoryMap == NULL) {\r
-    return ;\r
+    return;\r
   }\r
 \r
   SortMemoryMap (MemoryMap, UefiMemoryMapSize, mUefiDescriptorSize);\r
   MergeMemoryMapForNotPresentEntry (MemoryMap, &UefiMemoryMapSize, mUefiDescriptorSize);\r
 \r
   mUefiMemoryMapSize = UefiMemoryMapSize;\r
-  mUefiMemoryMap = AllocateCopyPool (UefiMemoryMapSize, MemoryMap);\r
+  mUefiMemoryMap     = AllocateCopyPool (UefiMemoryMapSize, MemoryMap);\r
   ASSERT (mUefiMemoryMap != NULL);\r
 \r
   gBS->FreePool (MemoryMap);\r
+\r
+  //\r
+  // Get additional information from GCD memory map.\r
+  //\r
+  GetGcdMemoryMap ();\r
+\r
+  //\r
+  // Get UEFI memory attributes table.\r
+  //\r
+  GetUefiMemoryAttributesTable ();\r
 }\r
 \r
 /**\r
@@ -1049,40 +1448,94 @@ SetUefiMemMapAttributes (
   VOID\r
   )\r
 {\r
-  EFI_STATUS            Status;\r
-  EFI_MEMORY_DESCRIPTOR *MemoryMap;\r
-  UINTN                 MemoryMapEntryCount;\r
-  UINTN                 Index;\r
+  EFI_STATUS             Status;\r
+  EFI_MEMORY_DESCRIPTOR  *MemoryMap;\r
+  UINTN                  MemoryMapEntryCount;\r
+  UINTN                  Index;\r
+  EFI_MEMORY_DESCRIPTOR  *Entry;\r
 \r
   DEBUG ((DEBUG_INFO, "SetUefiMemMapAttributes\n"));\r
 \r
-  if (mUefiMemoryMap == NULL) {\r
-    DEBUG ((DEBUG_INFO, "UefiMemoryMap - NULL\n"));\r
-    return ;\r
+  if (mUefiMemoryMap != NULL) {\r
+    MemoryMapEntryCount = mUefiMemoryMapSize/mUefiDescriptorSize;\r
+    MemoryMap           = mUefiMemoryMap;\r
+    for (Index = 0; Index < MemoryMapEntryCount; Index++) {\r
+      if (IsUefiPageNotPresent (MemoryMap)) {\r
+        Status = SmmSetMemoryAttributes (\r
+                   MemoryMap->PhysicalStart,\r
+                   EFI_PAGES_TO_SIZE ((UINTN)MemoryMap->NumberOfPages),\r
+                   EFI_MEMORY_RP\r
+                   );\r
+        DEBUG ((\r
+          DEBUG_INFO,\r
+          "UefiMemory protection: 0x%lx - 0x%lx %r\n",\r
+          MemoryMap->PhysicalStart,\r
+          MemoryMap->PhysicalStart + (UINT64)EFI_PAGES_TO_SIZE ((UINTN)MemoryMap->NumberOfPages),\r
+          Status\r
+          ));\r
+      }\r
+\r
+      MemoryMap = NEXT_MEMORY_DESCRIPTOR (MemoryMap, mUefiDescriptorSize);\r
+    }\r
   }\r
 \r
-  MemoryMapEntryCount = mUefiMemoryMapSize/mUefiDescriptorSize;\r
-  MemoryMap = mUefiMemoryMap;\r
-  for (Index = 0; Index < MemoryMapEntryCount; Index++) {\r
-    if (IsUefiPageNotPresent(MemoryMap)) {\r
+  //\r
+  // Do not free mUefiMemoryMap, it will be checked in IsSmmCommBufferForbiddenAddress().\r
+  //\r
+\r
+  //\r
+  // Set untested memory as not present.\r
+  //\r
+  if (mGcdMemSpace != NULL) {\r
+    for (Index = 0; Index < mGcdMemNumberOfDesc; Index++) {\r
       Status = SmmSetMemoryAttributes (\r
-                 MemoryMap->PhysicalStart,\r
-                 EFI_PAGES_TO_SIZE((UINTN)MemoryMap->NumberOfPages),\r
+                 mGcdMemSpace[Index].BaseAddress,\r
+                 mGcdMemSpace[Index].Length,\r
                  EFI_MEMORY_RP\r
                  );\r
       DEBUG ((\r
         DEBUG_INFO,\r
-        "UefiMemory protection: 0x%lx - 0x%lx %r\n",\r
-        MemoryMap->PhysicalStart,\r
-        MemoryMap->PhysicalStart + (UINT64)EFI_PAGES_TO_SIZE((UINTN)MemoryMap->NumberOfPages),\r
+        "GcdMemory protection: 0x%lx - 0x%lx %r\n",\r
+        mGcdMemSpace[Index].BaseAddress,\r
+        mGcdMemSpace[Index].BaseAddress + mGcdMemSpace[Index].Length,\r
         Status\r
         ));\r
     }\r
-    MemoryMap = NEXT_MEMORY_DESCRIPTOR(MemoryMap, mUefiDescriptorSize);\r
   }\r
 \r
   //\r
-  // Do free mUefiMemoryMap, it will be checked in IsSmmCommBufferForbiddenAddress().\r
+  // Do not free mGcdMemSpace, it will be checked in IsSmmCommBufferForbiddenAddress().\r
+  //\r
+\r
+  //\r
+  // Set UEFI runtime memory with EFI_MEMORY_RO as not present.\r
+  //\r
+  if (mUefiMemoryAttributesTable != NULL) {\r
+    Entry = (EFI_MEMORY_DESCRIPTOR *)(mUefiMemoryAttributesTable + 1);\r
+    for (Index = 0; Index < mUefiMemoryAttributesTable->NumberOfEntries; Index++) {\r
+      if ((Entry->Type == EfiRuntimeServicesCode) || (Entry->Type == EfiRuntimeServicesData)) {\r
+        if ((Entry->Attribute & EFI_MEMORY_RO) != 0) {\r
+          Status = SmmSetMemoryAttributes (\r
+                     Entry->PhysicalStart,\r
+                     EFI_PAGES_TO_SIZE ((UINTN)Entry->NumberOfPages),\r
+                     EFI_MEMORY_RP\r
+                     );\r
+          DEBUG ((\r
+            DEBUG_INFO,\r
+            "UefiMemoryAttribute protection: 0x%lx - 0x%lx %r\n",\r
+            Entry->PhysicalStart,\r
+            Entry->PhysicalStart + (UINT64)EFI_PAGES_TO_SIZE ((UINTN)Entry->NumberOfPages),\r
+            Status\r
+            ));\r
+        }\r
+      }\r
+\r
+      Entry = NEXT_MEMORY_DESCRIPTOR (Entry, mUefiMemoryAttributesTable->DescriptorSize);\r
+    }\r
+  }\r
+\r
+  //\r
+  // Do not free mUefiMemoryAttributesTable, it will be checked in IsSmmCommBufferForbiddenAddress().\r
   //\r
 }\r
 \r
@@ -1099,24 +1552,353 @@ IsSmmCommBufferForbiddenAddress (
   IN UINT64  Address\r
   )\r
 {\r
-  EFI_MEMORY_DESCRIPTOR *MemoryMap;\r
-  UINTN                 MemoryMapEntryCount;\r
-  UINTN                 Index;\r
+  EFI_MEMORY_DESCRIPTOR  *MemoryMap;\r
+  UINTN                  MemoryMapEntryCount;\r
+  UINTN                  Index;\r
+  EFI_MEMORY_DESCRIPTOR  *Entry;\r
+\r
+  if (mUefiMemoryMap != NULL) {\r
+    MemoryMap           = mUefiMemoryMap;\r
+    MemoryMapEntryCount = mUefiMemoryMapSize/mUefiDescriptorSize;\r
+    for (Index = 0; Index < MemoryMapEntryCount; Index++) {\r
+      if (IsUefiPageNotPresent (MemoryMap)) {\r
+        if ((Address >= MemoryMap->PhysicalStart) &&\r
+            (Address < MemoryMap->PhysicalStart + EFI_PAGES_TO_SIZE ((UINTN)MemoryMap->NumberOfPages)))\r
+        {\r
+          return TRUE;\r
+        }\r
+      }\r
 \r
-  if (mUefiMemoryMap == NULL) {\r
-    return FALSE;\r
+      MemoryMap = NEXT_MEMORY_DESCRIPTOR (MemoryMap, mUefiDescriptorSize);\r
+    }\r
   }\r
 \r
-  MemoryMap = mUefiMemoryMap;\r
-  MemoryMapEntryCount = mUefiMemoryMapSize/mUefiDescriptorSize;\r
-  for (Index = 0; Index < MemoryMapEntryCount; Index++) {\r
-    if (IsUefiPageNotPresent (MemoryMap)) {\r
-      if ((Address >= MemoryMap->PhysicalStart) &&\r
-          (Address < MemoryMap->PhysicalStart + EFI_PAGES_TO_SIZE((UINTN)MemoryMap->NumberOfPages)) ) {\r
+  if (mGcdMemSpace != NULL) {\r
+    for (Index = 0; Index < mGcdMemNumberOfDesc; Index++) {\r
+      if ((Address >= mGcdMemSpace[Index].BaseAddress) &&\r
+          (Address < mGcdMemSpace[Index].BaseAddress + mGcdMemSpace[Index].Length))\r
+      {\r
         return TRUE;\r
       }\r
     }\r
-    MemoryMap = NEXT_MEMORY_DESCRIPTOR(MemoryMap, mUefiDescriptorSize);\r
   }\r
+\r
+  if (mUefiMemoryAttributesTable != NULL) {\r
+    Entry = (EFI_MEMORY_DESCRIPTOR *)(mUefiMemoryAttributesTable + 1);\r
+    for (Index = 0; Index < mUefiMemoryAttributesTable->NumberOfEntries; Index++) {\r
+      if ((Entry->Type == EfiRuntimeServicesCode) || (Entry->Type == EfiRuntimeServicesData)) {\r
+        if ((Entry->Attribute & EFI_MEMORY_RO) != 0) {\r
+          if ((Address >= Entry->PhysicalStart) &&\r
+              (Address < Entry->PhysicalStart + LShiftU64 (Entry->NumberOfPages, EFI_PAGE_SHIFT)))\r
+          {\r
+            return TRUE;\r
+          }\r
+\r
+          Entry = NEXT_MEMORY_DESCRIPTOR (Entry, mUefiMemoryAttributesTable->DescriptorSize);\r
+        }\r
+      }\r
+    }\r
+  }\r
+\r
   return FALSE;\r
 }\r
+\r
+/**\r
+  This function set given attributes of the memory region specified by\r
+  BaseAddress and Length.\r
+\r
+  @param  This              The EDKII_SMM_MEMORY_ATTRIBUTE_PROTOCOL instance.\r
+  @param  BaseAddress       The physical address that is the start address of\r
+                            a memory region.\r
+  @param  Length            The size in bytes of the memory region.\r
+  @param  Attributes        The bit mask of attributes to set for the memory\r
+                            region.\r
+\r
+  @retval EFI_SUCCESS           The attributes were set for the memory region.\r
+  @retval EFI_INVALID_PARAMETER Length is zero.\r
+                                Attributes specified an illegal combination of\r
+                                attributes that cannot be set together.\r
+  @retval EFI_UNSUPPORTED       The processor does not support one or more\r
+                                bytes of the memory resource range specified\r
+                                by BaseAddress and Length.\r
+                                The bit mask of attributes is not supported for\r
+                                the memory resource range specified by\r
+                                BaseAddress and Length.\r
+\r
+**/\r
+EFI_STATUS\r
+EFIAPI\r
+EdkiiSmmSetMemoryAttributes (\r
+  IN  EDKII_SMM_MEMORY_ATTRIBUTE_PROTOCOL  *This,\r
+  IN  EFI_PHYSICAL_ADDRESS                 BaseAddress,\r
+  IN  UINT64                               Length,\r
+  IN  UINT64                               Attributes\r
+  )\r
+{\r
+  return SmmSetMemoryAttributes (BaseAddress, Length, Attributes);\r
+}\r
+\r
+/**\r
+  This function clears given attributes of the memory region specified by\r
+  BaseAddress and Length.\r
+\r
+  @param  This              The EDKII_SMM_MEMORY_ATTRIBUTE_PROTOCOL instance.\r
+  @param  BaseAddress       The physical address that is the start address of\r
+                            a memory region.\r
+  @param  Length            The size in bytes of the memory region.\r
+  @param  Attributes        The bit mask of attributes to clear for the memory\r
+                            region.\r
+\r
+  @retval EFI_SUCCESS           The attributes were cleared for the memory region.\r
+  @retval EFI_INVALID_PARAMETER Length is zero.\r
+                                Attributes specified an illegal combination of\r
+                                attributes that cannot be cleared together.\r
+  @retval EFI_UNSUPPORTED       The processor does not support one or more\r
+                                bytes of the memory resource range specified\r
+                                by BaseAddress and Length.\r
+                                The bit mask of attributes is not supported for\r
+                                the memory resource range specified by\r
+                                BaseAddress and Length.\r
+\r
+**/\r
+EFI_STATUS\r
+EFIAPI\r
+EdkiiSmmClearMemoryAttributes (\r
+  IN  EDKII_SMM_MEMORY_ATTRIBUTE_PROTOCOL  *This,\r
+  IN  EFI_PHYSICAL_ADDRESS                 BaseAddress,\r
+  IN  UINT64                               Length,\r
+  IN  UINT64                               Attributes\r
+  )\r
+{\r
+  return SmmClearMemoryAttributes (BaseAddress, Length, Attributes);\r
+}\r
+\r
+/**\r
+  This function retrieves the attributes of the memory region specified by\r
+  BaseAddress and Length. If different attributes are got from different part\r
+  of the memory region, EFI_NO_MAPPING will be returned.\r
+\r
+  @param  This              The EDKII_SMM_MEMORY_ATTRIBUTE_PROTOCOL instance.\r
+  @param  BaseAddress       The physical address that is the start address of\r
+                            a memory region.\r
+  @param  Length            The size in bytes of the memory region.\r
+  @param  Attributes        Pointer to attributes returned.\r
+\r
+  @retval EFI_SUCCESS           The attributes got for the memory region.\r
+  @retval EFI_INVALID_PARAMETER Length is zero.\r
+                                Attributes is NULL.\r
+  @retval EFI_NO_MAPPING        Attributes are not consistent cross the memory\r
+                                region.\r
+  @retval EFI_UNSUPPORTED       The processor does not support one or more\r
+                                bytes of the memory resource range specified\r
+                                by BaseAddress and Length.\r
+\r
+**/\r
+EFI_STATUS\r
+EFIAPI\r
+EdkiiSmmGetMemoryAttributes (\r
+  IN  EDKII_SMM_MEMORY_ATTRIBUTE_PROTOCOL  *This,\r
+  IN  EFI_PHYSICAL_ADDRESS                 BaseAddress,\r
+  IN  UINT64                               Length,\r
+  OUT UINT64                               *Attributes\r
+  )\r
+{\r
+  EFI_PHYSICAL_ADDRESS  Address;\r
+  UINT64                *PageEntry;\r
+  UINT64                MemAttr;\r
+  PAGE_ATTRIBUTE        PageAttr;\r
+  INT64                 Size;\r
+  UINTN                 PageTableBase;\r
+  BOOLEAN               EnablePML5Paging;\r
+  IA32_CR4              Cr4;\r
+\r
+  if ((Length < SIZE_4KB) || (Attributes == NULL)) {\r
+    return EFI_INVALID_PARAMETER;\r
+  }\r
+\r
+  Size    = (INT64)Length;\r
+  MemAttr = (UINT64)-1;\r
+\r
+  PageTableBase    = AsmReadCr3 () & PAGING_4K_ADDRESS_MASK_64;\r
+  Cr4.UintN        = AsmReadCr4 ();\r
+  EnablePML5Paging = (BOOLEAN)(Cr4.Bits.LA57 == 1);\r
+\r
+  do {\r
+    PageEntry = GetPageTableEntry (PageTableBase, EnablePML5Paging, BaseAddress, &PageAttr);\r
+    if ((PageEntry == NULL) || (PageAttr == PageNone)) {\r
+      return EFI_UNSUPPORTED;\r
+    }\r
+\r
+    //\r
+    // If the memory range is cross page table boundary, make sure they\r
+    // share the same attribute. Return EFI_NO_MAPPING if not.\r
+    //\r
+    *Attributes = GetAttributesFromPageEntry (PageEntry);\r
+    if ((MemAttr != (UINT64)-1) && (*Attributes != MemAttr)) {\r
+      return EFI_NO_MAPPING;\r
+    }\r
+\r
+    switch (PageAttr) {\r
+      case Page4K:\r
+        Address      = *PageEntry & ~mAddressEncMask & PAGING_4K_ADDRESS_MASK_64;\r
+        Size        -= (SIZE_4KB - (BaseAddress - Address));\r
+        BaseAddress += (SIZE_4KB - (BaseAddress - Address));\r
+        break;\r
+\r
+      case Page2M:\r
+        Address      = *PageEntry & ~mAddressEncMask & PAGING_2M_ADDRESS_MASK_64;\r
+        Size        -= SIZE_2MB - (BaseAddress - Address);\r
+        BaseAddress += SIZE_2MB - (BaseAddress - Address);\r
+        break;\r
+\r
+      case Page1G:\r
+        Address      = *PageEntry & ~mAddressEncMask & PAGING_1G_ADDRESS_MASK_64;\r
+        Size        -= SIZE_1GB - (BaseAddress - Address);\r
+        BaseAddress += SIZE_1GB - (BaseAddress - Address);\r
+        break;\r
+\r
+      default:\r
+        return EFI_UNSUPPORTED;\r
+    }\r
+\r
+    MemAttr = *Attributes;\r
+  } while (Size > 0);\r
+\r
+  return EFI_SUCCESS;\r
+}\r
+\r
+/**\r
+  Prevent the memory pages used for SMM page table from been overwritten.\r
+**/\r
+VOID\r
+EnablePageTableProtection (\r
+  VOID\r
+  )\r
+{\r
+  PAGE_TABLE_POOL       *HeadPool;\r
+  PAGE_TABLE_POOL       *Pool;\r
+  UINT64                PoolSize;\r
+  EFI_PHYSICAL_ADDRESS  Address;\r
+  UINTN                 PageTableBase;\r
+\r
+  if (mPageTablePool == NULL) {\r
+    return;\r
+  }\r
+\r
+  PageTableBase = AsmReadCr3 () & PAGING_4K_ADDRESS_MASK_64;\r
+\r
+  //\r
+  // ConvertMemoryPageAttributes might update mPageTablePool. It's safer to\r
+  // remember original one in advance.\r
+  //\r
+  HeadPool = mPageTablePool;\r
+  Pool     = HeadPool;\r
+  do {\r
+    Address  = (EFI_PHYSICAL_ADDRESS)(UINTN)Pool;\r
+    PoolSize = Pool->Offset + EFI_PAGES_TO_SIZE (Pool->FreePages);\r
+    //\r
+    // Set entire pool including header, used-memory and left free-memory as ReadOnly in SMM page table.\r
+    //\r
+    ConvertMemoryPageAttributes (PageTableBase, m5LevelPagingNeeded, Address, PoolSize, EFI_MEMORY_RO, TRUE, NULL, NULL);\r
+    Pool = Pool->NextPool;\r
+  } while (Pool != HeadPool);\r
+}\r
+\r
+/**\r
+  Return whether memory used by SMM page table need to be set as Read Only.\r
+\r
+  @retval TRUE  Need to set SMM page table as Read Only.\r
+  @retval FALSE Do not set SMM page table as Read Only.\r
+**/\r
+BOOLEAN\r
+IfReadOnlyPageTableNeeded (\r
+  VOID\r
+  )\r
+{\r
+  //\r
+  // Don't mark page table memory as read-only if\r
+  //  - no restriction on access to non-SMRAM memory; or\r
+  //  - SMM heap guard feature enabled; or\r
+  //      BIT2: SMM page guard enabled\r
+  //      BIT3: SMM pool guard enabled\r
+  //  - SMM profile feature enabled\r
+  //\r
+  if (!IsRestrictedMemoryAccess () ||\r
+      ((PcdGet8 (PcdHeapGuardPropertyMask) & (BIT3 | BIT2)) != 0) ||\r
+      FeaturePcdGet (PcdCpuSmmProfileEnable))\r
+  {\r
+    if (sizeof (UINTN) == sizeof (UINT64)) {\r
+      //\r
+      // Restriction on access to non-SMRAM memory and heap guard could not be enabled at the same time.\r
+      //\r
+      ASSERT (\r
+        !(IsRestrictedMemoryAccess () &&\r
+          (PcdGet8 (PcdHeapGuardPropertyMask) & (BIT3 | BIT2)) != 0)\r
+        );\r
+\r
+      //\r
+      // Restriction on access to non-SMRAM memory and SMM profile could not be enabled at the same time.\r
+      //\r
+      ASSERT (!(IsRestrictedMemoryAccess () && FeaturePcdGet (PcdCpuSmmProfileEnable)));\r
+    }\r
+\r
+    return FALSE;\r
+  }\r
+\r
+  return TRUE;\r
+}\r
+\r
+/**\r
+  This function sets memory attribute for page table.\r
+**/\r
+VOID\r
+SetPageTableAttributes (\r
+  VOID\r
+  )\r
+{\r
+  BOOLEAN  CetEnabled;\r
+\r
+  if (!IfReadOnlyPageTableNeeded ()) {\r
+    return;\r
+  }\r
+\r
+  DEBUG ((DEBUG_INFO, "SetPageTableAttributes\n"));\r
+\r
+  //\r
+  // Disable write protection, because we need mark page table to be write protected.\r
+  // We need *write* page table memory, to mark itself to be *read only*.\r
+  //\r
+  CetEnabled = ((AsmReadCr4 () & CR4_CET_ENABLE) != 0) ? TRUE : FALSE;\r
+  if (CetEnabled) {\r
+    //\r
+    // CET must be disabled if WP is disabled.\r
+    //\r
+    DisableCet ();\r
+  }\r
+\r
+  AsmWriteCr0 (AsmReadCr0 () & ~CR0_WP);\r
+\r
+  // Set memory used by page table as Read Only.\r
+  DEBUG ((DEBUG_INFO, "Start...\n"));\r
+  EnablePageTableProtection ();\r
+\r
+  //\r
+  // Enable write protection, after page table attribute updated.\r
+  //\r
+  AsmWriteCr0 (AsmReadCr0 () | CR0_WP);\r
+  mIsReadOnlyPageTable = TRUE;\r
+\r
+  //\r
+  // Flush TLB after mark all page table pool as read only.\r
+  //\r
+  FlushTlbForAll ();\r
+\r
+  if (CetEnabled) {\r
+    //\r
+    // re-enable CET.\r
+    //\r
+    EnableCet ();\r
+  }\r
+\r
+  return;\r
+}\r