1 files changed, 49 insertions, 0 deletions
diff --git a/MdeModulePkg/Universal/Acpi/BootScriptExecutorDxe/ScriptExecute.c b/MdeModulePkg/Universal/Acpi/BootScriptExecutorDxe/ScriptExecute.c
index 6c1c304a15..3fbdf5cd83 100644
--- a/MdeModulePkg/Universal/Acpi/BootScriptExecutorDxe/ScriptExecute.c
+++ b/MdeModulePkg/Universal/Acpi/BootScriptExecutorDxe/ScriptExecute.c
@@ -209,6 +209,47 @@ S3BootScriptExecutorEntryFunction (
 }
 
 /**
+  Register image to memory profile.
+
+  @param FileName       File name of the image.
+  @param ImageBase      Image base address.
+  @param ImageSize      Image size.
+  @param FileType       File type of the image.
+
+**/
+VOID
+RegisterMemoryProfileImage (
+  IN EFI_GUID                       *FileName,
+  IN PHYSICAL_ADDRESS               ImageBase,
+  IN UINT64                         ImageSize,
+  IN EFI_FV_FILETYPE                FileType
+  )
+{
+  EFI_STATUS                        Status;
+  EDKII_MEMORY_PROFILE_PROTOCOL     *ProfileProtocol;
+  MEDIA_FW_VOL_FILEPATH_DEVICE_PATH *FilePath;
+  UINT8                             TempBuffer[sizeof (MEDIA_FW_VOL_FILEPATH_DEVICE_PATH) + sizeof (EFI_DEVICE_PATH_PROTOCOL)];
+
+  if ((PcdGet8 (PcdMemoryProfilePropertyMask) & BIT0) != 0) {
+
+    FilePath = (MEDIA_FW_VOL_FILEPATH_DEVICE_PATH *)TempBuffer;
+    Status = gBS->LocateProtocol (&gEdkiiMemoryProfileGuid, NULL, (VOID **) &ProfileProtocol);
+    if (!EFI_ERROR (Status)) {
+      EfiInitializeFwVolDevicepathNode (FilePath, FileName);
+      SetDevicePathEndNode (FilePath + 1);
+
+      Status = ProfileProtocol->RegisterImage (
+                                  ProfileProtocol,
+                                  (EFI_DEVICE_PATH_PROTOCOL *) FilePath,
+                                  ImageBase,
+                                  ImageSize,
+                                  FileType
+                                  );
+    }
+  }
+}
+
+/**
   This is the Event notification function to reload BootScriptExecutor image
   to RESERVED mem and save it to LockBox.
   
@@ -302,6 +343,14 @@ ReadyToLockEventNotify (
   // Flush the instruction cache so the image data is written before we execute it
   //
   InvalidateInstructionCacheRange ((VOID *)(UINTN)ImageContext.ImageAddress, (UINTN)ImageContext.ImageSize);
+
+  RegisterMemoryProfileImage (
+    &gEfiCallerIdGuid,
+    ImageContext.ImageAddress,
+    ImageContext.ImageSize,
+    EFI_FV_FILETYPE_DRIVER
+    );
+
   Status = ((EFI_IMAGE_ENTRY_POINT)(UINTN)(ImageContext.EntryPoint)) (NewImageHandle, gST);
   ASSERT_EFI_ERROR (Status);