/** @file
  Support routines for SMRAM profile.

  Copyright (c) 2014 - 2018, Intel Corporation. All rights reserved.<BR>
  SPDX-License-Identifier: BSD-2-Clause-Patent

**/

#include "PiSmmCore.h"

#define IS_SMRAM_PROFILE_ENABLED        ((PcdGet8 (PcdMemoryProfilePropertyMask) & BIT1) != 0)
#define IS_UEFI_MEMORY_PROFILE_ENABLED  ((PcdGet8 (PcdMemoryProfilePropertyMask) & BIT0) != 0)

#define GET_OCCUPIED_SIZE(ActualSize, Alignment) \
  ((ActualSize) + (((Alignment) - ((ActualSize) & ((Alignment) - 1))) & ((Alignment) - 1)))

typedef struct {
  UINT32                    Signature;
  MEMORY_PROFILE_CONTEXT    Context;
  LIST_ENTRY                *DriverInfoList;
} MEMORY_PROFILE_CONTEXT_DATA;

typedef struct {
  UINT32                        Signature;
  MEMORY_PROFILE_DRIVER_INFO    DriverInfo;
  LIST_ENTRY                    *AllocInfoList;
  CHAR8                         *PdbString;
  LIST_ENTRY                    Link;
} MEMORY_PROFILE_DRIVER_INFO_DATA;

typedef struct {
  UINT32                       Signature;
  MEMORY_PROFILE_ALLOC_INFO    AllocInfo;
  CHAR8                        *ActionString;
  LIST_ENTRY                   Link;
} MEMORY_PROFILE_ALLOC_INFO_DATA;

//
// When free memory less than 4 pages, dump it.
//
#define SMRAM_INFO_DUMP_PAGE_THRESHOLD  4

GLOBAL_REMOVE_IF_UNREFERENCED MEMORY_PROFILE_FREE_MEMORY  mSmramFreeMemory = {
  {
    MEMORY_PROFILE_FREE_MEMORY_SIGNATURE,
    sizeof (MEMORY_PROFILE_FREE_MEMORY),
    MEMORY_PROFILE_FREE_MEMORY_REVISION
  },
  0,
  0
};

GLOBAL_REMOVE_IF_UNREFERENCED LIST_ENTRY                   mImageQueue          = INITIALIZE_LIST_HEAD_VARIABLE (mImageQueue);
GLOBAL_REMOVE_IF_UNREFERENCED MEMORY_PROFILE_CONTEXT_DATA  mSmramProfileContext = {
  MEMORY_PROFILE_CONTEXT_SIGNATURE,
  {
    {
      MEMORY_PROFILE_CONTEXT_SIGNATURE,
      sizeof (MEMORY_PROFILE_CONTEXT),
      MEMORY_PROFILE_CONTEXT_REVISION
    },
    0,
    0,
    { 0 },
    { 0 },
    0,
    0,
    0
  },
  &mImageQueue,
};
GLOBAL_REMOVE_IF_UNREFERENCED MEMORY_PROFILE_CONTEXT_DATA  *mSmramProfileContextPtr = NULL;

GLOBAL_REMOVE_IF_UNREFERENCED BOOLEAN                   mSmramReadyToLock;
GLOBAL_REMOVE_IF_UNREFERENCED BOOLEAN                   mSmramProfileGettingStatus   = FALSE;
GLOBAL_REMOVE_IF_UNREFERENCED BOOLEAN                   mSmramProfileRecordingEnable = MEMORY_PROFILE_RECORDING_DISABLE;
GLOBAL_REMOVE_IF_UNREFERENCED EFI_DEVICE_PATH_PROTOCOL  *mSmramProfileDriverPath;
GLOBAL_REMOVE_IF_UNREFERENCED UINTN                     mSmramProfileDriverPathSize;

/**
  Dump SMRAM information.

**/
VOID
DumpSmramInfo (
  VOID
  );

/**
  Get memory profile data.

  @param[in]      This              The EDKII_SMM_MEMORY_PROFILE_PROTOCOL instance.
  @param[in, out] ProfileSize       On entry, points to the size in bytes of the ProfileBuffer.
                                    On return, points to the size of the data returned in ProfileBuffer.
  @param[out]     ProfileBuffer     Profile buffer.

  @return EFI_SUCCESS               Get the memory profile data successfully.
  @return EFI_UNSUPPORTED           Memory profile is unsupported.
  @return EFI_BUFFER_TO_SMALL       The ProfileSize is too small for the resulting data.
                                    ProfileSize is updated with the size required.

**/
EFI_STATUS
EFIAPI
SmramProfileProtocolGetData (
  IN     EDKII_SMM_MEMORY_PROFILE_PROTOCOL  *This,
  IN OUT UINT64                             *ProfileSize,
  OUT VOID                                  *ProfileBuffer
  );

/**
  Register image to memory profile.

  @param[in] This               The EDKII_SMM_MEMORY_PROFILE_PROTOCOL instance.
  @param[in] FilePath           File path of the image.
  @param[in] ImageBase          Image base address.
  @param[in] ImageSize          Image size.
  @param[in] FileType           File type of the image.

  @return EFI_SUCCESS           Register successfully.
  @return EFI_UNSUPPORTED       Memory profile is unsupported,
                                or memory profile for the image is not required.
  @return EFI_OUT_OF_RESOURCE   No enough resource for this register.

**/
EFI_STATUS
EFIAPI
SmramProfileProtocolRegisterImage (
  IN EDKII_SMM_MEMORY_PROFILE_PROTOCOL  *This,
  IN EFI_DEVICE_PATH_PROTOCOL           *FilePath,
  IN PHYSICAL_ADDRESS                   ImageBase,
  IN UINT64                             ImageSize,
  IN EFI_FV_FILETYPE                    FileType
  );

/**
  Unregister image from memory profile.

  @param[in] This               The EDKII_SMM_MEMORY_PROFILE_PROTOCOL instance.
  @param[in] FilePath           File path of the image.
  @param[in] ImageBase          Image base address.
  @param[in] ImageSize          Image size.

  @return EFI_SUCCESS           Unregister successfully.
  @return EFI_UNSUPPORTED       Memory profile is unsupported,
                                or memory profile for the image is not required.
  @return EFI_NOT_FOUND         The image is not found.

**/
EFI_STATUS
EFIAPI
SmramProfileProtocolUnregisterImage (
  IN EDKII_SMM_MEMORY_PROFILE_PROTOCOL  *This,
  IN EFI_DEVICE_PATH_PROTOCOL           *FilePath,
  IN PHYSICAL_ADDRESS                   ImageBase,
  IN UINT64                             ImageSize
  );

/**
  Get memory profile recording state.

  @param[in]  This              The EDKII_SMM_MEMORY_PROFILE_PROTOCOL instance.
  @param[out] RecordingState    Recording state.

  @return EFI_SUCCESS           Memory profile recording state is returned.
  @return EFI_UNSUPPORTED       Memory profile is unsupported.
  @return EFI_INVALID_PARAMETER RecordingState is NULL.

**/
EFI_STATUS
EFIAPI
SmramProfileProtocolGetRecordingState (
  IN EDKII_SMM_MEMORY_PROFILE_PROTOCOL  *This,
  OUT BOOLEAN                           *RecordingState
  );

/**
  Set memory profile recording state.

  @param[in] This               The EDKII_SMM_MEMORY_PROFILE_PROTOCOL instance.
  @param[in] RecordingState     Recording state.

  @return EFI_SUCCESS           Set memory profile recording state successfully.
  @return EFI_UNSUPPORTED       Memory profile is unsupported.

**/
EFI_STATUS
EFIAPI
SmramProfileProtocolSetRecordingState (
  IN EDKII_SMM_MEMORY_PROFILE_PROTOCOL  *This,
  IN BOOLEAN                            RecordingState
  );

/**
  Record memory profile of multilevel caller.

  @param[in] This               The EDKII_SMM_MEMORY_PROFILE_PROTOCOL instance.
  @param[in] CallerAddress      Address of caller.
  @param[in] Action             Memory profile action.
  @param[in] MemoryType         Memory type.
                                EfiMaxMemoryType means the MemoryType is unknown.
  @param[in] Buffer             Buffer address.
  @param[in] Size               Buffer size.
  @param[in] ActionString       String for memory profile action.
                                Only needed for user defined allocate action.

  @return EFI_SUCCESS           Memory profile is updated.
  @return EFI_UNSUPPORTED       Memory profile is unsupported,
                                or memory profile for the image is not required,
                                or memory profile for the memory type is not required.
  @return EFI_ACCESS_DENIED     It is during memory profile data getting.
  @return EFI_ABORTED           Memory profile recording is not enabled.
  @return EFI_OUT_OF_RESOURCES  No enough resource to update memory profile for allocate action.
  @return EFI_NOT_FOUND         No matched allocate info found for free action.

**/
EFI_STATUS
EFIAPI
SmramProfileProtocolRecord (
  IN EDKII_SMM_MEMORY_PROFILE_PROTOCOL  *This,
  IN PHYSICAL_ADDRESS                   CallerAddress,
  IN MEMORY_PROFILE_ACTION              Action,
  IN EFI_MEMORY_TYPE                    MemoryType,
  IN VOID                               *Buffer,
  IN UINTN                              Size,
  IN CHAR8                              *ActionString OPTIONAL
  );

GLOBAL_REMOVE_IF_UNREFERENCED EDKII_SMM_MEMORY_PROFILE_PROTOCOL  mSmmProfileProtocol = {
  SmramProfileProtocolGetData,
  SmramProfileProtocolRegisterImage,
  SmramProfileProtocolUnregisterImage,
  SmramProfileProtocolGetRecordingState,
  SmramProfileProtocolSetRecordingState,
  SmramProfileProtocolRecord,
};

/**
  Return SMRAM profile context.

  @return SMRAM profile context.

**/
MEMORY_PROFILE_CONTEXT_DATA *
GetSmramProfileContext (
  VOID
  )
{
  return mSmramProfileContextPtr;
}

/**
  Retrieves and returns the Subsystem of a PE/COFF image that has been loaded into system memory.
  If Pe32Data is NULL, then ASSERT().

  @param Pe32Data   The pointer to the PE/COFF image that is loaded in system memory.

  @return The Subsystem of the PE/COFF image.

**/
UINT16
InternalPeCoffGetSubsystem (
  IN VOID  *Pe32Data
  )
{
  EFI_IMAGE_OPTIONAL_HEADER_PTR_UNION  Hdr;
  EFI_IMAGE_DOS_HEADER                 *DosHdr;
  UINT16                               Magic;

  ASSERT (Pe32Data != NULL);

  DosHdr = (EFI_IMAGE_DOS_HEADER *)Pe32Data;
  if (DosHdr->e_magic == EFI_IMAGE_DOS_SIGNATURE) {
    //
    // DOS image header is present, so read the PE header after the DOS image header.
    //
    Hdr.Pe32 = (EFI_IMAGE_NT_HEADERS32 *)((UINTN)Pe32Data + (UINTN)((DosHdr->e_lfanew) & 0x0ffff));
  } else {
    //
    // DOS image header is not present, so PE header is at the image base.
    //
    Hdr.Pe32 = (EFI_IMAGE_NT_HEADERS32 *)Pe32Data;
  }

  if (Hdr.Te->Signature == EFI_TE_IMAGE_HEADER_SIGNATURE) {
    return Hdr.Te->Subsystem;
  } else if (Hdr.Pe32->Signature == EFI_IMAGE_NT_SIGNATURE) {
    Magic = Hdr.Pe32->OptionalHeader.Magic;
    if (Magic == EFI_IMAGE_NT_OPTIONAL_HDR32_MAGIC) {
      return Hdr.Pe32->OptionalHeader.Subsystem;
    } else if (Magic == EFI_IMAGE_NT_OPTIONAL_HDR64_MAGIC) {
      return Hdr.Pe32Plus->OptionalHeader.Subsystem;
    }
  }

  return 0x0000;
}

/**
  Retrieves and returns a pointer to the entry point to a PE/COFF image that has been loaded
  into system memory with the PE/COFF Loader Library functions.

  Retrieves the entry point to the PE/COFF image specified by Pe32Data and returns this entry
  point in EntryPoint.  If the entry point could not be retrieved from the PE/COFF image, then
  return RETURN_INVALID_PARAMETER.  Otherwise return RETURN_SUCCESS.
  If Pe32Data is NULL, then ASSERT().
  If EntryPoint is NULL, then ASSERT().

  @param  Pe32Data                  The pointer to the PE/COFF image that is loaded in system memory.
  @param  EntryPoint                The pointer to entry point to the PE/COFF image to return.

  @retval RETURN_SUCCESS            EntryPoint was returned.
  @retval RETURN_INVALID_PARAMETER  The entry point could not be found in the PE/COFF image.

**/
RETURN_STATUS
InternalPeCoffGetEntryPoint (
  IN  VOID  *Pe32Data,
  OUT VOID  **EntryPoint
  )
{
  EFI_IMAGE_DOS_HEADER                 *DosHdr;
  EFI_IMAGE_OPTIONAL_HEADER_PTR_UNION  Hdr;

  ASSERT (Pe32Data   != NULL);
  ASSERT (EntryPoint != NULL);

  DosHdr = (EFI_IMAGE_DOS_HEADER *)Pe32Data;
  if (DosHdr->e_magic == EFI_IMAGE_DOS_SIGNATURE) {
    //
    // DOS image header is present, so read the PE header after the DOS image header.
    //
    Hdr.Pe32 = (EFI_IMAGE_NT_HEADERS32 *)((UINTN)Pe32Data + (UINTN)((DosHdr->e_lfanew) & 0x0ffff));
  } else {
    //
    // DOS image header is not present, so PE header is at the image base.
    //
    Hdr.Pe32 = (EFI_IMAGE_NT_HEADERS32 *)Pe32Data;
  }

  //
  // Calculate the entry point relative to the start of the image.
  // AddressOfEntryPoint is common for PE32 & PE32+
  //
  if (Hdr.Te->Signature == EFI_TE_IMAGE_HEADER_SIGNATURE) {
    *EntryPoint = (VOID *)((UINTN)Pe32Data + (UINTN)(Hdr.Te->AddressOfEntryPoint & 0x0ffffffff) + sizeof (EFI_TE_IMAGE_HEADER) - Hdr.Te->StrippedSize);
    return RETURN_SUCCESS;
  } else if (Hdr.Pe32->Signature == EFI_IMAGE_NT_SIGNATURE) {
    *EntryPoint = (VOID *)((UINTN)Pe32Data + (UINTN)(Hdr.Pe32->OptionalHeader.AddressOfEntryPoint & 0x0ffffffff));
    return RETURN_SUCCESS;
  }

  return RETURN_UNSUPPORTED;
}

