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