[PATCH v5 3/7] MdeModulePkg: Invoke all ExitBootServicesCallback instances at ExitBootServices

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

The protocol's intent is to allow drivers to install callbacks that can
modify the memory map at ExitBootServices time, so that any changes will
lead to the EFI_INVALID_PARAMETER error. This error is specified to require
the EBS caller to call GetMemoryMap again if it already had.

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>
---
 MdeModulePkg/Core/Dxe/DxeMain.inf       |  1 +
 MdeModulePkg/Core/Dxe/DxeMain/DxeMain.c | 62 ++++++++++++++++++++
 2 files changed, 63 insertions(+)

diff --git a/MdeModulePkg/Core/Dxe/DxeMain.inf b/MdeModulePkg/Core/Dxe/DxeMain.inf
index e4bca89577..bdd9cf8222 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
+  gEdkiiExitBootServicesCallbackProtocolGuid    ## CONSUMES
 
   # Arch Protocols
   gEfiBdsArchProtocolGuid                       ## CONSUMES
diff --git a/MdeModulePkg/Core/Dxe/DxeMain/DxeMain.c b/MdeModulePkg/Core/Dxe/DxeMain/DxeMain.c
index 5733f0c8ec..8cf7d6bcbf 100644
--- a/MdeModulePkg/Core/Dxe/DxeMain/DxeMain.c
+++ b/MdeModulePkg/Core/Dxe/DxeMain/DxeMain.c
@@ -6,6 +6,7 @@ SPDX-License-Identifier: BSD-2-Clause-Patent
 
 **/
 
+#include <Protocol/ExitBootServicesCallback.h>
 #include "DxeMain.h"
 
 //
@@ -744,6 +745,54 @@ CalculateEfiHdrCrc (
   Hdr->CRC32 = Crc;
 }
 
+/**
+  Invokes TerminateMemoryMapPrehook from every instance of the
+  EdkiiExitBootServicesProtocol.
+**/
+STATIC
+EFI_STATUS
+InvokeTerminateMemoryMapPrehooks (
+  VOID
+  )
+{
+  UINTN       NoHandles;
+  UINTN       Index;
+  EFI_HANDLE  *HandleBuffer;
+  EFI_STATUS  Status;
+  EDKII_EXIT_BOOT_SERVICES_CALLBACK_PROTOCOL *Callback;
+
+  Status = gBS->LocateHandleBuffer (
+                  ByProtocol,
+                  &gEdkiiExitBootServicesCallbackProtocolGuid,
+                  NULL,
+                  &NoHandles,
+                  &HandleBuffer
+                  );
+  if (EFI_ERROR (Status) && NoHandles == 0) {
+    return Status;
+  }
+
+  for (Index = 0; Index < NoHandles; Index++) {
+    Status = gBS->HandleProtocol (
+                    HandleBuffer[Index],
+                    &gEdkiiExitBootServicesCallbackProtocolGuid,
+                    (VOID **)&Callback
+                    );
+    if (EFI_ERROR (Status)) {
+      continue;
+    }
+
+    Status = Callback->TerminateMemoryMapPrehook(Callback);
+    if (EFI_ERROR (Status) || Status == EFI_WARN_STALE_DATA) {
+      goto done;
+    }
+  }
+
+done:
+  FreePool(HandleBuffer);
+  return Status;
+}
+
 /**
   Terminates all boot services.
 
@@ -768,6 +817,19 @@ CoreExitBootServices (
   //
   gTimer->SetTimerPeriod (gTimer, 0);
 
+  //
+  // Invoke all protocols installed for ExitBootServices prior to
+  // CoreTerminateMemoryMap.
+  //
+  Status = InvokeTerminateMemoryMapPrehooks();
+  if (EFI_ERROR (Status)) {
+    //
+    // Notify other drivers that ExitBootServices failed
+    //
+    CoreNotifySignalList (&gEventExitBootServicesFailedGuid);
+    return Status;
+  }
+
   //
   // Terminate memory services if the MapKey matches
   //
-- 
2.38.0.rc1.362.ged0d419d3c-goog