/**
  Build driver info.

  @param ContextData    Memory profile context.
  @param FileName       File name of the image.
  @param ImageBase      Image base address.
  @param ImageSize      Image size.
  @param EntryPoint     Entry point of the image.
  @param ImageSubsystem Image subsystem of the image.
  @param FileType       File type of the image.

  @return Pointer to memory profile driver info.

**/
MEMORY_PROFILE_DRIVER_INFO_DATA *
BuildDriverInfo (
  IN MEMORY_PROFILE_CONTEXT_DATA  *ContextData,
  IN EFI_GUID                     *FileName,
  IN PHYSICAL_ADDRESS             ImageBase,
  IN UINT64                       ImageSize,
  IN PHYSICAL_ADDRESS             EntryPoint,
  IN UINT16                       ImageSubsystem,
  IN EFI_FV_FILETYPE              FileType
  )
{
  EFI_STATUS                       Status;
  MEMORY_PROFILE_DRIVER_INFO       *DriverInfo;
  MEMORY_PROFILE_DRIVER_INFO_DATA  *DriverInfoData;
  VOID                             *EntryPointInImage;
  CHAR8                            *PdbString;
  UINTN                            PdbSize;
  UINTN                            PdbOccupiedSize;

  PdbSize         = 0;
  PdbOccupiedSize = 0;
  PdbString       = NULL;
  if (ImageBase != 0) {
    PdbString = PeCoffLoaderGetPdbPointer ((VOID *)(UINTN)ImageBase);
    if (PdbString != NULL) {
      PdbSize         = AsciiStrSize (PdbString);
      PdbOccupiedSize = GET_OCCUPIED_SIZE (PdbSize, sizeof (UINT64));
    }
  }

  //
  // Use SmmInternalAllocatePool() that will not update profile for this AllocatePool action.
  //
  Status = SmmInternalAllocatePool (
             EfiRuntimeServicesData,
             sizeof (*DriverInfoData) + sizeof (LIST_ENTRY) + PdbSize,
             (VOID **)&DriverInfoData
             );
  if (EFI_ERROR (Status)) {
    return NULL;
  }

  ASSERT (DriverInfoData != NULL);

  ZeroMem (DriverInfoData, sizeof (*DriverInfoData));

  DriverInfo                   = &DriverInfoData->DriverInfo;
  DriverInfoData->Signature    = MEMORY_PROFILE_DRIVER_INFO_SIGNATURE;
  DriverInfo->Header.Signature = MEMORY_PROFILE_DRIVER_INFO_SIGNATURE;
  DriverInfo->Header.Length    = (UINT16)(sizeof (MEMORY_PROFILE_DRIVER_INFO) + PdbOccupiedSize);
  DriverInfo->Header.Revision  = MEMORY_PROFILE_DRIVER_INFO_REVISION;
  if (FileName != NULL) {
    CopyMem (&DriverInfo->FileName, FileName, sizeof (EFI_GUID));
  }

  DriverInfo->ImageBase      = ImageBase;
  DriverInfo->ImageSize      = ImageSize;
  DriverInfo->EntryPoint     = EntryPoint;
  DriverInfo->ImageSubsystem = ImageSubsystem;
  if ((EntryPoint != 0) && ((EntryPoint < ImageBase) || (EntryPoint >= (ImageBase + ImageSize)))) {
    //
    // If the EntryPoint is not in the range of image buffer, it should come from emulation environment.
    // So patch ImageBuffer here to align the EntryPoint.
    //
    Status = InternalPeCoffGetEntryPoint ((VOID *)(UINTN)ImageBase, &EntryPointInImage);
    ASSERT_EFI_ERROR (Status);
    DriverInfo->ImageBase = ImageBase + EntryPoint - (PHYSICAL_ADDRESS)(UINTN)EntryPointInImage;
  }

  DriverInfo->FileType          = FileType;
  DriverInfoData->AllocInfoList = (LIST_ENTRY *)(DriverInfoData + 1);
  InitializeListHead (DriverInfoData->AllocInfoList);
  DriverInfo->CurrentUsage     = 0;
  DriverInfo->PeakUsage        = 0;
  DriverInfo->AllocRecordCount = 0;
  if (PdbSize != 0) {
    DriverInfo->PdbStringOffset = (UINT16)sizeof (MEMORY_PROFILE_DRIVER_INFO);
    DriverInfoData->PdbString   = (CHAR8 *)(DriverInfoData->AllocInfoList + 1);
    CopyMem (DriverInfoData->PdbString, PdbString, PdbSize);
  } else {
    DriverInfo->PdbStringOffset = 0;
    DriverInfoData->PdbString   = NULL;
  }

  InsertTailList (ContextData->DriverInfoList, &DriverInfoData->Link);
  ContextData->Context.ImageCount++;
  ContextData->Context.TotalImageSize += DriverInfo->ImageSize;

  return DriverInfoData;
}

/**
  Register image to DXE.

  @param FileName       File name of the image.
  @param ImageBase      Image base address.
  @param ImageSize      Image size.
  @param FileType       File type of the image.

**/
VOID
RegisterImageToDxe (
  IN EFI_GUID          *FileName,
  IN PHYSICAL_ADDRESS  ImageBase,
  IN UINT64            ImageSize,
  IN EFI_FV_FILETYPE   FileType
  )
{
  EFI_STATUS                         Status;
  EDKII_MEMORY_PROFILE_PROTOCOL      *ProfileProtocol;
  MEDIA_FW_VOL_FILEPATH_DEVICE_PATH  *FilePath;
  UINT8                              TempBuffer[sizeof (MEDIA_FW_VOL_FILEPATH_DEVICE_PATH) + sizeof (EFI_DEVICE_PATH_PROTOCOL)];

  if (IS_UEFI_MEMORY_PROFILE_ENABLED) {
    FilePath = (MEDIA_FW_VOL_FILEPATH_DEVICE_PATH *)TempBuffer;
    Status   = gBS->LocateProtocol (&gEdkiiMemoryProfileGuid, NULL, (VOID **)&ProfileProtocol);
    if (!EFI_ERROR (Status)) {
      EfiInitializeFwVolDevicepathNode (FilePath, FileName);
      SetDevicePathEndNode (FilePath + 1);

      Status = ProfileProtocol->RegisterImage (
                                  ProfileProtocol,
                                  (EFI_DEVICE_PATH_PROTOCOL *)FilePath,
                                  ImageBase,
                                  ImageSize,
                                  FileType
                                  );
    }
  }
}

/**
  Unregister image from DXE.

  @param FileName       File name of the image.
  @param ImageBase      Image base address.
  @param ImageSize      Image size.

**/
VOID
UnregisterImageFromDxe (
  IN EFI_GUID          *FileName,
  IN PHYSICAL_ADDRESS  ImageBase,
  IN UINT64            ImageSize
  )
{
  EFI_STATUS                         Status;
  EDKII_MEMORY_PROFILE_PROTOCOL      *ProfileProtocol;
  MEDIA_FW_VOL_FILEPATH_DEVICE_PATH  *FilePath;
  UINT8                              TempBuffer[sizeof (MEDIA_FW_VOL_FILEPATH_DEVICE_PATH) + sizeof (EFI_DEVICE_PATH_PROTOCOL)];

  if (IS_UEFI_MEMORY_PROFILE_ENABLED) {
    FilePath = (MEDIA_FW_VOL_FILEPATH_DEVICE_PATH *)TempBuffer;
    Status   = gBS->LocateProtocol (&gEdkiiMemoryProfileGuid, NULL, (VOID *)&ProfileProtocol);
    if (!EFI_ERROR (Status)) {
      EfiInitializeFwVolDevicepathNode (FilePath, FileName);
      SetDevicePathEndNode (FilePath + 1);

      Status = ProfileProtocol->UnregisterImage (
                                  ProfileProtocol,
                                  (EFI_DEVICE_PATH_PROTOCOL *)FilePath,
                                  ImageBase,
                                  ImageSize
                                  );
    }
  }
}

/**
  Return if record for this driver is needed..

  @param DriverFilePath     Driver file path.

  @retval TRUE              Record for this driver is needed.
  @retval FALSE             Record for this driver is not needed.

**/
BOOLEAN
NeedRecordThisDriver (
  IN EFI_DEVICE_PATH_PROTOCOL  *DriverFilePath
  )
{
  EFI_DEVICE_PATH_PROTOCOL  *TmpDevicePath;
  EFI_DEVICE_PATH_PROTOCOL  *DevicePathInstance;
  UINTN                     DevicePathSize;
  UINTN                     FilePathSize;

  if (!IsDevicePathValid (mSmramProfileDriverPath, mSmramProfileDriverPathSize)) {
    //
    // Invalid Device Path means record all.
    //
    return TRUE;
  }

  //
  // Record FilePath without end node.
  //
  FilePathSize = GetDevicePathSize (DriverFilePath) - sizeof (EFI_DEVICE_PATH_PROTOCOL);

  DevicePathInstance = mSmramProfileDriverPath;
  do {
    //
    // Find End node (it might be END_ENTIRE or END_INSTANCE)
    //
    TmpDevicePath = DevicePathInstance;
    while (!IsDevicePathEndType (TmpDevicePath)) {
      TmpDevicePath = NextDevicePathNode (TmpDevicePath);
    }

    //
    // Do not compare END node
    //
    DevicePathSize = (UINTN)TmpDevicePath - (UINTN)DevicePathInstance;
    if ((FilePathSize == DevicePathSize) &&
        (CompareMem (DriverFilePath, DevicePathInstance, DevicePathSize) == 0))
    {
      return TRUE;
    }

    //
    // Get next instance
    //
    DevicePathInstance = (EFI_DEVICE_PATH_PROTOCOL *)((UINTN)DevicePathInstance + DevicePathSize + DevicePathNodeLength (TmpDevicePath));
  } while (DevicePathSubType (TmpDevicePath) != END_ENTIRE_DEVICE_PATH_SUBTYPE);

  return FALSE;
}

/**
  Register SMM Core to SMRAM profile.

  @param ContextData    SMRAM profile context.

  @retval TRUE          Register success.
  @retval FALSE         Register fail.

**/
BOOLEAN
RegisterSmmCore (
  IN MEMORY_PROFILE_CONTEXT_DATA  *ContextData
  )
{
  MEMORY_PROFILE_DRIVER_INFO_DATA    *DriverInfoData;
  PHYSICAL_ADDRESS                   ImageBase;
  UINT8                              TempBuffer[sizeof (MEDIA_FW_VOL_FILEPATH_DEVICE_PATH) + sizeof (EFI_DEVICE_PATH_PROTOCOL)];
  MEDIA_FW_VOL_FILEPATH_DEVICE_PATH  *FilePath;

  FilePath = (MEDIA_FW_VOL_FILEPATH_DEVICE_PATH *)TempBuffer;
  EfiInitializeFwVolDevicepathNode (FilePath, &gEfiCallerIdGuid);
  SetDevicePathEndNode (FilePath + 1);

  if (!NeedRecordThisDriver ((EFI_DEVICE_PATH_PROTOCOL *)FilePath)) {
    return FALSE;
  }

  ImageBase      = gSmmCorePrivate->PiSmmCoreImageBase;
  DriverInfoData = BuildDriverInfo (
                     ContextData,
                     &gEfiCallerIdGuid,
                     ImageBase,
                     gSmmCorePrivate->PiSmmCoreImageSize,
                     gSmmCorePrivate->PiSmmCoreEntryPoint,
                     InternalPeCoffGetSubsystem ((VOID *)(UINTN)ImageBase),
                     EFI_FV_FILETYPE_SMM_CORE
                     );
  if (DriverInfoData == NULL) {
    return FALSE;
  }

  return TRUE;
}

/**
  Initialize SMRAM profile.

**/
VOID
SmramProfileInit (
  VOID
  )
{
  MEMORY_PROFILE_CONTEXT_DATA  *SmramProfileContext;

  RegisterImageToDxe (
    &gEfiCallerIdGuid,
    gSmmCorePrivate->PiSmmCoreImageBase,
    gSmmCorePrivate->PiSmmCoreImageSize,
    EFI_FV_FILETYPE_SMM_CORE
    );

  if (!IS_SMRAM_PROFILE_ENABLED) {
    return;
  }

  SmramProfileContext = GetSmramProfileContext ();
  if (SmramProfileContext != NULL) {
    return;
  }

  mSmramProfileGettingStatus = FALSE;
  if ((PcdGet8 (PcdMemoryProfilePropertyMask) & BIT7) != 0) {
    mSmramProfileRecordingEnable = MEMORY_PROFILE_RECORDING_DISABLE;
  } else {
    mSmramProfileRecordingEnable = MEMORY_PROFILE_RECORDING_ENABLE;
  }

  mSmramProfileDriverPathSize = PcdGetSize (PcdMemoryProfileDriverPath);
  mSmramProfileDriverPath     = AllocateCopyPool (mSmramProfileDriverPathSize, PcdGetPtr (PcdMemoryProfileDriverPath));
  mSmramProfileContextPtr     = &mSmramProfileContext;

  RegisterSmmCore (&mSmramProfileContext);

  DEBUG ((DEBUG_INFO, "SmramProfileInit SmramProfileContext - 0x%x\n", &mSmramProfileContext));
}

/**
  Install SMRAM profile protocol.

**/
VOID
SmramProfileInstallProtocol (
  VOID
  )
{
  EFI_HANDLE  Handle;
  EFI_STATUS  Status;

  if (!IS_SMRAM_PROFILE_ENABLED) {
    return;
  }

  Handle = NULL;
  Status = SmmInstallProtocolInterface (
             &Handle,
             &gEdkiiSmmMemoryProfileGuid,
             EFI_NATIVE_INTERFACE,
             &mSmmProfileProtocol
             );
  ASSERT_EFI_ERROR (Status);
}

/**
  Get the GUID file name from the file path.

  @param FilePath  File path.

  @return The GUID file name from the file path.

**/
EFI_GUID *
GetFileNameFromFilePath (
  IN EFI_DEVICE_PATH_PROTOCOL  *FilePath
  )
{
  MEDIA_FW_VOL_FILEPATH_DEVICE_PATH  *ThisFilePath;
  EFI_GUID                           *FileName;

  FileName = NULL;
  if (FilePath != NULL) {
    ThisFilePath = (MEDIA_FW_VOL_FILEPATH_DEVICE_PATH *)FilePath;
    while (!IsDevicePathEnd (ThisFilePath)) {
      FileName = EfiGetNameGuidFromFwVolDevicePathNode (ThisFilePath);
      if (FileName != NULL) {
        break;
      }

      ThisFilePath = (MEDIA_FW_VOL_FILEPATH_DEVICE_PATH *)NextDevicePathNode (ThisFilePath);
    }
  }

  return FileName;
}

