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   ASSERT (AllocInfoData != NULL);
 842   AllocInfo = &AllocInfoData->AllocInfo;
 843   AllocInfoData->Signature      = MEMORY_PROFILE_ALLOC_INFO_SIGNATURE;
 844   AllocInfo->Header.Signature   = MEMORY_PROFILE_ALLOC_INFO_SIGNATURE;
 845   AllocInfo->Header.Length      = sizeof (MEMORY_PROFILE_ALLOC_INFO);
 846   AllocInfo->Header.Revision    = MEMORY_PROFILE_ALLOC_INFO_REVISION;
 847   AllocInfo->CallerAddress      = CallerAddress;
 848   AllocInfo->SequenceId         = ContextData->Context.SequenceCount;
 849   AllocInfo->Action             = Action;
 850   AllocInfo->MemoryType         = MemoryType;
 851   AllocInfo->Buffer             = (PHYSICAL_ADDRESS) (UINTN) Buffer;
 852   AllocInfo->Size               = Size;
 853
 854   InsertTailList (DriverInfoData->AllocInfoList, &AllocInfoData->Link);
 855
 856   ProfileMemoryIndex = GetProfileMemoryIndex (MemoryType);
 857
 858   DriverInfo = &DriverInfoData->DriverInfo;
 859   DriverInfo->CurrentUsage += Size;
 860   if (DriverInfo->PeakUsage < DriverInfo->CurrentUsage) {
 861     DriverInfo->PeakUsage = DriverInfo->CurrentUsage;
 862   }
 863   DriverInfo->CurrentUsageByType[ProfileMemoryIndex] += Size;
 864   if (DriverInfo->PeakUsageByType[ProfileMemoryIndex] < DriverInfo->CurrentUsageByType[ProfileMemoryIndex]) {
 865     DriverInfo->PeakUsageByType[ProfileMemoryIndex] = DriverInfo->CurrentUsageByType[ProfileMemoryIndex];
 866   }
 867   DriverInfo->AllocRecordCount ++;
 868
 869   Context = &ContextData->Context;
 870   Context->CurrentTotalUsage += Size;
 871   if (Context->PeakTotalUsage < Context->CurrentTotalUsage) {
 872     Context->PeakTotalUsage = Context->CurrentTotalUsage;
 873   }
 874   Context->CurrentTotalUsageByType[ProfileMemoryIndex] += Size;
 875   if (Context->PeakTotalUsageByType[ProfileMemoryIndex] < Context->CurrentTotalUsageByType[ProfileMemoryIndex]) {
 876     Context->PeakTotalUsageByType[ProfileMemoryIndex] = Context->CurrentTotalUsageByType[ProfileMemoryIndex];
 877   }
 878   Context->SequenceCount ++;
 879
 880   return TRUE;
 881 }
 882
 883 /**
 884   Get memory profile alloc info from memory profile
 885
 886   @param DriverInfoData     Driver info
 887   @param Action             This Free action
 888   @param Size               Buffer size
 889   @param Buffer             Buffer address
 890
 891   @return Pointer to memory profile alloc info.
 892 **/
 893 MEMORY_PROFILE_ALLOC_INFO_DATA *
 894 GetMemoryProfileAllocInfoFromAddress (
 895   IN MEMORY_PROFILE_DRIVER_INFO_DATA    *DriverInfoData,
 896   IN MEMORY_PROFILE_ACTION              Action,
 897   IN UINTN                              Size,
 898   IN VOID                               *Buffer
 899   )
 900 {
 901   LIST_ENTRY                        *AllocInfoList;
 902   LIST_ENTRY                        *AllocLink;
 903   MEMORY_PROFILE_ALLOC_INFO         *AllocInfo;
 904   MEMORY_PROFILE_ALLOC_INFO_DATA    *AllocInfoData;
 905
 906   AllocInfoList = DriverInfoData->AllocInfoList;
 907
 908   for (AllocLink = AllocInfoList->ForwardLink;
 909        AllocLink != AllocInfoList;
 910        AllocLink = AllocLink->ForwardLink) {
 911     AllocInfoData = CR (
 912                       AllocLink,
 913                       MEMORY_PROFILE_ALLOC_INFO_DATA,
 914                       Link,
 915                       MEMORY_PROFILE_ALLOC_INFO_SIGNATURE
 916                       );
 917     AllocInfo = &AllocInfoData->AllocInfo;
 918     if (AllocInfo->Action != Action) {
 919       continue;
 920     }
 921     switch (Action) {
 922       case MemoryProfileActionAllocatePages:
 923         if ((AllocInfo->Buffer <= (PHYSICAL_ADDRESS) (UINTN) Buffer) &&
 924             ((AllocInfo->Buffer + AllocInfo->Size) >= ((PHYSICAL_ADDRESS) (UINTN) Buffer + Size))) {
 925           return AllocInfoData;
 926         }
 927         break;
 928       case MemoryProfileActionAllocatePool:
 929         if (AllocInfo->Buffer == (PHYSICAL_ADDRESS) (UINTN) Buffer) {
 930           return AllocInfoData;
 931         }
 932         break;
 933       default:
 934         ASSERT (FALSE);
 935         break;
 936     }
 937   }
 938
 939   return NULL;
 940 }
 941
 942 /**
 943   Update memory profile Free information.
 944
 945   @param CallerAddress  Address of caller who call Free.
 946   @param Action         This Free action.
 947   @param Size           Buffer size.
 948   @param Buffer         Buffer address.
 949
 950   @retval TRUE          Profile udpate success.
 951   @retval FALSE         Profile update fail.
 952
 953 **/
 954 BOOLEAN
 955 CoreUpdateProfileFree (
 956   IN PHYSICAL_ADDRESS       CallerAddress,
 957   IN MEMORY_PROFILE_ACTION  Action,
 958   IN UINTN                  Size,
 959   IN VOID                   *Buffer
 960   )
 961 {
 962   MEMORY_PROFILE_CONTEXT           *Context;
 963   MEMORY_PROFILE_DRIVER_INFO       *DriverInfo;
 964   MEMORY_PROFILE_ALLOC_INFO        *AllocInfo;
 965   MEMORY_PROFILE_CONTEXT_DATA      *ContextData;
 966   MEMORY_PROFILE_DRIVER_INFO_DATA  *DriverInfoData;
 967   LIST_ENTRY                       *DriverLink;
 968   LIST_ENTRY                       *DriverInfoList;
 969   MEMORY_PROFILE_DRIVER_INFO_DATA  *ThisDriverInfoData;
 970   MEMORY_PROFILE_ALLOC_INFO_DATA   *AllocInfoData;
 971   EFI_MEMORY_TYPE                  ProfileMemoryIndex;
 972
 973   ContextData = GetMemoryProfileContext ();
 974   if (ContextData == NULL) {
 975     return FALSE;
 976   }
 977
 978   DriverInfoData = GetMemoryProfileDriverInfoFromAddress (ContextData, CallerAddress);
 979   ASSERT (DriverInfoData != NULL);
 980
 981   switch (Action) {
 982     case MemoryProfileActionFreePages:
 983       AllocInfoData = GetMemoryProfileAllocInfoFromAddress (DriverInfoData, MemoryProfileActionAllocatePages, Size, Buffer);
 984       break;
 985     case MemoryProfileActionFreePool:
 986       AllocInfoData = GetMemoryProfileAllocInfoFromAddress (DriverInfoData, MemoryProfileActionAllocatePool, 0, Buffer);
 987       break;
 988     default:
 989       ASSERT (FALSE);
 990       AllocInfoData = NULL;
 991       break;
 992   }
 993   if (AllocInfoData == NULL) {
 994     //
 995     // Legal case, because driver A might free memory allocated by driver B, by some protocol.
 996     //
 997     DriverInfoList = ContextData->DriverInfoList;
 998
 999     for (DriverLink = DriverInfoList->ForwardLink;
1000          DriverLink != DriverInfoList;
1001          DriverLink = DriverLink->ForwardLink) {
1002       ThisDriverInfoData = CR (
1003                              DriverLink,
1004                              MEMORY_PROFILE_DRIVER_INFO_DATA,
1005                              Link,
1006                              MEMORY_PROFILE_DRIVER_INFO_SIGNATURE
1007                              );
1008       switch (Action) {
1009         case MemoryProfileActionFreePages:
1010           AllocInfoData = GetMemoryProfileAllocInfoFromAddress (ThisDriverInfoData, MemoryProfileActionAllocatePages, Size, Buffer);
1011           break;
1012         case MemoryProfileActionFreePool:
1013           AllocInfoData = GetMemoryProfileAllocInfoFromAddress (ThisDriverInfoData, MemoryProfileActionAllocatePool, 0, Buffer);
1014           break;
1015         default:
1016           ASSERT (FALSE);
1017           AllocInfoData = NULL;
1018           break;
1019       }
1020       if (AllocInfoData != NULL) {
1021         DriverInfoData = ThisDriverInfoData;
1022         break;
1023       }
1024     }
1025
1026     if (AllocInfoData == NULL) {
1027       //
1028       // No matched allocate operation is found for this free operation.
1029       // It is because the specified memory type allocate operation has been
1030       // filtered by CoreNeedRecordProfile(), but free operations have no
1031       // memory type information, they can not be filtered by CoreNeedRecordProfile().
1032       // Then, they will be filtered here.
1033       //
1034       return FALSE;
1035     }
1036   }
1037
1038   Context = &ContextData->Context;
1039   DriverInfo = &DriverInfoData->DriverInfo;
1040   AllocInfo = &AllocInfoData->AllocInfo;
1041
1042   ProfileMemoryIndex = GetProfileMemoryIndex (AllocInfo->MemoryType);
1043
1044   Context->CurrentTotalUsage -= AllocInfo->Size;
1045   Context->CurrentTotalUsageByType[ProfileMemoryIndex] -= AllocInfo->Size;
1046
1047   DriverInfo->CurrentUsage -= AllocInfo->Size;
1048   DriverInfo->CurrentUsageByType[ProfileMemoryIndex] -= AllocInfo->Size;
1049   DriverInfo->AllocRecordCount --;
1050
1051   RemoveEntryList (&AllocInfoData->Link);
1052
1053   if (Action == MemoryProfileActionFreePages) {
1054     if (AllocInfo->Buffer != (PHYSICAL_ADDRESS) (UINTN) Buffer) {
1055       CoreUpdateProfileAllocate (
1056         AllocInfo->CallerAddress,
1057         MemoryProfileActionAllocatePages,
1058         AllocInfo->MemoryType,
1059         (UINTN) ((PHYSICAL_ADDRESS) (UINTN) Buffer - AllocInfo->Buffer),
1060         (VOID *) (UINTN) AllocInfo->Buffer
1061         );
1062     }
1063     if (AllocInfo->Buffer + AllocInfo->Size != ((PHYSICAL_ADDRESS) (UINTN) Buffer + Size)) {
1064       CoreUpdateProfileAllocate (
1065         AllocInfo->CallerAddress,
1066         MemoryProfileActionAllocatePages,
1067         AllocInfo->MemoryType,
1068         (UINTN) ((AllocInfo->Buffer + AllocInfo->Size) - ((PHYSICAL_ADDRESS) (UINTN) Buffer + Size)),
1069         (VOID *) ((UINTN) Buffer + Size)
1070         );
1071     }
1072   }
1073
1074   //
1075   // Use CoreInternalFreePool() that will not update profile for this FreePool action.
1076   //
1077   CoreInternalFreePool (AllocInfoData);
1078
1079   return TRUE;
1080 }
1081
1082 /**
1083   Update memory profile information.
1084
1085   @param CallerAddress  Address of caller who call Allocate or Free.
1086   @param Action         This Allocate or Free action.
1087   @param MemoryType     Memory type.
1088   @param Size           Buffer size.
1089   @param Buffer         Buffer address.
1090
1091   @retval TRUE          Profile udpate success.
1092   @retval FALSE         Profile update fail.
1093
1094 **/
1095 BOOLEAN
1096 CoreUpdateProfile (
1097   IN PHYSICAL_ADDRESS       CallerAddress,
1098   IN MEMORY_PROFILE_ACTION  Action,
1099   IN EFI_MEMORY_TYPE        MemoryType, // Valid for AllocatePages/AllocatePool
1100   IN UINTN                  Size,       // Valid for AllocatePages/FreePages/AllocatePool
1101   IN VOID                   *Buffer
1102   )
1103 {
1104   MEMORY_PROFILE_CONTEXT_DATA   *ContextData;
1105
1106   if (!IS_UEFI_MEMORY_PROFILE_ENABLED) {
1107     return FALSE;
1108   }
1109
1110   if (!mMemoryProfileRecordingStatus) {
1111     return FALSE;
1112   }
1113
1114   //
1115   // Free operations have no memory type information, so skip the check.
1116   //
1117   if ((Action == MemoryProfileActionAllocatePages) || (Action == MemoryProfileActionAllocatePool)) {
1118     //
1119     // Only record limited MemoryType.
1120     //
1121     if (!CoreNeedRecordProfile (MemoryType)) {
1122       return FALSE;
1123     }
1124   }
1125
1126   ContextData = GetMemoryProfileContext ();
1127   if (ContextData == NULL) {
1128     return FALSE;
1129   }
1130
1131   switch (Action) {
1132     case MemoryProfileActionAllocatePages:
1133       CoreUpdateProfileAllocate (CallerAddress, Action, MemoryType, Size, Buffer);
1134       break;
1135     case MemoryProfileActionFreePages:
1136       CoreUpdateProfileFree (CallerAddress, Action, Size, Buffer);
1137       break;
1138     case MemoryProfileActionAllocatePool:
1139       CoreUpdateProfileAllocate (CallerAddress, Action, MemoryType, Size, Buffer);
1140       break;
1141     case MemoryProfileActionFreePool:
1142       CoreUpdateProfileFree (CallerAddress, Action, 0, Buffer);
1143       break;
1144     default:
1145       ASSERT (FALSE);
1146       break;
1147   }
1148   return TRUE;
1149 }
1150
1151 ////////////////////
1152
1153 /**
1154   Get memory profile data size.
1155
1156   @return Memory profile data size.
1157
1158 **/
1159 UINTN
1160 MemoryProfileGetDataSize (
1161   VOID
1162   )
1163 {
1164   MEMORY_PROFILE_CONTEXT_DATA       *ContextData;
1165   MEMORY_PROFILE_DRIVER_INFO_DATA   *DriverInfoData;
1166   LIST_ENTRY                        *DriverInfoList;
1167   LIST_ENTRY                        *DriverLink;
1168   UINTN                             TotalSize;
1169
1170
1171   ContextData = GetMemoryProfileContext ();
1172   if (ContextData == NULL) {
1173     return 0;
1174   }
1175
1176   TotalSize = sizeof (MEMORY_PROFILE_CONTEXT);
1177   TotalSize += sizeof (MEMORY_PROFILE_DRIVER_INFO) * (UINTN) ContextData->Context.ImageCount;
1178
1179   DriverInfoList = ContextData->DriverInfoList;
1180   for (DriverLink = DriverInfoList->ForwardLink;
1181        DriverLink != DriverInfoList;
1182        DriverLink = DriverLink->ForwardLink) {
1183     DriverInfoData = CR (
1184                        DriverLink,
1185                        MEMORY_PROFILE_DRIVER_INFO_DATA,
1186                        Link,
1187                        MEMORY_PROFILE_DRIVER_INFO_SIGNATURE
1188                        );
1189     TotalSize += sizeof (MEMORY_PROFILE_ALLOC_INFO) * (UINTN) DriverInfoData->DriverInfo.AllocRecordCount;
1190   }
1191
1192   return TotalSize;
1193 }
1194
1195 /**
1196   Copy memory profile data.
1197
1198   @param ProfileBuffer  The buffer to hold memory profile data.
1199
1200 **/
1201 VOID
1202 MemoryProfileCopyData (
1203   IN VOID   *ProfileBuffer
1204   )
1205 {
1206   MEMORY_PROFILE_CONTEXT            *Context;
1207   MEMORY_PROFILE_DRIVER_INFO        *DriverInfo;
1208   MEMORY_PROFILE_ALLOC_INFO         *AllocInfo;
1209   MEMORY_PROFILE_CONTEXT_DATA       *ContextData;
1210   MEMORY_PROFILE_DRIVER_INFO_DATA   *DriverInfoData;
1211   MEMORY_PROFILE_ALLOC_INFO_DATA    *AllocInfoData;
1212   LIST_ENTRY                        *DriverInfoList;
1213   LIST_ENTRY                        *DriverLink;
1214   LIST_ENTRY                        *AllocInfoList;
1215   LIST_ENTRY                        *AllocLink;
1216
1217   ContextData = GetMemoryProfileContext ();
1218   if (ContextData == NULL) {
1219     return ;
1220   }
1221
1222   Context = ProfileBuffer;
1223   CopyMem (Context, &ContextData->Context, sizeof (MEMORY_PROFILE_CONTEXT));
1224   DriverInfo = (MEMORY_PROFILE_DRIVER_INFO *) (Context + 1);
1225
1226   DriverInfoList = ContextData->DriverInfoList;
1227   for (DriverLink = DriverInfoList->ForwardLink;
1228        DriverLink != DriverInfoList;
1229        DriverLink = DriverLink->ForwardLink) {
1230     DriverInfoData = CR (
1231                        DriverLink,
1232                        MEMORY_PROFILE_DRIVER_INFO_DATA,
1233                        Link,
1234                        MEMORY_PROFILE_DRIVER_INFO_SIGNATURE
1235                        );
1236     CopyMem (DriverInfo, &DriverInfoData->DriverInfo, sizeof (MEMORY_PROFILE_DRIVER_INFO));
1237     AllocInfo = (MEMORY_PROFILE_ALLOC_INFO *) (DriverInfo + 1);
1238
1239     AllocInfoList = DriverInfoData->AllocInfoList;
1240     for (AllocLink = AllocInfoList->ForwardLink;
1241          AllocLink != AllocInfoList;
1242          AllocLink = AllocLink->ForwardLink) {
1243       AllocInfoData = CR (
1244                         AllocLink,
1245                         MEMORY_PROFILE_ALLOC_INFO_DATA,
1246                         Link,
1247                         MEMORY_PROFILE_ALLOC_INFO_SIGNATURE
1248                         );
1249       CopyMem (AllocInfo, &AllocInfoData->AllocInfo, sizeof (MEMORY_PROFILE_ALLOC_INFO));
1250       AllocInfo += 1;
1251     }
1252
1253     DriverInfo = (MEMORY_PROFILE_DRIVER_INFO *) ((UINTN) (DriverInfo + 1) + sizeof (MEMORY_PROFILE_ALLOC_INFO) * (UINTN) DriverInfo->AllocRecordCount);
1254   }
1255 }
1256
1257 /**
1258   Get memory profile data.
1259
1260   @param[in]      This              The EDKII_MEMORY_PROFILE_PROTOCOL instance.
1261   @param[in, out] ProfileSize       On entry, points to the size in bytes of the ProfileBuffer.
1262                                     On return, points to the size of the data returned in ProfileBuffer.
1263   @param[out]     ProfileBuffer     Profile buffer.
1264
1265   @return EFI_SUCCESS               Get the memory profile data successfully.
1266   @return EFI_BUFFER_TO_SMALL       The ProfileSize is too small for the resulting data.
1267                                     ProfileSize is updated with the size required.
1268
1269 **/
1270 EFI_STATUS
1271 EFIAPI
1272 ProfileProtocolGetData (
1273   IN     EDKII_MEMORY_PROFILE_PROTOCOL  *This,
1274   IN OUT UINT64                         *ProfileSize,
1275      OUT VOID                           *ProfileBuffer
1276   )
1277 {
1278   UINTN                                 Size;
1279   MEMORY_PROFILE_CONTEXT_DATA           *ContextData;
1280   BOOLEAN                               MemoryProfileRecordingStatus;
1281
1282   ContextData = GetMemoryProfileContext ();
1283   if (ContextData == NULL) {
1284     return EFI_UNSUPPORTED;
1285   }
1286
1287   MemoryProfileRecordingStatus = mMemoryProfileRecordingStatus;
1288   mMemoryProfileRecordingStatus = FALSE;
1289
1290   Size = MemoryProfileGetDataSize ();
1291
1292   if (*ProfileSize < Size) {
1293     *ProfileSize = Size;
1294     mMemoryProfileRecordingStatus = MemoryProfileRecordingStatus;
1295     return EFI_BUFFER_TOO_SMALL;
1296   }
1297
1298   *ProfileSize = Size;
1299   MemoryProfileCopyData (ProfileBuffer);
1300
1301   mMemoryProfileRecordingStatus = MemoryProfileRecordingStatus;
1302   return EFI_SUCCESS;
1303 }
1304
1305 /**
1306   Register image to memory profile.
1307
1308   @param[in] This               The EDKII_MEMORY_PROFILE_PROTOCOL instance.
1309   @param[in] FilePath           File path of the image.
1310   @param[in] ImageBase          Image base address.
1311   @param[in] ImageSize          Image size.
1312   @param[in] FileType           File type of the image.
1313
1314   @return EFI_SUCCESS           Register success.
1315   @return EFI_OUT_OF_RESOURCE   No enough resource for this register.
1316
1317 **/
1318 EFI_STATUS
1319 EFIAPI
1320 ProfileProtocolRegisterImage (
1321   IN EDKII_MEMORY_PROFILE_PROTOCOL  *This,
1322   IN EFI_DEVICE_PATH_PROTOCOL       *FilePath,
1323   IN PHYSICAL_ADDRESS               ImageBase,
1324   IN UINT64                         ImageSize,
1325   IN EFI_FV_FILETYPE                FileType
1326   )
1327 {
1328   EFI_STATUS                        Status;
1329   LOADED_IMAGE_PRIVATE_DATA         DriverEntry;
1330   VOID                              *EntryPointInImage;
1331
1332   ZeroMem (&DriverEntry, sizeof (DriverEntry));
1333   DriverEntry.Info.FilePath = FilePath;
1334   DriverEntry.ImageContext.ImageAddress = ImageBase;
1335   DriverEntry.ImageContext.ImageSize = ImageSize;
1336   Status = InternalPeCoffGetEntryPoint ((VOID *) (UINTN) ImageBase, &EntryPointInImage);
1337   ASSERT_EFI_ERROR (Status);
1338   DriverEntry.ImageContext.EntryPoint = (PHYSICAL_ADDRESS) (UINTN) EntryPointInImage;
1339   DriverEntry.ImageContext.ImageType = InternalPeCoffGetSubsystem ((VOID *) (UINTN) ImageBase);
1340
1341   return RegisterMemoryProfileImage (&DriverEntry, FileType) ? EFI_SUCCESS: EFI_OUT_OF_RESOURCES;
1342 }
1343
1344 /**
1345   Unregister image from memory profile.
1346
1347   @param[in] This               The EDKII_MEMORY_PROFILE_PROTOCOL instance.
1348   @param[in] FilePath           File path of the image.
1349   @param[in] ImageBase          Image base address.
1350   @param[in] ImageSize          Image size.
1351
1352   @return EFI_SUCCESS           Unregister success.
1353   @return EFI_NOT_FOUND         The image is not found.
1354
1355 **/
1356 EFI_STATUS
1357 EFIAPI
1358 ProfileProtocolUnregisterImage (
1359   IN EDKII_MEMORY_PROFILE_PROTOCOL  *This,
1360   IN EFI_DEVICE_PATH_PROTOCOL       *FilePath,
1361   IN PHYSICAL_ADDRESS               ImageBase,
1362   IN UINT64                         ImageSize
1363   )
1364 {
1365   EFI_STATUS                        Status;
1366   LOADED_IMAGE_PRIVATE_DATA         DriverEntry;
1367   VOID                              *EntryPointInImage;
1368
1369   ZeroMem (&DriverEntry, sizeof (DriverEntry));
1370   DriverEntry.Info.FilePath = FilePath;
1371   DriverEntry.ImageContext.ImageAddress = ImageBase;
1372   DriverEntry.ImageContext.ImageSize = ImageSize;
1373   Status = InternalPeCoffGetEntryPoint ((VOID *) (UINTN) ImageBase, &EntryPointInImage);
1374   ASSERT_EFI_ERROR (Status);
1375   DriverEntry.ImageContext.EntryPoint = (PHYSICAL_ADDRESS) (UINTN) EntryPointInImage;
1376
1377   return UnregisterMemoryProfileImage (&DriverEntry) ? EFI_SUCCESS: EFI_NOT_FOUND;
1378 }
1379
1380 ////////////////////