[PATCH v2 06/15] OvmfPkg/MemEncryptSevLib: Add an interface to retrieve the encryption mask

["Lendacky, Thomas" <thomas.lendacky@amd.com>]

From: Tom Lendacky <thomas.lendacky@amd.com>

BZ: https://bugzilla.tianocore.org/show_bug.cgi?id=3108

To ensure that we always use a validated encryption mask for an SEV-ES
guest, create a new interface in the MemEncryptSevLib library to return
the encryption mask. This can be used in place of the multiple locations
where CPUID is used to retrieve the value (which would require validation
again) and allows the validated mask to be returned.

The PEI phase will use the value from the SEV-ES work area. Since the
SEV-ES work area isn't valid in the DXE phase, the DXE phase will use the
PcdPteMemoryEncryptionAddressOrMask PCD which is set during PEI.

Cc: Jordan Justen <jordan.l.justen@intel.com>
Cc: Laszlo Ersek <lersek@redhat.com>
Cc: Ard Biesheuvel <ard.biesheuvel@arm.com>
Cc: Rebecca Cran <rebecca@bsdio.com>
Cc: Peter Grehan <grehan@freebsd.org>
Cc: Anthony Perard <anthony.perard@citrix.com>
Cc: Julien Grall <julien@xen.org>
Cc: Brijesh Singh <brijesh.singh@amd.com>
Signed-off-by: Tom Lendacky <thomas.lendacky@amd.com>
---
 OvmfPkg/AmdSev/AmdSevX64.dsc                                                               |   4 +-
 OvmfPkg/Bhyve/BhyveX64.dsc                                                                 |   4 +-
 OvmfPkg/OvmfPkgIa32.dsc                                                                    |   4 +-
 OvmfPkg/OvmfPkgIa32X64.dsc                                                                 |   4 +-
 OvmfPkg/OvmfPkgX64.dsc                                                                     |   4 +-
 OvmfPkg/OvmfXen.dsc                                                                        |   3 +-
 OvmfPkg/Library/BaseMemEncryptSevLib/{BaseMemEncryptSevLib.inf => DxeMemEncryptSevLib.inf} |  15 +-
 OvmfPkg/Library/BaseMemEncryptSevLib/PeiMemEncryptSevLib.inf                               |  56 +++++++
 OvmfPkg/Include/Library/MemEncryptSevLib.h                                                 |  12 ++
 OvmfPkg/Library/BaseMemEncryptSevLib/DxeMemEncryptSevLibInternal.c                         | 145 ++++++++++++++++++
 OvmfPkg/Library/BaseMemEncryptSevLib/MemEncryptSevLibInternal.c                            |  94 +-----------
 OvmfPkg/Library/BaseMemEncryptSevLib/PeiMemEncryptSevLibInternal.c                         | 159 ++++++++++++++++++++
 12 files changed, 400 insertions(+), 104 deletions(-)

diff --git a/OvmfPkg/AmdSev/AmdSevX64.dsc b/OvmfPkg/AmdSev/AmdSevX64.dsc
index bb7697eb324b..d39436e7afc9 100644
--- a/OvmfPkg/AmdSev/AmdSevX64.dsc
+++ b/OvmfPkg/AmdSev/AmdSevX64.dsc
@@ -164,7 +164,7 @@ [LibraryClasses]
   QemuFwCfgSimpleParserLib|OvmfPkg/Library/QemuFwCfgSimpleParserLib/QemuFwCfgSimpleParserLib.inf
   VirtioLib|OvmfPkg/Library/VirtioLib/VirtioLib.inf
   LoadLinuxLib|OvmfPkg/Library/LoadLinuxLib/LoadLinuxLib.inf
-  MemEncryptSevLib|OvmfPkg/Library/BaseMemEncryptSevLib/BaseMemEncryptSevLib.inf
+  MemEncryptSevLib|OvmfPkg/Library/BaseMemEncryptSevLib/DxeMemEncryptSevLib.inf
   LockBoxLib|OvmfPkg/Library/LockBoxLib/LockBoxBaseLib.inf
   CustomizedDisplayLib|MdeModulePkg/Library/CustomizedDisplayLib/CustomizedDisplayLib.inf
   FrameBufferBltLib|MdeModulePkg/Library/FrameBufferBltLib/FrameBufferBltLib.inf
@@ -285,6 +285,8 @@ [LibraryClasses.common.PEIM]
   Tpm2DeviceLib|SecurityPkg/Library/Tpm2DeviceLibDTpm/Tpm2DeviceLibDTpm.inf
 !endif
 