/**
  Register SMM image to SMRAM profile.

  @param DriverEntry    SMM image info.
  @param RegisterToDxe  Register image to DXE.

  @return EFI_SUCCESS           Register successfully.
  @return EFI_UNSUPPORTED       Memory profile is unsupported,
                                or memory profile for the image is not required.
  @return EFI_OUT_OF_RESOURCES  No enough resource for this register.

**/
EFI_STATUS
RegisterSmramProfileImage (
  IN EFI_SMM_DRIVER_ENTRY  *DriverEntry,
  IN BOOLEAN               RegisterToDxe
  )
{
  MEMORY_PROFILE_CONTEXT_DATA        *ContextData;
  MEMORY_PROFILE_DRIVER_INFO_DATA    *DriverInfoData;
  UINT8                              TempBuffer[sizeof (MEDIA_FW_VOL_FILEPATH_DEVICE_PATH) + sizeof (EFI_DEVICE_PATH_PROTOCOL)];
  MEDIA_FW_VOL_FILEPATH_DEVICE_PATH  *FilePath;

  if (RegisterToDxe) {
    RegisterImageToDxe (
      &DriverEntry->FileName,
      DriverEntry->ImageBuffer,
      EFI_PAGES_TO_SIZE (DriverEntry->NumberOfPage),
      EFI_FV_FILETYPE_SMM
      );
  }

  if (!IS_SMRAM_PROFILE_ENABLED) {
    return EFI_UNSUPPORTED;
  }

  FilePath = (MEDIA_FW_VOL_FILEPATH_DEVICE_PATH *)TempBuffer;
  EfiInitializeFwVolDevicepathNode (FilePath, &DriverEntry->FileName);
  SetDevicePathEndNode (FilePath + 1);

  if (!NeedRecordThisDriver ((EFI_DEVICE_PATH_PROTOCOL *)FilePath)) {
    return EFI_UNSUPPORTED;
  }

  ContextData = GetSmramProfileContext ();
  if (ContextData == NULL) {
    return EFI_UNSUPPORTED;
  }

  DriverInfoData = BuildDriverInfo (
                     ContextData,
                     &DriverEntry->FileName,
                     DriverEntry->ImageBuffer,
                     EFI_PAGES_TO_SIZE (DriverEntry->NumberOfPage),
                     DriverEntry->ImageEntryPoint,
                     InternalPeCoffGetSubsystem ((VOID *)(UINTN)DriverEntry->ImageBuffer),
                     EFI_FV_FILETYPE_SMM
                     );
  if (DriverInfoData == NULL) {
    return EFI_OUT_OF_RESOURCES;
  }

  return EFI_SUCCESS;
}

/**
  Search image from memory profile.

  @param ContextData    Memory profile context.
  @param FileName       Image file name.
  @param Address        Image Address.

  @return Pointer to memory profile driver info.

**/
MEMORY_PROFILE_DRIVER_INFO_DATA *
GetMemoryProfileDriverInfoByFileNameAndAddress (
  IN MEMORY_PROFILE_CONTEXT_DATA  *ContextData,
  IN EFI_GUID                     *FileName,
  IN PHYSICAL_ADDRESS             Address
  )
{
  MEMORY_PROFILE_DRIVER_INFO       *DriverInfo;
  MEMORY_PROFILE_DRIVER_INFO_DATA  *DriverInfoData;
  LIST_ENTRY                       *DriverLink;
  LIST_ENTRY                       *DriverInfoList;

  DriverInfoList = ContextData->DriverInfoList;

  for (DriverLink = DriverInfoList->ForwardLink;
       DriverLink != DriverInfoList;
       DriverLink = DriverLink->ForwardLink)
  {
    DriverInfoData = CR (
                       DriverLink,
                       MEMORY_PROFILE_DRIVER_INFO_DATA,
                       Link,
                       MEMORY_PROFILE_DRIVER_INFO_SIGNATURE
                       );
    DriverInfo = &DriverInfoData->DriverInfo;
    if ((CompareGuid (&DriverInfo->FileName, FileName)) &&
        (Address >= DriverInfo->ImageBase) &&
        (Address < (DriverInfo->ImageBase + DriverInfo->ImageSize)))
    {
      return DriverInfoData;
    }
  }

  return NULL;
}

/**
  Search image from memory profile.
  It will return image, if (Address >= ImageBuffer) AND (Address < ImageBuffer + ImageSize)

  @param ContextData    Memory profile context.
  @param Address        Image or Function address.

  @return Pointer to memory profile driver info.

**/
MEMORY_PROFILE_DRIVER_INFO_DATA *
GetMemoryProfileDriverInfoFromAddress (
  IN MEMORY_PROFILE_CONTEXT_DATA  *ContextData,
  IN PHYSICAL_ADDRESS             Address
  )
{
  MEMORY_PROFILE_DRIVER_INFO       *DriverInfo;
  MEMORY_PROFILE_DRIVER_INFO_DATA  *DriverInfoData;
  LIST_ENTRY                       *DriverLink;
  LIST_ENTRY                       *DriverInfoList;

  DriverInfoList = ContextData->DriverInfoList;

  for (DriverLink = DriverInfoList->ForwardLink;
       DriverLink != DriverInfoList;
       DriverLink = DriverLink->ForwardLink)
  {
    DriverInfoData = CR (
                       DriverLink,
                       MEMORY_PROFILE_DRIVER_INFO_DATA,
                       Link,
                       MEMORY_PROFILE_DRIVER_INFO_SIGNATURE
                       );
    DriverInfo = &DriverInfoData->DriverInfo;
    if ((Address >= DriverInfo->ImageBase) &&
        (Address < (DriverInfo->ImageBase + DriverInfo->ImageSize)))
    {
      return DriverInfoData;
    }
  }

  return NULL;
}

/**
  Unregister image from SMRAM profile.

  @param DriverEntry        SMM image info.
  @param UnregisterFromDxe  Unregister image from DXE.

  @return EFI_SUCCESS           Unregister successfully.
  @return EFI_UNSUPPORTED       Memory profile is unsupported,
                                or memory profile for the image is not required.
  @return EFI_NOT_FOUND         The image is not found.

**/
EFI_STATUS
UnregisterSmramProfileImage (
  IN EFI_SMM_DRIVER_ENTRY  *DriverEntry,
  IN BOOLEAN               UnregisterFromDxe
  )
{
  EFI_STATUS                         Status;
  MEMORY_PROFILE_CONTEXT_DATA        *ContextData;
  MEMORY_PROFILE_DRIVER_INFO_DATA    *DriverInfoData;
  EFI_GUID                           *FileName;
  PHYSICAL_ADDRESS                   ImageAddress;
  VOID                               *EntryPointInImage;
  UINT8                              TempBuffer[sizeof (MEDIA_FW_VOL_FILEPATH_DEVICE_PATH) + sizeof (EFI_DEVICE_PATH_PROTOCOL)];
  MEDIA_FW_VOL_FILEPATH_DEVICE_PATH  *FilePath;

  if (UnregisterFromDxe) {
    UnregisterImageFromDxe (
      &DriverEntry->FileName,
      DriverEntry->ImageBuffer,
      EFI_PAGES_TO_SIZE (DriverEntry->NumberOfPage)
      );
  }

  if (!IS_SMRAM_PROFILE_ENABLED) {
    return EFI_UNSUPPORTED;
  }

  FilePath = (MEDIA_FW_VOL_FILEPATH_DEVICE_PATH *)TempBuffer;
  EfiInitializeFwVolDevicepathNode (FilePath, &DriverEntry->FileName);
  SetDevicePathEndNode (FilePath + 1);

  if (!NeedRecordThisDriver ((EFI_DEVICE_PATH_PROTOCOL *)FilePath)) {
    return EFI_UNSUPPORTED;
  }

  ContextData = GetSmramProfileContext ();
  if (ContextData == NULL) {
    return EFI_UNSUPPORTED;
  }

  DriverInfoData = NULL;
  FileName       = &DriverEntry->FileName;
  ImageAddress   = DriverEntry->ImageBuffer;
  if ((DriverEntry->ImageEntryPoint < ImageAddress) || (DriverEntry->ImageEntryPoint >= (ImageAddress + EFI_PAGES_TO_SIZE (DriverEntry->NumberOfPage)))) {
    //
    // If the EntryPoint is not in the range of image buffer, it should come from emulation environment.
    // So patch ImageAddress here to align the EntryPoint.
    //
    Status = InternalPeCoffGetEntryPoint ((VOID *)(UINTN)ImageAddress, &EntryPointInImage);
    ASSERT_EFI_ERROR (Status);
    ImageAddress = ImageAddress + (UINTN)DriverEntry->ImageEntryPoint - (UINTN)EntryPointInImage;
  }

  if (FileName != NULL) {
    DriverInfoData = GetMemoryProfileDriverInfoByFileNameAndAddress (ContextData, FileName, ImageAddress);
  }

  if (DriverInfoData == NULL) {
    DriverInfoData = GetMemoryProfileDriverInfoFromAddress (ContextData, ImageAddress);
  }

  if (DriverInfoData == NULL) {
    return EFI_NOT_FOUND;
  }

  ContextData->Context.TotalImageSize -= DriverInfoData->DriverInfo.ImageSize;

  // Keep the ImageBase for RVA calculation in Application.
  // DriverInfoData->DriverInfo.ImageBase = 0;
  DriverInfoData->DriverInfo.ImageSize = 0;

  if (DriverInfoData->DriverInfo.PeakUsage == 0) {
    ContextData->Context.ImageCount--;
    RemoveEntryList (&DriverInfoData->Link);
    //
    // Use SmmInternalFreePool() that will not update profile for this FreePool action.
    //
    SmmInternalFreePool (DriverInfoData);
  }

  return EFI_SUCCESS;
}

/**
  Return if this memory type needs to be recorded into memory profile.
  Only need to record EfiRuntimeServicesCode and EfiRuntimeServicesData for SMRAM profile.

  @param MemoryType     Memory type.

  @retval TRUE          This memory type need to be recorded.
  @retval FALSE         This memory type need not to be recorded.

**/
BOOLEAN
SmmCoreNeedRecordProfile (
  IN EFI_MEMORY_TYPE  MemoryType
  )
{
  UINT64  TestBit;

  if ((MemoryType != EfiRuntimeServicesCode) &&
      (MemoryType != EfiRuntimeServicesData))
  {
    return FALSE;
  }

  TestBit = LShiftU64 (1, MemoryType);

  if ((PcdGet64 (PcdMemoryProfileMemoryType) & TestBit) != 0) {
    return TRUE;
  } else {
    return FALSE;
  }
}

/**
  Convert EFI memory type to profile memory index. The rule is:
  If BIOS memory type (0 ~ EfiMaxMemoryType - 1), ProfileMemoryIndex = MemoryType.
  As SMRAM profile is only to record EfiRuntimeServicesCode and EfiRuntimeServicesData,
  so return input memory type directly.

  @param MemoryType     Memory type.

  @return EFI memory type as profile memory index.

**/
EFI_MEMORY_TYPE
GetProfileMemoryIndex (
  IN EFI_MEMORY_TYPE  MemoryType
  )
{
  return MemoryType;
}

/**
  Update SMRAM profile FreeMemoryPages information

  @param ContextData    Memory profile context.

**/
VOID
SmramProfileUpdateFreePages (
  IN MEMORY_PROFILE_CONTEXT_DATA  *ContextData
  )
{
  LIST_ENTRY      *Node;
  FREE_PAGE_LIST  *Pages;
  LIST_ENTRY      *FreePageList;
  UINTN           NumberOfPages;

  NumberOfPages = 0;
  FreePageList  = &mSmmMemoryMap;
  for (Node = FreePageList->BackLink;
       Node != FreePageList;
       Node = Node->BackLink)
  {
    Pages          = BASE_CR (Node, FREE_PAGE_LIST, Link);
    NumberOfPages += Pages->NumberOfPages;
  }

  mSmramFreeMemory.TotalFreeMemoryPages = NumberOfPages;

  if (NumberOfPages <= SMRAM_INFO_DUMP_PAGE_THRESHOLD) {
    DumpSmramInfo ();
  }
}

