]> git.proxmox.com Git - mirror_edk2.git/blobdiff - OvmfPkg/Library/BaseMemEncryptSevLib/X64/VirtualMemory.c
OvmfPkg/MemEncryptSevLib: rewrap to 79 characters width
[mirror_edk2.git] / OvmfPkg / Library / BaseMemEncryptSevLib / X64 / VirtualMemory.c
index dbaad7766dbe6b5d99d7f2d47507d5e6ffb45036..65b8babaac44b6ef3fe38cb0d29bd66f37393114 100644 (file)
@@ -2,18 +2,18 @@
 \r
   Virtual Memory Management Services to set or clear the memory encryption bit\r
 \r
-Copyright (c) 2006 - 2016, Intel Corporation. All rights reserved.<BR>\r
-Copyright (c) 2017, AMD Incorporated. All rights reserved.<BR>\r
+  Copyright (c) 2006 - 2016, Intel Corporation. All rights reserved.<BR>\r
+  Copyright (c) 2017, AMD Incorporated. All rights reserved.<BR>\r
 \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
+  This program and the accompanying materials are licensed and made available\r
+  under the terms and conditions of the BSD License which accompanies this\r
+  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
+  THE PROGRAM IS DISTRIBUTED UNDER THE BSD LICENSE ON AN "AS IS" BASIS, WITHOUT\r
+  WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.\r
 \r
-Code is derived from MdeModulePkg/Core/DxeIplPeim/X64/VirtualMemory.c\r
+  Code is derived from MdeModulePkg/Core/DxeIplPeim/X64/VirtualMemory.c\r
 \r
 **/\r
 \r
@@ -25,6 +25,7 @@ Code is derived from MdeModulePkg/Core/DxeIplPeim/X64/VirtualMemory.c
 \r
 STATIC BOOLEAN mAddressEncMaskChecked = FALSE;\r
 STATIC UINT64  mAddressEncMask;\r
