[PATCH v3 4/4] MdeModulePkg/Universal/Acpi/BootScriptExecutorDxe: 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.

This module updates the under-4GB page tables configured by the S3-Resume
code in UefiCpuPkg/Universal/Acpi/S3Resume2Pei. The mask is saved at module
start (ScriptExecute.c), and applied when tables are expanded on-demand by
page-faults above 4GB's (SetIdtEntry.c).

CC: Jeff Fan <jeff.fan@intel.com>
Cc: Feng Tian <feng.tian@intel.com>
Cc: Star Zeng <star.zeng@intel.com>
Cc: Laszlo Ersek <lersek@redhat.com>
Cc: Brijesh Singh <brijesh.singh@amd.com>
Contributed-under: TianoCore Contribution Agreement 1.0
Signed-off-by: Leo Duran <leo.duran@amd.com>
---
 .../Acpi/BootScriptExecutorDxe/BootScriptExecutorDxe.inf |  2 ++
 .../Universal/Acpi/BootScriptExecutorDxe/ScriptExecute.c |  7 +++++++
 .../Universal/Acpi/BootScriptExecutorDxe/ScriptExecute.h |  2 ++
 .../Acpi/BootScriptExecutorDxe/X64/SetIdtEntry.c         | 16 +++++++++++-----
 4 files changed, 22 insertions(+), 5 deletions(-)

diff --git a/MdeModulePkg/Universal/Acpi/BootScriptExecutorDxe/BootScriptExecutorDxe.inf b/MdeModulePkg/Universal/Acpi/BootScriptExecutorDxe/BootScriptExecutorDxe.inf
index 7cd38cf..29af7f5 100644
--- a/MdeModulePkg/Universal/Acpi/BootScriptExecutorDxe/BootScriptExecutorDxe.inf
+++ b/MdeModulePkg/Universal/Acpi/BootScriptExecutorDxe/BootScriptExecutorDxe.inf
@@ -5,6 +5,7 @@
 # depends on any PEI or DXE service.
 #
 # 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 of the BSD License
@@ -85,6 +86,7 @@
   gEfiMdeModulePkgTokenSpaceGuid.PcdUse1GPageTable              ## SOMETIMES_CONSUMES
   gEfiMdeModulePkgTokenSpaceGuid.PcdMemoryProfilePropertyMask   ## CONSUMES
   gEfiMdeModulePkgTokenSpaceGuid.PcdAcpiS3Enable                ## CONSUMES
+  gEfiMdeModulePkgTokenSpaceGuid.PcdPteMemoryEncryptionAddressOrMask    ## CONSUMES
 
 [Depex]
   gEfiLockBoxProtocolGuid
diff --git a/MdeModulePkg/Universal/Acpi/BootScriptExecutorDxe/ScriptExecute.c b/MdeModulePkg/Universal/Acpi/BootScriptExecutorDxe/ScriptExecute.c
index f67fbca..5146b33 100644
--- a/MdeModulePkg/Universal/Acpi/BootScriptExecutorDxe/ScriptExecute.c
+++ b/MdeModulePkg/Universal/Acpi/BootScriptExecutorDxe/ScriptExecute.c
@@ -5,6 +5,7 @@
   in the entry point. The functionality is to interpret and restore the S3 boot script
 
 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 of the BSD License
@@ -23,6 +24,7 @@ EFI_GUID              mBootScriptExecutorImageGuid = {
 };
 
 BOOLEAN               mPage1GSupport = FALSE;
