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

@@ -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.

OvmfPkg/AmdSevDxe/AmdSevDxe.inf

@@ -52,6 +52,7 @@
 
 [Guids]
   gConfidentialComputingSevSnpBlobGuid
+  gEfiEventBeforeExitBootServicesGuid
 
 [Pcd]
   gUefiOvmfPkgTokenSpaceGuid.PcdOvmfHostBridgePciDevId