+  MemEncryptSevLib|OvmfPkg/Library/BaseMemEncryptSevLib/PeiMemEncryptSevLib.inf
+
 [LibraryClasses.common.DXE_CORE]
   HobLib|MdePkg/Library/DxeCoreHobLib/DxeCoreHobLib.inf
   DxeCoreEntryPoint|MdePkg/Library/DxeCoreEntryPoint/DxeCoreEntryPoint.inf
diff --git a/OvmfPkg/Bhyve/BhyveX64.dsc b/OvmfPkg/Bhyve/BhyveX64.dsc
index b93fe30ae4e0..33edf3d2d6b5 100644
--- a/OvmfPkg/Bhyve/BhyveX64.dsc
+++ b/OvmfPkg/Bhyve/BhyveX64.dsc
@@ -163,7 +163,7 @@ [LibraryClasses]
   QemuFwCfgS3Lib|OvmfPkg/Library/QemuFwCfgS3Lib/BaseQemuFwCfgS3LibNull.inf
   BhyveFwCtlLib|OvmfPkg/Library/BhyveFwCtlLib/BhyveFwCtlLib.inf
   VirtioLib|OvmfPkg/Library/VirtioLib/VirtioLib.inf
-  MemEncryptSevLib|OvmfPkg/Library/BaseMemEncryptSevLib/BaseMemEncryptSevLib.inf
+  MemEncryptSevLib|OvmfPkg/Library/BaseMemEncryptSevLib/DxeMemEncryptSevLib.inf
   LockBoxLib|OvmfPkg/Library/LockBoxLib/LockBoxBaseLib.inf
 
   CustomizedDisplayLib|MdeModulePkg/Library/CustomizedDisplayLib/CustomizedDisplayLib.inf
@@ -292,6 +292,8 @@ [LibraryClasses.common.PEIM]
   Tpm2DeviceLib|SecurityPkg/Library/Tpm2DeviceLibDTpm/Tpm2DeviceLibDTpm.inf
 !endif
 
+  MemEncryptSevLib|OvmfPkg/Library/BaseMemEncryptSevLib/PeiMemEncryptSevLib.inf
+
 [LibraryClasses.common.DXE_CORE]
   HobLib|MdePkg/Library/DxeCoreHobLib/DxeCoreHobLib.inf
   DxeCoreEntryPoint|MdePkg/Library/DxeCoreEntryPoint/DxeCoreEntryPoint.inf
diff --git a/OvmfPkg/OvmfPkgIa32.dsc b/OvmfPkg/OvmfPkgIa32.dsc
index 26a013ec353e..75c5f46a7786 100644
--- a/OvmfPkg/OvmfPkgIa32.dsc
+++ b/OvmfPkg/OvmfPkgIa32.dsc
@@ -170,7 +170,7 @@ [LibraryClasses]
   QemuFwCfgSimpleParserLib|OvmfPkg/Library/QemuFwCfgSimpleParserLib/QemuFwCfgSimpleParserLib.inf
   VirtioLib|OvmfPkg/Library/VirtioLib/VirtioLib.inf
   LoadLinuxLib|OvmfPkg/Library/LoadLinuxLib/LoadLinuxLib.inf
-  MemEncryptSevLib|OvmfPkg/Library/BaseMemEncryptSevLib/BaseMemEncryptSevLib.inf
+  MemEncryptSevLib|OvmfPkg/Library/BaseMemEncryptSevLib/DxeMemEncryptSevLib.inf
 !if $(SMM_REQUIRE) == FALSE
   LockBoxLib|OvmfPkg/Library/LockBoxLib/LockBoxBaseLib.inf
 !endif
@@ -310,6 +310,8 @@ [LibraryClasses.common.PEIM]
   Tpm2DeviceLib|SecurityPkg/Library/Tpm2DeviceLibDTpm/Tpm2DeviceLibDTpm.inf
 !endif
 
+  MemEncryptSevLib|OvmfPkg/Library/BaseMemEncryptSevLib/PeiMemEncryptSevLib.inf
+
 [LibraryClasses.common.DXE_CORE]
   HobLib|MdePkg/Library/DxeCoreHobLib/DxeCoreHobLib.inf
   DxeCoreEntryPoint|MdePkg/Library/DxeCoreEntryPoint/DxeCoreEntryPoint.inf
