]> git.proxmox.com Git - mirror_edk2.git/blobdiff - MdeModulePkg/Core/PiSmmCore/SmramProfileRecord.c
MdeModulePkg PiSmmCore: Retrieve Smram base address from system table
[mirror_edk2.git] / MdeModulePkg / Core / PiSmmCore / SmramProfileRecord.c
index 57b34a018dc40e1a82dbbb6b14cf560e6887f996..dda9f12a7dcd612644f5ee3ab5338114b6fcbfb3 100644 (file)
 #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
@@ -72,7 +78,10 @@ GLOBAL_REMOVE_IF_UNREFERENCED MEMORY_PROFILE_CONTEXT_DATA mSmramProfileContext =
 GLOBAL_REMOVE_IF_UNREFERENCED MEMORY_PROFILE_CONTEXT_DATA *mSmramProfileContextPtr = NULL;\r
 \r
 BOOLEAN mSmramReadyToLock;\r
-BOOLEAN mSmramProfileRecordingStatus = FALSE;\r
+BOOLEAN mSmramProfileGettingStatus = FALSE;\r
+BOOLEAN mSmramProfileRecordingEnable = MEMORY_PROFILE_RECORDING_DISABLE;\r
+EFI_DEVICE_PATH_PROTOCOL *mSmramProfileDriverPath;\r
+UINTN                    mSmramProfileDriverPathSize;\r
 \r
 /**\r
   Dump SMRAM infromation.\r
@@ -83,6 +92,155 @@ DumpSmramInfo (
   VOID\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
+/**\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
+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
@@ -239,7 +397,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
@@ -260,25 +417,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
@@ -302,6 +474,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
@@ -332,7 +512,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
@@ -371,7 +551,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
@@ -389,6 +569,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
@@ -405,15 +644,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
@@ -443,6 +683,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
@@ -452,7 +699,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
@@ -460,17 +714,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
@@ -478,10 +791,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
@@ -492,9 +803,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
@@ -507,10 +830,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
@@ -557,42 +880,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
@@ -632,10 +919,7 @@ GetMemoryProfileDriverInfoFromAddress (
     }\r
   }\r
 \r
-  //\r
-  // Should never come here.\r
-  //\r
-  return FindDummyImage (ContextData);\r
+  return NULL;\r
 }\r
 \r
 /**\r
@@ -644,11 +928,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
@@ -660,10 +946,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
@@ -673,9 +957,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
@@ -697,12 +993,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
@@ -714,7 +1011,7 @@ UnregisterSmramProfileImage (
     SmmInternalFreePool (DriverInfoData);\r
   }\r
 \r
-  return TRUE;\r
+  return EFI_SUCCESS;\r
 }\r
 \r
 /**\r
@@ -807,54 +1104,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
@@ -862,42 +1185,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
@@ -906,7 +1243,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
@@ -928,10 +1265,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
@@ -960,11 +1297,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
@@ -982,114 +1320,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
+      DriverInfo->CurrentUsage -= AllocInfo->Size;\r
+      DriverInfo->CurrentUsageByType[ProfileMemoryIndex] -= AllocInfo->Size;\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
-    }\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
@@ -1098,68 +1465,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
@@ -1192,17 +1582,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
@@ -1210,7 +1604,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
@@ -1222,7 +1615,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
@@ -1233,19 +1639,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
@@ -1291,6 +1698,9 @@ SmramProfileCopyData (
   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
@@ -1322,17 +1732,21 @@ SmramProfileCopyData (
                        Link,\r
                        MEMORY_PROFILE_DRIVER_INFO_SIGNATURE\r
                        );\r
-    if (*ProfileOffset < (Offset + sizeof (MEMORY_PROFILE_DRIVER_INFO))) {\r
-      if (RemainingSize >= sizeof (MEMORY_PROFILE_DRIVER_INFO)) {\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
-        RemainingSize -= sizeof (MEMORY_PROFILE_DRIVER_INFO);\r
-        ProfileBuffer = (UINT8 *) ProfileBuffer + 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 += sizeof (MEMORY_PROFILE_DRIVER_INFO);\r
+    Offset += DriverInfoData->DriverInfo.Header.Length;\r
 \r
     AllocInfoList = DriverInfoData->AllocInfoList;\r
     for (AllocLink = AllocInfoList->ForwardLink;\r
@@ -1344,17 +1758,21 @@ SmramProfileCopyData (
                         Link,\r
                         MEMORY_PROFILE_ALLOC_INFO_SIGNATURE\r
                         );\r
-      if (*ProfileOffset < (Offset + sizeof (MEMORY_PROFILE_ALLOC_INFO))) {\r
-        if (RemainingSize >= sizeof (MEMORY_PROFILE_ALLOC_INFO)) {\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
-          RemainingSize -= sizeof (MEMORY_PROFILE_ALLOC_INFO);\r
-          ProfileBuffer = (UINT8 *) ProfileBuffer + 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 += sizeof (MEMORY_PROFILE_ALLOC_INFO);\r
+      Offset += AllocInfoData->AllocInfo.Header.Length;\r
     }\r
   }\r
 \r
@@ -1370,14 +1788,16 @@ SmramProfileCopyData (
            Node = Node->BackLink) {\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
-            Index++;\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
@@ -1412,29 +1832,31 @@ SmramProfileCopyData (
     }\r
     Offset += sizeof (MEMORY_PROFILE_DESCRIPTOR);\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
-            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
+  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
-        Offset += sizeof (MEMORY_PROFILE_DESCRIPTOR);\r
       }\r
     }\r
   }\r
@@ -1484,6 +1906,241 @@ Done:
   *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
   SMRAM profile handler to get profile info.\r
 \r
@@ -1496,20 +2153,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
@@ -1527,15 +2184,15 @@ SmramProfileHandlerGetData (
   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
@@ -1564,7 +2221,7 @@ SmramProfileHandlerGetData (
   SmramProfileParameterGetData->Header.ReturnStatus = 0;\r
 \r
 Done:\r
-  mSmramProfileRecordingStatus = SmramProfileRecordingStatus;\r
+  mSmramProfileGettingStatus = SmramProfileGettingStatus;\r
 }\r
 \r
 /**\r
@@ -1580,15 +2237,15 @@ SmramProfileHandlerGetDataByOffset (
 {\r
   SMRAM_PROFILE_PARAMETER_GET_PROFILE_DATA_BY_OFFSET    SmramProfileGetDataByOffset;\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 (&SmramProfileGetDataByOffset, SmramProfileParameterGetDataByOffset, sizeof (SmramProfileGetDataByOffset));\r
@@ -1607,7 +2264,7 @@ SmramProfileHandlerGetDataByOffset (
   SmramProfileParameterGetDataByOffset->Header.ReturnStatus = 0;\r
 \r
 Done:\r
-  mSmramProfileRecordingStatus = SmramProfileRecordingStatus;\r
+  mSmramProfileGettingStatus = SmramProfileGettingStatus;\r
 }\r
 \r
 /**\r
@@ -1624,7 +2281,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
@@ -1634,8 +2290,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
@@ -1654,7 +2310,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
@@ -1664,8 +2319,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
@@ -1697,6 +2352,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
@@ -1775,6 +2431,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
@@ -1821,15 +2498,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
@@ -1846,7 +2523,7 @@ DumpSmramRange (
 \r
   DEBUG ((EFI_D_INFO, "======= SmramProfile end =======\n"));\r
 \r
-  mSmramProfileRecordingStatus = SmramProfileRecordingStatus;\r
+  mSmramProfileGettingStatus = SmramProfileGettingStatus;\r
 }\r
 \r
 /**\r
@@ -1863,15 +2540,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
@@ -1888,7 +2565,7 @@ DumpFreePagesList (
 \r
   DEBUG ((EFI_D_INFO, "======= SmramProfile end =======\n"));\r
 \r
-  mSmramProfileRecordingStatus = SmramProfileRecordingStatus;\r
+  mSmramProfileGettingStatus = SmramProfileGettingStatus;\r
 }\r
 \r
 /**\r
@@ -1900,59 +2577,93 @@ 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
+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 CHAR8 mUserDefinedActionString[] = {"UserDefined-0x80000000"};\r
+\r
 typedef struct {\r
   EFI_MEMORY_TYPE   MemoryType;\r
-  CHAR16            *MemoryTypeStr;\r
+  CHAR            *MemoryTypeStr;\r
 } PROFILE_MEMORY_TYPE_STRING;\r
 \r
 GLOBAL_REMOVE_IF_UNREFERENCED PROFILE_MEMORY_TYPE_STRING mMemoryTypeString[] = {\r
-  {EfiRuntimeServicesCode, L"EfiRuntimeServicesCode"},\r
-  {EfiRuntimeServicesData, L"EfiRuntimeServicesData"}\r
+  {EfiRuntimeServicesCode, "EfiRuntimeServicesCode"},\r
+  {EfiRuntimeServicesData, "EfiRuntimeServicesData"}\r
 };\r
 \r
 /**\r
@@ -1963,19 +2674,51 @@ GLOBAL_REMOVE_IF_UNREFERENCED PROFILE_MEMORY_TYPE_STRING mMemoryTypeString[] = {
   @return Pointer to string.\r
 \r
 **/\r
