[PATCH 2/3] MdeModulePkg/Universal/CapsulePei: Add support for PCD PcdPteMemoryEncryptionAddressOrMask

[Leo Duran <leo.duran@amd.com>]

This PCD holds the address mask for page table entries when memory
encryption is enabled on AMD processors supporting the Secure Encrypted
Virtualization (SEV) feature.

The mask is applied when 4GB tables are created (UefiCapsule.c), and when
the tables are expanded on-demand by page-faults above 4GB's (X64Entry.c).

Cc: Feng Tian <feng.tian@intel.com>
Cc: Star Zeng <star.zeng@intel.com>
Cc: Laszlo Ersek <lersek@redhat.com>
Contributed-under: TianoCore Contribution Agreement 1.0
Signed-off-by: Brijesh Singh <brijesh.singh@amd.com>
Signed-off-by: Leo Duran <leo.duran@amd.com>
---
 MdeModulePkg/Universal/CapsulePei/CapsulePei.inf |  2 ++
 MdeModulePkg/Universal/CapsulePei/CapsuleX64.inf |  4 ++++
 MdeModulePkg/Universal/CapsulePei/UefiCapsule.c  |  9 +++++----
 MdeModulePkg/Universal/CapsulePei/X64/X64Entry.c | 10 ++++++----
 4 files changed, 17 insertions(+), 8 deletions(-)

diff --git a/MdeModulePkg/Universal/CapsulePei/CapsulePei.inf b/MdeModulePkg/Universal/CapsulePei/CapsulePei.inf
index d2ca0d0..c54bc21 100644
--- a/MdeModulePkg/Universal/CapsulePei/CapsulePei.inf
+++ b/MdeModulePkg/Universal/CapsulePei/CapsulePei.inf
@@ -7,6 +7,7 @@
 #  buffer overflow, integer overflow.
 #
 # Copyright (c) 2006 - 2016, Intel Corporation. All rights reserved.<BR>
+# Copyright (c) 2017, AMD Incorporated. All rights reserved.<BR>
 #
 # This program and the accompanying materials
 # are licensed and made available under the terms and conditions
@@ -76,6 +77,7 @@ [Ppis.IA32]
 [Pcd.IA32]
   gEfiMdeModulePkgTokenSpaceGuid.PcdCapsuleCoalesceFile     ## SOMETIMES_CONSUMES
   gEfiMdeModulePkgTokenSpaceGuid.PcdUse1GPageTable          ## SOMETIMES_CONSUMES
+  gEfiMdeModulePkgTokenSpaceGuid.PcdPteMemoryEncryptionAddressOrMask    ## CONSUMES
 
 [FeaturePcd.IA32]
   gEfiMdeModulePkgTokenSpaceGuid.PcdDxeIplSwitchToLongMode  ## CONSUMES
diff --git a/MdeModulePkg/Universal/CapsulePei/CapsuleX64.inf b/MdeModulePkg/Universal/CapsulePei/CapsuleX64.inf
index 8318eaa..cf8543b 100644
--- a/MdeModulePkg/Universal/CapsulePei/CapsuleX64.inf
+++ b/MdeModulePkg/Universal/CapsulePei/CapsuleX64.inf
@@ -10,6 +10,7 @@
 #  buffer overflow, integer overflow.
 #
 # Copyright (c) 2011 - 2015, Intel Corporation. All rights reserved.<BR>
+# Copyright (c) 2017, AMD Incorporated. All rights reserved.<BR>
 #
 # This program and the accompanying materials
 # are licensed and made available under the terms and conditions
@@ -53,6 +54,9 @@ [LibraryClasses]
   CpuExceptionHandlerLib
   DebugAgentLib
 
+[Pcd]
+  gEfiMdeModulePkgTokenSpaceGuid.PcdPteMemoryEncryptionAddressOrMask    ## CONSUMES
+
 [Depex]
   FALSE
 
diff --git a/MdeModulePkg/Universal/CapsulePei/UefiCapsule.c b/MdeModulePkg/Universal/CapsulePei/UefiCapsule.c
index 9ac9d22..77bc68c 100644
--- a/MdeModulePkg/Universal/CapsulePei/UefiCapsule.c
+++ b/MdeModulePkg/Universal/CapsulePei/UefiCapsule.c
@@ -2,6 +2,7 @@
   Capsule update PEIM for UEFI2.0
 
 Copyright (c) 2006 - 2016, Intel Corporation. All rights reserved.<BR>
+Copyright (c) 2017, AMD Incorporated. All rights reserved.<BR>
 
 This program and the accompanying materials
 are licensed and made available under the terms and conditions
