[PATCH v4 4/6] MdeModulePkg: DxeMain accepts all memory at EBS if needed

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

With the addition of the EfiUnacceptedMemory memory type, it is possible
the EFI-enlightened guests do not themselves support the new memory
type. This commit uses the new PcdEnableUnacceptedMemory to enable
unaccepted memory support before ExitBootServices is called by not
accepting all unaccepted memory at EBS.

The expected usage is to set the new Pcd with a protocol that is usable
by bootloaders and directly-booted OSes when they can determine that the
OS does indeed support unaccepted memory.

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>

Signed-off-by: Dionna Glaze <dionnaglaze@google.com>
---
 MdeModulePkg/Core/Dxe/DxeMain.h         | 10 +++
 MdeModulePkg/Core/Dxe/DxeMain.inf       |  2 +
 MdeModulePkg/Core/Dxe/DxeMain/DxeMain.c | 14 +++-
 MdeModulePkg/Core/Dxe/Mem/Page.c        | 87 ++++++++++++++++++++
 4 files changed, 112 insertions(+), 1 deletion(-)

diff --git a/MdeModulePkg/Core/Dxe/DxeMain.h b/MdeModulePkg/Core/Dxe/DxeMain.h
index 815a6b4bd8..ac943c87a3 100644
--- a/MdeModulePkg/Core/Dxe/DxeMain.h
+++ b/MdeModulePkg/Core/Dxe/DxeMain.h
@@ -2698,6 +2698,16 @@ CoreInitializeMemoryProtection (
   VOID
   );
 
