2 files changed, 40 insertions, 0 deletions
diff --git a/OvmfPkg/Library/SmmCpuFeaturesLib/SmmCpuFeaturesLib.c b/OvmfPkg/Library/SmmCpuFeaturesLib/SmmCpuFeaturesLib.c
index 13d929a983..59c319e01b 100644
--- a/OvmfPkg/Library/SmmCpuFeaturesLib/SmmCpuFeaturesLib.c
+++ b/OvmfPkg/Library/SmmCpuFeaturesLib/SmmCpuFeaturesLib.c
@@ -15,8 +15,10 @@
 #include <Library/BaseLib.h>
 #include <Library/BaseMemoryLib.h>
 #include <Library/DebugLib.h>
+#include <Library/MemEncryptSevLib.h>
 #include <Library/SmmCpuFeaturesLib.h>
 #include <Library/SmmServicesTableLib.h>
+#include <Library/UefiBootServicesTableLib.h>
 #include <PiSmm.h>
 #include <Register/QemuSmramSaveStateMap.h>
 
@@ -185,6 +187,42 @@ SmmCpuFeaturesSmmRelocationComplete (
   VOID
   )
 {
+  EFI_STATUS Status;
+  UINTN      MapPagesBase;
+  UINTN      MapPagesCount;
+
+  if (!MemEncryptSevIsEnabled ()) {
+    return;
+  }
+
+  //
+  // Now that SMBASE relocation is complete, re-encrypt the original SMRAM save
+  // state map's container pages, and release the pages to DXE. (The pages were
+  // allocated in PlatformPei.)
+  //
+  Status = MemEncryptSevLocateInitialSmramSaveStateMapPages (
+             &MapPagesBase,
+             &MapPagesCount
+             );
+  ASSERT_EFI_ERROR (Status);
+
+  Status = MemEncryptSevSetPageEncMask (
+             0,             // Cr3BaseAddress -- use current CR3
+             MapPagesBase,  // BaseAddress
+             MapPagesCount, // NumPages
+             TRUE           // Flush
+             );
+  if (EFI_ERROR (Status)) {
+    DEBUG ((DEBUG_ERROR, "%a: MemEncryptSevSetPageEncMask(): %r\n",
+      __FUNCTION__, Status));
+    ASSERT (FALSE);
+    CpuDeadLoop ();
+  }
+
+  ZeroMem ((VOID *)MapPagesBase, EFI_PAGES_TO_SIZE (MapPagesCount));
+
+  Status = gBS->FreePages (MapPagesBase, MapPagesCount);
+  ASSERT_EFI_ERROR (Status);
 }
 
 /**
diff --git a/OvmfPkg/Library/SmmCpuFeaturesLib/SmmCpuFeaturesLib.inf b/OvmfPkg/Library/SmmCpuFeaturesLib/SmmCpuFeaturesLib.inf
index 5184abbf21..7c2aaa890b 100644
--- a/OvmfPkg/Library/SmmCpuFeaturesLib/SmmCpuFeaturesLib.inf
+++ b/OvmfPkg/Library/SmmCpuFeaturesLib/SmmCpuFeaturesLib.inf
@@ -36,4 +36,6 @@
   BaseLib
   BaseMemoryLib
   DebugLib
+  MemEncryptSevLib
   SmmServicesTableLib
+  UefiBootServicesTableLib