MdeModulePkg/Core/Dxe/Mem/MemoryProfileRecord.c

   1 /** @file
   2   Support routines for UEFI memory profile.
   3
   4   Copyright (c) 2014 - 2016, Intel Corporation. All rights reserved.<BR>
   5   This program and the accompanying materials
   6   are licensed and made available under the terms and conditions of the BSD License
   7   which accompanies this distribution.  The full text of the license may be found at
   8   http://opensource.org/licenses/bsd-license.php.
   9
  10   THE PROGRAM IS DISTRIBUTED UNDER THE BSD LICENSE ON AN "AS IS" BASIS,
  11   WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
  12
  13 **/
  14
  15 #include "DxeMain.h"
  16 #include "Imem.h"
  17
  18 #define IS_UEFI_MEMORY_PROFILE_ENABLED ((PcdGet8 (PcdMemoryProfilePropertyMask) & BIT0) != 0)
  19
  20 #define GET_OCCUPIED_SIZE(ActualSize, Alignment) \
  21   ((ActualSize) + (((Alignment) - ((ActualSize) & ((Alignment) - 1))) & ((Alignment) - 1)))
  22
  23 typedef struct {
  24   UINT32                        Signature;
  25   MEMORY_PROFILE_CONTEXT        Context;
  26   LIST_ENTRY                    *DriverInfoList;
  27 } MEMORY_PROFILE_CONTEXT_DATA;
  28
  29 typedef struct {
  30   UINT32                        Signature;
  31   MEMORY_PROFILE_DRIVER_INFO    DriverInfo;
  32   LIST_ENTRY                    *AllocInfoList;
  33   CHAR8                         *PdbString;
  34   LIST_ENTRY                    Link;
  35 } MEMORY_PROFILE_DRIVER_INFO_DATA;
  36
  37 typedef struct {
  38   UINT32                        Signature;
  39   MEMORY_PROFILE_ALLOC_INFO     AllocInfo;
  40   CHAR8                         *ActionString;
  41   LIST_ENTRY                    Link;
  42 } MEMORY_PROFILE_ALLOC_INFO_DATA;
  43
  44
  45 GLOBAL_REMOVE_IF_UNREFERENCED LIST_ENTRY  mImageQueue = INITIALIZE_LIST_HEAD_VARIABLE (mImageQueue);
  46 GLOBAL_REMOVE_IF_UNREFERENCED MEMORY_PROFILE_CONTEXT_DATA mMemoryProfileContext = {
  47   MEMORY_PROFILE_CONTEXT_SIGNATURE,
  48   {
  49     {
  50       MEMORY_PROFILE_CONTEXT_SIGNATURE,
  51       sizeof (MEMORY_PROFILE_CONTEXT),
  52       MEMORY_PROFILE_CONTEXT_REVISION
  53     },
  54     0,
  55     0,
  56     {0},
  57     {0},
  58     0,
  59     0,
  60     0
  61   },
  62   &mImageQueue,
  63 };
  64 GLOBAL_REMOVE_IF_UNREFERENCED MEMORY_PROFILE_CONTEXT_DATA *mMemoryProfileContextPtr = NULL;
  65
  66 EFI_LOCK mMemoryProfileLock = EFI_INITIALIZE_LOCK_VARIABLE (TPL_NOTIFY);
  67 BOOLEAN mMemoryProfileGettingStatus = FALSE;
  68 BOOLEAN mMemoryProfileRecordingEnable = MEMORY_PROFILE_RECORDING_DISABLE;
  69 EFI_DEVICE_PATH_PROTOCOL *mMemoryProfileDriverPath;
  70 UINTN                    mMemoryProfileDriverPathSize;
  71
  72 /**
  73   Get memory profile data.
  74
  75   @param[in]      This              The EDKII_MEMORY_PROFILE_PROTOCOL instance.
  76   @param[in, out] ProfileSize       On entry, points to the size in bytes of the ProfileBuffer.
  77                                     On return, points to the size of the data returned in ProfileBuffer.
  78   @param[out]     ProfileBuffer     Profile buffer.
  79
  80   @return EFI_SUCCESS               Get the memory profile data successfully.
  81   @return EFI_UNSUPPORTED           Memory profile is unsupported.
  82   @return EFI_BUFFER_TO_SMALL       The ProfileSize is too small for the resulting data.
  83                                     ProfileSize is updated with the size required.
  84
  85 **/
  86 EFI_STATUS
  87 EFIAPI
  88 ProfileProtocolGetData (
  89   IN     EDKII_MEMORY_PROFILE_PROTOCOL  *This,
  90   IN OUT UINT64                         *ProfileSize,
  91      OUT VOID                           *ProfileBuffer
  92   );
  93
  94 /**
  95   Register image to memory profile.
  96
  97   @param[in] This               The EDKII_MEMORY_PROFILE_PROTOCOL instance.
  98   @param[in] FilePath           File path of the image.
  99   @param[in] ImageBase          Image base address.
 100   @param[in] ImageSize          Image size.
 101   @param[in] FileType           File type of the image.
 102
 103   @return EFI_SUCCESS           Register successfully.
 104   @return EFI_UNSUPPORTED       Memory profile is unsupported,
 105                                 or memory profile for the image is not required.
 106   @return EFI_OUT_OF_RESOURCE   No enough resource for this register.
 107
 108 **/
 109 EFI_STATUS
 110 EFIAPI
 111 ProfileProtocolRegisterImage (
 112   IN EDKII_MEMORY_PROFILE_PROTOCOL      *This,
 113   IN EFI_DEVICE_PATH_PROTOCOL           *FilePath,
 114   IN PHYSICAL_ADDRESS                   ImageBase,
 115   IN UINT64                             ImageSize,
 116   IN EFI_FV_FILETYPE                    FileType
 117   );
 118
 119 /**
 120   Unregister image from memory profile.
 121
 122   @param[in] This               The EDKII_MEMORY_PROFILE_PROTOCOL instance.
 123   @param[in] FilePath           File path of the image.
 124   @param[in] ImageBase          Image base address.
 125   @param[in] ImageSize          Image size.
 126
 127   @return EFI_SUCCESS           Unregister successfully.
 128   @return EFI_UNSUPPORTED       Memory profile is unsupported,
 129                                 or memory profile for the image is not required.
 130   @return EFI_NOT_FOUND         The image is not found.
 131
 132 **/
 133 EFI_STATUS
 134 EFIAPI
 135 ProfileProtocolUnregisterImage (
 136   IN EDKII_MEMORY_PROFILE_PROTOCOL      *This,
 137   IN EFI_DEVICE_PATH_PROTOCOL           *FilePath,
 138   IN PHYSICAL_ADDRESS                   ImageBase,
 139   IN UINT64                             ImageSize
 140   );
 141
 142 /**
 143   Get memory profile recording state.
 144
 145   @param[in]  This              The EDKII_MEMORY_PROFILE_PROTOCOL instance.
 146   @param[out] RecordingState    Recording state.
 147
 148   @return EFI_SUCCESS           Memory profile recording state is returned.
 149   @return EFI_UNSUPPORTED       Memory profile is unsupported.
 150   @return EFI_INVALID_PARAMETER RecordingState is NULL.
 151
 152 **/
 153 EFI_STATUS
 154 EFIAPI
 155 ProfileProtocolGetRecordingState (
 156   IN EDKII_MEMORY_PROFILE_PROTOCOL      *This,
 157   OUT BOOLEAN                           *RecordingState
 158   );
 159
 160 /**
 161   Set memory profile recording state.
 162
 163   @param[in] This               The EDKII_MEMORY_PROFILE_PROTOCOL instance.
 164   @param[in] RecordingState     Recording state.
 165
 166   @return EFI_SUCCESS           Set memory profile recording state successfully.
 167   @return EFI_UNSUPPORTED       Memory profile is unsupported.
 168
 169 **/
 170 EFI_STATUS
 171 EFIAPI
 172 ProfileProtocolSetRecordingState (
 173   IN EDKII_MEMORY_PROFILE_PROTOCOL      *This,
 174   IN BOOLEAN                            RecordingState
 175   );
 176
 177 /**
 178   Record memory profile of multilevel caller.
 179
 180   @param[in] This               The EDKII_MEMORY_PROFILE_PROTOCOL instance.
 181   @param[in] CallerAddress      Address of caller.
 182   @param[in] Action             Memory profile action.
 183   @param[in] MemoryType         Memory type.
 184                                 EfiMaxMemoryType means the MemoryType is unknown.
 185   @param[in] Buffer             Buffer address.
 186   @param[in] Size               Buffer size.
 187   @param[in] ActionString       String for memory profile action.
 188                                 Only needed for user defined allocate action.
 189
 190   @return EFI_SUCCESS           Memory profile is updated.
 191   @return EFI_UNSUPPORTED       Memory profile is unsupported,
 192                                 or memory profile for the image is not required,
 193                                 or memory profile for the memory type is not required.
 194   @return EFI_ACCESS_DENIED     It is during memory profile data getting.
 195   @return EFI_ABORTED           Memory profile recording is not enabled.
 196   @return EFI_OUT_OF_RESOURCES  No enough resource to update memory profile for allocate action.
 197   @return EFI_NOT_FOUND         No matched allocate info found for free action.
 198
 199 **/
 200 EFI_STATUS
 201 EFIAPI
 202 ProfileProtocolRecord (
 203   IN EDKII_MEMORY_PROFILE_PROTOCOL      *This,
 204   IN PHYSICAL_ADDRESS                   CallerAddress,
 205   IN MEMORY_PROFILE_ACTION              Action,
 206   IN EFI_MEMORY_TYPE                    MemoryType,
 207   IN VOID                               *Buffer,
 208   IN UINTN                              Size,
 209   IN CHAR8                              *ActionString OPTIONAL
 210   );
 211
 212 EDKII_MEMORY_PROFILE_PROTOCOL mProfileProtocol = {
 213   ProfileProtocolGetData,
 214   ProfileProtocolRegisterImage,
 215   ProfileProtocolUnregisterImage,
 216   ProfileProtocolGetRecordingState,
 217   ProfileProtocolSetRecordingState,
 218   ProfileProtocolRecord,
 219 };
 220
 221 /**
 222   Acquire lock on mMemoryProfileLock.
 223 **/
 224 VOID
 225 CoreAcquireMemoryProfileLock (
 226   VOID
 227   )
 228 {
 229   CoreAcquireLock (&mMemoryProfileLock);
 230 }
 231
 232 /**
 233   Release lock on mMemoryProfileLock.
 234 **/
 235 VOID
 236 CoreReleaseMemoryProfileLock (
 237   VOID
 238   )
 239 {
 240   CoreReleaseLock (&mMemoryProfileLock);
 241 }
 242
 243 /**
 244   Return memory profile context.
 245
 246   @return Memory profile context.
 247
 248 **/
 249 MEMORY_PROFILE_CONTEXT_DATA *
 250 GetMemoryProfileContext (
 251   VOID
 252   )
 253 {
 254   return mMemoryProfileContextPtr;
 255 }
 256
 257 /**
 258   Retrieves the magic value from the PE/COFF header.
 259
 260   @param Hdr    The buffer in which to return the PE32, PE32+, or TE header.
 261
 262   @return EFI_IMAGE_NT_OPTIONAL_HDR32_MAGIC - Image is PE32
 263   @return EFI_IMAGE_NT_OPTIONAL_HDR64_MAGIC - Image is PE32+
 264
 265 **/
 266 UINT16
 267 InternalPeCoffGetPeHeaderMagicValue (
 268   IN  EFI_IMAGE_OPTIONAL_HEADER_PTR_UNION  Hdr
 269   )
 270 {
 271   //
 272   // NOTE: Some versions of Linux ELILO for Itanium have an incorrect magic value
 273   //       in the PE/COFF Header.  If the MachineType is Itanium(IA64) and the
 274   //       Magic value in the OptionalHeader is  EFI_IMAGE_NT_OPTIONAL_HDR32_MAGIC
 275   //       then override the returned value to EFI_IMAGE_NT_OPTIONAL_HDR64_MAGIC
 276   //
 277   if (Hdr.Pe32->FileHeader.Machine == IMAGE_FILE_MACHINE_IA64 && Hdr.Pe32->OptionalHeader.Magic == EFI_IMAGE_NT_OPTIONAL_HDR32_MAGIC) {
 278     return EFI_IMAGE_NT_OPTIONAL_HDR64_MAGIC;
 279   }
 280   //
 281   // Return the magic value from the PC/COFF Optional Header
 282   //
 283   return Hdr.Pe32->OptionalHeader.Magic;
 284 }
 285
 286 /**
 287   Retrieves and returns the Subsystem of a PE/COFF image that has been loaded into system memory.
 288   If Pe32Data is NULL, then ASSERT().
 289
 290   @param Pe32Data   The pointer to the PE/COFF image that is loaded in system memory.
 291
 292   @return The Subsystem of the PE/COFF image.
 293
 294 **/
 295 UINT16
 296 InternalPeCoffGetSubsystem (
 297   IN VOID  *Pe32Data
 298   )
 299 {
 300   EFI_IMAGE_OPTIONAL_HEADER_PTR_UNION  Hdr;
 301   EFI_IMAGE_DOS_HEADER                 *DosHdr;
 302   UINT16                               Magic;
 303
 304   ASSERT (Pe32Data != NULL);
 305
 306   DosHdr = (EFI_IMAGE_DOS_HEADER *) Pe32Data;
 307   if (DosHdr->e_magic == EFI_IMAGE_DOS_SIGNATURE) {
 308     //
 309     // DOS image header is present, so read the PE header after the DOS image header.
 310     //
 311     Hdr.Pe32 = (EFI_IMAGE_NT_HEADERS32 *) ((UINTN) Pe32Data + (UINTN) ((DosHdr->e_lfanew) & 0x0ffff));
 312   } else {
 313     //
 314     // DOS image header is not present, so PE header is at the image base.
 315     //
 316     Hdr.Pe32 = (EFI_IMAGE_NT_HEADERS32 *) Pe32Data;
 317   }
 318
 319   if (Hdr.Te->Signature == EFI_TE_IMAGE_HEADER_SIGNATURE) {
 320     return Hdr.Te->Subsystem;
 321   } else if (Hdr.Pe32->Signature == EFI_IMAGE_NT_SIGNATURE)  {
 322     Magic = InternalPeCoffGetPeHeaderMagicValue (Hdr);
 323     if (Magic == EFI_IMAGE_NT_OPTIONAL_HDR32_MAGIC) {
 324       return Hdr.Pe32->OptionalHeader.Subsystem;
 325     } else if (Magic == EFI_IMAGE_NT_OPTIONAL_HDR64_MAGIC) {
 326       return Hdr.Pe32Plus->OptionalHeader.Subsystem;
 327     }
 328   }
 329
 330   return 0x0000;
 331 }
 332
 333 /**
 334   Retrieves and returns a pointer to the entry point to a PE/COFF image that has been loaded
 335   into system memory with the PE/COFF Loader Library functions.
 336
 337   Retrieves the entry point to the PE/COFF image specified by Pe32Data and returns this entry
 338   point in EntryPoint.  If the entry point could not be retrieved from the PE/COFF image, then
 339   return RETURN_INVALID_PARAMETER.  Otherwise return RETURN_SUCCESS.
 340   If Pe32Data is NULL, then ASSERT().
 341   If EntryPoint is NULL, then ASSERT().
 342
 343   @param  Pe32Data                  The pointer to the PE/COFF image that is loaded in system memory.
 344   @param  EntryPoint                The pointer to entry point to the PE/COFF image to return.
 345
 346   @retval RETURN_SUCCESS            EntryPoint was returned.
 347   @retval RETURN_INVALID_PARAMETER  The entry point could not be found in the PE/COFF image.
 348
 349 **/
 350 RETURN_STATUS
 351 InternalPeCoffGetEntryPoint (
 352   IN  VOID  *Pe32Data,
 353   OUT VOID  **EntryPoint
 354   )
 355 {
 356   EFI_IMAGE_DOS_HEADER                  *DosHdr;
 357   EFI_IMAGE_OPTIONAL_HEADER_PTR_UNION   Hdr;
 358
 359   ASSERT (Pe32Data   != NULL);
 360   ASSERT (EntryPoint != NULL);
 361
 362   DosHdr = (EFI_IMAGE_DOS_HEADER *) Pe32Data;
 363   if (DosHdr->e_magic == EFI_IMAGE_DOS_SIGNATURE) {
 364     //
 365     // DOS image header is present, so read the PE header after the DOS image header.
 366     //
 367     Hdr.Pe32 = (EFI_IMAGE_NT_HEADERS32 *) ((UINTN) Pe32Data + (UINTN) ((DosHdr->e_lfanew) & 0x0ffff));
 368   } else {
 369     //
 370     // DOS image header is not present, so PE header is at the image base.
 371     //
 372     Hdr.Pe32 = (EFI_IMAGE_NT_HEADERS32 *) Pe32Data;
 373   }
 374
 375   //
 376   // Calculate the entry point relative to the start of the image.
 377   // AddressOfEntryPoint is common for PE32 & PE32+
 378   //
 379   if (Hdr.Te->Signature == EFI_TE_IMAGE_HEADER_SIGNATURE) {
 380     *EntryPoint = (VOID *) ((UINTN) Pe32Data + (UINTN) (Hdr.Te->AddressOfEntryPoint & 0x0ffffffff) + sizeof (EFI_TE_IMAGE_HEADER) - Hdr.Te->StrippedSize);
 381     return RETURN_SUCCESS;
 382   } else if (Hdr.Pe32->Signature == EFI_IMAGE_NT_SIGNATURE) {
 383     *EntryPoint = (VOID *) ((UINTN) Pe32Data + (UINTN) (Hdr.Pe32->OptionalHeader.AddressOfEntryPoint & 0x0ffffffff));
 384     return RETURN_SUCCESS;
 385   }
 386
 387   return RETURN_UNSUPPORTED;
 388 }
 389
 390 /**
 391   Build driver info.
 392
 393   @param ContextData    Memory profile context.
 394   @param FileName       File name of the image.
 395   @param ImageBase      Image base address.
 396   @param ImageSize      Image size.
 397   @param EntryPoint     Entry point of the image.
 398   @param ImageSubsystem Image subsystem of the image.
 399   @param FileType       File type of the image.
 400
 401   @return Pointer to memory profile driver info.
 402
 403 **/
 404 MEMORY_PROFILE_DRIVER_INFO_DATA *
 405 BuildDriverInfo (
 406   IN MEMORY_PROFILE_CONTEXT_DATA    *ContextData,
 407   IN EFI_GUID                       *FileName,
 408   IN PHYSICAL_ADDRESS               ImageBase,
 409   IN UINT64                         ImageSize,
 410   IN PHYSICAL_ADDRESS               EntryPoint,
 411   IN UINT16                         ImageSubsystem,
 412   IN EFI_FV_FILETYPE                FileType
 413   )
 414 {
 415   EFI_STATUS                        Status;
 416   MEMORY_PROFILE_DRIVER_INFO        *DriverInfo;
 417   MEMORY_PROFILE_DRIVER_INFO_DATA   *DriverInfoData;
 418   VOID                              *EntryPointInImage;
 419   CHAR8                             *PdbString;
 420   UINTN                             PdbSize;
 421   UINTN                             PdbOccupiedSize;
 422
 423   PdbSize = 0;
 424   PdbOccupiedSize = 0;
 425   PdbString = NULL;
 426   if (ImageBase != 0) {
 427     PdbString = PeCoffLoaderGetPdbPointer ((VOID*) (UINTN) ImageBase);
 428     if (PdbString != NULL) {
 429       PdbSize = AsciiStrSize (PdbString);
 430       PdbOccupiedSize = GET_OCCUPIED_SIZE (PdbSize, sizeof (UINT64));
 431     }
 432   }
 433
 434   //
 435   // Use CoreInternalAllocatePool() that will not update profile for this AllocatePool action.
 436   //
 437   Status = CoreInternalAllocatePool (
 438              EfiBootServicesData,
 439              sizeof (*DriverInfoData) + sizeof (LIST_ENTRY) + PdbSize,
 440              (VOID **) &DriverInfoData
 441              );
 442   if (EFI_ERROR (Status)) {
 443     return NULL;
 444   }
 445
 446   ZeroMem (DriverInfoData, sizeof (*DriverInfoData));
 447
 448   DriverInfo = &DriverInfoData->DriverInfo;
 449   DriverInfoData->Signature = MEMORY_PROFILE_DRIVER_INFO_SIGNATURE;
 450   DriverInfo->Header.Signature = MEMORY_PROFILE_DRIVER_INFO_SIGNATURE;
 451   DriverInfo->Header.Length = (UINT16) (sizeof (MEMORY_PROFILE_DRIVER_INFO) + PdbOccupiedSize);
 452   DriverInfo->Header.Revision = MEMORY_PROFILE_DRIVER_INFO_REVISION;
 453   if (FileName != NULL) {
 454     CopyMem (&DriverInfo->FileName, FileName, sizeof (EFI_GUID));
 455   }
 456   DriverInfo->ImageBase = ImageBase;
 457   DriverInfo->ImageSize = ImageSize;
 458   DriverInfo->EntryPoint = EntryPoint;
 459   DriverInfo->ImageSubsystem = ImageSubsystem;
 460   if ((EntryPoint != 0) && ((EntryPoint < ImageBase) || (EntryPoint >= (ImageBase + ImageSize)))) {
 461     //
 462     // If the EntryPoint is not in the range of image buffer, it should come from emulation environment.
 463     // So patch ImageBuffer here to align the EntryPoint.
 464     //
 465     Status = InternalPeCoffGetEntryPoint ((VOID *) (UINTN) ImageBase, &EntryPointInImage);
 466     ASSERT_EFI_ERROR (Status);
 467     DriverInfo->ImageBase = ImageBase + EntryPoint - (PHYSICAL_ADDRESS) (UINTN) EntryPointInImage;
 468   }
 469   DriverInfo->FileType = FileType;
 470   DriverInfoData->AllocInfoList = (LIST_ENTRY *) (DriverInfoData + 1);
 471   InitializeListHead (DriverInfoData->AllocInfoList);
 472   DriverInfo->CurrentUsage = 0;
 473   DriverInfo->PeakUsage = 0;
 474   DriverInfo->AllocRecordCount = 0;
 475   if (PdbSize != 0) {
 476     DriverInfo->PdbStringOffset = (UINT16) sizeof (MEMORY_PROFILE_DRIVER_INFO);
 477     DriverInfoData->PdbString = (CHAR8 *) (DriverInfoData->AllocInfoList + 1);
 478     CopyMem (DriverInfoData->PdbString, PdbString, PdbSize);
 479   } else {
 480     DriverInfo->PdbStringOffset = 0;
 481     DriverInfoData->PdbString = NULL;
 482   }
 483
 484   InsertTailList (ContextData->DriverInfoList, &DriverInfoData->Link);
 485   ContextData->Context.ImageCount ++;
 486   ContextData->Context.TotalImageSize += DriverInfo->ImageSize;
 487
 488   return DriverInfoData;
 489 }
 490
 491 /**
 492   Return if record for this driver is needed..
 493
 494   @param DriverFilePath     Driver file path.
 495
 496   @retval TRUE              Record for this driver is needed.
 497   @retval FALSE             Record for this driver is not needed.
 498
 499 **/
 500 BOOLEAN
 501 NeedRecordThisDriver (
 502   IN EFI_DEVICE_PATH_PROTOCOL       *DriverFilePath
 503   )
 504 {
 505   EFI_DEVICE_PATH_PROTOCOL          *TmpDevicePath;
 506   EFI_DEVICE_PATH_PROTOCOL          *DevicePathInstance;
 507   UINTN                             DevicePathSize;
 508   UINTN                             FilePathSize;
 509
 510   if (!IsDevicePathValid (mMemoryProfileDriverPath, mMemoryProfileDriverPathSize)) {
 511     //
 512     // Invalid Device Path means record all.
 513     //
 514     return TRUE;
 515   }
 516
 517   //
 518   // Record FilePath without END node.
 519   //
 520   FilePathSize = GetDevicePathSize (DriverFilePath) - sizeof(EFI_DEVICE_PATH_PROTOCOL);
 521
 522   DevicePathInstance = mMemoryProfileDriverPath;
 523   do {
 524     //
 525     // Find END node (it might be END_ENTIRE or END_INSTANCE).
 526     //
 527     TmpDevicePath = DevicePathInstance;
 528     while (!IsDevicePathEndType (TmpDevicePath)) {
 529       TmpDevicePath = NextDevicePathNode (TmpDevicePath);
 530     }
 531
 532     //
 533     // Do not compare END node.
 534     //
 535     DevicePathSize = (UINTN)TmpDevicePath - (UINTN)DevicePathInstance;
 536     if ((FilePathSize == DevicePathSize) &&
 537         (CompareMem (DriverFilePath, DevicePathInstance, DevicePathSize) == 0)) {
 538       return TRUE;
 539     }
 540
 541     //
 542     // Get next instance.
 543     //
 544     DevicePathInstance = (EFI_DEVICE_PATH_PROTOCOL *)((UINTN)DevicePathInstance + DevicePathSize + DevicePathNodeLength(TmpDevicePath));
 545   } while (DevicePathSubType (TmpDevicePath) != END_ENTIRE_DEVICE_PATH_SUBTYPE);
 546
 547   return FALSE;
 548 }
 549
 550 /**
 551   Register DXE Core to memory profile.
 552
 553   @param HobStart       The start address of the HOB.
 554   @param ContextData    Memory profile context.
 555
 556   @retval TRUE      Register success.
 557   @retval FALSE     Register fail.
 558
 559 **/
 560 BOOLEAN
 561 RegisterDxeCore (
 562   IN VOID                           *HobStart,
 563   IN MEMORY_PROFILE_CONTEXT_DATA    *ContextData
 564   )
 565 {
 566   EFI_PEI_HOB_POINTERS              DxeCoreHob;
 567   MEMORY_PROFILE_DRIVER_INFO_DATA   *DriverInfoData;
 568   PHYSICAL_ADDRESS                  ImageBase;
 569   UINT8                             TempBuffer[sizeof(MEDIA_FW_VOL_FILEPATH_DEVICE_PATH) + sizeof(EFI_DEVICE_PATH_PROTOCOL)];
 570   MEDIA_FW_VOL_FILEPATH_DEVICE_PATH *FilePath;
 571
 572   ASSERT (ContextData != NULL);
 573
 574   //
 575   // Searching for image hob
 576   //
 577   DxeCoreHob.Raw          = HobStart;
 578   while ((DxeCoreHob.Raw = GetNextHob (EFI_HOB_TYPE_MEMORY_ALLOCATION, DxeCoreHob.Raw)) != NULL) {
 579     if (CompareGuid (&DxeCoreHob.MemoryAllocationModule->MemoryAllocationHeader.Name, &gEfiHobMemoryAllocModuleGuid)) {
 580       //
 581       // Find Dxe Core HOB
 582       //
 583       break;
 584     }
 585     DxeCoreHob.Raw = GET_NEXT_HOB (DxeCoreHob);
 586   }
 587   ASSERT (DxeCoreHob.Raw != NULL);
 588
 589   FilePath = (MEDIA_FW_VOL_FILEPATH_DEVICE_PATH *) TempBuffer;
 590   EfiInitializeFwVolDevicepathNode (FilePath, &DxeCoreHob.MemoryAllocationModule->ModuleName);
 591   SetDevicePathEndNode (FilePath + 1);
 592
 593   if (!NeedRecordThisDriver ((EFI_DEVICE_PATH_PROTOCOL *) FilePath)) {
 594     return FALSE;
 595   }
 596
 597   ImageBase = DxeCoreHob.MemoryAllocationModule->MemoryAllocationHeader.MemoryBaseAddress;
 598   DriverInfoData = BuildDriverInfo (
 599                      ContextData,
 600                      &DxeCoreHob.MemoryAllocationModule->ModuleName,
 601                      ImageBase,
 602                      DxeCoreHob.MemoryAllocationModule->MemoryAllocationHeader.MemoryLength,
 603                      DxeCoreHob.MemoryAllocationModule->EntryPoint,
 604                      InternalPeCoffGetSubsystem ((VOID *) (UINTN) ImageBase),
 605                      EFI_FV_FILETYPE_DXE_CORE
 606                      );
 607   if (DriverInfoData == NULL) {
 608     return FALSE;
 609   }
 610
 611   return TRUE;
 612 }
 613
 614 /**
 615   Initialize memory profile.
 616
 617   @param HobStart   The start address of the HOB.
 618
 619 **/
 620 VOID
 621 MemoryProfileInit (
 622   IN VOID   *HobStart
 623   )
 624 {
 625   MEMORY_PROFILE_CONTEXT_DATA   *ContextData;
 626
 627   if (!IS_UEFI_MEMORY_PROFILE_ENABLED) {
 628     return;
 629   }
 630
 631   ContextData = GetMemoryProfileContext ();
 632   if (ContextData != NULL) {
 633     return;
 634   }
 635
 636   mMemoryProfileGettingStatus = FALSE;
 637   if ((PcdGet8 (PcdMemoryProfilePropertyMask) & BIT7) != 0) {
 638     mMemoryProfileRecordingEnable = MEMORY_PROFILE_RECORDING_DISABLE;
 639   } else {
 640     mMemoryProfileRecordingEnable = MEMORY_PROFILE_RECORDING_ENABLE;
 641   }
 642   mMemoryProfileDriverPathSize = PcdGetSize (PcdMemoryProfileDriverPath);
 643   mMemoryProfileDriverPath = AllocateCopyPool (mMemoryProfileDriverPathSize, PcdGetPtr (PcdMemoryProfileDriverPath));
 644   mMemoryProfileContextPtr = &mMemoryProfileContext;
 645
 646   RegisterDxeCore (HobStart, &mMemoryProfileContext);
 647
 648   DEBUG ((EFI_D_INFO, "MemoryProfileInit MemoryProfileContext - 0x%x\n", &mMemoryProfileContext));
 649 }
 650
 651 /**
 652   Install memory profile protocol.
 653
 654 **/
 655 VOID
 656 MemoryProfileInstallProtocol (
 657   VOID
 658   )
 659 {
 660   EFI_HANDLE    Handle;
 661   EFI_STATUS    Status;
 662
 663   if (!IS_UEFI_MEMORY_PROFILE_ENABLED) {
 664     return;
 665   }
 666
 667   Handle = NULL;
 668   Status = CoreInstallMultipleProtocolInterfaces (
 669              &Handle,
 670              &gEdkiiMemoryProfileGuid,
 671              &mProfileProtocol,
 672              NULL
 673              );
 674   ASSERT_EFI_ERROR (Status);
 675 }
 676
 677 /**
 678   Get the GUID file name from the file path.
 679
 680   @param FilePath  File path.
 681
 682   @return The GUID file name from the file path.
 683
 684 **/
 685 EFI_GUID *
 686 GetFileNameFromFilePath (
 687   IN EFI_DEVICE_PATH_PROTOCOL   *FilePath
 688   )
 689 {
 690   MEDIA_FW_VOL_FILEPATH_DEVICE_PATH     *ThisFilePath;
 691   EFI_GUID                              *FileName;
 692
 693   FileName = NULL;
 694   if (FilePath != NULL) {
 695     ThisFilePath = (MEDIA_FW_VOL_FILEPATH_DEVICE_PATH *) FilePath;
 696     while (!IsDevicePathEnd (ThisFilePath)) {
 697       FileName = EfiGetNameGuidFromFwVolDevicePathNode (ThisFilePath);
 698       if (FileName != NULL) {
 699         break;
 700       }
 701       ThisFilePath = (MEDIA_FW_VOL_FILEPATH_DEVICE_PATH *) NextDevicePathNode (ThisFilePath);
 702     }
 703   }
 704
 705   return FileName;
 706 }
 707
 708 /**
 709   Register image to memory profile.
 710
 711   @param DriverEntry    Image info.
 712   @param FileType       Image file type.
 713
 714   @return EFI_SUCCESS           Register successfully.
 715   @return EFI_UNSUPPORTED       Memory profile is unsupported,
 716                                 or memory profile for the image is not required.
 717   @return EFI_OUT_OF_RESOURCES  No enough resource for this register.
 718
 719 **/
 720 EFI_STATUS
 721 RegisterMemoryProfileImage (
 722   IN LOADED_IMAGE_PRIVATE_DATA  *DriverEntry,
 723   IN EFI_FV_FILETYPE            FileType
 724   )
 725 {
 726   MEMORY_PROFILE_CONTEXT_DATA       *ContextData;
 727   MEMORY_PROFILE_DRIVER_INFO_DATA   *DriverInfoData;
 728
 729   if (!IS_UEFI_MEMORY_PROFILE_ENABLED) {
 730     return EFI_UNSUPPORTED;
 731   }
 732
 733   if (!NeedRecordThisDriver (DriverEntry->Info.FilePath)) {
 734     return EFI_UNSUPPORTED;
 735   }
 736
 737   ContextData = GetMemoryProfileContext ();
 738   if (ContextData == NULL) {
 739     return EFI_UNSUPPORTED;
 740   }
 741
 742   DriverInfoData = BuildDriverInfo (
 743                      ContextData,
 744                      GetFileNameFromFilePath (DriverEntry->Info.FilePath),
 745                      DriverEntry->ImageContext.ImageAddress,
 746                      DriverEntry->ImageContext.ImageSize,
 747                      DriverEntry->ImageContext.EntryPoint,
 748                      DriverEntry->ImageContext.ImageType,
 749                      FileType
 750                      );
 751   if (DriverInfoData == NULL) {
 752     return EFI_OUT_OF_RESOURCES;
 753   }
 754
 755   return EFI_SUCCESS;
 756 }
 757
 758 /**
 759   Search image from memory profile.
 760
 761   @param ContextData    Memory profile context.
 762   @param FileName       Image file name.
 763   @param Address        Image Address.
 764
 765   @return Pointer to memory profile driver info.
 766
 767 **/
 768 MEMORY_PROFILE_DRIVER_INFO_DATA *
 769 GetMemoryProfileDriverInfoByFileNameAndAddress (
 770   IN MEMORY_PROFILE_CONTEXT_DATA    *ContextData,
 771   IN EFI_GUID                       *FileName,
 772   IN PHYSICAL_ADDRESS               Address
 773   )
 774 {
 775   MEMORY_PROFILE_DRIVER_INFO        *DriverInfo;
 776   MEMORY_PROFILE_DRIVER_INFO_DATA   *DriverInfoData;
 777   LIST_ENTRY                        *DriverLink;
 778   LIST_ENTRY                        *DriverInfoList;
 779
 780   DriverInfoList = ContextData->DriverInfoList;
 781
 782   for (DriverLink = DriverInfoList->ForwardLink;
 783        DriverLink != DriverInfoList;
 784        DriverLink = DriverLink->ForwardLink) {
 785     DriverInfoData = CR (
 786                        DriverLink,
 787                        MEMORY_PROFILE_DRIVER_INFO_DATA,
 788                        Link,
 789                        MEMORY_PROFILE_DRIVER_INFO_SIGNATURE
 790                        );
 791     DriverInfo = &DriverInfoData->DriverInfo;
 792     if ((CompareGuid (&DriverInfo->FileName, FileName)) &&
 793         (Address >= DriverInfo->ImageBase) &&
 794         (Address < (DriverInfo->ImageBase + DriverInfo->ImageSize))) {
 795       return DriverInfoData;
 796     }
 797   }
 798
 799   return NULL;
 800 }
 801
 802 /**
 803   Search image from memory profile.
 804   It will return image, if (Address >= ImageBuffer) AND (Address < ImageBuffer + ImageSize).
 805
 806   @param ContextData    Memory profile context.
 807   @param Address        Image or Function address.
 808
 809   @return Pointer to memory profile driver info.
 810
 811 **/
 812 MEMORY_PROFILE_DRIVER_INFO_DATA *
 813 GetMemoryProfileDriverInfoFromAddress (
 814   IN MEMORY_PROFILE_CONTEXT_DATA    *ContextData,
 815   IN PHYSICAL_ADDRESS               Address
 816   )
 817 {
 818   MEMORY_PROFILE_DRIVER_INFO        *DriverInfo;
 819   MEMORY_PROFILE_DRIVER_INFO_DATA   *DriverInfoData;
 820   LIST_ENTRY                        *DriverLink;
 821   LIST_ENTRY                        *DriverInfoList;
 822
 823   DriverInfoList = ContextData->DriverInfoList;
 824
 825   for (DriverLink = DriverInfoList->ForwardLink;
 826        DriverLink != DriverInfoList;
 827        DriverLink = DriverLink->ForwardLink) {
 828     DriverInfoData = CR (
 829                        DriverLink,
 830                        MEMORY_PROFILE_DRIVER_INFO_DATA,
 831                        Link,
 832                        MEMORY_PROFILE_DRIVER_INFO_SIGNATURE
 833                        );
 834     DriverInfo = &DriverInfoData->DriverInfo;
 835     if ((Address >= DriverInfo->ImageBase) &&
 836         (Address < (DriverInfo->ImageBase + DriverInfo->ImageSize))) {
 837       return DriverInfoData;
 838     }
 839   }
 840
 841   return NULL;
 842 }
 843
 844 /**
 845   Unregister image from memory profile.
 846
 847   @param DriverEntry    Image info.
 848
 849   @return EFI_SUCCESS           Unregister successfully.
 850   @return EFI_UNSUPPORTED       Memory profile is unsupported,
 851                                 or memory profile for the image is not required.
 852   @return EFI_NOT_FOUND         The image is not found.
 853
 854 **/
 855 EFI_STATUS
 856 UnregisterMemoryProfileImage (
 857   IN LOADED_IMAGE_PRIVATE_DATA      *DriverEntry
 858   )
 859 {
 860   EFI_STATUS                        Status;
 861   MEMORY_PROFILE_CONTEXT_DATA       *ContextData;
 862   MEMORY_PROFILE_DRIVER_INFO_DATA   *DriverInfoData;
 863   EFI_GUID                          *FileName;
 864   PHYSICAL_ADDRESS                  ImageAddress;
 865   VOID                              *EntryPointInImage;
 866
 867   if (!IS_UEFI_MEMORY_PROFILE_ENABLED) {
 868     return EFI_UNSUPPORTED;
 869   }
 870
 871   if (!NeedRecordThisDriver (DriverEntry->Info.FilePath)) {
 872     return EFI_UNSUPPORTED;
 873   }
 874
 875   ContextData = GetMemoryProfileContext ();
 876   if (ContextData == NULL) {
 877     return EFI_UNSUPPORTED;
 878   }
 879
 880   DriverInfoData = NULL;
 881   FileName = GetFileNameFromFilePath (DriverEntry->Info.FilePath);
 882   ImageAddress = DriverEntry->ImageContext.ImageAddress;
 883   if ((DriverEntry->ImageContext.EntryPoint < ImageAddress) || (DriverEntry->ImageContext.EntryPoint >= (ImageAddress + DriverEntry->ImageContext.ImageSize))) {
 884     //
 885     // If the EntryPoint is not in the range of image buffer, it should come from emulation environment.
 886     // So patch ImageAddress here to align the EntryPoint.
 887     //
 888     Status = InternalPeCoffGetEntryPoint ((VOID *) (UINTN) ImageAddress, &EntryPointInImage);
 889     ASSERT_EFI_ERROR (Status);
 890     ImageAddress = ImageAddress + (UINTN) DriverEntry->ImageContext.EntryPoint - (UINTN) EntryPointInImage;
 891   }
 892   if (FileName != NULL) {
 893     DriverInfoData = GetMemoryProfileDriverInfoByFileNameAndAddress (ContextData, FileName, ImageAddress);
 894   }
 895   if (DriverInfoData == NULL) {
 896     DriverInfoData = GetMemoryProfileDriverInfoFromAddress (ContextData, ImageAddress);
 897   }
 898   if (DriverInfoData == NULL) {
 899     return EFI_NOT_FOUND;
 900   }
 901
 902   ContextData->Context.TotalImageSize -= DriverInfoData->DriverInfo.ImageSize;
 903
 904   // Keep the ImageBase for RVA calculation in Application.
 905   //DriverInfoData->DriverInfo.ImageBase = 0;
 906   DriverInfoData->DriverInfo.ImageSize = 0;
 907
 908   if (DriverInfoData->DriverInfo.PeakUsage == 0) {
 909     ContextData->Context.ImageCount --;
 910     RemoveEntryList (&DriverInfoData->Link);
 911     //
 912     // Use CoreInternalFreePool() that will not update profile for this FreePool action.
 913     //
 914     CoreInternalFreePool (DriverInfoData, NULL);
 915   }
 916
 917   return EFI_SUCCESS;
 918 }
 919
 920 /**
 921   Return if this memory type needs to be recorded into memory profile.
 922   If BIOS memory type (0 ~ EfiMaxMemoryType - 1), it checks bit (1 << MemoryType).
 923   If OS memory type (0x80000000 ~ 0xFFFFFFFF), it checks bit63 - 0x8000000000000000.
 924   If OEM memory type (0x70000000 ~ 0x7FFFFFFF), it checks bit62 - 0x4000000000000000.
 925
 926   @param MemoryType     Memory type.
 927
 928   @retval TRUE          This memory type need to be recorded.
 929   @retval FALSE         This memory type need not to be recorded.
 930
 931 **/
 932 BOOLEAN
 933 CoreNeedRecordProfile (
 934   IN EFI_MEMORY_TYPE    MemoryType
 935   )
 936 {
 937   UINT64 TestBit;
 938
 939   if ((UINT32) MemoryType >= MEMORY_TYPE_OS_RESERVED_MIN) {
 940     TestBit = BIT63;
 941   } else if ((UINT32) MemoryType >= MEMORY_TYPE_OEM_RESERVED_MIN) {
 942     TestBit = BIT62;
 943   } else {
 944     TestBit = LShiftU64 (1, MemoryType);
 945   }
 946
 947   if ((PcdGet64 (PcdMemoryProfileMemoryType) & TestBit) != 0) {
 948     return TRUE;
 949   } else {
 950     return FALSE;
 951   }
 952 }
 953
 954 /**
 955   Convert EFI memory type to profile memory index. The rule is:
 956   If BIOS memory type (0 ~ EfiMaxMemoryType - 1), ProfileMemoryIndex = MemoryType.
 957   If OS memory type (0x80000000 ~ 0xFFFFFFFF), ProfileMemoryIndex = EfiMaxMemoryType.
 958   If OEM memory type (0x70000000 ~ 0x7FFFFFFF), ProfileMemoryIndex = EfiMaxMemoryType + 1.
 959
 960   @param MemoryType     Memory type.
 961
 962   @return Profile memory index.
 963
 964 **/
 965 UINTN
 966 GetProfileMemoryIndex (
 967   IN EFI_MEMORY_TYPE    MemoryType
 968   )
 969 {
 970   if ((UINT32) MemoryType >= MEMORY_TYPE_OS_RESERVED_MIN) {
 971     return EfiMaxMemoryType;
 972   } else if ((UINT32) MemoryType >= MEMORY_TYPE_OEM_RESERVED_MIN) {
 973     return EfiMaxMemoryType + 1;
 974   } else {
 975     return MemoryType;
 976   }
 977 }
 978
 979 /**
 980   Update memory profile Allocate information.
 981
 982   @param CallerAddress  Address of caller who call Allocate.
 983   @param Action         This Allocate action.
 984   @param MemoryType     Memory type.
 985   @param Size           Buffer size.
 986   @param Buffer         Buffer address.
 987   @param ActionString   String for memory profile action.
 988
 989   @return EFI_SUCCESS           Memory profile is updated.
 990   @return EFI_UNSUPPORTED       Memory profile is unsupported,
 991                                 or memory profile for the image is not required.
 992   @return EFI_OUT_OF_RESOURCES  No enough resource to update memory profile for allocate action.
 993
 994 **/
 995 EFI_STATUS
 996 CoreUpdateProfileAllocate (
 997   IN PHYSICAL_ADDRESS       CallerAddress,
 998   IN MEMORY_PROFILE_ACTION  Action,
 999   IN EFI_MEMORY_TYPE        MemoryType,
1000   IN UINTN                  Size,
1001   IN VOID                   *Buffer,
1002   IN CHAR8                  *ActionString OPTIONAL
1003   )
1004 {
1005   EFI_STATUS                        Status;
1006   MEMORY_PROFILE_CONTEXT            *Context;
1007   MEMORY_PROFILE_DRIVER_INFO        *DriverInfo;
1008   MEMORY_PROFILE_ALLOC_INFO         *AllocInfo;
1009   MEMORY_PROFILE_CONTEXT_DATA       *ContextData;
1010   MEMORY_PROFILE_DRIVER_INFO_DATA   *DriverInfoData;
1011   MEMORY_PROFILE_ALLOC_INFO_DATA    *AllocInfoData;
1012   UINTN                             ProfileMemoryIndex;
1013   MEMORY_PROFILE_ACTION             BasicAction;
1014   UINTN                             ActionStringSize;
1015   UINTN                             ActionStringOccupiedSize;
1016
1017   BasicAction = Action & MEMORY_PROFILE_ACTION_BASIC_MASK;
1018
1019   ContextData = GetMemoryProfileContext ();
1020   if (ContextData == NULL) {
1021     return EFI_UNSUPPORTED;
1022   }
1023
1024   DriverInfoData = GetMemoryProfileDriverInfoFromAddress (ContextData, CallerAddress);
1025   if (DriverInfoData == NULL) {
1026     return EFI_UNSUPPORTED;
1027   }
1028
1029   ActionStringSize = 0;
1030   ActionStringOccupiedSize = 0;
1031   if (ActionString != NULL) {
1032     ActionStringSize = AsciiStrSize (ActionString);
1033     ActionStringOccupiedSize = GET_OCCUPIED_SIZE (ActionStringSize, sizeof (UINT64));
1034   }
1035
1036   //
1037   // Use CoreInternalAllocatePool() that will not update profile for this AllocatePool action.
1038   //
1039   AllocInfoData = NULL;
1040   Status = CoreInternalAllocatePool (
1041              EfiBootServicesData,
1042              sizeof (*AllocInfoData) + ActionStringSize,
1043              (VOID **) &AllocInfoData
1044              );
1045   if (EFI_ERROR (Status)) {
1046     return EFI_OUT_OF_RESOURCES;
1047   }
1048   ASSERT (AllocInfoData != NULL);
1049
1050   //
1051   // Only update SequenceCount if and only if it is basic action.
1052   //
1053   if (Action == BasicAction) {
1054     ContextData->Context.SequenceCount ++;
1055   }
1056
1057   AllocInfo = &AllocInfoData->AllocInfo;
1058   AllocInfoData->Signature      = MEMORY_PROFILE_ALLOC_INFO_SIGNATURE;
1059   AllocInfo->Header.Signature   = MEMORY_PROFILE_ALLOC_INFO_SIGNATURE;
1060   AllocInfo->Header.Length      = (UINT16) (sizeof (MEMORY_PROFILE_ALLOC_INFO) + ActionStringOccupiedSize);
1061   AllocInfo->Header.Revision    = MEMORY_PROFILE_ALLOC_INFO_REVISION;
1062   AllocInfo->CallerAddress      = CallerAddress;
1063   AllocInfo->SequenceId         = ContextData->Context.SequenceCount;
1064   AllocInfo->Action             = Action;
1065   AllocInfo->MemoryType         = MemoryType;
1066   AllocInfo->Buffer             = (PHYSICAL_ADDRESS) (UINTN) Buffer;
1067   AllocInfo->Size               = Size;
1068   if (ActionString != NULL) {
1069     AllocInfo->ActionStringOffset = (UINT16) sizeof (MEMORY_PROFILE_ALLOC_INFO);
1070     AllocInfoData->ActionString = (CHAR8 *) (AllocInfoData + 1);
1071     CopyMem (AllocInfoData->ActionString, ActionString, ActionStringSize);
1072   } else {
1073     AllocInfo->ActionStringOffset = 0;
1074     AllocInfoData->ActionString = NULL;
1075   }
1076
1077   InsertTailList (DriverInfoData->AllocInfoList, &AllocInfoData->Link);
1078
1079   Context = &ContextData->Context;
1080   DriverInfo = &DriverInfoData->DriverInfo;
1081   DriverInfo->AllocRecordCount ++;
1082
1083   //
1084   // Update summary if and only if it is basic action.
1085   //
1086   if (Action == BasicAction) {
1087     ProfileMemoryIndex = GetProfileMemoryIndex (MemoryType);
1088
1089     DriverInfo->CurrentUsage += Size;
1090     if (DriverInfo->PeakUsage < DriverInfo->CurrentUsage) {
1091       DriverInfo->PeakUsage = DriverInfo->CurrentUsage;
1092     }
1093     DriverInfo->CurrentUsageByType[ProfileMemoryIndex] += Size;
1094     if (DriverInfo->PeakUsageByType[ProfileMemoryIndex] < DriverInfo->CurrentUsageByType[ProfileMemoryIndex]) {
1095       DriverInfo->PeakUsageByType[ProfileMemoryIndex] = DriverInfo->CurrentUsageByType[ProfileMemoryIndex];
1096     }
1097
1098     Context->CurrentTotalUsage += Size;
1099     if (Context->PeakTotalUsage < Context->CurrentTotalUsage) {
1100       Context->PeakTotalUsage = Context->CurrentTotalUsage;
1101     }
1102     Context->CurrentTotalUsageByType[ProfileMemoryIndex] += Size;
1103     if (Context->PeakTotalUsageByType[ProfileMemoryIndex] < Context->CurrentTotalUsageByType[ProfileMemoryIndex]) {
1104       Context->PeakTotalUsageByType[ProfileMemoryIndex] = Context->CurrentTotalUsageByType[ProfileMemoryIndex];
1105     }
1106   }
1107
1108   return EFI_SUCCESS;
1109 }
1110
1111 /**
1112   Get memory profile alloc info from memory profile.
1113
1114   @param DriverInfoData     Driver info.
1115   @param BasicAction        This Free basic action.
1116   @param Size               Buffer size.
1117   @param Buffer             Buffer address.
1118
1119   @return Pointer to memory profile alloc info.
1120
1121 **/
1122 MEMORY_PROFILE_ALLOC_INFO_DATA *
1123 GetMemoryProfileAllocInfoFromAddress (
1124   IN MEMORY_PROFILE_DRIVER_INFO_DATA    *DriverInfoData,
1125   IN MEMORY_PROFILE_ACTION              BasicAction,
1126   IN UINTN                              Size,
1127   IN VOID                               *Buffer
1128   )
1129 {
1130   LIST_ENTRY                        *AllocInfoList;
1131   LIST_ENTRY                        *AllocLink;
1132   MEMORY_PROFILE_ALLOC_INFO         *AllocInfo;
1133   MEMORY_PROFILE_ALLOC_INFO_DATA    *AllocInfoData;
1134
1135   AllocInfoList = DriverInfoData->AllocInfoList;
1136
1137   for (AllocLink = AllocInfoList->ForwardLink;
1138        AllocLink != AllocInfoList;
1139        AllocLink = AllocLink->ForwardLink) {
1140     AllocInfoData = CR (
1141                       AllocLink,
1142                       MEMORY_PROFILE_ALLOC_INFO_DATA,
1143                       Link,
1144                       MEMORY_PROFILE_ALLOC_INFO_SIGNATURE
1145                       );
1146     AllocInfo = &AllocInfoData->AllocInfo;
1147     if ((AllocInfo->Action & MEMORY_PROFILE_ACTION_BASIC_MASK) != BasicAction) {
1148       continue;
1149     }
1150     switch (BasicAction) {
1151       case MemoryProfileActionAllocatePages:
1152         if ((AllocInfo->Buffer <= (PHYSICAL_ADDRESS) (UINTN) Buffer) &&
1153             ((AllocInfo->Buffer + AllocInfo->Size) >= ((PHYSICAL_ADDRESS) (UINTN) Buffer + Size))) {
1154           return AllocInfoData;
1155         }
1156         break;
1157       case MemoryProfileActionAllocatePool:
1158         if (AllocInfo->Buffer == (PHYSICAL_ADDRESS) (UINTN) Buffer) {
1159           return AllocInfoData;
1160         }
1161         break;
1162       default:
1163         ASSERT (FALSE);
1164         break;
1165     }
1166   }
1167
1168   return NULL;
1169 }
1170
1171 /**
1172   Update memory profile Free information.
1173
1174   @param CallerAddress  Address of caller who call Free.
1175   @param Action         This Free action.
1176   @param Size           Buffer size.
1177   @param Buffer         Buffer address.
1178
1179   @return EFI_SUCCESS           Memory profile is updated.
1180   @return EFI_UNSUPPORTED       Memory profile is unsupported.
1181   @return EFI_NOT_FOUND         No matched allocate info found for free action.
1182
1183 **/
1184 EFI_STATUS
1185 CoreUpdateProfileFree (
1186   IN PHYSICAL_ADDRESS       CallerAddress,
1187   IN MEMORY_PROFILE_ACTION  Action,
1188   IN UINTN                  Size,
1189   IN VOID                   *Buffer
1190   )
1191 {
1192   MEMORY_PROFILE_CONTEXT           *Context;
1193   MEMORY_PROFILE_DRIVER_INFO       *DriverInfo;
1194   MEMORY_PROFILE_ALLOC_INFO        *AllocInfo;
1195   MEMORY_PROFILE_CONTEXT_DATA      *ContextData;
1196   MEMORY_PROFILE_DRIVER_INFO_DATA  *DriverInfoData;
1197   LIST_ENTRY                       *DriverLink;
1198   LIST_ENTRY                       *DriverInfoList;
1199   MEMORY_PROFILE_DRIVER_INFO_DATA  *ThisDriverInfoData;
1200   MEMORY_PROFILE_ALLOC_INFO_DATA   *AllocInfoData;
1201   UINTN                            ProfileMemoryIndex;
1202   MEMORY_PROFILE_ACTION            BasicAction;
1203   BOOLEAN                          Found;
1204
1205   BasicAction = Action & MEMORY_PROFILE_ACTION_BASIC_MASK;
1206
1207   ContextData = GetMemoryProfileContext ();
1208   if (ContextData == NULL) {
1209     return EFI_UNSUPPORTED;
1210   }
1211
1212   DriverInfoData = GetMemoryProfileDriverInfoFromAddress (ContextData, CallerAddress);
1213
1214   //
1215   // Do not return if DriverInfoData == NULL here,
1216   // because driver A might free memory allocated by driver B.
1217   //
1218
1219   //
1220   // Need use do-while loop to find all possible records,
1221   // because one address might be recorded multiple times.
1222   //
1223   Found = FALSE;
1224   AllocInfoData = NULL;
1225   do {
1226     if (DriverInfoData != NULL) {
1227       switch (BasicAction) {
1228         case MemoryProfileActionFreePages:
1229           AllocInfoData = GetMemoryProfileAllocInfoFromAddress (DriverInfoData, MemoryProfileActionAllocatePages, Size, Buffer);
1230           break;
1231         case MemoryProfileActionFreePool:
1232           AllocInfoData = GetMemoryProfileAllocInfoFromAddress (DriverInfoData, MemoryProfileActionAllocatePool, 0, Buffer);
1233           break;
1234         default:
1235           ASSERT (FALSE);
1236           AllocInfoData = NULL;
1237           break;
1238       }
1239     }
1240     if (AllocInfoData == NULL) {
1241       //
1242       // Legal case, because driver A might free memory allocated by driver B, by some protocol.
1243       //
1244       DriverInfoList = ContextData->DriverInfoList;
1245
1246       for (DriverLink = DriverInfoList->ForwardLink;
1247            DriverLink != DriverInfoList;
1248            DriverLink = DriverLink->ForwardLink) {
1249         ThisDriverInfoData = CR (
1250                                DriverLink,
1251                                MEMORY_PROFILE_DRIVER_INFO_DATA,
1252                                Link,
1253                                MEMORY_PROFILE_DRIVER_INFO_SIGNATURE
1254                                );
1255         switch (BasicAction) {
1256           case MemoryProfileActionFreePages:
1257             AllocInfoData = GetMemoryProfileAllocInfoFromAddress (ThisDriverInfoData, MemoryProfileActionAllocatePages, Size, Buffer);
1258             break;
1259           case MemoryProfileActionFreePool:
1260             AllocInfoData = GetMemoryProfileAllocInfoFromAddress (ThisDriverInfoData, MemoryProfileActionAllocatePool, 0, Buffer);
1261             break;
1262           default:
1263             ASSERT (FALSE);
1264             AllocInfoData = NULL;
1265             break;
1266         }
1267         if (AllocInfoData != NULL) {
1268           DriverInfoData = ThisDriverInfoData;
1269           break;
1270         }
1271       }
1272
1273       if (AllocInfoData == NULL) {
1274         //
1275         // If (!Found), no matched allocate info is found for this free action.
1276         // It is because the specified memory type allocate actions have been filtered by
1277         // CoreNeedRecordProfile(), but free actions may have no memory type information,
1278         // they can not be filtered by CoreNeedRecordProfile(). Then, they will be
1279         // filtered here.
1280         //
1281         // If (Found), it is normal exit path.
1282         return (Found ? EFI_SUCCESS : EFI_NOT_FOUND);
1283       }
1284     }
1285
1286     Found = TRUE;
1287
1288     Context = &ContextData->Context;
1289     DriverInfo = &DriverInfoData->DriverInfo;
1290     AllocInfo = &AllocInfoData->AllocInfo;
1291
1292     DriverInfo->AllocRecordCount --;
1293     //
1294     // Update summary if and only if it is basic action.
1295     //
1296     if (AllocInfo->Action == (AllocInfo->Action & MEMORY_PROFILE_ACTION_BASIC_MASK)) {
1297       ProfileMemoryIndex = GetProfileMemoryIndex (AllocInfo->MemoryType);
1298
1299       Context->CurrentTotalUsage -= AllocInfo->Size;
1300       Context->CurrentTotalUsageByType[ProfileMemoryIndex] -= AllocInfo->Size;
1301
1302       DriverInfo->CurrentUsage -= AllocInfo->Size;
1303       DriverInfo->CurrentUsageByType[ProfileMemoryIndex] -= AllocInfo->Size;
1304     }
1305
1306     RemoveEntryList (&AllocInfoData->Link);
1307
1308     if (BasicAction == MemoryProfileActionFreePages) {
1309       if (AllocInfo->Buffer != (PHYSICAL_ADDRESS) (UINTN) Buffer) {
1310         CoreUpdateProfileAllocate (
1311           AllocInfo->CallerAddress,
1312           AllocInfo->Action,
1313           AllocInfo->MemoryType,
1314           (UINTN) ((PHYSICAL_ADDRESS) (UINTN) Buffer - AllocInfo->Buffer),
1315           (VOID *) (UINTN) AllocInfo->Buffer,
1316           AllocInfoData->ActionString
1317           );
1318       }
1319       if (AllocInfo->Buffer + AllocInfo->Size != ((PHYSICAL_ADDRESS) (UINTN) Buffer + Size)) {
1320         CoreUpdateProfileAllocate (
1321           AllocInfo->CallerAddress,
1322           AllocInfo->Action,
1323           AllocInfo->MemoryType,
1324           (UINTN) ((AllocInfo->Buffer + AllocInfo->Size) - ((PHYSICAL_ADDRESS) (UINTN) Buffer + Size)),
1325           (VOID *) ((UINTN) Buffer + Size),
1326           AllocInfoData->ActionString
1327           );
1328       }
1329     }
1330
1331     //
1332     // Use CoreInternalFreePool() that will not update profile for this FreePool action.
1333     //
1334     CoreInternalFreePool (AllocInfoData, NULL);
1335   } while (TRUE);
1336 }
1337
1338 /**
1339   Update memory profile information.
1340
1341   @param CallerAddress  Address of caller who call Allocate or Free.
1342   @param Action         This Allocate or Free action.
1343   @param MemoryType     Memory type.
1344                         EfiMaxMemoryType means the MemoryType is unknown.
1345   @param Size           Buffer size.
1346   @param Buffer         Buffer address.
1347   @param ActionString   String for memory profile action.
1348                         Only needed for user defined allocate action.
1349
1350   @return EFI_SUCCESS           Memory profile is updated.
1351   @return EFI_UNSUPPORTED       Memory profile is unsupported,
1352                                 or memory profile for the image is not required,
1353                                 or memory profile for the memory type is not required.
1354   @return EFI_ACCESS_DENIED     It is during memory profile data getting.
1355   @return EFI_ABORTED           Memory profile recording is not enabled.
1356   @return EFI_OUT_OF_RESOURCES  No enough resource to update memory profile for allocate action.
1357   @return EFI_NOT_FOUND         No matched allocate info found for free action.
1358
1359 **/
1360 EFI_STATUS
1361 EFIAPI
1362 CoreUpdateProfile (
1363   IN PHYSICAL_ADDRESS       CallerAddress,
1364   IN MEMORY_PROFILE_ACTION  Action,
1365   IN EFI_MEMORY_TYPE        MemoryType,
1366   IN UINTN                  Size,       // Valid for AllocatePages/FreePages/AllocatePool
1367   IN VOID                   *Buffer,
1368   IN CHAR8                  *ActionString OPTIONAL
1369   )
1370 {
1371   EFI_STATUS                    Status;
1372   MEMORY_PROFILE_CONTEXT_DATA   *ContextData;
1373   MEMORY_PROFILE_ACTION         BasicAction;
1374
1375   if (!IS_UEFI_MEMORY_PROFILE_ENABLED) {
1376     return EFI_UNSUPPORTED;
1377   }
1378
1379   if (mMemoryProfileGettingStatus) {
1380     return EFI_ACCESS_DENIED;
1381   }
1382
1383   if (!mMemoryProfileRecordingEnable) {
1384     return EFI_ABORTED;
1385   }
1386
1387   //
1388   // Get the basic action to know how to process the record
1389   //
1390   BasicAction = Action & MEMORY_PROFILE_ACTION_BASIC_MASK;
1391
1392   //
1393   // EfiMaxMemoryType means the MemoryType is unknown.
1394   //
1395   if (MemoryType != EfiMaxMemoryType) {
1396     //
1397     // Only record limited MemoryType.
1398     //
1399     if (!CoreNeedRecordProfile (MemoryType)) {
1400       return EFI_UNSUPPORTED;
1401     }
1402   }
1403
1404   ContextData = GetMemoryProfileContext ();
1405   if (ContextData == NULL) {
1406     return EFI_UNSUPPORTED;
1407   }
1408
1409   CoreAcquireMemoryProfileLock ();
1410   switch (BasicAction) {
1411     case MemoryProfileActionAllocatePages:
1412       Status = CoreUpdateProfileAllocate (CallerAddress, Action, MemoryType, Size, Buffer, ActionString);
1413       break;
1414     case MemoryProfileActionFreePages:
1415       Status = CoreUpdateProfileFree (CallerAddress, Action, Size, Buffer);
1416       break;
1417     case MemoryProfileActionAllocatePool:
1418       Status = CoreUpdateProfileAllocate (CallerAddress, Action, MemoryType, Size, Buffer, ActionString);
1419       break;
1420     case MemoryProfileActionFreePool:
1421       Status = CoreUpdateProfileFree (CallerAddress, Action, 0, Buffer);
1422       break;
1423     default:
1424       ASSERT (FALSE);
1425       Status = EFI_UNSUPPORTED;
1426       break;
1427   }
1428   CoreReleaseMemoryProfileLock ();
1429
1430   return Status;
1431 }
1432
1433 ////////////////////
1434
1435 /**
1436   Get memory profile data size.
1437
1438   @return Memory profile data size.
1439
1440 **/
1441 UINTN
1442 MemoryProfileGetDataSize (
1443   VOID
1444   )
1445 {
1446   MEMORY_PROFILE_CONTEXT_DATA       *ContextData;
1447   MEMORY_PROFILE_DRIVER_INFO_DATA   *DriverInfoData;
1448   MEMORY_PROFILE_ALLOC_INFO_DATA    *AllocInfoData;
1449   LIST_ENTRY                        *DriverInfoList;
1450   LIST_ENTRY                        *DriverLink;
1451   LIST_ENTRY                        *AllocInfoList;
1452   LIST_ENTRY                        *AllocLink;
1453   UINTN                             TotalSize;
1454
1455
1456   ContextData = GetMemoryProfileContext ();
1457   if (ContextData == NULL) {
1458     return 0;
1459   }
1460
1461   TotalSize = sizeof (MEMORY_PROFILE_CONTEXT);
1462
1463   DriverInfoList = ContextData->DriverInfoList;
1464   for (DriverLink = DriverInfoList->ForwardLink;
1465        DriverLink != DriverInfoList;
1466        DriverLink = DriverLink->ForwardLink) {
1467     DriverInfoData = CR (
1468                        DriverLink,
1469                        MEMORY_PROFILE_DRIVER_INFO_DATA,
1470                        Link,
1471                        MEMORY_PROFILE_DRIVER_INFO_SIGNATURE
1472                        );
1473     TotalSize += DriverInfoData->DriverInfo.Header.Length;
1474
1475     AllocInfoList = DriverInfoData->AllocInfoList;
1476     for (AllocLink = AllocInfoList->ForwardLink;
1477          AllocLink != AllocInfoList;
1478          AllocLink = AllocLink->ForwardLink) {
1479       AllocInfoData = CR (
1480                         AllocLink,
1481                         MEMORY_PROFILE_ALLOC_INFO_DATA,
1482                         Link,
1483                         MEMORY_PROFILE_ALLOC_INFO_SIGNATURE
1484                         );
1485       TotalSize += AllocInfoData->AllocInfo.Header.Length;
1486     }
1487   }
1488
1489   return TotalSize;
1490 }
1491
1492 /**
1493   Copy memory profile data.
1494
1495   @param ProfileBuffer  The buffer to hold memory profile data.
1496
1497 **/
1498 VOID
1499 MemoryProfileCopyData (
1500   IN VOID   *ProfileBuffer
1501   )
1502 {
1503   MEMORY_PROFILE_CONTEXT            *Context;
1504   MEMORY_PROFILE_DRIVER_INFO        *DriverInfo;
1505   MEMORY_PROFILE_ALLOC_INFO         *AllocInfo;
1506   MEMORY_PROFILE_CONTEXT_DATA       *ContextData;
1507   MEMORY_PROFILE_DRIVER_INFO_DATA   *DriverInfoData;
1508   MEMORY_PROFILE_ALLOC_INFO_DATA    *AllocInfoData;
1509   LIST_ENTRY                        *DriverInfoList;
1510   LIST_ENTRY                        *DriverLink;
1511   LIST_ENTRY                        *AllocInfoList;
1512   LIST_ENTRY                        *AllocLink;
1513   UINTN                             PdbSize;
1514   UINTN                             ActionStringSize;
1515
1516   ContextData = GetMemoryProfileContext ();
1517   if (ContextData == NULL) {
1518     return ;
1519   }
1520
1521   Context = ProfileBuffer;
1522   CopyMem (Context, &ContextData->Context, sizeof (MEMORY_PROFILE_CONTEXT));
1523   DriverInfo = (MEMORY_PROFILE_DRIVER_INFO *) (Context + 1);
1524
1525   DriverInfoList = ContextData->DriverInfoList;
1526   for (DriverLink = DriverInfoList->ForwardLink;
1527        DriverLink != DriverInfoList;
1528        DriverLink = DriverLink->ForwardLink) {
1529     DriverInfoData = CR (
1530                        DriverLink,
1531                        MEMORY_PROFILE_DRIVER_INFO_DATA,
1532                        Link,
1533                        MEMORY_PROFILE_DRIVER_INFO_SIGNATURE
1534                        );
1535     CopyMem (DriverInfo, &DriverInfoData->DriverInfo, sizeof (MEMORY_PROFILE_DRIVER_INFO));
1536     if (DriverInfo->PdbStringOffset != 0) {
1537       PdbSize = AsciiStrSize (DriverInfoData->PdbString);
1538       CopyMem ((VOID *) ((UINTN) DriverInfo + DriverInfo->PdbStringOffset), DriverInfoData->PdbString, PdbSize);
1539     }
1540     AllocInfo = (MEMORY_PROFILE_ALLOC_INFO *) ((UINTN) DriverInfo + DriverInfo->Header.Length);
1541
1542     AllocInfoList = DriverInfoData->AllocInfoList;
1543     for (AllocLink = AllocInfoList->ForwardLink;
1544          AllocLink != AllocInfoList;
1545          AllocLink = AllocLink->ForwardLink) {
1546       AllocInfoData = CR (
1547                         AllocLink,
1548                         MEMORY_PROFILE_ALLOC_INFO_DATA,
1549                         Link,
1550                         MEMORY_PROFILE_ALLOC_INFO_SIGNATURE
1551                         );
1552       CopyMem (AllocInfo, &AllocInfoData->AllocInfo, sizeof (MEMORY_PROFILE_ALLOC_INFO));
1553       if (AllocInfo->ActionStringOffset != 0) {
1554         ActionStringSize = AsciiStrSize (AllocInfoData->ActionString);
1555         CopyMem ((VOID *) ((UINTN) AllocInfo + AllocInfo->ActionStringOffset), AllocInfoData->ActionString, ActionStringSize);
1556       }
1557       AllocInfo = (MEMORY_PROFILE_ALLOC_INFO *) ((UINTN) AllocInfo + AllocInfo->Header.Length);
1558     }
1559
1560     DriverInfo = (MEMORY_PROFILE_DRIVER_INFO *)  AllocInfo;
1561   }
1562 }
1563
1564 /**
1565   Get memory profile data.
1566
1567   @param[in]      This              The EDKII_MEMORY_PROFILE_PROTOCOL instance.
1568   @param[in, out] ProfileSize       On entry, points to the size in bytes of the ProfileBuffer.
1569                                     On return, points to the size of the data returned in ProfileBuffer.
1570   @param[out]     ProfileBuffer     Profile buffer.
1571
1572   @return EFI_SUCCESS               Get the memory profile data successfully.
1573   @return EFI_UNSUPPORTED           Memory profile is unsupported.
1574   @return EFI_BUFFER_TO_SMALL       The ProfileSize is too small for the resulting data.
1575                                     ProfileSize is updated with the size required.
1576
1577 **/
1578 EFI_STATUS
1579 EFIAPI
1580 ProfileProtocolGetData (
1581   IN     EDKII_MEMORY_PROFILE_PROTOCOL  *This,
1582   IN OUT UINT64                         *ProfileSize,
1583      OUT VOID                           *ProfileBuffer
1584   )
1585 {
1586   UINTN                                 Size;
1587   MEMORY_PROFILE_CONTEXT_DATA           *ContextData;
1588   BOOLEAN                               MemoryProfileGettingStatus;
1589
1590   ContextData = GetMemoryProfileContext ();
1591   if (ContextData == NULL) {
1592     return EFI_UNSUPPORTED;
1593   }
1594
1595   MemoryProfileGettingStatus = mMemoryProfileGettingStatus;
1596   mMemoryProfileGettingStatus = TRUE;
1597
1598   Size = MemoryProfileGetDataSize ();
1599
1600   if (*ProfileSize < Size) {
1601     *ProfileSize = Size;
1602     mMemoryProfileGettingStatus = MemoryProfileGettingStatus;
1603     return EFI_BUFFER_TOO_SMALL;
1604   }
1605
1606   *ProfileSize = Size;
1607   MemoryProfileCopyData (ProfileBuffer);
1608
1609   mMemoryProfileGettingStatus = MemoryProfileGettingStatus;
1610   return EFI_SUCCESS;
1611 }
1612
1613 /**
1614   Register image to memory profile.
1615
1616   @param[in] This               The EDKII_MEMORY_PROFILE_PROTOCOL instance.
1617   @param[in] FilePath           File path of the image.
1618   @param[in] ImageBase          Image base address.
1619   @param[in] ImageSize          Image size.
1620   @param[in] FileType           File type of the image.
1621
1622   @return EFI_SUCCESS           Register successfully.
1623   @return EFI_UNSUPPORTED       Memory profile is unsupported,
1624                                 or memory profile for the image is not required.
1625   @return EFI_OUT_OF_RESOURCES  No enough resource for this register.
1626
1627 **/
1628 EFI_STATUS
1629 EFIAPI
1630 ProfileProtocolRegisterImage (
1631   IN EDKII_MEMORY_PROFILE_PROTOCOL  *This,
1632   IN EFI_DEVICE_PATH_PROTOCOL       *FilePath,
1633   IN PHYSICAL_ADDRESS               ImageBase,
1634   IN UINT64                         ImageSize,
1635   IN EFI_FV_FILETYPE                FileType
1636   )
1637 {
1638   EFI_STATUS                        Status;
1639   LOADED_IMAGE_PRIVATE_DATA         DriverEntry;
1640   VOID                              *EntryPointInImage;
1641
1642   ZeroMem (&DriverEntry, sizeof (DriverEntry));
1643   DriverEntry.Info.FilePath = FilePath;
1644   DriverEntry.ImageContext.ImageAddress = ImageBase;
1645   DriverEntry.ImageContext.ImageSize = ImageSize;
1646   Status = InternalPeCoffGetEntryPoint ((VOID *) (UINTN) ImageBase, &EntryPointInImage);
1647   ASSERT_EFI_ERROR (Status);
1648   DriverEntry.ImageContext.EntryPoint = (PHYSICAL_ADDRESS) (UINTN) EntryPointInImage;
1649   DriverEntry.ImageContext.ImageType = InternalPeCoffGetSubsystem ((VOID *) (UINTN) ImageBase);
1650
1651   return RegisterMemoryProfileImage (&DriverEntry, FileType);
1652 }
1653
1654 /**
1655   Unregister image from memory profile.
1656
1657   @param[in] This               The EDKII_MEMORY_PROFILE_PROTOCOL instance.
1658   @param[in] FilePath           File path of the image.
1659   @param[in] ImageBase          Image base address.
1660   @param[in] ImageSize          Image size.
1661
1662   @return EFI_SUCCESS           Unregister successfully.
1663   @return EFI_UNSUPPORTED       Memory profile is unsupported,
1664                                 or memory profile for the image is not required.
1665   @return EFI_NOT_FOUND         The image is not found.
1666
1667 **/
1668 EFI_STATUS
1669 EFIAPI
1670 ProfileProtocolUnregisterImage (
1671   IN EDKII_MEMORY_PROFILE_PROTOCOL  *This,
1672   IN EFI_DEVICE_PATH_PROTOCOL       *FilePath,
1673   IN PHYSICAL_ADDRESS               ImageBase,
1674   IN UINT64                         ImageSize
1675   )
1676 {
1677   EFI_STATUS                        Status;
1678   LOADED_IMAGE_PRIVATE_DATA         DriverEntry;
1679   VOID                              *EntryPointInImage;
1680
1681   ZeroMem (&DriverEntry, sizeof (DriverEntry));
1682   DriverEntry.Info.FilePath = FilePath;
1683   DriverEntry.ImageContext.ImageAddress = ImageBase;
1684   DriverEntry.ImageContext.ImageSize = ImageSize;
1685   Status = InternalPeCoffGetEntryPoint ((VOID *) (UINTN) ImageBase, &EntryPointInImage);
1686   ASSERT_EFI_ERROR (Status);
1687   DriverEntry.ImageContext.EntryPoint = (PHYSICAL_ADDRESS) (UINTN) EntryPointInImage;
1688
1689   return UnregisterMemoryProfileImage (&DriverEntry);
1690 }
1691
1692 /**
1693   Get memory profile recording state.
1694
1695   @param[in]  This              The EDKII_MEMORY_PROFILE_PROTOCOL instance.
1696   @param[out] RecordingState    Recording state.
1697
1698   @return EFI_SUCCESS           Memory profile recording state is returned.
1699   @return EFI_UNSUPPORTED       Memory profile is unsupported.
1700   @return EFI_INVALID_PARAMETER RecordingState is NULL.
1701
1702 **/
1703 EFI_STATUS
1704 EFIAPI
1705 ProfileProtocolGetRecordingState (
1706   IN EDKII_MEMORY_PROFILE_PROTOCOL      *This,
1707   OUT BOOLEAN                           *RecordingState
1708   )
1709 {
1710   MEMORY_PROFILE_CONTEXT_DATA       *ContextData;
1711
1712   ContextData = GetMemoryProfileContext ();
1713   if (ContextData == NULL) {
1714     return EFI_UNSUPPORTED;
1715   }
1716
1717   if (RecordingState == NULL) {
1718     return EFI_INVALID_PARAMETER;
1719   }
1720   *RecordingState = mMemoryProfileRecordingEnable;
1721   return EFI_SUCCESS;
1722 }
1723
1724 /**
1725   Set memory profile recording state.
1726
1727   @param[in] This               The EDKII_MEMORY_PROFILE_PROTOCOL instance.
1728   @param[in] RecordingState     Recording state.
1729
1730   @return EFI_SUCCESS           Set memory profile recording state successfully.
1731   @return EFI_UNSUPPORTED       Memory profile is unsupported.
1732
1733 **/
1734 EFI_STATUS
1735 EFIAPI
1736 ProfileProtocolSetRecordingState (
1737   IN EDKII_MEMORY_PROFILE_PROTOCOL      *This,
1738   IN BOOLEAN                            RecordingState
1739   )
1740 {
1741   MEMORY_PROFILE_CONTEXT_DATA       *ContextData;
1742
1743   ContextData = GetMemoryProfileContext ();
1744   if (ContextData == NULL) {
1745     return EFI_UNSUPPORTED;
1746   }
1747
1748   mMemoryProfileRecordingEnable = RecordingState;
1749   return EFI_SUCCESS;
1750 }
1751
1752 /**
1753   Record memory profile of multilevel caller.
1754
1755   @param[in] This               The EDKII_MEMORY_PROFILE_PROTOCOL instance.
1756   @param[in] CallerAddress      Address of caller.
1757   @param[in] Action             Memory profile action.
1758   @param[in] MemoryType         Memory type.
1759                                 EfiMaxMemoryType means the MemoryType is unknown.
1760   @param[in] Buffer             Buffer address.
1761   @param[in] Size               Buffer size.
1762   @param[in] ActionString       String for memory profile action.
1763                                 Only needed for user defined allocate action.
1764
1765   @return EFI_SUCCESS           Memory profile is updated.
1766   @return EFI_UNSUPPORTED       Memory profile is unsupported,
1767                                 or memory profile for the image is not required,
1768                                 or memory profile for the memory type is not required.
1769   @return EFI_ACCESS_DENIED     It is during memory profile data getting.
1770   @return EFI_ABORTED           Memory profile recording is not enabled.
1771   @return EFI_OUT_OF_RESOURCES  No enough resource to update memory profile for allocate action.
1772   @return EFI_NOT_FOUND         No matched allocate info found for free action.
1773
1774 **/
1775 EFI_STATUS
1776 EFIAPI
1777 ProfileProtocolRecord (
1778   IN EDKII_MEMORY_PROFILE_PROTOCOL      *This,
1779   IN PHYSICAL_ADDRESS                   CallerAddress,
1780   IN MEMORY_PROFILE_ACTION              Action,
1781   IN EFI_MEMORY_TYPE                    MemoryType,
1782   IN VOID                               *Buffer,
1783   IN UINTN                              Size,
1784   IN CHAR8                              *ActionString OPTIONAL
1785   )
1786 {
1787   return CoreUpdateProfile (CallerAddress, Action, MemoryType, Size, Buffer, ActionString);
1788 }
1789
1790 ////////////////////