diff --git a/OvmfPkg/OvmfPkgIa32X64.dsc b/OvmfPkg/OvmfPkgIa32X64.dsc
index 10579fe46c5b..8693248b4ea0 100644
--- a/OvmfPkg/OvmfPkgIa32X64.dsc
+++ b/OvmfPkg/OvmfPkgIa32X64.dsc
@@ -174,7 +174,7 @@ [LibraryClasses]
   QemuFwCfgSimpleParserLib|OvmfPkg/Library/QemuFwCfgSimpleParserLib/QemuFwCfgSimpleParserLib.inf
   VirtioLib|OvmfPkg/Library/VirtioLib/VirtioLib.inf
   LoadLinuxLib|OvmfPkg/Library/LoadLinuxLib/LoadLinuxLib.inf
-  MemEncryptSevLib|OvmfPkg/Library/BaseMemEncryptSevLib/BaseMemEncryptSevLib.inf
+  MemEncryptSevLib|OvmfPkg/Library/BaseMemEncryptSevLib/DxeMemEncryptSevLib.inf
 !if $(SMM_REQUIRE) == FALSE
   LockBoxLib|OvmfPkg/Library/LockBoxLib/LockBoxBaseLib.inf
 !endif
@@ -314,6 +314,8 @@ [LibraryClasses.common.PEIM]
   Tpm2DeviceLib|SecurityPkg/Library/Tpm2DeviceLibDTpm/Tpm2DeviceLibDTpm.inf
 !endif
 
+  MemEncryptSevLib|OvmfPkg/Library/BaseMemEncryptSevLib/PeiMemEncryptSevLib.inf
+
 [LibraryClasses.common.DXE_CORE]
   HobLib|MdePkg/Library/DxeCoreHobLib/DxeCoreHobLib.inf
   DxeCoreEntryPoint|MdePkg/Library/DxeCoreEntryPoint/DxeCoreEntryPoint.inf
diff --git a/OvmfPkg/OvmfPkgX64.dsc b/OvmfPkg/OvmfPkgX64.dsc
index c9235e48ad62..799468e08e92 100644
--- a/OvmfPkg/OvmfPkgX64.dsc
+++ b/OvmfPkg/OvmfPkgX64.dsc
@@ -174,7 +174,7 @@ [LibraryClasses]
   QemuFwCfgSimpleParserLib|OvmfPkg/Library/QemuFwCfgSimpleParserLib/QemuFwCfgSimpleParserLib.inf
   VirtioLib|OvmfPkg/Library/VirtioLib/VirtioLib.inf
   LoadLinuxLib|OvmfPkg/Library/LoadLinuxLib/LoadLinuxLib.inf
-  MemEncryptSevLib|OvmfPkg/Library/BaseMemEncryptSevLib/BaseMemEncryptSevLib.inf
+  MemEncryptSevLib|OvmfPkg/Library/BaseMemEncryptSevLib/DxeMemEncryptSevLib.inf
 !if $(SMM_REQUIRE) == FALSE
   LockBoxLib|OvmfPkg/Library/LockBoxLib/LockBoxBaseLib.inf
 !endif
@@ -314,6 +314,8 @@ [LibraryClasses.common.PEIM]
   Tpm2DeviceLib|SecurityPkg/Library/Tpm2DeviceLibDTpm/Tpm2DeviceLibDTpm.inf
 !endif
 
+  MemEncryptSevLib|OvmfPkg/Library/BaseMemEncryptSevLib/PeiMemEncryptSevLib.inf
+
 [LibraryClasses.common.DXE_CORE]
   HobLib|MdePkg/Library/DxeCoreHobLib/DxeCoreHobLib.inf
   DxeCoreEntryPoint|MdePkg/Library/DxeCoreEntryPoint/DxeCoreEntryPoint.inf
diff --git a/OvmfPkg/OvmfXen.dsc b/OvmfPkg/OvmfXen.dsc
index 12b7a87ee877..c13388ba5a05 100644
--- a/OvmfPkg/OvmfXen.dsc
+++ b/OvmfPkg/OvmfXen.dsc
@@ -161,7 +161,7 @@ [LibraryClasses]
   SerializeVariablesLib|OvmfPkg/Library/SerializeVariablesLib/SerializeVariablesLib.inf
   QemuFwCfgLib|OvmfPkg/Library/QemuFwCfgLib/QemuFwCfgDxeLib.inf
   QemuLoadImageLib|OvmfPkg/Library/GenericQemuLoadImageLib/GenericQemuLoadImageLib.inf
