Home | History | Annotate | Download | only in Mem
      1 /** @file
      2   Support routines for UEFI memory profile.
      3 
      4   Copyright (c) 2014 - 2015, 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);
    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);
   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 ////////////////////
   1390