MdeModulePkg/Core/Dxe/Mem/MemoryProfileRecord.c

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