From a44830727ac6bc6e2851f3a2432b177b3168cb56 Mon Sep 17 00:00:00 2001
From: Wei6 Xu <wei6.xu@intel.com>
Date: Tue, 7 May 2024 01:32:48 +0800
Subject: [PATCH] StandaloneMmPkg/Core: Install Loaded Image Protocol for MM
 drivers

Install Loaded Image Protocol into MM handle database for each MM
driver.
Change EFI_MM_DRIVER_ENTRY structure definition to hold the Loaded
Image Protocol data directly, instead a pointer to the protocol, to
avoid allocating pool for each MM driver.

Cc: Ard Biesheuvel <ardb+tianocore@kernel.org>
Cc: Sami Mujawar <sami.mujawar@arm.com>
Cc: Ray Ni <ray.ni@intel.com>
Cc: Jiaxin Wu <jiaxin.wu@intel.com>
Signed-off-by: Wei6 Xu <wei6.xu@intel.com>
---
 StandaloneMmPkg/Core/Dispatcher.c       | 32 +++++++++++++++++++++++++
 StandaloneMmPkg/Core/StandaloneMmCore.h |  2 +-
 2 files changed, 33 insertions(+), 1 deletion(-)

diff --git a/StandaloneMmPkg/Core/Dispatcher.c b/StandaloneMmPkg/Core/Dispatcher.c
index 31ba236df1..dd8f79a2d2 100644
--- a/StandaloneMmPkg/Core/Dispatcher.c
+++ b/StandaloneMmPkg/Core/Dispatcher.c
@@ -185,6 +185,31 @@ MmLoadImage (
   DriverEntry->ImageBuffer     = DstBuffer;
   DriverEntry->NumberOfPage    = PageCount;
 
+  //
+  // Fill in the remaining fields of the Loaded Image Protocol instance.
+  // Note: ImageBase is an SMRAM address that can not be accessed outside of SMRAM if SMRAM window is closed.
+  //
+  DriverEntry->LoadedImage.Revision     = EFI_LOADED_IMAGE_PROTOCOL_REVISION;
+  DriverEntry->LoadedImage.ParentHandle = NULL;
+  DriverEntry->LoadedImage.SystemTable  = NULL;
+  DriverEntry->LoadedImage.DeviceHandle = NULL;
+  DriverEntry->LoadedImage.FilePath     = NULL;
+
+  DriverEntry->LoadedImage.ImageBase     = (VOID *)(UINTN)DriverEntry->ImageBuffer;
+  DriverEntry->LoadedImage.ImageSize     = ImageContext.ImageSize;
+  DriverEntry->LoadedImage.ImageCodeType = EfiRuntimeServicesCode;
+  DriverEntry->LoadedImage.ImageDataType = EfiRuntimeServicesData;
+
+  //
+  // Install Loaded Image protocol into MM handle database for the MM Driver
+  //
+  MmInstallProtocolInterface (
+    &DriverEntry->ImageHandle,
+    &gEfiLoadedImageProtocolGuid,
+    EFI_NATIVE_INTERFACE,
+    &DriverEntry->LoadedImage
+    );
+
   //
   // Print the load address and the PDB file name if it is available
   //
@@ -430,6 +455,13 @@ MmDispatcher (
       if (EFI_ERROR (Status)) {
         DEBUG ((DEBUG_INFO, "StartImage Status - %r\n", Status));
         MmFreePages (DriverEntry->ImageBuffer, DriverEntry->NumberOfPage);
+        if (DriverEntry->ImageHandle != NULL) {
+          MmUninstallProtocolInterface (
+            DriverEntry->ImageHandle,
+            &gEfiLoadedImageProtocolGuid,
+            &DriverEntry->LoadedImage
+            );
+        }
       }
     }
 
diff --git a/StandaloneMmPkg/Core/StandaloneMmCore.h b/StandaloneMmPkg/Core/StandaloneMmCore.h
index ae6c840f0c..c08d1d76c0 100644
--- a/StandaloneMmPkg/Core/StandaloneMmCore.h
+++ b/StandaloneMmPkg/Core/StandaloneMmCore.h
@@ -85,7 +85,7 @@ typedef struct {
   BOOLEAN                       DepexProtocolError;
 
   EFI_HANDLE                    ImageHandle;
-  EFI_LOADED_IMAGE_PROTOCOL     *LoadedImage;
+  EFI_LOADED_IMAGE_PROTOCOL     LoadedImage;
   //
   // Image EntryPoint in MMRAM
   //