MdeModulePkg/Core/Dxe/Mem/MemoryProfileRecord.c

   1 /** @file
   2   Support routines for UEFI memory profile.
   3
   4   Copyright (c) 2014, 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
  17 #define IS_UEFI_MEMORY_PROFILE_ENABLED ((PcdGet8 (PcdMemoryProfilePropertyMask) & BIT0) != 0)
  18
  19 typedef struct {
  20   UINT32                        Signature;
  21   MEMORY_PROFILE_CONTEXT        Context;
  22   LIST_ENTRY                    *DriverInfoList;
  23 } MEMORY_PROFILE_CONTEXT_DATA;
  24
  25 typedef struct {
  26   UINT32                        Signature;
  27   MEMORY_PROFILE_DRIVER_INFO    DriverInfo;
  28   LIST_ENTRY                    *AllocInfoList;
  29   LIST_ENTRY                    Link;
  30 } MEMORY_PROFILE_DRIVER_INFO_DATA;
  31
  32 typedef struct {
  33   UINT32                        Signature;
  34   MEMORY_PROFILE_ALLOC_INFO     AllocInfo;
  35   LIST_ENTRY                    Link;
  36 } MEMORY_PROFILE_ALLOC_INFO_DATA;
  37
  38
  39 GLOBAL_REMOVE_IF_UNREFERENCED LIST_ENTRY  mImageQueue = INITIALIZE_LIST_HEAD_VARIABLE (mImageQueue);
  40 GLOBAL_REMOVE_IF_UNREFERENCED MEMORY_PROFILE_CONTEXT_DATA mMemoryProfileContext = {
  41   MEMORY_PROFILE_CONTEXT_SIGNATURE,
  42   {
  43     {
  44       MEMORY_PROFILE_CONTEXT_SIGNATURE,
  45       sizeof (MEMORY_PROFILE_CONTEXT),
  46       MEMORY_PROFILE_CONTEXT_REVISION
  47     },
  48     0,
  49     0,
  50     {0},
  51     {0},
  52     0,
  53     0,
  54     0
  55   },
  56   &mImageQueue,
  57 };
  58 GLOBAL_REMOVE_IF_UNREFERENCED MEMORY_PROFILE_CONTEXT_DATA *mMemoryProfileContextPtr = NULL;
  59
  60 BOOLEAN mMemoryProfileRecordingStatus = FALSE;
  61
  62 /**
  63   Get memory profile data.
  64
  65   @param[in]      This              The EDKII_MEMORY_PROFILE_PROTOCOL instance.
  66   @param[in, out] ProfileSize       On entry, points to the size in bytes of the ProfileBuffer.
  67                                     On return, points to the size of the data returned in ProfileBuffer.
  68   @param[out]     ProfileBuffer     Profile buffer.
  69
  70   @return EFI_SUCCESS               Get the memory profile data successfully.
  71   @return EFI_BUFFER_TO_SMALL       The ProfileSize is too small for the resulting data.
  72                                     ProfileSize is updated with the size required.
  73
  74 **/
  75 EFI_STATUS
  76 EFIAPI
  77 ProfileProtocolGetData (
  78   IN     EDKII_MEMORY_PROFILE_PROTOCOL  *This,
  79   IN OUT UINT64                         *ProfileSize,
  80      OUT VOID                           *ProfileBuffer
  81   );
  82
  83 /**
  84   Register image to memory profile.
  85
  86   @param[in] This               The EDKII_MEMORY_PROFILE_PROTOCOL instance.
  87   @param[in] FilePath           File path of the image.
  88   @param[in] ImageBase          Image base address.
  89   @param[in] ImageSize          Image size.
  90   @param[in] FileType           File type of the image.
  91
  92   @return EFI_SUCCESS           Register success.
  93   @return EFI_OUT_OF_RESOURCE   No enough resource for this register.
  94
  95 **/
  96 EFI_STATUS
  97 EFIAPI
  98 ProfileProtocolRegisterImage (
  99   IN EDKII_MEMORY_PROFILE_PROTOCOL      *This,
 100   IN EFI_DEVICE_PATH_PROTOCOL           *FilePath,
 101   IN PHYSICAL_ADDRESS                   ImageBase,
 102   IN UINT64                             ImageSize,
 103   IN EFI_FV_FILETYPE                    FileType
 104   );
 105
 106 /**
 107   Unregister image from memory profile.
 108
 109   @param[in] This               The EDKII_MEMORY_PROFILE_PROTOCOL instance.
 110   @param[in] FilePath           File path of the image.
 111   @param[in] ImageBase          Image base address.
 112   @param[in] ImageSize          Image size.
 113
 114   @return EFI_SUCCESS           Unregister success.
 115   @return EFI_NOT_FOUND         The image is not found.
 116
 117 **/
 118 EFI_STATUS
 119 EFIAPI
 120 ProfileProtocolUnregisterImage (
 121   IN EDKII_MEMORY_PROFILE_PROTOCOL      *This,
 122   IN EFI_DEVICE_PATH_PROTOCOL           *FilePath,
 123   IN PHYSICAL_ADDRESS                   ImageBase,
 124   IN UINT64                             ImageSize
 125   );
 126
 127 EDKII_MEMORY_PROFILE_PROTOCOL mProfileProtocol = {
 128   ProfileProtocolGetData,
 129   ProfileProtocolRegisterImage,
 130   ProfileProtocolUnregisterImage
 131 };
 132
 133 /**
 134   Return memory profile context.
 135
 136   @return Memory profile context.
 137
 138 **/
 139 MEMORY_PROFILE_CONTEXT_DATA *
 140 GetMemoryProfileContext (
 141   VOID
 142   )
 143 {
 144   return mMemoryProfileContextPtr;
 145 }
 146
 147 /**
 148   Retrieves the magic value from the PE/COFF header.
 149
 150   @param Hdr    The buffer in which to return the PE32, PE32+, or TE header.
 151
 152   @return EFI_IMAGE_NT_OPTIONAL_HDR32_MAGIC - Image is PE32
 153   @return EFI_IMAGE_NT_OPTIONAL_HDR64_MAGIC - Image is PE32+
 154
 155 **/
 156 UINT16
 157 InternalPeCoffGetPeHeaderMagicValue (
 158   IN  EFI_IMAGE_OPTIONAL_HEADER_PTR_UNION  Hdr
 159   )
 160 {
 161   //
 162   // NOTE: Some versions of Linux ELILO for Itanium have an incorrect magic value
 163   //       in the PE/COFF Header.  If the MachineType is Itanium(IA64) and the
 164   //       Magic value in the OptionalHeader is  EFI_IMAGE_NT_OPTIONAL_HDR32_MAGIC
 165   //       then override the returned value to EFI_IMAGE_NT_OPTIONAL_HDR64_MAGIC
 166   //
 167   if (Hdr.Pe32->FileHeader.Machine == IMAGE_FILE_MACHINE_IA64 && Hdr.Pe32->OptionalHeader.Magic == EFI_IMAGE_NT_OPTIONAL_HDR32_MAGIC) {
 168     return EFI_IMAGE_NT_OPTIONAL_HDR64_MAGIC;
 169   }
 170   //
 171   // Return the magic value from the PC/COFF Optional Header
 172   //
 173   return Hdr.Pe32->OptionalHeader.Magic;
 174 }
 175
 176 /**
 177   Retrieves and returns the Subsystem of a PE/COFF image that has been loaded into system memory.
 178   If Pe32Data is NULL, then ASSERT().
 179
 180   @param Pe32Data   The pointer to the PE/COFF image that is loaded in system memory.
 181
 182   @return The Subsystem of the PE/COFF image.
 183
 184 **/
 185 UINT16
 186 InternalPeCoffGetSubsystem (
 187   IN VOID  *Pe32Data
 188   )
 189 {
 190   EFI_IMAGE_OPTIONAL_HEADER_PTR_UNION  Hdr;
 191   EFI_IMAGE_DOS_HEADER                 *DosHdr;
 192   UINT16                               Magic;
 193
 194   ASSERT (Pe32Data != NULL);
 195
 196   DosHdr = (EFI_IMAGE_DOS_HEADER *) Pe32Data;
 197   if (DosHdr->e_magic == EFI_IMAGE_DOS_SIGNATURE) {
 198     //
 199     // DOS image header is present, so read the PE header after the DOS image header.
 200     //
 201     Hdr.Pe32 = (EFI_IMAGE_NT_HEADERS32 *) ((UINTN) Pe32Data + (UINTN) ((DosHdr->e_lfanew) & 0x0ffff));
 202   } else {
 203     //
 204     // DOS image header is not present, so PE header is at the image base.
 205     //
 206     Hdr.Pe32 = (EFI_IMAGE_NT_HEADERS32 *) Pe32Data;
 207   }
 208
 209   if (Hdr.Te->Signature == EFI_TE_IMAGE_HEADER_SIGNATURE) {
 210     return Hdr.Te->Subsystem;
 211   } else if (Hdr.Pe32->Signature == EFI_IMAGE_NT_SIGNATURE)  {
 212     Magic = InternalPeCoffGetPeHeaderMagicValue (Hdr);
 213     if (Magic == EFI_IMAGE_NT_OPTIONAL_HDR32_MAGIC) {
 214       return Hdr.Pe32->OptionalHeader.Subsystem;
 215     } else if (Magic == EFI_IMAGE_NT_OPTIONAL_HDR64_MAGIC) {
 216       return Hdr.Pe32Plus->OptionalHeader.Subsystem;
 217     }
 218   }
 219
 220   return 0x0000;
 221 }
 222
 223 /**
 224   Retrieves and returns a pointer to the entry point to a PE/COFF image that has been loaded
 225   into system memory with the PE/COFF Loader Library functions.
 226
 227   Retrieves the entry point to the PE/COFF image specified by Pe32Data and returns this entry
 228   point in EntryPoint.  If the entry point could not be retrieved from the PE/COFF image, then
 229   return RETURN_INVALID_PARAMETER.  Otherwise return RETURN_SUCCESS.
 230   If Pe32Data is NULL, then ASSERT().
 231   If EntryPoint is NULL, then ASSERT().
 232
 233   @param  Pe32Data                  The pointer to the PE/COFF image that is loaded in system memory.
 234   @param  EntryPoint                The pointer to entry point to the PE/COFF image to return.
 235
 236   @retval RETURN_SUCCESS            EntryPoint was returned.
 237   @retval RETURN_INVALID_PARAMETER  The entry point could not be found in the PE/COFF image.
 238
 239 **/
 240 RETURN_STATUS
 241 InternalPeCoffGetEntryPoint (
 242   IN  VOID  *Pe32Data,
 243   OUT VOID  **EntryPoint
 244   )
 245 {
 246   EFI_IMAGE_DOS_HEADER                  *DosHdr;
 247   EFI_IMAGE_OPTIONAL_HEADER_PTR_UNION   Hdr;
 248
 249   ASSERT (Pe32Data   != NULL);
 250   ASSERT (EntryPoint != NULL);
 251
 252   DosHdr = (EFI_IMAGE_DOS_HEADER *) Pe32Data;
 253   if (DosHdr->e_magic == EFI_IMAGE_DOS_SIGNATURE) {
 254     //
 255     // DOS image header is present, so read the PE header after the DOS image header.
 256     //
 257     Hdr.Pe32 = (EFI_IMAGE_NT_HEADERS32 *) ((UINTN) Pe32Data + (UINTN) ((DosHdr->e_lfanew) & 0x0ffff));
 258   } else {
 259     //
 260     // DOS image header is not present, so PE header is at the image base.
 261     //
 262     Hdr.Pe32 = (EFI_IMAGE_NT_HEADERS32 *) Pe32Data;
 263   }
 264
 265   //
 266   // Calculate the entry point relative to the start of the image.
 267   // AddressOfEntryPoint is common for PE32 & PE32+
 268   //
 269   if (Hdr.Te->Signature == EFI_TE_IMAGE_HEADER_SIGNATURE) {
 270     *EntryPoint = (VOID *) ((UINTN) Pe32Data + (UINTN) (Hdr.Te->AddressOfEntryPoint & 0x0ffffffff) + sizeof (EFI_TE_IMAGE_HEADER) - Hdr.Te->StrippedSize);
 271     return RETURN_SUCCESS;
 272   } else if (Hdr.Pe32->Signature == EFI_IMAGE_NT_SIGNATURE) {
 273     *EntryPoint = (VOID *) ((UINTN) Pe32Data + (UINTN) (Hdr.Pe32->OptionalHeader.AddressOfEntryPoint & 0x0ffffffff));
 274     return RETURN_SUCCESS;
 275   }
 276
 277   return RETURN_UNSUPPORTED;
 278 }
 279
 280 /**
 281   Build driver info.
 282
 283   @param ContextData    Memory profile context.
 284   @param FileName       File name of the image.
 285   @param ImageBase      Image base address.
 286   @param ImageSize      Image size.
 287   @param EntryPoint     Entry point of the image.
 288   @param ImageSubsystem Image subsystem of the image.
 289   @param FileType       File type of the image.
 290
 291   @return Pointer to memory profile driver info.
 292
 293 **/
 294 MEMORY_PROFILE_DRIVER_INFO_DATA *
 295 BuildDriverInfo (
 296   IN MEMORY_PROFILE_CONTEXT_DATA    *ContextData,
 297   IN EFI_GUID                       *FileName,
 298   IN PHYSICAL_ADDRESS               ImageBase,
 299   IN UINT64                         ImageSize,
 300   IN PHYSICAL_ADDRESS               EntryPoint,
 301   IN UINT16                         ImageSubsystem,
 302   IN EFI_FV_FILETYPE                FileType
 303   )
 304 {
 305   EFI_STATUS                        Status;
 306   MEMORY_PROFILE_DRIVER_INFO        *DriverInfo;
 307   MEMORY_PROFILE_DRIVER_INFO_DATA   *DriverInfoData;
 308   VOID                              *EntryPointInImage;
 309
 310   //
 311   // Use CoreInternalAllocatePool() that will not update profile for this AllocatePool action.
 312   //
 313   Status = CoreInternalAllocatePool (
 314              EfiBootServicesData,
 315              sizeof (*DriverInfoData) + sizeof (LIST_ENTRY),
 316              (VOID **) &DriverInfoData
 317              );
 318   if (EFI_ERROR (Status)) {
 319     return NULL;
 320   }
 321
 322   ZeroMem (DriverInfoData, sizeof (*DriverInfoData));
 323
 324   DriverInfo = &DriverInfoData->DriverInfo;
 325   DriverInfoData->Signature = MEMORY_PROFILE_DRIVER_INFO_SIGNATURE;
 326   DriverInfo->Header.Signature = MEMORY_PROFILE_DRIVER_INFO_SIGNATURE;
 327   DriverInfo->Header.Length = sizeof (MEMORY_PROFILE_DRIVER_INFO);
 328   DriverInfo->Header.Revision = MEMORY_PROFILE_DRIVER_INFO_REVISION;
 329   if (FileName != NULL) {
 330     CopyMem (&DriverInfo->FileName, FileName, sizeof (EFI_GUID));
 331   }
 332   DriverInfo->ImageBase = ImageBase;
 333   DriverInfo->ImageSize = ImageSize;
 334   DriverInfo->EntryPoint = EntryPoint;
 335   DriverInfo->ImageSubsystem = ImageSubsystem;
 336   if ((EntryPoint != 0) && ((EntryPoint < ImageBase) || (EntryPoint >= (ImageBase + ImageSize)))) {
 337     //
 338     // If the EntryPoint is not in the range of image buffer, it should come from emulation environment.
 339     // So patch ImageBuffer here to align the EntryPoint.
 340     //
 341     Status = InternalPeCoffGetEntryPoint ((VOID *) (UINTN) ImageBase, &EntryPointInImage);
 342     ASSERT_EFI_ERROR (Status);
 343     DriverInfo->ImageBase = ImageBase + EntryPoint - (PHYSICAL_ADDRESS) (UINTN) EntryPointInImage;
 344   }
 345   DriverInfo->FileType = FileType;
 346   DriverInfoData->AllocInfoList = (LIST_ENTRY *) (DriverInfoData + 1);
 347   InitializeListHead (DriverInfoData->AllocInfoList);
 348   DriverInfo->CurrentUsage = 0;
 349   DriverInfo->PeakUsage = 0;
 350   DriverInfo->AllocRecordCount = 0;
 351
 352   InsertTailList (ContextData->DriverInfoList, &DriverInfoData->Link);
 353   ContextData->Context.ImageCount ++;
 354   ContextData->Context.TotalImageSize += DriverInfo->ImageSize;
 355
 356   return DriverInfoData;
 357 }
 358
 359 /**
 360   Register DXE Core to memory profile.
 361
 362   @param HobStart       The start address of the HOB.
 363   @param ContextData    Memory profile context.
 364
 365   @retval TRUE      Register success.
 366   @retval FALSE     Register fail.
 367
 368 **/
 369 BOOLEAN
 370 RegisterDxeCore (
 371   IN VOID                           *HobStart,
 372   IN MEMORY_PROFILE_CONTEXT_DATA    *ContextData
 373   )
 374 {
 375   EFI_PEI_HOB_POINTERS              DxeCoreHob;
 376   MEMORY_PROFILE_DRIVER_INFO_DATA   *DriverInfoData;
 377   PHYSICAL_ADDRESS                  ImageBase;
 378
 379   ASSERT (ContextData != NULL);
 380
 381   //
 382   // Searching for image hob
 383   //
 384   DxeCoreHob.Raw          = HobStart;
 385   while ((DxeCoreHob.Raw = GetNextHob (EFI_HOB_TYPE_MEMORY_ALLOCATION, DxeCoreHob.Raw)) != NULL) {
 386     if (CompareGuid (&DxeCoreHob.MemoryAllocationModule->MemoryAllocationHeader.Name, &gEfiHobMemoryAllocModuleGuid)) {
 387       //
 388       // Find Dxe Core HOB
 389       //
 390       break;
 391     }
 392     DxeCoreHob.Raw = GET_NEXT_HOB (DxeCoreHob);
 393   }
 394   ASSERT (DxeCoreHob.Raw != NULL);
 395
 396   ImageBase = DxeCoreHob.MemoryAllocationModule->MemoryAllocationHeader.MemoryBaseAddress;
 397   DriverInfoData = BuildDriverInfo (
 398                      ContextData,
 399                      &DxeCoreHob.MemoryAllocationModule->ModuleName,
 400                      ImageBase,
 401                      DxeCoreHob.MemoryAllocationModule->MemoryAllocationHeader.MemoryLength,
 402                      DxeCoreHob.MemoryAllocationModule->EntryPoint,
 403                      InternalPeCoffGetSubsystem ((VOID *) (UINTN) ImageBase),
 404                      EFI_FV_FILETYPE_DXE_CORE
 405                      );
 406   if (DriverInfoData == NULL) {
 407     return FALSE;
 408   }
 409
 410   return TRUE;
 411 }
 412
 413 /**
 414   Initialize memory profile.
 415
 416   @param HobStart   The start address of the HOB.
 417
 418 **/
 419 VOID
 420 MemoryProfileInit (
 421   IN VOID   *HobStart
 422   )
 423 {
 424   MEMORY_PROFILE_CONTEXT_DATA   *ContextData;
 425
 426   if (!IS_UEFI_MEMORY_PROFILE_ENABLED) {
 427     return;
 428   }
 429
 430   ContextData = GetMemoryProfileContext ();
 431   if (ContextData != NULL) {
 432     return;
 433   }
 434
 435   mMemoryProfileRecordingStatus = TRUE;
 436   mMemoryProfileContextPtr = &mMemoryProfileContext;
 437
 438   RegisterDxeCore (HobStart, &mMemoryProfileContext);
 439
 440   DEBUG ((EFI_D_INFO, "MemoryProfileInit MemoryProfileContext - 0x%x\n", &mMemoryProfileContext));
 441 }
 442
 443 /**
 444   Install memory profile protocol.
 445
 446 **/
 447 VOID
 448 MemoryProfileInstallProtocol (
 449   VOID
 450   )
 451 {
 452   EFI_HANDLE    Handle;
 453   EFI_STATUS    Status;
 454
 455   if (!IS_UEFI_MEMORY_PROFILE_ENABLED) {
 456     return;
 457   }
 458
 459   Handle = NULL;
 460   Status = CoreInstallMultipleProtocolInterfaces (
 461              &Handle,
 462              &gEdkiiMemoryProfileGuid,
 463              &mProfileProtocol,
 464              NULL
 465              );
 466   ASSERT_EFI_ERROR (Status);
 467 }
 468
 469 /**
 470   Get the GUID file name from the file path.
 471
 472   @param FilePath  File path.
 473
 474   @return The GUID file name from the file path.
 475
 476 **/
 477 EFI_GUID *
 478 GetFileNameFromFilePath (
 479   IN EFI_DEVICE_PATH_PROTOCOL   *FilePath
 480   )
 481 {
 482   MEDIA_FW_VOL_FILEPATH_DEVICE_PATH     *ThisFilePath;
 483   EFI_GUID                              *FileName;
 484
 485   FileName = NULL;
 486   ThisFilePath = (MEDIA_FW_VOL_FILEPATH_DEVICE_PATH *) FilePath;
 487   while (!IsDevicePathEnd (ThisFilePath)) {
 488     FileName = EfiGetNameGuidFromFwVolDevicePathNode (ThisFilePath);
 489     if (FileName != NULL) {
 490       break;
 491     }
 492     ThisFilePath = (MEDIA_FW_VOL_FILEPATH_DEVICE_PATH *) NextDevicePathNode (ThisFilePath);
 493   }
 494
 495   return FileName;
 496 }
 497
 498 /**
 499   Register image to memory profile.
 500
 501   @param DriverEntry    Image info.
 502   @param FileType       Image file type.
 503
 504   @retval TRUE          Register success.
 505   @retval FALSE         Register fail.
 506
 507 **/
 508 BOOLEAN
 509 RegisterMemoryProfileImage (
 510   IN LOADED_IMAGE_PRIVATE_DATA  *DriverEntry,
 511   IN EFI_FV_FILETYPE            FileType
 512   )
 513 {
 514   MEMORY_PROFILE_CONTEXT_DATA       *ContextData;
 515   MEMORY_PROFILE_DRIVER_INFO_DATA   *DriverInfoData;
 516
 517   if (!IS_UEFI_MEMORY_PROFILE_ENABLED) {
 518     return FALSE;
 519   }
 520
 521   ContextData = GetMemoryProfileContext ();
 522   if (ContextData == NULL) {
 523     return FALSE;
 524   }
 525
 526   DriverInfoData = BuildDriverInfo (
 527                      ContextData,
 528                      GetFileNameFromFilePath (DriverEntry->Info.FilePath),
 529                      DriverEntry->ImageContext.ImageAddress,
 530                      DriverEntry->ImageContext.ImageSize,
 531                      DriverEntry->ImageContext.EntryPoint,
 532                      DriverEntry->ImageContext.ImageType,
 533                      FileType
 534                      );
 535   if (DriverInfoData == NULL) {
 536     return FALSE;
 537   }
 538
 539   return TRUE;
 540 }
 541
 542 /**
 543   Search image from memory profile.
 544
 545   @param ContextData    Memory profile context.
 546   @param FileName       Image file name.
 547   @param Address        Image Address.
 548
 549   @return Pointer to memory profile driver info.
 550
 551 **/
 552 MEMORY_PROFILE_DRIVER_INFO_DATA *
 553 GetMemoryProfileDriverInfoByFileNameAndAddress (
 554   IN MEMORY_PROFILE_CONTEXT_DATA    *ContextData,
 555   IN EFI_GUID                       *FileName,
 556   IN PHYSICAL_ADDRESS               Address
 557   )
 558 {
 559   MEMORY_PROFILE_DRIVER_INFO        *DriverInfo;
 560   MEMORY_PROFILE_DRIVER_INFO_DATA   *DriverInfoData;
 561   LIST_ENTRY                        *DriverLink;
 562   LIST_ENTRY                        *DriverInfoList;
 563
 564   DriverInfoList = ContextData->DriverInfoList;
 565
 566   for (DriverLink = DriverInfoList->ForwardLink;
 567        DriverLink != DriverInfoList;
 568        DriverLink = DriverLink->ForwardLink) {
 569     DriverInfoData = CR (
 570                        DriverLink,
 571                        MEMORY_PROFILE_DRIVER_INFO_DATA,
 572                        Link,
 573                        MEMORY_PROFILE_DRIVER_INFO_SIGNATURE
 574                        );
 575     DriverInfo = &DriverInfoData->DriverInfo;
 576     if ((CompareGuid (&DriverInfo->FileName, FileName)) &&
 577         (Address >= DriverInfo->ImageBase) &&
 578         (Address < (DriverInfo->ImageBase + DriverInfo->ImageSize))) {
 579       return DriverInfoData;
 580     }
 581   }
 582
 583   return NULL;
 584 }
 585
 586 /**
 587   Search dummy image from memory profile.
 588
 589   @param ContextData    Memory profile context.
 590
 591   @return Pointer to memory profile driver info.
 592
 593 **/
 594 MEMORY_PROFILE_DRIVER_INFO_DATA *
 595 FindDummyImage (
 596   IN MEMORY_PROFILE_CONTEXT_DATA    *ContextData
 597   )
 598 {
 599   MEMORY_PROFILE_DRIVER_INFO_DATA   *DriverInfoData;
 600   LIST_ENTRY                        *DriverLink;
 601   LIST_ENTRY                        *DriverInfoList;
 602
 603   DriverInfoList = ContextData->DriverInfoList;
 604
 605   for (DriverLink = DriverInfoList->ForwardLink;
 606        DriverLink != DriverInfoList;
 607        DriverLink = DriverLink->ForwardLink) {
 608     DriverInfoData = CR (
 609                    DriverLink,
 610                    MEMORY_PROFILE_DRIVER_INFO_DATA,
 611                    Link,
 612                    MEMORY_PROFILE_DRIVER_INFO_SIGNATURE
 613                    );
 614     if (CompareGuid (&gZeroGuid, &DriverInfoData->DriverInfo.FileName)) {
 615       return DriverInfoData;
 616     }
 617   }
 618
 619   return BuildDriverInfo (ContextData, &gZeroGuid, 0, 0, 0, 0, 0);
 620 }
 621
 622 /**
 623   Search image from memory profile.
 624   It will return image, if (Address >= ImageBuffer) AND (Address < ImageBuffer + ImageSize)
 625
 626   @param ContextData    Memory profile context.
 627   @param Address        Image or Function address.
 628
 629   @return Pointer to memory profile driver info.
 630
 631 **/
 632 MEMORY_PROFILE_DRIVER_INFO_DATA *
 633 GetMemoryProfileDriverInfoFromAddress (
 634   IN MEMORY_PROFILE_CONTEXT_DATA    *ContextData,
 635   IN PHYSICAL_ADDRESS               Address
 636   )
 637 {
 638   MEMORY_PROFILE_DRIVER_INFO        *DriverInfo;
 639   MEMORY_PROFILE_DRIVER_INFO_DATA   *DriverInfoData;
 640   LIST_ENTRY                        *DriverLink;
 641   LIST_ENTRY                        *DriverInfoList;
 642
 643   DriverInfoList = ContextData->DriverInfoList;
 644
 645   for (DriverLink = DriverInfoList->ForwardLink;
 646        DriverLink != DriverInfoList;
 647        DriverLink = DriverLink->ForwardLink) {
 648     DriverInfoData = CR (
 649                        DriverLink,
 650                        MEMORY_PROFILE_DRIVER_INFO_DATA,
 651                        Link,
 652                        MEMORY_PROFILE_DRIVER_INFO_SIGNATURE
 653                        );
 654     DriverInfo = &DriverInfoData->DriverInfo;
 655     if ((Address >= DriverInfo->ImageBase) &&
 656         (Address < (DriverInfo->ImageBase + DriverInfo->ImageSize))) {
 657       return DriverInfoData;
 658     }
 659   }
 660
 661   //
 662   // Should never come here.
 663   //
 664   return FindDummyImage (ContextData);
 665 }
 666
 667 /**
 668   Unregister image from memory profile.
 669
 670   @param DriverEntry    Image info.
 671
 672   @retval TRUE          Unregister success.
 673   @retval FALSE         Unregister fail.
 674
 675 **/
 676 BOOLEAN
 677 UnregisterMemoryProfileImage (
 678   IN LOADED_IMAGE_PRIVATE_DATA      *DriverEntry
 679   )
 680 {
 681   EFI_STATUS                        Status;
 682   MEMORY_PROFILE_CONTEXT_DATA       *ContextData;
 683   MEMORY_PROFILE_DRIVER_INFO_DATA   *DriverInfoData;
 684   EFI_GUID                          *FileName;
 685   PHYSICAL_ADDRESS                  ImageAddress;
 686   VOID                              *EntryPointInImage;
 687
 688   if (!IS_UEFI_MEMORY_PROFILE_ENABLED) {
 689     return FALSE;
 690   }
 691
 692   ContextData = GetMemoryProfileContext ();
 693   if (ContextData == NULL) {
 694     return FALSE;
 695   }
 696
 697   DriverInfoData = NULL;
 698   FileName = GetFileNameFromFilePath (DriverEntry->Info.FilePath);
 699   ImageAddress = DriverEntry->ImageContext.ImageAddress;
 700   if ((DriverEntry->ImageContext.EntryPoint < ImageAddress) || (DriverEntry->ImageContext.EntryPoint >= (ImageAddress + DriverEntry->ImageContext.ImageSize))) {
 701     //
 702     // If the EntryPoint is not in the range of image buffer, it should come from emulation environment.
 703     // So patch ImageAddress here to align the EntryPoint.
 704     //
 705     Status = InternalPeCoffGetEntryPoint ((VOID *) (UINTN) ImageAddress, &EntryPointInImage);
 706     ASSERT_EFI_ERROR (Status);
 707     ImageAddress = ImageAddress + (UINTN) DriverEntry->ImageContext.EntryPoint - (UINTN) EntryPointInImage;
 708   }
 709   if (FileName != NULL) {
 710     DriverInfoData = GetMemoryProfileDriverInfoByFileNameAndAddress (ContextData, FileName, ImageAddress);
 711   }
 712   if (DriverInfoData == NULL) {
 713     DriverInfoData = GetMemoryProfileDriverInfoFromAddress (ContextData, ImageAddress);
 714   }
 715   if (DriverInfoData == NULL) {
 716     return FALSE;
 717   }
 718
 719   ContextData->Context.TotalImageSize -= DriverInfoData->DriverInfo.ImageSize;
 720
 721   DriverInfoData->DriverInfo.ImageBase = 0;
 722   DriverInfoData->DriverInfo.ImageSize = 0;
 723
 724   if (DriverInfoData->DriverInfo.PeakUsage == 0) {
 725     ContextData->Context.ImageCount --;
 726     RemoveEntryList (&DriverInfoData->Link);
 727     //
 728     // Use CoreInternalFreePool() that will not update profile for this FreePool action.
 729     //
 730     CoreInternalFreePool (DriverInfoData);
 731   }
 732
 733   return TRUE;
 734 }
 735
 736 /**
 737   Return if this memory type needs to be recorded into memory profile.
 738   If BIOS memory type (0 ~ EfiMaxMemoryType), it checks bit (1 << MemoryType).
 739   If OS memory type (0x80000000 ~ 0xFFFFFFFF), it checks bit63 - 0x8000000000000000.
 740
 741   @param MemoryType     Memory type.
 742
 743   @retval TRUE          This memory type need to be recorded.
 744   @retval FALSE         This memory type need not to be recorded.
 745
 746 **/
 747 BOOLEAN
 748 CoreNeedRecordProfile (
 749   IN EFI_MEMORY_TYPE    MemoryType
 750   )
 751 {
 752   UINT64 TestBit;
 753
 754   if ((UINT32) MemoryType >= 0x80000000) {
 755     TestBit = BIT63;
 756   } else {
 757     TestBit = LShiftU64 (1, MemoryType);
 758   }
 759
 760   if ((PcdGet64 (PcdMemoryProfileMemoryType) & TestBit) != 0) {
 761     return TRUE;
 762   } else {
 763     return FALSE;
 764   }
 765 }
 766
 767 /**
 768   Convert EFI memory type to profile memory index. The rule is:
 769   If BIOS memory type (0 ~ EfiMaxMemoryType), ProfileMemoryIndex = MemoryType.
 770   If OS memory type (0x80000000 ~ 0xFFFFFFFF), ProfileMemoryIndex = EfiMaxMemoryType.
 771
 772   @param MemoryType     Memory type.
 773
 774   @return EFI memory type as profile memory index.
 775
 776 **/
 777 EFI_MEMORY_TYPE
 778 GetProfileMemoryIndex (
 779   IN EFI_MEMORY_TYPE    MemoryType
 780   )
 781 {
 782   if ((UINT32) MemoryType >= 0x80000000) {
 783     return EfiMaxMemoryType;
 784   } else {
 785     return MemoryType;
 786   }
 787 }
 788
 789 /**
 790   Update memory profile Allocate information.
 791
 792   @param CallerAddress  Address of caller who call Allocate.
 793   @param Action         This Allocate action.
 794   @param MemoryType     Memory type.
 795   @param Size           Buffer size.
 796   @param Buffer         Buffer address.
 797
 798   @retval TRUE          Profile udpate success.
 799   @retval FALSE         Profile update fail.
 800
 801 **/
 802 BOOLEAN
 803 CoreUpdateProfileAllocate (
 804   IN PHYSICAL_ADDRESS       CallerAddress,
 805   IN MEMORY_PROFILE_ACTION  Action,
 806   IN EFI_MEMORY_TYPE        MemoryType,
 807   IN UINTN                  Size,
 808   IN VOID                   *Buffer
 809   )
 810 {
 811   EFI_STATUS                        Status;
 812   MEMORY_PROFILE_CONTEXT           *Context;
 813   MEMORY_PROFILE_DRIVER_INFO       *DriverInfo;
 814   MEMORY_PROFILE_ALLOC_INFO        *AllocInfo;
 815   MEMORY_PROFILE_CONTEXT_DATA       *ContextData;
 816   MEMORY_PROFILE_DRIVER_INFO_DATA   *DriverInfoData;
 817   MEMORY_PROFILE_ALLOC_INFO_DATA    *AllocInfoData;
 818   EFI_MEMORY_TYPE                   ProfileMemoryIndex;
 819
 820   AllocInfoData = NULL;
 821
 822   ContextData = GetMemoryProfileContext ();
 823   if (ContextData == NULL) {
 824     return FALSE;
 825   }
 826
 827   DriverInfoData = GetMemoryProfileDriverInfoFromAddress (ContextData, CallerAddress);
 828   ASSERT (DriverInfoData != NULL);
 829
 830   //
 831   // Use CoreInternalAllocatePool() that will not update profile for this AllocatePool action.
 832   //
 833   Status = CoreInternalAllocatePool (
 834              EfiBootServicesData,
 835              sizeof (*AllocInfoData),
 836              (VOID **) &AllocInfoData
 837              );
 838   if (EFI_ERROR (Status)) {
 839     return FALSE;
 840   }
 841   AllocInfo = &AllocInfoData->AllocInfo;
 842   AllocInfoData->Signature      = MEMORY_PROFILE_ALLOC_INFO_SIGNATURE;
 843   AllocInfo->Header.Signature   = MEMORY_PROFILE_ALLOC_INFO_SIGNATURE;
 844   AllocInfo->Header.Length      = sizeof (MEMORY_PROFILE_ALLOC_INFO);
 845   AllocInfo->Header.Revision    = MEMORY_PROFILE_ALLOC_INFO_REVISION;
 846   AllocInfo->CallerAddress      = CallerAddress;
 847   AllocInfo->SequenceId         = ContextData->Context.SequenceCount;
 848   AllocInfo->Action             = Action;
 849   AllocInfo->MemoryType         = MemoryType;
 850   AllocInfo->Buffer             = (PHYSICAL_ADDRESS) (UINTN) Buffer;
 851   AllocInfo->Size               = Size;
 852
 853   InsertTailList (DriverInfoData->AllocInfoList, &AllocInfoData->Link);
 854
 855   ProfileMemoryIndex = GetProfileMemoryIndex (MemoryType);
 856
 857   DriverInfo = &DriverInfoData->DriverInfo;
 858   DriverInfo->CurrentUsage += Size;
 859   if (DriverInfo->PeakUsage < DriverInfo->CurrentUsage) {
 860     DriverInfo->PeakUsage = DriverInfo->CurrentUsage;
 861   }
 862   DriverInfo->CurrentUsageByType[ProfileMemoryIndex] += Size;
 863   if (DriverInfo->PeakUsageByType[ProfileMemoryIndex] < DriverInfo->CurrentUsageByType[ProfileMemoryIndex]) {
 864     DriverInfo->PeakUsageByType[ProfileMemoryIndex] = DriverInfo->CurrentUsageByType[ProfileMemoryIndex];
 865   }
 866   DriverInfo->AllocRecordCount ++;
 867
 868   Context = &ContextData->Context;
 869   Context->CurrentTotalUsage += Size;
 870   if (Context->PeakTotalUsage < Context->CurrentTotalUsage) {
 871     Context->PeakTotalUsage = Context->CurrentTotalUsage;
 872   }
 873   Context->CurrentTotalUsageByType[ProfileMemoryIndex] += Size;
 874   if (Context->PeakTotalUsageByType[ProfileMemoryIndex] < Context->CurrentTotalUsageByType[ProfileMemoryIndex]) {
 875     Context->PeakTotalUsageByType[ProfileMemoryIndex] = Context->CurrentTotalUsageByType[ProfileMemoryIndex];
 876   }
 877   Context->SequenceCount ++;
 878
 879   return TRUE;
 880 }
 881
 882 /**
 883   Get memory profile alloc info from memory profile
 884
 885   @param DriverInfoData     Driver info
 886   @param Action             This Free action
 887   @param Size               Buffer size
 888   @param Buffer             Buffer address
 889
 890   @return Pointer to memory profile alloc info.
 891 **/
 892 MEMORY_PROFILE_ALLOC_INFO_DATA *
 893 GetMemoryProfileAllocInfoFromAddress (
 894   IN MEMORY_PROFILE_DRIVER_INFO_DATA    *DriverInfoData,
 895   IN MEMORY_PROFILE_ACTION              Action,
 896   IN UINTN                              Size,
 897   IN VOID                               *Buffer
 898   )
 899 {
 900   LIST_ENTRY                        *AllocInfoList;
 901   LIST_ENTRY                        *AllocLink;
 902   MEMORY_PROFILE_ALLOC_INFO         *AllocInfo;
 903   MEMORY_PROFILE_ALLOC_INFO_DATA    *AllocInfoData;
 904
 905   AllocInfoList = DriverInfoData->AllocInfoList;
 906
 907   for (AllocLink = AllocInfoList->ForwardLink;
 908        AllocLink != AllocInfoList;
 909        AllocLink = AllocLink->ForwardLink) {
 910     AllocInfoData = CR (
 911                       AllocLink,
 912                       MEMORY_PROFILE_ALLOC_INFO_DATA,
 913                       Link,
 914                       MEMORY_PROFILE_ALLOC_INFO_SIGNATURE
 915                       );
 916     AllocInfo = &AllocInfoData->AllocInfo;
 917     if (AllocInfo->Action != Action) {
 918       continue;
 919     }
 920     switch (Action) {
 921       case MemoryProfileActionAllocatePages:
 922         if ((AllocInfo->Buffer <= (PHYSICAL_ADDRESS) (UINTN) Buffer) &&
 923             ((AllocInfo->Buffer + AllocInfo->Size) >= ((PHYSICAL_ADDRESS) (UINTN) Buffer + Size))) {
 924           return AllocInfoData;
 925         }
 926         break;
 927       case MemoryProfileActionAllocatePool:
 928         if (AllocInfo->Buffer == (PHYSICAL_ADDRESS) (UINTN) Buffer) {
 929           return AllocInfoData;
 930         }
 931         break;
 932       default:
 933         ASSERT (FALSE);
 934         break;
 935     }
 936   }
 937
 938   return NULL;
 939 }
 940
 941 /**
 942   Update memory profile Free information.
 943
 944   @param CallerAddress  Address of caller who call Free.
 945   @param Action         This Free action.
 946   @param Size           Buffer size.
 947   @param Buffer         Buffer address.
 948
 949   @retval TRUE          Profile udpate success.
 950   @retval FALSE         Profile update fail.
 951
 952 **/
 953 BOOLEAN
 954 CoreUpdateProfileFree (
 955   IN PHYSICAL_ADDRESS       CallerAddress,
 956   IN MEMORY_PROFILE_ACTION  Action,
 957   IN UINTN                  Size,
 958   IN VOID                   *Buffer
 959   )
 960 {
 961   MEMORY_PROFILE_CONTEXT           *Context;
 962   MEMORY_PROFILE_DRIVER_INFO       *DriverInfo;
 963   MEMORY_PROFILE_ALLOC_INFO        *AllocInfo;
 964   MEMORY_PROFILE_CONTEXT_DATA      *ContextData;
 965   MEMORY_PROFILE_DRIVER_INFO_DATA  *DriverInfoData;
 966   LIST_ENTRY                       *DriverLink;
 967   LIST_ENTRY                       *DriverInfoList;
 968   MEMORY_PROFILE_DRIVER_INFO_DATA  *ThisDriverInfoData;
 969   MEMORY_PROFILE_ALLOC_INFO_DATA   *AllocInfoData;
 970   EFI_MEMORY_TYPE                  ProfileMemoryIndex;
 971
 972   ContextData = GetMemoryProfileContext ();
 973   if (ContextData == NULL) {
 974     return FALSE;
 975   }
 976
 977   DriverInfoData = GetMemoryProfileDriverInfoFromAddress (ContextData, CallerAddress);
 978   ASSERT (DriverInfoData != NULL);
 979
 980   switch (Action) {
 981     case MemoryProfileActionFreePages:
 982       AllocInfoData = GetMemoryProfileAllocInfoFromAddress (DriverInfoData, MemoryProfileActionAllocatePages, Size, Buffer);
 983       break;
 984     case MemoryProfileActionFreePool:
 985       AllocInfoData = GetMemoryProfileAllocInfoFromAddress (DriverInfoData, MemoryProfileActionAllocatePool, 0, Buffer);
 986       break;
 987     default:
 988       ASSERT (FALSE);
 989       AllocInfoData = NULL;
 990       break;
 991   }
 992   if (AllocInfoData == NULL) {
 993     //
 994     // Legal case, because driver A might free memory allocated by driver B, by some protocol.
 995     //
 996     DriverInfoList = ContextData->DriverInfoList;
 997
 998     for (DriverLink = DriverInfoList->ForwardLink;
 999          DriverLink != DriverInfoList;
1000          DriverLink = DriverLink->ForwardLink) {
1001       ThisDriverInfoData = CR (
1002                              DriverLink,
1003                              MEMORY_PROFILE_DRIVER_INFO_DATA,
1004                              Link,
1005                              MEMORY_PROFILE_DRIVER_INFO_SIGNATURE
1006                              );
1007       switch (Action) {
1008         case MemoryProfileActionFreePages:
1009           AllocInfoData = GetMemoryProfileAllocInfoFromAddress (ThisDriverInfoData, MemoryProfileActionAllocatePages, Size, Buffer);
1010           break;
1011         case MemoryProfileActionFreePool:
1012           AllocInfoData = GetMemoryProfileAllocInfoFromAddress (ThisDriverInfoData, MemoryProfileActionAllocatePool, 0, Buffer);
1013           break;
1014         default:
1015           ASSERT (FALSE);
1016           AllocInfoData = NULL;
1017           break;
1018       }
1019       if (AllocInfoData != NULL) {
1020         DriverInfoData = ThisDriverInfoData;
1021         break;
1022       }
1023     }
1024
1025     if (AllocInfoData == NULL) {
1026       //
1027       // No matched allocate operation is found for this free operation.
1028       // It is because the specified memory type allocate operation has been
1029       // filtered by CoreNeedRecordProfile(), but free operations have no
1030       // memory type information, they can not be filtered by CoreNeedRecordProfile().
1031       // Then, they will be filtered here.
1032       //
1033       return FALSE;
1034     }
1035   }
1036
1037   Context = &ContextData->Context;
1038   DriverInfo = &DriverInfoData->DriverInfo;
1039   AllocInfo = &AllocInfoData->AllocInfo;
1040
1041   ProfileMemoryIndex = GetProfileMemoryIndex (AllocInfo->MemoryType);
1042
1043   Context->CurrentTotalUsage -= AllocInfo->Size;
1044   Context->CurrentTotalUsageByType[ProfileMemoryIndex] -= AllocInfo->Size;
1045
1046   DriverInfo->CurrentUsage -= AllocInfo->Size;
1047   DriverInfo->CurrentUsageByType[ProfileMemoryIndex] -= AllocInfo->Size;
1048   DriverInfo->AllocRecordCount --;
1049
1050   RemoveEntryList (&AllocInfoData->Link);
1051
1052   if (Action == MemoryProfileActionFreePages) {
1053     if (AllocInfo->Buffer != (PHYSICAL_ADDRESS) (UINTN) Buffer) {
1054       CoreUpdateProfileAllocate (
1055         AllocInfo->CallerAddress,
1056         MemoryProfileActionAllocatePages,
1057         AllocInfo->MemoryType,
1058         (UINTN) ((PHYSICAL_ADDRESS) (UINTN) Buffer - AllocInfo->Buffer),
1059         (VOID *) (UINTN) AllocInfo->Buffer
1060         );
1061     }
1062     if (AllocInfo->Buffer + AllocInfo->Size != ((PHYSICAL_ADDRESS) (UINTN) Buffer + Size)) {
1063       CoreUpdateProfileAllocate (
1064         AllocInfo->CallerAddress,
1065         MemoryProfileActionAllocatePages,
1066         AllocInfo->MemoryType,
1067         (UINTN) ((AllocInfo->Buffer + AllocInfo->Size) - ((PHYSICAL_ADDRESS) (UINTN) Buffer + Size)),
1068         (VOID *) ((UINTN) Buffer + Size)
1069         );
1070     }
1071   }
1072
1073   //
1074   // Use CoreInternalFreePool() that will not update profile for this FreePool action.
1075   //
1076   CoreInternalFreePool (AllocInfoData);
1077
1078   return TRUE;
1079 }
1080
1081 /**
1082   Update memory profile information.
1083
1084   @param CallerAddress  Address of caller who call Allocate or Free.
1085   @param Action         This Allocate or Free action.
1086   @param MemoryType     Memory type.
1087   @param Size           Buffer size.
1088   @param Buffer         Buffer address.
1089
1090   @retval TRUE          Profile udpate success.
1091   @retval FALSE         Profile update fail.
1092
1093 **/
1094 BOOLEAN
1095 CoreUpdateProfile (
1096   IN PHYSICAL_ADDRESS       CallerAddress,
1097   IN MEMORY_PROFILE_ACTION  Action,
1098   IN EFI_MEMORY_TYPE        MemoryType, // Valid for AllocatePages/AllocatePool
1099   IN UINTN                  Size,       // Valid for AllocatePages/FreePages/AllocatePool
1100   IN VOID                   *Buffer
1101   )
1102 {
1103   MEMORY_PROFILE_CONTEXT_DATA   *ContextData;
1104
1105   if (!IS_UEFI_MEMORY_PROFILE_ENABLED) {
1106     return FALSE;
1107   }
1108
1109   if (!mMemoryProfileRecordingStatus) {
1110     return FALSE;
1111   }
1112
1113   //
1114   // Free operations have no memory type information, so skip the check.
1115   //
1116   if ((Action == MemoryProfileActionAllocatePages) || (Action == MemoryProfileActionAllocatePool)) {
1117     //
1118     // Only record limited MemoryType.
1119     //
1120     if (!CoreNeedRecordProfile (MemoryType)) {
1121       return FALSE;
1122     }
1123   }
1124
1125   ContextData = GetMemoryProfileContext ();
1126   if (ContextData == NULL) {
1127     return FALSE;
1128   }
1129
1130   switch (Action) {
1131     case MemoryProfileActionAllocatePages:
1132       CoreUpdateProfileAllocate (CallerAddress, Action, MemoryType, Size, Buffer);
1133       break;
1134     case MemoryProfileActionFreePages:
1135       CoreUpdateProfileFree (CallerAddress, Action, Size, Buffer);
1136       break;
1137     case MemoryProfileActionAllocatePool:
1138       CoreUpdateProfileAllocate (CallerAddress, Action, MemoryType, Size, Buffer);
1139       break;
1140     case MemoryProfileActionFreePool:
1141       CoreUpdateProfileFree (CallerAddress, Action, 0, Buffer);
1142       break;
1143     default:
1144       ASSERT (FALSE);
1145       break;
1146   }
1147   return TRUE;
1148 }
1149
1150 ////////////////////
1151
1152 /**
1153   Get memory profile data size.
1154
1155   @return Memory profile data size.
1156
1157 **/
1158 UINTN
1159 MemoryProfileGetDataSize (
1160   VOID
1161   )
1162 {
1163   MEMORY_PROFILE_CONTEXT_DATA       *ContextData;
1164   MEMORY_PROFILE_DRIVER_INFO_DATA   *DriverInfoData;
1165   LIST_ENTRY                        *DriverInfoList;
1166   LIST_ENTRY                        *DriverLink;
1167   UINTN                             TotalSize;
1168
1169
1170   ContextData = GetMemoryProfileContext ();
1171   if (ContextData == NULL) {
1172     return 0;
1173   }
1174
1175   TotalSize = sizeof (MEMORY_PROFILE_CONTEXT);
1176   TotalSize += sizeof (MEMORY_PROFILE_DRIVER_INFO) * (UINTN) ContextData->Context.ImageCount;
1177
1178   DriverInfoList = ContextData->DriverInfoList;
1179   for (DriverLink = DriverInfoList->ForwardLink;
1180        DriverLink != DriverInfoList;
1181        DriverLink = DriverLink->ForwardLink) {
1182     DriverInfoData = CR (
1183                        DriverLink,
1184                        MEMORY_PROFILE_DRIVER_INFO_DATA,
1185                        Link,
1186                        MEMORY_PROFILE_DRIVER_INFO_SIGNATURE
1187                        );
1188     TotalSize += sizeof (MEMORY_PROFILE_ALLOC_INFO) * (UINTN) DriverInfoData->DriverInfo.AllocRecordCount;
1189   }
1190
1191   return TotalSize;
1192 }
1193
1194 /**
1195   Copy memory profile data.
1196
1197   @param ProfileBuffer  The buffer to hold memory profile data.
1198
1199 **/
1200 VOID
1201 MemoryProfileCopyData (
1202   IN VOID   *ProfileBuffer
1203   )
1204 {
1205   MEMORY_PROFILE_CONTEXT            *Context;
1206   MEMORY_PROFILE_DRIVER_INFO        *DriverInfo;
1207   MEMORY_PROFILE_ALLOC_INFO         *AllocInfo;
1208   MEMORY_PROFILE_CONTEXT_DATA       *ContextData;
1209   MEMORY_PROFILE_DRIVER_INFO_DATA   *DriverInfoData;
1210   MEMORY_PROFILE_ALLOC_INFO_DATA    *AllocInfoData;
1211   LIST_ENTRY                        *DriverInfoList;
1212   LIST_ENTRY                        *DriverLink;
1213   LIST_ENTRY                        *AllocInfoList;
1214   LIST_ENTRY                        *AllocLink;
1215
1216   ContextData = GetMemoryProfileContext ();
1217   if (ContextData == NULL) {
1218     return ;
1219   }
1220
1221   Context = ProfileBuffer;
1222   CopyMem (Context, &ContextData->Context, sizeof (MEMORY_PROFILE_CONTEXT));
1223   DriverInfo = (MEMORY_PROFILE_DRIVER_INFO *) (Context + 1);
1224
1225   DriverInfoList = ContextData->DriverInfoList;
1226   for (DriverLink = DriverInfoList->ForwardLink;
1227        DriverLink != DriverInfoList;
1228        DriverLink = DriverLink->ForwardLink) {
1229     DriverInfoData = CR (
1230                        DriverLink,
1231                        MEMORY_PROFILE_DRIVER_INFO_DATA,
1232                        Link,
1233                        MEMORY_PROFILE_DRIVER_INFO_SIGNATURE
1234                        );
1235     CopyMem (DriverInfo, &DriverInfoData->DriverInfo, sizeof (MEMORY_PROFILE_DRIVER_INFO));
1236     AllocInfo = (MEMORY_PROFILE_ALLOC_INFO *) (DriverInfo + 1);
1237
1238     AllocInfoList = DriverInfoData->AllocInfoList;
1239     for (AllocLink = AllocInfoList->ForwardLink;
1240          AllocLink != AllocInfoList;
1241          AllocLink = AllocLink->ForwardLink) {
1242       AllocInfoData = CR (
1243                         AllocLink,
1244                         MEMORY_PROFILE_ALLOC_INFO_DATA,
1245                         Link,
1246                         MEMORY_PROFILE_ALLOC_INFO_SIGNATURE
1247                         );
1248       CopyMem (AllocInfo, &AllocInfoData->AllocInfo, sizeof (MEMORY_PROFILE_ALLOC_INFO));
1249       AllocInfo += 1;
1250     }
1251
1252     DriverInfo = (MEMORY_PROFILE_DRIVER_INFO *) ((UINTN) (DriverInfo + 1) + sizeof (MEMORY_PROFILE_ALLOC_INFO) * (UINTN) DriverInfo->AllocRecordCount);
1253   }
1254 }
1255
1256 /**
1257   Get memory profile data.
1258
1259   @param[in]      This              The EDKII_MEMORY_PROFILE_PROTOCOL instance.
1260   @param[in, out] ProfileSize       On entry, points to the size in bytes of the ProfileBuffer.
1261                                     On return, points to the size of the data returned in ProfileBuffer.
1262   @param[out]     ProfileBuffer     Profile buffer.
1263
1264   @return EFI_SUCCESS               Get the memory profile data successfully.
1265   @return EFI_BUFFER_TO_SMALL       The ProfileSize is too small for the resulting data.
1266                                     ProfileSize is updated with the size required.
1267
1268 **/
1269 EFI_STATUS
1270 EFIAPI
1271 ProfileProtocolGetData (
1272   IN     EDKII_MEMORY_PROFILE_PROTOCOL  *This,
1273   IN OUT UINT64                         *ProfileSize,
1274      OUT VOID                           *ProfileBuffer
1275   )
1276 {
1277   UINTN                                 Size;
1278   MEMORY_PROFILE_CONTEXT_DATA           *ContextData;
1279   BOOLEAN                               MemoryProfileRecordingStatus;
1280
1281   ContextData = GetMemoryProfileContext ();
1282   if (ContextData == NULL) {
1283     return EFI_UNSUPPORTED;
1284   }
1285
1286   MemoryProfileRecordingStatus = mMemoryProfileRecordingStatus;
1287   mMemoryProfileRecordingStatus = FALSE;
1288
1289   Size = MemoryProfileGetDataSize ();
1290
1291   if (*ProfileSize < Size) {
1292     *ProfileSize = Size;
1293     mMemoryProfileRecordingStatus = MemoryProfileRecordingStatus;
1294     return EFI_BUFFER_TOO_SMALL;
1295   }
1296
1297   *ProfileSize = Size;
1298   MemoryProfileCopyData (ProfileBuffer);
1299
1300   mMemoryProfileRecordingStatus = MemoryProfileRecordingStatus;
1301   return EFI_SUCCESS;
1302 }
1303
1304 /**
1305   Register image to memory profile.
1306
1307   @param[in] This               The EDKII_MEMORY_PROFILE_PROTOCOL instance.
1308   @param[in] FilePath           File path of the image.
1309   @param[in] ImageBase          Image base address.
1310   @param[in] ImageSize          Image size.
1311   @param[in] FileType           File type of the image.
1312
1313   @return EFI_SUCCESS           Register success.
1314   @return EFI_OUT_OF_RESOURCE   No enough resource for this register.
1315
1316 **/
1317 EFI_STATUS
1318 EFIAPI
1319 ProfileProtocolRegisterImage (
1320   IN EDKII_MEMORY_PROFILE_PROTOCOL  *This,
1321   IN EFI_DEVICE_PATH_PROTOCOL       *FilePath,
1322   IN PHYSICAL_ADDRESS               ImageBase,
1323   IN UINT64                         ImageSize,
1324   IN EFI_FV_FILETYPE                FileType
1325   )
1326 {
1327   EFI_STATUS                        Status;
1328   LOADED_IMAGE_PRIVATE_DATA         DriverEntry;
1329   VOID                              *EntryPointInImage;
1330
1331   ZeroMem (&DriverEntry, sizeof (DriverEntry));
1332   DriverEntry.Info.FilePath = FilePath;
1333   DriverEntry.ImageContext.ImageAddress = ImageBase;
1334   DriverEntry.ImageContext.ImageSize = ImageSize;
1335   Status = InternalPeCoffGetEntryPoint ((VOID *) (UINTN) ImageBase, &EntryPointInImage);
1336   ASSERT_EFI_ERROR (Status);
1337   DriverEntry.ImageContext.EntryPoint = (PHYSICAL_ADDRESS) (UINTN) EntryPointInImage;
1338   DriverEntry.ImageContext.ImageType = InternalPeCoffGetSubsystem ((VOID *) (UINTN) ImageBase);
1339
1340   return RegisterMemoryProfileImage (&DriverEntry, FileType) ? EFI_SUCCESS: EFI_OUT_OF_RESOURCES;
1341 }
1342
1343 /**
1344   Unregister image from memory profile.
1345
1346   @param[in] This               The EDKII_MEMORY_PROFILE_PROTOCOL instance.
1347   @param[in] FilePath           File path of the image.
1348   @param[in] ImageBase          Image base address.
1349   @param[in] ImageSize          Image size.
1350
1351   @return EFI_SUCCESS           Unregister success.
1352   @return EFI_NOT_FOUND         The image is not found.
1353
1354 **/
1355 EFI_STATUS
1356 EFIAPI
1357 ProfileProtocolUnregisterImage (
1358   IN EDKII_MEMORY_PROFILE_PROTOCOL  *This,
1359   IN EFI_DEVICE_PATH_PROTOCOL       *FilePath,
1360   IN PHYSICAL_ADDRESS               ImageBase,
1361   IN UINT64                         ImageSize
1362   )
1363 {
1364   EFI_STATUS                        Status;
1365   LOADED_IMAGE_PRIVATE_DATA         DriverEntry;
1366   VOID                              *EntryPointInImage;
1367
1368   ZeroMem (&DriverEntry, sizeof (DriverEntry));
1369   DriverEntry.Info.FilePath = FilePath;
1370   DriverEntry.ImageContext.ImageAddress = ImageBase;
1371   DriverEntry.ImageContext.ImageSize = ImageSize;
1372   Status = InternalPeCoffGetEntryPoint ((VOID *) (UINTN) ImageBase, &EntryPointInImage);
1373   ASSERT_EFI_ERROR (Status);
1374   DriverEntry.ImageContext.EntryPoint = (PHYSICAL_ADDRESS) (UINTN) EntryPointInImage;
1375
1376   return UnregisterMemoryProfileImage (&DriverEntry) ? EFI_SUCCESS: EFI_NOT_FOUND;
1377 }
1378
1379 ////////////////////