/**
  Update SMRAM profile Allocate information.

  @param CallerAddress  Address of caller who call Allocate.
  @param Action         This Allocate action.
  @param MemoryType     Memory type.
  @param Size           Buffer size.
  @param Buffer         Buffer address.
  @param ActionString   String for memory profile action.

  @return EFI_SUCCESS           Memory profile is updated.
  @return EFI_UNSUPPORTED       Memory profile is unsupported,
                                or memory profile for the image is not required.
  @return EFI_OUT_OF_RESOURCES  No enough resource to update memory profile for allocate action.

**/
EFI_STATUS
SmmCoreUpdateProfileAllocate (
  IN PHYSICAL_ADDRESS       CallerAddress,
  IN MEMORY_PROFILE_ACTION  Action,
  IN EFI_MEMORY_TYPE        MemoryType,
  IN UINTN                  Size,
  IN VOID                   *Buffer,
  IN CHAR8                  *ActionString OPTIONAL
  )
{
  EFI_STATUS                       Status;
  MEMORY_PROFILE_CONTEXT           *Context;
  MEMORY_PROFILE_DRIVER_INFO       *DriverInfo;
  MEMORY_PROFILE_ALLOC_INFO        *AllocInfo;
  MEMORY_PROFILE_CONTEXT_DATA      *ContextData;
  MEMORY_PROFILE_DRIVER_INFO_DATA  *DriverInfoData;
  MEMORY_PROFILE_ALLOC_INFO_DATA   *AllocInfoData;
  EFI_MEMORY_TYPE                  ProfileMemoryIndex;
  MEMORY_PROFILE_ACTION            BasicAction;
  UINTN                            ActionStringSize;
  UINTN                            ActionStringOccupiedSize;

  BasicAction = Action & MEMORY_PROFILE_ACTION_BASIC_MASK;

  ContextData = GetSmramProfileContext ();
  if (ContextData == NULL) {
    return EFI_UNSUPPORTED;
  }

  DriverInfoData = GetMemoryProfileDriverInfoFromAddress (ContextData, CallerAddress);
  if (DriverInfoData == NULL) {
    return EFI_UNSUPPORTED;
  }

  ActionStringSize         = 0;
  ActionStringOccupiedSize = 0;
  if (ActionString != NULL) {
    ActionStringSize         = AsciiStrSize (ActionString);
    ActionStringOccupiedSize = GET_OCCUPIED_SIZE (ActionStringSize, sizeof (UINT64));
  }

  //
  // Use SmmInternalAllocatePool() that will not update profile for this AllocatePool action.
  //
  AllocInfoData = NULL;
  Status        = SmmInternalAllocatePool (
                    EfiRuntimeServicesData,
                    sizeof (*AllocInfoData) + ActionStringSize,
                    (VOID **)&AllocInfoData
                    );
  if (EFI_ERROR (Status)) {
    return EFI_OUT_OF_RESOURCES;
  }

  ASSERT (AllocInfoData != NULL);

  //
  // Only update SequenceCount if and only if it is basic action.
  //
  if (Action == BasicAction) {
    ContextData->Context.SequenceCount++;
  }

  AllocInfo                   = &AllocInfoData->AllocInfo;
  AllocInfoData->Signature    = MEMORY_PROFILE_ALLOC_INFO_SIGNATURE;
  AllocInfo->Header.Signature = MEMORY_PROFILE_ALLOC_INFO_SIGNATURE;
  AllocInfo->Header.Length    = (UINT16)(sizeof (MEMORY_PROFILE_ALLOC_INFO) + ActionStringOccupiedSize);
  AllocInfo->Header.Revision  = MEMORY_PROFILE_ALLOC_INFO_REVISION;
  AllocInfo->CallerAddress    = CallerAddress;
  AllocInfo->SequenceId       = ContextData->Context.SequenceCount;
  AllocInfo->Action           = Action;
  AllocInfo->MemoryType       = MemoryType;
  AllocInfo->Buffer           = (PHYSICAL_ADDRESS)(UINTN)Buffer;
  AllocInfo->Size             = Size;
  if (ActionString != NULL) {
    AllocInfo->ActionStringOffset = (UINT16)sizeof (MEMORY_PROFILE_ALLOC_INFO);
    AllocInfoData->ActionString   = (CHAR8 *)(AllocInfoData + 1);
    CopyMem (AllocInfoData->ActionString, ActionString, ActionStringSize);
  } else {
    AllocInfo->ActionStringOffset = 0;
    AllocInfoData->ActionString   = NULL;
  }

  InsertTailList (DriverInfoData->AllocInfoList, &AllocInfoData->Link);

  Context    = &ContextData->Context;
  DriverInfo = &DriverInfoData->DriverInfo;
  DriverInfo->AllocRecordCount++;

  //
  // Update summary if and only if it is basic action.
  //
  if (Action == BasicAction) {
    ProfileMemoryIndex = GetProfileMemoryIndex (MemoryType);

    DriverInfo->CurrentUsage += Size;
    if (DriverInfo->PeakUsage < DriverInfo->CurrentUsage) {
      DriverInfo->PeakUsage = DriverInfo->CurrentUsage;
    }

    DriverInfo->CurrentUsageByType[ProfileMemoryIndex] += Size;
    if (DriverInfo->PeakUsageByType[ProfileMemoryIndex] < DriverInfo->CurrentUsageByType[ProfileMemoryIndex]) {
      DriverInfo->PeakUsageByType[ProfileMemoryIndex] = DriverInfo->CurrentUsageByType[ProfileMemoryIndex];
    }

    Context->CurrentTotalUsage += Size;
    if (Context->PeakTotalUsage < Context->CurrentTotalUsage) {
      Context->PeakTotalUsage = Context->CurrentTotalUsage;
    }

    Context->CurrentTotalUsageByType[ProfileMemoryIndex] += Size;
    if (Context->PeakTotalUsageByType[ProfileMemoryIndex] < Context->CurrentTotalUsageByType[ProfileMemoryIndex]) {
      Context->PeakTotalUsageByType[ProfileMemoryIndex] = Context->CurrentTotalUsageByType[ProfileMemoryIndex];
    }

    SmramProfileUpdateFreePages (ContextData);
  }

  return EFI_SUCCESS;
}

/**
  Get memory profile alloc info from memory profile

  @param DriverInfoData     Driver info
  @param BasicAction        This Free basic action
  @param Size               Buffer size
  @param Buffer             Buffer address

  @return Pointer to memory profile alloc info.
**/
MEMORY_PROFILE_ALLOC_INFO_DATA *
GetMemoryProfileAllocInfoFromAddress (
  IN MEMORY_PROFILE_DRIVER_INFO_DATA  *DriverInfoData,
  IN MEMORY_PROFILE_ACTION            BasicAction,
  IN UINTN                            Size,
  IN VOID                             *Buffer
  )
{
  LIST_ENTRY                      *AllocInfoList;
  LIST_ENTRY                      *AllocLink;
  MEMORY_PROFILE_ALLOC_INFO       *AllocInfo;
  MEMORY_PROFILE_ALLOC_INFO_DATA  *AllocInfoData;

  AllocInfoList = DriverInfoData->AllocInfoList;

  for (AllocLink = AllocInfoList->ForwardLink;
       AllocLink != AllocInfoList;
       AllocLink = AllocLink->ForwardLink)
  {
    AllocInfoData = CR (
                      AllocLink,
                      MEMORY_PROFILE_ALLOC_INFO_DATA,
                      Link,
                      MEMORY_PROFILE_ALLOC_INFO_SIGNATURE
                      );
    AllocInfo = &AllocInfoData->AllocInfo;
    if ((AllocInfo->Action & MEMORY_PROFILE_ACTION_BASIC_MASK) != BasicAction) {
      continue;
    }

    switch (BasicAction) {
      case MemoryProfileActionAllocatePages:
        if ((AllocInfo->Buffer <= (PHYSICAL_ADDRESS)(UINTN)Buffer) &&
            ((AllocInfo->Buffer + AllocInfo->Size) >= ((PHYSICAL_ADDRESS)(UINTN)Buffer + Size)))
        {
          return AllocInfoData;
        }

        break;
      case MemoryProfileActionAllocatePool:
        if (AllocInfo->Buffer == (PHYSICAL_ADDRESS)(UINTN)Buffer) {
          return AllocInfoData;
        }

        break;
      default:
        ASSERT (FALSE);
        break;
    }
  }

  return NULL;
}

/**
  Update SMRAM profile Free information.

  @param CallerAddress  Address of caller who call Free.
  @param Action         This Free action.
  @param Size           Buffer size.
  @param Buffer         Buffer address.

  @return EFI_SUCCESS           Memory profile is updated.
  @return EFI_UNSUPPORTED       Memory profile is unsupported.
  @return EFI_NOT_FOUND         No matched allocate info found for free action.

**/
EFI_STATUS
SmmCoreUpdateProfileFree (
  IN PHYSICAL_ADDRESS       CallerAddress,
  IN MEMORY_PROFILE_ACTION  Action,
  IN UINTN                  Size,
  IN VOID                   *Buffer
  )
{
  MEMORY_PROFILE_CONTEXT           *Context;
  MEMORY_PROFILE_DRIVER_INFO       *DriverInfo;
  MEMORY_PROFILE_ALLOC_INFO        *AllocInfo;
  MEMORY_PROFILE_CONTEXT_DATA      *ContextData;
  MEMORY_PROFILE_DRIVER_INFO_DATA  *DriverInfoData;
  LIST_ENTRY                       *DriverLink;
  LIST_ENTRY                       *DriverInfoList;
  MEMORY_PROFILE_DRIVER_INFO_DATA  *ThisDriverInfoData;
  MEMORY_PROFILE_ALLOC_INFO_DATA   *AllocInfoData;
  EFI_MEMORY_TYPE                  ProfileMemoryIndex;
  MEMORY_PROFILE_ACTION            BasicAction;
  BOOLEAN                          Found;

  BasicAction = Action & MEMORY_PROFILE_ACTION_BASIC_MASK;

  ContextData = GetSmramProfileContext ();
  if (ContextData == NULL) {
    return EFI_UNSUPPORTED;
  }

  DriverInfoData = GetMemoryProfileDriverInfoFromAddress (ContextData, CallerAddress);

  //
  // Do not return if DriverInfoData == NULL here,
  // because driver A might free memory allocated by driver B.
  //

  //
  // Need use do-while loop to find all possible record,
  // because one address might be recorded multiple times.
  //
  Found         = FALSE;
  AllocInfoData = NULL;
  do {
    if (DriverInfoData != NULL) {
      switch (BasicAction) {
        case MemoryProfileActionFreePages:
          AllocInfoData = GetMemoryProfileAllocInfoFromAddress (DriverInfoData, MemoryProfileActionAllocatePages, Size, Buffer);
          break;
        case MemoryProfileActionFreePool:
          AllocInfoData = GetMemoryProfileAllocInfoFromAddress (DriverInfoData, MemoryProfileActionAllocatePool, 0, Buffer);
          break;
        default:
          ASSERT (FALSE);
          AllocInfoData = NULL;
          break;
      }
    }

    if (AllocInfoData == NULL) {
      //
      // Legal case, because driver A might free memory allocated by driver B, by some protocol.
      //
      DriverInfoList = ContextData->DriverInfoList;

      for (DriverLink = DriverInfoList->ForwardLink;
           DriverLink != DriverInfoList;
           DriverLink = DriverLink->ForwardLink)
      {
        ThisDriverInfoData = CR (
                               DriverLink,
                               MEMORY_PROFILE_DRIVER_INFO_DATA,
                               Link,
                               MEMORY_PROFILE_DRIVER_INFO_SIGNATURE
                               );
        switch (BasicAction) {
          case MemoryProfileActionFreePages:
            AllocInfoData = GetMemoryProfileAllocInfoFromAddress (ThisDriverInfoData, MemoryProfileActionAllocatePages, Size, Buffer);
            break;
          case MemoryProfileActionFreePool:
            AllocInfoData = GetMemoryProfileAllocInfoFromAddress (ThisDriverInfoData, MemoryProfileActionAllocatePool, 0, Buffer);
            break;
          default:
            ASSERT (FALSE);
            AllocInfoData = NULL;
            break;
        }

        if (AllocInfoData != NULL) {
          DriverInfoData = ThisDriverInfoData;
          break;
        }
      }

      if (AllocInfoData == NULL) {
        //
        // If (!Found), no matched allocate info is found for this free action.
        // It is because the specified memory type allocate actions have been filtered by
        // CoreNeedRecordProfile(), but free actions have no memory type information,
        // they can not be filtered by CoreNeedRecordProfile(). Then, they will be
        // filtered here.
        //
        // If (Found), it is normal exit path.
        return (Found ? EFI_SUCCESS : EFI_NOT_FOUND);
      }
    }

    ASSERT (DriverInfoData != NULL);
    ASSERT (AllocInfoData != NULL);

    Found = TRUE;

    Context    = &ContextData->Context;
    DriverInfo = &DriverInfoData->DriverInfo;
    AllocInfo  = &AllocInfoData->AllocInfo;

    DriverInfo->AllocRecordCount--;
    //
    // Update summary if and only if it is basic action.
    //
    if (AllocInfo->Action == (AllocInfo->Action & MEMORY_PROFILE_ACTION_BASIC_MASK)) {
      ProfileMemoryIndex = GetProfileMemoryIndex (AllocInfo->MemoryType);

      Context->CurrentTotalUsage                           -= AllocInfo->Size;
      Context->CurrentTotalUsageByType[ProfileMemoryIndex] -= AllocInfo->Size;

      DriverInfo->CurrentUsage                           -= AllocInfo->Size;
      DriverInfo->CurrentUsageByType[ProfileMemoryIndex] -= AllocInfo->Size;
    }

    RemoveEntryList (&AllocInfoData->Link);

    if (BasicAction == MemoryProfileActionFreePages) {
      if (AllocInfo->Buffer != (PHYSICAL_ADDRESS)(UINTN)Buffer) {
        SmmCoreUpdateProfileAllocate (
          AllocInfo->CallerAddress,
          AllocInfo->Action,
          AllocInfo->MemoryType,
          (UINTN)((PHYSICAL_ADDRESS)(UINTN)Buffer - AllocInfo->Buffer),
          (VOID *)(UINTN)AllocInfo->Buffer,
          AllocInfoData->ActionString
          );
      }

      if (AllocInfo->Buffer + AllocInfo->Size != ((PHYSICAL_ADDRESS)(UINTN)Buffer + Size)) {
        SmmCoreUpdateProfileAllocate (
          AllocInfo->CallerAddress,
          AllocInfo->Action,
          AllocInfo->MemoryType,
          (UINTN)((AllocInfo->Buffer + AllocInfo->Size) - ((PHYSICAL_ADDRESS)(UINTN)Buffer + Size)),
          (VOID *)((UINTN)Buffer + Size),
          AllocInfoData->ActionString
          );
      }
    }

    //
    // Use SmmInternalFreePool() that will not update profile for this FreePool action.
    //
    SmmInternalFreePool (AllocInfoData);
  } while (TRUE);
}

/**
  Update SMRAM profile information.

  @param CallerAddress  Address of caller who call Allocate or Free.
  @param Action         This Allocate or Free action.
  @param MemoryType     Memory type.
                        EfiMaxMemoryType means the MemoryType is unknown.
  @param Size           Buffer size.
  @param Buffer         Buffer address.
  @param ActionString   String for memory profile action.
                        Only needed for user defined allocate action.

  @return EFI_SUCCESS           Memory profile is updated.
  @return EFI_UNSUPPORTED       Memory profile is unsupported,
                                or memory profile for the image is not required,
                                or memory profile for the memory type is not required.
  @return EFI_ACCESS_DENIED     It is during memory profile data getting.
  @return EFI_ABORTED           Memory profile recording is not enabled.
  @return EFI_OUT_OF_RESOURCES  No enough resource to update memory profile for allocate action.
  @return EFI_NOT_FOUND         No matched allocate info found for free action.

**/
EFI_STATUS
EFIAPI
SmmCoreUpdateProfile (
  IN PHYSICAL_ADDRESS       CallerAddress,
  IN MEMORY_PROFILE_ACTION  Action,
  IN EFI_MEMORY_TYPE        MemoryType, // Valid for AllocatePages/AllocatePool
  IN UINTN                  Size,       // Valid for AllocatePages/FreePages/AllocatePool
  IN VOID                   *Buffer,
  IN CHAR8                  *ActionString OPTIONAL
  )
{
  EFI_STATUS                   Status;
  MEMORY_PROFILE_CONTEXT_DATA  *ContextData;
  MEMORY_PROFILE_ACTION        BasicAction;

  if (!IS_SMRAM_PROFILE_ENABLED) {
    return EFI_UNSUPPORTED;
  }

  if (mSmramProfileGettingStatus) {
    return EFI_ACCESS_DENIED;
  }

  if (!mSmramProfileRecordingEnable) {
    return EFI_ABORTED;
  }

  //
  // Get the basic action to know how to process the record
  //
  BasicAction = Action & MEMORY_PROFILE_ACTION_BASIC_MASK;

  //
  // Free operations have no memory type information, so skip the check.
  //
  if ((BasicAction == MemoryProfileActionAllocatePages) || (BasicAction == MemoryProfileActionAllocatePool)) {
    //
    // Only record limited MemoryType.
    //
    if (!SmmCoreNeedRecordProfile (MemoryType)) {
      return EFI_UNSUPPORTED;
    }
  }

  ContextData = GetSmramProfileContext ();
  if (ContextData == NULL) {
    return EFI_UNSUPPORTED;
  }

  switch (BasicAction) {
    case MemoryProfileActionAllocatePages:
      Status = SmmCoreUpdateProfileAllocate (CallerAddress, Action, MemoryType, Size, Buffer, ActionString);
      break;
    case MemoryProfileActionFreePages:
      Status = SmmCoreUpdateProfileFree (CallerAddress, Action, Size, Buffer);
      break;
    case MemoryProfileActionAllocatePool:
      Status = SmmCoreUpdateProfileAllocate (CallerAddress, Action, MemoryType, Size, Buffer, ActionString);
      break;
    case MemoryProfileActionFreePool:
      Status = SmmCoreUpdateProfileFree (CallerAddress, Action, 0, Buffer);
      break;
    default:
      ASSERT (FALSE);
      Status = EFI_UNSUPPORTED;
      break;
  }

  return Status;
}

