]> git.proxmox.com Git - mirror_edk2.git/blobdiff - MdeModulePkg/Core/PiSmmCore/SmramProfileRecord.c
MdeModulePkg: remove PE/COFF header workaround for ELILO on IPF
[mirror_edk2.git] / MdeModulePkg / Core / PiSmmCore / SmramProfileRecord.c
index e119fadb67e5cd3669972e226bf104dd5f29f54f..e8d60c2d2357e43b3c815d902fd84da64337543c 100644 (file)
@@ -1,7 +1,7 @@
 /** @file\r
   Support routines for SMRAM profile.\r
 \r
-  Copyright (c) 2014, Intel Corporation. All rights reserved.<BR>\r
+  Copyright (c) 2014 - 2018, Intel Corporation. All rights reserved.<BR>\r
   This program and the accompanying materials\r
   are licensed and made available under the terms and conditions of the BSD License\r
   which accompanies this distribution.  The full text of the license may be found at\r
 #include "PiSmmCore.h"\r
 \r
 #define IS_SMRAM_PROFILE_ENABLED ((PcdGet8 (PcdMemoryProfilePropertyMask) & BIT1) != 0)\r
+#define IS_UEFI_MEMORY_PROFILE_ENABLED ((PcdGet8 (PcdMemoryProfilePropertyMask) & BIT0) != 0)\r
+\r
+#define GET_OCCUPIED_SIZE(ActualSize, Alignment) \\r
+  ((ActualSize) + (((Alignment) - ((ActualSize) & ((Alignment) - 1))) & ((Alignment) - 1)))\r
 \r
 typedef struct {\r
   UINT32                        Signature;\r
@@ -26,12 +30,14 @@ typedef struct {
   UINT32                        Signature;\r
   MEMORY_PROFILE_DRIVER_INFO    DriverInfo;\r
   LIST_ENTRY                    *AllocInfoList;\r
+  CHAR8                         *PdbString;\r
   LIST_ENTRY                    Link;\r
 } MEMORY_PROFILE_DRIVER_INFO_DATA;\r
 \r
 typedef struct {\r
   UINT32                        Signature;\r
   MEMORY_PROFILE_ALLOC_INFO     AllocInfo;\r
+  CHAR8                         *ActionString;\r
   LIST_ENTRY                    Link;\r
 } MEMORY_PROFILE_ALLOC_INFO_DATA;\r
 \r
@@ -69,52 +75,184 @@ GLOBAL_REMOVE_IF_UNREFERENCED MEMORY_PROFILE_CONTEXT_DATA mSmramProfileContext =
   },\r
   &mImageQueue,\r
 };\r
-GLOBAL_REMOVE_IF_UNREFERENCED MEMORY_PROFILE_CONTEXT_DATA *mSmramProfileContextPtr;\r
+GLOBAL_REMOVE_IF_UNREFERENCED MEMORY_PROFILE_CONTEXT_DATA *mSmramProfileContextPtr = NULL;\r
 \r
-BOOLEAN mSmramReadyToLock;\r
-BOOLEAN mSmramProfileRecordingStatus = FALSE;\r
+GLOBAL_REMOVE_IF_UNREFERENCED BOOLEAN mSmramReadyToLock;\r
+GLOBAL_REMOVE_IF_UNREFERENCED BOOLEAN mSmramProfileGettingStatus = FALSE;\r
+GLOBAL_REMOVE_IF_UNREFERENCED BOOLEAN mSmramProfileRecordingEnable = MEMORY_PROFILE_RECORDING_DISABLE;\r
+GLOBAL_REMOVE_IF_UNREFERENCED EFI_DEVICE_PATH_PROTOCOL *mSmramProfileDriverPath;\r
+GLOBAL_REMOVE_IF_UNREFERENCED UINTN                    mSmramProfileDriverPathSize;\r
 \r
 /**\r
-  Return SMRAM profile context.\r
-\r
-  @return SMRAM profile context.\r
+  Dump SMRAM infromation.\r
 \r
 **/\r
-MEMORY_PROFILE_CONTEXT_DATA *\r
-GetSmramProfileContext (\r
+VOID\r
+DumpSmramInfo (\r
   VOID\r
-  )\r
-{\r
-  return mSmramProfileContextPtr;\r
-}\r
+  );\r
 \r
 /**\r
-  Retrieves the magic value from the PE/COFF header.\r
+  Get memory profile data.\r
 \r
-  @param Hdr    The buffer in which to return the PE32, PE32+, or TE header.\r
+  @param[in]      This              The EDKII_SMM_MEMORY_PROFILE_PROTOCOL instance.\r
+  @param[in, out] ProfileSize       On entry, points to the size in bytes of the ProfileBuffer.\r
+                                    On return, points to the size of the data returned in ProfileBuffer.\r
+  @param[out]     ProfileBuffer     Profile buffer.\r
 \r
-  @return EFI_IMAGE_NT_OPTIONAL_HDR32_MAGIC - Image is PE32\r
-  @return EFI_IMAGE_NT_OPTIONAL_HDR64_MAGIC - Image is PE32+\r
+  @return EFI_SUCCESS               Get the memory profile data successfully.\r
+  @return EFI_UNSUPPORTED           Memory profile is unsupported.\r
+  @return EFI_BUFFER_TO_SMALL       The ProfileSize is too small for the resulting data.\r
+                                    ProfileSize is updated with the size required.\r
 \r
 **/\r