-  MemEncryptSevLib|OvmfPkg/Library/BaseMemEncryptSevLib/BaseMemEncryptSevLib.inf
+  MemEncryptSevLib|OvmfPkg/Library/BaseMemEncryptSevLib/DxeMemEncryptSevLib.inf
   LockBoxLib|OvmfPkg/Library/LockBoxLib/LockBoxBaseLib.inf
   CustomizedDisplayLib|MdeModulePkg/Library/CustomizedDisplayLib/CustomizedDisplayLib.inf
   FrameBufferBltLib|MdeModulePkg/Library/FrameBufferBltLib/FrameBufferBltLib.inf
@@ -273,6 +273,7 @@ [LibraryClasses.common.PEIM]
   QemuFwCfgS3Lib|OvmfPkg/Library/QemuFwCfgS3Lib/PeiQemuFwCfgS3LibFwCfg.inf
   PcdLib|MdePkg/Library/PeiPcdLib/PeiPcdLib.inf
   QemuFwCfgLib|OvmfPkg/Library/QemuFwCfgLib/QemuFwCfgPeiLib.inf
+  MemEncryptSevLib|OvmfPkg/Library/BaseMemEncryptSevLib/PeiMemEncryptSevLib.inf
 
 [LibraryClasses.common.DXE_CORE]
   HobLib|MdePkg/Library/DxeCoreHobLib/DxeCoreHobLib.inf
diff --git a/OvmfPkg/Library/BaseMemEncryptSevLib/BaseMemEncryptSevLib.inf b/OvmfPkg/Library/BaseMemEncryptSevLib/DxeMemEncryptSevLib.inf
similarity index 68%
rename from OvmfPkg/Library/BaseMemEncryptSevLib/BaseMemEncryptSevLib.inf
rename to OvmfPkg/Library/BaseMemEncryptSevLib/DxeMemEncryptSevLib.inf
index 7c44d0952815..837db0876184 100644
--- a/OvmfPkg/Library/BaseMemEncryptSevLib/BaseMemEncryptSevLib.inf
+++ b/OvmfPkg/Library/BaseMemEncryptSevLib/DxeMemEncryptSevLib.inf
@@ -1,7 +1,7 @@
 ## @file
 #  Library provides the helper functions for SEV guest
 #
-# Copyright (c) 2017 Advanced Micro Devices. All rights reserved.<BR>
+# Copyright (c) 2017 - 2020, Advanced Micro Devices. All rights reserved.<BR>
 #
 #  SPDX-License-Identifier: BSD-2-Clause-Patent
 #
@@ -10,11 +10,11 @@
 
 [Defines]
   INF_VERSION                    = 1.25
-  BASE_NAME                      = MemEncryptSevLib
+  BASE_NAME                      = DxeMemEncryptSevLib
   FILE_GUID                      = c1594631-3888-4be4-949f-9c630dbc842b
   MODULE_TYPE                    = BASE
   VERSION_STRING                 = 1.0
-  LIBRARY_CLASS                  = MemEncryptSevLib|PEIM DXE_DRIVER DXE_RUNTIME_DRIVER DXE_SMM_DRIVER UEFI_DRIVER
+  LIBRARY_CLASS                  = MemEncryptSevLib|DXE_DRIVER DXE_RUNTIME_DRIVER DXE_SMM_DRIVER UEFI_DRIVER
 
 #
 # The following information is for reference only and not required by the build
@@ -29,15 +29,17 @@ [Packages]
   OvmfPkg/OvmfPkg.dec
   UefiCpuPkg/UefiCpuPkg.dec
 
+[Sources]
+  DxeMemEncryptSevLibInternal.c
+  MemEncryptSevLibInternal.c
+
 [Sources.X64]
-  MemEncryptSevLibInternal.c
   X64/MemEncryptSevLib.c
   X64/VirtualMemory.c
   X64/VirtualMemory.h
 
 [Sources.IA32]
   Ia32/MemEncryptSevLib.c
-  MemEncryptSevLibInternal.c
 
 [LibraryClasses]
   BaseLib
@@ -49,3 +51,6 @@ [LibraryClasses]
 
 [FeaturePcd]
   gUefiOvmfPkgTokenSpaceGuid.PcdSmmSmramRequire