/**
  SMRAM profile ready to lock callback function.

**/
VOID
SmramProfileReadyToLock (
  VOID
  )
{
  if (!IS_SMRAM_PROFILE_ENABLED) {
    return;
  }

  DEBUG ((DEBUG_INFO, "SmramProfileReadyToLock\n"));
  mSmramReadyToLock = TRUE;
}

////////////////////

/**
  Get SMRAM profile data size.

  @return SMRAM profile data size.

**/
UINTN
SmramProfileGetDataSize (
  VOID
  )
{
  MEMORY_PROFILE_CONTEXT_DATA      *ContextData;
  MEMORY_PROFILE_DRIVER_INFO_DATA  *DriverInfoData;
  MEMORY_PROFILE_ALLOC_INFO_DATA   *AllocInfoData;
  LIST_ENTRY                       *DriverInfoList;
  LIST_ENTRY                       *DriverLink;
  LIST_ENTRY                       *AllocInfoList;
  LIST_ENTRY                       *AllocLink;
  UINTN                            TotalSize;
  LIST_ENTRY                       *Node;
  LIST_ENTRY                       *FreePageList;
  LIST_ENTRY                       *FreePoolList;
  FREE_POOL_HEADER                 *Pool;
  UINTN                            PoolListIndex;
  UINTN                            Index;
  UINTN                            SmmPoolTypeIndex;

  ContextData = GetSmramProfileContext ();
  if (ContextData == NULL) {
    return 0;
  }

  TotalSize = sizeof (MEMORY_PROFILE_CONTEXT);

  DriverInfoList = ContextData->DriverInfoList;
  for (DriverLink = DriverInfoList->ForwardLink;
       DriverLink != DriverInfoList;
       DriverLink = DriverLink->ForwardLink)
  {
    DriverInfoData = CR (
                       DriverLink,
                       MEMORY_PROFILE_DRIVER_INFO_DATA,
                       Link,
                       MEMORY_PROFILE_DRIVER_INFO_SIGNATURE
                       );
    TotalSize += DriverInfoData->DriverInfo.Header.Length;

    AllocInfoList = DriverInfoData->AllocInfoList;
    for (AllocLink = AllocInfoList->ForwardLink;
         AllocLink != AllocInfoList;
         AllocLink = AllocLink->ForwardLink)
    {
      AllocInfoData = CR (
                        AllocLink,
                        MEMORY_PROFILE_ALLOC_INFO_DATA,
                        Link,
                        MEMORY_PROFILE_ALLOC_INFO_SIGNATURE
                        );
      TotalSize += AllocInfoData->AllocInfo.Header.Length;
    }
  }

  Index        = 0;
  FreePageList = &mSmmMemoryMap;
  for (Node = FreePageList->BackLink;
       Node != FreePageList;
       Node = Node->BackLink)
  {
    Index++;
  }

  for (SmmPoolTypeIndex = 0; SmmPoolTypeIndex < SmmPoolTypeMax; SmmPoolTypeIndex++) {
    for (PoolListIndex = 0; PoolListIndex < MAX_POOL_INDEX; PoolListIndex++) {
      FreePoolList = &mSmmPoolLists[SmmPoolTypeIndex][PoolListIndex];
      for (Node = FreePoolList->BackLink;
           Node != FreePoolList;
           Node = Node->BackLink)
      {
        Pool = BASE_CR (Node, FREE_POOL_HEADER, Link);
        if (Pool->Header.Available) {
          Index++;
        }
      }
    }
  }

  TotalSize += (sizeof (MEMORY_PROFILE_FREE_MEMORY) + Index * sizeof (MEMORY_PROFILE_DESCRIPTOR));
  TotalSize += (sizeof (MEMORY_PROFILE_MEMORY_RANGE) + mFullSmramRangeCount * sizeof (MEMORY_PROFILE_DESCRIPTOR));

  return TotalSize;
}

/**
  Copy SMRAM profile data.

  @param ProfileBuffer  The buffer to hold SMRAM profile data.
  @param ProfileSize    On input, profile buffer size.
                        On output, actual profile data size copied.
  @param ProfileOffset  On input, profile buffer offset to copy.
                        On output, next time profile buffer offset to copy.

**/
VOID
SmramProfileCopyData (
  OUT VOID       *ProfileBuffer,
  IN OUT UINT64  *ProfileSize,
  IN OUT UINT64  *ProfileOffset
  )
{
  MEMORY_PROFILE_CONTEXT           *Context;
  MEMORY_PROFILE_DRIVER_INFO       *DriverInfo;
  MEMORY_PROFILE_ALLOC_INFO        *AllocInfo;
  MEMORY_PROFILE_CONTEXT_DATA      *ContextData;
  MEMORY_PROFILE_DRIVER_INFO_DATA  *DriverInfoData;
  MEMORY_PROFILE_ALLOC_INFO_DATA   *AllocInfoData;
  LIST_ENTRY                       *DriverInfoList;
  LIST_ENTRY                       *DriverLink;
  LIST_ENTRY                       *AllocInfoList;
  LIST_ENTRY                       *AllocLink;
  LIST_ENTRY                       *Node;
  FREE_PAGE_LIST                   *Pages;
  LIST_ENTRY                       *FreePageList;
  LIST_ENTRY                       *FreePoolList;
  FREE_POOL_HEADER                 *Pool;
  UINTN                            PoolListIndex;
  UINT32                           Index;
  MEMORY_PROFILE_FREE_MEMORY       *FreeMemory;
  MEMORY_PROFILE_MEMORY_RANGE      *MemoryRange;
  MEMORY_PROFILE_DESCRIPTOR        *MemoryProfileDescriptor;
  UINT64                           Offset;
  UINT64                           RemainingSize;
  UINTN                            PdbSize;
  UINTN                            ActionStringSize;
  UINTN                            SmmPoolTypeIndex;

  ContextData = GetSmramProfileContext ();
  if (ContextData == NULL) {
    return;
  }

  RemainingSize = *ProfileSize;
  Offset        = 0;

  if (*ProfileOffset < sizeof (MEMORY_PROFILE_CONTEXT)) {
    if (RemainingSize >= sizeof (MEMORY_PROFILE_CONTEXT)) {
      Context = ProfileBuffer;
      CopyMem (Context, &ContextData->Context, sizeof (MEMORY_PROFILE_CONTEXT));
      RemainingSize -= sizeof (MEMORY_PROFILE_CONTEXT);
      ProfileBuffer  = (UINT8 *)ProfileBuffer + sizeof (MEMORY_PROFILE_CONTEXT);
    } else {
      goto Done;
    }
  }

  Offset += sizeof (MEMORY_PROFILE_CONTEXT);

  DriverInfoList = ContextData->DriverInfoList;
  for (DriverLink = DriverInfoList->ForwardLink;
       DriverLink != DriverInfoList;
       DriverLink = DriverLink->ForwardLink)
  {
    DriverInfoData = CR (
                       DriverLink,
                       MEMORY_PROFILE_DRIVER_INFO_DATA,
                       Link,
                       MEMORY_PROFILE_DRIVER_INFO_SIGNATURE
                       );
    if (*ProfileOffset < (Offset + DriverInfoData->DriverInfo.Header.Length)) {
      if (RemainingSize >= DriverInfoData->DriverInfo.Header.Length) {
        DriverInfo = ProfileBuffer;
        CopyMem (DriverInfo, &DriverInfoData->DriverInfo, sizeof (MEMORY_PROFILE_DRIVER_INFO));
        if (DriverInfo->PdbStringOffset != 0) {
          PdbSize = AsciiStrSize (DriverInfoData->PdbString);
          CopyMem ((VOID *)((UINTN)DriverInfo + DriverInfo->PdbStringOffset), DriverInfoData->PdbString, PdbSize);
        }

        RemainingSize -= DriverInfo->Header.Length;
        ProfileBuffer  = (UINT8 *)ProfileBuffer + DriverInfo->Header.Length;
      } else {
        goto Done;
      }
    }

    Offset += DriverInfoData->DriverInfo.Header.Length;

    AllocInfoList = DriverInfoData->AllocInfoList;
    for (AllocLink = AllocInfoList->ForwardLink;
         AllocLink != AllocInfoList;
         AllocLink = AllocLink->ForwardLink)
    {
      AllocInfoData = CR (
                        AllocLink,
                        MEMORY_PROFILE_ALLOC_INFO_DATA,
                        Link,
                        MEMORY_PROFILE_ALLOC_INFO_SIGNATURE
                        );
      if (*ProfileOffset < (Offset + AllocInfoData->AllocInfo.Header.Length)) {
        if (RemainingSize >= AllocInfoData->AllocInfo.Header.Length) {
          AllocInfo = ProfileBuffer;
          CopyMem (AllocInfo, &AllocInfoData->AllocInfo, sizeof (MEMORY_PROFILE_ALLOC_INFO));
          if (AllocInfo->ActionStringOffset) {
            ActionStringSize = AsciiStrSize (AllocInfoData->ActionString);
            CopyMem ((VOID *)((UINTN)AllocInfo + AllocInfo->ActionStringOffset), AllocInfoData->ActionString, ActionStringSize);
          }

          RemainingSize -= AllocInfo->Header.Length;
          ProfileBuffer  = (UINT8 *)ProfileBuffer + AllocInfo->Header.Length;
        } else {
          goto Done;
        }
      }

      Offset += AllocInfoData->AllocInfo.Header.Length;
    }
  }

  if (*ProfileOffset < (Offset + sizeof (MEMORY_PROFILE_FREE_MEMORY))) {
    if (RemainingSize >= sizeof (MEMORY_PROFILE_FREE_MEMORY)) {
      FreeMemory = ProfileBuffer;
      CopyMem (FreeMemory, &mSmramFreeMemory, sizeof (MEMORY_PROFILE_FREE_MEMORY));
      Index        = 0;
      FreePageList = &mSmmMemoryMap;
      for (Node = FreePageList->BackLink;
           Node != FreePageList;
           Node = Node->BackLink)
      {
        Index++;
      }

      for (SmmPoolTypeIndex = 0; SmmPoolTypeIndex < SmmPoolTypeMax; SmmPoolTypeIndex++) {
        for (PoolListIndex = 0; PoolListIndex < MAX_POOL_INDEX; PoolListIndex++) {
          FreePoolList = &mSmmPoolLists[SmmPoolTypeIndex][MAX_POOL_INDEX - PoolListIndex - 1];
          for (Node = FreePoolList->BackLink;
               Node != FreePoolList;
               Node = Node->BackLink)
          {
            Pool = BASE_CR (Node, FREE_POOL_HEADER, Link);
            if (Pool->Header.Available) {
              Index++;
            }
          }
        }
      }

      FreeMemory->FreeMemoryEntryCount = Index;

      RemainingSize -= sizeof (MEMORY_PROFILE_FREE_MEMORY);
      ProfileBuffer  = (UINT8 *)ProfileBuffer + sizeof (MEMORY_PROFILE_FREE_MEMORY);
    } else {
      goto Done;
    }
  }

  Offset      += sizeof (MEMORY_PROFILE_FREE_MEMORY);
  FreePageList = &mSmmMemoryMap;
  for (Node = FreePageList->BackLink;
       Node != FreePageList;
       Node = Node->BackLink)
  {
    if (*ProfileOffset < (Offset + sizeof (MEMORY_PROFILE_DESCRIPTOR))) {
      if (RemainingSize >= sizeof (MEMORY_PROFILE_DESCRIPTOR)) {
        Pages                                     = BASE_CR (Node, FREE_PAGE_LIST, Link);
        MemoryProfileDescriptor                   = ProfileBuffer;
        MemoryProfileDescriptor->Header.Signature = MEMORY_PROFILE_DESCRIPTOR_SIGNATURE;
        MemoryProfileDescriptor->Header.Length    = sizeof (MEMORY_PROFILE_DESCRIPTOR);
        MemoryProfileDescriptor->Header.Revision  = MEMORY_PROFILE_DESCRIPTOR_REVISION;
        MemoryProfileDescriptor->Address          = (PHYSICAL_ADDRESS)(UINTN)Pages;
        MemoryProfileDescriptor->Size             = EFI_PAGES_TO_SIZE (Pages->NumberOfPages);

        RemainingSize -= sizeof (MEMORY_PROFILE_DESCRIPTOR);
        ProfileBuffer  = (UINT8 *)ProfileBuffer + sizeof (MEMORY_PROFILE_DESCRIPTOR);
      } else {
        goto Done;
      }
    }

    Offset += sizeof (MEMORY_PROFILE_DESCRIPTOR);
  }

  for (SmmPoolTypeIndex = 0; SmmPoolTypeIndex < SmmPoolTypeMax; SmmPoolTypeIndex++) {
    for (PoolListIndex = 0; PoolListIndex < MAX_POOL_INDEX; PoolListIndex++) {
      FreePoolList = &mSmmPoolLists[SmmPoolTypeIndex][MAX_POOL_INDEX - PoolListIndex - 1];
      for (Node = FreePoolList->BackLink;
           Node != FreePoolList;
           Node = Node->BackLink)
      {
        Pool = BASE_CR (Node, FREE_POOL_HEADER, Link);
        if (Pool->Header.Available) {
          if (*ProfileOffset < (Offset + sizeof (MEMORY_PROFILE_DESCRIPTOR))) {
            if (RemainingSize >= sizeof (MEMORY_PROFILE_DESCRIPTOR)) {
              MemoryProfileDescriptor                   = ProfileBuffer;
              MemoryProfileDescriptor->Header.Signature = MEMORY_PROFILE_DESCRIPTOR_SIGNATURE;
              MemoryProfileDescriptor->Header.Length    = sizeof (MEMORY_PROFILE_DESCRIPTOR);
              MemoryProfileDescriptor->Header.Revision  = MEMORY_PROFILE_DESCRIPTOR_REVISION;
              MemoryProfileDescriptor->Address          = (PHYSICAL_ADDRESS)(UINTN)Pool;
              MemoryProfileDescriptor->Size             = Pool->Header.Size;

              RemainingSize -= sizeof (MEMORY_PROFILE_DESCRIPTOR);
              ProfileBuffer  = (UINT8 *)ProfileBuffer + sizeof (MEMORY_PROFILE_DESCRIPTOR);
            } else {
              goto Done;
            }
          }

          Offset += sizeof (MEMORY_PROFILE_DESCRIPTOR);
        }
      }
    }
  }

  if (*ProfileOffset < (Offset + sizeof (MEMORY_PROFILE_MEMORY_RANGE))) {
    if (RemainingSize >= sizeof (MEMORY_PROFILE_MEMORY_RANGE)) {
      MemoryRange                   = ProfileBuffer;
      MemoryRange->Header.Signature = MEMORY_PROFILE_MEMORY_RANGE_SIGNATURE;
      MemoryRange->Header.Length    = sizeof (MEMORY_PROFILE_MEMORY_RANGE);
      MemoryRange->Header.Revision  = MEMORY_PROFILE_MEMORY_RANGE_REVISION;
      MemoryRange->MemoryRangeCount = (UINT32)mFullSmramRangeCount;

      RemainingSize -= sizeof (MEMORY_PROFILE_MEMORY_RANGE);
      ProfileBuffer  = (UINT8 *)ProfileBuffer + sizeof (MEMORY_PROFILE_MEMORY_RANGE);
    } else {
      goto Done;
    }
  }

  Offset += sizeof (MEMORY_PROFILE_MEMORY_RANGE);
  for (Index = 0; Index < mFullSmramRangeCount; Index++) {
    if (*ProfileOffset < (Offset + sizeof (MEMORY_PROFILE_DESCRIPTOR))) {
      if (RemainingSize >= sizeof (MEMORY_PROFILE_DESCRIPTOR)) {
        MemoryProfileDescriptor                   = ProfileBuffer;
        MemoryProfileDescriptor->Header.Signature = MEMORY_PROFILE_DESCRIPTOR_SIGNATURE;
        MemoryProfileDescriptor->Header.Length    = sizeof (MEMORY_PROFILE_DESCRIPTOR);
        MemoryProfileDescriptor->Header.Revision  = MEMORY_PROFILE_DESCRIPTOR_REVISION;
        MemoryProfileDescriptor->Address          = mFullSmramRanges[Index].PhysicalStart;
        MemoryProfileDescriptor->Size             = mFullSmramRanges[Index].PhysicalSize;

        RemainingSize -= sizeof (MEMORY_PROFILE_DESCRIPTOR);
        ProfileBuffer  = (UINT8 *)ProfileBuffer + sizeof (MEMORY_PROFILE_DESCRIPTOR);
      } else {
        goto Done;
      }
    }

    Offset += sizeof (MEMORY_PROFILE_DESCRIPTOR);
  }

Done:
  //
  // On output, actual profile data size copied.
  //
  *ProfileSize -= RemainingSize;
  //
  // On output, next time profile buffer offset to copy.
  //
  *ProfileOffset = Offset;
}