+UINT64                mPteMemoryEncryptionAddressOrMask = 0;
 
 /**
   Entry function of Boot script exector. This function will be executed in
@@ -440,6 +442,11 @@ BootScriptExecutorEntryPoint (
         }
       }
 
+      //
+      // Save OrMask to fix-up Address field in page-fault handler
+      //
+      mPteMemoryEncryptionAddressOrMask = PcdGet64 (PcdPteMemoryEncryptionAddressOrMask);
+
       BufferSize = sizeof (BOOT_SCRIPT_EXECUTOR_VARIABLE);
 
       BootScriptExecutorBuffer = 0xFFFFFFFF;
diff --git a/MdeModulePkg/Universal/Acpi/BootScriptExecutorDxe/ScriptExecute.h b/MdeModulePkg/Universal/Acpi/BootScriptExecutorDxe/ScriptExecute.h
index 772347a..cb88549 100644
--- a/MdeModulePkg/Universal/Acpi/BootScriptExecutorDxe/ScriptExecute.h
+++ b/MdeModulePkg/Universal/Acpi/BootScriptExecutorDxe/ScriptExecute.h
@@ -5,6 +5,7 @@
   in the entry point. The functionality is to interpret and restore the S3 boot script 
   
 Copyright (c) 2006 - 2014, 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
@@ -87,5 +88,6 @@ SetIdtEntry (
 extern UINT32 AsmFixAddress16;
 extern UINT32 AsmJmpAddr32;
 extern BOOLEAN mPage1GSupport;
+extern UINT64 mPteMemoryEncryptionAddressOrMask;
 
 #endif //_BOOT_SCRIPT_EXECUTOR_H_
diff --git a/MdeModulePkg/Universal/Acpi/BootScriptExecutorDxe/X64/SetIdtEntry.c b/MdeModulePkg/Universal/Acpi/BootScriptExecutorDxe/X64/SetIdtEntry.c
index 6674560..37724f0 100644
--- a/MdeModulePkg/Universal/Acpi/BootScriptExecutorDxe/X64/SetIdtEntry.c
+++ b/MdeModulePkg/Universal/Acpi/BootScriptExecutorDxe/X64/SetIdtEntry.c
@@ -4,6 +4,8 @@
   Set a IDT entry for interrupt vector 3 for debug purpose for x64 platform
 
 Copyright (c) 2006 - 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 of the BSD License
@@ -25,6 +27,10 @@ WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
 #define IA32_PG_RW                  BIT1
 #define IA32_PG_PS                  BIT7
 
+#define PAGING_4K_ADDRESS_MASK_64 0x000FFFFFFFFFF000ull
+#define PAGING_2M_ADDRESS_MASK_64 0x000FFFFFFFE00000ull
+#define PAGING_1G_ADDRESS_MASK_64 0x000FFFFFC0000000ull
+
 UINT64                              mPhyMask;
 VOID                                *mOriginalHandler;
 UINTN                               mPageFaultBuffer;
@@ -207,7 +213,7 @@ AcquirePage (
   //
   // Link & Record the current uplink.
   //
-  *Uplink = Address | IA32_PG_P | IA32_PG_RW;
+  *Uplink = Address | (mPteMemoryEncryptionAddressOrMask & PAGING_4K_ADDRESS_MASK_64) | IA32_PG_P | IA32_PG_RW;
   mPageFaultUplink[mPageFaultIndex] = Uplink;
 
   mPageFaultIndex = (mPageFaultIndex + 1) % EXTRA_PAGE_TABLE_PAGES;
@@ -245,19 +251,19 @@ PageFaultHandler (
   if ((PageTable[PTIndex] & IA32_PG_P) == 0) {
     AcquirePage (&PageTable[PTIndex]);
   }
-  PageTable = (UINT64*)(UINTN)(PageTable[PTIndex] & mPhyMask);
+  PageTable = (UINT64*)(UINTN)((PageTable[PTIndex] & ~(mPteMemoryEncryptionAddressOrMask & PAGING_4K_ADDRESS_MASK_64)) & mPhyMask);
   PTIndex = BitFieldRead64 (PFAddress, 30, 38);
   // PDPTE
   if (mPage1GSupport) {
-    PageTable[PTIndex] = (PFAddress & ~((1ull << 30) - 1)) | IA32_PG_P | IA32_PG_RW | IA32_PG_PS;
+    PageTable[PTIndex] = ((PFAddress | (mPteMemoryEncryptionAddressOrMask & PAGING_1G_ADDRESS_MASK_64)) & ~((1ull << 30) - 1)) | IA32_PG_P | IA32_PG_RW | IA32_PG_PS;
   } else {
     if ((PageTable[PTIndex] & IA32_PG_P) == 0) {
       AcquirePage (&PageTable[PTIndex]);
     }
-    PageTable = (UINT64*)(UINTN)(PageTable[PTIndex] & mPhyMask);
+    PageTable = (UINT64*)(UINTN)((PageTable[PTIndex] & ~(mPteMemoryEncryptionAddressOrMask & PAGING_4K_ADDRESS_MASK_64)) & mPhyMask);
     PTIndex = BitFieldRead64 (PFAddress, 21, 29);
     // PD
-    PageTable[PTIndex] = (PFAddress & ~((1ull << 21) - 1)) | IA32_PG_P | IA32_PG_RW | IA32_PG_PS;
+    PageTable[PTIndex] = ((PFAddress | (mPteMemoryEncryptionAddressOrMask & PAGING_2M_ADDRESS_MASK_64)) & ~((1ull << 21) - 1)) | IA32_PG_P | IA32_PG_RW | IA32_PG_PS;
   }
 
   return TRUE;
-- 
2.7.4