+
+[Pcd]
+  gEfiMdeModulePkgTokenSpaceGuid.PcdPteMemoryEncryptionAddressOrMask
diff --git a/OvmfPkg/Library/BaseMemEncryptSevLib/PeiMemEncryptSevLib.inf b/OvmfPkg/Library/BaseMemEncryptSevLib/PeiMemEncryptSevLib.inf
new file mode 100644
index 000000000000..7c29d14039d6
--- /dev/null
+++ b/OvmfPkg/Library/BaseMemEncryptSevLib/PeiMemEncryptSevLib.inf
@@ -0,0 +1,56 @@
+## @file
+#  Library provides the helper functions for SEV guest
+#
+# Copyright (c) 2020 Advanced Micro Devices. All rights reserved.<BR>
+#
+#  SPDX-License-Identifier: BSD-2-Clause-Patent
+#
+#
+##
+
+[Defines]
+  INF_VERSION                    = 1.25
+  BASE_NAME                      = PeiMemEncryptSevLib
+  FILE_GUID                      = 15d9a694-3d2a-4184-9672-ba55c3070e07
+  MODULE_TYPE                    = BASE
+  VERSION_STRING                 = 1.0
+  LIBRARY_CLASS                  = MemEncryptSevLib|PEIM
+
+#
+# The following information is for reference only and not required by the build
+# tools.
+#
+# VALID_ARCHITECTURES           = IA32 X64
+#
+
+[Packages]
+  MdeModulePkg/MdeModulePkg.dec
+  MdePkg/MdePkg.dec
+  OvmfPkg/OvmfPkg.dec
+  UefiCpuPkg/UefiCpuPkg.dec
+
+[Sources]
+  PeiMemEncryptSevLibInternal.c
+  MemEncryptSevLibInternal.c
+
+[Sources.X64]
+  X64/MemEncryptSevLib.c
+  X64/VirtualMemory.c
+  X64/VirtualMemory.h
+
+[Sources.IA32]
+  Ia32/MemEncryptSevLib.c
+
+[LibraryClasses]
+  BaseLib
+  CacheMaintenanceLib
+  CpuLib
+  DebugLib
+  MemoryAllocationLib
+  PcdLib
+
+[FeaturePcd]
+  gUefiOvmfPkgTokenSpaceGuid.PcdSmmSmramRequire
+
+[FixedPcd]
+  gUefiCpuPkgTokenSpaceGuid.PcdSevEsWorkAreaBase
diff --git a/OvmfPkg/Include/Library/MemEncryptSevLib.h b/OvmfPkg/Include/Library/MemEncryptSevLib.h
index a2c70aa550fe..872abe6725dc 100644
--- a/OvmfPkg/Include/Library/MemEncryptSevLib.h
+++ b/OvmfPkg/Include/Library/MemEncryptSevLib.h
@@ -135,4 +135,16 @@ MemEncryptSevLocateInitialSmramSaveStateMapPages (
   OUT UINTN *BaseAddress,
   OUT UINTN *NumberOfPages
   );
+
+/**
+  Returns the SEV encryption mask.
+
+  @return  The SEV pagetable encryption mask
+**/
+UINT64
+EFIAPI
+MemEncryptSevGetEncryptionMask (
+  VOID
+  );
+
 #endif // _MEM_ENCRYPT_SEV_LIB_H_
