[PATCH 2/4] 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 adds a dynamic Pcd that can be set to enable
unaccepted memory support before ExitBootServices is called.

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 +++++++++++++++++++++++++
 MdeModulePkg/MdeModulePkg.dec           |  6 ++
 MdeModulePkg/MdeModulePkg.uni           |  6 ++
 OvmfPkg/AmdSev/AmdSevX64.dsc            |  1 +
 OvmfPkg/Bhyve/BhyveX64.dsc              |  2 +
 OvmfPkg/CloudHv/CloudHvX64.dsc          |  2 +
 OvmfPkg/IntelTdx/IntelTdxX64.dsc        |  2 +
 OvmfPkg/OvmfPkgIa32X64.dsc              |  2 +
 OvmfPkg/OvmfPkgX64.dsc                  |  2 +
 OvmfPkg/OvmfXen.dsc                     |  2 +
 13 files changed, 137 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.
diff --git a/MdeModulePkg/MdeModulePkg.dec b/MdeModulePkg/MdeModulePkg.dec
index 58e6ab0048..dd07b3725a 100644
--- a/MdeModulePkg/MdeModulePkg.dec
+++ b/MdeModulePkg/MdeModulePkg.dec
@@ -2102,6 +2102,12 @@
   # @Prompt The shared bit mask when Intel Tdx is enabled.
   gEfiMdeModulePkgTokenSpaceGuid.PcdTdxSharedBitMask|0x0|UINT64|0x10000025
 
+  ## Indicates if the memory map may include unaccepted memory after ExitBootServices().<BR><BR>
+  #   TRUE  - The memory map may include unaccepted memory after ExitBootServices().<BR>
+  #   FALSE - The memory map may not include unaccepted memory after ExitBootServices().<BR>
+  # @Prompt Support unaccepted memory type.
+  gEfiMdeModulePkgTokenSpaceGuid.PcdEnableUnacceptedMemory|FALSE|BOOLEAN|0x10000026
+
 [PcdsPatchableInModule]
   ## Specify memory size with page number for PEI code when
   #  Loading Module at Fixed Address feature is enabled.
diff --git a/MdeModulePkg/MdeModulePkg.uni b/MdeModulePkg/MdeModulePkg.uni
index 33ce9f6198..fde57da123 100644
--- a/MdeModulePkg/MdeModulePkg.uni
+++ b/MdeModulePkg/MdeModulePkg.uni
@@ -1338,3 +1338,9 @@
 #string STR_gEfiMdeModulePkgTokenSpaceGuid_PcdPcieResizableBarSupport_HELP #language en-US "Indicates if the PCIe Resizable BAR Capability Supported.<BR><BR>\n"
                                                                                             "TRUE  - PCIe Resizable BAR Capability is supported.<BR>\n"
                                                                                             "FALSE - PCIe Resizable BAR Capability is not supported.<BR>"
+
+#string STR_gEfiMdeModulePkgTokenSpaceGuid_PcdEnableUnacceptedMemory_PROMPT #language en-US "Support unaccepted memory type"
+#string STR_gEfiMdeModulePkgTokenSpaceGuid_PcdEnableUnacceptedMemory_HELP #language en-US "Indicates if the memory map may include unaccepted memory "
+                                                                                          "after ExitBootServices().<BR><BR>\n"
+                                                                                          "TRUE  - The memory map may include unaccepted memory after ExitBootServices().<BR>\n"
+                                                                                          "FALSE - The memory map may not include unaccepted memory after ExitBootServices().<BR>\n"
diff --git a/OvmfPkg/AmdSev/AmdSevX64.dsc b/OvmfPkg/AmdSev/AmdSevX64.dsc
index 90e8a213ef..23086748c5 100644
--- a/OvmfPkg/AmdSev/AmdSevX64.dsc
+++ b/OvmfPkg/AmdSev/AmdSevX64.dsc
@@ -526,6 +526,7 @@
 
   # Set ConfidentialComputing defaults
   gEfiMdePkgTokenSpaceGuid.PcdConfidentialComputingGuestAttr|0
+  gEfiMdeModulePkgTokenSpaceGuid.PcdEnableUnacceptedMemory|FALSE
 
 !include OvmfPkg/OvmfTpmPcds.dsc.inc
 
diff --git a/OvmfPkg/Bhyve/BhyveX64.dsc b/OvmfPkg/Bhyve/BhyveX64.dsc
index 475b88b21a..004be8b019 100644
--- a/OvmfPkg/Bhyve/BhyveX64.dsc
+++ b/OvmfPkg/Bhyve/BhyveX64.dsc
@@ -559,6 +559,8 @@
   # Set Tdx shared bit mask
   gEfiMdeModulePkgTokenSpaceGuid.PcdTdxSharedBitMask|0x0
 
