]> git.proxmox.com Git - mirror_edk2.git/blobdiff - MdeModulePkg/Core/Dxe/Mem/MemoryProfileRecord.c
MdeModulePkg: Apply uncrustify changes
[mirror_edk2.git] / MdeModulePkg / Core / Dxe / Mem / MemoryProfileRecord.c
index 6626e101592ece266fd1fa613e42a73d93326f16..00e33b707d4b95f28992f99c4c83ae3f7de0a904 100644 (file)
@@ -1,44 +1,42 @@
 /** @file\r
   Support routines for UEFI memory profile.\r
 \r
-  Copyright (c) 2014 - 2016, Intel Corporation. All rights reserved.<BR>\r
-  This program and the accompanying materials\r
-  are licensed and made available under the terms and conditions of the BSD License\r
-  which accompanies this distribution.  The full text of the license may be found at\r
-  http://opensource.org/licenses/bsd-license.php.\r
-\r
-  THE PROGRAM IS DISTRIBUTED UNDER THE BSD LICENSE ON AN "AS IS" BASIS,\r
-  WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.\r
+  Copyright (c) 2014 - 2018, Intel Corporation. All rights reserved.<BR>\r
+  SPDX-License-Identifier: BSD-2-Clause-Patent\r
 \r
 **/\r
 \r
 #include "DxeMain.h"\r
 #include "Imem.h"\r
 \r
-#define IS_UEFI_MEMORY_PROFILE_ENABLED ((PcdGet8 (PcdMemoryProfilePropertyMask) & BIT0) != 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
-  MEMORY_PROFILE_CONTEXT        Context;\r
-  LIST_ENTRY                    *DriverInfoList;\r
+  UINT32                    Signature;\r
+  MEMORY_PROFILE_CONTEXT    Context;\r
+  LIST_ENTRY                *DriverInfoList;\r
 } MEMORY_PROFILE_CONTEXT_DATA;\r
 \r
 typedef struct {\r
   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
-  LIST_ENTRY                    Link;\r
+  UINT32                       Signature;\r
+  MEMORY_PROFILE_ALLOC_INFO    AllocInfo;\r
+  CHAR8                        *ActionString;\r
+  LIST_ENTRY                   Link;\r
 } MEMORY_PROFILE_ALLOC_INFO_DATA;\r
 \r