diff --git a/OvmfPkg/Library/BaseMemEncryptSevLib/DxeMemEncryptSevLibInternal.c b/OvmfPkg/Library/BaseMemEncryptSevLib/DxeMemEncryptSevLibInternal.c
new file mode 100644
index 000000000000..2816f859a0c4
--- /dev/null
+++ b/OvmfPkg/Library/BaseMemEncryptSevLib/DxeMemEncryptSevLibInternal.c
@@ -0,0 +1,145 @@
+/** @file
+
+  Secure Encrypted Virtualization (SEV) library helper function
+
+  Copyright (c) 2017 - 2020, AMD Incorporated. All rights reserved.<BR>
+
+  SPDX-License-Identifier: BSD-2-Clause-Patent
+
+**/
+
+#include <Library/BaseLib.h>
+#include <Library/DebugLib.h>
+#include <Library/MemEncryptSevLib.h>
+#include <Library/PcdLib.h>
+#include <Register/Amd/Cpuid.h>
+#include <Register/Amd/Msr.h>
+#include <Register/Cpuid.h>
+#include <Uefi/UefiBaseType.h>
+
+STATIC BOOLEAN mSevStatus = FALSE;
+STATIC BOOLEAN mSevEsStatus = FALSE;
+STATIC BOOLEAN mSevStatusChecked = FALSE;
+
+STATIC UINT64  mSevEncryptionMask = 0;
+STATIC BOOLEAN mSevEncryptionMaskSaved = FALSE;
+
+/**
+  Reads and sets the status of SEV features.
+
+  **/
+STATIC
+VOID
+EFIAPI
+InternalMemEncryptSevStatus (
+  VOID
+  )
+{
+  UINT32                            RegEax;
+  MSR_SEV_STATUS_REGISTER           Msr;
+  CPUID_MEMORY_ENCRYPTION_INFO_EAX  Eax;
+  BOOLEAN                           ReadSevMsr;
+  UINT64                            EncryptionMask;
+
+  ReadSevMsr = FALSE;
+
+  EncryptionMask = PcdGet64 (PcdPteMemoryEncryptionAddressOrMask);
+  if (EncryptionMask != 0) {
+    //
+    // The MSR has been read before, so it is safe to read it again and avoid
+    // having to validate the CPUID information.
+    //
+    ReadSevMsr = TRUE;
+  } else {
+    //
+    // Check if memory encryption leaf exist
+    //
+    AsmCpuid (CPUID_EXTENDED_FUNCTION, &RegEax, NULL, NULL, NULL);
+    if (RegEax >= CPUID_MEMORY_ENCRYPTION_INFO) {
+      //
+      // CPUID Fn8000_001F[EAX] Bit 1 (Sev supported)
+      //
+      AsmCpuid (CPUID_MEMORY_ENCRYPTION_INFO, &Eax.Uint32, NULL, NULL, NULL);
+
+      if (Eax.Bits.SevBit) {
+        ReadSevMsr = TRUE;
+      }
+    }
+  }
+
+  if (ReadSevMsr) {
+    //
+    // Check MSR_0xC0010131 Bit 0 (Sev Enabled)
+    //
+    Msr.Uint32 = AsmReadMsr32 (MSR_SEV_STATUS);
+    if (Msr.Bits.SevBit) {
+      mSevStatus = TRUE;
+    }
+
+    //
+    // Check MSR_0xC0010131 Bit 1 (Sev-Es Enabled)
+    //
+    if (Msr.Bits.SevEsBit) {
+      mSevEsStatus = TRUE;
+    }
+  }
+
+  mSevStatusChecked = TRUE;
+}
+
+/**
+  Returns a boolean to indicate whether SEV-ES is enabled.
+
+  @retval TRUE           SEV-ES is enabled
+  @retval FALSE          SEV-ES is not enabled
+**/
+BOOLEAN
+EFIAPI
+MemEncryptSevEsIsEnabled (
+  VOID
+  )
+{
+  if (!mSevStatusChecked) {
+    InternalMemEncryptSevStatus ();
+  }
+
+  return mSevEsStatus;
+}
+
+/**
+  Returns a boolean to indicate whether SEV is enabled.
+
+  @retval TRUE           SEV is enabled
+  @retval FALSE          SEV is not enabled
+**/
+BOOLEAN
+EFIAPI
+MemEncryptSevIsEnabled (
+  VOID
+  )
+{
+  if (!mSevStatusChecked) {
+    InternalMemEncryptSevStatus ();
+  }
+
+  return mSevStatus;
+}
+
+/**
+  Returns the SEV encryption mask.
+
+  @return  The SEV pagtable encryption mask
+**/
+UINT64
+EFIAPI
+MemEncryptSevGetEncryptionMask (
+  VOID
+  )
+{
+  if (!mSevEncryptionMaskSaved) {
+    mSevEncryptionMask = PcdGet64 (PcdPteMemoryEncryptionAddressOrMask);
+    mSevEncryptionMaskSaved = TRUE;
+  }
+
+  return mSevEncryptionMask;
+}
diff --git a/OvmfPkg/Library/BaseMemEncryptSevLib/MemEncryptSevLibInternal.c b/OvmfPkg/Library/BaseMemEncryptSevLib/MemEncryptSevLibInternal.c
index 02b8eb225d81..b4a9f464e268 100644
--- a/OvmfPkg/Library/BaseMemEncryptSevLib/MemEncryptSevLibInternal.c
+++ b/OvmfPkg/Library/BaseMemEncryptSevLib/MemEncryptSevLibInternal.c
@@ -2,7 +2,7 @@
 
   Secure Encrypted Virtualization (SEV) library helper function
 
-  Copyright (c) 2017, AMD Incorporated. All rights reserved.<BR>
+  Copyright (c) 2017 - 2020, AMD Incorporated. All rights reserved.<BR>
 
   SPDX-License-Identifier: BSD-2-Clause-Patent
 
@@ -12,102 +12,10 @@
 #include <Library/DebugLib.h>
 #include <Library/MemEncryptSevLib.h>
 #include <Library/PcdLib.h>
