From 1c6ce560f042f70f4a8d7a26652ceab68b5df8d1 Mon Sep 17 00:00:00 2001
From: Crystal Lee <CrystalLee@ami.com>
Date: Mon, 11 Nov 2024 15:49:47 +0800
Subject: [PATCH] ArmPkg: Add a protocol notification for
 gEfiDxeMmReadyToLockProtocolGuid

Add and install a protocol notification handler for
gEfiDxeMmReadyToLockProtocolGuid, in order to to trigger MMI
handler to install gEfiMmReadyToLockProtocolGuid protocol in Mm.

Signed-off-by: Crystal Lee <CrystalLee@ami.com>
---
 .../MmCommunicationDxe/MmCommunication.c      | 38 +++++++++++++++++++
 .../MmCommunicationDxe/MmCommunication.inf    |  1 +
 2 files changed, 39 insertions(+)

diff --git a/ArmPkg/Drivers/MmCommunicationDxe/MmCommunication.c b/ArmPkg/Drivers/MmCommunicationDxe/MmCommunication.c
index 3c325d175a..2d5e362259 100644
--- a/ArmPkg/Drivers/MmCommunicationDxe/MmCommunication.c
+++ b/ArmPkg/Drivers/MmCommunicationDxe/MmCommunication.c
@@ -14,6 +14,7 @@
 #include <Library/HobLib.h>
 #include <Library/PcdLib.h>
 #include <Library/UefiBootServicesTableLib.h>
+#include <Library/UefiLib.h>
 #include <Library/UefiRuntimeServicesTableLib.h>
 
 #include <Protocol/MmCommunication2.h>
@@ -553,6 +554,34 @@ STATIC EFI_GUID *CONST  mGuidedEventGuid[] = {
 
 STATIC EFI_EVENT  mGuidedEvent[ARRAY_SIZE (mGuidedEventGuid)];
 
+/**
+  Event notification that is fired when ReadyToLockProtocol is signaled.
+
+  @param  Event                 The Event that is being processed, not used.
+  @param  Context               Event Context, not used.
+
+**/
+VOID
+EFIAPI
+MMReadyToLockProtocolNotify (
+  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, &gEfiDxeMmReadyToLockProtocolGuid);
+  Header.MessageLength = 1;
+  Header.Data[0]       = 0;
+
+  Size = sizeof (Header);
+  MmCommunication2Communicate (&mMmCommunication2, &Header, &Header, &Size);
+}
+
 /**
   Event notification that is fired when GUIDed Event Group is signaled.
 
@@ -680,6 +709,15 @@ MmCommunication2Initialize (
                   );
   ASSERT_EFI_ERROR (Status);
 
+  Status = EfiNamedEventListen (
+             &gEfiDxeMmReadyToLockProtocolGuid,
+             TPL_NOTIFY,
+             MMReadyToLockProtocolNotify,
+             NULL,
+             NULL
+             );
+  ASSERT_EFI_ERROR (Status);
+
   for (Index = 0; Index < ARRAY_SIZE (mGuidedEventGuid); Index++) {
     Status = gBS->CreateEventEx (
                     EVT_NOTIFY_SIGNAL,
diff --git a/ArmPkg/Drivers/MmCommunicationDxe/MmCommunication.inf b/ArmPkg/Drivers/MmCommunicationDxe/MmCommunication.inf
index 1cd0a25e69..6a57f0c79f 100644
--- a/ArmPkg/Drivers/MmCommunicationDxe/MmCommunication.inf
+++ b/ArmPkg/Drivers/MmCommunicationDxe/MmCommunication.inf
@@ -41,6 +41,7 @@
   UefiDriverEntryPoint
 
 [Protocols]
+  gEfiDxeMmReadyToLockProtocolGuid              ## UNDEFINED # SmiHandlerRegister
   gEfiMmCommunication2ProtocolGuid              ## PRODUCES
 
 [Guids]