-\r
-GLOBAL_REMOVE_IF_UNREFERENCED LIST_ENTRY  mImageQueue = INITIALIZE_LIST_HEAD_VARIABLE (mImageQueue);\r
-GLOBAL_REMOVE_IF_UNREFERENCED MEMORY_PROFILE_CONTEXT_DATA mMemoryProfileContext = {\r
+GLOBAL_REMOVE_IF_UNREFERENCED LIST_ENTRY                   mImageQueue           = INITIALIZE_LIST_HEAD_VARIABLE (mImageQueue);\r
+GLOBAL_REMOVE_IF_UNREFERENCED MEMORY_PROFILE_CONTEXT_DATA  mMemoryProfileContext = {\r
   MEMORY_PROFILE_CONTEXT_SIGNATURE,\r
   {\r
     {\r
@@ -48,17 +46,21 @@ GLOBAL_REMOVE_IF_UNREFERENCED MEMORY_PROFILE_CONTEXT_DATA mMemoryProfileContext
     },\r
     0,\r
     0,\r
-    {0},\r
-    {0},\r
+    { 0 },\r
+    { 0 },\r
     0,\r
     0,\r
     0\r
   },\r
   &mImageQueue,\r
 };\r
-GLOBAL_REMOVE_IF_UNREFERENCED MEMORY_PROFILE_CONTEXT_DATA *mMemoryProfileContextPtr = NULL;\r
+GLOBAL_REMOVE_IF_UNREFERENCED MEMORY_PROFILE_CONTEXT_DATA  *mMemoryProfileContextPtr = NULL;\r
 \r
-BOOLEAN mMemoryProfileRecordingStatus = FALSE;\r
+GLOBAL_REMOVE_IF_UNREFERENCED EFI_LOCK                  mMemoryProfileLock            = EFI_INITIALIZE_LOCK_VARIABLE (TPL_NOTIFY);\r
+GLOBAL_REMOVE_IF_UNREFERENCED BOOLEAN                   mMemoryProfileGettingStatus   = FALSE;\r
+GLOBAL_REMOVE_IF_UNREFERENCED BOOLEAN                   mMemoryProfileRecordingEnable = MEMORY_PROFILE_RECORDING_DISABLE;\r
+GLOBAL_REMOVE_IF_UNREFERENCED EFI_DEVICE_PATH_PROTOCOL  *mMemoryProfileDriverPath;\r
+GLOBAL_REMOVE_IF_UNREFERENCED UINTN                     mMemoryProfileDriverPathSize;\r
 \r
 /**\r
   Get memory profile data.\r
@@ -67,9 +69,10 @@ BOOLEAN mMemoryProfileRecordingStatus = FALSE;
   @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
+\r
   @return EFI_SUCCESS               Get the memory profile data successfully.\r
-  @return EFI_BUFFER_TO_SMALL       The ProfileSize is too small for the resulting data. \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
@@ -78,7 +81,7 @@ EFIAPI
 ProfileProtocolGetData (\r
   IN     EDKII_MEMORY_PROFILE_PROTOCOL  *This,\r
   IN OUT UINT64                         *ProfileSize,\r
-     OUT VOID                           *ProfileBuffer\r
+  OUT VOID                              *ProfileBuffer\r
   );\r
 \r
 /**\r
@@ -90,18 +93,20 @@ ProfileProtocolGetData (
   @param[in] ImageSize          Image size.\r
   @param[in] FileType           File type of the image.\r
 \r
-  @return EFI_SUCCESS           Register success.\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
 ProfileProtocolRegisterImage (\r
-  IN EDKII_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
+  IN EDKII_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
@@ -112,66 +117,134 @@ ProfileProtocolRegisterImage (
   @param[in] ImageBase          Image base address.\r
   @param[in] ImageSize          Image size.\r
 \r
-  @return EFI_SUCCESS           Unregister success.\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
 ProfileProtocolUnregisterImage (\r
-  IN EDKII_MEMORY_PROFILE_PROTOCOL      *This,\r
-  IN EFI_DEVICE_PATH_PROTOCOL           *FilePath,\r
-  IN PHYSICAL_ADDRESS                   ImageBase,\r
-  IN UINT64                             ImageSize\r
+  IN EDKII_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_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
+ProfileProtocolGetRecordingState (\r
+  IN EDKII_MEMORY_PROFILE_PROTOCOL  *This,\r
+  OUT BOOLEAN                       *RecordingState\r
   );\r
 \r
-EDKII_MEMORY_PROFILE_PROTOCOL mProfileProtocol = {\r
+/**\r
+  Set memory profile recording state.\r
+\r
+  @param[in] This               The EDKII_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
+ProfileProtocolSetRecordingState (\r
+  IN EDKII_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_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
+ProfileProtocolRecord (\r
+  IN EDKII_MEMORY_PROFILE_PROTOCOL  *This,\r
+  IN PHYSICAL_ADDRESS               CallerAddress,\r
+  IN MEMORY_PROFILE_ACTION          Action,\r
+  IN EFI_MEMORY_TYPE                MemoryType,\r
+  IN VOID                           *Buffer,\r
+  IN UINTN                          Size,\r
+  IN CHAR8                          *ActionString OPTIONAL\r
+  );\r
+\r
+GLOBAL_REMOVE_IF_UNREFERENCED EDKII_MEMORY_PROFILE_PROTOCOL  mProfileProtocol = {\r
   ProfileProtocolGetData,\r
   ProfileProtocolRegisterImage,\r
-  ProfileProtocolUnregisterImage\r
+  ProfileProtocolUnregisterImage,\r
+  ProfileProtocolGetRecordingState,\r
+  ProfileProtocolSetRecordingState,\r
+  ProfileProtocolRecord,\r
 };\r
 \r
 /**\r
-  Return memory profile context.\r
-\r
-  @return Memory profile context.\r
-\r
+  Acquire lock on mMemoryProfileLock.\r
 **/\r
-MEMORY_PROFILE_CONTEXT_DATA *\r
-GetMemoryProfileContext (\r
+VOID\r
+CoreAcquireMemoryProfileLock (\r
   VOID\r
   )\r
 {\r
-  return mMemoryProfileContextPtr;\r
+  CoreAcquireLock (&mMemoryProfileLock);\r
 }\r
 \r
 /**\r
-  Retrieves the magic value from the PE/COFF header.\r
+  Release lock on mMemoryProfileLock.\r
+**/\r
+VOID\r
+CoreReleaseMemoryProfileLock (\r
+  VOID\r
+  )\r
+{\r
+  CoreReleaseLock (&mMemoryProfileLock);\r
+}\r
 \r
-  @param Hdr    The buffer in which to return the PE32, PE32+, or TE header.\r
+/**\r
+  Return memory profile context.\r
 \r
-  @return EFI_IMAGE_NT_OPTIONAL_HDR32_MAGIC - Image is PE32\r
-  @return EFI_IMAGE_NT_OPTIONAL_HDR64_MAGIC - Image is PE32+\r
+  @return Memory profile context.\r
 \r
 **/\r
-UINT16\r
-InternalPeCoffGetPeHeaderMagicValue (\r
-  IN  EFI_IMAGE_OPTIONAL_HEADER_PTR_UNION  Hdr\r
+MEMORY_PROFILE_CONTEXT_DATA *\r
+GetMemoryProfileContext (\r
+  VOID\r
   )\r
 {\r
-  //\r
-  // NOTE: Some versions of Linux ELILO for Itanium have an incorrect magic value\r
-  //       in the PE/COFF Header.  If the MachineType is Itanium(IA64) and the\r
-  //       Magic value in the OptionalHeader is  EFI_IMAGE_NT_OPTIONAL_HDR32_MAGIC\r
-  //       then override the returned value to EFI_IMAGE_NT_OPTIONAL_HDR64_MAGIC\r
-  //\r
-  if (Hdr.Pe32->FileHeader.Machine == IMAGE_FILE_MACHINE_IA64 && Hdr.Pe32->OptionalHeader.Magic == EFI_IMAGE_NT_OPTIONAL_HDR32_MAGIC) {\r
-    return EFI_IMAGE_NT_OPTIONAL_HDR64_MAGIC;\r
-  }\r
-  //\r
-  // Return the magic value from the PC/COFF Optional Header\r
-  //\r
-  return Hdr.Pe32->OptionalHeader.Magic;\r
+  return mMemoryProfileContextPtr;\r
 }\r
 \r
 /**\r
@@ -194,23 +267,23 @@ InternalPeCoffGetSubsystem (
 \r
   ASSERT (Pe32Data != NULL);\r
 \r
-  DosHdr = (EFI_IMAGE_DOS_HEADER *) Pe32Data;\r
+  DosHdr = (EFI_IMAGE_DOS_HEADER *)Pe32Data;\r
   if (DosHdr->e_magic == EFI_IMAGE_DOS_SIGNATURE) {\r
     //\r
     // DOS image header is present, so read the PE header after the DOS image header.\r
     //\r
-    Hdr.Pe32 = (EFI_IMAGE_NT_HEADERS32 *) ((UINTN) Pe32Data + (UINTN) ((DosHdr->e_lfanew) & 0x0ffff));\r
+    Hdr.Pe32 = (EFI_IMAGE_NT_HEADERS32 *)((UINTN)Pe32Data + (UINTN)((DosHdr->e_lfanew) & 0x0ffff));\r
   } else {\r
     //\r
     // DOS image header is not present, so PE header is at the image base.\r
     //\r
-    Hdr.Pe32 = (EFI_IMAGE_NT_HEADERS32 *) Pe32Data;\r
+    Hdr.Pe32 = (EFI_IMAGE_NT_HEADERS32 *)Pe32Data;\r
   }\r
 \r
   if (Hdr.Te->Signature == EFI_TE_IMAGE_HEADER_SIGNATURE) {\r
     return Hdr.Te->Subsystem;\r
-  } else if (Hdr.Pe32->Signature == EFI_IMAGE_NT_SIGNATURE)  {\r
-    Magic = InternalPeCoffGetPeHeaderMagicValue (Hdr);\r
+  } else if (Hdr.Pe32->Signature == EFI_IMAGE_NT_SIGNATURE) {\r
+    Magic = Hdr.Pe32->OptionalHeader.Magic;\r
     if (Magic == EFI_IMAGE_NT_OPTIONAL_HDR32_MAGIC) {\r
       return Hdr.Pe32->OptionalHeader.Subsystem;\r
     } else if (Magic == EFI_IMAGE_NT_OPTIONAL_HDR64_MAGIC) {\r
@@ -244,23 +317,23 @@ InternalPeCoffGetEntryPoint (
   OUT VOID  **EntryPoint\r
   )\r
 {\r
-  EFI_IMAGE_DOS_HEADER                  *DosHdr;\r
-  EFI_IMAGE_OPTIONAL_HEADER_PTR_UNION   Hdr;\r
+  EFI_IMAGE_DOS_HEADER                 *DosHdr;\r
+  EFI_IMAGE_OPTIONAL_HEADER_PTR_UNION  Hdr;\r
 \r
   ASSERT (Pe32Data   != NULL);\r
   ASSERT (EntryPoint != NULL);\r
 \r
-  DosHdr = (EFI_IMAGE_DOS_HEADER *) Pe32Data;\r
+  DosHdr = (EFI_IMAGE_DOS_HEADER *)Pe32Data;\r
   if (DosHdr->e_magic == EFI_IMAGE_DOS_SIGNATURE) {\r
     //\r
     // DOS image header is present, so read the PE header after the DOS image header.\r
     //\r
-    Hdr.Pe32 = (EFI_IMAGE_NT_HEADERS32 *) ((UINTN) Pe32Data + (UINTN) ((DosHdr->e_lfanew) & 0x0ffff));\r
+    Hdr.Pe32 = (EFI_IMAGE_NT_HEADERS32 *)((UINTN)Pe32Data + (UINTN)((DosHdr->e_lfanew) & 0x0ffff));\r
   } else {\r
     //\r
     // DOS image header is not present, so PE header is at the image base.\r
     //\r
-    Hdr.Pe32 = (EFI_IMAGE_NT_HEADERS32 *) Pe32Data;\r
+    Hdr.Pe32 = (EFI_IMAGE_NT_HEADERS32 *)Pe32Data;\r
   }\r
 \r
   //\r
@@ -268,10 +341,10 @@ InternalPeCoffGetEntryPoint (
   // AddressOfEntryPoint is common for PE32 & PE32+\r
   //\r
   if (Hdr.Te->Signature == EFI_TE_IMAGE_HEADER_SIGNATURE) {\r
-    *EntryPoint = (VOID *) ((UINTN) Pe32Data + (UINTN) (Hdr.Te->AddressOfEntryPoint & 0x0ffffffff) + sizeof (EFI_TE_IMAGE_HEADER) - Hdr.Te->StrippedSize);\r
+    *EntryPoint = (VOID *)((UINTN)Pe32Data + (UINTN)(Hdr.Te->AddressOfEntryPoint & 0x0ffffffff) + sizeof (EFI_TE_IMAGE_HEADER) - Hdr.Te->StrippedSize);\r
     return RETURN_SUCCESS;\r
   } else if (Hdr.Pe32->Signature == EFI_IMAGE_NT_SIGNATURE) {\r
-    *EntryPoint = (VOID *) ((UINTN) Pe32Data + (UINTN) (Hdr.Pe32->OptionalHeader.AddressOfEntryPoint & 0x0ffffffff));\r
+    *EntryPoint = (VOID *)((UINTN)Pe32Data + (UINTN)(Hdr.Pe32->OptionalHeader.AddressOfEntryPoint & 0x0ffffffff));\r
     return RETURN_SUCCESS;\r
   }\r
 \r
@@ -294,69 +367,155 @@ InternalPeCoffGetEntryPoint (
 **/\r
 MEMORY_PROFILE_DRIVER_INFO_DATA *\r
 BuildDriverInfo (\r
-  IN MEMORY_PROFILE_CONTEXT_DATA    *ContextData,\r
-  IN EFI_GUID                       *FileName,\r
-  IN PHYSICAL_ADDRESS               ImageBase,\r
-  IN UINT64                         ImageSize,\r
-  IN PHYSICAL_ADDRESS               EntryPoint,\r
-  IN UINT16                         ImageSubsystem,\r
-  IN EFI_FV_FILETYPE                FileType\r
+  IN MEMORY_PROFILE_CONTEXT_DATA  *ContextData,\r
+  IN EFI_GUID                     *FileName,\r
+  IN PHYSICAL_ADDRESS             ImageBase,\r
+  IN UINT64                       ImageSize,\r
+  IN PHYSICAL_ADDRESS             EntryPoint,\r
+  IN UINT16                       ImageSubsystem,\r
+  IN EFI_FV_FILETYPE              FileType\r
   )\r
 {\r
-  EFI_STATUS                        Status;\r
-  MEMORY_PROFILE_DRIVER_INFO        *DriverInfo;\r
-  MEMORY_PROFILE_DRIVER_INFO_DATA   *DriverInfoData;\r
-  VOID                              *EntryPointInImage;\r
+  EFI_STATUS                       Status;\r
+  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 CoreInternalAllocatePool() that will not update profile for this AllocatePool action.\r
   //\r
   Status = CoreInternalAllocatePool (\r
              EfiBootServicesData,\r
-             sizeof (*DriverInfoData) + sizeof (LIST_ENTRY),\r
-             (VOID **) &DriverInfoData\r
+             sizeof (*DriverInfoData) + sizeof (LIST_ENTRY) + PdbSize,\r
+             (VOID **)&DriverInfoData\r
              );\r
   if (EFI_ERROR (Status)) {\r
     return NULL;\r
   }\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                   = &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.Revision = MEMORY_PROFILE_DRIVER_INFO_REVISION;\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
   }\r
-  DriverInfo->ImageBase = ImageBase;\r
-  DriverInfo->ImageSize = ImageSize;\r
-  DriverInfo->EntryPoint = EntryPoint;\r
+\r
+  DriverInfo->ImageBase      = ImageBase;\r
+  DriverInfo->ImageSize      = ImageSize;\r
+  DriverInfo->EntryPoint     = EntryPoint;\r
   DriverInfo->ImageSubsystem = ImageSubsystem;\r
   if ((EntryPoint != 0) && ((EntryPoint < ImageBase) || (EntryPoint >= (ImageBase + ImageSize)))) {\r
     //\r
     // If the EntryPoint is not in the range of image buffer, it should come from emulation environment.\r
     // So patch ImageBuffer here to align the EntryPoint.\r
     //\r
-    Status = InternalPeCoffGetEntryPoint ((VOID *) (UINTN) ImageBase, &EntryPointInImage);\r
+    Status = InternalPeCoffGetEntryPoint ((VOID *)(UINTN)ImageBase, &EntryPointInImage);\r
     ASSERT_EFI_ERROR (Status);\r
-    DriverInfo->ImageBase = ImageBase + EntryPoint - (PHYSICAL_ADDRESS) (UINTN) EntryPointInImage;\r
+    DriverInfo->ImageBase = ImageBase + EntryPoint - (PHYSICAL_ADDRESS)(UINTN)EntryPointInImage;\r
   }\r
-  DriverInfo->FileType = FileType;\r
-  DriverInfoData->AllocInfoList = (LIST_ENTRY *) (DriverInfoData + 1);\r
+\r
+  DriverInfo->FileType          = FileType;\r
+  DriverInfoData->AllocInfoList = (LIST_ENTRY *)(DriverInfoData + 1);\r
   InitializeListHead (DriverInfoData->AllocInfoList);\r
-  DriverInfo->CurrentUsage = 0;\r
-  DriverInfo->PeakUsage = 0;\r
+  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
+  ContextData->Context.ImageCount++;\r
   ContextData->Context.TotalImageSize += DriverInfo->ImageSize;\r
 \r
   return DriverInfoData;\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 (mMemoryProfileDriverPath, mMemoryProfileDriverPathSize)) {\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 = mMemoryProfileDriverPath;\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
+    {\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 DXE Core to memory profile.\r
 \r
@@ -369,20 +528,22 @@ BuildDriverInfo (
 **/\r
 BOOLEAN\r
 RegisterDxeCore (\r
-  IN VOID                           *HobStart,\r
-  IN MEMORY_PROFILE_CONTEXT_DATA    *ContextData\r
+  IN VOID                         *HobStart,\r
+  IN MEMORY_PROFILE_CONTEXT_DATA  *ContextData\r
   )\r
 {\r
-  EFI_PEI_HOB_POINTERS              DxeCoreHob;\r
-  MEMORY_PROFILE_DRIVER_INFO_DATA   *DriverInfoData;\r
-  PHYSICAL_ADDRESS                  ImageBase;\r
+  EFI_PEI_HOB_POINTERS               DxeCoreHob;\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
 \r
   //\r
   // Searching for image hob\r
   //\r
-  DxeCoreHob.Raw          = HobStart;\r
+  DxeCoreHob.Raw = HobStart;\r
   while ((DxeCoreHob.Raw = GetNextHob (EFI_HOB_TYPE_MEMORY_ALLOCATION, DxeCoreHob.Raw)) != NULL) {\r
     if (CompareGuid (&DxeCoreHob.MemoryAllocationModule->MemoryAllocationHeader.Name, &gEfiHobMemoryAllocModuleGuid)) {\r
       //\r
@@ -390,18 +551,28 @@ RegisterDxeCore (
       //\r
       break;\r
     }\r
+\r
     DxeCoreHob.Raw = GET_NEXT_HOB (DxeCoreHob);\r
   }\r
+\r
   ASSERT (DxeCoreHob.Raw != NULL);\r
 \r
-  ImageBase = DxeCoreHob.MemoryAllocationModule->MemoryAllocationHeader.MemoryBaseAddress;\r
+  FilePath = (MEDIA_FW_VOL_FILEPATH_DEVICE_PATH *)TempBuffer;\r
+  EfiInitializeFwVolDevicepathNode (FilePath, &DxeCoreHob.MemoryAllocationModule->ModuleName);\r
+  SetDevicePathEndNode (FilePath + 1);\r
+\r
+  if (!NeedRecordThisDriver ((EFI_DEVICE_PATH_PROTOCOL *)FilePath)) {\r
+    return FALSE;\r
+  }\r
+\r
+  ImageBase      = DxeCoreHob.MemoryAllocationModule->MemoryAllocationHeader.MemoryBaseAddress;\r
   DriverInfoData = BuildDriverInfo (\r
                      ContextData,\r
                      &DxeCoreHob.MemoryAllocationModule->ModuleName,\r
                      ImageBase,\r
                      DxeCoreHob.MemoryAllocationModule->MemoryAllocationHeader.MemoryLength,\r
                      DxeCoreHob.MemoryAllocationModule->EntryPoint,\r
-                     InternalPeCoffGetSubsystem ((VOID *) (UINTN) ImageBase),\r
+                     InternalPeCoffGetSubsystem ((VOID *)(UINTN)ImageBase),\r
                      EFI_FV_FILETYPE_DXE_CORE\r
                      );\r
   if (DriverInfoData == NULL) {\r
@@ -419,10 +590,10 @@ RegisterDxeCore (
 **/\r
 VOID\r
 MemoryProfileInit (\r
-  IN VOID   *HobStart\r
+  IN VOID  *HobStart\r
   )\r
 {\r
-  MEMORY_PROFILE_CONTEXT_DATA   *ContextData;\r
+  MEMORY_PROFILE_CONTEXT_DATA  *ContextData;\r
 \r
   if (!IS_UEFI_MEMORY_PROFILE_ENABLED) {\r
     return;\r
@@ -433,12 +604,20 @@ MemoryProfileInit (
     return;\r
   }\r
 \r
-  mMemoryProfileRecordingStatus = TRUE;\r
-  mMemoryProfileContextPtr = &mMemoryProfileContext;\r
+  mMemoryProfileGettingStatus = FALSE;\r
+  if ((PcdGet8 (PcdMemoryProfilePropertyMask) & BIT7) != 0) {\r
+    mMemoryProfileRecordingEnable = MEMORY_PROFILE_RECORDING_DISABLE;\r
+  } else {\r
+    mMemoryProfileRecordingEnable = MEMORY_PROFILE_RECORDING_ENABLE;\r
+  }\r
+\r
+  mMemoryProfileDriverPathSize = PcdGetSize (PcdMemoryProfileDriverPath);\r
+  mMemoryProfileDriverPath     = AllocateCopyPool (mMemoryProfileDriverPathSize, PcdGetPtr (PcdMemoryProfileDriverPath));\r
+  mMemoryProfileContextPtr     = &mMemoryProfileContext;\r
 \r
   RegisterDxeCore (HobStart, &mMemoryProfileContext);\r
 \r
-  DEBUG ((EFI_D_INFO, "MemoryProfileInit MemoryProfileContext - 0x%x\n", &mMemoryProfileContext));\r
+  DEBUG ((DEBUG_INFO, "MemoryProfileInit MemoryProfileContext - 0x%x\n", &mMemoryProfileContext));\r
 }\r
 \r
 /**\r
@@ -450,8 +629,8 @@ MemoryProfileInstallProtocol (
   VOID\r
   )\r
 {\r
-  EFI_HANDLE    Handle;\r
-  EFI_STATUS    Status;\r
+  EFI_HANDLE  Handle;\r
+  EFI_STATUS  Status;\r
 \r
   if (!IS_UEFI_MEMORY_PROFILE_ENABLED) {\r
     return;\r
@@ -477,21 +656,22 @@ MemoryProfileInstallProtocol (
 **/\r
 EFI_GUID *\r
 GetFileNameFromFilePath (\r
-  IN EFI_DEVICE_PATH_PROTOCOL   *FilePath\r
+  IN EFI_DEVICE_PATH_PROTOCOL  *FilePath\r
   )\r
 {\r
-  MEDIA_FW_VOL_FILEPATH_DEVICE_PATH     *ThisFilePath;\r
-  EFI_GUID                              *FileName;\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
+    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
+      ThisFilePath = (MEDIA_FW_VOL_FILEPATH_DEVICE_PATH *)NextDevicePathNode (ThisFilePath);\r
     }\r
   }\r
 \r
@@ -504,26 +684,32 @@ GetFileNameFromFilePath (
   @param DriverEntry    Image info.\r
   @param FileType       Image file type.\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
 RegisterMemoryProfileImage (\r
   IN LOADED_IMAGE_PRIVATE_DATA  *DriverEntry,\r
   IN EFI_FV_FILETYPE            FileType\r
   )\r
 {\r
-  MEMORY_PROFILE_CONTEXT_DATA       *ContextData;\r
-  MEMORY_PROFILE_DRIVER_INFO_DATA   *DriverInfoData;\r
+  MEMORY_PROFILE_CONTEXT_DATA      *ContextData;\r
+  MEMORY_PROFILE_DRIVER_INFO_DATA  *DriverInfoData;\r
 \r
   if (!IS_UEFI_MEMORY_PROFILE_ENABLED) {\r
-    return FALSE;\r
+    return EFI_UNSUPPORTED;\r
+  }\r
+\r
+  if (!NeedRecordThisDriver (DriverEntry->Info.FilePath)) {\r
+    return EFI_UNSUPPORTED;\r
   }\r
 \r
   ContextData = GetMemoryProfileContext ();\r
   if (ContextData == NULL) {\r
-    return FALSE;\r
+    return EFI_UNSUPPORTED;\r
   }\r
 \r
   DriverInfoData = BuildDriverInfo (\r
@@ -536,10 +722,10 @@ RegisterMemoryProfileImage (
                      FileType\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
@@ -554,21 +740,22 @@ RegisterMemoryProfileImage (
 **/\r
 MEMORY_PROFILE_DRIVER_INFO_DATA *\r
 GetMemoryProfileDriverInfoByFileNameAndAddress (\r
-  IN MEMORY_PROFILE_CONTEXT_DATA    *ContextData,\r
-  IN EFI_GUID                       *FileName,\r
-  IN PHYSICAL_ADDRESS               Address\r
+  IN MEMORY_PROFILE_CONTEXT_DATA  *ContextData,\r
+  IN EFI_GUID                     *FileName,\r
+  IN PHYSICAL_ADDRESS             Address\r
   )\r
 {\r
-  MEMORY_PROFILE_DRIVER_INFO        *DriverInfo;\r
-  MEMORY_PROFILE_DRIVER_INFO_DATA   *DriverInfoData;\r
-  LIST_ENTRY                        *DriverLink;\r
-  LIST_ENTRY                        *DriverInfoList;\r
+  MEMORY_PROFILE_DRIVER_INFO       *DriverInfo;\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
+       DriverLink = DriverLink->ForwardLink)\r
+  {\r
     DriverInfoData = CR (\r
                        DriverLink,\r
                        MEMORY_PROFILE_DRIVER_INFO_DATA,\r
@@ -578,7 +765,8 @@ GetMemoryProfileDriverInfoByFileNameAndAddress (
     DriverInfo = &DriverInfoData->DriverInfo;\r
     if ((CompareGuid (&DriverInfo->FileName, FileName)) &&\r
         (Address >= DriverInfo->ImageBase) &&\r
-        (Address < (DriverInfo->ImageBase + DriverInfo->ImageSize))) {\r
+        (Address < (DriverInfo->ImageBase + DriverInfo->ImageSize)))\r
+    {\r
       return DriverInfoData;\r
     }\r
   }\r
@@ -586,45 +774,9 @@ GetMemoryProfileDriverInfoByFileNameAndAddress (
   return NULL;\r
 }\r
 \r
-/**\r
-  Search dummy image from memory 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
+  It will return image, if (Address >= ImageBuffer) AND (Address < ImageBuffer + ImageSize).\r
 \r
   @param ContextData    Memory profile context.\r
   @param Address        Image or Function address.\r
@@ -634,20 +786,21 @@ FindDummyImage (
 **/\r
 MEMORY_PROFILE_DRIVER_INFO_DATA *\r
 GetMemoryProfileDriverInfoFromAddress (\r
-  IN MEMORY_PROFILE_CONTEXT_DATA    *ContextData,\r
-  IN PHYSICAL_ADDRESS               Address\r
+  IN MEMORY_PROFILE_CONTEXT_DATA  *ContextData,\r
+  IN PHYSICAL_ADDRESS             Address\r
   )\r
 {\r
-  MEMORY_PROFILE_DRIVER_INFO        *DriverInfo;\r
-  MEMORY_PROFILE_DRIVER_INFO_DATA   *DriverInfoData;\r
-  LIST_ENTRY                        *DriverLink;\r
-  LIST_ENTRY                        *DriverInfoList;\r
+  MEMORY_PROFILE_DRIVER_INFO       *DriverInfo;\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
+       DriverLink = DriverLink->ForwardLink)\r
+  {\r
     DriverInfoData = CR (\r
                        DriverLink,\r
                        MEMORY_PROFILE_DRIVER_INFO_DATA,\r
@@ -656,15 +809,13 @@ GetMemoryProfileDriverInfoFromAddress (
                        );\r
     DriverInfo = &DriverInfoData->DriverInfo;\r
     if ((Address >= DriverInfo->ImageBase) &&\r
-        (Address < (DriverInfo->ImageBase + DriverInfo->ImageSize))) {\r
+        (Address < (DriverInfo->ImageBase + DriverInfo->ImageSize)))\r
+    {\r
       return DriverInfoData;\r
     }\r
   }\r
 \r
-  //\r
-  // Should never come here.\r
-  //\r
-  return FindDummyImage (ContextData);\r
+  return NULL;\r
 }\r
 \r
 /**\r
@@ -672,60 +823,70 @@ GetMemoryProfileDriverInfoFromAddress (
 \r
   @param DriverEntry    Image info.\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
 UnregisterMemoryProfileImage (\r
-  IN LOADED_IMAGE_PRIVATE_DATA      *DriverEntry\r
+  IN LOADED_IMAGE_PRIVATE_DATA  *DriverEntry\r
   )\r
 {\r
-  EFI_STATUS                        Status;\r
-  MEMORY_PROFILE_CONTEXT_DATA       *ContextData;\r
-  MEMORY_PROFILE_DRIVER_INFO_DATA   *DriverInfoData;\r
-  EFI_GUID                          *FileName;\r
-  PHYSICAL_ADDRESS                  ImageAddress;\r
-  VOID                              *EntryPointInImage;\r
+  EFI_STATUS                       Status;\r
+  MEMORY_PROFILE_CONTEXT_DATA      *ContextData;\r
+  MEMORY_PROFILE_DRIVER_INFO_DATA  *DriverInfoData;\r
+  EFI_GUID                         *FileName;\r
+  PHYSICAL_ADDRESS                 ImageAddress;\r
+  VOID                             *EntryPointInImage;\r
 \r
   if (!IS_UEFI_MEMORY_PROFILE_ENABLED) {\r
-    return FALSE;\r
+    return EFI_UNSUPPORTED;\r
+  }\r
+\r
+  if (!NeedRecordThisDriver (DriverEntry->Info.FilePath)) {\r
+    return EFI_UNSUPPORTED;\r
   }\r
 \r
   ContextData = GetMemoryProfileContext ();\r
   if (ContextData == NULL) {\r
-    return FALSE;\r
+    return EFI_UNSUPPORTED;\r
   }\r
 \r
   DriverInfoData = NULL;\r
-  FileName = GetFileNameFromFilePath (DriverEntry->Info.FilePath);\r
-  ImageAddress = DriverEntry->ImageContext.ImageAddress;\r
+  FileName       = GetFileNameFromFilePath (DriverEntry->Info.FilePath);\r
+  ImageAddress   = DriverEntry->ImageContext.ImageAddress;\r
   if ((DriverEntry->ImageContext.EntryPoint < ImageAddress) || (DriverEntry->ImageContext.EntryPoint >= (ImageAddress + DriverEntry->ImageContext.ImageSize))) {\r
     //\r
     // If the EntryPoint is not in the range of image buffer, it should come from emulation environment.\r
     // So patch ImageAddress here to align the EntryPoint.\r
     //\r
-    Status = InternalPeCoffGetEntryPoint ((VOID *) (UINTN) ImageAddress, &EntryPointInImage);\r
+    Status = InternalPeCoffGetEntryPoint ((VOID *)(UINTN)ImageAddress, &EntryPointInImage);\r
     ASSERT_EFI_ERROR (Status);\r
-    ImageAddress = ImageAddress + (UINTN) DriverEntry->ImageContext.EntryPoint - (UINTN) EntryPointInImage;\r
+    ImageAddress = ImageAddress + (UINTN)DriverEntry->ImageContext.EntryPoint - (UINTN)EntryPointInImage;\r
   }\r
+\r
   if (FileName != NULL) {\r
     DriverInfoData = GetMemoryProfileDriverInfoByFileNameAndAddress (ContextData, FileName, ImageAddress);\r
   }\r
+\r
   if (DriverInfoData == NULL) {\r
     DriverInfoData = GetMemoryProfileDriverInfoFromAddress (ContextData, ImageAddress);\r
   }\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
-    ContextData->Context.ImageCount --;\r
+    ContextData->Context.ImageCount--;\r
     RemoveEntryList (&DriverInfoData->Link);\r
     //\r
     // Use CoreInternalFreePool() that will not update profile for this FreePool action.\r
@@ -733,7 +894,7 @@ UnregisterMemoryProfileImage (
     CoreInternalFreePool (DriverInfoData, NULL);\r
   }\r
 \r
-  return TRUE;\r
+  return EFI_SUCCESS;\r
 }\r
 \r
 /**\r
@@ -750,14 +911,14 @@ UnregisterMemoryProfileImage (
 **/\r
 BOOLEAN\r
 CoreNeedRecordProfile (\r
-  IN EFI_MEMORY_TYPE    MemoryType\r
+  IN EFI_MEMORY_TYPE  MemoryType\r
   )\r
 {\r
-  UINT64 TestBit;\r
+  UINT64  TestBit;\r
 \r
-  if ((UINT32) MemoryType >= MEMORY_TYPE_OS_RESERVED_MIN) {\r
+  if ((UINT32)MemoryType >= MEMORY_TYPE_OS_RESERVED_MIN) {\r
     TestBit = BIT63;\r
-  } else if ((UINT32) MemoryType >= MEMORY_TYPE_OEM_RESERVED_MIN) {\r
+  } else if ((UINT32)MemoryType >= MEMORY_TYPE_OEM_RESERVED_MIN) {\r
     TestBit = BIT62;\r
   } else {\r
     TestBit = LShiftU64 (1, MemoryType);\r
@@ -783,12 +944,12 @@ CoreNeedRecordProfile (
 **/\r
 UINTN\r
 GetProfileMemoryIndex (\r
-  IN EFI_MEMORY_TYPE    MemoryType\r
+  IN EFI_MEMORY_TYPE  MemoryType\r
   )\r
 {\r
-  if ((UINT32) MemoryType >= MEMORY_TYPE_OS_RESERVED_MIN) {\r
+  if ((UINT32)MemoryType >= MEMORY_TYPE_OS_RESERVED_MIN) {\r
     return EfiMaxMemoryType;\r
-  } else if ((UINT32) MemoryType >= MEMORY_TYPE_OEM_RESERVED_MIN) {\r
+  } else if ((UINT32)MemoryType >= MEMORY_TYPE_OEM_RESERVED_MIN) {\r
     return EfiMaxMemoryType + 1;\r
   } else {\r
     return MemoryType;\r
@@ -803,120 +964,163 @@ GetProfileMemoryIndex (
   @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
 CoreUpdateProfileAllocate (\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
+  EFI_STATUS                       Status;\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
-  UINTN                             ProfileMemoryIndex;\r
+  MEMORY_PROFILE_CONTEXT_DATA      *ContextData;\r
+  MEMORY_PROFILE_DRIVER_INFO_DATA  *DriverInfoData;\r
+  MEMORY_PROFILE_ALLOC_INFO_DATA   *AllocInfoData;\r
+  UINTN                            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 = GetMemoryProfileContext ();\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 CoreInternalAllocatePool() that will not update profile for this AllocatePool action.\r
   //\r
-  Status = CoreInternalAllocatePool (\r
-             EfiBootServicesData,\r
-             sizeof (*AllocInfoData),\r
-             (VOID **) &AllocInfoData\r
-             );\r
+  AllocInfoData = NULL;\r
+  Status        = CoreInternalAllocatePool (\r
+                    EfiBootServicesData,\r
+                    sizeof (*AllocInfoData) + ActionStringSize,\r
+                    (VOID **)&AllocInfoData\r
+                    );\r
   if (EFI_ERROR (Status)) {\r
-    return FALSE;\r
+    return EFI_OUT_OF_RESOURCES;\r
   }\r
+\r
   ASSERT (AllocInfoData != NULL);\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.Revision    = MEMORY_PROFILE_ALLOC_INFO_REVISION;\r
-  AllocInfo->CallerAddress      = CallerAddress;\r
-  AllocInfo->SequenceId         = ContextData->Context.SequenceCount;\r
-  AllocInfo->Action             = Action;\r
-  AllocInfo->MemoryType         = MemoryType;\r
-  AllocInfo->Buffer             = (PHYSICAL_ADDRESS) (UINTN) Buffer;\r
-  AllocInfo->Size               = Size;\r
 \r
-  InsertTailList (DriverInfoData->AllocInfoList, &AllocInfoData->Link);\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    = (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
+  AllocInfo->Action           = Action;\r
+  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
-  ProfileMemoryIndex = GetProfileMemoryIndex (MemoryType);\r
+  InsertTailList (DriverInfoData->AllocInfoList, &AllocInfoData->Link);\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
+  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
+\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
+\r
+    Context->CurrentTotalUsageByType[ProfileMemoryIndex] += Size;\r
+    if (Context->PeakTotalUsageByType[ProfileMemoryIndex] < Context->CurrentTotalUsageByType[ProfileMemoryIndex]) {\r
+      Context->PeakTotalUsageByType[ProfileMemoryIndex] = Context->CurrentTotalUsageByType[ProfileMemoryIndex];\r
+    }\r
   }\r
-  Context->SequenceCount ++;\r
 \r
-  return TRUE;\r
+  return EFI_SUCCESS;\r
 }\r
 \r
 /**\r
-  Get memory profile alloc info from memory profile\r
+  Get memory profile alloc info from memory profile.\r
 \r
-  @param DriverInfoData     Driver info\r
-  @param Action             This Free action\r
-  @param Size               Buffer size\r
-  @param Buffer             Buffer address\r
+  @param DriverInfoData     Driver info.\r
+  @param BasicAction        This Free basic action.\r
+  @param Size               Buffer size.\r
+  @param Buffer             Buffer address.\r
 \r
   @return Pointer to memory profile alloc info.\r
+\r
 **/\r
 MEMORY_PROFILE_ALLOC_INFO_DATA *\r
 GetMemoryProfileAllocInfoFromAddress (\r
-  IN MEMORY_PROFILE_DRIVER_INFO_DATA    *DriverInfoData,\r
-  IN MEMORY_PROFILE_ACTION              Action,\r
-  IN UINTN                              Size,\r
-  IN VOID                               *Buffer\r
+  IN MEMORY_PROFILE_DRIVER_INFO_DATA  *DriverInfoData,\r
+  IN MEMORY_PROFILE_ACTION            BasicAction,\r
+  IN UINTN                            Size,\r
+  IN VOID                             *Buffer\r
   )\r
 {\r
-  LIST_ENTRY                        *AllocInfoList;\r
-  LIST_ENTRY                        *AllocLink;\r
-  MEMORY_PROFILE_ALLOC_INFO         *AllocInfo;\r
-  MEMORY_PROFILE_ALLOC_INFO_DATA    *AllocInfoData;\r
+  LIST_ENTRY                      *AllocInfoList;\r
+  LIST_ENTRY                      *AllocLink;\r
+  MEMORY_PROFILE_ALLOC_INFO       *AllocInfo;\r
+  MEMORY_PROFILE_ALLOC_INFO_DATA  *AllocInfoData;\r
 \r
   AllocInfoList = DriverInfoData->AllocInfoList;\r
 \r
   for (AllocLink = AllocInfoList->ForwardLink;\r
        AllocLink != AllocInfoList;\r
-       AllocLink = AllocLink->ForwardLink) {\r
+       AllocLink = AllocLink->ForwardLink)\r
+  {\r
     AllocInfoData = CR (\r
                       AllocLink,\r
                       MEMORY_PROFILE_ALLOC_INFO_DATA,\r
@@ -924,20 +1128,24 @@ 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
+\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
+        if ((AllocInfo->Buffer <= (PHYSICAL_ADDRESS)(UINTN)Buffer) &&\r
+            ((AllocInfo->Buffer + AllocInfo->Size) >= ((PHYSICAL_ADDRESS)(UINTN)Buffer + Size)))\r
+        {\r
           return AllocInfoData;\r
         }\r
+\r
         break;\r
       case MemoryProfileActionAllocatePool:\r
-        if (AllocInfo->Buffer == (PHYSICAL_ADDRESS) (UINTN) Buffer) {\r
+        if (AllocInfo->Buffer == (PHYSICAL_ADDRESS)(UINTN)Buffer) {\r
           return AllocInfoData;\r
         }\r
+\r
         break;\r
       default:\r
         ASSERT (FALSE);\r
@@ -956,11 +1164,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
 CoreUpdateProfileFree (\r
   IN PHYSICAL_ADDRESS       CallerAddress,\r
   IN MEMORY_PROFILE_ACTION  Action,\r
@@ -978,114 +1187,147 @@ CoreUpdateProfileFree (
   MEMORY_PROFILE_DRIVER_INFO_DATA  *ThisDriverInfoData;\r
   MEMORY_PROFILE_ALLOC_INFO_DATA   *AllocInfoData;\r
   UINTN                            ProfileMemoryIndex;\r
+  MEMORY_PROFILE_ACTION            BasicAction;\r
+  BOOLEAN                          Found;\r
+\r
+  BasicAction = Action & MEMORY_PROFILE_ACTION_BASIC_MASK;\r
 \r
   ContextData = GetMemoryProfileContext ();\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 records,\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
+      {\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
+\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 may 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
-      CoreUpdateProfileAllocate (\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
-      CoreUpdateProfileAllocate (\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 CoreInternalFreePool() that will not update profile for this FreePool action.\r
-  //\r
-  CoreInternalFreePool (AllocInfoData, NULL);\r
+    RemoveEntryList (&AllocInfoData->Link);\r
+\r
+    if (BasicAction == MemoryProfileActionFreePages) {\r
+      if (AllocInfo->Buffer != (PHYSICAL_ADDRESS)(UINTN)Buffer) {\r
+        CoreUpdateProfileAllocate (\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
 \r
-  return TRUE;\r
+      if (AllocInfo->Buffer + AllocInfo->Size != ((PHYSICAL_ADDRESS)(UINTN)Buffer + Size)) {\r
+        CoreUpdateProfileAllocate (\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
+    //\r
+    // Use CoreInternalFreePool() that will not update profile for this FreePool action.\r
+    //\r
+    CoreInternalFreePool (AllocInfoData, NULL);\r
+  } while (TRUE);\r
 }\r
 \r
 /**\r
@@ -1094,62 +1336,94 @@ CoreUpdateProfileFree (
   @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
 CoreUpdateProfile (\r
   IN PHYSICAL_ADDRESS       CallerAddress,\r
   IN MEMORY_PROFILE_ACTION  Action,\r
   IN EFI_MEMORY_TYPE        MemoryType,\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
-  MEMORY_PROFILE_CONTEXT_DATA   *ContextData;\r
+  EFI_STATUS                   Status;\r
+  MEMORY_PROFILE_CONTEXT_DATA  *ContextData;\r
+  MEMORY_PROFILE_ACTION        BasicAction;\r
 \r
   if (!IS_UEFI_MEMORY_PROFILE_ENABLED) {\r
-    return FALSE;\r
+    return EFI_UNSUPPORTED;\r
   }\r
 \r
-  if (!mMemoryProfileRecordingStatus) {\r
-    return FALSE;\r
+  if (mMemoryProfileGettingStatus) {\r
+    return EFI_ACCESS_DENIED;\r
+  }\r
+\r
+  if (!mMemoryProfileRecordingEnable) {\r
+    return EFI_ABORTED;\r
   }\r
 \r
   //\r
-  // Only record limited MemoryType.\r
+  // Get the basic action to know how to process the record\r
   //\r
-  if (!CoreNeedRecordProfile (MemoryType)) {\r
-    return FALSE;\r
+  BasicAction = Action & MEMORY_PROFILE_ACTION_BASIC_MASK;\r
+\r
+  //\r
+  // EfiMaxMemoryType means the MemoryType is unknown.\r
+  //\r
+  if (MemoryType != EfiMaxMemoryType) {\r
+    //\r
+    // Only record limited MemoryType.\r
+    //\r
+    if (!CoreNeedRecordProfile (MemoryType)) {\r
+      return EFI_UNSUPPORTED;\r
+    }\r
   }\r
 \r
   ContextData = GetMemoryProfileContext ();\r
   if (ContextData == NULL) {\r
-    return FALSE;\r
+    return EFI_UNSUPPORTED;\r
   }\r
 \r
-  switch (Action) {\r
+  CoreAcquireMemoryProfileLock ();\r
+  switch (BasicAction) {\r
     case MemoryProfileActionAllocatePages:\r
-      CoreUpdateProfileAllocate (CallerAddress, Action, MemoryType, Size, Buffer);\r
+      Status = CoreUpdateProfileAllocate (CallerAddress, Action, MemoryType, Size, Buffer, ActionString);\r
       break;\r
     case MemoryProfileActionFreePages:\r
-      CoreUpdateProfileFree (CallerAddress, Action, Size, Buffer);\r
+      Status = CoreUpdateProfileFree (CallerAddress, Action, Size, Buffer);\r
       break;\r
     case MemoryProfileActionAllocatePool:\r
-      CoreUpdateProfileAllocate (CallerAddress, Action, MemoryType, Size, Buffer);\r
+      Status = CoreUpdateProfileAllocate (CallerAddress, Action, MemoryType, Size, Buffer, ActionString);\r
       break;\r
     case MemoryProfileActionFreePool:\r
-      CoreUpdateProfileFree (CallerAddress, Action, 0, Buffer);\r
+      Status = CoreUpdateProfileFree (CallerAddress, Action, 0, Buffer);\r
       break;\r
     default:\r
       ASSERT (FALSE);\r
+      Status = EFI_UNSUPPORTED;\r
       break;\r
   }\r
-  return TRUE;\r
+\r
+  CoreReleaseMemoryProfileLock ();\r
+\r
+  return Status;\r
 }\r
 \r
 ////////////////////\r
@@ -1165,12 +1439,14 @@ MemoryProfileGetDataSize (
   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
-\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
 \r
   ContextData = GetMemoryProfileContext ();\r
   if (ContextData == NULL) {\r
@@ -1178,19 +1454,33 @@ MemoryProfileGetDataSize (
   }\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
        DriverLink != DriverInfoList;\r
-       DriverLink = DriverLink->ForwardLink) {\r
+       DriverLink = DriverLink->ForwardLink)\r
+  {\r
     DriverInfoData = CR (\r
                        DriverLink,\r
                        MEMORY_PROFILE_DRIVER_INFO_DATA,\r
                        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
+    {\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
   return TotalSize;\r
@@ -1204,33 +1494,36 @@ MemoryProfileGetDataSize (
 **/\r
 VOID\r
 MemoryProfileCopyData (\r
-  IN VOID   *ProfileBuffer\r
+  IN VOID  *ProfileBuffer\r
   )\r
 {\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
-  LIST_ENTRY                        *DriverInfoList;\r
-  LIST_ENTRY                        *DriverLink;\r
-  LIST_ENTRY                        *AllocInfoList;\r
-  LIST_ENTRY                        *AllocLink;\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
+  LIST_ENTRY                       *DriverInfoList;\r
+  LIST_ENTRY                       *DriverLink;\r
+  LIST_ENTRY                       *AllocInfoList;\r
+  LIST_ENTRY                       *AllocLink;\r
+  UINTN                            PdbSize;\r
+  UINTN                            ActionStringSize;\r
 \r
   ContextData = GetMemoryProfileContext ();\r
   if (ContextData == NULL) {\r
-    return ;\r
+    return;\r
   }\r
 \r
   Context = ProfileBuffer;\r
   CopyMem (Context, &ContextData->Context, sizeof (MEMORY_PROFILE_CONTEXT));\r
-  DriverInfo = (MEMORY_PROFILE_DRIVER_INFO *) (Context + 1);\r
+  DriverInfo = (MEMORY_PROFILE_DRIVER_INFO *)(Context + 1);\r
 \r
   DriverInfoList = ContextData->DriverInfoList;\r
   for (DriverLink = DriverInfoList->ForwardLink;\r
        DriverLink != DriverInfoList;\r
-       DriverLink = DriverLink->ForwardLink) {\r
+       DriverLink = DriverLink->ForwardLink)\r
+  {\r
     DriverInfoData = CR (\r
                        DriverLink,\r
                        MEMORY_PROFILE_DRIVER_INFO_DATA,\r
@@ -1238,12 +1531,18 @@ MemoryProfileCopyData (
                        MEMORY_PROFILE_DRIVER_INFO_SIGNATURE\r
                        );\r
     CopyMem (DriverInfo, &DriverInfoData->DriverInfo, sizeof (MEMORY_PROFILE_DRIVER_INFO));\r
-    AllocInfo = (MEMORY_PROFILE_ALLOC_INFO *) (DriverInfo + 1);\r
+    if (DriverInfo->PdbStringOffset != 0) {\r
+      PdbSize = AsciiStrSize (DriverInfoData->PdbString);\r
+      CopyMem ((VOID *)((UINTN)DriverInfo + DriverInfo->PdbStringOffset), DriverInfoData->PdbString, PdbSize);\r
+    }\r
+\r
+    AllocInfo = (MEMORY_PROFILE_ALLOC_INFO *)((UINTN)DriverInfo + DriverInfo->Header.Length);\r
 \r
     AllocInfoList = DriverInfoData->AllocInfoList;\r
     for (AllocLink = AllocInfoList->ForwardLink;\r
          AllocLink != AllocInfoList;\r
-         AllocLink = AllocLink->ForwardLink) {\r
+         AllocLink = AllocLink->ForwardLink)\r
+    {\r
       AllocInfoData = CR (\r
                         AllocLink,\r
                         MEMORY_PROFILE_ALLOC_INFO_DATA,\r
@@ -1251,10 +1550,15 @@ MemoryProfileCopyData (
                         MEMORY_PROFILE_ALLOC_INFO_SIGNATURE\r
                         );\r
       CopyMem (AllocInfo, &AllocInfoData->AllocInfo, sizeof (MEMORY_PROFILE_ALLOC_INFO));\r
-      AllocInfo += 1;\r
+      if (AllocInfo->ActionStringOffset != 0) {\r
+        ActionStringSize = AsciiStrSize (AllocInfoData->ActionString);\r
+        CopyMem ((VOID *)((UINTN)AllocInfo + AllocInfo->ActionStringOffset), AllocInfoData->ActionString, ActionStringSize);\r
+      }\r
+\r
+      AllocInfo = (MEMORY_PROFILE_ALLOC_INFO *)((UINTN)AllocInfo + AllocInfo->Header.Length);\r
     }\r
 \r
-    DriverInfo = (MEMORY_PROFILE_DRIVER_INFO *) ((UINTN) (DriverInfo + 1) + sizeof (MEMORY_PROFILE_ALLOC_INFO) * (UINTN) DriverInfo->AllocRecordCount);\r
+    DriverInfo = (MEMORY_PROFILE_DRIVER_INFO *)AllocInfo;\r
   }\r
 }\r
 \r
@@ -1265,9 +1569,10 @@ MemoryProfileCopyData (
   @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
+\r
   @return EFI_SUCCESS               Get the memory profile data successfully.\r
-  @return EFI_BUFFER_TO_SMALL       The ProfileSize is too small for the resulting data. \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
@@ -1276,33 +1581,33 @@ EFIAPI
 ProfileProtocolGetData (\r
   IN     EDKII_MEMORY_PROFILE_PROTOCOL  *This,\r
   IN OUT UINT64                         *ProfileSize,\r
-     OUT VOID                           *ProfileBuffer\r
+  OUT VOID                              *ProfileBuffer\r
   )\r
 {\r
-  UINTN                                 Size;\r
-  MEMORY_PROFILE_CONTEXT_DATA           *ContextData;\r
-  BOOLEAN                               MemoryProfileRecordingStatus;\r
+  UINTN                        Size;\r
+  MEMORY_PROFILE_CONTEXT_DATA  *ContextData;\r
+  BOOLEAN                      MemoryProfileGettingStatus;\r
 \r
   ContextData = GetMemoryProfileContext ();\r
   if (ContextData == NULL) {\r
     return EFI_UNSUPPORTED;\r
   }\r
 \r
-  MemoryProfileRecordingStatus = mMemoryProfileRecordingStatus;\r
-  mMemoryProfileRecordingStatus = FALSE;\r
+  MemoryProfileGettingStatus  = mMemoryProfileGettingStatus;\r
+  mMemoryProfileGettingStatus = TRUE;\r
 \r
   Size = MemoryProfileGetDataSize ();\r
 \r
   if (*ProfileSize < Size) {\r
-    *ProfileSize = Size;\r
-    mMemoryProfileRecordingStatus = MemoryProfileRecordingStatus;\r
+    *ProfileSize                = Size;\r
+    mMemoryProfileGettingStatus = MemoryProfileGettingStatus;\r
     return EFI_BUFFER_TOO_SMALL;\r
   }\r
 \r
   *ProfileSize = Size;\r
   MemoryProfileCopyData (ProfileBuffer);\r
 \r
-  mMemoryProfileRecordingStatus = MemoryProfileRecordingStatus;\r
+  mMemoryProfileGettingStatus = MemoryProfileGettingStatus;\r
   return EFI_SUCCESS;\r
 }\r
 \r
@@ -1315,8 +1620,10 @@ ProfileProtocolGetData (
   @param[in] ImageSize          Image size.\r
   @param[in] FileType           File type of the image.\r
 \r
-  @return EFI_SUCCESS           Register success.\r
-  @return EFI_OUT_OF_RESOURCE   No enough resource for this register.\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
@@ -1329,20 +1636,20 @@ ProfileProtocolRegisterImage (
   IN EFI_FV_FILETYPE                FileType\r
   )\r
 {\r
-  EFI_STATUS                        Status;\r
-  LOADED_IMAGE_PRIVATE_DATA         DriverEntry;\r
-  VOID                              *EntryPointInImage;\r
+  EFI_STATUS                 Status;\r
+  LOADED_IMAGE_PRIVATE_DATA  DriverEntry;\r
+  VOID                       *EntryPointInImage;\r
 \r
   ZeroMem (&DriverEntry, sizeof (DriverEntry));\r
-  DriverEntry.Info.FilePath = FilePath;\r
+  DriverEntry.Info.FilePath             = FilePath;\r
   DriverEntry.ImageContext.ImageAddress = ImageBase;\r
-  DriverEntry.ImageContext.ImageSize = ImageSize;\r
-  Status = InternalPeCoffGetEntryPoint ((VOID *) (UINTN) ImageBase, &EntryPointInImage);\r
+  DriverEntry.ImageContext.ImageSize    = ImageSize;\r
+  Status                                = InternalPeCoffGetEntryPoint ((VOID *)(UINTN)ImageBase, &EntryPointInImage);\r
   ASSERT_EFI_ERROR (Status);\r
-  DriverEntry.ImageContext.EntryPoint = (PHYSICAL_ADDRESS) (UINTN) EntryPointInImage;\r
-  DriverEntry.ImageContext.ImageType = InternalPeCoffGetSubsystem ((VOID *) (UINTN) ImageBase);\r
+  DriverEntry.ImageContext.EntryPoint = (PHYSICAL_ADDRESS)(UINTN)EntryPointInImage;\r
+  DriverEntry.ImageContext.ImageType  = InternalPeCoffGetSubsystem ((VOID *)(UINTN)ImageBase);\r
 \r
-  return RegisterMemoryProfileImage (&DriverEntry, FileType) ? EFI_SUCCESS: EFI_OUT_OF_RESOURCES;\r
+  return RegisterMemoryProfileImage (&DriverEntry, FileType);\r
 }\r
 \r
 /**\r
@@ -1353,7 +1660,9 @@ ProfileProtocolRegisterImage (
   @param[in] ImageBase          Image base address.\r
   @param[in] ImageSize          Image size.\r
 \r
-  @return EFI_SUCCESS           Unregister success.\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
@@ -1366,19 +1675,118 @@ ProfileProtocolUnregisterImage (
   IN UINT64                         ImageSize\r
   )\r
 {\r
-  EFI_STATUS                        Status;\r
-  LOADED_IMAGE_PRIVATE_DATA         DriverEntry;\r
-  VOID                              *EntryPointInImage;\r
+  EFI_STATUS                 Status;\r
+  LOADED_IMAGE_PRIVATE_DATA  DriverEntry;\r
+  VOID                       *EntryPointInImage;\r
 \r
   ZeroMem (&DriverEntry, sizeof (DriverEntry));\r
-  DriverEntry.Info.FilePath = FilePath;\r
+  DriverEntry.Info.FilePath             = FilePath;\r
   DriverEntry.ImageContext.ImageAddress = ImageBase;\r
-  DriverEntry.ImageContext.ImageSize = ImageSize;\r
-  Status = InternalPeCoffGetEntryPoint ((VOID *) (UINTN) ImageBase, &EntryPointInImage);\r
+  DriverEntry.ImageContext.ImageSize    = ImageSize;\r
+  Status                                = InternalPeCoffGetEntryPoint ((VOID *)(UINTN)ImageBase, &EntryPointInImage);\r
   ASSERT_EFI_ERROR (Status);\r
-  DriverEntry.ImageContext.EntryPoint = (PHYSICAL_ADDRESS) (UINTN) EntryPointInImage;\r
+  DriverEntry.ImageContext.EntryPoint = (PHYSICAL_ADDRESS)(UINTN)EntryPointInImage;\r
+\r
+  return UnregisterMemoryProfileImage (&DriverEntry);\r
+}\r
 \r
-  return UnregisterMemoryProfileImage (&DriverEntry) ? EFI_SUCCESS: EFI_NOT_FOUND;\r
+/**\r
+  Get memory profile recording state.\r
+\r
+  @param[in]  This              The EDKII_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
+ProfileProtocolGetRecordingState (\r
+  IN EDKII_MEMORY_PROFILE_PROTOCOL  *This,\r
+  OUT BOOLEAN                       *RecordingState\r
+  )\r
+{\r
+  MEMORY_PROFILE_CONTEXT_DATA  *ContextData;\r
+\r
+  ContextData = GetMemoryProfileContext ();\r
+  if (ContextData == NULL) {\r
+    return EFI_UNSUPPORTED;\r
+  }\r
+\r
+  if (RecordingState == NULL) {\r
+    return EFI_INVALID_PARAMETER;\r
+  }\r
+\r
+  *RecordingState = mMemoryProfileRecordingEnable;\r
+  return EFI_SUCCESS;\r
+}\r
+\r
+/**\r
+  Set memory profile recording state.\r
+\r
+  @param[in] This               The EDKII_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
+ProfileProtocolSetRecordingState (\r
+  IN EDKII_MEMORY_PROFILE_PROTOCOL  *This,\r
+  IN BOOLEAN                        RecordingState\r
+  )\r
+{\r
+  MEMORY_PROFILE_CONTEXT_DATA  *ContextData;\r
+\r
+  ContextData = GetMemoryProfileContext ();\r
+  if (ContextData == NULL) {\r
+    return EFI_UNSUPPORTED;\r
+  }\r
+\r
+  mMemoryProfileRecordingEnable = RecordingState;\r
+  return EFI_SUCCESS;\r
+}\r
+\r
+/**\r
+  Record memory profile of multilevel caller.\r
+\r
+  @param[in] This               The EDKII_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
+ProfileProtocolRecord (\r
+  IN EDKII_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 CoreUpdateProfile (CallerAddress, Action, MemoryType, Size, Buffer, ActionString);\r
 }\r
 \r
 ////////////////////\r