-CHAR16 *\r
+CHAR8 *\r
 ProfileMemoryTypeToStr (\r
   IN EFI_MEMORY_TYPE    MemoryType\r
   )\r
 {\r
   UINTN     Index;\r
-  for (Index = 0; Index < sizeof (mMemoryTypeString) / sizeof (mMemoryTypeString[0]); 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 L"UnexpectedMemoryType";\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
@@ -1999,7 +2742,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
@@ -2007,8 +2750,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
@@ -2019,8 +2762,8 @@ DumpSmramProfile (
   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], ProfileMemoryTypeToStr (TypeIndex)));\r
-      DEBUG ((EFI_D_INFO, "  PeakTotalUsage[0x%02x]     - 0x%016lx (%s)\n", TypeIndex, Context->PeakTotalUsageByType[TypeIndex], ProfileMemoryTypeToStr (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
@@ -2050,8 +2793,8 @@ DumpSmramProfile (
     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], ProfileMemoryTypeToStr (TypeIndex)));\r
-        DEBUG ((EFI_D_INFO, "    PeakUsage[0x%02x]        - 0x%016lx (%s)\n", TypeIndex, DriverInfo->PeakUsageByType[TypeIndex], ProfileMemoryTypeToStr (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
@@ -2070,8 +2813,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
@@ -2079,7 +2830,7 @@ DumpSmramProfile (
 \r
   DEBUG ((EFI_D_INFO, "======= SmramProfile end =======\n"));\r
 \r
-  mSmramProfileRecordingStatus = SmramProfileRecordingStatus;\r
+  mSmramProfileGettingStatus = SmramProfileGettingStatus;\r
 }\r
 \r
 /**\r