+STATIC PAGE_TABLE_POOL   *mPageTablePool = NULL;\r
 \r
 typedef enum {\r
    SetCBit,\r
@@ -62,10 +63,129 @@ GetMemEncryptionAddressMask (
   return mAddressEncMask;\r
 }\r
 \r
+/**\r
+  Initialize a buffer pool for page table use only.\r
+\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\r
+  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. Usually this won't\r
+  happen often in practice.\r
+\r
+  @param[in] 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
+STATIC\r
+BOOLEAN\r
+InitializePageTablePool (\r
+  IN  UINTN                           PoolPages\r
+  )\r
+{\r
+  VOID                      *Buffer;\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
+  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 than once to allocate memory for page tables.\r
+\r
+  Allocates the number of 4KB pages and returns a pointer to the allocated\r
+  buffer. The buffer returned is aligned on a 4KB boundary.\r
+\r
+  If Pages is 0, then NULL is returned.\r
+  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
+STATIC\r
+VOID *\r
+EFIAPI\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
+    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
+  DEBUG ((\r
+    DEBUG_VERBOSE,\r
+    "%a:%a: Buffer=0x%Lx Pages=%ld\n",\r
+    gEfiCallerBaseName,\r
+    __FUNCTION__,\r
+    Buffer,\r
+    Pages\r
+    ));\r
+\r
+  return Buffer;\r
+}\r
+\r
+\r
 /**\r
   Split 2M page to 4K.\r
 \r
-  @param[in]      PhysicalAddress       Start physical address the 2M page covered.\r
+  @param[in]      PhysicalAddress       Start physical address the 2M page\r
+                                        covered.\r
   @param[in, out] PageEntry2M           Pointer to 2M page entry.\r
   @param[in]      StackBase             Stack base address.\r
   @param[in]      StackSize             Stack size.\r
@@ -85,7 +205,7 @@ Split2MPageTo4K (
   PAGE_TABLE_4K_ENTRY               *PageTableEntry, *PageTableEntry1;\r
   UINT64                            AddressEncMask;\r
 \r
-  PageTableEntry = AllocatePages(1);\r
+  PageTableEntry = AllocatePageTableMemory(1);\r
 \r
   PageTableEntry1 = PageTableEntry;\r
 \r
@@ -95,14 +215,19 @@ Split2MPageTo4K (
   ASSERT (*PageEntry2M & AddressEncMask);\r
 \r
   PhysicalAddress4K = PhysicalAddress;\r
-  for (IndexOfPageTableEntries = 0; IndexOfPageTableEntries < 512; IndexOfPageTableEntries++, PageTableEntry++, PhysicalAddress4K += SIZE_4KB) {\r
+  for (IndexOfPageTableEntries = 0;\r
+       IndexOfPageTableEntries < 512;\r
+       (IndexOfPageTableEntries++,\r
+        PageTableEntry++,\r
+        PhysicalAddress4K += SIZE_4KB)) {\r
     //\r
     // Fill in the Page Table entries\r
     //\r
     PageTableEntry->Uint64 = (UINT64) PhysicalAddress4K | AddressEncMask;\r
     PageTableEntry->Bits.ReadWrite = 1;\r
     PageTableEntry->Bits.Present = 1;\r
-    if ((PhysicalAddress4K >= StackBase) && (PhysicalAddress4K < StackBase + StackSize)) {\r
+    if ((PhysicalAddress4K >= StackBase) &&\r
+        (PhysicalAddress4K < StackBase + StackSize)) {\r
       //\r
       // Set Nx bit for stack.\r
       //\r
@@ -113,13 +238,188 @@ Split2MPageTo4K (
   //\r
   // Fill in 2M page entry.\r
   //\r
-  *PageEntry2M = (UINT64) (UINTN) PageTableEntry1 | IA32_PG_P | IA32_PG_RW | AddressEncMask;\r
+  *PageEntry2M = ((UINT64)(UINTN)PageTableEntry1 |\r
+                  IA32_PG_P | IA32_PG_RW | AddressEncMask);\r
 }\r
 \r
+/**\r
+  Set one page of page table pool memory to be read-only.\r
+\r
+  @param[in] PageTableBase    Base address of page table (CR3).\r
+  @param[in] Address          Start address of a page to be set as read-only.\r
+  @param[in] Level4Paging     Level 4 paging flag.\r
+\r
+**/\r
+STATIC\r
+VOID\r
+SetPageTablePoolReadOnly (\r
+  IN  UINTN                             PageTableBase,\r
+  IN  EFI_PHYSICAL_ADDRESS              Address,\r
+  IN  BOOLEAN                           Level4Paging\r
+  )\r
+{\r
+  UINTN                 Index;\r
+  UINTN                 EntryIndex;\r
+  UINT64                AddressEncMask;\r
+  EFI_PHYSICAL_ADDRESS  PhysicalAddress;\r
+  UINT64                *PageTable;\r
+  UINT64                *NewPageTable;\r
+  UINT64                PageAttr;\r
+  UINT64                LevelSize[5];\r
+  UINT64                LevelMask[5];\r
+  UINTN                 LevelShift[5];\r
+  UINTN                 Level;\r
+  UINT64                PoolUnitSize;\r
+\r
+  ASSERT (PageTableBase != 0);\r
+\r
+  //\r
+  // Since the page table is always from page table pool, which is always\r
+  // located at the boundary of PcdPageTablePoolAlignment, we just need to\r
+  // set the whole pool unit to be read-only.\r
+  //\r
+  Address = Address & PAGE_TABLE_POOL_ALIGN_MASK;\r
+\r
+  LevelShift[1] = PAGING_L1_ADDRESS_SHIFT;\r
+  LevelShift[2] = PAGING_L2_ADDRESS_SHIFT;\r
+  LevelShift[3] = PAGING_L3_ADDRESS_SHIFT;\r
+  LevelShift[4] = PAGING_L4_ADDRESS_SHIFT;\r
+\r
+  LevelMask[1] = PAGING_4K_ADDRESS_MASK_64;\r
+  LevelMask[2] = PAGING_2M_ADDRESS_MASK_64;\r
+  LevelMask[3] = PAGING_1G_ADDRESS_MASK_64;\r
+  LevelMask[4] = PAGING_1G_ADDRESS_MASK_64;\r
+\r
+  LevelSize[1] = SIZE_4KB;\r
+  LevelSize[2] = SIZE_2MB;\r
+  LevelSize[3] = SIZE_1GB;\r
+  LevelSize[4] = SIZE_512GB;\r
+\r
+  AddressEncMask  = GetMemEncryptionAddressMask() &\r
+                    PAGING_1G_ADDRESS_MASK_64;\r
+  PageTable       = (UINT64 *)(UINTN)PageTableBase;\r
+  PoolUnitSize    = PAGE_TABLE_POOL_UNIT_SIZE;\r
+\r
+  for (Level = (Level4Paging) ? 4 : 3; Level > 0; --Level) {\r
+    Index = ((UINTN)RShiftU64 (Address, LevelShift[Level]));\r
+    Index &= PAGING_PAE_INDEX_MASK;\r
+\r
+    PageAttr = PageTable[Index];\r
+    if ((PageAttr & IA32_PG_PS) == 0) {\r
+      //\r
+      // Go to next level of table.\r
+      //\r
+      PageTable = (UINT64 *)(UINTN)(PageAttr & ~AddressEncMask &\r
+                                    PAGING_4K_ADDRESS_MASK_64);\r
+      continue;\r
+    }\r
+\r
+    if (PoolUnitSize >= LevelSize[Level]) {\r
+      //\r
+      // Clear R/W bit if current page granularity is not larger than pool unit\r
+      // size.\r
+      //\r
+      if ((PageAttr & IA32_PG_RW) != 0) {\r
+        while (PoolUnitSize > 0) {\r
+          //\r
+          // PAGE_TABLE_POOL_UNIT_SIZE and PAGE_TABLE_POOL_ALIGNMENT are fit in\r
+          // one page (2MB). Then we don't need to update attributes for pages\r
+          // crossing page directory. ASSERT below is for that purpose.\r
+          //\r
+          ASSERT (Index < EFI_PAGE_SIZE/sizeof (UINT64));\r
+\r
+          PageTable[Index] &= ~(UINT64)IA32_PG_RW;\r
+          PoolUnitSize    -= LevelSize[Level];\r
+\r
+          ++Index;\r
+        }\r
+      }\r
+\r
+      break;\r
+\r
+    } else {\r
+      //\r
+      // The smaller granularity of page must be needed.\r
+      //\r
+      ASSERT (Level > 1);\r
+\r
+      NewPageTable = AllocatePageTableMemory (1);\r
+      ASSERT (NewPageTable != NULL);\r
+\r
+      PhysicalAddress = PageAttr & LevelMask[Level];\r
+      for (EntryIndex = 0;\r
+            EntryIndex < EFI_PAGE_SIZE/sizeof (UINT64);\r
+            ++EntryIndex) {\r
+        NewPageTable[EntryIndex] = PhysicalAddress  | AddressEncMask |\r
+                                   IA32_PG_P | IA32_PG_RW;\r
+        if (Level > 2) {\r
+          NewPageTable[EntryIndex] |= IA32_PG_PS;\r
+        }\r
+        PhysicalAddress += LevelSize[Level - 1];\r
+      }\r
+\r
+      PageTable[Index] = (UINT64)(UINTN)NewPageTable | AddressEncMask |\r
+                                        IA32_PG_P | IA32_PG_RW;\r
+      PageTable = NewPageTable;\r
+    }\r
+  }\r
+}\r
+\r
+/**\r
+  Prevent the memory pages used for page table from been overwritten.\r
+\r
+  @param[in] PageTableBase    Base address of page table (CR3).\r
+  @param[in] Level4Paging     Level 4 paging flag.\r
+\r
+**/\r
+STATIC\r
+VOID\r
+EnablePageTableProtection (\r
+  IN  UINTN     PageTableBase,\r
+  IN  BOOLEAN   Level4Paging\r
+  )\r
+{\r
+  PAGE_TABLE_POOL         *HeadPool;\r
+  PAGE_TABLE_POOL         *Pool;\r
+  UINT64                  PoolSize;\r
+  EFI_PHYSICAL_ADDRESS    Address;\r
+\r
+  if (mPageTablePool == NULL) {\r
+    return;\r
+  }\r
+\r
+  //\r
+  // SetPageTablePoolReadOnly 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
+    //\r
+    // The size of one pool must be multiple of PAGE_TABLE_POOL_UNIT_SIZE,\r
+    // which is one of page size of the processor (2MB by default). Let's apply\r
+    // the protection to them one by one.\r
+    //\r
+    while (PoolSize > 0) {\r
+      SetPageTablePoolReadOnly(PageTableBase, Address, Level4Paging);\r
+      Address   += PAGE_TABLE_POOL_UNIT_SIZE;\r
+      PoolSize  -= PAGE_TABLE_POOL_UNIT_SIZE;\r
+    }\r
+\r
+    Pool = Pool->NextPool;\r
+  } while (Pool != HeadPool);\r
+\r
+}\r
+\r
+\r
 /**\r
   Split 1G page to 2M.\r
 \r
-  @param[in]      PhysicalAddress       Start physical address the 1G page covered.\r
+  @param[in]      PhysicalAddress       Start physical address the 1G page\r
+                                        covered.\r
   @param[in, out] PageEntry1G           Pointer to 1G page entry.\r
   @param[in]      StackBase             Stack base address.\r
   @param[in]      StackSize             Stack size.\r
@@ -139,7 +439,7 @@ Split1GPageTo2M (
   PAGE_TABLE_ENTRY                  *PageDirectoryEntry;\r
   UINT64                            AddressEncMask;\r
 \r
-  PageDirectoryEntry = AllocatePages(1);\r
+  PageDirectoryEntry = AllocatePageTableMemory(1);\r
 \r
   AddressEncMask = GetMemEncryptionAddressMask ();\r
   ASSERT (PageDirectoryEntry != NULL);\r
@@ -147,15 +447,26 @@ Split1GPageTo2M (
   //\r
   // Fill in 1G page entry.\r
   //\r
-  *PageEntry1G = (UINT64) (UINTN) PageDirectoryEntry | IA32_PG_P | IA32_PG_RW | AddressEncMask;\r
+  *PageEntry1G = ((UINT64)(UINTN)PageDirectoryEntry |\r
+                  IA32_PG_P | IA32_PG_RW | AddressEncMask);\r
 \r
   PhysicalAddress2M = PhysicalAddress;\r
-  for (IndexOfPageDirectoryEntries = 0; IndexOfPageDirectoryEntries < 512; IndexOfPageDirectoryEntries++, PageDirectoryEntry++, PhysicalAddress2M += SIZE_2MB) {\r
-    if ((PhysicalAddress2M < StackBase + StackSize) && ((PhysicalAddress2M + SIZE_2MB) > StackBase)) {\r
+  for (IndexOfPageDirectoryEntries = 0;\r
+       IndexOfPageDirectoryEntries < 512;\r
+       (IndexOfPageDirectoryEntries++,\r
+        PageDirectoryEntry++,\r
+        PhysicalAddress2M += SIZE_2MB)) {\r
+    if ((PhysicalAddress2M < StackBase + StackSize) &&\r
+        ((PhysicalAddress2M + SIZE_2MB) > StackBase)) {\r
       //\r
       // Need to split this 2M page that covers stack range.\r
       //\r
-      Split2MPageTo4K (PhysicalAddress2M, (UINT64 *) PageDirectoryEntry, StackBase, StackSize);\r
+      Split2MPageTo4K (\r
+        PhysicalAddress2M,\r
+        (UINT64 *)PageDirectoryEntry,\r
+        StackBase,\r
+        StackSize\r
+        );\r
     } else {\r
       //\r
       // Fill in the Page Directory entries\r
@@ -195,8 +506,50 @@ SetOrClearCBit(
 }\r
 \r
 /**\r
-  This function either sets or clears memory encryption bit for the memory region\r
-  specified by PhysicalAddress and length from the current page table context.\r
+ Check the WP status in CR0 register. This bit is used to lock or unlock write\r
+ access to pages marked as read-only.\r
+\r
+  @retval TRUE    Write protection is enabled.\r
+  @retval FALSE   Write protection is disabled.\r
+**/\r
+STATIC\r
+BOOLEAN\r
+IsReadOnlyPageWriteProtected (\r
+  VOID\r
+  )\r
+{\r
+  return ((AsmReadCr0 () & BIT16) != 0);\r
+}\r
+\r
+\r
+/**\r
+ Disable Write Protect on pages marked as read-only.\r
+**/\r
+STATIC\r
+VOID\r
+DisableReadOnlyPageWriteProtect (\r
+  VOID\r
+  )\r
+{\r
+  AsmWriteCr0 (AsmReadCr0() & ~BIT16);\r
+}\r
+\r
+/**\r
+ Enable Write Protect on pages marked as read-only.\r
+**/\r
+VOID\r
+EnableReadOnlyPageWriteProtect (\r
+  VOID\r
+  )\r
+{\r
+  AsmWriteCr0 (AsmReadCr0() | BIT16);\r
+}\r
+\r
+\r
+/**\r
+  This function either sets or clears memory encryption bit for the memory\r
+  region specified by PhysicalAddress and length from the current page table\r
+  context.\r
 \r
   The function iterates through the physicalAddress one page at a time, and set\r
   or clears the memory encryption mask in the page table. If it encounters\r
@@ -212,11 +565,11 @@ SetOrClearCBit(
   @param[in]  Flush                   Flush the caches before applying the\r
                                       encryption mask\r
 \r
-  @retval RETURN_SUCCESS              The attributes were cleared for the memory\r
-                                      region.\r
+  @retval RETURN_SUCCESS              The attributes were cleared for the\r
+                                      memory region.\r
   @retval RETURN_INVALID_PARAMETER    Number of pages is zero.\r
-  @retval RETURN_UNSUPPORTED          Setting the memory encyrption attribute is\r
-                                      not supported\r
+  @retval RETURN_UNSUPPORTED          Setting the memory encyrption attribute\r
+                                      is not supported\r
 **/\r
 \r
 STATIC\r