+  gEfiMdeModulePkgTokenSpaceGuid.PcdEnableUnacceptedMemory|FALSE
+
   gEfiSecurityPkgTokenSpaceGuid.PcdOptionRomImageVerificationPolicy|0x00
 
   # MdeModulePkg resolution sets up the system display resolution
diff --git a/OvmfPkg/CloudHv/CloudHvX64.dsc b/OvmfPkg/CloudHv/CloudHvX64.dsc
index 10b16104ac..41f43a2631 100644
--- a/OvmfPkg/CloudHv/CloudHvX64.dsc
+++ b/OvmfPkg/CloudHv/CloudHvX64.dsc
@@ -618,6 +618,8 @@
   # Set Tdx shared bit mask
   gEfiMdeModulePkgTokenSpaceGuid.PcdTdxSharedBitMask|0x0
 
+  gEfiMdeModulePkgTokenSpaceGuid.PcdEnableUnacceptedMemory|FALSE
+
   # Set SEV-ES defaults
   gEfiMdeModulePkgTokenSpaceGuid.PcdGhcbBase|0
   gEfiMdeModulePkgTokenSpaceGuid.PcdGhcbSize|0
diff --git a/OvmfPkg/IntelTdx/IntelTdxX64.dsc b/OvmfPkg/IntelTdx/IntelTdxX64.dsc
index c0c1a15b09..55b6a2a845 100644
--- a/OvmfPkg/IntelTdx/IntelTdxX64.dsc
+++ b/OvmfPkg/IntelTdx/IntelTdxX64.dsc
@@ -514,6 +514,8 @@
   # Set Tdx shared bit mask
   gEfiMdeModulePkgTokenSpaceGuid.PcdTdxSharedBitMask|0x0
 
+  gEfiMdeModulePkgTokenSpaceGuid.PcdEnableUnacceptedMemory|FALSE
+
   # Set SEV-ES defaults
   gEfiMdeModulePkgTokenSpaceGuid.PcdGhcbBase|0
   gEfiMdeModulePkgTokenSpaceGuid.PcdGhcbSize|0
diff --git a/OvmfPkg/OvmfPkgIa32X64.dsc b/OvmfPkg/OvmfPkgIa32X64.dsc
index af566b953f..aebe1c3192 100644
--- a/OvmfPkg/OvmfPkgIa32X64.dsc
+++ b/OvmfPkg/OvmfPkgIa32X64.dsc
@@ -655,6 +655,8 @@
   gEfiMdeModulePkgTokenSpaceGuid.PcdPteMemoryEncryptionAddressOrMask|0x0
   gEfiMdeModulePkgTokenSpaceGuid.PcdTdxSharedBitMask|0x0
 
+  gEfiMdeModulePkgTokenSpaceGuid.PcdEnableUnacceptedMemory|FALSE
+
   # Set SEV-ES defaults
   gEfiMdeModulePkgTokenSpaceGuid.PcdGhcbBase|0
   gEfiMdeModulePkgTokenSpaceGuid.PcdGhcbSize|0
diff --git a/OvmfPkg/OvmfPkgX64.dsc b/OvmfPkg/OvmfPkgX64.dsc
index f39d9cd117..6e4418388e 100644
--- a/OvmfPkg/OvmfPkgX64.dsc
+++ b/OvmfPkg/OvmfPkgX64.dsc
@@ -679,6 +679,8 @@
   # Set Tdx shared bit mask
   gEfiMdeModulePkgTokenSpaceGuid.PcdTdxSharedBitMask|0x0
 
+  gEfiMdeModulePkgTokenSpaceGuid.PcdEnableUnacceptedMemory|FALSE
+
   # Set SEV-ES defaults
   gEfiMdeModulePkgTokenSpaceGuid.PcdGhcbBase|0
   gEfiMdeModulePkgTokenSpaceGuid.PcdGhcbSize|0
diff --git a/OvmfPkg/OvmfXen.dsc b/OvmfPkg/OvmfXen.dsc
index 58a7c97cdd..0f57e22a2b 100644
--- a/OvmfPkg/OvmfXen.dsc
+++ b/OvmfPkg/OvmfXen.dsc
@@ -505,6 +505,8 @@
   # Set Tdx shared bit mask
   gEfiMdeModulePkgTokenSpaceGuid.PcdTdxSharedBitMask|0x0
 
+  gEfiMdeModulePkgTokenSpaceGuid.PcdEnableUnacceptedMemory|FALSE
+
   gEfiSecurityPkgTokenSpaceGuid.PcdOptionRomImageVerificationPolicy|0x00
 
 ################################################################################
-- 
2.37.3.998.g577e59143f-goog