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

Retrieves the MM Core image info from Memory Allocation HOB reported
by MM IPL. Then install Loaded Image Protocol for MM Core with the
image info from HOB.

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/StandaloneMmCore.c | 85 +++++++++++++++++++++++++
 1 file changed, 85 insertions(+)

diff --git a/StandaloneMmPkg/Core/StandaloneMmCore.c b/StandaloneMmPkg/Core/StandaloneMmCore.c
index b70e340c04..c6f8886750 100644
--- a/StandaloneMmPkg/Core/StandaloneMmCore.c
+++ b/StandaloneMmPkg/Core/StandaloneMmCore.c
@@ -317,6 +317,86 @@ MmEndOfDxeHandler (
   return Status;
 }
 
+/**
+  Install LoadedImage protocol for MM Core.
+
+**/
+VOID
+MmCoreInstallLoadedImage (
+  VOID
+  )
+{
+  EFI_STATUS                 Status;
+  EFI_PHYSICAL_ADDRESS       MmCoreImageBaseAddress;
+  UINT64                     MmCoreImageLength;
+  EFI_PEI_HOB_POINTERS       Hob;
+  EFI_LOADED_IMAGE_PROTOCOL  *LoadedImage;
+  EFI_HANDLE                 ImageHandle;
+
+  //
+  // Searching for Memory Allocation HOB
+  //
+  Hob.Raw = GetHobList ();
+  while ((Hob.Raw = GetNextHob (EFI_HOB_TYPE_MEMORY_ALLOCATION, Hob.Raw)) != NULL) {
+    //
+    // Find MM Core HOB
+    //
+    if (CompareGuid (
+          &Hob.MemoryAllocationModule->MemoryAllocationHeader.Name,
+          &gEfiHobMemoryAllocModuleGuid
+          ))
+    {
+      if (CompareGuid (&Hob.MemoryAllocationModule->ModuleName, &gEfiCallerIdGuid)) {
+        break;
+      }
+    }
+
+    Hob.Raw = GET_NEXT_HOB (Hob);
+  }
+
+  if (Hob.Raw == NULL) {
+    return;
+  }
+
+  MmCoreImageBaseAddress = Hob.MemoryAllocationModule->MemoryAllocationHeader.MemoryBaseAddress;
+  MmCoreImageLength      = Hob.MemoryAllocationModule->MemoryAllocationHeader.MemoryLength;
+
+  //
+  // Allocate a Loaded Image Protocol in MM
+  //
+  LoadedImage = AllocatePool (sizeof (EFI_LOADED_IMAGE_PROTOCOL));
+  ASSERT (LoadedImage != NULL);
+  if (LoadedImage == NULL) {
+    return;
+  }
+
+  ZeroMem (LoadedImage, sizeof (EFI_LOADED_IMAGE_PROTOCOL));
+
+  //
+  // Fill in the remaining fields of the Loaded Image Protocol instance.
+  //
+  LoadedImage->Revision     = EFI_LOADED_IMAGE_PROTOCOL_REVISION;
+  LoadedImage->ParentHandle = NULL;
+  LoadedImage->SystemTable  = NULL;
+
+  LoadedImage->ImageBase     = (VOID *)(UINTN)MmCoreImageBaseAddress;
+  LoadedImage->ImageSize     = MmCoreImageLength;
+  LoadedImage->ImageCodeType = EfiRuntimeServicesCode;
+  LoadedImage->ImageDataType = EfiRuntimeServicesData;
+
+  //
+  // Create a new image handle in the MM handle database for the MM Core
+  //
+  ImageHandle = NULL;
+  Status      = MmInstallProtocolInterface (
+                  &ImageHandle,
+                  &gEfiLoadedImageProtocolGuid,
+                  EFI_NATIVE_INTERFACE,
+                  LoadedImage
+                  );
+  ASSERT_EFI_ERROR (Status);
+}
+
 /**
   The main entry point to MM Foundation.
 
@@ -657,6 +737,11 @@ StandaloneMmMain (
     DEBUG ((DEBUG_INFO, "MmiHandlerRegister - GUID %g - Status %d\n", mMmCoreMmiHandlers[Index].HandlerType, Status));
   }
 
+  //
+  // Install Loaded Image Protocol form MM Core
+  //
+  MmCoreInstallLoadedImage ();
+
   DEBUG ((DEBUG_INFO, "MmMain Done!\n"));
 
   return EFI_SUCCESS;