-#include <Register/Amd/Cpuid.h>
-#include <Register/Amd/Msr.h>
-#include <Register/Cpuid.h>
 #include <Register/QemuSmramSaveStateMap.h>
 #include <Register/SmramSaveStateMap.h>
 #include <Uefi/UefiBaseType.h>
 
-STATIC BOOLEAN mSevStatus = FALSE;
-STATIC BOOLEAN mSevEsStatus = FALSE;
-STATIC BOOLEAN mSevStatusChecked = FALSE;
-
-/**
-  Reads and sets the status of SEV features.
-
-  **/
-STATIC
-VOID
-EFIAPI
-InternalMemEncryptSevStatus (
-  VOID
-  )
-{
-  UINT32                            RegEax;
-  MSR_SEV_STATUS_REGISTER           Msr;
-  CPUID_MEMORY_ENCRYPTION_INFO_EAX  Eax;
-
-  //
-  // Check if memory encryption leaf exist
-  //
-  AsmCpuid (CPUID_EXTENDED_FUNCTION, &RegEax, NULL, NULL, NULL);
-  if (RegEax >= CPUID_MEMORY_ENCRYPTION_INFO) {
-    //
-    // CPUID Fn8000_001F[EAX] Bit 1 (Sev supported)
-    //
-    AsmCpuid (CPUID_MEMORY_ENCRYPTION_INFO, &Eax.Uint32, NULL, NULL, NULL);
-
-    if (Eax.Bits.SevBit) {
-      //
-      // Check MSR_0xC0010131 Bit 0 (Sev Enabled)
-      //
-      Msr.Uint32 = AsmReadMsr32 (MSR_SEV_STATUS);
-      if (Msr.Bits.SevBit) {
-        mSevStatus = TRUE;
-      }
-
-      //
-      // Check MSR_0xC0010131 Bit 1 (Sev-Es Enabled)
-      //
-      if (Msr.Bits.SevEsBit) {
-        mSevEsStatus = TRUE;
-      }
-    }
-  }
-
-  mSevStatusChecked = TRUE;
-}
-
-/**
-  Returns a boolean to indicate whether SEV-ES is enabled.
-
-  @retval TRUE           SEV-ES is enabled
-  @retval FALSE          SEV-ES is not enabled
-**/
-BOOLEAN
-EFIAPI
-MemEncryptSevEsIsEnabled (
-  VOID
-  )
-{
-  if (!mSevStatusChecked) {
-    InternalMemEncryptSevStatus ();
-  }
-
-  return mSevEsStatus;
-}
-
-/**
-  Returns a boolean to indicate whether SEV is enabled.
-
-  @retval TRUE           SEV is enabled
-  @retval FALSE          SEV is not enabled
-**/
-BOOLEAN
-EFIAPI
-MemEncryptSevIsEnabled (
-  VOID
-  )
-{
-  if (!mSevStatusChecked) {
-    InternalMemEncryptSevStatus ();
-  }
-
-  return mSevStatus;
-}
-
-
 /**
   Locate the page range that covers the initial (pre-SMBASE-relocation) SMRAM
   Save State Map.
diff --git a/OvmfPkg/Library/BaseMemEncryptSevLib/PeiMemEncryptSevLibInternal.c b/OvmfPkg/Library/BaseMemEncryptSevLib/PeiMemEncryptSevLibInternal.c
new file mode 100644
index 000000000000..e2fd109d120f
--- /dev/null
+++ b/OvmfPkg/Library/BaseMemEncryptSevLib/PeiMemEncryptSevLibInternal.c
@@ -0,0 +1,159 @@
+/** @file
+
+  Secure Encrypted Virtualization (SEV) library helper function
+
+  Copyright (c) 2020, AMD Incorporated. All rights reserved.<BR>
+
+  SPDX-License-Identifier: BSD-2-Clause-Patent
+
+**/
+
+#include <Library/BaseLib.h>
+#include <Library/DebugLib.h>
+#include <Library/MemEncryptSevLib.h>
+#include <Library/PcdLib.h>
+#include <Register/Amd/Cpuid.h>
+#include <Register/Amd/Msr.h>
+#include <Register/Cpuid.h>
+#include <Uefi/UefiBaseType.h>
+
+STATIC BOOLEAN mSevStatus = FALSE;
+STATIC BOOLEAN mSevEsStatus = FALSE;
+STATIC BOOLEAN mSevStatusChecked = FALSE;
+
+STATIC UINT64  mSevEncryptionMask = 0;
+STATIC BOOLEAN mSevEncryptionMaskSaved = FALSE;
+
+/**
+  Reads and sets the status of SEV features.
+
+  **/
+STATIC
+VOID
+EFIAPI
+InternalMemEncryptSevStatus (
+  VOID
+  )
+{
+  UINT32                            RegEax;
+  MSR_SEV_STATUS_REGISTER           Msr;
+  CPUID_MEMORY_ENCRYPTION_INFO_EAX  Eax;
+  BOOLEAN                           ReadSevMsr;
+  SEC_SEV_ES_WORK_AREA              *SevEsWorkArea;
+
+  ReadSevMsr = FALSE;
+
+  SevEsWorkArea = (SEC_SEV_ES_WORK_AREA *) FixedPcdGet32 (PcdSevEsWorkAreaBase);
+  if (SevEsWorkArea != NULL && SevEsWorkArea->EncryptionMask != 0) {
+    //
+    // The MSR has been read before, so it is safe to read it again and avoid
+    // having to validate the CPUID information.
+    //
+    ReadSevMsr = TRUE;
+  } else {
+    //
+    // Check if memory encryption leaf exist
+    //
+    AsmCpuid (CPUID_EXTENDED_FUNCTION, &RegEax, NULL, NULL, NULL);
+    if (RegEax >= CPUID_MEMORY_ENCRYPTION_INFO) {
+      //
+      // CPUID Fn8000_001F[EAX] Bit 1 (Sev supported)
+      //
+      AsmCpuid (CPUID_MEMORY_ENCRYPTION_INFO, &Eax.Uint32, NULL, NULL, NULL);
+
+      if (Eax.Bits.SevBit) {
+        ReadSevMsr = TRUE;
+      }
+    }
+  }
+
+  if (ReadSevMsr) {
+    //
+    // Check MSR_0xC0010131 Bit 0 (Sev Enabled)
+    //
+    Msr.Uint32 = AsmReadMsr32 (MSR_SEV_STATUS);
+    if (Msr.Bits.SevBit) {
+      mSevStatus = TRUE;
+    }
+
+    //
+    // Check MSR_0xC0010131 Bit 1 (Sev-Es Enabled)
+    //
+    if (Msr.Bits.SevEsBit) {
+      mSevEsStatus = TRUE;
+    }
+  }
+
+  mSevStatusChecked = TRUE;
+}
+
+/**
+  Returns a boolean to indicate whether SEV-ES is enabled.
+
+  @retval TRUE           SEV-ES is enabled
+  @retval FALSE          SEV-ES is not enabled
+**/
+BOOLEAN
+EFIAPI
+MemEncryptSevEsIsEnabled (
+  VOID
+  )
+{
+  if (!mSevStatusChecked) {
+    InternalMemEncryptSevStatus ();
+  }
+
+  return mSevEsStatus;
+}
+
+/**
+  Returns a boolean to indicate whether SEV is enabled.
+
+  @retval TRUE           SEV is enabled
+  @retval FALSE          SEV is not enabled
+**/
+BOOLEAN
+EFIAPI
+MemEncryptSevIsEnabled (
+  VOID
+  )
+{
+  if (!mSevStatusChecked) {
+    InternalMemEncryptSevStatus ();
+  }
+
+  return mSevStatus;
+}
+
+/**
+  Returns the SEV encryption mask.
+
+  @return  The SEV pagtable encryption mask
+**/
+UINT64
+EFIAPI
+MemEncryptSevGetEncryptionMask (
+  VOID
+  )
+{
+  if (!mSevEncryptionMaskSaved) {
+    SEC_SEV_ES_WORK_AREA  *SevEsWorkArea;
+
+    SevEsWorkArea = (SEC_SEV_ES_WORK_AREA *) FixedPcdGet32 (PcdSevEsWorkAreaBase);
+    if (SevEsWorkArea != NULL) {
+      mSevEncryptionMask = SevEsWorkArea->EncryptionMask;
+    } else {
+      CPUID_MEMORY_ENCRYPTION_INFO_EBX  Ebx;
+
+      //
+      // CPUID Fn8000_001F[EBX] Bit 0:5 (memory encryption bit position)
+      //
+      AsmCpuid (CPUID_MEMORY_ENCRYPTION_INFO, NULL, &Ebx.Uint32, NULL, NULL);
+      mSevEncryptionMask = LShiftU64 (1, Ebx.Bits.PtePosBits);
+    }
+
+    mSevEncryptionMaskSaved = TRUE;
+  }
+
+  return mSevEncryptionMask;
+}
-- 
2.30.0