-UINT16\r
-InternalPeCoffGetPeHeaderMagicValue (\r
-  IN  EFI_IMAGE_OPTIONAL_HEADER_PTR_UNION  Hdr\r
+EFI_STATUS\r
+EFIAPI\r
+SmramProfileProtocolGetData (\r
+  IN     EDKII_SMM_MEMORY_PROFILE_PROTOCOL  *This,\r
+  IN OUT UINT64                             *ProfileSize,\r
+     OUT VOID                               *ProfileBuffer\r
+  );\r
+\r
+/**\r
+  Register image to memory profile.\r
+\r
+  @param[in] This               The EDKII_SMM_MEMORY_PROFILE_PROTOCOL instance.\r
+  @param[in] FilePath           File path of the image.\r
+  @param[in] ImageBase          Image base address.\r
+  @param[in] ImageSize          Image size.\r
+  @param[in] FileType           File type of the image.\r
+\r
+  @return EFI_SUCCESS           Register successfully.\r
+  @return EFI_UNSUPPORTED       Memory profile is unsupported,\r
+                                or memory profile for the image is not required.\r
+  @return EFI_OUT_OF_RESOURCE   No enough resource for this register.\r
+\r
+**/\r
+EFI_STATUS\r
+EFIAPI\r
+SmramProfileProtocolRegisterImage (\r
+  IN EDKII_SMM_MEMORY_PROFILE_PROTOCOL  *This,\r
+  IN EFI_DEVICE_PATH_PROTOCOL           *FilePath,\r
+  IN PHYSICAL_ADDRESS                   ImageBase,\r
+  IN UINT64                             ImageSize,\r
+  IN EFI_FV_FILETYPE                    FileType\r
+  );\r
+\r
+/**\r
+  Unregister image from memory profile.\r
+\r
+  @param[in] This               The EDKII_SMM_MEMORY_PROFILE_PROTOCOL instance.\r
+  @param[in] FilePath           File path of the image.\r
+  @param[in] ImageBase          Image base address.\r
+  @param[in] ImageSize          Image size.\r
+\r
+  @return EFI_SUCCESS           Unregister successfully.\r
+  @return EFI_UNSUPPORTED       Memory profile is unsupported,\r
+                                or memory profile for the image is not required.\r
+  @return EFI_NOT_FOUND         The image is not found.\r
+\r
+**/\r
+EFI_STATUS\r
+EFIAPI\r
+SmramProfileProtocolUnregisterImage (\r
+  IN EDKII_SMM_MEMORY_PROFILE_PROTOCOL  *This,\r
+  IN EFI_DEVICE_PATH_PROTOCOL           *FilePath,\r
+  IN PHYSICAL_ADDRESS                   ImageBase,\r
+  IN UINT64                             ImageSize\r
+  );\r
+\r
+/**\r
+  Get memory profile recording state.\r
+\r
+  @param[in]  This              The EDKII_SMM_MEMORY_PROFILE_PROTOCOL instance.\r
+  @param[out] RecordingState    Recording state.\r
+\r
+  @return EFI_SUCCESS           Memory profile recording state is returned.\r
+  @return EFI_UNSUPPORTED       Memory profile is unsupported.\r
+  @return EFI_INVALID_PARAMETER RecordingState is NULL.\r
+\r
+**/\r
+EFI_STATUS\r
+EFIAPI\r
+SmramProfileProtocolGetRecordingState (\r
+  IN EDKII_SMM_MEMORY_PROFILE_PROTOCOL  *This,\r
+  OUT BOOLEAN                           *RecordingState\r
+  );\r
+\r
+/**\r
+  Set memory profile recording state.\r
+\r
+  @param[in] This               The EDKII_SMM_MEMORY_PROFILE_PROTOCOL instance.\r
+  @param[in] RecordingState     Recording state.\r
+\r
+  @return EFI_SUCCESS           Set memory profile recording state successfully.\r
+  @return EFI_UNSUPPORTED       Memory profile is unsupported.\r
+\r
+**/\r
+EFI_STATUS\r
+EFIAPI\r
+SmramProfileProtocolSetRecordingState (\r
+  IN EDKII_SMM_MEMORY_PROFILE_PROTOCOL  *This,\r
+  IN BOOLEAN                            RecordingState\r
+  );\r
+\r
+/**\r
+  Record memory profile of multilevel caller.\r
+\r
+  @param[in] This               The EDKII_SMM_MEMORY_PROFILE_PROTOCOL instance.\r
+  @param[in] CallerAddress      Address of caller.\r
+  @param[in] Action             Memory profile action.\r
+  @param[in] MemoryType         Memory type.\r
+                                EfiMaxMemoryType means the MemoryType is unknown.\r
+  @param[in] Buffer             Buffer address.\r
+  @param[in] Size               Buffer size.\r
+  @param[in] ActionString       String for memory profile action.\r
+                                Only needed for user defined allocate action.\r
+\r
+  @return EFI_SUCCESS           Memory profile is updated.\r
+  @return EFI_UNSUPPORTED       Memory profile is unsupported,\r
+                                or memory profile for the image is not required,\r
+                                or memory profile for the memory type is not required.\r
+  @return EFI_ACCESS_DENIED     It is during memory profile data getting.\r
+  @return EFI_ABORTED           Memory profile recording is not enabled.\r
+  @return EFI_OUT_OF_RESOURCES  No enough resource to update memory profile for allocate action.\r
+  @return EFI_NOT_FOUND         No matched allocate info found for free action.\r
+\r
+**/\r
+EFI_STATUS\r
+EFIAPI\r
+SmramProfileProtocolRecord (\r
+  IN EDKII_SMM_MEMORY_PROFILE_PROTOCOL  *This,\r
+  IN PHYSICAL_ADDRESS                   CallerAddress,\r
+  IN MEMORY_PROFILE_ACTION              Action,\r
+  IN EFI_MEMORY_TYPE                    MemoryType,\r
+  IN VOID                               *Buffer,\r
+  IN UINTN                              Size,\r
+  IN CHAR8                              *ActionString OPTIONAL\r
+  );\r
+\r
+GLOBAL_REMOVE_IF_UNREFERENCED EDKII_SMM_MEMORY_PROFILE_PROTOCOL mSmmProfileProtocol = {\r
+  SmramProfileProtocolGetData,\r
+  SmramProfileProtocolRegisterImage,\r
+  SmramProfileProtocolUnregisterImage,\r
+  SmramProfileProtocolGetRecordingState,\r
+  SmramProfileProtocolSetRecordingState,\r
+  SmramProfileProtocolRecord,\r
+};\r
+\r
+/**\r
+  Return SMRAM profile context.\r
+\r
+  @return SMRAM profile context.\r
+\r
+**/\r
+MEMORY_PROFILE_CONTEXT_DATA *\r
+GetSmramProfileContext (\r
+  VOID\r
   )\r
 {\r
-  //\r
-  // NOTE: Some versions of Linux ELILO for Itanium have an incorrect magic value\r
-  //       in the PE/COFF Header.  If the MachineType is Itanium(IA64) and the\r
-  //       Magic value in the OptionalHeader is  EFI_IMAGE_NT_OPTIONAL_HDR32_MAGIC\r
-  //       then override the returned value to EFI_IMAGE_NT_OPTIONAL_HDR64_MAGIC\r
-  //\r
-  if (Hdr.Pe32->FileHeader.Machine == IMAGE_FILE_MACHINE_IA64 && Hdr.Pe32->OptionalHeader.Magic == EFI_IMAGE_NT_OPTIONAL_HDR32_MAGIC) {\r
-    return EFI_IMAGE_NT_OPTIONAL_HDR64_MAGIC;\r
-  }\r
-  //\r
-  // Return the magic value from the PC/COFF Optional Header\r
-  //\r
-  return Hdr.Pe32->OptionalHeader.Magic;\r
+  return mSmramProfileContextPtr;\r
 }\r
 \r
 /**\r
@@ -153,7 +291,7 @@ InternalPeCoffGetSubsystem (
   if (Hdr.Te->Signature == EFI_TE_IMAGE_HEADER_SIGNATURE) {\r
     return Hdr.Te->Subsystem;\r
   } else if (Hdr.Pe32->Signature == EFI_IMAGE_NT_SIGNATURE)  {\r
-    Magic = InternalPeCoffGetPeHeaderMagicValue (Hdr);\r
+    Magic = Hdr.Pe32->OptionalHeader.Magic;\r
     if (Magic == EFI_IMAGE_NT_OPTIONAL_HDR32_MAGIC) {\r
       return Hdr.Pe32->OptionalHeader.Subsystem;\r
     } else if (Magic == EFI_IMAGE_NT_OPTIONAL_HDR64_MAGIC) {\r
@@ -230,7 +368,6 @@ InternalPeCoffGetEntryPoint (
   @param ImageSize      Image size.\r
   @param EntryPoint     Entry point of the image.\r
   @param ImageSubsystem Image subsystem of the image.\r
-\r
   @param FileType       File type of the image.\r
 \r
   @return Pointer to memory profile driver info.\r
@@ -251,25 +388,40 @@ BuildDriverInfo (
   MEMORY_PROFILE_DRIVER_INFO        *DriverInfo;\r
   MEMORY_PROFILE_DRIVER_INFO_DATA   *DriverInfoData;\r
   VOID                              *EntryPointInImage;\r
+  CHAR8                             *PdbString;\r
+  UINTN                             PdbSize;\r
+  UINTN                             PdbOccupiedSize;\r
+\r
+  PdbSize = 0;\r
+  PdbOccupiedSize = 0;\r
+  PdbString = NULL;\r
+  if (ImageBase != 0) {\r
+    PdbString = PeCoffLoaderGetPdbPointer ((VOID*) (UINTN) ImageBase);\r
+    if (PdbString != NULL) {\r
+      PdbSize = AsciiStrSize (PdbString);\r
+      PdbOccupiedSize = GET_OCCUPIED_SIZE (PdbSize, sizeof (UINT64));\r
+    }\r
+  }\r
 \r
   //\r
   // Use SmmInternalAllocatePool() that will not update profile for this AllocatePool action.\r
   //\r
   Status = SmmInternalAllocatePool (\r
              EfiRuntimeServicesData,\r
-             sizeof (*DriverInfoData) + sizeof (LIST_ENTRY),\r
+             sizeof (*DriverInfoData) + sizeof (LIST_ENTRY) + PdbSize,\r
              (VOID **) &DriverInfoData\r
              );\r
   if (EFI_ERROR (Status)) {\r
     return NULL;\r
   }\r
+  ASSERT (DriverInfoData != NULL);\r
 \r
   ZeroMem (DriverInfoData, sizeof (*DriverInfoData));\r
 \r
   DriverInfo = &DriverInfoData->DriverInfo;\r
   DriverInfoData->Signature = MEMORY_PROFILE_DRIVER_INFO_SIGNATURE;\r
   DriverInfo->Header.Signature = MEMORY_PROFILE_DRIVER_INFO_SIGNATURE;\r
-  DriverInfo->Header.Length = sizeof (MEMORY_PROFILE_DRIVER_INFO);\r
+  DriverInfo->Header.Length = (UINT16) (sizeof (MEMORY_PROFILE_DRIVER_INFO) + PdbOccupiedSize);\r
   DriverInfo->Header.Revision = MEMORY_PROFILE_DRIVER_INFO_REVISION;\r
   if (FileName != NULL) {\r
     CopyMem (&DriverInfo->FileName, FileName, sizeof (EFI_GUID));\r
@@ -293,6 +445,14 @@ BuildDriverInfo (
   DriverInfo->CurrentUsage = 0;\r
   DriverInfo->PeakUsage = 0;\r
   DriverInfo->AllocRecordCount = 0;\r
+  if (PdbSize != 0) {\r
+    DriverInfo->PdbStringOffset = (UINT16) sizeof (MEMORY_PROFILE_DRIVER_INFO);\r
+    DriverInfoData->PdbString = (CHAR8 *) (DriverInfoData->AllocInfoList + 1);\r
+    CopyMem (DriverInfoData->PdbString, PdbString, PdbSize);\r
+  } else {\r
+    DriverInfo->PdbStringOffset = 0;\r
+    DriverInfoData->PdbString = NULL;\r
+  }\r
 \r
   InsertTailList (ContextData->DriverInfoList, &DriverInfoData->Link);\r
   ContextData->Context.ImageCount ++;\r
@@ -323,7 +483,7 @@ RegisterImageToDxe (
   MEDIA_FW_VOL_FILEPATH_DEVICE_PATH *FilePath;\r
   UINT8                             TempBuffer[sizeof (MEDIA_FW_VOL_FILEPATH_DEVICE_PATH) + sizeof (EFI_DEVICE_PATH_PROTOCOL)];\r
 \r
-  if (IS_SMRAM_PROFILE_ENABLED) {\r
+  if (IS_UEFI_MEMORY_PROFILE_ENABLED) {\r
 \r
     FilePath = (MEDIA_FW_VOL_FILEPATH_DEVICE_PATH *)TempBuffer;\r
     Status = gBS->LocateProtocol (&gEdkiiMemoryProfileGuid, NULL, (VOID **) &ProfileProtocol);\r
@@ -362,7 +522,7 @@ UnregisterImageFromDxe (
   MEDIA_FW_VOL_FILEPATH_DEVICE_PATH *FilePath;\r
   UINT8                             TempBuffer[sizeof (MEDIA_FW_VOL_FILEPATH_DEVICE_PATH) + sizeof (EFI_DEVICE_PATH_PROTOCOL)];\r
 \r
-  if (IS_SMRAM_PROFILE_ENABLED) {\r
+  if (IS_UEFI_MEMORY_PROFILE_ENABLED) {\r
 \r
     FilePath = (MEDIA_FW_VOL_FILEPATH_DEVICE_PATH *)TempBuffer;\r
     Status = gBS->LocateProtocol (&gEdkiiMemoryProfileGuid, NULL, (VOID *) &ProfileProtocol);\r
@@ -380,6 +540,65 @@ UnregisterImageFromDxe (
   }\r
 }\r
 \r
+/**\r
+  Return if record for this driver is needed..\r
+\r
+  @param DriverFilePath     Driver file path.\r
+\r
+  @retval TRUE              Record for this driver is needed.\r
+  @retval FALSE             Record for this driver is not needed.\r
+\r
+**/\r
+BOOLEAN\r
+NeedRecordThisDriver (\r
+  IN EFI_DEVICE_PATH_PROTOCOL       *DriverFilePath\r
+  )\r
+{\r
+  EFI_DEVICE_PATH_PROTOCOL    *TmpDevicePath;\r
+  EFI_DEVICE_PATH_PROTOCOL    *DevicePathInstance;\r
+  UINTN                       DevicePathSize;\r
+  UINTN                       FilePathSize;\r
+\r
+  if (!IsDevicePathValid (mSmramProfileDriverPath, mSmramProfileDriverPathSize)) {\r
+    //\r
+    // Invalid Device Path means record all.\r
+    //\r
+    return TRUE;\r
+  }\r
+\r
+  //\r
+  // Record FilePath without end node.\r
+  //\r
+  FilePathSize = GetDevicePathSize (DriverFilePath) - sizeof(EFI_DEVICE_PATH_PROTOCOL);\r
+\r
+  DevicePathInstance = mSmramProfileDriverPath;\r
+  do {\r
+    //\r
+    // Find End node (it might be END_ENTIRE or END_INSTANCE)\r
+    //\r
+    TmpDevicePath = DevicePathInstance;\r
+    while (!IsDevicePathEndType (TmpDevicePath)) {\r
+      TmpDevicePath = NextDevicePathNode (TmpDevicePath);\r
+    }\r
+\r
+    //\r
+    // Do not compare END node\r
+    //\r
+    DevicePathSize = (UINTN)TmpDevicePath - (UINTN)DevicePathInstance;\r
+    if ((FilePathSize == DevicePathSize) &&\r
+        (CompareMem (DriverFilePath, DevicePathInstance, DevicePathSize) == 0)) {\r
+      return TRUE;\r
+    }\r
+\r
+    //\r
+    // Get next instance\r
+    //\r
+    DevicePathInstance = (EFI_DEVICE_PATH_PROTOCOL *)((UINTN)DevicePathInstance + DevicePathSize + DevicePathNodeLength(TmpDevicePath));\r
+  } while (DevicePathSubType (TmpDevicePath) != END_ENTIRE_DEVICE_PATH_SUBTYPE);\r
+\r
+  return FALSE;\r
+}\r
+\r
 /**\r
   Register SMM Core to SMRAM profile.\r
 \r
@@ -396,15 +615,16 @@ RegisterSmmCore (
 {\r
   MEMORY_PROFILE_DRIVER_INFO_DATA   *DriverInfoData;\r
   PHYSICAL_ADDRESS                  ImageBase;\r
+  UINT8                             TempBuffer[sizeof(MEDIA_FW_VOL_FILEPATH_DEVICE_PATH) + sizeof(EFI_DEVICE_PATH_PROTOCOL)];\r
+  MEDIA_FW_VOL_FILEPATH_DEVICE_PATH *FilePath;\r
 \r
-  ASSERT (ContextData != NULL);\r
+  FilePath = (MEDIA_FW_VOL_FILEPATH_DEVICE_PATH *) TempBuffer;\r
+  EfiInitializeFwVolDevicepathNode (FilePath, &gEfiCallerIdGuid);\r
+  SetDevicePathEndNode (FilePath + 1);\r
 \r
-  RegisterImageToDxe (\r
-    &gEfiCallerIdGuid,\r
-    gSmmCorePrivate->PiSmmCoreImageBase,\r
-    gSmmCorePrivate->PiSmmCoreImageSize,\r
-    EFI_FV_FILETYPE_SMM_CORE\r
-    );\r
+  if (!NeedRecordThisDriver ((EFI_DEVICE_PATH_PROTOCOL *) FilePath)) {\r
+    return FALSE;\r
+  }\r
 \r
   ImageBase = gSmmCorePrivate->PiSmmCoreImageBase;\r
   DriverInfoData = BuildDriverInfo (\r
@@ -434,6 +654,13 @@ SmramProfileInit (
 {\r
   MEMORY_PROFILE_CONTEXT_DATA *SmramProfileContext;\r
 \r
+  RegisterImageToDxe (\r
+    &gEfiCallerIdGuid,\r
+    gSmmCorePrivate->PiSmmCoreImageBase,\r
+    gSmmCorePrivate->PiSmmCoreImageSize,\r
+    EFI_FV_FILETYPE_SMM_CORE\r
+    );\r
+\r
   if (!IS_SMRAM_PROFILE_ENABLED) {\r
     return;\r
   }\r
@@ -443,7 +670,14 @@ SmramProfileInit (
     return;\r
   }\r
 \r
-  mSmramProfileRecordingStatus = TRUE;\r
+  mSmramProfileGettingStatus = FALSE;\r
+  if ((PcdGet8 (PcdMemoryProfilePropertyMask) & BIT7) != 0) {\r
+    mSmramProfileRecordingEnable = MEMORY_PROFILE_RECORDING_DISABLE;\r
+  } else {\r
+    mSmramProfileRecordingEnable = MEMORY_PROFILE_RECORDING_ENABLE;\r
+  }\r
+  mSmramProfileDriverPathSize = PcdGetSize (PcdMemoryProfileDriverPath);\r
+  mSmramProfileDriverPath = AllocateCopyPool (mSmramProfileDriverPathSize, PcdGetPtr (PcdMemoryProfileDriverPath));\r
   mSmramProfileContextPtr = &mSmramProfileContext;\r
 \r
   RegisterSmmCore (&mSmramProfileContext);\r
@@ -451,17 +685,76 @@ SmramProfileInit (
   DEBUG ((EFI_D_INFO, "SmramProfileInit SmramProfileContext - 0x%x\n", &mSmramProfileContext));\r
 }\r
 \r
+/**\r
+  Install SMRAM profile protocol.\r
+\r
+**/\r
+VOID\r
+SmramProfileInstallProtocol (\r
+  VOID\r
+  )\r
+{\r
+  EFI_HANDLE    Handle;\r
+  EFI_STATUS    Status;\r
+\r
+  if (!IS_SMRAM_PROFILE_ENABLED) {\r
+    return;\r
+  }\r
+\r
+  Handle = NULL;\r
+  Status = SmmInstallProtocolInterface (\r
+             &Handle,\r
+             &gEdkiiSmmMemoryProfileGuid,\r
+             EFI_NATIVE_INTERFACE,\r
+             &mSmmProfileProtocol\r
+             );\r
+  ASSERT_EFI_ERROR (Status);\r
+}\r
+\r
+/**\r
+  Get the GUID file name from the file path.\r
+\r
+  @param FilePath  File path.\r
+\r
+  @return The GUID file name from the file path.\r
+\r
+**/\r
+EFI_GUID *\r
+GetFileNameFromFilePath (\r
+  IN EFI_DEVICE_PATH_PROTOCOL   *FilePath\r
+  )\r
+{\r
+  MEDIA_FW_VOL_FILEPATH_DEVICE_PATH     *ThisFilePath;\r
+  EFI_GUID                              *FileName;\r
+\r
+  FileName = NULL;\r
+  if (FilePath != NULL) {\r
+    ThisFilePath = (MEDIA_FW_VOL_FILEPATH_DEVICE_PATH *) FilePath;\r
+    while (!IsDevicePathEnd (ThisFilePath)) {\r
+      FileName = EfiGetNameGuidFromFwVolDevicePathNode (ThisFilePath);\r
+      if (FileName != NULL) {\r
+        break;\r
+      }\r
+      ThisFilePath = (MEDIA_FW_VOL_FILEPATH_DEVICE_PATH *) NextDevicePathNode (ThisFilePath);\r
+    }\r
+  }\r
+\r
+  return FileName;\r
+}\r
+\r
 /**\r
   Register SMM image to SMRAM profile.\r
 \r
   @param DriverEntry    SMM image info.\r
   @param RegisterToDxe  Register image to DXE.\r
 \r
-  @retval TRUE          Register success.\r
-  @retval FALSE         Register fail.\r
+  @return EFI_SUCCESS           Register successfully.\r
+  @return EFI_UNSUPPORTED       Memory profile is unsupported,\r
+                                or memory profile for the image is not required.\r
+  @return EFI_OUT_OF_RESOURCES  No enough resource for this register.\r
 \r
 **/\r
-BOOLEAN\r
+EFI_STATUS\r
 RegisterSmramProfileImage (\r
   IN EFI_SMM_DRIVER_ENTRY   *DriverEntry,\r
   IN BOOLEAN                RegisterToDxe\r
@@ -469,10 +762,8 @@ RegisterSmramProfileImage (
 {\r
   MEMORY_PROFILE_CONTEXT_DATA       *ContextData;\r
   MEMORY_PROFILE_DRIVER_INFO_DATA   *DriverInfoData;\r
-\r
-  if (!IS_SMRAM_PROFILE_ENABLED) {\r
-    return FALSE;\r
-  }\r
+  UINT8                             TempBuffer[sizeof(MEDIA_FW_VOL_FILEPATH_DEVICE_PATH) + sizeof(EFI_DEVICE_PATH_PROTOCOL)];\r
+  MEDIA_FW_VOL_FILEPATH_DEVICE_PATH *FilePath;\r
 \r
   if (RegisterToDxe) {\r
     RegisterImageToDxe (\r
@@ -483,9 +774,21 @@ RegisterSmramProfileImage (
       );\r
   }\r
 \r
+  if (!IS_SMRAM_PROFILE_ENABLED) {\r
+    return EFI_UNSUPPORTED;\r
+  }\r
+\r
+  FilePath = (MEDIA_FW_VOL_FILEPATH_DEVICE_PATH *) TempBuffer;\r
+  EfiInitializeFwVolDevicepathNode (FilePath, &DriverEntry->FileName);\r
+  SetDevicePathEndNode (FilePath + 1);\r
+\r
+  if (!NeedRecordThisDriver ((EFI_DEVICE_PATH_PROTOCOL *) FilePath)) {\r
+    return EFI_UNSUPPORTED;\r
+  }\r
+\r
   ContextData = GetSmramProfileContext ();\r
   if (ContextData == NULL) {\r
-    return FALSE;\r
+    return EFI_UNSUPPORTED;\r
   }\r
 \r
   DriverInfoData = BuildDriverInfo (\r
@@ -498,10 +801,10 @@ RegisterSmramProfileImage (
                      EFI_FV_FILETYPE_SMM\r
                      );\r
   if (DriverInfoData == NULL) {\r
-    return FALSE;\r
+    return EFI_OUT_OF_RESOURCES;\r
   }\r
 \r
-  return TRUE;\r
+  return EFI_SUCCESS;\r
 }\r
 \r
 /**\r
@@ -548,42 +851,6 @@ GetMemoryProfileDriverInfoByFileNameAndAddress (
   return NULL;\r
 }\r
 \r
-/**\r
-  Search dummy image from SMRAM profile.\r
-\r
-  @param ContextData    Memory profile context.\r
-\r
-  @return Pointer to memory profile driver info.\r
-\r
-**/\r
-MEMORY_PROFILE_DRIVER_INFO_DATA *\r
-FindDummyImage (\r
-  IN MEMORY_PROFILE_CONTEXT_DATA    *ContextData\r
-  )\r
-{\r
-  MEMORY_PROFILE_DRIVER_INFO_DATA   *DriverInfoData;\r
-  LIST_ENTRY                        *DriverLink;\r
-  LIST_ENTRY                        *DriverInfoList;\r
-\r
-  DriverInfoList = ContextData->DriverInfoList;\r
-\r
-  for (DriverLink = DriverInfoList->ForwardLink;\r
-       DriverLink != DriverInfoList;\r
-       DriverLink = DriverLink->ForwardLink) {\r
-    DriverInfoData = CR (\r
-                   DriverLink,\r
-                   MEMORY_PROFILE_DRIVER_INFO_DATA,\r
-                   Link,\r
-                   MEMORY_PROFILE_DRIVER_INFO_SIGNATURE\r
-                   );\r
-    if (CompareGuid (&gZeroGuid, &DriverInfoData->DriverInfo.FileName)) {\r
-      return DriverInfoData;\r
-    }\r
-  }\r
-\r
-  return BuildDriverInfo (ContextData, &gZeroGuid, 0, 0, 0, 0, 0);\r
-}\r
-\r
 /**\r
   Search image from memory profile.\r
   It will return image, if (Address >= ImageBuffer) AND (Address < ImageBuffer + ImageSize)\r
@@ -623,10 +890,7 @@ GetMemoryProfileDriverInfoFromAddress (
     }\r
   }\r
 \r
-  //\r
-  // Should never come here.\r
-  //\r
-  return FindDummyImage (ContextData);\r
+  return NULL;\r
 }\r
 \r
 /**\r
@@ -635,11 +899,13 @@ GetMemoryProfileDriverInfoFromAddress (
   @param DriverEntry        SMM image info.\r
   @param UnregisterFromDxe  Unregister image from DXE.\r
 \r
-  @retval TRUE              Unregister success.\r
-  @retval FALSE             Unregister fail.\r
+  @return EFI_SUCCESS           Unregister successfully.\r
+  @return EFI_UNSUPPORTED       Memory profile is unsupported,\r
+                                or memory profile for the image is not required.\r
+  @return EFI_NOT_FOUND         The image is not found.\r
 \r
 **/\r
-BOOLEAN\r
+EFI_STATUS\r
 UnregisterSmramProfileImage (\r
   IN EFI_SMM_DRIVER_ENTRY  *DriverEntry,\r
   IN BOOLEAN               UnregisterFromDxe\r
@@ -651,10 +917,8 @@ UnregisterSmramProfileImage (
   EFI_GUID                          *FileName;\r
   PHYSICAL_ADDRESS                  ImageAddress;\r
   VOID                              *EntryPointInImage;\r
-\r
-  if (!IS_SMRAM_PROFILE_ENABLED) {\r
-    return FALSE;\r
-  }\r
+  UINT8                             TempBuffer[sizeof(MEDIA_FW_VOL_FILEPATH_DEVICE_PATH) + sizeof(EFI_DEVICE_PATH_PROTOCOL)];\r
+  MEDIA_FW_VOL_FILEPATH_DEVICE_PATH *FilePath;\r
 \r
   if (UnregisterFromDxe) {\r
     UnregisterImageFromDxe (\r
@@ -664,9 +928,21 @@ UnregisterSmramProfileImage (
       );\r
   }\r
 \r
+  if (!IS_SMRAM_PROFILE_ENABLED) {\r
+    return EFI_UNSUPPORTED;\r
+  }\r
+\r
+  FilePath = (MEDIA_FW_VOL_FILEPATH_DEVICE_PATH *) TempBuffer;\r
+  EfiInitializeFwVolDevicepathNode (FilePath, &DriverEntry->FileName);\r
+  SetDevicePathEndNode (FilePath + 1);\r
+\r
+  if (!NeedRecordThisDriver ((EFI_DEVICE_PATH_PROTOCOL *) FilePath)) {\r
+    return EFI_UNSUPPORTED;\r
+  }\r
+\r
   ContextData = GetSmramProfileContext ();\r
   if (ContextData == NULL) {\r
-    return FALSE;\r
+    return EFI_UNSUPPORTED;\r
   }\r
 \r
   DriverInfoData = NULL;\r
@@ -688,12 +964,13 @@ UnregisterSmramProfileImage (
     DriverInfoData = GetMemoryProfileDriverInfoFromAddress (ContextData, ImageAddress);\r
   }\r
   if (DriverInfoData == NULL) {\r
-    return FALSE;\r
+    return EFI_NOT_FOUND;\r
   }\r
 \r
   ContextData->Context.TotalImageSize -= DriverInfoData->DriverInfo.ImageSize;\r
 \r
-  DriverInfoData->DriverInfo.ImageBase = 0;\r
+  // Keep the ImageBase for RVA calculation in Application.\r
+  //DriverInfoData->DriverInfo.ImageBase = 0;\r
   DriverInfoData->DriverInfo.ImageSize = 0;\r
 \r
   if (DriverInfoData->DriverInfo.PeakUsage == 0) {\r
@@ -705,13 +982,12 @@ UnregisterSmramProfileImage (
     SmmInternalFreePool (DriverInfoData);\r
   }\r
 \r
-  return TRUE;\r
+  return EFI_SUCCESS;\r
 }\r
 \r
 /**\r
   Return if this memory type needs to be recorded into memory profile.\r
-  If BIOS memory type (0 ~ EfiMaxMemoryType), it checks bit (1 << MemoryType).\r
-  If OS memory type (0x80000000 ~ 0xFFFFFFFF), it checks bit63 - 0x8000000000000000.\r
+  Only need to record EfiRuntimeServicesCode and EfiRuntimeServicesData for SMRAM profile.\r
 \r
   @param MemoryType     Memory type.\r
 \r
@@ -726,12 +1002,13 @@ SmmCoreNeedRecordProfile (
 {\r
   UINT64 TestBit;\r
 \r
-  if ((UINT32) MemoryType >= 0x80000000) {\r
-    TestBit = BIT63;\r
-  } else {\r
-    TestBit = LShiftU64 (1, MemoryType);\r
+  if (MemoryType != EfiRuntimeServicesCode &&\r
+      MemoryType != EfiRuntimeServicesData) {\r
+    return FALSE;\r
   }\r
 \r
+  TestBit = LShiftU64 (1, MemoryType);\r
+\r
   if ((PcdGet64 (PcdMemoryProfileMemoryType) & TestBit) != 0) {\r
     return TRUE;\r
   } else {\r
@@ -741,8 +1018,9 @@ SmmCoreNeedRecordProfile (
 \r
 /**\r
   Convert EFI memory type to profile memory index. The rule is:\r
-  If BIOS memory type (0 ~ EfiMaxMemoryType), ProfileMemoryIndex = MemoryType.\r
-  If OS memory type (0x80000000 ~ 0xFFFFFFFF), ProfileMemoryIndex = EfiMaxMemoryType.\r
+  If BIOS memory type (0 ~ EfiMaxMemoryType - 1), ProfileMemoryIndex = MemoryType.\r
+  As SMRAM profile is only to record EfiRuntimeServicesCode and EfiRuntimeServicesData,\r
+  so return input memory type directly.\r
 \r
   @param MemoryType     Memory type.\r
 \r
@@ -754,11 +1032,7 @@ GetProfileMemoryIndex (
   IN EFI_MEMORY_TYPE    MemoryType\r
   )\r
 {\r
-  if ((UINT32) MemoryType >= 0x80000000) {\r
-    return EfiMaxMemoryType;\r
-  } else {\r
-    return MemoryType;\r
-  }\r
+  return MemoryType;\r
 }\r
 \r
 /**\r
@@ -801,54 +1075,80 @@ SmramProfileUpdateFreePages (
   @param MemoryType     Memory type.\r
   @param Size           Buffer size.\r
   @param Buffer         Buffer address.\r
+  @param ActionString   String for memory profile action.\r
 \r
-  @retval TRUE          Profile udpate success.\r
-  @retval FALSE         Profile update fail.\r
+  @return EFI_SUCCESS           Memory profile is updated.\r
+  @return EFI_UNSUPPORTED       Memory profile is unsupported,\r
+                                or memory profile for the image is not required.\r
+  @return EFI_OUT_OF_RESOURCES  No enough resource to update memory profile for allocate action.\r
 \r
 **/\r
-BOOLEAN\r
+EFI_STATUS\r
 SmmCoreUpdateProfileAllocate (\r
   IN PHYSICAL_ADDRESS       CallerAddress,\r
   IN MEMORY_PROFILE_ACTION  Action,\r
   IN EFI_MEMORY_TYPE        MemoryType,\r
   IN UINTN                  Size,\r
-  IN VOID                   *Buffer\r
+  IN VOID                   *Buffer,\r
+  IN CHAR8                  *ActionString OPTIONAL\r
   )\r
 {\r
   EFI_STATUS                        Status;\r
-  MEMORY_PROFILE_CONTEXT           *Context;\r
-  MEMORY_PROFILE_DRIVER_INFO       *DriverInfo;\r
-  MEMORY_PROFILE_ALLOC_INFO        *AllocInfo;\r
+  MEMORY_PROFILE_CONTEXT            *Context;\r
+  MEMORY_PROFILE_DRIVER_INFO        *DriverInfo;\r
+  MEMORY_PROFILE_ALLOC_INFO         *AllocInfo;\r
   MEMORY_PROFILE_CONTEXT_DATA       *ContextData;\r
   MEMORY_PROFILE_DRIVER_INFO_DATA   *DriverInfoData;\r
   MEMORY_PROFILE_ALLOC_INFO_DATA    *AllocInfoData;\r
   EFI_MEMORY_TYPE                   ProfileMemoryIndex;\r
+  MEMORY_PROFILE_ACTION             BasicAction;\r
+  UINTN                             ActionStringSize;\r
+  UINTN                             ActionStringOccupiedSize;\r
 \r
-  AllocInfoData = NULL;\r
+  BasicAction = Action & MEMORY_PROFILE_ACTION_BASIC_MASK;\r
 \r
   ContextData = GetSmramProfileContext ();\r
   if (ContextData == NULL) {\r
-    return FALSE;\r
+    return EFI_UNSUPPORTED;\r
   }\r
 \r
   DriverInfoData = GetMemoryProfileDriverInfoFromAddress (ContextData, CallerAddress);\r
-  ASSERT (DriverInfoData != NULL);\r
+  if (DriverInfoData == NULL) {\r
+    return EFI_UNSUPPORTED;\r
+  }\r
+\r
+  ActionStringSize = 0;\r
+  ActionStringOccupiedSize = 0;\r
+  if (ActionString != NULL) {\r
+    ActionStringSize = AsciiStrSize (ActionString);\r
+    ActionStringOccupiedSize = GET_OCCUPIED_SIZE (ActionStringSize, sizeof (UINT64));\r
+  }\r
 \r
   //\r
   // Use SmmInternalAllocatePool() that will not update profile for this AllocatePool action.\r
   //\r
+  AllocInfoData = NULL;\r
   Status = SmmInternalAllocatePool (\r
              EfiRuntimeServicesData,\r
-             sizeof (*AllocInfoData),\r
+             sizeof (*AllocInfoData) + ActionStringSize,\r
              (VOID **) &AllocInfoData\r
              );\r
   if (EFI_ERROR (Status)) {\r
-    return FALSE;\r
+    return EFI_OUT_OF_RESOURCES;\r
   }\r
+  ASSERT (AllocInfoData != NULL);\r
+\r
+  //\r
+  // Only update SequenceCount if and only if it is basic action.\r
+  //\r
+  if (Action == BasicAction) {\r
+    ContextData->Context.SequenceCount ++;\r
+  }\r
+\r
   AllocInfo = &AllocInfoData->AllocInfo;\r
   AllocInfoData->Signature      = MEMORY_PROFILE_ALLOC_INFO_SIGNATURE;\r
   AllocInfo->Header.Signature   = MEMORY_PROFILE_ALLOC_INFO_SIGNATURE;\r
-  AllocInfo->Header.Length      = sizeof (MEMORY_PROFILE_ALLOC_INFO);\r
+  AllocInfo->Header.Length      = (UINT16) (sizeof (MEMORY_PROFILE_ALLOC_INFO) + ActionStringOccupiedSize);\r
   AllocInfo->Header.Revision    = MEMORY_PROFILE_ALLOC_INFO_REVISION;\r
   AllocInfo->CallerAddress      = CallerAddress;\r
   AllocInfo->SequenceId         = ContextData->Context.SequenceCount;\r
@@ -856,42 +1156,56 @@ SmmCoreUpdateProfileAllocate (
   AllocInfo->MemoryType         = MemoryType;\r
   AllocInfo->Buffer             = (PHYSICAL_ADDRESS) (UINTN) Buffer;\r
   AllocInfo->Size               = Size;\r
+  if (ActionString != NULL) {\r
+    AllocInfo->ActionStringOffset = (UINT16) sizeof (MEMORY_PROFILE_ALLOC_INFO);\r
+    AllocInfoData->ActionString = (CHAR8 *) (AllocInfoData + 1);\r
+    CopyMem (AllocInfoData->ActionString, ActionString, ActionStringSize);\r
+  } else {\r
+    AllocInfo->ActionStringOffset = 0;\r
+    AllocInfoData->ActionString = NULL;\r
+  }\r
 \r
   InsertTailList (DriverInfoData->AllocInfoList, &AllocInfoData->Link);\r
 \r
-  ProfileMemoryIndex = GetProfileMemoryIndex (MemoryType);\r
-\r
+  Context = &ContextData->Context;\r
   DriverInfo = &DriverInfoData->DriverInfo;\r
-  DriverInfo->CurrentUsage += Size;\r
-  if (DriverInfo->PeakUsage < DriverInfo->CurrentUsage) {\r
-    DriverInfo->PeakUsage = DriverInfo->CurrentUsage;\r
-  }\r
-  DriverInfo->CurrentUsageByType[ProfileMemoryIndex] += Size;\r
-  if (DriverInfo->PeakUsageByType[ProfileMemoryIndex] < DriverInfo->CurrentUsageByType[ProfileMemoryIndex]) {\r
-    DriverInfo->PeakUsageByType[ProfileMemoryIndex] = DriverInfo->CurrentUsageByType[ProfileMemoryIndex];\r
-  }\r
   DriverInfo->AllocRecordCount ++;\r
 \r
-  Context = &ContextData->Context;\r
-  Context->CurrentTotalUsage += Size;\r
-  if (Context->PeakTotalUsage < Context->CurrentTotalUsage) {\r
-    Context->PeakTotalUsage = Context->CurrentTotalUsage;\r
-  }\r
-  Context->CurrentTotalUsageByType[ProfileMemoryIndex] += Size;\r
-  if (Context->PeakTotalUsageByType[ProfileMemoryIndex] < Context->CurrentTotalUsageByType[ProfileMemoryIndex]) {\r
-    Context->PeakTotalUsageByType[ProfileMemoryIndex] = Context->CurrentTotalUsageByType[ProfileMemoryIndex];\r
+  //\r
+  // Update summary if and only if it is basic action.\r
+  //\r
+  if (Action == BasicAction) {\r
+    ProfileMemoryIndex = GetProfileMemoryIndex (MemoryType);\r
+\r
+    DriverInfo->CurrentUsage += Size;\r
+    if (DriverInfo->PeakUsage < DriverInfo->CurrentUsage) {\r
+      DriverInfo->PeakUsage = DriverInfo->CurrentUsage;\r
+    }\r
+    DriverInfo->CurrentUsageByType[ProfileMemoryIndex] += Size;\r
+    if (DriverInfo->PeakUsageByType[ProfileMemoryIndex] < DriverInfo->CurrentUsageByType[ProfileMemoryIndex]) {\r
+      DriverInfo->PeakUsageByType[ProfileMemoryIndex] = DriverInfo->CurrentUsageByType[ProfileMemoryIndex];\r
+    }\r
+\r
+    Context->CurrentTotalUsage += Size;\r
+    if (Context->PeakTotalUsage < Context->CurrentTotalUsage) {\r
+      Context->PeakTotalUsage = Context->CurrentTotalUsage;\r
+    }\r
+    Context->CurrentTotalUsageByType[ProfileMemoryIndex] += Size;\r
+    if (Context->PeakTotalUsageByType[ProfileMemoryIndex] < Context->CurrentTotalUsageByType[ProfileMemoryIndex]) {\r
+      Context->PeakTotalUsageByType[ProfileMemoryIndex] = Context->CurrentTotalUsageByType[ProfileMemoryIndex];\r
+    }\r
+\r
+    SmramProfileUpdateFreePages (ContextData);\r
   }\r
-  Context->SequenceCount ++;\r
 \r
-  SmramProfileUpdateFreePages (ContextData);\r
-  return TRUE;\r
+  return EFI_SUCCESS;\r
 }\r
 \r
 /**\r
   Get memory profile alloc info from memory profile\r
 \r
   @param DriverInfoData     Driver info\r
-  @param Action             This Free action\r
+  @param BasicAction        This Free basic action\r
   @param Size               Buffer size\r
   @param Buffer             Buffer address\r
 \r
@@ -900,7 +1214,7 @@ SmmCoreUpdateProfileAllocate (
 MEMORY_PROFILE_ALLOC_INFO_DATA *\r
 GetMemoryProfileAllocInfoFromAddress (\r
   IN MEMORY_PROFILE_DRIVER_INFO_DATA    *DriverInfoData,\r
-  IN MEMORY_PROFILE_ACTION              Action,\r
+  IN MEMORY_PROFILE_ACTION              BasicAction,\r
   IN UINTN                              Size,\r
   IN VOID                               *Buffer\r
   )\r
@@ -922,10 +1236,10 @@ GetMemoryProfileAllocInfoFromAddress (
                       MEMORY_PROFILE_ALLOC_INFO_SIGNATURE\r
                       );\r
     AllocInfo = &AllocInfoData->AllocInfo;\r
-    if (AllocInfo->Action != Action) {\r
+    if ((AllocInfo->Action & MEMORY_PROFILE_ACTION_BASIC_MASK) != BasicAction) {\r
       continue;\r
     }\r
-    switch (Action) {\r
+    switch (BasicAction) {\r
       case MemoryProfileActionAllocatePages:\r
         if ((AllocInfo->Buffer <= (PHYSICAL_ADDRESS) (UINTN) Buffer) &&\r
             ((AllocInfo->Buffer + AllocInfo->Size) >= ((PHYSICAL_ADDRESS) (UINTN) Buffer + Size))) {\r
@@ -954,11 +1268,12 @@ GetMemoryProfileAllocInfoFromAddress (
   @param Size           Buffer size.\r
   @param Buffer         Buffer address.\r
 \r
-  @retval TRUE          Profile udpate success.\r
-  @retval FALSE         Profile update fail.\r
+  @return EFI_SUCCESS           Memory profile is updated.\r
+  @return EFI_UNSUPPORTED       Memory profile is unsupported.\r
+  @return EFI_NOT_FOUND         No matched allocate info found for free action.\r
 \r
 **/\r
-BOOLEAN\r
+EFI_STATUS\r
 SmmCoreUpdateProfileFree (\r
   IN PHYSICAL_ADDRESS       CallerAddress,\r
   IN MEMORY_PROFILE_ACTION  Action,\r
@@ -976,114 +1291,143 @@ SmmCoreUpdateProfileFree (
   MEMORY_PROFILE_DRIVER_INFO_DATA  *ThisDriverInfoData;\r
   MEMORY_PROFILE_ALLOC_INFO_DATA   *AllocInfoData;\r
   EFI_MEMORY_TYPE                  ProfileMemoryIndex;\r
+  MEMORY_PROFILE_ACTION            BasicAction;\r
+  BOOLEAN                          Found;\r
+\r
+  BasicAction = Action & MEMORY_PROFILE_ACTION_BASIC_MASK;\r
 \r
   ContextData = GetSmramProfileContext ();\r
   if (ContextData == NULL) {\r
-    return FALSE;\r
+    return EFI_UNSUPPORTED;\r
   }\r
 \r
   DriverInfoData = GetMemoryProfileDriverInfoFromAddress (ContextData, CallerAddress);\r
-  ASSERT (DriverInfoData != NULL);\r
 \r
-  switch (Action) {\r
-    case MemoryProfileActionFreePages:\r
-      AllocInfoData = GetMemoryProfileAllocInfoFromAddress (DriverInfoData, MemoryProfileActionAllocatePages, Size, Buffer);\r
-      break;\r
-    case MemoryProfileActionFreePool:\r
-      AllocInfoData = GetMemoryProfileAllocInfoFromAddress (DriverInfoData, MemoryProfileActionAllocatePool, 0, Buffer);\r
-      break;\r
-    default:\r
-      ASSERT (FALSE);\r
-      AllocInfoData = NULL;\r
-      break;\r
-  }\r
-  if (AllocInfoData == NULL) {\r
-    //\r
-    // Legal case, because driver A might free memory allocated by driver B, by some protocol.\r
-    //\r
-    DriverInfoList = ContextData->DriverInfoList;\r
-\r
-    for (DriverLink = DriverInfoList->ForwardLink;\r
-         DriverLink != DriverInfoList;\r
-         DriverLink = DriverLink->ForwardLink) {\r
-      ThisDriverInfoData = CR (\r
-                             DriverLink,\r
-                             MEMORY_PROFILE_DRIVER_INFO_DATA,\r
-                             Link,\r
-                             MEMORY_PROFILE_DRIVER_INFO_SIGNATURE\r
-                             );\r
-      switch (Action) {\r
+  //\r
+  // Do not return if DriverInfoData == NULL here,\r
+  // because driver A might free memory allocated by driver B.\r
+  //\r
+\r
+  //\r
+  // Need use do-while loop to find all possible record,\r
+  // because one address might be recorded multiple times.\r
+  //\r
+  Found = FALSE;\r
+  AllocInfoData = NULL;\r
+  do {\r
+    if (DriverInfoData != NULL) {\r
+      switch (BasicAction) {\r
         case MemoryProfileActionFreePages:\r
-          AllocInfoData = GetMemoryProfileAllocInfoFromAddress (ThisDriverInfoData, MemoryProfileActionAllocatePages, Size, Buffer);\r
+          AllocInfoData = GetMemoryProfileAllocInfoFromAddress (DriverInfoData, MemoryProfileActionAllocatePages, Size, Buffer);\r
           break;\r
         case MemoryProfileActionFreePool:\r
-          AllocInfoData = GetMemoryProfileAllocInfoFromAddress (ThisDriverInfoData, MemoryProfileActionAllocatePool, 0, Buffer);\r
+          AllocInfoData = GetMemoryProfileAllocInfoFromAddress (DriverInfoData, MemoryProfileActionAllocatePool, 0, Buffer);\r
           break;\r
         default:\r
           ASSERT (FALSE);\r
           AllocInfoData = NULL;\r
           break;\r
       }\r
-      if (AllocInfoData != NULL) {\r
-        DriverInfoData = ThisDriverInfoData;\r
-        break;\r
-      }\r
     }\r
-\r
     if (AllocInfoData == NULL) {\r
       //\r
-      // No matched allocate operation is found for this free operation.\r
-      // It is because the specified memory type allocate operation has been\r
-      // filtered by CoreNeedRecordProfile(), but free operations have no\r
-      // memory type information, they can not be filtered by CoreNeedRecordProfile().\r
-      // Then, they will be filtered here.\r
+      // Legal case, because driver A might free memory allocated by driver B, by some protocol.\r
       //\r
-      return FALSE;\r
+      DriverInfoList = ContextData->DriverInfoList;\r
+\r
+      for (DriverLink = DriverInfoList->ForwardLink;\r
+           DriverLink != DriverInfoList;\r
+           DriverLink = DriverLink->ForwardLink) {\r
+        ThisDriverInfoData = CR (\r
+                               DriverLink,\r
+                               MEMORY_PROFILE_DRIVER_INFO_DATA,\r
+                               Link,\r
+                               MEMORY_PROFILE_DRIVER_INFO_SIGNATURE\r
+                               );\r
+        switch (BasicAction) {\r
+          case MemoryProfileActionFreePages:\r
+            AllocInfoData = GetMemoryProfileAllocInfoFromAddress (ThisDriverInfoData, MemoryProfileActionAllocatePages, Size, Buffer);\r
+            break;\r
+          case MemoryProfileActionFreePool:\r
+            AllocInfoData = GetMemoryProfileAllocInfoFromAddress (ThisDriverInfoData, MemoryProfileActionAllocatePool, 0, Buffer);\r
+            break;\r
+          default:\r
+            ASSERT (FALSE);\r
+            AllocInfoData = NULL;\r
+            break;\r
+        }\r
+        if (AllocInfoData != NULL) {\r
+          DriverInfoData = ThisDriverInfoData;\r
+          break;\r
+        }\r
+      }\r
+\r
+      if (AllocInfoData == NULL) {\r
+        //\r
+        // If (!Found), no matched allocate info is found for this free action.\r
+        // It is because the specified memory type allocate actions have been filtered by\r
+        // CoreNeedRecordProfile(), but free actions have no memory type information,\r
+        // they can not be filtered by CoreNeedRecordProfile(). Then, they will be\r
+        // filtered here.\r
+        //\r
+        // If (Found), it is normal exit path.\r
+        return (Found ? EFI_SUCCESS : EFI_NOT_FOUND);\r
+      }\r
     }\r
-  }\r
 \r
-  Context = &ContextData->Context;\r
-  DriverInfo = &DriverInfoData->DriverInfo;\r
-  AllocInfo = &AllocInfoData->AllocInfo;\r
+    ASSERT (DriverInfoData != NULL);\r
+    ASSERT (AllocInfoData != NULL);\r
 \r
-  ProfileMemoryIndex = GetProfileMemoryIndex (AllocInfo->MemoryType);\r
+    Found = TRUE;\r
 \r
-  Context->CurrentTotalUsage -= AllocInfo->Size;\r
-  Context->CurrentTotalUsageByType[ProfileMemoryIndex] -= AllocInfo->Size;\r
+    Context = &ContextData->Context;\r
+    DriverInfo = &DriverInfoData->DriverInfo;\r
+    AllocInfo = &AllocInfoData->AllocInfo;\r
 \r
-  DriverInfo->CurrentUsage -= AllocInfo->Size;\r
-  DriverInfo->CurrentUsageByType[ProfileMemoryIndex] -= AllocInfo->Size;\r
-  DriverInfo->AllocRecordCount --;\r
+    DriverInfo->AllocRecordCount --;\r
+    //\r
+    // Update summary if and only if it is basic action.\r
+    //\r
+    if (AllocInfo->Action == (AllocInfo->Action & MEMORY_PROFILE_ACTION_BASIC_MASK)) {\r
+      ProfileMemoryIndex = GetProfileMemoryIndex (AllocInfo->MemoryType);\r
 \r
-  RemoveEntryList (&AllocInfoData->Link);\r
+      Context->CurrentTotalUsage -= AllocInfo->Size;\r
+      Context->CurrentTotalUsageByType[ProfileMemoryIndex] -= AllocInfo->Size;\r
 \r
-  if (Action == MemoryProfileActionFreePages) {\r
-    if (AllocInfo->Buffer != (PHYSICAL_ADDRESS) (UINTN) Buffer) {\r
-      SmmCoreUpdateProfileAllocate (\r
-        AllocInfo->CallerAddress,\r
-        MemoryProfileActionAllocatePages,\r
-        AllocInfo->MemoryType,\r
-        (UINTN) ((PHYSICAL_ADDRESS) (UINTN) Buffer - AllocInfo->Buffer),\r
-        (VOID *) (UINTN) AllocInfo->Buffer\r
-        );\r
-    }\r
-    if (AllocInfo->Buffer + AllocInfo->Size != ((PHYSICAL_ADDRESS) (UINTN) Buffer + Size)) {\r
-      SmmCoreUpdateProfileAllocate (\r
-        AllocInfo->CallerAddress,\r
-        MemoryProfileActionAllocatePages,\r
-        AllocInfo->MemoryType,\r
-        (UINTN) ((AllocInfo->Buffer + AllocInfo->Size) - ((PHYSICAL_ADDRESS) (UINTN) Buffer + Size)),\r
-        (VOID *) ((UINTN) Buffer + Size)\r
-        );\r
+      DriverInfo->CurrentUsage -= AllocInfo->Size;\r
+      DriverInfo->CurrentUsageByType[ProfileMemoryIndex] -= AllocInfo->Size;\r
     }\r
-  }\r
 \r
-  //\r
-  // Use SmmInternalFreePool() that will not update profile for this FreePool action.\r
-  //\r
-  SmmInternalFreePool (AllocInfoData);\r
+    RemoveEntryList (&AllocInfoData->Link);\r
+\r
+    if (BasicAction == MemoryProfileActionFreePages) {\r
+      if (AllocInfo->Buffer != (PHYSICAL_ADDRESS) (UINTN) Buffer) {\r
+        SmmCoreUpdateProfileAllocate (\r
+          AllocInfo->CallerAddress,\r
+          AllocInfo->Action,\r
+          AllocInfo->MemoryType,\r
+          (UINTN) ((PHYSICAL_ADDRESS) (UINTN) Buffer - AllocInfo->Buffer),\r
+          (VOID *) (UINTN) AllocInfo->Buffer,\r
+          AllocInfoData->ActionString\r
+          );\r
+      }\r
+      if (AllocInfo->Buffer + AllocInfo->Size != ((PHYSICAL_ADDRESS) (UINTN) Buffer + Size)) {\r
+        SmmCoreUpdateProfileAllocate (\r
+          AllocInfo->CallerAddress,\r
+          AllocInfo->Action,\r
+          AllocInfo->MemoryType,\r
+          (UINTN) ((AllocInfo->Buffer + AllocInfo->Size) - ((PHYSICAL_ADDRESS) (UINTN) Buffer + Size)),\r
+          (VOID *) ((UINTN) Buffer + Size),\r
+          AllocInfoData->ActionString\r
+          );\r
+      }\r
+    }\r
 \r
-  return TRUE;\r
+    //\r
+    // Use SmmInternalFreePool() that will not update profile for this FreePool action.\r
+    //\r
+    SmmInternalFreePool (AllocInfoData);\r
+  } while (TRUE);\r
 }\r
 \r
 /**\r
@@ -1092,68 +1436,91 @@ SmmCoreUpdateProfileFree (
   @param CallerAddress  Address of caller who call Allocate or Free.\r
   @param Action         This Allocate or Free action.\r
   @param MemoryType     Memory type.\r
+                        EfiMaxMemoryType means the MemoryType is unknown.\r
   @param Size           Buffer size.\r
   @param Buffer         Buffer address.\r
-\r
-  @retval TRUE          Profile udpate success.\r
-  @retval FALSE         Profile update fail.\r
+  @param ActionString   String for memory profile action.\r
+                        Only needed for user defined allocate action.\r
+\r
+  @return EFI_SUCCESS           Memory profile is updated.\r
+  @return EFI_UNSUPPORTED       Memory profile is unsupported,\r
+                                or memory profile for the image is not required,\r
+                                or memory profile for the memory type is not required.\r
+  @return EFI_ACCESS_DENIED     It is during memory profile data getting.\r
+  @return EFI_ABORTED           Memory profile recording is not enabled.\r
+  @return EFI_OUT_OF_RESOURCES  No enough resource to update memory profile for allocate action.\r
+  @return EFI_NOT_FOUND         No matched allocate info found for free action.\r
 \r
 **/\r
-BOOLEAN\r
+EFI_STATUS\r
+EFIAPI\r
 SmmCoreUpdateProfile (\r
   IN PHYSICAL_ADDRESS       CallerAddress,\r
   IN MEMORY_PROFILE_ACTION  Action,\r
   IN EFI_MEMORY_TYPE        MemoryType, // Valid for AllocatePages/AllocatePool\r
   IN UINTN                  Size,       // Valid for AllocatePages/FreePages/AllocatePool\r
-  IN VOID                   *Buffer\r
+  IN VOID                   *Buffer,\r
+  IN CHAR8                  *ActionString OPTIONAL\r
   )\r
 {\r
+  EFI_STATUS                    Status;\r
   MEMORY_PROFILE_CONTEXT_DATA   *ContextData;\r
+  MEMORY_PROFILE_ACTION         BasicAction;\r
 \r
   if (!IS_SMRAM_PROFILE_ENABLED) {\r
-    return FALSE;\r
+    return EFI_UNSUPPORTED;\r
   }\r
 \r
-  if (!mSmramProfileRecordingStatus) {\r
-    return FALSE;\r
+  if (mSmramProfileGettingStatus) {\r
+    return EFI_ACCESS_DENIED;\r
+  }\r
+\r
+  if (!mSmramProfileRecordingEnable) {\r
+    return EFI_ABORTED;\r
   }\r
 \r
+  //\r
+  // Get the basic action to know how to process the record\r
+  //\r
+  BasicAction = Action & MEMORY_PROFILE_ACTION_BASIC_MASK;\r
+\r
   //\r
   // Free operations have no memory type information, so skip the check.\r
   //\r
-  if ((Action == MemoryProfileActionAllocatePages) || (Action == MemoryProfileActionAllocatePool)) {\r
+  if ((BasicAction == MemoryProfileActionAllocatePages) || (BasicAction == MemoryProfileActionAllocatePool)) {\r
     //\r
     // Only record limited MemoryType.\r
     //\r
     if (!SmmCoreNeedRecordProfile (MemoryType)) {\r
-      return FALSE;\r
+      return EFI_UNSUPPORTED;\r
     }\r
   }\r
 \r
   ContextData = GetSmramProfileContext ();\r
   if (ContextData == NULL) {\r
-    return FALSE;\r
+    return EFI_UNSUPPORTED;\r
   }\r
 \r
-  switch (Action) {\r
+  switch (BasicAction) {\r
     case MemoryProfileActionAllocatePages:\r
-      SmmCoreUpdateProfileAllocate (CallerAddress, Action, MemoryType, Size, Buffer);\r
+      Status = SmmCoreUpdateProfileAllocate (CallerAddress, Action, MemoryType, Size, Buffer, ActionString);\r
       break;\r
     case MemoryProfileActionFreePages:\r
-      SmmCoreUpdateProfileFree (CallerAddress, Action, Size, Buffer);\r
+      Status = SmmCoreUpdateProfileFree (CallerAddress, Action, Size, Buffer);\r
       break;\r
     case MemoryProfileActionAllocatePool:\r
-      SmmCoreUpdateProfileAllocate (CallerAddress, Action, MemoryType, Size, Buffer);\r
+      Status = SmmCoreUpdateProfileAllocate (CallerAddress, Action, MemoryType, Size, Buffer, ActionString);\r
       break;\r
     case MemoryProfileActionFreePool:\r
-      SmmCoreUpdateProfileFree (CallerAddress, Action, 0, Buffer);\r
+      Status = SmmCoreUpdateProfileFree (CallerAddress, Action, 0, Buffer);\r
       break;\r
     default:\r
       ASSERT (FALSE);\r
+      Status = EFI_UNSUPPORTED;\r
       break;\r
   }\r
 \r
-  return TRUE;\r
+  return Status;\r
 }\r
 \r
 /**\r
@@ -1175,61 +1542,6 @@ SmramProfileReadyToLock (
 \r
 ////////////////////\r
 \r
-/**\r
-  This function check if the address is in SMRAM.\r
-\r
-  @param Buffer  the buffer address to be checked.\r
-  @param Length  the buffer length to be checked.\r
-\r
-  @retval TRUE  this address is in SMRAM.\r
-  @retval FALSE this address is NOT in SMRAM.\r
-\r
-**/\r
-BOOLEAN\r
-InternalIsAddressInSmram (\r
-  IN PHYSICAL_ADDRESS   Buffer,\r
-  IN UINT64             Length\r
-  )\r
-{\r
-  UINTN  Index;\r
-\r
-  for (Index = 0; Index < mFullSmramRangeCount; Index ++) {\r
-    if (((Buffer >= mFullSmramRanges[Index].CpuStart) && (Buffer < mFullSmramRanges[Index].CpuStart + mFullSmramRanges[Index].PhysicalSize)) ||\r
-        ((mFullSmramRanges[Index].CpuStart >= Buffer) && (mFullSmramRanges[Index].CpuStart < Buffer + Length))) {\r
-      return TRUE;\r
-    }\r
-  }\r
-\r
-  return FALSE;\r
-}\r
-\r
-/**\r
-  This function check if the address refered by Buffer and Length is valid.\r
-\r
-  @param Buffer  the buffer address to be checked.\r
-  @param Length  the buffer length to be checked.\r
-\r
-  @retval TRUE  this address is valid.\r
-  @retval FALSE this address is NOT valid.\r
-**/\r
-BOOLEAN\r
-InternalIsAddressValid (\r
-  IN UINTN                 Buffer,\r
-  IN UINTN                 Length\r
-  )\r
-{\r
-  if (Buffer > (MAX_ADDRESS - Length)) {\r
-    //\r
-    // Overflow happen\r
-    //\r
-    return FALSE;\r
-  }\r
-  if (InternalIsAddressInSmram ((PHYSICAL_ADDRESS) Buffer, (UINT64)Length)) {\r
-    return FALSE;\r
-  }\r
-  return TRUE;\r
-}\r
-\r
 /**\r
   Get SMRAM profile data size.\r
 \r
@@ -1241,17 +1553,21 @@ SmramProfileGetDataSize (
   VOID\r
   )\r
 {\r
-  MEMORY_PROFILE_CONTEXT_DATA      *ContextData;\r
-  MEMORY_PROFILE_DRIVER_INFO_DATA  *DriverInfoData;\r
-  LIST_ENTRY                      *DriverInfoList;\r
-  LIST_ENTRY                      *DriverLink;\r
-  UINTN                           TotalSize;\r
-  LIST_ENTRY                      *Node;\r
-  LIST_ENTRY                      *FreePageList;\r
-  LIST_ENTRY                      *FreePoolList;\r
-  FREE_POOL_HEADER                *Pool;\r
-  UINTN                           PoolListIndex;\r
-  UINTN                           Index;\r
+  MEMORY_PROFILE_CONTEXT_DATA       *ContextData;\r
+  MEMORY_PROFILE_DRIVER_INFO_DATA   *DriverInfoData;\r
+  MEMORY_PROFILE_ALLOC_INFO_DATA    *AllocInfoData;\r
+  LIST_ENTRY                        *DriverInfoList;\r
+  LIST_ENTRY                        *DriverLink;\r
+  LIST_ENTRY                        *AllocInfoList;\r
+  LIST_ENTRY                        *AllocLink;\r
+  UINTN                             TotalSize;\r
+  LIST_ENTRY                        *Node;\r
+  LIST_ENTRY                        *FreePageList;\r
+  LIST_ENTRY                        *FreePoolList;\r
+  FREE_POOL_HEADER                  *Pool;\r
+  UINTN                             PoolListIndex;\r
+  UINTN                             Index;\r
+  UINTN                             SmmPoolTypeIndex;\r
 \r
   ContextData = GetSmramProfileContext ();\r
   if (ContextData == NULL) {\r
@@ -1259,7 +1575,6 @@ SmramProfileGetDataSize (
   }\r
 \r
   TotalSize = sizeof (MEMORY_PROFILE_CONTEXT);\r
-  TotalSize += sizeof (MEMORY_PROFILE_DRIVER_INFO) * (UINTN) ContextData->Context.ImageCount;\r
 \r
   DriverInfoList = ContextData->DriverInfoList;\r
   for (DriverLink = DriverInfoList->ForwardLink;\r
@@ -1271,7 +1586,20 @@ SmramProfileGetDataSize (
                    Link,\r
                    MEMORY_PROFILE_DRIVER_INFO_SIGNATURE\r
                    );\r
-    TotalSize += sizeof (MEMORY_PROFILE_ALLOC_INFO) * (UINTN) DriverInfoData->DriverInfo.AllocRecordCount;\r
+    TotalSize += DriverInfoData->DriverInfo.Header.Length;\r
+\r
+    AllocInfoList = DriverInfoData->AllocInfoList;\r
+    for (AllocLink = AllocInfoList->ForwardLink;\r
+         AllocLink != AllocInfoList;\r
+         AllocLink = AllocLink->ForwardLink) {\r
+      AllocInfoData = CR (\r
+                        AllocLink,\r
+                        MEMORY_PROFILE_ALLOC_INFO_DATA,\r
+                        Link,\r
+                        MEMORY_PROFILE_ALLOC_INFO_SIGNATURE\r
+                        );\r
+      TotalSize += AllocInfoData->AllocInfo.Header.Length;\r
+    }\r
   }\r
 \r
 \r
@@ -1282,19 +1610,20 @@ SmramProfileGetDataSize (
        Node = Node->BackLink) {\r
     Index++;\r
   }\r
-  for (PoolListIndex = 0; PoolListIndex < MAX_POOL_INDEX; PoolListIndex++) {\r
-    FreePoolList = &mSmmPoolLists[PoolListIndex];\r
-    for (Node = FreePoolList->BackLink;\r
-         Node != FreePoolList;\r
-         Node = Node->BackLink) {\r
-      Pool = BASE_CR (Node, FREE_POOL_HEADER, Link);\r
-      if (Pool->Header.Available) {\r
-        Index++;\r
+  for (SmmPoolTypeIndex = 0; SmmPoolTypeIndex < SmmPoolTypeMax; SmmPoolTypeIndex++) {\r
+    for (PoolListIndex = 0; PoolListIndex < MAX_POOL_INDEX; PoolListIndex++) {\r
+      FreePoolList = &mSmmPoolLists[SmmPoolTypeIndex][PoolListIndex];\r
+      for (Node = FreePoolList->BackLink;\r
+           Node != FreePoolList;\r
+           Node = Node->BackLink) {\r
+        Pool = BASE_CR (Node, FREE_POOL_HEADER, Link);\r
+        if (Pool->Header.Available) {\r
+          Index++;\r
+        }\r
       }\r
     }\r
   }\r
 \r
-\r
   TotalSize += (sizeof (MEMORY_PROFILE_FREE_MEMORY) + Index * sizeof (MEMORY_PROFILE_DESCRIPTOR));\r
   TotalSize += (sizeof (MEMORY_PROFILE_MEMORY_RANGE) + mFullSmramRangeCount * sizeof (MEMORY_PROFILE_DESCRIPTOR));\r
 \r
@@ -1305,11 +1634,17 @@ SmramProfileGetDataSize (
   Copy SMRAM profile data.\r
 \r
   @param ProfileBuffer  The buffer to hold SMRAM profile data.\r
+  @param ProfileSize    On input, profile buffer size.\r
+                        On output, actual profile data size copied.\r
+  @param ProfileOffset  On input, profile buffer offset to copy.\r
+                        On output, next time profile buffer offset to copy.\r
 \r
 **/\r
 VOID\r
 SmramProfileCopyData (\r
-  IN VOID   *ProfileBuffer\r
+  OUT VOID      *ProfileBuffer,\r
+  IN OUT UINT64 *ProfileSize,\r
+  IN OUT UINT64 *ProfileOffset\r
   )\r
 {\r
   MEMORY_PROFILE_CONTEXT           *Context;\r
@@ -1332,15 +1667,31 @@ SmramProfileCopyData (
   MEMORY_PROFILE_FREE_MEMORY      *FreeMemory;\r
   MEMORY_PROFILE_MEMORY_RANGE     *MemoryRange;\r
   MEMORY_PROFILE_DESCRIPTOR       *MemoryProfileDescriptor;\r
+  UINT64                          Offset;\r
+  UINT64                          RemainingSize;\r
+  UINTN                           PdbSize;\r
+  UINTN                           ActionStringSize;\r
+  UINTN                           SmmPoolTypeIndex;\r
 \r
   ContextData = GetSmramProfileContext ();\r
   if (ContextData == NULL) {\r
     return ;\r
   }\r
 \r
-  Context = ProfileBuffer;\r
-  CopyMem (Context, &ContextData->Context, sizeof (MEMORY_PROFILE_CONTEXT));\r
-  DriverInfo = (MEMORY_PROFILE_DRIVER_INFO *) (Context + 1);\r
+  RemainingSize = *ProfileSize;\r
+  Offset = 0;\r
+\r
+  if (*ProfileOffset < sizeof (MEMORY_PROFILE_CONTEXT)) {\r
+    if (RemainingSize >= sizeof (MEMORY_PROFILE_CONTEXT)) {\r
+      Context = ProfileBuffer;\r
+      CopyMem (Context, &ContextData->Context, sizeof (MEMORY_PROFILE_CONTEXT));\r
+      RemainingSize -= sizeof (MEMORY_PROFILE_CONTEXT);\r
+      ProfileBuffer = (UINT8 *) ProfileBuffer + sizeof (MEMORY_PROFILE_CONTEXT);\r
+    } else {\r
+      goto Done;\r
+    }\r
+  }\r
+  Offset += sizeof (MEMORY_PROFILE_CONTEXT);\r
 \r
   DriverInfoList = ContextData->DriverInfoList;\r
   for (DriverLink = DriverInfoList->ForwardLink;\r
@@ -1352,8 +1703,21 @@ SmramProfileCopyData (
                        Link,\r
                        MEMORY_PROFILE_DRIVER_INFO_SIGNATURE\r
                        );\r
-    CopyMem (DriverInfo, &DriverInfoData->DriverInfo, sizeof (MEMORY_PROFILE_DRIVER_INFO));\r
-    AllocInfo = (MEMORY_PROFILE_ALLOC_INFO *) (DriverInfo + 1);\r
+    if (*ProfileOffset < (Offset + DriverInfoData->DriverInfo.Header.Length)) {\r
+      if (RemainingSize >= DriverInfoData->DriverInfo.Header.Length) {\r
+        DriverInfo = ProfileBuffer;\r
+        CopyMem (DriverInfo, &DriverInfoData->DriverInfo, sizeof (MEMORY_PROFILE_DRIVER_INFO));\r
+        if (DriverInfo->PdbStringOffset != 0) {\r
+          PdbSize = AsciiStrSize (DriverInfoData->PdbString);\r
+          CopyMem ((VOID *) ((UINTN) DriverInfo + DriverInfo->PdbStringOffset), DriverInfoData->PdbString, PdbSize);\r
+        }\r
+        RemainingSize -= DriverInfo->Header.Length;\r
+        ProfileBuffer = (UINT8 *) ProfileBuffer + DriverInfo->Header.Length;\r
+      } else {\r
+        goto Done;\r
+      }\r
+    }\r
+    Offset += DriverInfoData->DriverInfo.Header.Length;\r
 \r
     AllocInfoList = DriverInfoData->AllocInfoList;\r
     for (AllocLink = AllocInfoList->ForwardLink;\r
@@ -1365,64 +1729,387 @@ SmramProfileCopyData (
                         Link,\r
                         MEMORY_PROFILE_ALLOC_INFO_SIGNATURE\r
                         );\r
-      CopyMem (AllocInfo, &AllocInfoData->AllocInfo, sizeof (MEMORY_PROFILE_ALLOC_INFO));\r
-      AllocInfo += 1;\r
+      if (*ProfileOffset < (Offset + AllocInfoData->AllocInfo.Header.Length)) {\r
+        if (RemainingSize >= AllocInfoData->AllocInfo.Header.Length) {\r
+          AllocInfo = ProfileBuffer;\r
+          CopyMem (AllocInfo, &AllocInfoData->AllocInfo, sizeof (MEMORY_PROFILE_ALLOC_INFO));\r
+          if (AllocInfo->ActionStringOffset) {\r
+            ActionStringSize = AsciiStrSize (AllocInfoData->ActionString);\r
+            CopyMem ((VOID *) ((UINTN) AllocInfo + AllocInfo->ActionStringOffset), AllocInfoData->ActionString, ActionStringSize);\r
+          }\r
+          RemainingSize -= AllocInfo->Header.Length;\r
+          ProfileBuffer = (UINT8 *) ProfileBuffer + AllocInfo->Header.Length;\r
+        } else {\r
+          goto Done;\r
+        }\r
+      }\r
+      Offset += AllocInfoData->AllocInfo.Header.Length;\r
     }\r
-\r
-    DriverInfo = (MEMORY_PROFILE_DRIVER_INFO *) ((UINTN) (DriverInfo + 1) + sizeof (MEMORY_PROFILE_ALLOC_INFO) * (UINTN) DriverInfo->AllocRecordCount);\r
   }\r
 \r
 \r
-  FreeMemory = (MEMORY_PROFILE_FREE_MEMORY *) DriverInfo;\r
-  CopyMem (FreeMemory, &mSmramFreeMemory, sizeof (MEMORY_PROFILE_FREE_MEMORY));\r
-  MemoryProfileDescriptor = (MEMORY_PROFILE_DESCRIPTOR *) (FreeMemory + 1);\r
-  Index = 0;\r
+  if (*ProfileOffset < (Offset + sizeof (MEMORY_PROFILE_FREE_MEMORY))) {\r
+    if (RemainingSize >= sizeof (MEMORY_PROFILE_FREE_MEMORY)) {\r
+      FreeMemory = ProfileBuffer;\r
+      CopyMem (FreeMemory, &mSmramFreeMemory, sizeof (MEMORY_PROFILE_FREE_MEMORY));\r
+      Index = 0;\r
+      FreePageList = &mSmmMemoryMap;\r
+      for (Node = FreePageList->BackLink;\r
+           Node != FreePageList;\r
+           Node = Node->BackLink) {\r
+        Index++;\r
+      }\r
+      for (SmmPoolTypeIndex = 0; SmmPoolTypeIndex < SmmPoolTypeMax; SmmPoolTypeIndex++) {\r
+        for (PoolListIndex = 0; PoolListIndex < MAX_POOL_INDEX; PoolListIndex++) {\r
+          FreePoolList = &mSmmPoolLists[SmmPoolTypeIndex][MAX_POOL_INDEX - PoolListIndex - 1];\r
+          for (Node = FreePoolList->BackLink;\r
+               Node != FreePoolList;\r
+               Node = Node->BackLink) {\r
+            Pool = BASE_CR (Node, FREE_POOL_HEADER, Link);\r
+            if (Pool->Header.Available) {\r
+              Index++;\r
+            }\r
+          }\r
+        }\r
+      }\r
+      FreeMemory->FreeMemoryEntryCount = Index;\r
+\r
+      RemainingSize -= sizeof (MEMORY_PROFILE_FREE_MEMORY);\r
+      ProfileBuffer = (UINT8 *) ProfileBuffer + sizeof (MEMORY_PROFILE_FREE_MEMORY);\r
+    } else {\r
+      goto Done;\r
+    }\r
+  }\r
+  Offset += sizeof (MEMORY_PROFILE_FREE_MEMORY);\r
   FreePageList = &mSmmMemoryMap;\r
   for (Node = FreePageList->BackLink;\r
        Node != FreePageList;\r
        Node = Node->BackLink) {\r
-    Pages = BASE_CR (Node, FREE_PAGE_LIST, Link);\r
-    MemoryProfileDescriptor->Header.Signature = MEMORY_PROFILE_DESCRIPTOR_SIGNATURE;\r
-    MemoryProfileDescriptor->Header.Length = sizeof (MEMORY_PROFILE_DESCRIPTOR);\r
-    MemoryProfileDescriptor->Header.Revision = MEMORY_PROFILE_DESCRIPTOR_REVISION;\r
-    MemoryProfileDescriptor->Address = (PHYSICAL_ADDRESS) (UINTN) Pages;\r
-    MemoryProfileDescriptor->Size = EFI_PAGES_TO_SIZE (Pages->NumberOfPages);\r
-    MemoryProfileDescriptor++;\r
-    Index++;\r
-  }\r
-  for (PoolListIndex = 0; PoolListIndex < MAX_POOL_INDEX; PoolListIndex++) {\r
-    FreePoolList = &mSmmPoolLists[MAX_POOL_INDEX - PoolListIndex - 1];\r
-    for (Node = FreePoolList->BackLink;\r
-         Node != FreePoolList;\r
-         Node = Node->BackLink) {\r
-      Pool = BASE_CR (Node, FREE_POOL_HEADER, Link);\r
-      if (Pool->Header.Available) {\r
+    if (*ProfileOffset < (Offset + sizeof (MEMORY_PROFILE_DESCRIPTOR))) {\r
+      if (RemainingSize >= sizeof (MEMORY_PROFILE_DESCRIPTOR)) {\r
+        Pages = BASE_CR (Node, FREE_PAGE_LIST, Link);\r
+        MemoryProfileDescriptor = ProfileBuffer;\r
         MemoryProfileDescriptor->Header.Signature = MEMORY_PROFILE_DESCRIPTOR_SIGNATURE;\r
         MemoryProfileDescriptor->Header.Length = sizeof (MEMORY_PROFILE_DESCRIPTOR);\r
         MemoryProfileDescriptor->Header.Revision = MEMORY_PROFILE_DESCRIPTOR_REVISION;\r
-        MemoryProfileDescriptor->Address = (PHYSICAL_ADDRESS) (UINTN) Pool;\r
-        MemoryProfileDescriptor->Size = Pool->Header.Size;\r
-        MemoryProfileDescriptor++;\r
-        Index++;\r
+        MemoryProfileDescriptor->Address = (PHYSICAL_ADDRESS) (UINTN) Pages;\r
+        MemoryProfileDescriptor->Size = EFI_PAGES_TO_SIZE (Pages->NumberOfPages);\r
+\r
+        RemainingSize -= sizeof (MEMORY_PROFILE_DESCRIPTOR);\r
+        ProfileBuffer = (UINT8 *) ProfileBuffer + sizeof (MEMORY_PROFILE_DESCRIPTOR);\r
+      } else {\r
+        goto Done;\r
+      }\r
+    }\r
+    Offset += sizeof (MEMORY_PROFILE_DESCRIPTOR);\r
+  }\r
+  for (SmmPoolTypeIndex = 0; SmmPoolTypeIndex < SmmPoolTypeMax; SmmPoolTypeIndex++) {\r
+    for (PoolListIndex = 0; PoolListIndex < MAX_POOL_INDEX; PoolListIndex++) {\r
+      FreePoolList = &mSmmPoolLists[SmmPoolTypeIndex][MAX_POOL_INDEX - PoolListIndex - 1];\r
+      for (Node = FreePoolList->BackLink;\r
+           Node != FreePoolList;\r
+           Node = Node->BackLink) {\r
+        Pool = BASE_CR (Node, FREE_POOL_HEADER, Link);\r
+        if (Pool->Header.Available) {\r
+          if (*ProfileOffset < (Offset + sizeof (MEMORY_PROFILE_DESCRIPTOR))) {\r
+            if (RemainingSize >= sizeof (MEMORY_PROFILE_DESCRIPTOR)) {\r
+              MemoryProfileDescriptor = ProfileBuffer;\r
+              MemoryProfileDescriptor->Header.Signature = MEMORY_PROFILE_DESCRIPTOR_SIGNATURE;\r
+              MemoryProfileDescriptor->Header.Length = sizeof (MEMORY_PROFILE_DESCRIPTOR);\r
+              MemoryProfileDescriptor->Header.Revision = MEMORY_PROFILE_DESCRIPTOR_REVISION;\r
+              MemoryProfileDescriptor->Address = (PHYSICAL_ADDRESS) (UINTN) Pool;\r
+              MemoryProfileDescriptor->Size = Pool->Header.Size;\r
+\r
+              RemainingSize -= sizeof (MEMORY_PROFILE_DESCRIPTOR);\r
+              ProfileBuffer = (UINT8 *) ProfileBuffer + sizeof (MEMORY_PROFILE_DESCRIPTOR);\r
+            } else {\r
+              goto Done;\r
+            }\r
+          }\r
+          Offset += sizeof (MEMORY_PROFILE_DESCRIPTOR);\r
+        }\r
       }\r
     }\r
   }\r
-  FreeMemory->FreeMemoryEntryCount = Index;\r
 \r
-  MemoryRange = (MEMORY_PROFILE_MEMORY_RANGE *) MemoryProfileDescriptor;\r
-  MemoryRange->Header.Signature = MEMORY_PROFILE_MEMORY_RANGE_SIGNATURE;\r
-  MemoryRange->Header.Length = sizeof (MEMORY_PROFILE_MEMORY_RANGE);\r
-  MemoryRange->Header.Revision = MEMORY_PROFILE_MEMORY_RANGE_REVISION;\r
-  MemoryRange->MemoryRangeCount = (UINT32) mFullSmramRangeCount;\r
-  MemoryProfileDescriptor = (MEMORY_PROFILE_DESCRIPTOR *) (MemoryRange + 1);\r
+  if (*ProfileOffset < (Offset + sizeof (MEMORY_PROFILE_MEMORY_RANGE))) {\r
+    if (RemainingSize >= sizeof (MEMORY_PROFILE_MEMORY_RANGE)) {\r
+      MemoryRange = ProfileBuffer;\r
+      MemoryRange->Header.Signature = MEMORY_PROFILE_MEMORY_RANGE_SIGNATURE;\r
+      MemoryRange->Header.Length = sizeof (MEMORY_PROFILE_MEMORY_RANGE);\r
+      MemoryRange->Header.Revision = MEMORY_PROFILE_MEMORY_RANGE_REVISION;\r
+      MemoryRange->MemoryRangeCount = (UINT32) mFullSmramRangeCount;\r
+\r
+      RemainingSize -= sizeof (MEMORY_PROFILE_MEMORY_RANGE);\r
+      ProfileBuffer = (UINT8 *) ProfileBuffer + sizeof (MEMORY_PROFILE_MEMORY_RANGE);\r
+    } else {\r
+      goto Done;\r
+    }\r
+  }\r
+  Offset += sizeof (MEMORY_PROFILE_MEMORY_RANGE);\r
   for (Index = 0; Index < mFullSmramRangeCount; Index++) {\r
-    MemoryProfileDescriptor->Header.Signature = MEMORY_PROFILE_DESCRIPTOR_SIGNATURE;\r
-    MemoryProfileDescriptor->Header.Length = sizeof (MEMORY_PROFILE_DESCRIPTOR);\r
-    MemoryProfileDescriptor->Header.Revision = MEMORY_PROFILE_DESCRIPTOR_REVISION;\r
-    MemoryProfileDescriptor->Address = mFullSmramRanges[Index].PhysicalStart;\r
-    MemoryProfileDescriptor->Size = mFullSmramRanges[Index].PhysicalSize;\r
-    MemoryProfileDescriptor++; \r
+    if (*ProfileOffset < (Offset + sizeof (MEMORY_PROFILE_DESCRIPTOR))) {\r
+      if (RemainingSize >= sizeof (MEMORY_PROFILE_DESCRIPTOR)) {\r
+        MemoryProfileDescriptor = ProfileBuffer;\r
+        MemoryProfileDescriptor->Header.Signature = MEMORY_PROFILE_DESCRIPTOR_SIGNATURE;\r
+        MemoryProfileDescriptor->Header.Length = sizeof (MEMORY_PROFILE_DESCRIPTOR);\r
+        MemoryProfileDescriptor->Header.Revision = MEMORY_PROFILE_DESCRIPTOR_REVISION;\r
+        MemoryProfileDescriptor->Address = mFullSmramRanges[Index].PhysicalStart;\r
+        MemoryProfileDescriptor->Size = mFullSmramRanges[Index].PhysicalSize;\r
+\r
+        RemainingSize -= sizeof (MEMORY_PROFILE_DESCRIPTOR);\r
+        ProfileBuffer = (UINT8 *) ProfileBuffer + sizeof (MEMORY_PROFILE_DESCRIPTOR);\r
+      } else {\r
+        goto Done;\r
+      }\r
+    }\r
+    Offset += sizeof (MEMORY_PROFILE_DESCRIPTOR);\r
   }\r
+\r
+Done:\r
+  //\r
+  // On output, actual profile data size copied.\r
+  //\r
+  *ProfileSize -= RemainingSize;\r
+  //\r
+  // On output, next time profile buffer offset to copy.\r
+  //\r
+  *ProfileOffset = Offset;\r
+}\r
+\r
+/**\r
+  Get memory profile data.\r
+\r
+  @param[in]      This              The EDKII_SMM_MEMORY_PROFILE_PROTOCOL instance.\r
+  @param[in, out] ProfileSize       On entry, points to the size in bytes of the ProfileBuffer.\r
+                                    On return, points to the size of the data returned in ProfileBuffer.\r
+  @param[out]     ProfileBuffer     Profile buffer.\r
+\r
+  @return EFI_SUCCESS               Get the memory profile data successfully.\r
+  @return EFI_UNSUPPORTED           Memory profile is unsupported.\r
+  @return EFI_BUFFER_TO_SMALL       The ProfileSize is too small for the resulting data.\r
+                                    ProfileSize is updated with the size required.\r
+\r
+**/\r
+EFI_STATUS\r
+EFIAPI\r
+SmramProfileProtocolGetData (\r
+  IN     EDKII_SMM_MEMORY_PROFILE_PROTOCOL  *This,\r
+  IN OUT UINT64                             *ProfileSize,\r
+     OUT VOID                               *ProfileBuffer\r
+  )\r
+{\r
+  UINT64                                Size;\r
+  UINT64                                Offset;\r
+  MEMORY_PROFILE_CONTEXT_DATA           *ContextData;\r
+  BOOLEAN                               SmramProfileGettingStatus;\r
+\r
+  ContextData = GetSmramProfileContext ();\r
+  if (ContextData == NULL) {\r
+    return EFI_UNSUPPORTED;\r
+  }\r
+\r
+  SmramProfileGettingStatus = mSmramProfileGettingStatus;\r
+  mSmramProfileGettingStatus = TRUE;\r
+\r
+  Size = SmramProfileGetDataSize ();\r
+\r
+  if (*ProfileSize < Size) {\r
+    *ProfileSize = Size;\r
+    mSmramProfileGettingStatus = SmramProfileGettingStatus;\r
+    return EFI_BUFFER_TOO_SMALL;\r
+  }\r
+\r
+  Offset = 0;\r
+  SmramProfileCopyData (ProfileBuffer, &Size, &Offset);\r
+  *ProfileSize = Size;\r
+\r
+  mSmramProfileGettingStatus = SmramProfileGettingStatus;\r
+  return EFI_SUCCESS;\r
+}\r
+\r
+/**\r
+  Register image to memory profile.\r
+\r
+  @param[in] This               The EDKII_SMM_MEMORY_PROFILE_PROTOCOL instance.\r
+  @param[in] FilePath           File path of the image.\r
+  @param[in] ImageBase          Image base address.\r
+  @param[in] ImageSize          Image size.\r
+  @param[in] FileType           File type of the image.\r
+\r
+  @return EFI_SUCCESS           Register successfully.\r
+  @return EFI_UNSUPPORTED       Memory profile is unsupported,\r
+                                or memory profile for the image is not required.\r
+  @return EFI_OUT_OF_RESOURCES  No enough resource for this register.\r
+\r
+**/\r
+EFI_STATUS\r
+EFIAPI\r
+SmramProfileProtocolRegisterImage (\r
+  IN EDKII_SMM_MEMORY_PROFILE_PROTOCOL  *This,\r
+  IN EFI_DEVICE_PATH_PROTOCOL           *FilePath,\r
+  IN PHYSICAL_ADDRESS                   ImageBase,\r
+  IN UINT64                             ImageSize,\r
+  IN EFI_FV_FILETYPE                    FileType\r
+  )\r
+{\r
+  EFI_STATUS                        Status;\r
+  EFI_SMM_DRIVER_ENTRY              DriverEntry;\r
+  VOID                              *EntryPointInImage;\r
+  EFI_GUID                          *Name;\r
+\r
+  ZeroMem (&DriverEntry, sizeof (DriverEntry));\r
+  Name = GetFileNameFromFilePath (FilePath);\r
+  if (Name != NULL) {\r
+    CopyMem (&DriverEntry.FileName, Name, sizeof (EFI_GUID));\r
+  }\r
+  DriverEntry.ImageBuffer = ImageBase;\r
+  DriverEntry.NumberOfPage = EFI_SIZE_TO_PAGES ((UINTN) ImageSize);\r
+  Status = InternalPeCoffGetEntryPoint ((VOID *) (UINTN) DriverEntry.ImageBuffer, &EntryPointInImage);\r
+  ASSERT_EFI_ERROR (Status);\r
+  DriverEntry.ImageEntryPoint = (PHYSICAL_ADDRESS) (UINTN) EntryPointInImage;\r
+\r
+  return RegisterSmramProfileImage (&DriverEntry, FALSE);\r
+}\r
+\r
+/**\r
+  Unregister image from memory profile.\r
+\r
+  @param[in] This               The EDKII_SMM_MEMORY_PROFILE_PROTOCOL instance.\r
+  @param[in] FilePath           File path of the image.\r
+  @param[in] ImageBase          Image base address.\r
+  @param[in] ImageSize          Image size.\r
+\r
+  @return EFI_SUCCESS           Unregister successfully.\r
+  @return EFI_UNSUPPORTED       Memory profile is unsupported,\r
+                                or memory profile for the image is not required.\r
+  @return EFI_NOT_FOUND         The image is not found.\r
+\r
+**/\r
+EFI_STATUS\r
+EFIAPI\r
+SmramProfileProtocolUnregisterImage (\r
+  IN EDKII_SMM_MEMORY_PROFILE_PROTOCOL  *This,\r
+  IN EFI_DEVICE_PATH_PROTOCOL           *FilePath,\r
+  IN PHYSICAL_ADDRESS                   ImageBase,\r
+  IN UINT64                             ImageSize\r
+  )\r
+{\r
+  EFI_STATUS                        Status;\r
+  EFI_SMM_DRIVER_ENTRY              DriverEntry;\r
+  VOID                              *EntryPointInImage;\r
+  EFI_GUID                          *Name;\r
+\r
+  ZeroMem (&DriverEntry, sizeof (DriverEntry));\r
+  Name = GetFileNameFromFilePath (FilePath);\r
+  if (Name != NULL) {\r
+    CopyMem (&DriverEntry.FileName, Name, sizeof (EFI_GUID));\r
+  }\r
+  DriverEntry.ImageBuffer = ImageBase;\r
+  DriverEntry.NumberOfPage = EFI_SIZE_TO_PAGES ((UINTN) ImageSize);\r
+  Status = InternalPeCoffGetEntryPoint ((VOID *) (UINTN) DriverEntry.ImageBuffer, &EntryPointInImage);\r
+  ASSERT_EFI_ERROR (Status);\r
+  DriverEntry.ImageEntryPoint = (PHYSICAL_ADDRESS) (UINTN) EntryPointInImage;\r
+\r
+  return UnregisterSmramProfileImage (&DriverEntry, FALSE);\r
+}\r
+\r
+/**\r
+  Get memory profile recording state.\r
+\r
+  @param[in]  This              The EDKII_SMM_MEMORY_PROFILE_PROTOCOL instance.\r
+  @param[out] RecordingState    Recording state.\r
+\r
+  @return EFI_SUCCESS           Memory profile recording state is returned.\r
+  @return EFI_UNSUPPORTED       Memory profile is unsupported.\r
+  @return EFI_INVALID_PARAMETER RecordingState is NULL.\r
+\r
+**/\r
+EFI_STATUS\r
+EFIAPI\r
+SmramProfileProtocolGetRecordingState (\r
+  IN EDKII_SMM_MEMORY_PROFILE_PROTOCOL  *This,\r
+  OUT BOOLEAN                           *RecordingState\r
+  )\r
+{\r
+  MEMORY_PROFILE_CONTEXT_DATA           *ContextData;\r
+\r
+  ContextData = GetSmramProfileContext ();\r
+  if (ContextData == NULL) {\r
+    return EFI_UNSUPPORTED;\r
+  }\r
+\r
+  if (RecordingState == NULL) {\r
+    return EFI_INVALID_PARAMETER;\r
+  }\r
+  *RecordingState = mSmramProfileRecordingEnable;\r
+  return EFI_SUCCESS;\r
+}\r
+\r
+/**\r
+  Set memory profile recording state.\r
+\r
+  @param[in] This               The EDKII_SMM_MEMORY_PROFILE_PROTOCOL instance.\r
+  @param[in] RecordingState     Recording state.\r
+\r
+  @return EFI_SUCCESS           Set memory profile recording state successfully.\r
+  @return EFI_UNSUPPORTED       Memory profile is unsupported.\r
+\r
+**/\r
+EFI_STATUS\r
+EFIAPI\r
+SmramProfileProtocolSetRecordingState (\r
+  IN EDKII_SMM_MEMORY_PROFILE_PROTOCOL  *This,\r
+  IN BOOLEAN                            RecordingState\r
+  )\r
+{\r
+  MEMORY_PROFILE_CONTEXT_DATA           *ContextData;\r
+\r
+  ContextData = GetSmramProfileContext ();\r
+  if (ContextData == NULL) {\r
+    return EFI_UNSUPPORTED;\r
+  }\r
+\r
+  mSmramProfileRecordingEnable = RecordingState;\r
+  return EFI_SUCCESS;\r
+}\r
+\r
+/**\r
+  Record memory profile of multilevel caller.\r
+\r
+  @param[in] This               The EDKII_SMM_MEMORY_PROFILE_PROTOCOL instance.\r
+  @param[in] CallerAddress      Address of caller.\r
+  @param[in] Action             Memory profile action.\r
+  @param[in] MemoryType         Memory type.\r
+                                EfiMaxMemoryType means the MemoryType is unknown.\r
+  @param[in] Buffer             Buffer address.\r
+  @param[in] Size               Buffer size.\r
+  @param[in] ActionString       String for memory profile action.\r
+                                Only needed for user defined allocate action.\r
+\r
+  @return EFI_SUCCESS           Memory profile is updated.\r
+  @return EFI_UNSUPPORTED       Memory profile is unsupported,\r
+                                or memory profile for the image is not required,\r
+                                or memory profile for the memory type is not required.\r
+  @return EFI_ACCESS_DENIED     It is during memory profile data getting.\r
+  @return EFI_ABORTED           Memory profile recording is not enabled.\r
+  @return EFI_OUT_OF_RESOURCES  No enough resource to update memory profile for allocate action.\r
+  @return EFI_NOT_FOUND         No matched allocate info found for free action.\r
+\r
+**/\r
+EFI_STATUS\r
+EFIAPI\r
+SmramProfileProtocolRecord (\r
+  IN EDKII_SMM_MEMORY_PROFILE_PROTOCOL  *This,\r
+  IN PHYSICAL_ADDRESS                   CallerAddress,\r
+  IN MEMORY_PROFILE_ACTION              Action,\r
+  IN EFI_MEMORY_TYPE                    MemoryType,\r
+  IN VOID                               *Buffer,\r
+  IN UINTN                              Size,\r
+  IN CHAR8                              *ActionString OPTIONAL\r
+  )\r
+{\r
+  return SmmCoreUpdateProfile (CallerAddress, Action, MemoryType, Size, Buffer, ActionString);\r
 }\r
 \r
 /**\r
@@ -1437,20 +2124,20 @@ SmramProfileHandlerGetInfo (
   )\r
 {\r
   MEMORY_PROFILE_CONTEXT_DATA   *ContextData;\r
-  BOOLEAN                       SmramProfileRecordingStatus;\r
+  BOOLEAN                       SmramProfileGettingStatus;\r
 \r
   ContextData = GetSmramProfileContext ();\r
   if (ContextData == NULL) {\r
     return ;\r
   }\r
 \r
-  SmramProfileRecordingStatus = mSmramProfileRecordingStatus;\r
-  mSmramProfileRecordingStatus = FALSE;\r
+  SmramProfileGettingStatus = mSmramProfileGettingStatus;\r
+  mSmramProfileGettingStatus = TRUE;\r
 \r
   SmramProfileParameterGetInfo->ProfileSize = SmramProfileGetDataSize();\r
   SmramProfileParameterGetInfo->Header.ReturnStatus = 0;\r
 \r
-  mSmramProfileRecordingStatus = SmramProfileRecordingStatus;\r
+  mSmramProfileGettingStatus = SmramProfileGettingStatus;\r
 }\r
 \r
 /**\r
@@ -1465,17 +2152,18 @@ SmramProfileHandlerGetData (
   )\r
 {\r
   UINT64                                    ProfileSize;\r
+  UINT64                                    ProfileOffset;\r
   SMRAM_PROFILE_PARAMETER_GET_PROFILE_DATA  SmramProfileGetData;\r
   MEMORY_PROFILE_CONTEXT_DATA               *ContextData;\r
-  BOOLEAN                                   SmramProfileRecordingStatus;\r
+  BOOLEAN                                   SmramProfileGettingStatus;\r
 \r
   ContextData = GetSmramProfileContext ();\r
   if (ContextData == NULL) {\r
     return ;\r
   }\r
 \r
-  SmramProfileRecordingStatus = mSmramProfileRecordingStatus;\r
-  mSmramProfileRecordingStatus = FALSE;\r
+  SmramProfileGettingStatus = mSmramProfileGettingStatus;\r
+  mSmramProfileGettingStatus = TRUE;\r
 \r
 \r
   CopyMem (&SmramProfileGetData, SmramProfileParameterGetData, sizeof (SmramProfileGetData));\r
@@ -1485,7 +2173,7 @@ SmramProfileHandlerGetData (
   //\r
   // Sanity check\r
   //\r
-  if (!InternalIsAddressValid ((UINTN) SmramProfileGetData.ProfileBuffer, (UINTN) ProfileSize)) {\r
+  if (!SmmIsBufferOutsideSmmValid ((UINTN) SmramProfileGetData.ProfileBuffer, (UINTN) ProfileSize)) {\r
     DEBUG ((EFI_D_ERROR, "SmramProfileHandlerGetData: SMM ProfileBuffer in SMRAM or overflow!\n"));\r
     SmramProfileParameterGetData->ProfileSize = ProfileSize;\r
     SmramProfileParameterGetData->Header.ReturnStatus = (UINT64) (INT64) (INTN) EFI_ACCESS_DENIED;\r
@@ -1498,12 +2186,56 @@ SmramProfileHandlerGetData (
     goto Done;\r
   }\r
 \r
+  ProfileOffset = 0;\r
+  SmramProfileCopyData ((VOID *) (UINTN) SmramProfileGetData.ProfileBuffer, &ProfileSize, &ProfileOffset);\r
   SmramProfileParameterGetData->ProfileSize = ProfileSize;\r
-  SmramProfileCopyData ((VOID *) (UINTN) SmramProfileGetData.ProfileBuffer);\r
   SmramProfileParameterGetData->Header.ReturnStatus = 0;\r
 \r
 Done:\r
-  mSmramProfileRecordingStatus = SmramProfileRecordingStatus;\r
+  mSmramProfileGettingStatus = SmramProfileGettingStatus;\r
+}\r
+\r
+/**\r
+  SMRAM profile handler to get profile data by offset.\r
+\r
+  @param SmramProfileParameterGetDataByOffset   The parameter of SMM profile get data by offset.\r
+\r
+**/\r
+VOID\r
+SmramProfileHandlerGetDataByOffset (\r
+  IN SMRAM_PROFILE_PARAMETER_GET_PROFILE_DATA_BY_OFFSET     *SmramProfileParameterGetDataByOffset\r
+  )\r
+{\r
+  SMRAM_PROFILE_PARAMETER_GET_PROFILE_DATA_BY_OFFSET    SmramProfileGetDataByOffset;\r
+  MEMORY_PROFILE_CONTEXT_DATA                           *ContextData;\r
+  BOOLEAN                                               SmramProfileGettingStatus;\r
+\r
+  ContextData = GetSmramProfileContext ();\r
+  if (ContextData == NULL) {\r
+    return ;\r
+  }\r
+\r
+  SmramProfileGettingStatus = mSmramProfileGettingStatus;\r
+  mSmramProfileGettingStatus = TRUE;\r
+\r
+\r
+  CopyMem (&SmramProfileGetDataByOffset, SmramProfileParameterGetDataByOffset, sizeof (SmramProfileGetDataByOffset));\r
+\r
+  //\r
+  // Sanity check\r
+  //\r
+  if (!SmmIsBufferOutsideSmmValid ((UINTN) SmramProfileGetDataByOffset.ProfileBuffer, (UINTN) SmramProfileGetDataByOffset.ProfileSize)) {\r
+    DEBUG ((EFI_D_ERROR, "SmramProfileHandlerGetDataByOffset: SMM ProfileBuffer in SMRAM or overflow!\n"));\r
+    SmramProfileParameterGetDataByOffset->Header.ReturnStatus = (UINT64) (INT64) (INTN) EFI_ACCESS_DENIED;\r
+    goto Done;\r
+  }\r
+\r
+  SmramProfileCopyData ((VOID *) (UINTN) SmramProfileGetDataByOffset.ProfileBuffer, &SmramProfileGetDataByOffset.ProfileSize, &SmramProfileGetDataByOffset.ProfileOffset);\r
+  CopyMem (SmramProfileParameterGetDataByOffset, &SmramProfileGetDataByOffset, sizeof (SmramProfileGetDataByOffset));\r
+  SmramProfileParameterGetDataByOffset->Header.ReturnStatus = 0;\r
+\r
+Done:\r
+  mSmramProfileGettingStatus = SmramProfileGettingStatus;\r
 }\r
 \r
 /**\r
@@ -1520,7 +2252,6 @@ SmramProfileHandlerRegisterImage (
   EFI_STATUS                        Status;\r
   EFI_SMM_DRIVER_ENTRY              DriverEntry;\r
   VOID                              *EntryPointInImage;\r
-  BOOLEAN                           Ret;\r
 \r
   ZeroMem (&DriverEntry, sizeof (DriverEntry));\r
   CopyMem (&DriverEntry.FileName, &SmramProfileParameterRegisterImage->FileName, sizeof(EFI_GUID));\r
@@ -1530,8 +2261,8 @@ SmramProfileHandlerRegisterImage (
   ASSERT_EFI_ERROR (Status);\r
   DriverEntry.ImageEntryPoint = (PHYSICAL_ADDRESS) (UINTN) EntryPointInImage;\r
 \r
-  Ret = RegisterSmramProfileImage (&DriverEntry, FALSE);\r
-  if (Ret) {\r
+  Status = RegisterSmramProfileImage (&DriverEntry, FALSE);\r
+  if (!EFI_ERROR (Status)) {\r
     SmramProfileParameterRegisterImage->Header.ReturnStatus = 0;\r
   }\r
 }\r
@@ -1550,7 +2281,6 @@ SmramProfileHandlerUnregisterImage (
   EFI_STATUS                        Status;\r
   EFI_SMM_DRIVER_ENTRY              DriverEntry;\r
   VOID                              *EntryPointInImage;\r
-  BOOLEAN                           Ret;\r
 \r
   ZeroMem (&DriverEntry, sizeof (DriverEntry));\r
   CopyMem (&DriverEntry.FileName, &SmramProfileParameterUnregisterImage->FileName, sizeof (EFI_GUID));\r
@@ -1560,8 +2290,8 @@ SmramProfileHandlerUnregisterImage (
   ASSERT_EFI_ERROR (Status);\r
   DriverEntry.ImageEntryPoint = (PHYSICAL_ADDRESS) (UINTN) EntryPointInImage;\r
 \r
-  Ret = UnregisterSmramProfileImage (&DriverEntry, FALSE);\r
-  if (Ret) {\r
+  Status = UnregisterSmramProfileImage (&DriverEntry, FALSE);\r
+  if (!EFI_ERROR (Status)) {\r
     SmramProfileParameterUnregisterImage->Header.ReturnStatus = 0;\r
   }\r
 }\r
@@ -1593,6 +2323,7 @@ SmramProfileHandler (
 {\r
   SMRAM_PROFILE_PARAMETER_HEADER           *SmramProfileParameterHeader;\r
   UINTN                                    TempCommBufferSize;\r
+  SMRAM_PROFILE_PARAMETER_RECORDING_STATE  *ParameterRecordingState;\r
 \r
   DEBUG ((EFI_D_ERROR, "SmramProfileHandler Enter\n"));\r
 \r
@@ -1610,7 +2341,7 @@ SmramProfileHandler (
     return EFI_SUCCESS;\r
   }\r
 \r
-  if (mSmramReadyToLock && !InternalIsAddressValid ((UINTN)CommBuffer, TempCommBufferSize)) {\r
+  if (mSmramReadyToLock && !SmmIsBufferOutsideSmmValid ((UINTN)CommBuffer, TempCommBufferSize)) {\r
     DEBUG ((EFI_D_ERROR, "SmramProfileHandler: SMM communication buffer in SMRAM or overflow!\n"));\r
     return EFI_SUCCESS;\r
   }\r
@@ -1641,6 +2372,14 @@ SmramProfileHandler (
     }\r
     SmramProfileHandlerGetData ((SMRAM_PROFILE_PARAMETER_GET_PROFILE_DATA *) (UINTN) CommBuffer);\r
     break;\r
+  case SMRAM_PROFILE_COMMAND_GET_PROFILE_DATA_BY_OFFSET:\r
+    DEBUG ((EFI_D_ERROR, "SmramProfileHandlerGetDataByOffset\n"));\r
+    if (TempCommBufferSize != sizeof (SMRAM_PROFILE_PARAMETER_GET_PROFILE_DATA_BY_OFFSET)) {\r
+      DEBUG ((EFI_D_ERROR, "SmramProfileHandler: SMM communication buffer size invalid!\n"));\r
+      return EFI_SUCCESS;\r
+    }\r
+    SmramProfileHandlerGetDataByOffset ((SMRAM_PROFILE_PARAMETER_GET_PROFILE_DATA_BY_OFFSET *) (UINTN) CommBuffer);\r
+    break;\r
   case SMRAM_PROFILE_COMMAND_REGISTER_IMAGE:\r
     DEBUG ((EFI_D_ERROR, "SmramProfileHandlerRegisterImage\n"));\r
     if (TempCommBufferSize != sizeof (SMRAM_PROFILE_PARAMETER_REGISTER_IMAGE)) {\r
@@ -1663,6 +2402,27 @@ SmramProfileHandler (
     }\r
     SmramProfileHandlerUnregisterImage ((SMRAM_PROFILE_PARAMETER_UNREGISTER_IMAGE *) (UINTN) CommBuffer);\r
     break;\r
+  case SMRAM_PROFILE_COMMAND_GET_RECORDING_STATE:\r
+    DEBUG ((EFI_D_ERROR, "SmramProfileHandlerGetRecordingState\n"));\r
+    if (TempCommBufferSize != sizeof (SMRAM_PROFILE_PARAMETER_RECORDING_STATE)) {\r
+      DEBUG ((EFI_D_ERROR, "SmramProfileHandler: SMM communication buffer size invalid!\n"));\r
+      return EFI_SUCCESS;\r
+    }\r
+    ParameterRecordingState = (SMRAM_PROFILE_PARAMETER_RECORDING_STATE *) (UINTN) CommBuffer;\r
+    ParameterRecordingState->RecordingState = mSmramProfileRecordingEnable;\r
+    ParameterRecordingState->Header.ReturnStatus = 0;\r
+    break;\r
+  case SMRAM_PROFILE_COMMAND_SET_RECORDING_STATE:\r
+    DEBUG ((EFI_D_ERROR, "SmramProfileHandlerSetRecordingState\n"));\r
+    if (TempCommBufferSize != sizeof (SMRAM_PROFILE_PARAMETER_RECORDING_STATE)) {\r
+      DEBUG ((EFI_D_ERROR, "SmramProfileHandler: SMM communication buffer size invalid!\n"));\r
+      return EFI_SUCCESS;\r
+    }\r
+    ParameterRecordingState = (SMRAM_PROFILE_PARAMETER_RECORDING_STATE *) (UINTN) CommBuffer;\r
+    mSmramProfileRecordingEnable = ParameterRecordingState->RecordingState;\r
+    ParameterRecordingState->Header.ReturnStatus = 0;\r
+    break;\r
+\r
   default:\r
     break;\r
   }\r
@@ -1709,15 +2469,15 @@ DumpSmramRange (
 {\r
   UINTN                         Index;\r
   MEMORY_PROFILE_CONTEXT_DATA   *ContextData;\r
-  BOOLEAN                       SmramProfileRecordingStatus;\r
+  BOOLEAN                       SmramProfileGettingStatus;\r
 \r
   ContextData = GetSmramProfileContext ();\r
   if (ContextData == NULL) {\r
     return ;\r
   }\r
 \r
-  SmramProfileRecordingStatus = mSmramProfileRecordingStatus;\r
-  mSmramProfileRecordingStatus = FALSE;\r
+  SmramProfileGettingStatus = mSmramProfileGettingStatus;\r
+  mSmramProfileGettingStatus = TRUE;\r
 \r
   DEBUG ((EFI_D_INFO, "FullSmramRange address - 0x%08x\n", mFullSmramRanges));\r
 \r
@@ -1734,7 +2494,7 @@ DumpSmramRange (
 \r
   DEBUG ((EFI_D_INFO, "======= SmramProfile end =======\n"));\r
 \r
-  mSmramProfileRecordingStatus = SmramProfileRecordingStatus;\r
+  mSmramProfileGettingStatus = SmramProfileGettingStatus;\r
 }\r
 \r
 /**\r
@@ -1751,15 +2511,15 @@ DumpFreePagesList (
   FREE_PAGE_LIST                *Pages;\r
   UINTN                         Index;\r
   MEMORY_PROFILE_CONTEXT_DATA   *ContextData;\r
-  BOOLEAN                       SmramProfileRecordingStatus;\r
+  BOOLEAN                       SmramProfileGettingStatus;\r
 \r
   ContextData = GetSmramProfileContext ();\r
   if (ContextData == NULL) {\r
     return ;\r
   }\r
 \r
-  SmramProfileRecordingStatus = mSmramProfileRecordingStatus;\r
-  mSmramProfileRecordingStatus = FALSE;\r
+  SmramProfileGettingStatus = mSmramProfileGettingStatus;\r
+  mSmramProfileGettingStatus = TRUE;\r
 \r
   DEBUG ((EFI_D_INFO, "======= SmramProfile begin =======\n"));\r
 \r
@@ -1776,7 +2536,7 @@ DumpFreePagesList (
 \r
   DEBUG ((EFI_D_INFO, "======= SmramProfile end =======\n"));\r
 \r
-  mSmramProfileRecordingStatus = SmramProfileRecordingStatus;\r
+  mSmramProfileGettingStatus = SmramProfileGettingStatus;\r
 }\r
 \r
 /**\r
@@ -1788,69 +2548,147 @@ DumpFreePoolList (
   VOID\r
   )\r
 {\r
-  LIST_ENTRY                     *FreePoolList;\r
-  LIST_ENTRY                     *Node;\r
-  FREE_POOL_HEADER               *Pool;\r
-  UINTN                          Index;\r
-  UINTN                          PoolListIndex;\r
+  LIST_ENTRY                    *FreePoolList;\r
+  LIST_ENTRY                    *Node;\r
+  FREE_POOL_HEADER              *Pool;\r
+  UINTN                         Index;\r
+  UINTN                         PoolListIndex;\r
   MEMORY_PROFILE_CONTEXT_DATA   *ContextData;\r
-  BOOLEAN                       SmramProfileRecordingStatus;\r
+  BOOLEAN                       SmramProfileGettingStatus;\r
+  UINTN                         SmmPoolTypeIndex;\r
 \r
   ContextData = GetSmramProfileContext ();\r
   if (ContextData == NULL) {\r
     return ;\r
   }\r
 \r
-  SmramProfileRecordingStatus = mSmramProfileRecordingStatus;\r
-  mSmramProfileRecordingStatus = FALSE;\r
+  SmramProfileGettingStatus = mSmramProfileGettingStatus;\r
+  mSmramProfileGettingStatus = TRUE;\r
 \r
-  DEBUG ((EFI_D_INFO, "======= SmramProfile begin =======\n"));\r
+  DEBUG ((DEBUG_INFO, "======= SmramProfile begin =======\n"));\r
 \r
-  for (PoolListIndex = 0; PoolListIndex < MAX_POOL_INDEX; PoolListIndex++) {\r
-    DEBUG ((EFI_D_INFO, "FreePoolList (%d):\n", PoolListIndex));\r
-    FreePoolList = &mSmmPoolLists[PoolListIndex];\r
-    for (Node = FreePoolList->BackLink, Index = 0;\r
-         Node != FreePoolList;\r
-         Node = Node->BackLink, Index++) {\r
-      Pool = BASE_CR (Node, FREE_POOL_HEADER, Link);\r
-      DEBUG ((EFI_D_INFO, "  Index - 0x%x\n", Index));\r
-      DEBUG ((EFI_D_INFO, "    PhysicalStart - 0x%016lx\n", (PHYSICAL_ADDRESS) (UINTN) Pool));\r
-      DEBUG ((EFI_D_INFO, "    Size          - 0x%08x\n", Pool->Header.Size));\r
-      DEBUG ((EFI_D_INFO, "    Available     - 0x%02x\n", Pool->Header.Available));\r
+  for (SmmPoolTypeIndex = 0; SmmPoolTypeIndex < SmmPoolTypeMax; SmmPoolTypeIndex++) {\r
+    for (PoolListIndex = 0; PoolListIndex < MAX_POOL_INDEX; PoolListIndex++) {\r
+      DEBUG ((DEBUG_INFO, "FreePoolList(%d)(%d):\n", SmmPoolTypeIndex, PoolListIndex));\r
+      FreePoolList = &mSmmPoolLists[SmmPoolTypeIndex][PoolListIndex];\r
+      for (Node = FreePoolList->BackLink, Index = 0;\r
+           Node != FreePoolList;\r
+           Node = Node->BackLink, Index++) {\r
+        Pool = BASE_CR (Node, FREE_POOL_HEADER, Link);\r
+        DEBUG ((DEBUG_INFO, "  Index - 0x%x\n", Index));\r
+        DEBUG ((DEBUG_INFO, "    PhysicalStart - 0x%016lx\n", (PHYSICAL_ADDRESS) (UINTN) Pool));\r
+        DEBUG ((DEBUG_INFO, "    Size          - 0x%08x\n", Pool->Header.Size));\r
+        DEBUG ((DEBUG_INFO, "    Available     - 0x%02x\n", Pool->Header.Available));\r
+      }\r
     }\r
   }\r
 \r
-  DEBUG ((EFI_D_INFO, "======= SmramProfile end =======\n"));\r
+  DEBUG ((DEBUG_INFO, "======= SmramProfile end =======\n"));\r
 \r
-  mSmramProfileRecordingStatus = SmramProfileRecordingStatus;\r
+  mSmramProfileGettingStatus = SmramProfileGettingStatus;\r
 }\r
 \r
-GLOBAL_REMOVE_IF_UNREFERENCED CHAR16 *mActionString[] = {\r
-  L"Unknown",\r
-  L"AllocatePages",\r
-  L"FreePages",\r
-  L"AllocatePool",\r
-  L"FreePool",\r
+GLOBAL_REMOVE_IF_UNREFERENCED CHAR8 *mSmmActionString[] = {\r
+  "SmmUnknown",\r
+  "gSmst->SmmAllocatePages",\r
+  "gSmst->SmmFreePages",\r
+  "gSmst->SmmAllocatePool",\r
+  "gSmst->SmmFreePool",\r
+};\r
+\r
+typedef struct {\r
+  MEMORY_PROFILE_ACTION  Action;\r
+  CHAR8                 *String;\r
+} ACTION_STRING;\r
+\r
+GLOBAL_REMOVE_IF_UNREFERENCED ACTION_STRING mExtActionString[] = {\r
+  {MEMORY_PROFILE_ACTION_LIB_ALLOCATE_PAGES,                    "Lib:AllocatePages"},\r
+  {MEMORY_PROFILE_ACTION_LIB_ALLOCATE_RUNTIME_PAGES,            "Lib:AllocateRuntimePages"},\r
+  {MEMORY_PROFILE_ACTION_LIB_ALLOCATE_RESERVED_PAGES,           "Lib:AllocateReservedPages"},\r
+  {MEMORY_PROFILE_ACTION_LIB_FREE_PAGES,                        "Lib:FreePages"},\r
+  {MEMORY_PROFILE_ACTION_LIB_ALLOCATE_ALIGNED_PAGES,            "Lib:AllocateAlignedPages"},\r
+  {MEMORY_PROFILE_ACTION_LIB_ALLOCATE_ALIGNED_RUNTIME_PAGES,    "Lib:AllocateAlignedRuntimePages"},\r
+  {MEMORY_PROFILE_ACTION_LIB_ALLOCATE_ALIGNED_RESERVED_PAGES,   "Lib:AllocateAlignedReservedPages"},\r
+  {MEMORY_PROFILE_ACTION_LIB_FREE_ALIGNED_PAGES,                "Lib:FreeAlignedPages"},\r
+  {MEMORY_PROFILE_ACTION_LIB_ALLOCATE_POOL,                     "Lib:AllocatePool"},\r
+  {MEMORY_PROFILE_ACTION_LIB_ALLOCATE_RUNTIME_POOL,             "Lib:AllocateRuntimePool"},\r
+  {MEMORY_PROFILE_ACTION_LIB_ALLOCATE_RESERVED_POOL,            "Lib:AllocateReservedPool"},\r
+  {MEMORY_PROFILE_ACTION_LIB_FREE_POOL,                         "Lib:FreePool"},\r
+  {MEMORY_PROFILE_ACTION_LIB_ALLOCATE_ZERO_POOL,                "Lib:AllocateZeroPool"},\r
+  {MEMORY_PROFILE_ACTION_LIB_ALLOCATE_RUNTIME_ZERO_POOL,        "Lib:AllocateRuntimeZeroPool"},\r
+  {MEMORY_PROFILE_ACTION_LIB_ALLOCATE_RESERVED_ZERO_POOL,       "Lib:AllocateReservedZeroPool"},\r
+  {MEMORY_PROFILE_ACTION_LIB_ALLOCATE_COPY_POOL,                "Lib:AllocateCopyPool"},\r
+  {MEMORY_PROFILE_ACTION_LIB_ALLOCATE_RUNTIME_COPY_POOL,        "Lib:AllocateRuntimeCopyPool"},\r
+  {MEMORY_PROFILE_ACTION_LIB_ALLOCATE_RESERVED_COPY_POOL,       "Lib:AllocateReservedCopyPool"},\r
+  {MEMORY_PROFILE_ACTION_LIB_REALLOCATE_POOL,                   "Lib:ReallocatePool"},\r
+  {MEMORY_PROFILE_ACTION_LIB_REALLOCATE_RUNTIME_POOL,           "Lib:ReallocateRuntimePool"},\r
+  {MEMORY_PROFILE_ACTION_LIB_REALLOCATE_RESERVED_POOL,          "Lib:ReallocateReservedPool"},\r
 };\r
 \r
-GLOBAL_REMOVE_IF_UNREFERENCED CHAR16 *mMemoryTypeString[] = {\r
-  L"EfiReservedMemoryType",\r
-  L"EfiLoaderCode",\r
-  L"EfiLoaderData",\r
-  L"EfiBootServicesCode",\r
-  L"EfiBootServicesData",\r
-  L"EfiRuntimeServicesCode",\r
-  L"EfiRuntimeServicesData",\r
-  L"EfiConventionalMemory",\r
-  L"EfiUnusableMemory",\r
-  L"EfiACPIReclaimMemory",\r
-  L"EfiACPIMemoryNVS",\r
-  L"EfiMemoryMappedIO",\r
-  L"EfiMemoryMappedIOPortSpace",\r
-  L"EfiPalCode",\r
-  L"EfiOSReserved",\r
+typedef struct {\r
+  EFI_MEMORY_TYPE   MemoryType;\r
+  CHAR8             *MemoryTypeStr;\r
+} PROFILE_MEMORY_TYPE_STRING;\r
+\r
+GLOBAL_REMOVE_IF_UNREFERENCED PROFILE_MEMORY_TYPE_STRING mMemoryTypeString[] = {\r
+  {EfiRuntimeServicesCode, "EfiRuntimeServicesCode"},\r
+  {EfiRuntimeServicesData, "EfiRuntimeServicesData"}\r
 };\r
 \r
+/**\r
+  Memory type to string.\r
+\r
+  @param[in] MemoryType Memory type.\r
+\r
+  @return Pointer to string.\r
+\r
+**/\r
+CHAR8 *\r
+ProfileMemoryTypeToStr (\r
+  IN EFI_MEMORY_TYPE    MemoryType\r
+  )\r
+{\r
+  UINTN     Index;\r
+  for (Index = 0; Index < ARRAY_SIZE (mMemoryTypeString); Index++) {\r
+    if (mMemoryTypeString[Index].MemoryType == MemoryType) {\r
+      return mMemoryTypeString[Index].MemoryTypeStr;\r
+    }\r
+  }\r
+\r
+  return "UnexpectedMemoryType";\r
+}\r
+\r
+/**\r
+  Action to string.\r
+\r
+  @param[in] Action                     Profile action.\r
+\r
+  @return Pointer to string.\r
+\r
+**/\r
+CHAR8 *\r
+ProfileActionToStr (\r
+  IN MEMORY_PROFILE_ACTION  Action\r
+  )\r
+{\r
+  UINTN     Index;\r
+  UINTN     ActionStringCount;\r
+  CHAR8     **ActionString;\r
+\r
+  ActionString = mSmmActionString;\r
+  ActionStringCount = ARRAY_SIZE (mSmmActionString);\r
+\r
+  if ((UINTN) (UINT32) Action < ActionStringCount) {\r
+    return ActionString[Action];\r
+  }\r
+  for (Index = 0; Index < ARRAY_SIZE (mExtActionString); Index++) {\r
+    if (mExtActionString[Index].Action == Action) {\r
+      return mExtActionString[Index].String;\r
+    }\r
+  }\r
+\r
+  return ActionString[0];\r
+}\r
 \r
 /**\r
   Dump SMRAM profile.\r
@@ -1873,7 +2711,7 @@ DumpSmramProfile (
   LIST_ENTRY                        *AllocInfoList;\r
   UINTN                             AllocIndex;\r
   LIST_ENTRY                        *AllocLink;\r
-  BOOLEAN                           SmramProfileRecordingStatus;\r
+  BOOLEAN                           SmramProfileGettingStatus;\r
   UINTN                             TypeIndex;\r
 \r
   ContextData = GetSmramProfileContext ();\r
@@ -1881,8 +2719,8 @@ DumpSmramProfile (
     return ;\r
   }\r
 \r
-  SmramProfileRecordingStatus = mSmramProfileRecordingStatus;\r
-  mSmramProfileRecordingStatus = FALSE;\r
+  SmramProfileGettingStatus = mSmramProfileGettingStatus;\r
+  mSmramProfileGettingStatus = TRUE;\r
 \r
   Context = &ContextData->Context;\r
   DEBUG ((EFI_D_INFO, "======= SmramProfile begin =======\n"));\r
@@ -1890,11 +2728,11 @@ DumpSmramProfile (
 \r
   DEBUG ((EFI_D_INFO, "  CurrentTotalUsage     - 0x%016lx\n", Context->CurrentTotalUsage));\r
   DEBUG ((EFI_D_INFO, "  PeakTotalUsage        - 0x%016lx\n", Context->PeakTotalUsage));\r
-  for (TypeIndex = 0; TypeIndex <= EfiMaxMemoryType; TypeIndex++) {\r
+  for (TypeIndex = 0; TypeIndex < sizeof (Context->CurrentTotalUsageByType) / sizeof (Context->CurrentTotalUsageByType[0]); TypeIndex++) {\r
     if ((Context->CurrentTotalUsageByType[TypeIndex] != 0) ||\r
         (Context->PeakTotalUsageByType[TypeIndex] != 0)) {\r
-      DEBUG ((EFI_D_INFO, "  CurrentTotalUsage[0x%02x]  - 0x%016lx (%s)\n", TypeIndex, Context->CurrentTotalUsageByType[TypeIndex], mMemoryTypeString[TypeIndex]));\r
-      DEBUG ((EFI_D_INFO, "  PeakTotalUsage[0x%02x]     - 0x%016lx (%s)\n", TypeIndex, Context->PeakTotalUsageByType[TypeIndex], mMemoryTypeString[TypeIndex]));\r
+      DEBUG ((EFI_D_INFO, "  CurrentTotalUsage[0x%02x]  - 0x%016lx (%a)\n", TypeIndex, Context->CurrentTotalUsageByType[TypeIndex], ProfileMemoryTypeToStr (TypeIndex)));\r
+      DEBUG ((EFI_D_INFO, "  PeakTotalUsage[0x%02x]     - 0x%016lx (%a)\n", TypeIndex, Context->PeakTotalUsageByType[TypeIndex], ProfileMemoryTypeToStr (TypeIndex)));\r
     }\r
   }\r
   DEBUG ((EFI_D_INFO, "  TotalImageSize        - 0x%016lx\n", Context->TotalImageSize));\r
@@ -1921,11 +2759,11 @@ DumpSmramProfile (
     DEBUG ((EFI_D_INFO, "    FileType            - 0x%02x\n", DriverInfo->FileType));\r
     DEBUG ((EFI_D_INFO, "    CurrentUsage        - 0x%016lx\n", DriverInfo->CurrentUsage));\r
     DEBUG ((EFI_D_INFO, "    PeakUsage           - 0x%016lx\n", DriverInfo->PeakUsage));\r
-    for (TypeIndex = 0; TypeIndex <= EfiMaxMemoryType; TypeIndex++) {\r
+    for (TypeIndex = 0; TypeIndex < sizeof (DriverInfo->CurrentUsageByType) / sizeof (DriverInfo->CurrentUsageByType[0]); TypeIndex++) {\r
       if ((DriverInfo->CurrentUsageByType[TypeIndex] != 0) ||\r
           (DriverInfo->PeakUsageByType[TypeIndex] != 0)) {\r
-        DEBUG ((EFI_D_INFO, "    CurrentUsage[0x%02x]     - 0x%016lx (%s)\n", TypeIndex, DriverInfo->CurrentUsageByType[TypeIndex], mMemoryTypeString[TypeIndex]));\r
-        DEBUG ((EFI_D_INFO, "    PeakUsage[0x%02x]        - 0x%016lx (%s)\n", TypeIndex, DriverInfo->PeakUsageByType[TypeIndex], mMemoryTypeString[TypeIndex]));\r
+        DEBUG ((EFI_D_INFO, "    CurrentUsage[0x%02x]     - 0x%016lx (%a)\n", TypeIndex, DriverInfo->CurrentUsageByType[TypeIndex], ProfileMemoryTypeToStr (TypeIndex)));\r
+        DEBUG ((EFI_D_INFO, "    PeakUsage[0x%02x]        - 0x%016lx (%a)\n", TypeIndex, DriverInfo->PeakUsageByType[TypeIndex], ProfileMemoryTypeToStr (TypeIndex)));\r
       }\r
     }\r
     DEBUG ((EFI_D_INFO, "    AllocRecordCount    - 0x%08x\n", DriverInfo->AllocRecordCount));\r
@@ -1944,8 +2782,16 @@ DumpSmramProfile (
       DEBUG ((EFI_D_INFO, "    MEMORY_PROFILE_ALLOC_INFO (0x%x)\n", AllocIndex));\r
       DEBUG ((EFI_D_INFO, "      CallerAddress  - 0x%016lx (Offset: 0x%08x)\n", AllocInfo->CallerAddress, AllocInfo->CallerAddress - DriverInfo->ImageBase));\r
       DEBUG ((EFI_D_INFO, "      SequenceId     - 0x%08x\n", AllocInfo->SequenceId));\r
-      DEBUG ((EFI_D_INFO, "      Action         - 0x%08x (%s)\n", AllocInfo->Action, mActionString[(AllocInfo->Action < sizeof(mActionString)/sizeof(mActionString[0])) ? AllocInfo->Action : 0]));\r
-      DEBUG ((EFI_D_INFO, "      MemoryType     - 0x%08x\n", AllocInfo->MemoryType));\r
+      if ((AllocInfo->Action & MEMORY_PROFILE_ACTION_USER_DEFINED_MASK) != 0) {\r
+        if (AllocInfoData->ActionString != NULL) {\r
+          DEBUG ((EFI_D_INFO, "      Action         - 0x%08x (%a)\n", AllocInfo->Action, AllocInfoData->ActionString));\r
+        } else {\r
+          DEBUG ((EFI_D_INFO, "      Action         - 0x%08x (UserDefined-0x%08x)\n", AllocInfo->Action, AllocInfo->Action));\r
+        }\r
+      } else {\r
+        DEBUG ((EFI_D_INFO, "      Action         - 0x%08x (%a)\n", AllocInfo->Action, ProfileActionToStr (AllocInfo->Action)));\r
+      }\r
+      DEBUG ((EFI_D_INFO, "      MemoryType     - 0x%08x (%a)\n", AllocInfo->MemoryType, ProfileMemoryTypeToStr (AllocInfo->MemoryType)));\r
       DEBUG ((EFI_D_INFO, "      Buffer         - 0x%016lx\n", AllocInfo->Buffer));\r
       DEBUG ((EFI_D_INFO, "      Size           - 0x%016lx\n", AllocInfo->Size));\r
     }\r
@@ -1953,7 +2799,7 @@ DumpSmramProfile (
 \r
   DEBUG ((EFI_D_INFO, "======= SmramProfile end =======\n"));\r
 \r
-  mSmramProfileRecordingStatus = SmramProfileRecordingStatus;\r
+  mSmramProfileGettingStatus = SmramProfileGettingStatus;\r
 }\r
 \r
 /**\r