/**
  Get memory profile data.

  @param[in]      This              The EDKII_SMM_MEMORY_PROFILE_PROTOCOL instance.
  @param[in, out] ProfileSize       On entry, points to the size in bytes of the ProfileBuffer.
                                    On return, points to the size of the data returned in ProfileBuffer.
  @param[out]     ProfileBuffer     Profile buffer.

  @return EFI_SUCCESS               Get the memory profile data successfully.
  @return EFI_UNSUPPORTED           Memory profile is unsupported.
  @return EFI_BUFFER_TO_SMALL       The ProfileSize is too small for the resulting data.
                                    ProfileSize is updated with the size required.

**/
EFI_STATUS
EFIAPI
SmramProfileProtocolGetData (
  IN     EDKII_SMM_MEMORY_PROFILE_PROTOCOL  *This,
  IN OUT UINT64                             *ProfileSize,
  OUT VOID                                  *ProfileBuffer
  )
{
  UINT64                       Size;
  UINT64                       Offset;
  MEMORY_PROFILE_CONTEXT_DATA  *ContextData;
  BOOLEAN                      SmramProfileGettingStatus;

  ContextData = GetSmramProfileContext ();
  if (ContextData == NULL) {
    return EFI_UNSUPPORTED;
  }

  SmramProfileGettingStatus  = mSmramProfileGettingStatus;
  mSmramProfileGettingStatus = TRUE;

  Size = SmramProfileGetDataSize ();

  if (*ProfileSize < Size) {
    *ProfileSize               = Size;
    mSmramProfileGettingStatus = SmramProfileGettingStatus;
    return EFI_BUFFER_TOO_SMALL;
  }

  Offset = 0;
  SmramProfileCopyData (ProfileBuffer, &Size, &Offset);
  *ProfileSize = Size;

  mSmramProfileGettingStatus = SmramProfileGettingStatus;
  return EFI_SUCCESS;
}

/**
  Register image to memory profile.

  @param[in] This               The EDKII_SMM_MEMORY_PROFILE_PROTOCOL instance.
  @param[in] FilePath           File path of the image.
  @param[in] ImageBase          Image base address.
  @param[in] ImageSize          Image size.
  @param[in] FileType           File type of the image.

  @return EFI_SUCCESS           Register successfully.
  @return EFI_UNSUPPORTED       Memory profile is unsupported,
                                or memory profile for the image is not required.
  @return EFI_OUT_OF_RESOURCES  No enough resource for this register.

**/
EFI_STATUS
EFIAPI
SmramProfileProtocolRegisterImage (
  IN EDKII_SMM_MEMORY_PROFILE_PROTOCOL  *This,
  IN EFI_DEVICE_PATH_PROTOCOL           *FilePath,
  IN PHYSICAL_ADDRESS                   ImageBase,
  IN UINT64                             ImageSize,
  IN EFI_FV_FILETYPE                    FileType
  )
{
  EFI_STATUS            Status;
  EFI_SMM_DRIVER_ENTRY  DriverEntry;
  VOID                  *EntryPointInImage;
  EFI_GUID              *Name;

  ZeroMem (&DriverEntry, sizeof (DriverEntry));
  Name = GetFileNameFromFilePath (FilePath);
  if (Name != NULL) {
    CopyMem (&DriverEntry.FileName, Name, sizeof (EFI_GUID));
  }

  DriverEntry.ImageBuffer  = ImageBase;
  DriverEntry.NumberOfPage = EFI_SIZE_TO_PAGES ((UINTN)ImageSize);
  Status                   = InternalPeCoffGetEntryPoint ((VOID *)(UINTN)DriverEntry.ImageBuffer, &EntryPointInImage);
  ASSERT_EFI_ERROR (Status);
  DriverEntry.ImageEntryPoint = (PHYSICAL_ADDRESS)(UINTN)EntryPointInImage;

  return RegisterSmramProfileImage (&DriverEntry, FALSE);
}

/**
  Unregister image from memory profile.

  @param[in] This               The EDKII_SMM_MEMORY_PROFILE_PROTOCOL instance.
  @param[in] FilePath           File path of the image.
  @param[in] ImageBase          Image base address.
  @param[in] ImageSize          Image size.

  @return EFI_SUCCESS           Unregister successfully.
  @return EFI_UNSUPPORTED       Memory profile is unsupported,
                                or memory profile for the image is not required.
  @return EFI_NOT_FOUND         The image is not found.

**/
EFI_STATUS
EFIAPI
SmramProfileProtocolUnregisterImage (
  IN EDKII_SMM_MEMORY_PROFILE_PROTOCOL  *This,
  IN EFI_DEVICE_PATH_PROTOCOL           *FilePath,
  IN PHYSICAL_ADDRESS                   ImageBase,
  IN UINT64                             ImageSize
  )
{
  EFI_STATUS            Status;
  EFI_SMM_DRIVER_ENTRY  DriverEntry;
  VOID                  *EntryPointInImage;
  EFI_GUID              *Name;

  ZeroMem (&DriverEntry, sizeof (DriverEntry));
  Name = GetFileNameFromFilePath (FilePath);
  if (Name != NULL) {
    CopyMem (&DriverEntry.FileName, Name, sizeof (EFI_GUID));
  }

  DriverEntry.ImageBuffer  = ImageBase;
  DriverEntry.NumberOfPage = EFI_SIZE_TO_PAGES ((UINTN)ImageSize);
  Status                   = InternalPeCoffGetEntryPoint ((VOID *)(UINTN)DriverEntry.ImageBuffer, &EntryPointInImage);
  ASSERT_EFI_ERROR (Status);
  DriverEntry.ImageEntryPoint = (PHYSICAL_ADDRESS)(UINTN)EntryPointInImage;

  return UnregisterSmramProfileImage (&DriverEntry, FALSE);
}

/**
  Get memory profile recording state.

  @param[in]  This              The EDKII_SMM_MEMORY_PROFILE_PROTOCOL instance.
  @param[out] RecordingState    Recording state.

  @return EFI_SUCCESS           Memory profile recording state is returned.
  @return EFI_UNSUPPORTED       Memory profile is unsupported.
  @return EFI_INVALID_PARAMETER RecordingState is NULL.

**/
EFI_STATUS
EFIAPI
SmramProfileProtocolGetRecordingState (
  IN EDKII_SMM_MEMORY_PROFILE_PROTOCOL  *This,
  OUT BOOLEAN                           *RecordingState
  )
{
  MEMORY_PROFILE_CONTEXT_DATA  *ContextData;

  ContextData = GetSmramProfileContext ();
  if (ContextData == NULL) {
    return EFI_UNSUPPORTED;
  }

  if (RecordingState == NULL) {
    return EFI_INVALID_PARAMETER;
  }

  *RecordingState = mSmramProfileRecordingEnable;
  return EFI_SUCCESS;
}

/**
  Set memory profile recording state.

  @param[in] This               The EDKII_SMM_MEMORY_PROFILE_PROTOCOL instance.
  @param[in] RecordingState     Recording state.

  @return EFI_SUCCESS           Set memory profile recording state successfully.
  @return EFI_UNSUPPORTED       Memory profile is unsupported.

**/
EFI_STATUS
EFIAPI
SmramProfileProtocolSetRecordingState (
  IN EDKII_SMM_MEMORY_PROFILE_PROTOCOL  *This,
  IN BOOLEAN                            RecordingState
  )
{
  MEMORY_PROFILE_CONTEXT_DATA  *ContextData;

  ContextData = GetSmramProfileContext ();
  if (ContextData == NULL) {
    return EFI_UNSUPPORTED;
  }

  mSmramProfileRecordingEnable = RecordingState;
  return EFI_SUCCESS;
}

/**
  Record memory profile of multilevel caller.

  @param[in] This               The EDKII_SMM_MEMORY_PROFILE_PROTOCOL instance.
  @param[in] CallerAddress      Address of caller.
  @param[in] Action             Memory profile action.
  @param[in] MemoryType         Memory type.
                                EfiMaxMemoryType means the MemoryType is unknown.
  @param[in] Buffer             Buffer address.
  @param[in] Size               Buffer size.
  @param[in] ActionString       String for memory profile action.
                                Only needed for user defined allocate action.

  @return EFI_SUCCESS           Memory profile is updated.
  @return EFI_UNSUPPORTED       Memory profile is unsupported,
                                or memory profile for the image is not required,
                                or memory profile for the memory type is not required.
  @return EFI_ACCESS_DENIED     It is during memory profile data getting.
  @return EFI_ABORTED           Memory profile recording is not enabled.
  @return EFI_OUT_OF_RESOURCES  No enough resource to update memory profile for allocate action.
  @return EFI_NOT_FOUND         No matched allocate info found for free action.

**/
EFI_STATUS
EFIAPI
SmramProfileProtocolRecord (
  IN EDKII_SMM_MEMORY_PROFILE_PROTOCOL  *This,
  IN PHYSICAL_ADDRESS                   CallerAddress,
  IN MEMORY_PROFILE_ACTION              Action,
  IN EFI_MEMORY_TYPE                    MemoryType,
  IN VOID                               *Buffer,
  IN UINTN                              Size,
  IN CHAR8                              *ActionString OPTIONAL
  )
{
  return SmmCoreUpdateProfile (CallerAddress, Action, MemoryType, Size, Buffer, ActionString);
}

/**
  SMRAM profile handler to get profile info.

  @param SmramProfileParameterGetInfo The parameter of SMM profile get size.

**/
VOID
SmramProfileHandlerGetInfo (
  IN SMRAM_PROFILE_PARAMETER_GET_PROFILE_INFO  *SmramProfileParameterGetInfo
  )
{
  MEMORY_PROFILE_CONTEXT_DATA  *ContextData;
  BOOLEAN                      SmramProfileGettingStatus;

  ContextData = GetSmramProfileContext ();
  if (ContextData == NULL) {
    return;
  }

  SmramProfileGettingStatus  = mSmramProfileGettingStatus;
  mSmramProfileGettingStatus = TRUE;

  SmramProfileParameterGetInfo->ProfileSize         = SmramProfileGetDataSize ();
  SmramProfileParameterGetInfo->Header.ReturnStatus = 0;

  mSmramProfileGettingStatus = SmramProfileGettingStatus;
}

/**
  SMRAM profile handler to get profile data.

  @param SmramProfileParameterGetData The parameter of SMM profile get data.

**/
VOID
SmramProfileHandlerGetData (
  IN SMRAM_PROFILE_PARAMETER_GET_PROFILE_DATA  *SmramProfileParameterGetData
  )
{
  UINT64                                    ProfileSize;
  UINT64                                    ProfileOffset;
  SMRAM_PROFILE_PARAMETER_GET_PROFILE_DATA  SmramProfileGetData;
  MEMORY_PROFILE_CONTEXT_DATA               *ContextData;
  BOOLEAN                                   SmramProfileGettingStatus;

  ContextData = GetSmramProfileContext ();
  if (ContextData == NULL) {
    return;
  }

  SmramProfileGettingStatus  = mSmramProfileGettingStatus;
  mSmramProfileGettingStatus = TRUE;

  CopyMem (&SmramProfileGetData, SmramProfileParameterGetData, sizeof (SmramProfileGetData));

  ProfileSize = SmramProfileGetDataSize ();

  //
  // Sanity check
  //
  if (!SmmIsBufferOutsideSmmValid ((UINTN)SmramProfileGetData.ProfileBuffer, (UINTN)ProfileSize)) {
    DEBUG ((DEBUG_ERROR, "SmramProfileHandlerGetData: SMM ProfileBuffer in SMRAM or overflow!\n"));
    SmramProfileParameterGetData->ProfileSize         = ProfileSize;
    SmramProfileParameterGetData->Header.ReturnStatus = (UINT64)(INT64)(INTN)EFI_ACCESS_DENIED;
    goto Done;
  }

  if (SmramProfileGetData.ProfileSize < ProfileSize) {
    SmramProfileParameterGetData->ProfileSize         = ProfileSize;
    SmramProfileParameterGetData->Header.ReturnStatus = (UINT64)(INT64)(INTN)EFI_BUFFER_TOO_SMALL;
    goto Done;
  }

  ProfileOffset = 0;
  SmramProfileCopyData ((VOID *)(UINTN)SmramProfileGetData.ProfileBuffer, &ProfileSize, &ProfileOffset);
  SmramProfileParameterGetData->ProfileSize         = ProfileSize;
  SmramProfileParameterGetData->Header.ReturnStatus = 0;

Done:
  mSmramProfileGettingStatus = SmramProfileGettingStatus;
}