+/**
+   Accept and convert unaccepted memory to conventional memory if unaccepted
+   memory is not enabled and there is an implementation of MemoryAcceptProtocol
+   installed.
+ **/
+EFI_STATUS
+CoreResolveUnacceptedMemory (
+  VOID
+  );
+
 /**
   Install MemoryAttributesTable on memory allocation.
 
diff --git a/MdeModulePkg/Core/Dxe/DxeMain.inf b/MdeModulePkg/Core/Dxe/DxeMain.inf
index e4bca89577..deb8bb2ba8 100644
--- a/MdeModulePkg/Core/Dxe/DxeMain.inf
+++ b/MdeModulePkg/Core/Dxe/DxeMain.inf
@@ -153,6 +153,7 @@
   gEfiHiiPackageListProtocolGuid                ## SOMETIMES_PRODUCES
   gEfiSmmBase2ProtocolGuid                      ## SOMETIMES_CONSUMES
   gEdkiiPeCoffImageEmulatorProtocolGuid         ## SOMETIMES_CONSUMES
+  gEfiMemoryAcceptProtocolGuid                  ## SOMETIMES_CONSUMES
 
   # Arch Protocols
   gEfiBdsArchProtocolGuid                       ## CONSUMES
@@ -186,6 +187,7 @@
   gEfiMdeModulePkgTokenSpaceGuid.PcdHeapGuardPropertyMask                   ## CONSUMES
   gEfiMdeModulePkgTokenSpaceGuid.PcdCpuStackGuard                           ## CONSUMES
   gEfiMdeModulePkgTokenSpaceGuid.PcdFwVolDxeMaxEncapsulationDepth           ## CONSUMES
+  gEfiMdeModulePkgTokenSpaceGuid.PcdEnableUnacceptedMemory                  ## CONSUMES
 
 # [Hob]
 # RESOURCE_DESCRIPTOR   ## CONSUMES
diff --git a/MdeModulePkg/Core/Dxe/DxeMain/DxeMain.c b/MdeModulePkg/Core/Dxe/DxeMain/DxeMain.c
index 5733f0c8ec..8d1de32fe7 100644
--- a/MdeModulePkg/Core/Dxe/DxeMain/DxeMain.c
+++ b/MdeModulePkg/Core/Dxe/DxeMain/DxeMain.c
@@ -768,13 +768,25 @@ CoreExitBootServices (
   //
   gTimer->SetTimerPeriod (gTimer, 0);
 
+  //
+  // Accept all memory if unaccepted memory isn't enabled.
+  //
+  Status = CoreResolveUnacceptedMemory();
+  if (EFI_ERROR (Status)) {
+    //
+    // Notify other drivers that ExitBootServices failed
+    //
+    CoreNotifySignalList (&gEventExitBootServicesFailedGuid);
+    return Status;
+  }
+
   //
   // Terminate memory services if the MapKey matches
   //
   Status = CoreTerminateMemoryMap (MapKey);
   if (EFI_ERROR (Status)) {
     //
-    // Notify other drivers that ExitBootServices fail
+    // Notify other drivers that ExitBootServices failed
     //
     CoreNotifySignalList (&gEventExitBootServicesFailedGuid);
     return Status;
diff --git a/MdeModulePkg/Core/Dxe/Mem/Page.c b/MdeModulePkg/Core/Dxe/Mem/Page.c
index ffe79dcca9..cbebe62a28 100644
--- a/MdeModulePkg/Core/Dxe/Mem/Page.c
+++ b/MdeModulePkg/Core/Dxe/Mem/Page.c
@@ -9,6 +9,8 @@ SPDX-License-Identifier: BSD-2-Clause-Patent
 #include "DxeMain.h"
 #include "Imem.h"
 #include "HeapGuard.h"
+#include <Library/PcdLib.h>
+#include <Protocol/MemoryAccept.h>
 
 //
 // Entry for tracking the memory regions for each memory type to coalesce similar memory types
@@ -2118,6 +2120,91 @@ CoreFreePoolPages (
   CoreConvertPages (Memory, NumberOfPages, EfiConventionalMemory);
 }
 
+EFI_EVENT gExitBootServiceEvent = NULL;
+
+STATIC
+EFI_STATUS
+AcceptAllUnacceptedMemory (
+  IN EFI_MEMORY_ACCEPT_PROTOCOL *AcceptMemory
+  )
+{
+  EFI_GCD_MEMORY_SPACE_DESCRIPTOR  *AllDescMap;
+  UINTN                            NumEntries;
+  UINTN                            Index;
+  EFI_STATUS                       Status;
+
+  /*
+   * Get a copy of the memory space map to iterate over while
+   * changing the map.
+   */
+  Status = CoreGetMemorySpaceMap (&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 != EfiGcdMemoryTypeUnaccepted) {
+      continue;
+    }
+
+    Status = AcceptMemory->AcceptMemory (
+      AcceptMemory,
+      Desc->BaseAddress,
+      Desc->Length
+      );
+    if (EFI_ERROR(Status)) {
+      goto done;
+    }
+
+    Status = CoreRemoveMemorySpace(Desc->BaseAddress, Desc->Length);
+    if (EFI_ERROR(Status)) {
+      goto done;
+    }
+
+    Status = CoreAddMemorySpace (
+      EfiGcdMemoryTypeSystemMemory,
+      Desc->BaseAddress,
+      Desc->Length,
+      EFI_MEMORY_CPU_CRYPTO | EFI_MEMORY_XP | EFI_MEMORY_RO | EFI_MEMORY_RP
+      );
+    if (EFI_ERROR(Status)) {
+      goto done;
+    }
+  }
+
+done:
+  FreePool (AllDescMap);
+  return Status;
+}
+
+EFI_STATUS
+CoreResolveUnacceptedMemory (
+  VOID
+  )
+{
+  EFI_MEMORY_ACCEPT_PROTOCOL *AcceptMemory;
+  EFI_STATUS                 Status;
+
+  // No need to accept anything. Unaccepted memory is enabled.
+  if (PcdGetBool(PcdEnableUnacceptedMemory)) {
+    return EFI_SUCCESS;
+  }
+
+  Status = gBS->LocateProtocol (&gEfiMemoryAcceptProtocolGuid, NULL,
+    (VOID **)&AcceptMemory);
+  if (Status == EFI_NOT_FOUND) {
+    return EFI_SUCCESS;
+  }
+  if (Status != EFI_SUCCESS) {
+    DEBUG ((DEBUG_ERROR, "Error locating MemoryAcceptProtocol: %d\n", Status));
+    return Status;
+  }
+
+  return AcceptAllUnacceptedMemory(AcceptMemory);
+}
+
 /**
   Make sure the memory map is following all the construction rules,
   it is the last time to check memory map error before exit boot services.
-- 
2.37.3.998.g577e59143f-goog