[PATCH 1/3] OvmfPkg: Realize EfiMemoryAcceptProtocol in AmdSevDxe

["Dionna Glaze" <dionnaglaze@google.com>]

From: Sophia Wolf <phiawolf@google.com>

When a guest OS does not support unaccepted memory, the unaccepted
memory must be accepted before returning a memory map to the caller.

EfiMemoryAcceptProtocol is defined in MdePkg and is implemented /
Installed in AmdSevDxe for AMD SEV-SNP memory acceptance.

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>

Reviewed-by: Tom Lendacky <thomas.lendacky@amd.com>
Signed-off-by: Dionna Glaze <dionnaglaze@google.com>
---
 OvmfPkg/AmdSevDxe/AmdSevDxe.c                                      | 55 ++++++++++++++++++--
 OvmfPkg/AmdSevDxe/AmdSevDxe.inf                                    |  3 ++
 OvmfPkg/Library/BaseMemEncryptSevLib/X64/DxeSnpSystemRamValidate.c | 24 +++++++--
 3 files changed, 74 insertions(+), 8 deletions(-)

diff --git a/OvmfPkg/AmdSevDxe/AmdSevDxe.c b/OvmfPkg/AmdSevDxe/AmdSevDxe.c
index 662d3c4ccb..f7600c3c81 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 <Protocol/MemoryAccept.h>
 
 STATIC CONFIDENTIAL_COMPUTING_SNP_BLOB_LOCATION  mSnpBootDxeTable = {
   SIGNATURE_32 ('A',                                    'M', 'D', 'E'),
@@ -31,6 +32,40 @@ STATIC CONFIDENTIAL_COMPUTING_SNP_BLOB_LOCATION  mSnpBootDxeTable = {
   FixedPcdGet32 (PcdOvmfCpuidSize),
 };
 
+STATIC EFI_HANDLE  mAmdSevDxeHandle = NULL;
+
+#define IS_ALIGNED(x, y)  ((((x) & ((y) - 1)) == 0))
+
+STATIC
+EFI_STATUS
+EFIAPI
+AmdSevMemoryAccept (
+  IN EDKII_MEMORY_ACCEPT_PROTOCOL  *This,
+  IN EFI_PHYSICAL_ADDRESS          StartAddress,
+  IN UINTN                         Size
+  )
+{
+  //
+  // The StartAddress must be page-aligned, and the Size must be a positive
+  // multiple of SIZE_4KB. Use an assert instead of returning an erros since
+  // this is an EDK2-internal protocol.
+  //
+  ASSERT (IS_ALIGNED (StartAddress, SIZE_4KB));
+  ASSERT (IS_ALIGNED (Size, SIZE_4KB));
+  ASSERT (Size != 0);
+
+  MemEncryptSevSnpPreValidateSystemRam (
+    StartAddress,
+    EFI_SIZE_TO_PAGES (Size)
+    );
+
+  return EFI_SUCCESS;
+}
+
+STATIC EDKII_MEMORY_ACCEPT_PROTOCOL  mMemoryAcceptProtocol = {
+  AmdSevMemoryAccept
+};
+
 EFI_STATUS
 EFIAPI
 AmdSevDxeEntryPoint (
@@ -147,11 +182,23 @@ AmdSevDxeEntryPoint (
     }
   }
 
-  //
-  // 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.
-  //
   if (MemEncryptSevSnpIsEnabled ()) {
+    //
+    // Memory acceptance began being required in SEV-SNP, so install the
+    // memory accept protocol implementation for a SEV-SNP active guest.
+    //
+    Status = gBS->InstallProtocolInterface (
+                    &mAmdSevDxeHandle,
+                    &gEdkiiMemoryAcceptProtocolGuid,
+                    EFI_NATIVE_INTERFACE,
+                    &mMemoryAcceptProtocol
+                    );
+    ASSERT_EFI_ERROR (Status);
+
+    //
+    // 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.
+    //
     return gBS->InstallConfigurationTable (
                   &gConfidentialComputingSevSnpBlobGuid,
                   &mSnpBootDxeTable
diff --git a/OvmfPkg/AmdSevDxe/AmdSevDxe.inf b/OvmfPkg/AmdSevDxe/AmdSevDxe.inf
index 9acf860cf2..cd1b686c53 100644
--- a/OvmfPkg/AmdSevDxe/AmdSevDxe.inf
+++ b/OvmfPkg/AmdSevDxe/AmdSevDxe.inf
@@ -47,6 +47,9 @@
   gUefiOvmfPkgTokenSpaceGuid.PcdOvmfSnpSecretsBase
   gUefiOvmfPkgTokenSpaceGuid.PcdOvmfSnpSecretsSize
 
+[Protocols]
+  gEdkiiMemoryAcceptProtocolGuid
+
 [Guids]
   gConfidentialComputingSevSnpBlobGuid
 
diff --git a/OvmfPkg/Library/BaseMemEncryptSevLib/X64/DxeSnpSystemRamValidate.c b/OvmfPkg/Library/BaseMemEncryptSevLib/X64/DxeSnpSystemRamValidate.c
index d3a95e4913..cbcdd46f52 100644
--- a/OvmfPkg/Library/BaseMemEncryptSevLib/X64/DxeSnpSystemRamValidate.c
+++ b/OvmfPkg/Library/BaseMemEncryptSevLib/X64/DxeSnpSystemRamValidate.c
@@ -14,6 +14,7 @@
 #include <Library/MemEncryptSevLib.h>
 
 #include "SnpPageStateChange.h"
+#include "VirtualMemory.h"
 
 /**
   Pre-validate the system RAM when SEV-SNP is enabled in the guest VM.
@@ -29,12 +30,27 @@ MemEncryptSevSnpPreValidateSystemRam (
   IN UINTN             NumPages
   )
 {
+  EFI_STATUS  Status;
+
   if (!MemEncryptSevSnpIsEnabled ()) {
     return;
   }
 
-  //
-  // All the pre-validation must be completed in the PEI phase.
-  //
-  ASSERT (FALSE);
+  // DXE pre-validation may happen with the memory accept protocol.
+  // The protocol should only be called outside the prevalidated ranges
+  // that the PEI stage code explicitly skips. Specifically, only memory
+  // ranges that are classified as unaccepted.
+  if (BaseAddress >= SIZE_4GB) {
+    Status = InternalMemEncryptSevCreateIdentityMap1G (
+               0,
+               BaseAddress,
+               EFI_PAGES_TO_SIZE (NumPages)
+               );
+    if (EFI_ERROR (Status)) {
+      ASSERT (FALSE);
+      CpuDeadLoop ();
+    }
+  }
+
+  InternalSetPageState (BaseAddress, NumPages, SevSnpPagePrivate, TRUE);
 }
-- 
2.38.1.431.g37b22c650d-goog