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