]> git.proxmox.com Git - mirror_edk2.git/blobdiff - MdeModulePkg/Core/Dxe/Mem/MemoryProfileRecord.c
MdeModulePkg DxeCore/PiSmmCore: Add UEFI memory and SMRAM profile support.
[mirror_edk2.git] / MdeModulePkg / Core / Dxe / Mem / MemoryProfileRecord.c
diff --git a/MdeModulePkg/Core/Dxe/Mem/MemoryProfileRecord.c b/MdeModulePkg/Core/Dxe/Mem/MemoryProfileRecord.c
new file mode 100644 (file)
index 0000000..1602fdb
--- /dev/null
@@ -0,0 +1,1377 @@
+/** @file\r
+  Support routines for UEFI memory profile.\r
+\r
+  Copyright (c) 2014, 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
+\r
+**/\r
+\r
+#include "DxeMain.h"\r
+\r
+#define IS_UEFI_MEMORY_PROFILE_ENABLED ((PcdGet8 (PcdMemoryProfilePropertyMask) & BIT0) != 0)\r
+\r
+typedef struct {\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
+  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
+} 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
+  MEMORY_PROFILE_CONTEXT_SIGNATURE,\r
+  {\r
+    {\r
+      MEMORY_PROFILE_CONTEXT_SIGNATURE,\r
+      sizeof (MEMORY_PROFILE_CONTEXT),\r
+      MEMORY_PROFILE_CONTEXT_REVISION\r
+    },\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
+\r
+BOOLEAN mMemoryProfileRecordingStatus = FALSE;\r
+\r
+/**\r
+  Get memory profile data.\r
+\r
+  @param[in]      This              The EDKII_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_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
+ProfileProtocolGetData (\r
+  IN     EDKII_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_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 success.\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
+  );\r
+\r
+/**\r
+  Unregister image from memory profile.\r
+\r
+  @param[in] This               The EDKII_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 success.\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
+  );\r
+\r
+EDKII_MEMORY_PROFILE_PROTOCOL mProfileProtocol = {\r
+  ProfileProtocolGetData,\r
+  ProfileProtocolRegisterImage,\r
+  ProfileProtocolUnregisterImage\r
+};\r
+\r
+/**\r
+  Return memory profile context.\r
+\r
+  @return Memory profile context.\r
+\r
+**/\r
+MEMORY_PROFILE_CONTEXT_DATA *\r
+GetMemoryProfileContext (\r
+  VOID\r
+  )\r
+{\r
+  return mMemoryProfileContextPtr;\r
+}\r
+\r
+/**\r
+  Retrieves the magic value from the PE/COFF header.\r
+\r
+  @param Hdr    The buffer in which to return the PE32, PE32+, or TE header.\r
+\r
+  @return EFI_IMAGE_NT_OPTIONAL_HDR32_MAGIC - Image is PE32\r
+  @return EFI_IMAGE_NT_OPTIONAL_HDR64_MAGIC - Image is PE32+\r
+\r
+**/\r
+UINT16\r
+InternalPeCoffGetPeHeaderMagicValue (\r
+  IN  EFI_IMAGE_OPTIONAL_HEADER_PTR_UNION  Hdr\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
+}\r
+\r
+/**\r
+  Retrieves and returns the Subsystem of a PE/COFF image that has been loaded into system memory.\r
+  If Pe32Data is NULL, then ASSERT().\r
+\r
+  @param Pe32Data   The pointer to the PE/COFF image that is loaded in system memory.\r
+\r
+  @return The Subsystem of the PE/COFF image.\r
+\r
+**/\r
+UINT16\r
+InternalPeCoffGetSubsystem (\r
+  IN VOID  *Pe32Data\r
+  )\r
+{\r
+  EFI_IMAGE_OPTIONAL_HEADER_PTR_UNION  Hdr;\r
+  EFI_IMAGE_DOS_HEADER                 *DosHdr;\r
+  UINT16                               Magic;\r
+\r
+  ASSERT (Pe32Data != NULL);\r
+\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
+  } 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
+  }\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
+    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
+      return Hdr.Pe32Plus->OptionalHeader.Subsystem;\r
+    }\r
+  }\r
+\r
+  return 0x0000;\r
+}\r
+\r
+/**\r
+  Retrieves and returns a pointer to the entry point to a PE/COFF image that has been loaded\r
+  into system memory with the PE/COFF Loader Library functions.\r
+\r
+  Retrieves the entry point to the PE/COFF image specified by Pe32Data and returns this entry\r
+  point in EntryPoint.  If the entry point could not be retrieved from the PE/COFF image, then\r
+  return RETURN_INVALID_PARAMETER.  Otherwise return RETURN_SUCCESS.\r
+  If Pe32Data is NULL, then ASSERT().\r
+  If EntryPoint is NULL, then ASSERT().\r
+\r
+  @param  Pe32Data                  The pointer to the PE/COFF image that is loaded in system memory.\r
+  @param  EntryPoint                The pointer to entry point to the PE/COFF image to return.\r
+\r
+  @retval RETURN_SUCCESS            EntryPoint was returned.\r
+  @retval RETURN_INVALID_PARAMETER  The entry point could not be found in the PE/COFF image.\r
+\r
+**/\r
+RETURN_STATUS\r
+InternalPeCoffGetEntryPoint (\r
+  IN  VOID  *Pe32Data,\r
+  OUT VOID  **EntryPoint\r
+  )\r
+{\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
+  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
+  } 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
+  }\r
+\r
+  //\r
+  // Calculate the entry point relative to the start of the image.\r
+  // 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
+    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
+    return RETURN_SUCCESS;\r
+  }\r
+\r
+  return RETURN_UNSUPPORTED;\r
+}\r
+\r
+/**\r
+  Build driver info.\r
+\r
+  @param ContextData    Memory profile context.\r
+  @param FileName       File name of the image.\r
+  @param ImageBase      Image base address.\r
+  @param ImageSize      Image size.\r
+  @param EntryPoint     Entry point of the image.\r
+  @param ImageSubsystem Image subsystem of the image.\r
+  @param FileType       File type of the image.\r
+\r
+  @return Pointer to memory profile driver info.\r
+\r
+**/\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
+  )\r
+{\r
+  EFI_STATUS                        Status;\r
+  MEMORY_PROFILE_DRIVER_INFO        *DriverInfo;\r
+  MEMORY_PROFILE_DRIVER_INFO_DATA   *DriverInfoData;\r
+  VOID                              *EntryPointInImage;\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
+             );\r
+  if (EFI_ERROR (Status)) {\r
+    return NULL;\r
+  }\r
+\r
+  ZeroMem (DriverInfoData, sizeof (*DriverInfoData));\r
+\r
+  DriverInfo = &DriverInfoData->DriverInfo;\r
+  DriverInfoData->Signature = MEMORY_PROFILE_DRIVER_INFO_SIGNATURE;\r
+  DriverInfo->Header.Signature = MEMORY_PROFILE_DRIVER_INFO_SIGNATURE;\r
+  DriverInfo->Header.Length = sizeof (MEMORY_PROFILE_DRIVER_INFO);\r
+  DriverInfo->Header.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
+  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
+    ASSERT_EFI_ERROR (Status);\r
+    DriverInfo->ImageBase = ImageBase + EntryPoint - (PHYSICAL_ADDRESS) (UINTN) EntryPointInImage;\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->AllocRecordCount = 0;\r
+\r
+  InsertTailList (ContextData->DriverInfoList, &DriverInfoData->Link);\r
+  ContextData->Context.ImageCount ++;\r
+  ContextData->Context.TotalImageSize += DriverInfo->ImageSize;\r
+\r
+  return DriverInfoData;\r
+}\r
+\r
+/**\r
+  Register DXE Core to memory profile.\r
+\r
+  @param HobStart       The start address of the HOB.\r
+  @param ContextData    Memory profile context.\r
+\r
+  @retval TRUE      Register success.\r
+  @retval FALSE     Register fail.\r
+\r
+**/\r
+BOOLEAN\r
+RegisterDxeCore (\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
+\r
+  ASSERT (ContextData != NULL);\r
+\r
+  //\r
+  // Searching for image hob\r
+  //\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
+      // Find Dxe Core HOB\r
+      //\r
+      break;\r
+    }\r
+    DxeCoreHob.Raw = GET_NEXT_HOB (DxeCoreHob);\r
+  }\r
+  ASSERT (DxeCoreHob.Raw != NULL);\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
+                     EFI_FV_FILETYPE_DXE_CORE\r
+                     );\r
+  if (DriverInfoData == NULL) {\r
+    return FALSE;\r
+  }\r
+\r
+  return TRUE;\r
+}\r
+\r
+/**\r
+  Initialize memory profile.\r
+\r
+  @param HobStart   The start address of the HOB.\r
+\r
+**/\r
+VOID\r
+MemoryProfileInit (\r
+  IN VOID   *HobStart\r
+  )\r
+{\r
+  MEMORY_PROFILE_CONTEXT_DATA   *ContextData;\r
+\r
+  if (!IS_UEFI_MEMORY_PROFILE_ENABLED) {\r
+    return;\r
+  }\r
+\r
+  ContextData = GetMemoryProfileContext ();\r
+  if (ContextData != NULL) {\r
+    return;\r
+  }\r
+\r
+  mMemoryProfileRecordingStatus = TRUE;\r
+  mMemoryProfileContextPtr = &mMemoryProfileContext;\r
+\r
+  RegisterDxeCore (HobStart, &mMemoryProfileContext);\r
+\r
+  DEBUG ((EFI_D_INFO, "MemoryProfileInit MemoryProfileContext - 0x%x\n", &mMemoryProfileContext));\r
+}\r
+\r
+/**\r
+  Install memory profile protocol.\r
+\r
+**/\r
+VOID\r
+MemoryProfileInstallProtocol (\r
+  VOID\r
+  )\r
+{\r
+  EFI_HANDLE    Handle;\r
+  EFI_STATUS    Status;\r
+\r
+  if (!IS_UEFI_MEMORY_PROFILE_ENABLED) {\r
+    return;\r
+  }\r
+\r
+  Handle = NULL;\r
+  Status = CoreInstallMultipleProtocolInterfaces (\r
+             &Handle,\r
+             &gEdkiiMemoryProfileGuid,\r
+             &mProfileProtocol,\r
+             NULL\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
+  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
+  return FileName;\r
+}\r
+\r
+/**\r
+  Register image to memory profile.\r
+\r
+  @param DriverEntry    Image info.\r
+  @param FileType       Image file type.\r
+\r
+  @retval TRUE          Register success.\r
+  @retval FALSE         Register fail.\r
+\r
+**/\r
+BOOLEAN\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
+\r
+  if (!IS_UEFI_MEMORY_PROFILE_ENABLED) {\r
+    return FALSE;\r
+  }\r
+\r
+  ContextData = GetMemoryProfileContext ();\r
+  if (ContextData == NULL) {\r
+    return FALSE;\r
+  }\r
+\r
+  DriverInfoData = BuildDriverInfo (\r
+                     ContextData,\r
+                     GetFileNameFromFilePath (DriverEntry->Info.FilePath),\r
+                     DriverEntry->ImageContext.ImageAddress,\r
+                     DriverEntry->ImageContext.ImageSize,\r
+                     DriverEntry->ImageContext.EntryPoint,\r
+                     DriverEntry->ImageContext.ImageType,\r
+                     FileType\r
+                     );\r
+  if (DriverInfoData == NULL) {\r
+    return FALSE;\r
+  }\r
+\r
+  return TRUE;\r
+}\r
+\r
+/**\r
+  Search image from memory profile.\r
+\r
+  @param ContextData    Memory profile context.\r
+  @param FileName       Image file name.\r
+  @param Address        Image Address.\r
+\r
+  @return Pointer to memory profile driver info.\r
+\r
+**/\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
+  )\r
+{\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
+    DriverInfoData = CR (\r
+                       DriverLink,\r
+                       MEMORY_PROFILE_DRIVER_INFO_DATA,\r
+                       Link,\r
+                       MEMORY_PROFILE_DRIVER_INFO_SIGNATURE\r
+                       );\r
+    DriverInfo = &DriverInfoData->DriverInfo;\r
+    if ((CompareGuid (&DriverInfo->FileName, FileName)) &&\r
+        (Address >= DriverInfo->ImageBase) &&\r
+        (Address < (DriverInfo->ImageBase + DriverInfo->ImageSize))) {\r
+      return DriverInfoData;\r
+    }\r
+  }\r
+\r
+  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
+\r
+  @param ContextData    Memory profile context.\r
+  @param Address        Image or Function address.\r
+\r
+  @return Pointer to memory profile driver info.\r
+\r
+**/\r
+MEMORY_PROFILE_DRIVER_INFO_DATA *\r
+GetMemoryProfileDriverInfoFromAddress (\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
+\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
+    DriverInfo = &DriverInfoData->DriverInfo;\r
+    if ((Address >= DriverInfo->ImageBase) &&\r
+        (Address < (DriverInfo->ImageBase + DriverInfo->ImageSize))) {\r
+      return DriverInfoData;\r
+    }\r
+  }\r
+\r
+  //\r
+  // Should never come here.\r
+  //\r
+  return FindDummyImage (ContextData);\r
+}\r
+\r
+/**\r
+  Unregister image from memory profile.\r
+\r
+  @param DriverEntry    Image info.\r
+\r
+  @retval TRUE          Unregister success.\r
+  @retval FALSE         Unregister fail.\r
+\r
+**/\r
+BOOLEAN\r
+UnregisterMemoryProfileImage (\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
+\r
+  if (!IS_UEFI_MEMORY_PROFILE_ENABLED) {\r
+    return FALSE;\r
+  }\r
+\r
+  ContextData = GetMemoryProfileContext ();\r
+  if (ContextData == NULL) {\r
+    return FALSE;\r
+  }\r
+\r
+  DriverInfoData = NULL;\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
+    ASSERT_EFI_ERROR (Status);\r
+    ImageAddress = ImageAddress + (UINTN) DriverEntry->ImageContext.EntryPoint - (UINTN) EntryPointInImage;\r
+  }\r
+  if (FileName != NULL) {\r
+    DriverInfoData = GetMemoryProfileDriverInfoByFileNameAndAddress (ContextData, FileName, ImageAddress);\r
+  }\r
+  if (DriverInfoData == NULL) {\r
+    DriverInfoData = GetMemoryProfileDriverInfoFromAddress (ContextData, ImageAddress);\r
+  }\r
+  if (DriverInfoData == NULL) {\r
+    return FALSE;\r
+  }\r
+\r
+  ContextData->Context.TotalImageSize -= DriverInfoData->DriverInfo.ImageSize;\r
+\r
+  DriverInfoData->DriverInfo.ImageBase = 0;\r
+  DriverInfoData->DriverInfo.ImageSize = 0;\r
+\r
+  if (DriverInfoData->DriverInfo.PeakUsage == 0) {\r
+    ContextData->Context.ImageCount --;\r
+    RemoveEntryList (&DriverInfoData->Link);\r
+    //\r
+    // Use CoreInternalFreePool() that will not update profile for this FreePool action.\r
+    //\r
+    CoreInternalFreePool (DriverInfoData);\r
+  }\r
+\r
+  return TRUE;\r
+}\r
+\r
+/**\r
+  Return if this memory type needs to be recorded into memory profile.\r
+  If BIOS memory type (0 ~ EfiMaxMemoryType), it checks bit (1 << MemoryType).\r
+  If OS memory type (0x80000000 ~ 0xFFFFFFFF), it checks bit63 - 0x8000000000000000.\r
+\r
+  @param MemoryType     Memory type.\r
+\r
+  @retval TRUE          This memory type need to be recorded.\r
+  @retval FALSE         This memory type need not to be recorded.\r
+\r
+**/\r
+BOOLEAN\r
+CoreNeedRecordProfile (\r
+  IN EFI_MEMORY_TYPE    MemoryType\r
+  )\r
+{\r
+  UINT64 TestBit;\r
+\r
+  if ((UINT32) MemoryType >= 0x80000000) {\r
+    TestBit = BIT63;\r
+  } else {\r
+    TestBit = LShiftU64 (1, MemoryType);\r
+  }\r
+\r
+  if ((PcdGet64 (PcdMemoryProfileMemoryType) & TestBit) != 0) {\r
+    return TRUE;\r
+  } else {\r
+    return FALSE;\r
+  }\r
+}\r
+\r
+/**\r
+  Convert EFI memory type to profile memory index. The rule is:\r
+  If BIOS memory type (0 ~ EfiMaxMemoryType), ProfileMemoryIndex = MemoryType.\r
+  If OS memory type (0x80000000 ~ 0xFFFFFFFF), ProfileMemoryIndex = EfiMaxMemoryType.\r
+\r
+  @param MemoryType     Memory type.\r
+\r
+  @return EFI memory type as profile memory index.\r
+\r
+**/\r
+EFI_MEMORY_TYPE\r
+GetProfileMemoryIndex (\r
+  IN EFI_MEMORY_TYPE    MemoryType\r
+  )\r
+{\r
+  if ((UINT32) MemoryType >= 0x80000000) {\r
+    return EfiMaxMemoryType;\r
+  } else {\r
+    return MemoryType;\r
+  }\r
+}\r
+\r
+/**\r
+  Update memory profile Allocate information.\r
+\r
+  @param CallerAddress  Address of caller who call Allocate.\r
+  @param Action         This Allocate action.\r
+  @param MemoryType     Memory type.\r
+  @param Size           Buffer size.\r
+  @param Buffer         Buffer address.\r
+\r
+  @retval TRUE          Profile udpate success.\r
+  @retval FALSE         Profile update fail.\r
+\r
+**/\r
+BOOLEAN\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
+  )\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_DATA       *ContextData;\r
+  MEMORY_PROFILE_DRIVER_INFO_DATA   *DriverInfoData;\r
+  MEMORY_PROFILE_ALLOC_INFO_DATA    *AllocInfoData;\r
+  EFI_MEMORY_TYPE                   ProfileMemoryIndex;\r
+\r
+  ContextData = GetMemoryProfileContext ();\r
+  if (ContextData == NULL) {\r
+    return FALSE;\r
+  }\r
+\r
+  DriverInfoData = GetMemoryProfileDriverInfoFromAddress (ContextData, CallerAddress);\r
+  ASSERT (DriverInfoData != NULL);\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
+  if (EFI_ERROR (Status)) {\r
+    return FALSE;\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.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
+  ProfileMemoryIndex = GetProfileMemoryIndex (MemoryType);\r
+\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
+  Context->SequenceCount ++;\r
+\r
+  return TRUE;\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 Size               Buffer size\r
+  @param Buffer             Buffer address\r
+\r
+  @return Pointer to memory profile alloc info.\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
+  )\r
+{\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
+    AllocInfoData = CR (\r
+                      AllocLink,\r
+                      MEMORY_PROFILE_ALLOC_INFO_DATA,\r
+                      Link,\r
+                      MEMORY_PROFILE_ALLOC_INFO_SIGNATURE\r
+                      );\r
+    AllocInfo = &AllocInfoData->AllocInfo;\r
+    if (AllocInfo->Action != Action) {\r
+      continue;\r
+    }\r
+    switch (Action) {\r
+      case MemoryProfileActionAllocatePages:\r
+        if ((AllocInfo->Buffer <= (PHYSICAL_ADDRESS) (UINTN) Buffer) &&\r
+            ((AllocInfo->Buffer + AllocInfo->Size) >= ((PHYSICAL_ADDRESS) (UINTN) Buffer + Size))) {\r
+          return AllocInfoData;\r
+        }\r
+        break;\r
+      case MemoryProfileActionAllocatePool:\r
+        if (AllocInfo->Buffer == (PHYSICAL_ADDRESS) (UINTN) Buffer) {\r
+          return AllocInfoData;\r
+        }\r
+        break;\r
+      default:\r
+        ASSERT (FALSE);\r
+        break;\r
+    }\r
+  }\r
+\r
+  return NULL;\r
+}\r
+\r
+/**\r
+  Update memory profile Free information.\r
+\r
+  @param CallerAddress  Address of caller who call Free.\r
+  @param Action         This Free action.\r
+  @param Size           Buffer size.\r
+  @param Buffer         Buffer address.\r
+\r
+  @retval TRUE          Profile udpate success.\r
+  @retval FALSE         Profile update fail.\r
+\r
+**/\r
+BOOLEAN\r
+CoreUpdateProfileFree (\r
+  IN PHYSICAL_ADDRESS       CallerAddress,\r
+  IN MEMORY_PROFILE_ACTION  Action,\r
+  IN UINTN                  Size,\r
+  IN VOID                   *Buffer\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
+  LIST_ENTRY                       *DriverLink;\r
+  LIST_ENTRY                       *DriverInfoList;\r
+  MEMORY_PROFILE_DRIVER_INFO_DATA  *ThisDriverInfoData;\r
+  MEMORY_PROFILE_ALLOC_INFO_DATA   *AllocInfoData;\r
+  EFI_MEMORY_TYPE                  ProfileMemoryIndex;\r
+\r
+  ContextData = GetMemoryProfileContext ();\r
+  if (ContextData == NULL) {\r
+    return FALSE;\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
+        case MemoryProfileActionFreePages:\r
+          AllocInfoData = GetMemoryProfileAllocInfoFromAddress (ThisDriverInfoData, MemoryProfileActionAllocatePages, Size, Buffer);\r
+          break;\r
+        case MemoryProfileActionFreePool:\r
+          AllocInfoData = GetMemoryProfileAllocInfoFromAddress (ThisDriverInfoData, MemoryProfileActionAllocatePool, 0, Buffer);\r
+          break;\r
+        default:\r
+          ASSERT (FALSE);\r
+          AllocInfoData = NULL;\r
+          break;\r
+      }\r
+      if (AllocInfoData != NULL) {\r
+        DriverInfoData = ThisDriverInfoData;\r
+        break;\r
+      }\r
+    }\r
+\r
+    if (AllocInfoData == NULL) {\r
+      //\r
+      // 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
+      //\r
+      return FALSE;\r
+    }\r
+  }\r
+\r
+  Context = &ContextData->Context;\r
+  DriverInfo = &DriverInfoData->DriverInfo;\r
+  AllocInfo = &AllocInfoData->AllocInfo;\r
+\r
+  ProfileMemoryIndex = GetProfileMemoryIndex (AllocInfo->MemoryType);\r
+\r
+  Context->CurrentTotalUsage -= AllocInfo->Size;\r
+  Context->CurrentTotalUsageByType[ProfileMemoryIndex] -= AllocInfo->Size;\r
+\r
+  DriverInfo->CurrentUsage -= AllocInfo->Size;\r
+  DriverInfo->CurrentUsageByType[ProfileMemoryIndex] -= AllocInfo->Size;\r
+  DriverInfo->AllocRecordCount --;\r
+\r
+  RemoveEntryList (&AllocInfoData->Link);\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
+    }\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);\r
+\r
+  return TRUE;\r
+}\r
+\r
+/**\r
+  Update memory profile information.\r
+\r
+  @param CallerAddress  Address of caller who call Allocate or Free.\r
+  @param Action         This Allocate or Free action.\r
+  @param MemoryType     Memory type.\r
+  @param Size           Buffer size.\r
+  @param Buffer         Buffer address.\r
+\r
+  @retval TRUE          Profile udpate success.\r
+  @retval FALSE         Profile update fail.\r
+\r
+**/\r
+BOOLEAN\r
+CoreUpdateProfile (\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
+  )\r
+{\r
+  MEMORY_PROFILE_CONTEXT_DATA   *ContextData;\r
+\r
+  if (!IS_UEFI_MEMORY_PROFILE_ENABLED) {\r
+    return FALSE;\r
+  }\r
+\r
+  if (!mMemoryProfileRecordingStatus) {\r
+    return FALSE;\r
+  }\r
+\r
+  //\r
+  // Free operations have no memory type information, so skip the check.\r
+  //\r
+  if ((Action == MemoryProfileActionAllocatePages) || (Action == MemoryProfileActionAllocatePool)) {\r
+    //\r
+    // Only record limited MemoryType.\r
+    //\r
+    if (!CoreNeedRecordProfile (MemoryType)) {\r
+      return FALSE;\r
+    }\r
+  }\r
+\r
+  ContextData = GetMemoryProfileContext ();\r
+  if (ContextData == NULL) {\r
+    return FALSE;\r
+  }\r
+\r
+  switch (Action) {\r
+    case MemoryProfileActionAllocatePages:\r
+      CoreUpdateProfileAllocate (CallerAddress, Action, MemoryType, Size, Buffer);\r
+      break;\r
+    case MemoryProfileActionFreePages:\r
+      CoreUpdateProfileFree (CallerAddress, Action, Size, Buffer);\r
+      break;\r
+    case MemoryProfileActionAllocatePool:\r
+      CoreUpdateProfileAllocate (CallerAddress, Action, MemoryType, Size, Buffer);\r
+      break;\r
+    case MemoryProfileActionFreePool:\r
+      CoreUpdateProfileFree (CallerAddress, Action, 0, Buffer);\r
+      break;\r
+    default:\r
+      ASSERT (FALSE);\r
+      break;\r
+  }\r
+  return TRUE;\r
+}\r
+\r
+////////////////////\r
+\r
+/**\r
+  Get memory profile data size.\r
+\r
+  @return Memory profile data size.\r
+\r
+**/\r
+UINTN\r
+MemoryProfileGetDataSize (\r
+  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
+\r
+  ContextData = GetMemoryProfileContext ();\r
+  if (ContextData == NULL) {\r
+    return 0;\r
+  }\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
+    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
+  }\r
+\r
+  return TotalSize;\r
+}\r
+\r
+/**\r
+  Copy memory profile data.\r
+\r
+  @param ProfileBuffer  The buffer to hold memory profile data.\r
+\r
+**/\r
+VOID\r
+MemoryProfileCopyData (\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
+\r
+  ContextData = GetMemoryProfileContext ();\r
+  if (ContextData == NULL) {\r
+    return ;\r
+  }\r
+\r
+  Context = ProfileBuffer;\r
+  CopyMem (Context, &ContextData->Context, sizeof (MEMORY_PROFILE_CONTEXT));\r
+  DriverInfo = (MEMORY_PROFILE_DRIVER_INFO *) (Context + 1);\r
+\r
+  DriverInfoList = ContextData->DriverInfoList;\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
+    CopyMem (DriverInfo, &DriverInfoData->DriverInfo, sizeof (MEMORY_PROFILE_DRIVER_INFO));\r
+    AllocInfo = (MEMORY_PROFILE_ALLOC_INFO *) (DriverInfo + 1);\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
+      CopyMem (AllocInfo, &AllocInfoData->AllocInfo, sizeof (MEMORY_PROFILE_ALLOC_INFO));\r
+      AllocInfo += 1;\r
+    }\r
+\r
+    DriverInfo = (MEMORY_PROFILE_DRIVER_INFO *) ((UINTN) (DriverInfo + 1) + sizeof (MEMORY_PROFILE_ALLOC_INFO) * (UINTN) DriverInfo->AllocRecordCount);\r
+  }\r
+}\r
+\r
+/**\r
+  Get memory profile data.\r
+\r
+  @param[in]      This              The EDKII_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_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
+ProfileProtocolGetData (\r
+  IN     EDKII_MEMORY_PROFILE_PROTOCOL  *This,\r
+  IN OUT UINT64                         *ProfileSize,\r
+     OUT VOID                           *ProfileBuffer\r
+  )\r
+{\r
+  UINTN                                 Size;\r
+  MEMORY_PROFILE_CONTEXT_DATA           *ContextData;\r
+  BOOLEAN                               MemoryProfileRecordingStatus;\r
+\r
+  ContextData = GetMemoryProfileContext ();\r
+  if (ContextData == NULL) {\r
+    return EFI_UNSUPPORTED;\r
+  }\r
+\r
+  MemoryProfileRecordingStatus = mMemoryProfileRecordingStatus;\r
+  mMemoryProfileRecordingStatus = FALSE;\r
+\r
+  Size = MemoryProfileGetDataSize ();\r
+\r
+  if (*ProfileSize < Size) {\r
+    *ProfileSize = Size;\r
+    mMemoryProfileRecordingStatus = MemoryProfileRecordingStatus;\r
+    return EFI_BUFFER_TOO_SMALL;\r
+  }\r
+\r
+  *ProfileSize = Size;\r
+  MemoryProfileCopyData (ProfileBuffer);\r
+\r
+  mMemoryProfileRecordingStatus = MemoryProfileRecordingStatus;\r
+  return EFI_SUCCESS;\r
+}\r
+\r
+/**\r
+  Register image to memory profile.\r
+\r
+  @param[in] This               The EDKII_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 success.\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
+  )\r
+{\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.ImageContext.ImageAddress = ImageBase;\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
+\r
+  return RegisterMemoryProfileImage (&DriverEntry, FileType) ? EFI_SUCCESS: EFI_OUT_OF_RESOURCES;\r
+}\r
+\r
+/**\r
+  Unregister image from memory profile.\r
+\r
+  @param[in] This               The EDKII_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 success.\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
+  )\r
+{\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.ImageContext.ImageAddress = ImageBase;\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
+\r
+  return UnregisterMemoryProfileImage (&DriverEntry) ? EFI_SUCCESS: EFI_NOT_FOUND;\r
+}\r
+\r
+////////////////////\r