/**
  SMRAM profile handler to get profile data by offset.

  @param SmramProfileParameterGetDataByOffset   The parameter of SMM profile get data by offset.

**/
VOID
SmramProfileHandlerGetDataByOffset (
  IN SMRAM_PROFILE_PARAMETER_GET_PROFILE_DATA_BY_OFFSET  *SmramProfileParameterGetDataByOffset
  )
{
  SMRAM_PROFILE_PARAMETER_GET_PROFILE_DATA_BY_OFFSET  SmramProfileGetDataByOffset;
  MEMORY_PROFILE_CONTEXT_DATA                         *ContextData;
  BOOLEAN                                             SmramProfileGettingStatus;

  ContextData = GetSmramProfileContext ();
  if (ContextData == NULL) {
    return;
  }

  SmramProfileGettingStatus  = mSmramProfileGettingStatus;
  mSmramProfileGettingStatus = TRUE;

  CopyMem (&SmramProfileGetDataByOffset, SmramProfileParameterGetDataByOffset, sizeof (SmramProfileGetDataByOffset));

  //
  // Sanity check
  //
  if (!SmmIsBufferOutsideSmmValid ((UINTN)SmramProfileGetDataByOffset.ProfileBuffer, (UINTN)SmramProfileGetDataByOffset.ProfileSize)) {
    DEBUG ((DEBUG_ERROR, "SmramProfileHandlerGetDataByOffset: SMM ProfileBuffer in SMRAM or overflow!\n"));
    SmramProfileParameterGetDataByOffset->Header.ReturnStatus = (UINT64)(INT64)(INTN)EFI_ACCESS_DENIED;
    goto Done;
  }

  SmramProfileCopyData ((VOID *)(UINTN)SmramProfileGetDataByOffset.ProfileBuffer, &SmramProfileGetDataByOffset.ProfileSize, &SmramProfileGetDataByOffset.ProfileOffset);
  CopyMem (SmramProfileParameterGetDataByOffset, &SmramProfileGetDataByOffset, sizeof (SmramProfileGetDataByOffset));
  SmramProfileParameterGetDataByOffset->Header.ReturnStatus = 0;

Done:
  mSmramProfileGettingStatus = SmramProfileGettingStatus;
}

/**
  Dispatch function for a Software SMI handler.

  Caution: This function may receive untrusted input.
  Communicate buffer and buffer size are external input, so this function will do basic validation.

  @param DispatchHandle  The unique handle assigned to this handler by SmiHandlerRegister().
  @param Context         Points to an optional handler context which was specified when the
                         handler was registered.
  @param CommBuffer      A pointer to a collection of data in memory that will
                         be conveyed from a non-SMM environment into an SMM environment.
  @param CommBufferSize  The size of the CommBuffer.

  @retval EFI_SUCCESS Command is handled successfully.

**/
EFI_STATUS
EFIAPI
SmramProfileHandler (
  IN EFI_HANDLE  DispatchHandle,
  IN CONST VOID  *Context         OPTIONAL,
  IN OUT VOID    *CommBuffer      OPTIONAL,
  IN OUT UINTN   *CommBufferSize  OPTIONAL
  )
{
  SMRAM_PROFILE_PARAMETER_HEADER           *SmramProfileParameterHeader;
  UINTN                                    TempCommBufferSize;
  SMRAM_PROFILE_PARAMETER_RECORDING_STATE  *ParameterRecordingState;

  DEBUG ((DEBUG_ERROR, "SmramProfileHandler Enter\n"));

  //
  // If input is invalid, stop processing this SMI
  //
  if ((CommBuffer == NULL) || (CommBufferSize == NULL)) {
    return EFI_SUCCESS;
  }

  TempCommBufferSize = *CommBufferSize;

  if (TempCommBufferSize < sizeof (SMRAM_PROFILE_PARAMETER_HEADER)) {
    DEBUG ((DEBUG_ERROR, "SmramProfileHandler: SMM communication buffer size invalid!\n"));
    return EFI_SUCCESS;
  }

  if (mSmramReadyToLock && !SmmIsBufferOutsideSmmValid ((UINTN)CommBuffer, TempCommBufferSize)) {
    DEBUG ((DEBUG_ERROR, "SmramProfileHandler: SMM communication buffer in SMRAM or overflow!\n"));
    return EFI_SUCCESS;
  }

  SmramProfileParameterHeader = (SMRAM_PROFILE_PARAMETER_HEADER *)((UINTN)CommBuffer);

  SmramProfileParameterHeader->ReturnStatus = (UINT64)-1;

  if (GetSmramProfileContext () == NULL) {
    SmramProfileParameterHeader->ReturnStatus = (UINT64)(INT64)(INTN)EFI_UNSUPPORTED;
    return EFI_SUCCESS;
  }

  switch (SmramProfileParameterHeader->Command) {
    case SMRAM_PROFILE_COMMAND_GET_PROFILE_INFO:
      DEBUG ((DEBUG_ERROR, "SmramProfileHandlerGetInfo\n"));
      if (TempCommBufferSize != sizeof (SMRAM_PROFILE_PARAMETER_GET_PROFILE_INFO)) {
        DEBUG ((DEBUG_ERROR, "SmramProfileHandler: SMM communication buffer size invalid!\n"));
        return EFI_SUCCESS;
      }

      SmramProfileHandlerGetInfo ((SMRAM_PROFILE_PARAMETER_GET_PROFILE_INFO *)(UINTN)CommBuffer);
      break;
    case SMRAM_PROFILE_COMMAND_GET_PROFILE_DATA:
      DEBUG ((DEBUG_ERROR, "SmramProfileHandlerGetData\n"));
      if (TempCommBufferSize != sizeof (SMRAM_PROFILE_PARAMETER_GET_PROFILE_DATA)) {
        DEBUG ((DEBUG_ERROR, "SmramProfileHandler: SMM communication buffer size invalid!\n"));
        return EFI_SUCCESS;
      }

      SmramProfileHandlerGetData ((SMRAM_PROFILE_PARAMETER_GET_PROFILE_DATA *)(UINTN)CommBuffer);
      break;
    case SMRAM_PROFILE_COMMAND_GET_PROFILE_DATA_BY_OFFSET:
      DEBUG ((DEBUG_ERROR, "SmramProfileHandlerGetDataByOffset\n"));
      if (TempCommBufferSize != sizeof (SMRAM_PROFILE_PARAMETER_GET_PROFILE_DATA_BY_OFFSET)) {
        DEBUG ((DEBUG_ERROR, "SmramProfileHandler: SMM communication buffer size invalid!\n"));
        return EFI_SUCCESS;
      }

      SmramProfileHandlerGetDataByOffset ((SMRAM_PROFILE_PARAMETER_GET_PROFILE_DATA_BY_OFFSET *)(UINTN)CommBuffer);
      break;
    case SMRAM_PROFILE_COMMAND_GET_RECORDING_STATE:
      DEBUG ((DEBUG_ERROR, "SmramProfileHandlerGetRecordingState\n"));
      if (TempCommBufferSize != sizeof (SMRAM_PROFILE_PARAMETER_RECORDING_STATE)) {
        DEBUG ((DEBUG_ERROR, "SmramProfileHandler: SMM communication buffer size invalid!\n"));
        return EFI_SUCCESS;
      }

      ParameterRecordingState                      = (SMRAM_PROFILE_PARAMETER_RECORDING_STATE *)(UINTN)CommBuffer;
      ParameterRecordingState->RecordingState      = mSmramProfileRecordingEnable;
      ParameterRecordingState->Header.ReturnStatus = 0;
      break;
    case SMRAM_PROFILE_COMMAND_SET_RECORDING_STATE:
      DEBUG ((DEBUG_ERROR, "SmramProfileHandlerSetRecordingState\n"));
      if (TempCommBufferSize != sizeof (SMRAM_PROFILE_PARAMETER_RECORDING_STATE)) {
        DEBUG ((DEBUG_ERROR, "SmramProfileHandler: SMM communication buffer size invalid!\n"));
        return EFI_SUCCESS;
      }

      ParameterRecordingState                      = (SMRAM_PROFILE_PARAMETER_RECORDING_STATE *)(UINTN)CommBuffer;
      mSmramProfileRecordingEnable                 = ParameterRecordingState->RecordingState;
      ParameterRecordingState->Header.ReturnStatus = 0;
      break;

    //
    // Below 2 commands have been deprecated. They may not be (re-)used.
    //
    case SMRAM_PROFILE_COMMAND_DEPRECATED1:
    case SMRAM_PROFILE_COMMAND_DEPRECATED2:
      ASSERT (FALSE);
    //
    // Fall-through to the default (unrecognized command) case.
    //
    default:
      break;
  }

  DEBUG ((DEBUG_ERROR, "SmramProfileHandler Exit\n"));

  return EFI_SUCCESS;
}

/**
  Register SMRAM profile handler.

**/
VOID
RegisterSmramProfileHandler (
  VOID
  )
{
  EFI_STATUS  Status;
  EFI_HANDLE  DispatchHandle;

  if (!IS_SMRAM_PROFILE_ENABLED) {
    return;
  }

  Status = SmiHandlerRegister (
             SmramProfileHandler,
             &gEdkiiMemoryProfileGuid,
             &DispatchHandle
             );
  ASSERT_EFI_ERROR (Status);
}

////////////////////

/**
  Dump SMRAM range.

**/
VOID
DumpSmramRange (
  VOID
  )
{
  UINTN                        Index;
  MEMORY_PROFILE_CONTEXT_DATA  *ContextData;
  BOOLEAN                      SmramProfileGettingStatus;

  ContextData = GetSmramProfileContext ();
  if (ContextData == NULL) {
    return;
  }

  SmramProfileGettingStatus  = mSmramProfileGettingStatus;
  mSmramProfileGettingStatus = TRUE;

  DEBUG ((DEBUG_INFO, "FullSmramRange address - 0x%08x\n", mFullSmramRanges));

  DEBUG ((DEBUG_INFO, "======= SmramProfile begin =======\n"));

  DEBUG ((DEBUG_INFO, "FullSmramRange:\n"));
  for (Index = 0; Index < mFullSmramRangeCount; Index++) {
    DEBUG ((DEBUG_INFO, "  FullSmramRange (0x%x)\n", Index));
    DEBUG ((DEBUG_INFO, "    PhysicalStart      - 0x%016lx\n", mFullSmramRanges[Index].PhysicalStart));
    DEBUG ((DEBUG_INFO, "    CpuStart           - 0x%016lx\n", mFullSmramRanges[Index].CpuStart));
    DEBUG ((DEBUG_INFO, "    PhysicalSize       - 0x%016lx\n", mFullSmramRanges[Index].PhysicalSize));
    DEBUG ((DEBUG_INFO, "    RegionState        - 0x%016lx\n", mFullSmramRanges[Index].RegionState));
  }

  DEBUG ((DEBUG_INFO, "======= SmramProfile end =======\n"));

  mSmramProfileGettingStatus = SmramProfileGettingStatus;
}

/**
  Dump SMRAM free page list.

**/
VOID
DumpFreePagesList (
  VOID
  )
{
  LIST_ENTRY                   *FreePageList;
  LIST_ENTRY                   *Node;
  FREE_PAGE_LIST               *Pages;
  UINTN                        Index;
  MEMORY_PROFILE_CONTEXT_DATA  *ContextData;
  BOOLEAN                      SmramProfileGettingStatus;

  ContextData = GetSmramProfileContext ();
  if (ContextData == NULL) {
    return;
  }

  SmramProfileGettingStatus  = mSmramProfileGettingStatus;
  mSmramProfileGettingStatus = TRUE;

  DEBUG ((DEBUG_INFO, "======= SmramProfile begin =======\n"));

  DEBUG ((DEBUG_INFO, "FreePagesList:\n"));
  FreePageList = &mSmmMemoryMap;
  for (Node = FreePageList->BackLink, Index = 0;
       Node != FreePageList;
       Node = Node->BackLink, Index++)
  {
    Pages = BASE_CR (Node, FREE_PAGE_LIST, Link);
    DEBUG ((DEBUG_INFO, "  Index - 0x%x\n", Index));
    DEBUG ((DEBUG_INFO, "    PhysicalStart - 0x%016lx\n", (PHYSICAL_ADDRESS)(UINTN)Pages));
    DEBUG ((DEBUG_INFO, "    NumberOfPages - 0x%08x\n", Pages->NumberOfPages));
  }

  DEBUG ((DEBUG_INFO, "======= SmramProfile end =======\n"));

  mSmramProfileGettingStatus = SmramProfileGettingStatus;
}

/**
  Dump SMRAM free pool list.

**/
VOID
DumpFreePoolList (
  VOID
  )
{
  LIST_ENTRY                   *FreePoolList;
  LIST_ENTRY                   *Node;
  FREE_POOL_HEADER             *Pool;
  UINTN                        Index;
  UINTN                        PoolListIndex;
  MEMORY_PROFILE_CONTEXT_DATA  *ContextData;
  BOOLEAN                      SmramProfileGettingStatus;
  UINTN                        SmmPoolTypeIndex;

  ContextData = GetSmramProfileContext ();
  if (ContextData == NULL) {
    return;
  }

  SmramProfileGettingStatus  = mSmramProfileGettingStatus;
  mSmramProfileGettingStatus = TRUE;

  DEBUG ((DEBUG_INFO, "======= SmramProfile begin =======\n"));

  for (SmmPoolTypeIndex = 0; SmmPoolTypeIndex < SmmPoolTypeMax; SmmPoolTypeIndex++) {
    for (PoolListIndex = 0; PoolListIndex < MAX_POOL_INDEX; PoolListIndex++) {
      DEBUG ((DEBUG_INFO, "FreePoolList(%d)(%d):\n", SmmPoolTypeIndex, PoolListIndex));
      FreePoolList = &mSmmPoolLists[SmmPoolTypeIndex][PoolListIndex];
      for (Node = FreePoolList->BackLink, Index = 0;
           Node != FreePoolList;
           Node = Node->BackLink, Index++)
      {
        Pool = BASE_CR (Node, FREE_POOL_HEADER, Link);
        DEBUG ((DEBUG_INFO, "  Index - 0x%x\n", Index));
        DEBUG ((DEBUG_INFO, "    PhysicalStart - 0x%016lx\n", (PHYSICAL_ADDRESS)(UINTN)Pool));
        DEBUG ((DEBUG_INFO, "    Size          - 0x%08x\n", Pool->Header.Size));
        DEBUG ((DEBUG_INFO, "    Available     - 0x%02x\n", Pool->Header.Available));
      }
    }
  }

  DEBUG ((DEBUG_INFO, "======= SmramProfile end =======\n"));

  mSmramProfileGettingStatus = SmramProfileGettingStatus;
}

