]> git.proxmox.com Git - mirror_edk2.git/commitdiff
MdeModulePkg PiSmmCore: Enhance memory profile for memory leak detection
authorStar Zeng <star.zeng@intel.com>
Sat, 18 Jun 2016 02:14:40 +0000 (10:14 +0800)
committerStar Zeng <star.zeng@intel.com>
Fri, 1 Jul 2016 01:39:56 +0000 (09:39 +0800)
1. Produce SMM memory profile protocol.
2. Consume PcdMemoryProfilePropertyMask to support disable recording
at the start.
3. Consume PcdMemoryProfileDriverPath to control which drivers need
memory profile data.

Cc: Jiewen Yao <jiewen.yao@intel.com>
Contributed-under: TianoCore Contribution Agreement 1.0
Signed-off-by: Star Zeng <star.zeng@intel.com>
Reviewed-by: Jiewen Yao <jiewen.yao@intel.com>
MdeModulePkg/Core/PiSmmCore/Page.c
MdeModulePkg/Core/PiSmmCore/PiSmmCore.c
MdeModulePkg/Core/PiSmmCore/PiSmmCore.h
MdeModulePkg/Core/PiSmmCore/PiSmmCore.inf
MdeModulePkg/Core/PiSmmCore/Pool.c
MdeModulePkg/Core/PiSmmCore/SmramProfileRecord.c

