From 3ac296def15a0bd34b095d7009843e76fcdbf418 Mon Sep 17 00:00:00 2001
From: Hongbin1 Zhang <hongbin1.zhang@intel.com>
Date: Thu, 9 May 2024 16:52:52 +0800
Subject: [PATCH] StandaloneMmPkg/MmIpl: Install end of PEI notify PPI

Install end of PEI notify PPI for issue gEfiMmEndOfPeiProtocol
Handler in MM

Signed-off-by: Hongbin1 Zhang <hongbin1.zhang@intel.com>
Cc: Jiewen Yao <jiewen.yao@intel.com>
Cc: Ray Ni <ray.ni@intel.com>
Cc: Star Zeng <star.zeng@intel.com>
Cc: Jiaxin Wu <jiaxin.wu@intel.com>
Cc: Wei6 Xu <wei6.xu@intel.com>
Cc: Sami Mujawar <sami.mujawar@arm.com>
Cc: Ard Biesheuvel <ardb+tianocore@kernel.org>
Cc: Supreeth Venkatesh <supreeth.venkatesh@arm.com>
---
 .../StandaloneMmIplPei/StandaloneMmIplPei.c   | 53 +++++++++++++++++++
 .../StandaloneMmIplPei/StandaloneMmIplPei.h   | 19 +++++++
 .../StandaloneMmIplPei/StandaloneMmIplPei.inf |  4 ++
 3 files changed, 76 insertions(+)

diff --git a/StandaloneMmPkg/Drivers/StandaloneMmIplPei/StandaloneMmIplPei.c b/StandaloneMmPkg/Drivers/StandaloneMmIplPei/StandaloneMmIplPei.c
index d92c2659ad..90f3e46781 100644
--- a/StandaloneMmPkg/Drivers/StandaloneMmIplPei/StandaloneMmIplPei.c
+++ b/StandaloneMmPkg/Drivers/StandaloneMmIplPei/StandaloneMmIplPei.c
@@ -16,6 +16,12 @@ EFI_PEI_PPI_DESCRIPTOR  mPpiList = {
   &mMmCommunicationPpi
 };
 
+EFI_PEI_NOTIFY_DESCRIPTOR  mNotifyList = {
+  EFI_PEI_PPI_DESCRIPTOR_NOTIFY_CALLBACK | EFI_PEI_PPI_DESCRIPTOR_TERMINATE_LIST,
+  &gEfiEndOfPeiSignalPpiGuid,
+  EndOfPeiCallback
+};
+
 /**
   Communicates with a registered handler.
 
@@ -451,6 +457,47 @@ ExecuteMmCoreFromMmram (
   return Status;
 }
 
+/**
+  This is the callback function on end of PEI.
+
+  This callback is used for call MmEndOfPeiHandler in standalone MM core.
+
+  @param   PeiServices       General purpose services available to every PEIM.
+  @param   NotifyDescriptor  The notification structure this PEIM registered on install.
+  @param   Ppi               Pointer to the PPI data associated with this function.
+
+  @retval  EFI_SUCCESS       Exit boot services successfully.
+  @retval  Other             Exit boot services failed.
+**/
+EFI_STATUS
+EFIAPI
+EndOfPeiCallback (
+  IN  EFI_PEI_SERVICES           **PeiServices,
+  IN  EFI_PEI_NOTIFY_DESCRIPTOR  *NotifyDescriptor,
+  IN  VOID                       *Ppi
+  )
+{
+  EFI_MM_COMMUNICATE_HEADER  CommunicateHeader;
+  UINTN                      Size;
+  EFI_STATUS                 Status;
+
+  //
+  // Use Guid to initialize EFI_MM_COMMUNICATE_HEADER structure
+  //
+  CopyGuid (&CommunicateHeader.HeaderGuid, &gEfiMmEndOfPeiProtocol);
+  CommunicateHeader.MessageLength = 1;
+  CommunicateHeader.Data[0]       = 0;
+
+  //
+  // Generate the Software SMI and return the result
+  //
+  Size   = sizeof (CommunicateHeader);
+  Status = Communicate (NULL, &CommunicateHeader, &Size);
+  ASSERT_EFI_ERROR (Status);
+
+  return Status;
+}
+
 /**
   Build communication buffer HOB.
 
@@ -550,5 +597,11 @@ StandaloneMmIplPeiEntry (
   Status = PeiServicesInstallPpi (&mPpiList);
   ASSERT_EFI_ERROR (Status);
 
+  //
+  // Create end of pei callback to call MmEndOfPeiHandler
+  //
+  Status = PeiServicesNotifyPpi (&mNotifyList);
+  ASSERT_EFI_ERROR (Status);
+
   return EFI_SUCCESS;
 }
diff --git a/StandaloneMmPkg/Drivers/StandaloneMmIplPei/StandaloneMmIplPei.h b/StandaloneMmPkg/Drivers/StandaloneMmIplPei/StandaloneMmIplPei.h
index c385d98ce4..e9be3f5568 100644
--- a/StandaloneMmPkg/Drivers/StandaloneMmIplPei/StandaloneMmIplPei.h
+++ b/StandaloneMmPkg/Drivers/StandaloneMmIplPei/StandaloneMmIplPei.h
@@ -47,4 +47,23 @@ Communicate (
   IN OUT UINTN                           *CommSize
   );
 
+/**
+  This is the callback function on end of PEI.
+
+  This callback is used for call MmEndOfPeiHandler in standalone MM core.
+
+  @param   PeiServices       General purpose services available to every PEIM.
+  @param   NotifyDescriptor  The notification structure this PEIM registered on install.
+  @param   Ppi               Pointer to the PPI data associated with this function.
+  @retval  EFI_SUCCESS       Exit boot services successfully.
+  @retval  Other             Exit boot services failed.
+**/
+EFI_STATUS
+EFIAPI
+EndOfPeiCallback (
+  IN  EFI_PEI_SERVICES           **PeiServices,
+  IN  EFI_PEI_NOTIFY_DESCRIPTOR  *NotifyDescriptor,
+  IN  VOID                       *Ppi
+  );
+
 #endif
diff --git a/StandaloneMmPkg/Drivers/StandaloneMmIplPei/StandaloneMmIplPei.inf b/StandaloneMmPkg/Drivers/StandaloneMmIplPei/StandaloneMmIplPei.inf
index f462a10e35..9c748539c6 100644
--- a/StandaloneMmPkg/Drivers/StandaloneMmIplPei/StandaloneMmIplPei.inf
+++ b/StandaloneMmPkg/Drivers/StandaloneMmIplPei/StandaloneMmIplPei.inf
@@ -49,6 +49,10 @@
 [Ppis]
   gEfiPeiMmControlPpiGuid
   gEfiPeiMmCommunicationPpiGuid
+  gEfiEndOfPeiSignalPpiGuid
+
+[Protocols]
+  gEfiMmEndOfPeiProtocol
 
 [Pcd]
   gEfiMdeModulePkgTokenSpaceGuid.PcdMmCommBufferPages