2 files changed, 50 insertions, 2 deletions
diff --git a/ArmPkg/Drivers/MmCommunicationDxe/MmCommunication.c b/ArmPkg/Drivers/MmCommunicationDxe/MmCommunication.c
index d06dcc4d20..5f52afa0fe 100644
--- a/ArmPkg/Drivers/MmCommunicationDxe/MmCommunication.c
+++ b/ArmPkg/Drivers/MmCommunicationDxe/MmCommunication.c
@@ -259,6 +259,43 @@ GetMmCompatibility ()
   return Status;
 }
 
+STATIC EFI_GUID* CONST mGuidedEventGuid[] = {
+  &gEfiEndOfDxeEventGroupGuid,
+  &gEfiEventExitBootServicesGuid,
+  &gEfiEventReadyToBootGuid,
+};
+
+STATIC EFI_EVENT mGuidedEvent[ARRAY_SIZE (mGuidedEventGuid)];
+
+/**
+  Event notification that is fired when GUIDed Event Group is signaled.
+
+  @param  Event                 The Event that is being processed, not used.
+  @param  Context               Event Context, not used.
+
+**/
+STATIC
+VOID
+EFIAPI
+MmGuidedEventNotify (
+  IN EFI_EVENT  Event,
+  IN VOID       *Context
+  )
+{
+  EFI_MM_COMMUNICATE_HEADER   Header;
+  UINTN                       Size;
+
+  //
+  // Use Guid to initialize EFI_SMM_COMMUNICATE_HEADER structure
+  //
+  CopyGuid (&Header.HeaderGuid, Context);
+  Header.MessageLength = 1;
+  Header.Data[0] = 0;
+
+  Size = sizeof (Header);
+  MmCommunicationCommunicate (&mMmCommunication, &Header, &Size);
+}
+
 /**
   The Entry Point for MM Communication
 
@@ -281,6 +318,7 @@ MmCommunicationInitialize (
   )
 {
   EFI_STATUS                 Status;
+  UINTN                      Index;
 
   // Check if we can make the MM call
   Status = GetMmCompatibility ();
@@ -345,8 +383,13 @@ MmCommunicationInitialize (
                   NULL,
                   &mSetVirtualAddressMapEvent
                   );
-  if (Status == EFI_SUCCESS) {
-    return Status;
+  ASSERT_EFI_ERROR (Status);
+
+  for (Index = 0; Index < ARRAY_SIZE (mGuidedEventGuid); Index++) {
+    Status = gBS->CreateEventEx (EVT_NOTIFY_SIGNAL, TPL_CALLBACK,
+                    MmGuidedEventNotify, mGuidedEventGuid[Index],
+                    mGuidedEventGuid[Index], &mGuidedEvent[Index]);
+    ASSERT_EFI_ERROR (Status);
   }
 
   gBS->UninstallProtocolInterface (
diff --git a/ArmPkg/Drivers/MmCommunicationDxe/MmCommunication.inf b/ArmPkg/Drivers/MmCommunicationDxe/MmCommunication.inf
index 450fc9bb54..505228704e 100644
--- a/ArmPkg/Drivers/MmCommunicationDxe/MmCommunication.inf
+++ b/ArmPkg/Drivers/MmCommunicationDxe/MmCommunication.inf
@@ -42,6 +42,11 @@
 [Protocols]
   gEfiMmCommunicationProtocolGuid              ## PRODUCES
 
+[Guids]
+  gEfiEndOfDxeEventGroupGuid
+  gEfiEventExitBootServicesGuid
+  gEfiEventReadyToBootGuid
+
 [Pcd.common]
   gArmTokenSpaceGuid.PcdMmBufferBase
   gArmTokenSpaceGuid.PcdMmBufferSize