@@ -187,7 +188,7 @@ Create4GPageTables (
     //
     // Make a PML4 Entry
     //
-    PageMapLevel4Entry->Uint64 = (UINT64)(UINTN)PageDirectoryPointerEntry;
+    PageMapLevel4Entry->Uint64 = (UINT64)(UINTN)PageDirectoryPointerEntry | PcdGet64 (PcdPteMemoryEncryptionAddressOrMask);
     PageMapLevel4Entry->Bits.ReadWrite = 1;
     PageMapLevel4Entry->Bits.Present = 1;
 
@@ -198,7 +199,7 @@ Create4GPageTables (
         //
         // Fill in the Page Directory entries
         //
-        PageDirectory1GEntry->Uint64 = (UINT64)PageAddress;
+        PageDirectory1GEntry->Uint64 = (UINT64)PageAddress | PcdGet64 (PcdPteMemoryEncryptionAddressOrMask);
         PageDirectory1GEntry->Bits.ReadWrite = 1;
         PageDirectory1GEntry->Bits.Present = 1;
         PageDirectory1GEntry->Bits.MustBe1 = 1;
@@ -215,7 +216,7 @@ Create4GPageTables (
         //
         // Fill in a Page Directory Pointer Entries
         //
-        PageDirectoryPointerEntry->Uint64 = (UINT64)(UINTN)PageDirectoryEntry;
+        PageDirectoryPointerEntry->Uint64 = (UINT64)(UINTN)PageDirectoryEntry | PcdGet64 (PcdPteMemoryEncryptionAddressOrMask);
         PageDirectoryPointerEntry->Bits.ReadWrite = 1;
         PageDirectoryPointerEntry->Bits.Present = 1;
 
@@ -223,7 +224,7 @@ Create4GPageTables (
           //
           // Fill in the Page Directory entries
           //
-          PageDirectoryEntry->Uint64 = (UINT64)PageAddress;
+          PageDirectoryEntry->Uint64 = (UINT64)PageAddress | PcdGet64 (PcdPteMemoryEncryptionAddressOrMask);
           PageDirectoryEntry->Bits.ReadWrite = 1;
           PageDirectoryEntry->Bits.Present = 1;
           PageDirectoryEntry->Bits.MustBe1 = 1;
diff --git a/MdeModulePkg/Universal/CapsulePei/X64/X64Entry.c b/MdeModulePkg/Universal/CapsulePei/X64/X64Entry.c
index 5ad95d2..f3f68ee 100644
--- a/MdeModulePkg/Universal/CapsulePei/X64/X64Entry.c
+++ b/MdeModulePkg/Universal/CapsulePei/X64/X64Entry.c
@@ -2,6 +2,8 @@
   The X64 entrypoint is used to process capsule in long mode.
 
 Copyright (c) 2011 - 2016, Intel Corporation. All rights reserved.<BR>
+Copyright (c) 2017, AMD Incorporated. All rights reserved.<BR>
+
 This program and the accompanying materials
 are licensed and made available under the terms and conditions of the BSD License
 which accompanies this distribution.  The full text of the license may be found at
@@ -128,7 +130,7 @@ AcquirePage (
   //
   // Link & Record the current uplink.
   //
-  *Uplink = Address | IA32_PG_P | IA32_PG_RW;
+  *Uplink = Address | PcdGet64 (PcdPteMemoryEncryptionAddressOrMask) | IA32_PG_P | IA32_PG_RW;
   PageFaultContext->PageFaultUplink[PageFaultContext->PageFaultIndex] = Uplink;
 
   PageFaultContext->PageFaultIndex = (PageFaultContext->PageFaultIndex + 1) % EXTRA_PAGE_TABLE_PAGES;
@@ -183,7 +185,7 @@ PageFaultHandler (
   PTIndex = BitFieldRead64 (PFAddress, 30, 38);
   // PDPTE
   if (PageFaultContext->Page1GSupport) {
-    PageTable[PTIndex] = (PFAddress & ~((1ull << 30) - 1)) | IA32_PG_P | IA32_PG_RW | IA32_PG_PS;
+    PageTable[PTIndex] = (PFAddress & ~((1ull << 30) - 1)) | PcdGet64 (PcdPteMemoryEncryptionAddressOrMask) | IA32_PG_P | IA32_PG_RW | IA32_PG_PS;
   } else {
     if ((PageTable[PTIndex] & IA32_PG_P) == 0) {
       AcquirePage (PageFaultContext, &PageTable[PTIndex]);
@@ -191,7 +193,7 @@ PageFaultHandler (
     PageTable = (UINT64*)(UINTN)(PageTable[PTIndex] & PhyMask);
     PTIndex = BitFieldRead64 (PFAddress, 21, 29);
     // PD
-    PageTable[PTIndex] = (PFAddress & ~((1ull << 21) - 1)) | IA32_PG_P | IA32_PG_RW | IA32_PG_PS;
+    PageTable[PTIndex] = (PFAddress & ~((1ull << 21) - 1)) | PcdGet64 (PcdPteMemoryEncryptionAddressOrMask) | IA32_PG_P | IA32_PG_RW | IA32_PG_PS;
   }
 
   return NULL;
@@ -298,4 +300,4 @@ _ModuleEntryPoint (
   //
   ASSERT (FALSE);
   return EFI_SUCCESS;
-}
\ No newline at end of file
+}
-- 
1.9.1