UefiCpuPkg/PiSmmCpuDxeSmm: Add support for PCD PcdPteMemoryEncryptionAddressOrMask

[mirror_edk2.git] / UefiCpuPkg / PiSmmCpuDxeSmm / SmmCpuMemoryManagement.c

diff --git a/UefiCpuPkg/PiSmmCpuDxeSmm/SmmCpuMemoryManagement.c b/UefiCpuPkg/PiSmmCpuDxeSmm/SmmCpuMemoryManagement.c
index 13323d5409d555c881b0990b38461bdedceb8c05..a535389c26cebb6fe2ba068d930a4b671e4204c7 100644 (file)
--- a/UefiCpuPkg/PiSmmCpuDxeSmm/SmmCpuMemoryManagement.c
+++ b/UefiCpuPkg/PiSmmCpuDxeSmm/SmmCpuMemoryManagement.c
@@ -119,7 +119,7 @@ GetPageTableEntry (
       return NULL;\r
     }\r
 \r
-    L3PageTable = (UINT64 *)(UINTN)(L4PageTable[Index4] & PAGING_4K_ADDRESS_MASK_64);\r
+    L3PageTable = (UINT64 *)(UINTN)(L4PageTable[Index4] & ~mAddressEncMask & PAGING_4K_ADDRESS_MASK_64);\r
   } else {\r
     L3PageTable = (UINT64 *)GetPageTableBase ();\r
   }\r
@@ -133,7 +133,7 @@ GetPageTableEntry (
     return &L3PageTable[Index3];\r
   }\r
 \r
-  L2PageTable = (UINT64 *)(UINTN)(L3PageTable[Index3] & PAGING_4K_ADDRESS_MASK_64);\r
+  L2PageTable = (UINT64 *)(UINTN)(L3PageTable[Index3] & ~mAddressEncMask & PAGING_4K_ADDRESS_MASK_64);\r
   if (L2PageTable[Index2] == 0) {\r
     *PageAttribute = PageNone;\r
     return NULL;\r
@@ -145,7 +145,7 @@ GetPageTableEntry (
   }\r
 \r
   // 4k\r
-  L1PageTable = (UINT64 *)(UINTN)(L2PageTable[Index2] & PAGING_4K_ADDRESS_MASK_64);\r
+  L1PageTable = (UINT64 *)(UINTN)(L2PageTable[Index2] & ~mAddressEncMask & PAGING_4K_ADDRESS_MASK_64);\r
   if ((L1PageTable[Index1] == 0) && (Address != 0)) {\r
     *PageAttribute = PageNone;\r
     return NULL;\r
@@ -304,9 +304,9 @@ SplitPage (
       }\r
       BaseAddress = *PageEntry & PAGING_2M_ADDRESS_MASK_64;\r
       for (Index = 0; Index < SIZE_4KB / sizeof(UINT64); Index++) {\r
-        NewPageEntry[Index] = BaseAddress + SIZE_4KB * Index + ((*PageEntry) & PAGE_PROGATE_BITS);\r
+        NewPageEntry[Index] = (BaseAddress + SIZE_4KB * Index) | mAddressEncMask | ((*PageEntry) & PAGE_PROGATE_BITS);\r
       }\r
-      (*PageEntry) = (UINT64)(UINTN)NewPageEntry + PAGE_ATTRIBUTE_BITS;\r
+      (*PageEntry) = (UINT64)(UINTN)NewPageEntry | mAddressEncMask | PAGE_ATTRIBUTE_BITS;\r
       return RETURN_SUCCESS;\r
     } else {\r
       return RETURN_UNSUPPORTED;\r
@@ -325,9 +325,9 @@ SplitPage (
       }\r
       BaseAddress = *PageEntry & PAGING_1G_ADDRESS_MASK_64;\r
       for (Index = 0; Index < SIZE_4KB / sizeof(UINT64); Index++) {\r
-        NewPageEntry[Index] = BaseAddress + SIZE_2MB * Index + IA32_PG_PS + ((*PageEntry) & PAGE_PROGATE_BITS);\r
+        NewPageEntry[Index] = (BaseAddress + SIZE_2MB * Index) | mAddressEncMask | IA32_PG_PS | ((*PageEntry) & PAGE_PROGATE_BITS);\r
       }\r
-      (*PageEntry) = (UINT64)(UINTN)NewPageEntry + PAGE_ATTRIBUTE_BITS;\r
+      (*PageEntry) = (UINT64)(UINTN)NewPageEntry | mAddressEncMask | PAGE_ATTRIBUTE_BITS;\r
       return RETURN_SUCCESS;\r
     } else {\r
       return RETURN_UNSUPPORTED;\r