From a00e2e5513926a4031d96dc707113d13e1cdf814 Mon Sep 17 00:00:00 2001
From: Dionna Glaze <dionnaglaze@google.com>
Date: Thu, 26 Jan 2023 21:17:37 +0000
Subject: OvmfPkg: Add memory acceptance event in AmdSevDxe

The added behavior is to accept all unaccepted memory at
ExitBootServices if the behavior is not disabled. This allows safe
upgrades for OS loaders to affirm their support for the unaccepted
memory type.

Cc: Gerd Hoffmann <kraxel@redhat.com>
Cc: James Bottomley <jejb@linux.ibm.com>
Cc: Jiewen Yao <jiewen.yao@intel.com>
Cc: Tom Lendacky <thomas.lendacky@amd.com>
Cc: Ard Biesheuvel <ardb@kernel.org>
Cc: "Min M. Xu" <min.m.xu@intel.com>
Cc: Andrew Fish <afish@apple.com>
Cc: "Michael D. Kinney" <michael.d.kinney@intel.com>

Signed-off-by: Dionna Glaze <dionnaglaze@google.com>
Reviewed-by: Ard Biesheuvel <ardb@kernel.org>
---
 OvmfPkg/AmdSevDxe/AmdSevDxe.c   | 97 +++++++++++++++++++++++++++++++++++++++++
 OvmfPkg/AmdSevDxe/AmdSevDxe.inf |  1 +
 2 files changed, 98 insertions(+)

diff --git a/OvmfPkg/AmdSevDxe/AmdSevDxe.c b/OvmfPkg/AmdSevDxe/AmdSevDxe.c
index f7600c3c81..37d1a3ff55 100644
--- a/OvmfPkg/AmdSevDxe/AmdSevDxe.c
+++ b/OvmfPkg/AmdSevDxe/AmdSevDxe.c
@@ -20,6 +20,7 @@
 #include <Library/UefiBootServicesTableLib.h>
 #include <Guid/ConfidentialComputingSevSnpBlob.h>
 #include <Library/PcdLib.h>
+#include <Pi/PrePiDxeCis.h>
 #include <Protocol/MemoryAccept.h>
 
 STATIC CONFIDENTIAL_COMPUTING_SNP_BLOB_LOCATION  mSnpBootDxeTable = {
@@ -34,6 +35,10 @@ STATIC CONFIDENTIAL_COMPUTING_SNP_BLOB_LOCATION  mSnpBootDxeTable = {
 
 STATIC EFI_HANDLE  mAmdSevDxeHandle = NULL;
 
+STATIC BOOLEAN  mAcceptAllMemoryAtEBS = TRUE;
+
+STATIC EFI_EVENT  mAcceptAllMemoryEvent = NULL;
+
 #define IS_ALIGNED(x, y)  ((((x) & ((y) - 1)) == 0))
 
 STATIC
@@ -62,6 +67,82 @@ AmdSevMemoryAccept (
   return EFI_SUCCESS;
 }
 
+STATIC
+EFI_STATUS
+AcceptAllMemory (
+  VOID
+  )
+{
+  EFI_GCD_MEMORY_SPACE_DESCRIPTOR  *AllDescMap;
+  UINTN                            NumEntries;
+  UINTN                            Index;
+  EFI_STATUS                       Status;
+
+  DEBUG ((DEBUG_INFO, "Accepting all memory\n"));
+
+  /*
+   * Get a copy of the memory space map to iterate over while
+   * changing the map.
+   */
+  Status = gDS->GetMemorySpaceMap (&NumEntries, &AllDescMap);
+  if (EFI_ERROR (Status)) {
+    return Status;
+  }
+
+  for (Index = 0; Index < NumEntries; Index++) {
+    CONST EFI_GCD_MEMORY_SPACE_DESCRIPTOR  *Desc;
+
+    Desc = &AllDescMap[Index];
+    if (Desc->GcdMemoryType != EFI_GCD_MEMORY_TYPE_UNACCEPTED) {
+      continue;
+    }
+
+    Status = AmdSevMemoryAccept (
+               NULL,
+               Desc->BaseAddress,
+               Desc->Length
+               );
+    if (EFI_ERROR (Status)) {
+      break;
+    }
+
+    Status = gDS->RemoveMemorySpace (Desc->BaseAddress, Desc->Length);
+    if (EFI_ERROR (Status)) {
+      break;
+    }
+
+    Status = gDS->AddMemorySpace (
+                    EfiGcdMemoryTypeSystemMemory,
+                    Desc->BaseAddress,
+                    Desc->Length,
+                    EFI_MEMORY_CPU_CRYPTO | EFI_MEMORY_XP | EFI_MEMORY_RO | EFI_MEMORY_RP
+                    );
+    if (EFI_ERROR (Status)) {
+      break;
+    }
+  }
+
+  gBS->FreePool (AllDescMap);
+  return Status;
+}
+
+VOID
+EFIAPI
+ResolveUnacceptedMemory (
+  IN EFI_EVENT  Event,
+  IN VOID       *Context
+  )
+{
+  EFI_STATUS  Status;
+
+  if (!mAcceptAllMemoryAtEBS) {
+    return;
+  }
+
+  Status = AcceptAllMemory ();
+  ASSERT_EFI_ERROR (Status);
+}
+
 STATIC EDKII_MEMORY_ACCEPT_PROTOCOL  mMemoryAcceptProtocol = {
   AmdSevMemoryAccept
 };
@@ -195,6 +276,22 @@ AmdSevDxeEntryPoint (
                     );
     ASSERT_EFI_ERROR (Status);
 
+    // SEV-SNP support does not automatically imply unaccepted memory support,
+    // so make ExitBootServices accept all unaccepted memory if support is
+    // not communicated.
+    Status = gBS->CreateEventEx (
+                    EVT_NOTIFY_SIGNAL,
+                    TPL_CALLBACK,
+                    ResolveUnacceptedMemory,
+                    NULL,
+                    &gEfiEventBeforeExitBootServicesGuid,
+                    &mAcceptAllMemoryEvent
+                    );
+
+    if (EFI_ERROR (Status)) {
+      DEBUG ((DEBUG_ERROR, "AllowUnacceptedMemory event creation for EventBeforeExitBootServices failed.\n"));
+    }
+
     //
     // If its SEV-SNP active guest then install the CONFIDENTIAL_COMPUTING_SEV_SNP_BLOB.
     // It contains the location for both the Secrets and CPUID page.
diff --git a/OvmfPkg/AmdSevDxe/AmdSevDxe.inf b/OvmfPkg/AmdSevDxe/AmdSevDxe.inf
index cd1b686c53..5b443d45bc 100644
--- a/OvmfPkg/AmdSevDxe/AmdSevDxe.inf
+++ b/OvmfPkg/AmdSevDxe/AmdSevDxe.inf
@@ -52,6 +52,7 @@
 
 [Guids]
   gConfidentialComputingSevSnpBlobGuid
+  gEfiEventBeforeExitBootServicesGuid
 
 [Pcd]
   gUefiOvmfPkgTokenSpaceGuid.PcdOvmfHostBridgePciDevId
-- 
cgit v1.2.3