@@ -238,6 +591,8 @@ SetMemoryEncDec (
   PAGE_TABLE_4K_ENTRY            *PageTableEntry;\r
   UINT64                         PgTableMask;\r
   UINT64                         AddressEncMask;\r
+  BOOLEAN                        IsWpEnabled;\r
+  RETURN_STATUS                  Status;\r
 \r
   DEBUG ((\r
     DEBUG_VERBOSE,\r
@@ -267,13 +622,23 @@ SetMemoryEncDec (
 \r
   //\r
   // We are going to change the memory encryption attribute from C=0 -> C=1 or\r
-  // vice versa Flush the caches to ensure that data is written into memory with\r
-  // correct C-bit\r
+  // vice versa Flush the caches to ensure that data is written into memory\r
+  // with correct C-bit\r
   //\r
   if (CacheFlush) {\r
     WriteBackInvalidateDataCacheRange((VOID*) (UINTN)PhysicalAddress, Length);\r
   }\r
 \r
+  //\r
+  // Make sure that the page table is changeable.\r
+  //\r
+  IsWpEnabled = IsReadOnlyPageWriteProtected ();\r
+  if (IsWpEnabled) {\r
+    DisableReadOnlyPageWriteProtect ();\r
+  }\r
+\r
+  Status = EFI_SUCCESS;\r
+\r
   while (Length)\r
   {\r
     //\r
@@ -286,19 +651,32 @@ SetMemoryEncDec (
     PageMapLevel4Entry = (VOID*) (Cr3BaseAddress & ~PgTableMask);\r
     PageMapLevel4Entry += PML4_OFFSET(PhysicalAddress);\r
     if (!PageMapLevel4Entry->Bits.Present) {\r
-      DEBUG ((DEBUG_WARN,\r
-        "%a:%a ERROR bad PML4 for %lx\n", gEfiCallerBaseName, __FUNCTION__,\r
-        PhysicalAddress));\r
-      return RETURN_NO_MAPPING;\r
+      DEBUG ((\r
+        DEBUG_ERROR,\r
+        "%a:%a: bad PML4 for Physical=0x%Lx\n",\r
+        gEfiCallerBaseName,\r
+        __FUNCTION__,\r
+        PhysicalAddress\r
+        ));\r
+      Status = RETURN_NO_MAPPING;\r
+      goto Done;\r
     }\r
 \r
-    PageDirectory1GEntry = (VOID*) ((PageMapLevel4Entry->Bits.PageTableBaseAddress<<12) & ~PgTableMask);\r
+    PageDirectory1GEntry = (VOID *)(\r
+                             (PageMapLevel4Entry->Bits.PageTableBaseAddress <<\r
+                              12) & ~PgTableMask\r
+                             );\r
     PageDirectory1GEntry += PDP_OFFSET(PhysicalAddress);\r
     if (!PageDirectory1GEntry->Bits.Present) {\r
-      DEBUG ((DEBUG_WARN,\r
-        "%a:%a ERROR bad PDPE for %lx\n", gEfiCallerBaseName,\r
-         __FUNCTION__, PhysicalAddress));\r
-      return RETURN_NO_MAPPING;\r
+      DEBUG ((\r
+        DEBUG_ERROR,\r
+        "%a:%a: bad PDPE for Physical=0x%Lx\n",\r
+        gEfiCallerBaseName,\r
+        __FUNCTION__,\r
+        PhysicalAddress\r
+        ));\r
+      Status = RETURN_NO_MAPPING;\r
+      goto Done;\r
     }\r
 \r
     //\r
@@ -311,32 +689,56 @@ SetMemoryEncDec (
       //\r
       if (!(PhysicalAddress & (BIT30 - 1)) && Length >= BIT30) {\r
         SetOrClearCBit(&PageDirectory1GEntry->Uint64, Mode);\r
-        DEBUG ((DEBUG_VERBOSE,\r
-          "%a:%a Updated 1GB entry for %lx\n", gEfiCallerBaseName,\r
-          __FUNCTION__, PhysicalAddress));\r
+        DEBUG ((\r
+          DEBUG_VERBOSE,\r
+          "%a:%a: updated 1GB entry for Physical=0x%Lx\n",\r
+          gEfiCallerBaseName,\r
+          __FUNCTION__,\r
+          PhysicalAddress\r
+          ));\r
         PhysicalAddress += BIT30;\r
         Length -= BIT30;\r
       } else {\r
         //\r
         // We must split the page\r
         //\r
-        DEBUG ((DEBUG_VERBOSE,\r
-          "%a:%a Spliting 1GB page\n", gEfiCallerBaseName, __FUNCTION__));\r
-        Split1GPageTo2M(((UINT64)PageDirectory1GEntry->Bits.PageTableBaseAddress)<<30, (UINT64*) PageDirectory1GEntry, 0, 0);\r
+        DEBUG ((\r
+          DEBUG_VERBOSE,\r
+          "%a:%a: splitting 1GB page for Physical=0x%Lx\n",\r
+          gEfiCallerBaseName,\r
+          __FUNCTION__,\r
+          PhysicalAddress\r
+          ));\r
+        Split1GPageTo2M (\r
+          (UINT64)PageDirectory1GEntry->Bits.PageTableBaseAddress << 30,\r
+          (UINT64 *)PageDirectory1GEntry,\r
+          0,\r
+          0\r
+          );\r
         continue;\r
       }\r
     } else {\r
       //\r
       // Actually a PDP\r
       //\r
-      PageUpperDirectoryPointerEntry = (PAGE_MAP_AND_DIRECTORY_POINTER*) PageDirectory1GEntry;\r
-      PageDirectory2MEntry = (VOID*) ((PageUpperDirectoryPointerEntry->Bits.PageTableBaseAddress<<12) & ~PgTableMask);\r
+      PageUpperDirectoryPointerEntry =\r
+        (PAGE_MAP_AND_DIRECTORY_POINTER *)PageDirectory1GEntry;\r
+      PageDirectory2MEntry =\r
+        (VOID *)(\r
+          (PageUpperDirectoryPointerEntry->Bits.PageTableBaseAddress <<\r
+           12) & ~PgTableMask\r
+          );\r
       PageDirectory2MEntry += PDE_OFFSET(PhysicalAddress);\r
       if (!PageDirectory2MEntry->Bits.Present) {\r
-        DEBUG ((DEBUG_WARN,\r
-          "%a:%a ERROR bad PDE for %lx\n", gEfiCallerBaseName, __FUNCTION__,\r
-          PhysicalAddress));\r
-        return RETURN_NO_MAPPING;\r
+        DEBUG ((\r
+          DEBUG_ERROR,\r
+          "%a:%a: bad PDE for Physical=0x%Lx\n",\r
+          gEfiCallerBaseName,\r
+          __FUNCTION__,\r
+          PhysicalAddress\r
+          ));\r
+        Status = RETURN_NO_MAPPING;\r
+        goto Done;\r
       }\r
       //\r
       // If the MustBe1 bit is not a 1, it's not a 2MB entry\r
@@ -354,21 +756,40 @@ SetMemoryEncDec (
           //\r
           // We must split up this page into 4K pages\r
           //\r
-          DEBUG ((DEBUG_VERBOSE,\r
-            "%a:%a Spliting 2MB page at %lx\n", gEfiCallerBaseName,__FUNCTION__,\r
-            PhysicalAddress));\r
-          Split2MPageTo4K (((UINT64)PageDirectory2MEntry->Bits.PageTableBaseAddress) << 21, (UINT64*) PageDirectory2MEntry, 0, 0);\r
+          DEBUG ((\r
+            DEBUG_VERBOSE,\r
+            "%a:%a: splitting 2MB page for Physical=0x%Lx\n",\r
+            gEfiCallerBaseName,\r
+            __FUNCTION__,\r
+            PhysicalAddress\r
+            ));\r
+          Split2MPageTo4K (\r
+            (UINT64)PageDirectory2MEntry->Bits.PageTableBaseAddress << 21,\r
+            (UINT64 *)PageDirectory2MEntry,\r
+            0,\r
+            0\r
+            );\r
           continue;\r
         }\r
       } else {\r
-        PageDirectoryPointerEntry = (PAGE_MAP_AND_DIRECTORY_POINTER*) PageDirectory2MEntry;\r
-        PageTableEntry = (VOID*) (PageDirectoryPointerEntry->Bits.PageTableBaseAddress<<12 & ~PgTableMask);\r
+        PageDirectoryPointerEntry =\r
+          (PAGE_MAP_AND_DIRECTORY_POINTER *)PageDirectory2MEntry;\r
+        PageTableEntry =\r
+          (VOID *)(\r
+            (PageDirectoryPointerEntry->Bits.PageTableBaseAddress <<\r
+             12) & ~PgTableMask\r
+            );\r
         PageTableEntry += PTE_OFFSET(PhysicalAddress);\r
         if (!PageTableEntry->Bits.Present) {\r
-          DEBUG ((DEBUG_WARN,\r
-            "%a:%a ERROR bad PTE for %lx\n", gEfiCallerBaseName,\r
-            __FUNCTION__, PhysicalAddress));\r
-          return RETURN_NO_MAPPING;\r
+          DEBUG ((\r
+            DEBUG_ERROR,\r
+            "%a:%a: bad PTE for Physical=0x%Lx\n",\r
+            gEfiCallerBaseName,\r
+            __FUNCTION__,\r
+            PhysicalAddress\r
+            ));\r
+          Status = RETURN_NO_MAPPING;\r
+          goto Done;\r
         }\r
         SetOrClearCBit (&PageTableEntry->Uint64, Mode);\r
         PhysicalAddress += EFI_PAGE_SIZE;\r
@@ -377,12 +798,28 @@ SetMemoryEncDec (
     }\r
   }\r
 \r
+  //\r
+  // Protect the page table by marking the memory used for page table to be\r
+  // read-only.\r
+  //\r
+  if (IsWpEnabled) {\r
+    EnablePageTableProtection ((UINTN)PageMapLevel4Entry, TRUE);\r
+  }\r
+\r
   //\r
   // Flush TLB\r
   //\r
   CpuFlushTlb();\r
 \r
-  return RETURN_SUCCESS;\r
+Done:\r
+  //\r
+  // Restore page table write protection, if any.\r
+  //\r
+  if (IsWpEnabled) {\r
+    EnableReadOnlyPageWriteProtect ();\r
+  }\r
+\r
+  return Status;\r
 }\r
 \r
 /**\r
@@ -395,11 +832,11 @@ SetMemoryEncDec (
   @param[in]  Flush                   Flush the caches before applying the\r
                                       encryption mask\r
 \r
-  @retval RETURN_SUCCESS              The attributes were cleared for the memory\r
-                                      region.\r
+  @retval RETURN_SUCCESS              The attributes were cleared for the\r
+                                      memory region.\r
   @retval RETURN_INVALID_PARAMETER    Number of pages is zero.\r
-  @retval RETURN_UNSUPPORTED          Setting the memory encyrption attribute is\r
-                                      not supported\r
+  @retval RETURN_UNSUPPORTED          Setting the memory encyrption attribute\r
+                                      is not supported\r
 **/\r
 RETURN_STATUS\r
 EFIAPI\r
@@ -411,24 +848,30 @@ InternalMemEncryptSevSetMemoryDecrypted (
   )\r
 {\r
 \r
-  return SetMemoryEncDec (Cr3BaseAddress, PhysicalAddress, Length, ClearCBit, Flush);\r
+  return SetMemoryEncDec (\r
+           Cr3BaseAddress,\r
+           PhysicalAddress,\r
+           Length,\r
+           ClearCBit,\r
+           Flush\r
+           );\r
 }\r
 \r
 /**\r
   This function sets memory encryption bit for the memory region specified by\r
   PhysicalAddress and length from the current page table context.\r
 \r
-  @param[in]  PhysicalAddress         The physical address that is the start address\r
-                                      of a memory region.\r
+  @param[in]  PhysicalAddress         The physical address that is the start\r
+                                      address of a memory region.\r
   @param[in]  Length                  The length of memory region\r
   @param[in]  Flush                   Flush the caches before applying the\r
                                       encryption mask\r
 \r
-  @retval RETURN_SUCCESS              The attributes were cleared for the memory\r
-                                      region.\r
+  @retval RETURN_SUCCESS              The attributes were cleared for the\r
+                                      memory region.\r
   @retval RETURN_INVALID_PARAMETER    Number of pages is zero.\r
-  @retval RETURN_UNSUPPORTED          Setting the memory encyrption attribute is\r
-                                      not supported\r
+  @retval RETURN_UNSUPPORTED          Setting the memory encyrption attribute\r
+                                      is not supported\r
 **/\r
 RETURN_STATUS\r
 EFIAPI\r
@@ -439,5 +882,11 @@ InternalMemEncryptSevSetMemoryEncrypted (
   IN  BOOLEAN                 Flush\r
   )\r
 {\r
-  return SetMemoryEncDec (Cr3BaseAddress, PhysicalAddress, Length, SetCBit, Flush);\r
+  return SetMemoryEncDec (\r
+           Cr3BaseAddress,\r
+           PhysicalAddress,\r
+           Length,\r
+           SetCBit,\r
+           Flush\r
+           );\r
 }\r