GLOBAL_REMOVE_IF_UNREFERENCED CHAR8  *mSmmActionString[] = {
  "SmmUnknown",
  "gSmst->SmmAllocatePages",
  "gSmst->SmmFreePages",
  "gSmst->SmmAllocatePool",
  "gSmst->SmmFreePool",
};

typedef struct {
  MEMORY_PROFILE_ACTION    Action;
  CHAR8                    *String;
} ACTION_STRING;

GLOBAL_REMOVE_IF_UNREFERENCED ACTION_STRING  mExtActionString[] = {
  { MEMORY_PROFILE_ACTION_LIB_ALLOCATE_PAGES,                  "Lib:AllocatePages"                },
  { MEMORY_PROFILE_ACTION_LIB_ALLOCATE_RUNTIME_PAGES,          "Lib:AllocateRuntimePages"         },
  { MEMORY_PROFILE_ACTION_LIB_ALLOCATE_RESERVED_PAGES,         "Lib:AllocateReservedPages"        },
  { MEMORY_PROFILE_ACTION_LIB_FREE_PAGES,                      "Lib:FreePages"                    },
  { MEMORY_PROFILE_ACTION_LIB_ALLOCATE_ALIGNED_PAGES,          "Lib:AllocateAlignedPages"         },
  { MEMORY_PROFILE_ACTION_LIB_ALLOCATE_ALIGNED_RUNTIME_PAGES,  "Lib:AllocateAlignedRuntimePages"  },
  { MEMORY_PROFILE_ACTION_LIB_ALLOCATE_ALIGNED_RESERVED_PAGES, "Lib:AllocateAlignedReservedPages" },
  { MEMORY_PROFILE_ACTION_LIB_FREE_ALIGNED_PAGES,              "Lib:FreeAlignedPages"             },
  { MEMORY_PROFILE_ACTION_LIB_ALLOCATE_POOL,                   "Lib:AllocatePool"                 },
  { MEMORY_PROFILE_ACTION_LIB_ALLOCATE_RUNTIME_POOL,           "Lib:AllocateRuntimePool"          },
  { MEMORY_PROFILE_ACTION_LIB_ALLOCATE_RESERVED_POOL,          "Lib:AllocateReservedPool"         },
  { MEMORY_PROFILE_ACTION_LIB_FREE_POOL,                       "Lib:FreePool"                     },
  { MEMORY_PROFILE_ACTION_LIB_ALLOCATE_ZERO_POOL,              "Lib:AllocateZeroPool"             },
  { MEMORY_PROFILE_ACTION_LIB_ALLOCATE_RUNTIME_ZERO_POOL,      "Lib:AllocateRuntimeZeroPool"      },
  { MEMORY_PROFILE_ACTION_LIB_ALLOCATE_RESERVED_ZERO_POOL,     "Lib:AllocateReservedZeroPool"     },
  { MEMORY_PROFILE_ACTION_LIB_ALLOCATE_COPY_POOL,              "Lib:AllocateCopyPool"             },
  { MEMORY_PROFILE_ACTION_LIB_ALLOCATE_RUNTIME_COPY_POOL,      "Lib:AllocateRuntimeCopyPool"      },
  { MEMORY_PROFILE_ACTION_LIB_ALLOCATE_RESERVED_COPY_POOL,     "Lib:AllocateReservedCopyPool"     },
  { MEMORY_PROFILE_ACTION_LIB_REALLOCATE_POOL,                 "Lib:ReallocatePool"               },
  { MEMORY_PROFILE_ACTION_LIB_REALLOCATE_RUNTIME_POOL,         "Lib:ReallocateRuntimePool"        },
  { MEMORY_PROFILE_ACTION_LIB_REALLOCATE_RESERVED_POOL,        "Lib:ReallocateReservedPool"       },
};

typedef struct {
  EFI_MEMORY_TYPE    MemoryType;
  CHAR8              *MemoryTypeStr;
} PROFILE_MEMORY_TYPE_STRING;

GLOBAL_REMOVE_IF_UNREFERENCED PROFILE_MEMORY_TYPE_STRING  mMemoryTypeString[] = {
  { EfiRuntimeServicesCode, "EfiRuntimeServicesCode" },
  { EfiRuntimeServicesData, "EfiRuntimeServicesData" }
};

/**
  Memory type to string.

  @param[in] MemoryType Memory type.

  @return Pointer to string.

**/
CHAR8 *
ProfileMemoryTypeToStr (
  IN EFI_MEMORY_TYPE  MemoryType
  )
{
  UINTN  Index;

  for (Index = 0; Index < ARRAY_SIZE (mMemoryTypeString); Index++) {
    if (mMemoryTypeString[Index].MemoryType == MemoryType) {
      return mMemoryTypeString[Index].MemoryTypeStr;
    }
  }

  return "UnexpectedMemoryType";
}

/**
  Action to string.

  @param[in] Action                     Profile action.

  @return Pointer to string.

**/
CHAR8 *
ProfileActionToStr (
  IN MEMORY_PROFILE_ACTION  Action
  )
{
  UINTN  Index;
  UINTN  ActionStringCount;
  CHAR8  **ActionString;

  ActionString      = mSmmActionString;
  ActionStringCount = ARRAY_SIZE (mSmmActionString);

  if ((UINTN)(UINT32)Action < ActionStringCount) {
    return ActionString[Action];
  }

  for (Index = 0; Index < ARRAY_SIZE (mExtActionString); Index++) {
    if (mExtActionString[Index].Action == Action) {
      return mExtActionString[Index].String;
    }
  }

  return ActionString[0];
}

/**
  Dump SMRAM profile.

**/
VOID
DumpSmramProfile (
  VOID
  )
{
  MEMORY_PROFILE_CONTEXT           *Context;
  MEMORY_PROFILE_DRIVER_INFO       *DriverInfo;
  MEMORY_PROFILE_ALLOC_INFO        *AllocInfo;
  MEMORY_PROFILE_CONTEXT_DATA      *ContextData;
  MEMORY_PROFILE_DRIVER_INFO_DATA  *DriverInfoData;
  MEMORY_PROFILE_ALLOC_INFO_DATA   *AllocInfoData;
  LIST_ENTRY                       *SmramDriverInfoList;
  UINTN                            DriverIndex;
  LIST_ENTRY                       *DriverLink;
  LIST_ENTRY                       *AllocInfoList;
  UINTN                            AllocIndex;
  LIST_ENTRY                       *AllocLink;
  BOOLEAN                          SmramProfileGettingStatus;
  UINTN                            TypeIndex;

  ContextData = GetSmramProfileContext ();
  if (ContextData == NULL) {
    return;
  }

  SmramProfileGettingStatus  = mSmramProfileGettingStatus;
  mSmramProfileGettingStatus = TRUE;

  Context = &ContextData->Context;
  DEBUG ((DEBUG_INFO, "======= SmramProfile begin =======\n"));
  DEBUG ((DEBUG_INFO, "MEMORY_PROFILE_CONTEXT\n"));

  DEBUG ((DEBUG_INFO, "  CurrentTotalUsage     - 0x%016lx\n", Context->CurrentTotalUsage));
  DEBUG ((DEBUG_INFO, "  PeakTotalUsage        - 0x%016lx\n", Context->PeakTotalUsage));
  for (TypeIndex = 0; TypeIndex < sizeof (Context->CurrentTotalUsageByType) / sizeof (Context->CurrentTotalUsageByType[0]); TypeIndex++) {
    if ((Context->CurrentTotalUsageByType[TypeIndex] != 0) ||
        (Context->PeakTotalUsageByType[TypeIndex] != 0))
    {
      DEBUG ((DEBUG_INFO, "  CurrentTotalUsage[0x%02x]  - 0x%016lx (%a)\n", TypeIndex, Context->CurrentTotalUsageByType[TypeIndex], ProfileMemoryTypeToStr (TypeIndex)));
      DEBUG ((DEBUG_INFO, "  PeakTotalUsage[0x%02x]     - 0x%016lx (%a)\n", TypeIndex, Context->PeakTotalUsageByType[TypeIndex], ProfileMemoryTypeToStr (TypeIndex)));
    }
  }

  DEBUG ((DEBUG_INFO, "  TotalImageSize        - 0x%016lx\n", Context->TotalImageSize));
  DEBUG ((DEBUG_INFO, "  ImageCount            - 0x%08x\n", Context->ImageCount));
  DEBUG ((DEBUG_INFO, "  SequenceCount         - 0x%08x\n", Context->SequenceCount));

  SmramDriverInfoList = ContextData->DriverInfoList;
  for (DriverLink = SmramDriverInfoList->ForwardLink, DriverIndex = 0;
       DriverLink != SmramDriverInfoList;
       DriverLink = DriverLink->ForwardLink, DriverIndex++)
  {
    DriverInfoData = CR (
                       DriverLink,
                       MEMORY_PROFILE_DRIVER_INFO_DATA,
                       Link,
                       MEMORY_PROFILE_DRIVER_INFO_SIGNATURE
                       );
    DriverInfo = &DriverInfoData->DriverInfo;
    DEBUG ((DEBUG_INFO, "  MEMORY_PROFILE_DRIVER_INFO (0x%x)\n", DriverIndex));
    DEBUG ((DEBUG_INFO, "    FileName            - %g\n", &DriverInfo->FileName));
    DEBUG ((DEBUG_INFO, "    ImageBase           - 0x%016lx\n", DriverInfo->ImageBase));
    DEBUG ((DEBUG_INFO, "    ImageSize           - 0x%016lx\n", DriverInfo->ImageSize));
    DEBUG ((DEBUG_INFO, "    EntryPoint          - 0x%016lx\n", DriverInfo->EntryPoint));
    DEBUG ((DEBUG_INFO, "    ImageSubsystem      - 0x%04x\n", DriverInfo->ImageSubsystem));
    DEBUG ((DEBUG_INFO, "    FileType            - 0x%02x\n", DriverInfo->FileType));
    DEBUG ((DEBUG_INFO, "    CurrentUsage        - 0x%016lx\n", DriverInfo->CurrentUsage));
    DEBUG ((DEBUG_INFO, "    PeakUsage           - 0x%016lx\n", DriverInfo->PeakUsage));
    for (TypeIndex = 0; TypeIndex < sizeof (DriverInfo->CurrentUsageByType) / sizeof (DriverInfo->CurrentUsageByType[0]); TypeIndex++) {
      if ((DriverInfo->CurrentUsageByType[TypeIndex] != 0) ||
          (DriverInfo->PeakUsageByType[TypeIndex] != 0))
      {
        DEBUG ((DEBUG_INFO, "    CurrentUsage[0x%02x]     - 0x%016lx (%a)\n", TypeIndex, DriverInfo->CurrentUsageByType[TypeIndex], ProfileMemoryTypeToStr (TypeIndex)));
        DEBUG ((DEBUG_INFO, "    PeakUsage[0x%02x]        - 0x%016lx (%a)\n", TypeIndex, DriverInfo->PeakUsageByType[TypeIndex], ProfileMemoryTypeToStr (TypeIndex)));
      }
    }

    DEBUG ((DEBUG_INFO, "    AllocRecordCount    - 0x%08x\n", DriverInfo->AllocRecordCount));

    AllocInfoList = DriverInfoData->AllocInfoList;
    for (AllocLink = AllocInfoList->ForwardLink, AllocIndex = 0;
         AllocLink != AllocInfoList;
         AllocLink = AllocLink->ForwardLink, AllocIndex++)
    {
      AllocInfoData = CR (
                        AllocLink,
                        MEMORY_PROFILE_ALLOC_INFO_DATA,
                        Link,
                        MEMORY_PROFILE_ALLOC_INFO_SIGNATURE
                        );
      AllocInfo = &AllocInfoData->AllocInfo;
      DEBUG ((DEBUG_INFO, "    MEMORY_PROFILE_ALLOC_INFO (0x%x)\n", AllocIndex));
      DEBUG ((DEBUG_INFO, "      CallerAddress  - 0x%016lx (Offset: 0x%08x)\n", AllocInfo->CallerAddress, AllocInfo->CallerAddress - DriverInfo->ImageBase));
      DEBUG ((DEBUG_INFO, "      SequenceId     - 0x%08x\n", AllocInfo->SequenceId));
      if ((AllocInfo->Action & MEMORY_PROFILE_ACTION_USER_DEFINED_MASK) != 0) {
        if (AllocInfoData->ActionString != NULL) {
          DEBUG ((DEBUG_INFO, "      Action         - 0x%08x (%a)\n", AllocInfo->Action, AllocInfoData->ActionString));
        } else {
          DEBUG ((DEBUG_INFO, "      Action         - 0x%08x (UserDefined-0x%08x)\n", AllocInfo->Action, AllocInfo->Action));
        }
      } else {
        DEBUG ((DEBUG_INFO, "      Action         - 0x%08x (%a)\n", AllocInfo->Action, ProfileActionToStr (AllocInfo->Action)));
      }

      DEBUG ((DEBUG_INFO, "      MemoryType     - 0x%08x (%a)\n", AllocInfo->MemoryType, ProfileMemoryTypeToStr (AllocInfo->MemoryType)));
      DEBUG ((DEBUG_INFO, "      Buffer         - 0x%016lx\n", AllocInfo->Buffer));
      DEBUG ((DEBUG_INFO, "      Size           - 0x%016lx\n", AllocInfo->Size));
    }
  }

  DEBUG ((DEBUG_INFO, "======= SmramProfile end =======\n"));

  mSmramProfileGettingStatus = SmramProfileGettingStatus;
}

/**
  Dump SMRAM information.

**/
VOID
DumpSmramInfo (
  VOID
  )
{
  DEBUG_CODE_BEGIN ();
  if (IS_SMRAM_PROFILE_ENABLED) {
    DumpSmramProfile ();
    DumpFreePagesList ();
    DumpFreePoolList ();
    DumpSmramRange ();
  }

  DEBUG_CODE_END ();
}