index 9cc2a4cabc6efd265872d7f52614caf9d30ffc75..5c04e8c8bf4d46e07d51a81041da86b5504d1a7b 100644 (file)
@@ -1,7 +1,7 @@
 /** @file\r
   SMM Memory page management functions.\r
 \r
-  Copyright (c) 2009 - 2014, Intel Corporation. All rights reserved.<BR>\r
+  Copyright (c) 2009 - 2016, Intel Corporation. All rights reserved.<BR>\r
   This program and the accompanying materials are licensed and made available \r
   under the terms and conditions of the BSD License which accompanies this \r
   distribution.  The full text of the license may be found at        \r
@@ -226,7 +226,14 @@ SmmAllocatePages (
 \r
   Status = SmmInternalAllocatePages (Type, MemoryType, NumberOfPages, Memory);\r
   if (!EFI_ERROR (Status)) {\r
-    SmmCoreUpdateProfile ((EFI_PHYSICAL_ADDRESS) (UINTN) RETURN_ADDRESS (0), MemoryProfileActionAllocatePages, MemoryType, EFI_PAGES_TO_SIZE (NumberOfPages), (VOID *) (UINTN) *Memory);\r
+    SmmCoreUpdateProfile (\r
+      (EFI_PHYSICAL_ADDRESS) (UINTN) RETURN_ADDRESS (0),\r
+      MemoryProfileActionAllocatePages,\r
+      MemoryType,\r
+      EFI_PAGES_TO_SIZE (NumberOfPages),\r
+      (VOID *) (UINTN) *Memory,\r
+      NULL\r
+      );\r
   }\r
   return Status;\r
 }\r
@@ -344,7 +351,14 @@ SmmFreePages (
 \r
   Status = SmmInternalFreePages (Memory, NumberOfPages);\r
   if (!EFI_ERROR (Status)) {\r
-    SmmCoreUpdateProfile ((EFI_PHYSICAL_ADDRESS) (UINTN) RETURN_ADDRESS (0), MemoryProfileActionFreePages, 0, EFI_PAGES_TO_SIZE (NumberOfPages), (VOID *) (UINTN) Memory);\r
+    SmmCoreUpdateProfile (\r
+      (EFI_PHYSICAL_ADDRESS) (UINTN) RETURN_ADDRESS (0),\r
+      MemoryProfileActionFreePages,\r
+      EfiMaxMemoryType,\r
+      EFI_PAGES_TO_SIZE (NumberOfPages),\r
+      (VOID *) (UINTN) Memory,\r
+      NULL\r
+      );\r
   }\r
   return Status;\r
 }\r
index 7245f201fbdf78ca965208c2563f67918bc69509..551560b1fce1e92466e271065c83855f39d14565 100644 (file)
@@ -1,7 +1,7 @@
 /** @file\r
   SMM Core Main Entry Point\r
 \r
-  Copyright (c) 2009 - 2015, Intel Corporation. All rights reserved.<BR>\r
+  Copyright (c) 2009 - 2016, Intel Corporation. All rights reserved.<BR>\r
   This program and the accompanying materials are licensed and made available \r
   under the terms and conditions of the BSD License which accompanies this \r
   distribution.  The full text of the license may be found at        \r
@@ -632,6 +632,7 @@ SmmMain (
   }\r
 \r
   RegisterSmramProfileHandler ();\r
+  SmramProfileInstallProtocol ();\r
 \r
   SmmCoreInstallLoadedImage ();\r
 \r
index 0e9c92abef9a78111e12b92ca890a72aefcfb35b..000864d3aab028dba58bfe2a6ad3196631f45846 100644 (file)
@@ -2,7 +2,7 @@
   The internal header file includes the common header files, defines\r
   internal structure and functions used by SmmCore module.\r
 \r
-  Copyright (c) 2009 - 2015, Intel Corporation. All rights reserved.<BR>\r
+  Copyright (c) 2009 - 2016, Intel Corporation. All rights reserved.<BR>\r
   This program and the accompanying materials are licensed and made available \r
   under the terms and conditions of the BSD License which accompanies this \r
   distribution.  The full text of the license may be found at        \r
@@ -42,6 +42,7 @@
 #include <Library/BaseLib.h>\r
 #include <Library/BaseMemoryLib.h>\r
 #include <Library/PeCoffLib.h>\r
+#include <Library/PeCoffGetEntryPointLib.h>\r
 #include <Library/CacheMaintenanceLib.h>\r
 #include <Library/DebugLib.h>\r
 #include <Library/ReportStatusCodeLib.h>\r
@@ -884,17 +885,28 @@ SmramProfileInit (
   VOID\r
   );\r
 \r
+/**\r
+  Install SMRAM profile protocol.\r
+\r
+**/\r
+VOID\r
+SmramProfileInstallProtocol (\r
+  VOID\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 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
@@ -906,11 +918,13 @@ RegisterSmramProfileImage (
   @param DriverEntry        SMM image info.\r
   @param UnregisterToDxe    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 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                UnregisterToDxe\r
@@ -922,20 +936,31 @@ UnregisterSmramProfileImage (
   @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 EFI_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 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 CHAR8                  *ActionString OPTIONAL\r
   );\r
 \r
 /**\r
index 9c06b2ac2fd7e45540d5792cb08e636ab4fdef5a..f7e32c4e09a8443fb2103a318ddeeade37155ef6 100644 (file)
@@ -1,7 +1,7 @@
 ## @file\r
 # This module provide an SMM CIS compliant implementation of SMM Core.\r
 #\r
-# Copyright (c) 2009 - 2015, Intel Corporation. All rights reserved.<BR>\r
+# Copyright (c) 2009 - 2016, Intel Corporation. All rights reserved.<BR>\r
 #\r
 # This program and the accompanying materials\r
 # are licensed and made available under the terms and conditions of the BSD License\r
@@ -48,6 +48,7 @@
   BaseLib\r
   BaseMemoryLib\r
   PeCoffLib\r
+  PeCoffGetEntryPointLib\r
   CacheMaintenanceLib  \r
   DebugLib\r
   ReportStatusCodeLib\r
@@ -81,6 +82,7 @@
   gEfiMdeModulePkgTokenSpaceGuid.PcdLoadModuleAtFixAddressEnable        ## CONSUMES\r
   gEfiMdeModulePkgTokenSpaceGuid.PcdMemoryProfileMemoryType             ## CONSUMES\r
   gEfiMdeModulePkgTokenSpaceGuid.PcdMemoryProfilePropertyMask           ## CONSUMES\r
+  gEfiMdeModulePkgTokenSpaceGuid.PcdMemoryProfileDriverPath             ## CONSUMES\r
 \r
 [Guids]\r
   gAprioriGuid                                  ## SOMETIMES_CONSUMES   ## File\r
@@ -92,6 +94,8 @@
   ## SOMETIMES_CONSUMES   ## GUID # Locate protocol\r
   ## SOMETIMES_PRODUCES   ## GUID # SmiHandlerRegister\r
   gEdkiiMemoryProfileGuid\r
+  ## SOMETIMES_PRODUCES   ## GUID # Install protocol\r
+  gEdkiiSmmMemoryProfileGuid\r
   gZeroGuid                                     ## SOMETIMES_CONSUMES   ## GUID\r
 \r
 [UserExtensions.TianoCore."ExtraFiles"]\r
index 79847b4bb6d027b489361a7eff84ed2ad1c1defe..53ef3e07797387590aad169136f36b0704d66842 100644 (file)
@@ -1,7 +1,7 @@
 /** @file\r
   SMM Memory pool management functions.\r
 \r
-  Copyright (c) 2009 - 2015, Intel Corporation. All rights reserved.<BR>\r
+  Copyright (c) 2009 - 2016, Intel Corporation. All rights reserved.<BR>\r
   This program and the accompanying materials are licensed and made available \r
   under the terms and conditions of the BSD License which accompanies this \r
   distribution.  The full text of the license may be found at        \r
@@ -260,7 +260,14 @@ SmmAllocatePool (
 \r
   Status = SmmInternalAllocatePool (PoolType, Size, Buffer);\r
   if (!EFI_ERROR (Status)) {\r
-    SmmCoreUpdateProfile ((EFI_PHYSICAL_ADDRESS) (UINTN) RETURN_ADDRESS (0), MemoryProfileActionAllocatePool, PoolType, Size, *Buffer);\r
+    SmmCoreUpdateProfile (\r
+      (EFI_PHYSICAL_ADDRESS) (UINTN) RETURN_ADDRESS (0),\r
+      MemoryProfileActionAllocatePool,\r
+      PoolType,\r
+      Size,\r
+      *Buffer,\r
+      NULL\r
+      );\r
   }\r
   return Status;\r
 }\r
@@ -319,7 +326,14 @@ SmmFreePool (
 \r
   Status = SmmInternalFreePool (Buffer);\r
   if (!EFI_ERROR (Status)) {\r
-    SmmCoreUpdateProfile ((EFI_PHYSICAL_ADDRESS) (UINTN) RETURN_ADDRESS (0), MemoryProfileActionFreePool, 0, 0, Buffer);\r
+    SmmCoreUpdateProfile (\r
+      (EFI_PHYSICAL_ADDRESS) (UINTN) RETURN_ADDRESS (0),\r
+      MemoryProfileActionFreePool,\r
+      EfiMaxMemoryType,\r
+      0,\r
+      Buffer,\r
+      NULL\r
+      );\r
   }\r
   return Status;\r
 }\r
index 57b34a018dc40e1a82dbbb6b14cf560e6887f996..281e382d5c8102af3c0a512fc8d732c828a89717 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,13 +417,27 @@ 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
@@ -278,7 +449,7 @@ BuildDriverInfo (
   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 +473,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 +511,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 +550,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 +568,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 +643,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 +682,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 +698,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 +713,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 +790,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 +802,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 +829,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 +879,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 +918,7 @@ GetMemoryProfileDriverInfoFromAddress (
     }\r
   }\r
 \r
-  //\r
-  // Should never come here.\r
-  //\r
-  return FindDummyImage (ContextData);\r
+  return NULL;\r
 }\r
 \r
 /**\r
@@ -644,11 +927,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 +945,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 +956,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 +992,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 +1010,7 @@ UnregisterSmramProfileImage (
     SmmInternalFreePool (DriverInfoData);\r
   }\r
 \r
-  return TRUE;\r
+  return EFI_SUCCESS;\r
 }\r
 \r
 /**\r
@@ -807,54 +1103,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 +1184,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 +1242,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 +1264,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 +1296,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 +1319,140 @@ 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
-    }\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 (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
-  Context = &ContextData->Context;\r
-  DriverInfo = &DriverInfoData->DriverInfo;\r
-  AllocInfo = &AllocInfoData->AllocInfo;\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
-  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
@@ -1098,68 +1461,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 +1578,20 @@ 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
 \r
   ContextData = GetSmramProfileContext ();\r
   if (ContextData == NULL) {\r
@@ -1210,7 +1599,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 +1610,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
@@ -1291,6 +1692,8 @@ SmramProfileCopyData (
   MEMORY_PROFILE_DESCRIPTOR       *MemoryProfileDescriptor;\r
   UINT64                          Offset;\r
   UINT64                          RemainingSize;\r
+  UINTN                           PdbSize;\r
+  UINTN                           ActionStringSize;\r
 \r
   ContextData = GetSmramProfileContext ();\r
   if (ContextData == NULL) {\r
@@ -1322,17 +1725,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 +1751,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
@@ -1484,6 +1895,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 +2142,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 +2173,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 +2210,7 @@ SmramProfileHandlerGetData (
   SmramProfileParameterGetData->Header.ReturnStatus = 0;\r
 \r
 Done:\r
-  mSmramProfileRecordingStatus = SmramProfileRecordingStatus;\r
+  mSmramProfileGettingStatus = SmramProfileGettingStatus;\r
 }\r
 \r
 /**\r
@@ -1580,15 +2226,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 +2253,7 @@ SmramProfileHandlerGetDataByOffset (
   SmramProfileParameterGetDataByOffset->Header.ReturnStatus = 0;\r
 \r
 Done:\r
-  mSmramProfileRecordingStatus = SmramProfileRecordingStatus;\r
+  mSmramProfileGettingStatus = SmramProfileGettingStatus;\r
 }\r
 \r
 /**\r
@@ -1624,7 +2270,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 +2279,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 +2299,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 +2308,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 +2341,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 +2420,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 +2487,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 +2512,7 @@ DumpSmramRange (
 \r
   DEBUG ((EFI_D_INFO, "======= SmramProfile end =======\n"));\r
 \r
-  mSmramProfileRecordingStatus = SmramProfileRecordingStatus;\r
+  mSmramProfileGettingStatus = SmramProfileGettingStatus;\r
 }\r
 \r
 /**\r
@@ -1863,15 +2529,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 +2554,7 @@ DumpFreePagesList (
 \r
   DEBUG ((EFI_D_INFO, "======= SmramProfile end =======\n"));\r
 \r
-  mSmramProfileRecordingStatus = SmramProfileRecordingStatus;\r
+  mSmramProfileGettingStatus = SmramProfileGettingStatus;\r
 }\r
 \r
 /**\r
@@ -1900,21 +2566,21 @@ 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
 \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
@@ -1934,25 +2600,56 @@ DumpFreePoolList (
 \r
   DEBUG ((EFI_D_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,7 +2660,7 @@ 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
@@ -1975,7 +2672,39 @@ ProfileMemoryTypeToStr (
     }\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 = sizeof (mSmmActionString) / sizeof (mSmmActionString[0]);\r
+\r
+  if ((UINTN) (UINT32) Action < ActionStringCount) {\r
+    return ActionString[Action];\r
+  }\r
+  for (Index = 0; Index < sizeof (mExtActionString) / sizeof (mExtActionString[0]); 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 +2728,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 +2736,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 +2748,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 +2779,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 +2799,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 +2816,7 @@ DumpSmramProfile (
 \r
   DEBUG ((EFI_D_INFO, "======= SmramProfile end =======\n"));\r
 \r
-  mSmramProfileRecordingStatus = SmramProfileRecordingStatus;\r
+  mSmramProfileGettingStatus = SmramProfileGettingStatus;\r
 }\r
 \r
 /**\r