Core/PiSmmCore/SmiHandlerProfile.c

/** @file
  SMI handler profile support.

Copyright (c) 2017, Intel Corporation. All rights reserved.<BR>
SPDX-License-Identifier: BSD-2-Clause-Patent

**/

#include <PiSmm.h>
#include <Library/BaseLib.h>
#include <Library/BaseMemoryLib.h>
#include <Library/MemoryAllocationLib.h>
#include <Library/UefiBootServicesTableLib.h>
#include <Library/SmmServicesTableLib.h>
#include <Library/DebugLib.h>
#include <Library/PrintLib.h>
#include <Library/UefiLib.h>
#include <Library/DevicePathLib.h>
#include <Library/PeCoffGetEntryPointLib.h>
#include <Protocol/LoadedImage.h>
#include <Protocol/SmmAccess2.h>
#include <Protocol/SmmReadyToLock.h>
#include <Protocol/SmmEndOfDxe.h>

#include <Guid/SmiHandlerProfile.h>

#include "PiSmmCore.h"

#define GET_OCCUPIED_SIZE(ActualSize, Alignment) \
  ((ActualSize) + (((Alignment) - ((ActualSize) & ((Alignment) - 1))) & ((Alignment) - 1)))

typedef struct {
  EFI_GUID            FileGuid;
  PHYSICAL_ADDRESS    EntryPoint;
  PHYSICAL_ADDRESS    ImageBase;
  UINT64              ImageSize;
  UINT32              ImageRef;
  UINT16              PdbStringSize;
  CHAR8               *PdbString;
} IMAGE_STRUCT;

/**
  Register SMI handler profile handler.
**/
VOID
RegisterSmiHandlerProfileHandler(
  VOID
  );

/**
  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
  );

extern LIST_ENTRY  mSmiEntryList;
extern LIST_ENTRY  mHardwareSmiEntryList;
extern SMI_ENTRY   mRootSmiEntry;

extern SMI_HANDLER_PROFILE_PROTOCOL  mSmiHandlerProfile;

GLOBAL_REMOVE_IF_UNREFERENCED LIST_ENTRY      mHardwareSmiEntryList = INITIALIZE_LIST_HEAD_VARIABLE (mHardwareSmiEntryList);

GLOBAL_REMOVE_IF_UNREFERENCED LIST_ENTRY      mRootSmiEntryList = INITIALIZE_LIST_HEAD_VARIABLE (mRootSmiEntryList);

GLOBAL_REMOVE_IF_UNREFERENCED LIST_ENTRY      *mSmmCoreRootSmiEntryList = &mRootSmiEntryList;
GLOBAL_REMOVE_IF_UNREFERENCED LIST_ENTRY      *mSmmCoreSmiEntryList = &mSmiEntryList;
GLOBAL_REMOVE_IF_UNREFERENCED LIST_ENTRY      *mSmmCoreHardwareSmiEntryList = &mHardwareSmiEntryList;

GLOBAL_REMOVE_IF_UNREFERENCED IMAGE_STRUCT  *mImageStruct;
GLOBAL_REMOVE_IF_UNREFERENCED UINT32        mImageStructCountMax;
GLOBAL_REMOVE_IF_UNREFERENCED UINT32        mImageStructCount;

GLOBAL_REMOVE_IF_UNREFERENCED VOID   *mSmiHandlerProfileDatabase;
GLOBAL_REMOVE_IF_UNREFERENCED UINTN  mSmiHandlerProfileDatabaseSize;

GLOBAL_REMOVE_IF_UNREFERENCED UINTN  mSmmImageDatabaseSize;
GLOBAL_REMOVE_IF_UNREFERENCED UINTN  mSmmRootSmiDatabaseSize;
GLOBAL_REMOVE_IF_UNREFERENCED UINTN  mSmmSmiDatabaseSize;
GLOBAL_REMOVE_IF_UNREFERENCED UINTN  mSmmHardwareSmiDatabaseSize;

GLOBAL_REMOVE_IF_UNREFERENCED BOOLEAN  mSmiHandlerProfileRecordingStatus;

GLOBAL_REMOVE_IF_UNREFERENCED SMI_HANDLER_PROFILE_PROTOCOL  mSmiHandlerProfile = {
  SmiHandlerProfileRegisterHandler,
  SmiHandlerProfileUnregisterHandler,
};

/**
  This function dump raw data.

  @param  Data  raw data
  @param  Size  raw data size
**/
VOID
InternalDumpData (
  IN UINT8  *Data,
  IN UINTN  Size
  )
{
  UINTN  Index;
  for (Index = 0; Index < Size; Index++) {
    DEBUG ((DEBUG_INFO, "%02x ", (UINTN)Data[Index]));
  }
}

/**
  Get GUID name for an image.

  @param[in]  LoadedImage LoadedImage protocol.
  @param[out] Guid        Guid of the FFS
**/
VOID
GetDriverGuid (
  IN  EFI_LOADED_IMAGE_PROTOCOL  *LoadedImage,
  OUT EFI_GUID                   *Guid
  )
{
  EFI_GUID                    *FileName;

  FileName = NULL;
  if ((DevicePathType(LoadedImage->FilePath) == MEDIA_DEVICE_PATH) &&
      (DevicePathSubType(LoadedImage->FilePath) == MEDIA_PIWG_FW_FILE_DP)) {
    FileName = &((MEDIA_FW_VOL_FILEPATH_DEVICE_PATH *)LoadedImage->FilePath)->FvFileName;
  }
  if (FileName != NULL) {
    CopyGuid(Guid, FileName);
  } else {
    ZeroMem(Guid, sizeof(EFI_GUID));
  }
}

/**
  Add image structure.

  @param  ImageBase         image base
  @param  ImageSize         image size
  @param  EntryPoint        image entry point
  @param  Guid              FFS GUID of the image
  @param  PdbString         image PDB string
**/
VOID
AddImageStruct(
  IN PHYSICAL_ADDRESS       ImageBase,
  IN UINT64                 ImageSize,
  IN PHYSICAL_ADDRESS       EntryPoint,
  IN EFI_GUID               *Guid,
  IN CHAR8                  *PdbString
  )
{
  UINTN  PdbStringSize;

  if (mImageStructCount >= mImageStructCountMax) {
    ASSERT(FALSE);
    return;
  }

  CopyGuid(&mImageStruct[mImageStructCount].FileGuid, Guid);
  mImageStruct[mImageStructCount].ImageRef = mImageStructCount;
  mImageStruct[mImageStructCount].ImageBase = ImageBase;
  mImageStruct[mImageStructCount].ImageSize = ImageSize;
  mImageStruct[mImageStructCount].EntryPoint = EntryPoint;
  if (PdbString != NULL) {
    PdbStringSize = AsciiStrSize(PdbString);
    mImageStruct[mImageStructCount].PdbString = AllocateCopyPool (PdbStringSize, PdbString);
    if (mImageStruct[mImageStructCount].PdbString != NULL) {
      mImageStruct[mImageStructCount].PdbStringSize = (UINT16) PdbStringSize;
    }
  }

  mImageStructCount++;
}

/**
  return an image structure based upon image address.

  @param  Address  image address

  @return image structure
**/
IMAGE_STRUCT *
AddressToImageStruct(
  IN UINTN  Address
  )
{
  UINTN  Index;

  for (Index = 0; Index < mImageStructCount; Index++) {
    if ((Address >= mImageStruct[Index].ImageBase) &&
        (Address < mImageStruct[Index].ImageBase + mImageStruct[Index].ImageSize)) {
      return &mImageStruct[Index];
    }
  }
  return NULL;
}

/**
  return an image reference index based upon image address.

  @param  Address  image address

  @return image reference index
**/
UINT32
AddressToImageRef(
  IN UINTN  Address
  )
{
  IMAGE_STRUCT *ImageStruct;

  ImageStruct = AddressToImageStruct(Address);
  if (ImageStruct != NULL) {
    return ImageStruct->ImageRef;
  }
  return (UINT32)-1;
}

/**
  Collect SMM image information based upon loaded image protocol.
**/
VOID
GetSmmLoadedImage(
  VOID
  )
{
  EFI_STATUS                 Status;
  UINTN                      NoHandles;
  UINTN                      HandleBufferSize;
  EFI_HANDLE                 *HandleBuffer;
  UINTN                      Index;
  EFI_LOADED_IMAGE_PROTOCOL  *LoadedImage;
  CHAR16                     *PathStr;
  EFI_SMM_DRIVER_ENTRY       *LoadedImagePrivate;
  PHYSICAL_ADDRESS           EntryPoint;
  VOID                       *EntryPointInImage;
  EFI_GUID                   Guid;
  CHAR8                      *PdbString;
  PHYSICAL_ADDRESS           RealImageBase;

  HandleBufferSize = 0;
  HandleBuffer = NULL;
  Status = gSmst->SmmLocateHandle(
                    ByProtocol,
                    &gEfiLoadedImageProtocolGuid,
                    NULL,
                    &HandleBufferSize,
                    HandleBuffer
                    );
  if (Status != EFI_BUFFER_TOO_SMALL) {
    return;
  }
  HandleBuffer = AllocateZeroPool (HandleBufferSize);
  if (HandleBuffer == NULL) {
    return;
  }
  Status = gSmst->SmmLocateHandle(
                    ByProtocol,
                    &gEfiLoadedImageProtocolGuid,
                    NULL,
                    &HandleBufferSize,
                    HandleBuffer
                    );
  if (EFI_ERROR(Status)) {
    return;
  }

  NoHandles = HandleBufferSize/sizeof(EFI_HANDLE);
  mImageStructCountMax = (UINT32) NoHandles;
  mImageStruct = AllocateZeroPool(mImageStructCountMax * sizeof(IMAGE_STRUCT));
  if (mImageStruct == NULL) {
    goto Done;
  }

  for (Index = 0; Index < NoHandles; Index++) {
    Status = gSmst->SmmHandleProtocol(
                      HandleBuffer[Index],
                      &gEfiLoadedImageProtocolGuid,
                      (VOID **)&LoadedImage
                      );
    if (EFI_ERROR(Status)) {
      continue;
    }
    PathStr = ConvertDevicePathToText(LoadedImage->FilePath, TRUE, TRUE);
    GetDriverGuid(LoadedImage, &Guid);
    DEBUG ((DEBUG_INFO, "Image: %g ", &Guid));

    EntryPoint = 0;
    LoadedImagePrivate = BASE_CR(LoadedImage, EFI_SMM_DRIVER_ENTRY, SmmLoadedImage);
    RealImageBase = (UINTN)LoadedImage->ImageBase;
    if (LoadedImagePrivate->Signature == EFI_SMM_DRIVER_ENTRY_SIGNATURE) {
      EntryPoint = LoadedImagePrivate->ImageEntryPoint;
      if ((EntryPoint != 0) && ((EntryPoint < (UINTN)LoadedImage->ImageBase) || (EntryPoint >= ((UINTN)LoadedImage->ImageBase + LoadedImage->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(LoadedImage->ImageBase, &EntryPointInImage);
        ASSERT_EFI_ERROR(Status);
        RealImageBase = (UINTN)LoadedImage->ImageBase + EntryPoint - (UINTN)EntryPointInImage;
      }
    }
    DEBUG ((DEBUG_INFO, "(0x%lx - 0x%lx", RealImageBase, LoadedImage->ImageSize));
    if (EntryPoint != 0) {
      DEBUG ((DEBUG_INFO, ", EntryPoint:0x%lx", EntryPoint));
    }
    DEBUG ((DEBUG_INFO, ")\n"));

    if (RealImageBase != 0) {
      PdbString = PeCoffLoaderGetPdbPointer ((VOID*) (UINTN) RealImageBase);
      DEBUG ((DEBUG_INFO, "       pdb - %a\n", PdbString));
    } else {
      PdbString = NULL;
    }
    DEBUG ((DEBUG_INFO, "       (%s)\n", PathStr));

    AddImageStruct(RealImageBase, LoadedImage->ImageSize, EntryPoint, &Guid, PdbString);
  }

Done:
  FreePool(HandleBuffer);
  return;
}

/**
  Dump SMI child context.

  @param HandlerType  the handler type
  @param Context      the handler context
  @param ContextSize  the handler context size
**/
VOID
DumpSmiChildContext (
  IN EFI_GUID   *HandlerType,
  IN VOID       *Context,
  IN UINTN      ContextSize
  )
{
  CHAR16        *Str;

  if (CompareGuid (HandlerType, &gEfiSmmSwDispatch2ProtocolGuid)) {
    DEBUG ((DEBUG_INFO, "  SwSmi - 0x%lx\n", ((SMI_HANDLER_PROFILE_SW_REGISTER_CONTEXT *)Context)->SwSmiInputValue));
  } else if (CompareGuid (HandlerType, &gEfiSmmSxDispatch2ProtocolGuid)) {
    DEBUG ((DEBUG_INFO, "  SxType - 0x%x\n", ((EFI_SMM_SX_REGISTER_CONTEXT *)Context)->Type));
    DEBUG ((DEBUG_INFO, "  SxPhase - 0x%x\n", ((EFI_SMM_SX_REGISTER_CONTEXT *)Context)->Phase));
  } else if (CompareGuid (HandlerType, &gEfiSmmPowerButtonDispatch2ProtocolGuid)) {
    DEBUG ((DEBUG_INFO, "  PowerButtonPhase - 0x%x\n", ((EFI_SMM_POWER_BUTTON_REGISTER_CONTEXT *)Context)->Phase));
  } else if (CompareGuid (HandlerType, &gEfiSmmStandbyButtonDispatch2ProtocolGuid)) {
    DEBUG ((DEBUG_INFO, "  StandbyButtonPhase - 0x%x\n", ((EFI_SMM_STANDBY_BUTTON_REGISTER_CONTEXT *)Context)->Phase));
  } else if (CompareGuid (HandlerType, &gEfiSmmPeriodicTimerDispatch2ProtocolGuid)) {
    DEBUG ((DEBUG_INFO, "  PeriodicTimerPeriod - %ld\n", ((EFI_SMM_PERIODIC_TIMER_REGISTER_CONTEXT *)Context)->Period));
    DEBUG ((DEBUG_INFO, "  PeriodicTimerSmiTickInterval - %ld\n", ((EFI_SMM_PERIODIC_TIMER_REGISTER_CONTEXT *)Context)->SmiTickInterval));
  } else if (CompareGuid (HandlerType, &gEfiSmmGpiDispatch2ProtocolGuid)) {
    DEBUG ((DEBUG_INFO, "  GpiNum - 0x%lx\n", ((EFI_SMM_GPI_REGISTER_CONTEXT *)Context)->GpiNum));
  } else if (CompareGuid (HandlerType, &gEfiSmmIoTrapDispatch2ProtocolGuid)) {
    DEBUG ((DEBUG_INFO, "  IoTrapAddress - 0x%x\n", ((EFI_SMM_IO_TRAP_REGISTER_CONTEXT *)Context)->Address));
    DEBUG ((DEBUG_INFO, "  IoTrapLength - 0x%x\n", ((EFI_SMM_IO_TRAP_REGISTER_CONTEXT *)Context)->Length));
    DEBUG ((DEBUG_INFO, "  IoTrapType - 0x%x\n", ((EFI_SMM_IO_TRAP_REGISTER_CONTEXT *)Context)->Type));
  } else if (CompareGuid (HandlerType, &gEfiSmmUsbDispatch2ProtocolGuid)) {
    DEBUG ((DEBUG_INFO, "  UsbType - 0x%x\n", ((SMI_HANDLER_PROFILE_USB_REGISTER_CONTEXT *)Context)->Type));
    Str = ConvertDevicePathToText((EFI_DEVICE_PATH_PROTOCOL *)(((SMI_HANDLER_PROFILE_USB_REGISTER_CONTEXT *)Context) + 1), TRUE, TRUE);
    DEBUG ((DEBUG_INFO, "  UsbDevicePath - %s\n", Str));
    if (Str != NULL) {
      FreePool (Str);
    }
  } else {
    DEBUG ((DEBUG_INFO, "  Context - "));
    InternalDumpData (Context, ContextSize);
    DEBUG ((DEBUG_INFO, "\n"));
  }
}

/**
  Dump all SMI handlers associated with SmiEntry.

  @param SmiEntry  SMI entry.
**/
VOID
DumpSmiHandlerOnSmiEntry(
  IN SMI_ENTRY       *SmiEntry
  )
{
  LIST_ENTRY      *ListEntry;
  SMI_HANDLER     *SmiHandler;
  IMAGE_STRUCT    *ImageStruct;

  ListEntry = &SmiEntry->SmiHandlers;
  for (ListEntry = ListEntry->ForwardLink;
       ListEntry != &SmiEntry->SmiHandlers;
       ListEntry = ListEntry->ForwardLink) {
    SmiHandler = CR(ListEntry, SMI_HANDLER, Link, SMI_HANDLER_SIGNATURE);
    ImageStruct = AddressToImageStruct((UINTN)SmiHandler->Handler);
    if (ImageStruct != NULL) {
      DEBUG ((DEBUG_INFO, " Module - %g", &ImageStruct->FileGuid));
    }
    if ((ImageStruct != NULL) && (ImageStruct->PdbString[0] != 0)) {
      DEBUG ((DEBUG_INFO, " (Pdb - %a)", ImageStruct->PdbString));
    }
    DEBUG ((DEBUG_INFO, "\n"));
    if (SmiHandler->ContextSize != 0) {
      DumpSmiChildContext (&SmiEntry->HandlerType, SmiHandler->Context, SmiHandler->ContextSize);
    }
    DEBUG ((DEBUG_INFO, "  Handler - 0x%x", SmiHandler->Handler));
    if (ImageStruct != NULL) {
      DEBUG ((DEBUG_INFO, " <== RVA - 0x%x", (UINTN)SmiHandler->Handler - (UINTN) ImageStruct->ImageBase));
    }
    DEBUG ((DEBUG_INFO, "\n"));
    DEBUG ((DEBUG_INFO, "  CallerAddr - 0x%x", SmiHandler->CallerAddr));
    if (ImageStruct != NULL) {
      DEBUG ((DEBUG_INFO, " <== RVA - 0x%x", SmiHandler->CallerAddr - (UINTN) ImageStruct->ImageBase));
    }
    DEBUG ((DEBUG_INFO, "\n"));
  }

  return;
}

/**
  Dump all SMI entry on the list.

  @param SmiEntryList a list of SMI entry.
**/
VOID
DumpSmiEntryList(
  IN LIST_ENTRY      *SmiEntryList
  )
{
  LIST_ENTRY      *ListEntry;
  SMI_ENTRY       *SmiEntry;

  ListEntry = SmiEntryList;
  for (ListEntry = ListEntry->ForwardLink;
       ListEntry != SmiEntryList;
       ListEntry = ListEntry->ForwardLink) {
    SmiEntry = CR(ListEntry, SMI_ENTRY, AllEntries, SMI_ENTRY_SIGNATURE);
    DEBUG ((DEBUG_INFO, "SmiEntry - %g\n", &SmiEntry->HandlerType));
    DumpSmiHandlerOnSmiEntry(SmiEntry);
  }

  return;
}

/**
  SMM Ready To Lock event notification handler.

  This function collects all SMM image information and build SmiHandleProfile database,
  and register SmiHandlerProfile SMI handler.

  @param[in] Protocol   Points to the protocol's unique identifier.
  @param[in] Interface  Points to the interface instance.
  @param[in] Handle     The handle on which the interface was installed.

  @retval EFI_SUCCESS   Notification handler runs successfully.
**/
EFI_STATUS
EFIAPI
SmmReadyToLockInSmiHandlerProfile (
  IN CONST EFI_GUID  *Protocol,
  IN VOID            *Interface,
  IN EFI_HANDLE      Handle
  )
{
  //
  // Dump all image
  //
  DEBUG ((DEBUG_INFO, "##################\n"));
  DEBUG ((DEBUG_INFO, "# IMAGE DATABASE #\n"));
  DEBUG ((DEBUG_INFO, "##################\n"));
  GetSmmLoadedImage ();
  DEBUG ((DEBUG_INFO, "\n"));

  //
  // Dump SMI Handler
  //
  DEBUG ((DEBUG_INFO, "########################\n"));
  DEBUG ((DEBUG_INFO, "# SMI Handler DATABASE #\n"));
  DEBUG ((DEBUG_INFO, "########################\n"));

  DEBUG ((DEBUG_INFO, "# 1. ROOT SMI Handler #\n"));
  DEBUG_CODE (
    DumpSmiEntryList(mSmmCoreRootSmiEntryList);
  );

  DEBUG ((DEBUG_INFO, "# 2. GUID SMI Handler #\n"));
  DEBUG_CODE (
    DumpSmiEntryList(mSmmCoreSmiEntryList);
  );

  DEBUG ((DEBUG_INFO, "# 3. Hardware SMI Handler #\n"));
  DEBUG_CODE (
    DumpSmiEntryList(mSmmCoreHardwareSmiEntryList);
  );

  DEBUG ((DEBUG_INFO, "\n"));

  RegisterSmiHandlerProfileHandler();

  if (mImageStruct != NULL) {
    FreePool(mImageStruct);
  }

  return EFI_SUCCESS;
}

/**
  returns SMM image data base size.

  @return SMM image data base size.
**/
UINTN
GetSmmImageDatabaseSize(
  VOID
  )
{
  UINTN  Size;
  UINT32 Index;

  Size = 0;
  for (Index = 0; Index < mImageStructCount; Index++) {
    Size += sizeof(SMM_CORE_IMAGE_DATABASE_STRUCTURE) + GET_OCCUPIED_SIZE (mImageStruct[Index].PdbStringSize, sizeof (UINT64));
  }
  return Size;
}

/**
  returns all SMI handlers' size associated with SmiEntry.

  @param SmiEntry  SMI entry.

  @return all SMI handlers' size associated with SmiEntry.
**/
UINTN
GetSmmSmiHandlerSizeOnSmiEntry(
  IN SMI_ENTRY       *SmiEntry
  )
{
  LIST_ENTRY      *ListEntry;
  SMI_HANDLER     *SmiHandler;
  UINTN           Size;

  Size = 0;
  ListEntry = &SmiEntry->SmiHandlers;
  for (ListEntry = ListEntry->ForwardLink;
       ListEntry != &SmiEntry->SmiHandlers;
       ListEntry = ListEntry->ForwardLink) {
    SmiHandler = CR(ListEntry, SMI_HANDLER, Link, SMI_HANDLER_SIGNATURE);
    Size += sizeof(SMM_CORE_SMI_HANDLER_STRUCTURE) + GET_OCCUPIED_SIZE (SmiHandler->ContextSize, sizeof (UINT64));
  }

  return Size;
}

/**
  return all SMI handler database size on the SMI entry list.

  @param SmiEntryList a list of SMI entry.

  @return all SMI handler database size on the SMI entry list.
**/
UINTN
GetSmmSmiDatabaseSize(
  IN LIST_ENTRY      *SmiEntryList
  )
{
  LIST_ENTRY      *ListEntry;
  SMI_ENTRY       *SmiEntry;
  UINTN           Size;

  Size = 0;
  ListEntry = SmiEntryList;
  for (ListEntry = ListEntry->ForwardLink;
       ListEntry != SmiEntryList;
       ListEntry = ListEntry->ForwardLink) {
    SmiEntry = CR(ListEntry, SMI_ENTRY, AllEntries, SMI_ENTRY_SIGNATURE);
    Size += sizeof(SMM_CORE_SMI_DATABASE_STRUCTURE);
    Size += GetSmmSmiHandlerSizeOnSmiEntry(SmiEntry);
  }
  return Size;
}

/**
  return SMI handler profile database size.

  @return SMI handler profile database size.
**/
UINTN
GetSmiHandlerProfileDatabaseSize (
  VOID
  )
{
  mSmmImageDatabaseSize = GetSmmImageDatabaseSize();
  mSmmRootSmiDatabaseSize = GetSmmSmiDatabaseSize(mSmmCoreRootSmiEntryList);
  mSmmSmiDatabaseSize = GetSmmSmiDatabaseSize(mSmmCoreSmiEntryList);
  mSmmHardwareSmiDatabaseSize = GetSmmSmiDatabaseSize(mSmmCoreHardwareSmiEntryList);

  return mSmmImageDatabaseSize + mSmmSmiDatabaseSize + mSmmRootSmiDatabaseSize + mSmmHardwareSmiDatabaseSize;
}

/**
  get SMM image database.

  @param Data           The buffer to hold SMM image database
  @param ExpectedSize   The expected size of the SMM image database

  @return SMM image data base size.
**/
UINTN
GetSmmImageDatabaseData (
  IN OUT VOID  *Data,
  IN     UINTN ExpectedSize
  )
{
  SMM_CORE_IMAGE_DATABASE_STRUCTURE   *ImageStruct;
  UINTN                               Size;
  UINTN                               Index;

  ImageStruct = Data;
  Size = 0;
  for (Index = 0; Index < mImageStructCount; Index++) {
    if (Size >= ExpectedSize) {
      return 0;
    }
    if (sizeof(SMM_CORE_IMAGE_DATABASE_STRUCTURE) + GET_OCCUPIED_SIZE (mImageStruct[Index].PdbStringSize, sizeof (UINT64)) > ExpectedSize - Size) {
      return 0;
    }
    ImageStruct->Header.Signature = SMM_CORE_IMAGE_DATABASE_SIGNATURE;
    ImageStruct->Header.Length = (UINT32)(sizeof(SMM_CORE_IMAGE_DATABASE_STRUCTURE) + GET_OCCUPIED_SIZE (mImageStruct[Index].PdbStringSize, sizeof (UINT64)));
    ImageStruct->Header.Revision = SMM_CORE_IMAGE_DATABASE_REVISION;
    CopyGuid(&ImageStruct->FileGuid, &mImageStruct[Index].FileGuid);
    ImageStruct->ImageRef = mImageStruct[Index].ImageRef;
    ImageStruct->EntryPoint = mImageStruct[Index].EntryPoint;
    ImageStruct->ImageBase = mImageStruct[Index].ImageBase;
    ImageStruct->ImageSize = mImageStruct[Index].ImageSize;
    if (mImageStruct[Index].PdbStringSize != 0) {
      ImageStruct->PdbStringOffset = sizeof(SMM_CORE_IMAGE_DATABASE_STRUCTURE);
      CopyMem ((VOID *)((UINTN)ImageStruct + ImageStruct->PdbStringOffset), mImageStruct[Index].PdbString, mImageStruct[Index].PdbStringSize);
    } else {
      ImageStruct->PdbStringOffset = 0;
    }
    ImageStruct = (SMM_CORE_IMAGE_DATABASE_STRUCTURE *)((UINTN)ImageStruct + ImageStruct->Header.Length);
    Size += sizeof(SMM_CORE_IMAGE_DATABASE_STRUCTURE) + GET_OCCUPIED_SIZE (mImageStruct[Index].PdbStringSize, sizeof (UINT64));
  }

  if (ExpectedSize != Size) {
    return 0;
  }
  return Size;
}

/**
  get all SMI handler data associated with SmiEntry.

  @param SmiEntry       SMI entry.
  @param Data           The buffer to hold all SMI handler data
  @param MaxSize        The max size of the SMM image database
  @param Count          The count of the SMI handler.

  @return SMM image data base size.
**/
UINTN
GetSmmSmiHandlerDataOnSmiEntry(
  IN     SMI_ENTRY       *SmiEntry,
  IN OUT VOID            *Data,
  IN     UINTN           MaxSize,
     OUT UINT32          *Count
  )
{
  SMM_CORE_SMI_HANDLER_STRUCTURE   *SmiHandlerStruct;
  LIST_ENTRY                       *ListEntry;
  SMI_HANDLER                      *SmiHandler;
  UINTN                            Size;

  SmiHandlerStruct = Data;
  Size = 0;
  *Count = 0;
  ListEntry = &SmiEntry->SmiHandlers;
  for (ListEntry = ListEntry->ForwardLink;
       ListEntry != &SmiEntry->SmiHandlers;
       ListEntry = ListEntry->ForwardLink) {
    SmiHandler = CR(ListEntry, SMI_HANDLER, Link, SMI_HANDLER_SIGNATURE);
    if (Size >= MaxSize) {
      *Count = 0;
      return 0;
    }
    if (sizeof(SMM_CORE_SMI_HANDLER_STRUCTURE) + GET_OCCUPIED_SIZE (SmiHandler->ContextSize, sizeof (UINT64)) > MaxSize - Size) {
      *Count = 0;
      return 0;
    }
    SmiHandlerStruct->Length = (UINT32)(sizeof(SMM_CORE_SMI_HANDLER_STRUCTURE) + GET_OCCUPIED_SIZE (SmiHandler->ContextSize, sizeof (UINT64)));
    SmiHandlerStruct->CallerAddr = (UINTN)SmiHandler->CallerAddr;
    SmiHandlerStruct->Handler = (UINTN)SmiHandler->Handler;
    SmiHandlerStruct->ImageRef = AddressToImageRef((UINTN)SmiHandler->Handler);
    SmiHandlerStruct->ContextBufferSize = (UINT32)SmiHandler->ContextSize;
    if (SmiHandler->ContextSize != 0) {
      SmiHandlerStruct->ContextBufferOffset = sizeof(SMM_CORE_SMI_HANDLER_STRUCTURE);
      CopyMem ((UINT8 *)SmiHandlerStruct + SmiHandlerStruct->ContextBufferOffset, SmiHandler->Context, SmiHandler->ContextSize);
    } else {
      SmiHandlerStruct->ContextBufferOffset = 0;
    }
    Size += sizeof(SMM_CORE_SMI_HANDLER_STRUCTURE) + GET_OCCUPIED_SIZE (SmiHandler->ContextSize, sizeof (UINT64));
    SmiHandlerStruct = (SMM_CORE_SMI_HANDLER_STRUCTURE *)((UINTN)SmiHandlerStruct + SmiHandlerStruct->Length);
    *Count = *Count + 1;
  }

  return Size;
}

/**
  get all SMI handler database on the SMI entry list.

  @param SmiEntryList     a list of SMI entry.
  @param HandlerCategory  The handler category
  @param Data             The buffer to hold all SMI handler database
  @param ExpectedSize     The expected size of the SMM image database

  @return all SMI database size on the SMI entry list.
**/
UINTN
GetSmmSmiDatabaseData(
  IN     LIST_ENTRY      *SmiEntryList,
  IN     UINT32          HandlerCategory,
  IN OUT VOID            *Data,
  IN     UINTN           ExpectedSize
  )
{
  SMM_CORE_SMI_DATABASE_STRUCTURE   *SmiStruct;
  LIST_ENTRY                        *ListEntry;
  SMI_ENTRY                         *SmiEntry;
  UINTN                             Size;
  UINTN                             SmiHandlerSize;
  UINT32                            SmiHandlerCount;

  SmiStruct = Data;
  Size = 0;
  ListEntry = SmiEntryList;
  for (ListEntry = ListEntry->ForwardLink;
       ListEntry != SmiEntryList;
       ListEntry = ListEntry->ForwardLink) {
    SmiEntry = CR(ListEntry, SMI_ENTRY, AllEntries, SMI_ENTRY_SIGNATURE);
    if (Size >= ExpectedSize) {
      return 0;
    }
    if (sizeof(SMM_CORE_SMI_DATABASE_STRUCTURE) > ExpectedSize - Size) {
      return 0;
    }

    SmiStruct->Header.Signature = SMM_CORE_SMI_DATABASE_SIGNATURE;
    SmiStruct->Header.Length = sizeof(SMM_CORE_SMI_DATABASE_STRUCTURE);
    SmiStruct->Header.Revision = SMM_CORE_SMI_DATABASE_REVISION;
    SmiStruct->HandlerCategory = HandlerCategory;
    CopyGuid(&SmiStruct->HandlerType, &SmiEntry->HandlerType);
    Size += sizeof(SMM_CORE_SMI_DATABASE_STRUCTURE);
    SmiHandlerSize = GetSmmSmiHandlerDataOnSmiEntry(SmiEntry, (UINT8 *)SmiStruct + SmiStruct->Header.Length, ExpectedSize - Size, &SmiHandlerCount);
    SmiStruct->HandlerCount = SmiHandlerCount;
    Size += SmiHandlerSize;
    SmiStruct->Header.Length += (UINT32)SmiHandlerSize;
    SmiStruct = (VOID *)((UINTN)SmiStruct + SmiStruct->Header.Length);
  }
  if (ExpectedSize != Size) {
    return 0;
  }
  return Size;
}

/**
  Get SMI handler profile database.

  @param Data the buffer to hold SMI handler profile database

  @retval EFI_SUCCESS            the database is got.
  @retval EFI_INVALID_PARAMETER  the database size mismatch.
**/
EFI_STATUS
GetSmiHandlerProfileDatabaseData(
  IN OUT VOID *Data
  )
{
  UINTN  SmmImageDatabaseSize;
  UINTN  SmmSmiDatabaseSize;
  UINTN  SmmRootSmiDatabaseSize;
  UINTN  SmmHardwareSmiDatabaseSize;

  DEBUG((DEBUG_VERBOSE, "GetSmiHandlerProfileDatabaseData\n"));
  SmmImageDatabaseSize = GetSmmImageDatabaseData(Data, mSmmImageDatabaseSize);
  if (SmmImageDatabaseSize != mSmmImageDatabaseSize) {
    DEBUG((DEBUG_ERROR, "GetSmiHandlerProfileDatabaseData - SmmImageDatabaseSize mismatch!\n"));
    return EFI_INVALID_PARAMETER;
  }
  SmmRootSmiDatabaseSize = GetSmmSmiDatabaseData(mSmmCoreRootSmiEntryList, SmmCoreSmiHandlerCategoryRootHandler, (UINT8 *)Data + SmmImageDatabaseSize, mSmmRootSmiDatabaseSize);
  if (SmmRootSmiDatabaseSize != mSmmRootSmiDatabaseSize) {
    DEBUG((DEBUG_ERROR, "GetSmiHandlerProfileDatabaseData - SmmRootSmiDatabaseSize mismatch!\n"));
    return EFI_INVALID_PARAMETER;
  }
  SmmSmiDatabaseSize = GetSmmSmiDatabaseData(mSmmCoreSmiEntryList, SmmCoreSmiHandlerCategoryGuidHandler, (UINT8 *)Data + SmmImageDatabaseSize + mSmmRootSmiDatabaseSize, mSmmSmiDatabaseSize);
  if (SmmSmiDatabaseSize != mSmmSmiDatabaseSize) {
    DEBUG((DEBUG_ERROR, "GetSmiHandlerProfileDatabaseData - SmmSmiDatabaseSize mismatch!\n"));
    return EFI_INVALID_PARAMETER;
  }
  SmmHardwareSmiDatabaseSize = GetSmmSmiDatabaseData(mSmmCoreHardwareSmiEntryList, SmmCoreSmiHandlerCategoryHardwareHandler, (UINT8 *)Data + SmmImageDatabaseSize + SmmRootSmiDatabaseSize + SmmSmiDatabaseSize, mSmmHardwareSmiDatabaseSize);
  if (SmmHardwareSmiDatabaseSize != mSmmHardwareSmiDatabaseSize) {
    DEBUG((DEBUG_ERROR, "GetSmiHandlerProfileDatabaseData - SmmHardwareSmiDatabaseSize mismatch!\n"));
    return EFI_INVALID_PARAMETER;
  }

  return EFI_SUCCESS;
}

/**
  build SMI handler profile database.
**/
VOID
BuildSmiHandlerProfileDatabase(
  VOID
  )
{
  EFI_STATUS  Status;
  mSmiHandlerProfileDatabaseSize = GetSmiHandlerProfileDatabaseSize();
  mSmiHandlerProfileDatabase = AllocatePool(mSmiHandlerProfileDatabaseSize);
  if (mSmiHandlerProfileDatabase == NULL) {
    return;
  }
  Status = GetSmiHandlerProfileDatabaseData(mSmiHandlerProfileDatabase);
  if (EFI_ERROR(Status)) {
    FreePool(mSmiHandlerProfileDatabase);
    mSmiHandlerProfileDatabase = NULL;
  }
}

/**
  Copy SMI handler profile data.

  @param DataBuffer  The buffer to hold SMI handler profile data.
  @param DataSize    On input, data buffer size.
                     On output, actual data buffer size copied.
  @param DataOffset  On input, data buffer offset to copy.
                     On output, next time data buffer offset to copy.

**/
VOID
SmiHandlerProfileCopyData(
  OUT VOID      *DataBuffer,
  IN OUT UINT64 *DataSize,
  IN OUT UINT64 *DataOffset
  )
{
  if (*DataOffset >= mSmiHandlerProfileDatabaseSize) {
    *DataOffset = mSmiHandlerProfileDatabaseSize;
    return;
  }
  if (mSmiHandlerProfileDatabaseSize - *DataOffset < *DataSize) {
    *DataSize = mSmiHandlerProfileDatabaseSize - *DataOffset;
  }

  CopyMem(
    DataBuffer,
    (UINT8 *)mSmiHandlerProfileDatabase + *DataOffset,
    (UINTN)*DataSize
    );
  *DataOffset = *DataOffset + *DataSize;
}

/**
  SMI handler profile handler to get info.

  @param SmiHandlerProfileParameterGetInfo The parameter of SMI handler profile get info.

**/
VOID
SmiHandlerProfileHandlerGetInfo(
  IN SMI_HANDLER_PROFILE_PARAMETER_GET_INFO   *SmiHandlerProfileParameterGetInfo
  )
{
  BOOLEAN                       SmiHandlerProfileRecordingStatus;

  SmiHandlerProfileRecordingStatus = mSmiHandlerProfileRecordingStatus;
  mSmiHandlerProfileRecordingStatus = FALSE;

  SmiHandlerProfileParameterGetInfo->DataSize = mSmiHandlerProfileDatabaseSize;
  SmiHandlerProfileParameterGetInfo->Header.ReturnStatus = 0;

  mSmiHandlerProfileRecordingStatus = SmiHandlerProfileRecordingStatus;
}

/**
  SMI handler profile handler to get data by offset.

  @param SmiHandlerProfileParameterGetDataByOffset   The parameter of SMI handler profile get data by offset.

**/
VOID
SmiHandlerProfileHandlerGetDataByOffset(
  IN SMI_HANDLER_PROFILE_PARAMETER_GET_DATA_BY_OFFSET     *SmiHandlerProfileParameterGetDataByOffset
  )
{
  SMI_HANDLER_PROFILE_PARAMETER_GET_DATA_BY_OFFSET    SmiHandlerProfileGetDataByOffset;
  BOOLEAN                                             SmiHandlerProfileRecordingStatus;

  SmiHandlerProfileRecordingStatus = mSmiHandlerProfileRecordingStatus;
  mSmiHandlerProfileRecordingStatus = FALSE;

  CopyMem(&SmiHandlerProfileGetDataByOffset, SmiHandlerProfileParameterGetDataByOffset, sizeof(SmiHandlerProfileGetDataByOffset));

  //
  // Sanity check
  //
  if (!SmmIsBufferOutsideSmmValid((UINTN)SmiHandlerProfileGetDataByOffset.DataBuffer, (UINTN)SmiHandlerProfileGetDataByOffset.DataSize)) {
    DEBUG((DEBUG_ERROR, "SmiHandlerProfileHandlerGetDataByOffset: SMI handler profile get data in SMRAM or overflow!\n"));
    SmiHandlerProfileParameterGetDataByOffset->Header.ReturnStatus = (UINT64)(INT64)(INTN)EFI_ACCESS_DENIED;
    goto Done;
  }

  SmiHandlerProfileCopyData((VOID *)(UINTN)SmiHandlerProfileGetDataByOffset.DataBuffer, &SmiHandlerProfileGetDataByOffset.DataSize, &SmiHandlerProfileGetDataByOffset.DataOffset);
  CopyMem(SmiHandlerProfileParameterGetDataByOffset, &SmiHandlerProfileGetDataByOffset, sizeof(SmiHandlerProfileGetDataByOffset));
  SmiHandlerProfileParameterGetDataByOffset->Header.ReturnStatus = 0;

Done:
  mSmiHandlerProfileRecordingStatus = SmiHandlerProfileRecordingStatus;
}

/**
  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
SmiHandlerProfileHandler(
  IN EFI_HANDLE  DispatchHandle,
  IN CONST VOID  *Context         OPTIONAL,
  IN OUT VOID    *CommBuffer      OPTIONAL,
  IN OUT UINTN   *CommBufferSize  OPTIONAL
  )
{
  SMI_HANDLER_PROFILE_PARAMETER_HEADER           *SmiHandlerProfileParameterHeader;
  UINTN                                    TempCommBufferSize;

  DEBUG((DEBUG_ERROR, "SmiHandlerProfileHandler Enter\n"));

  if (mSmiHandlerProfileDatabase == NULL) {
    return EFI_SUCCESS;
  }

  //
  // If input is invalid, stop processing this SMI
  //
  if (CommBuffer == NULL || CommBufferSize == NULL) {
    return EFI_SUCCESS;
  }

  TempCommBufferSize = *CommBufferSize;

  if (TempCommBufferSize < sizeof(SMI_HANDLER_PROFILE_PARAMETER_HEADER)) {
    DEBUG((DEBUG_ERROR, "SmiHandlerProfileHandler: SMM communication buffer size invalid!\n"));
    return EFI_SUCCESS;
  }

  if (!SmmIsBufferOutsideSmmValid((UINTN)CommBuffer, TempCommBufferSize)) {
    DEBUG((DEBUG_ERROR, "SmiHandlerProfileHandler: SMM communication buffer in SMRAM or overflow!\n"));
    return EFI_SUCCESS;
  }

  SmiHandlerProfileParameterHeader = (SMI_HANDLER_PROFILE_PARAMETER_HEADER *)((UINTN)CommBuffer);
  SmiHandlerProfileParameterHeader->ReturnStatus = (UINT64)-1;

  switch (SmiHandlerProfileParameterHeader->Command) {
  case SMI_HANDLER_PROFILE_COMMAND_GET_INFO:
    DEBUG((DEBUG_ERROR, "SmiHandlerProfileHandlerGetInfo\n"));
    if (TempCommBufferSize != sizeof(SMI_HANDLER_PROFILE_PARAMETER_GET_INFO)) {
      DEBUG((DEBUG_ERROR, "SmiHandlerProfileHandler: SMM communication buffer size invalid!\n"));
      return EFI_SUCCESS;
    }
    SmiHandlerProfileHandlerGetInfo((SMI_HANDLER_PROFILE_PARAMETER_GET_INFO *)(UINTN)CommBuffer);
    break;
  case SMI_HANDLER_PROFILE_COMMAND_GET_DATA_BY_OFFSET:
    DEBUG((DEBUG_ERROR, "SmiHandlerProfileHandlerGetDataByOffset\n"));
    if (TempCommBufferSize != sizeof(SMI_HANDLER_PROFILE_PARAMETER_GET_DATA_BY_OFFSET)) {
      DEBUG((DEBUG_ERROR, "SmiHandlerProfileHandler: SMM communication buffer size invalid!\n"));
      return EFI_SUCCESS;
    }
    SmiHandlerProfileHandlerGetDataByOffset((SMI_HANDLER_PROFILE_PARAMETER_GET_DATA_BY_OFFSET *)(UINTN)CommBuffer);
    break;
  default:
    break;
  }

  DEBUG((DEBUG_ERROR, "SmiHandlerProfileHandler Exit\n"));

  return EFI_SUCCESS;
}

/**
  Register SMI handler profile handler.
**/
VOID
RegisterSmiHandlerProfileHandler (
  VOID
  )
{
  EFI_STATUS    Status;
  EFI_HANDLE    DispatchHandle;

  Status = gSmst->SmiHandlerRegister (
                    SmiHandlerProfileHandler,
                    &gSmiHandlerProfileGuid,
                    &DispatchHandle
                    );
  ASSERT_EFI_ERROR (Status);

  BuildSmiHandlerProfileDatabase();
}

/**
  Finds the SMI entry for the requested handler type.

  @param  HandlerType            The type of the interrupt
  @param  Create                 Create a new entry if not found

  @return SMI entry
**/
SMI_ENTRY  *
SmmCoreFindHardwareSmiEntry (
  IN EFI_GUID  *HandlerType,
  IN BOOLEAN   Create
  )
{
  LIST_ENTRY  *Link;
  SMI_ENTRY   *Item;
  SMI_ENTRY   *SmiEntry;

  //
  // Search the SMI entry list for the matching GUID
  //
  SmiEntry = NULL;
  for (Link = mHardwareSmiEntryList.ForwardLink;
       Link != &mHardwareSmiEntryList;
       Link = Link->ForwardLink) {

    Item = CR (Link, SMI_ENTRY, AllEntries, SMI_ENTRY_SIGNATURE);
    if (CompareGuid (&Item->HandlerType, HandlerType)) {
      //
      // This is the SMI entry
      //
      SmiEntry = Item;
      break;
    }
  }

  //
  // If the protocol entry was not found and Create is TRUE, then
  // allocate a new entry
  //
  if ((SmiEntry == NULL) && Create) {
    SmiEntry = AllocatePool (sizeof(SMI_ENTRY));
    if (SmiEntry != NULL) {
      //
      // Initialize new SMI entry structure
      //
      SmiEntry->Signature = SMI_ENTRY_SIGNATURE;
      CopyGuid ((VOID *)&SmiEntry->HandlerType, HandlerType);
      InitializeListHead (&SmiEntry->SmiHandlers);

      //
      // Add it to SMI entry list
      //
      InsertTailList (&mHardwareSmiEntryList, &SmiEntry->AllEntries);
    }
  }
  return SmiEntry;
}

/**
  Convert EFI_SMM_USB_REGISTER_CONTEXT to SMI_HANDLER_PROFILE_USB_REGISTER_CONTEXT.

  @param UsbContext                   A pointer to EFI_SMM_USB_REGISTER_CONTEXT
  @param UsbContextSize               The size of EFI_SMM_USB_REGISTER_CONTEXT in bytes
  @param SmiHandlerUsbContextSize     The size of SMI_HANDLER_PROFILE_USB_REGISTER_CONTEXT in bytes

  @return SmiHandlerUsbContext  A pointer to SMI_HANDLER_PROFILE_USB_REGISTER_CONTEXT
**/
SMI_HANDLER_PROFILE_USB_REGISTER_CONTEXT *
ConvertSmiHandlerUsbContext (
  IN EFI_SMM_USB_REGISTER_CONTEXT   *UsbContext,
  IN UINTN                          UsbContextSize,
  OUT UINTN                         *SmiHandlerUsbContextSize
  )
{
  UINTN                                     DevicePathSize;
  SMI_HANDLER_PROFILE_USB_REGISTER_CONTEXT  *SmiHandlerUsbContext;

  ASSERT (UsbContextSize == sizeof(EFI_SMM_USB_REGISTER_CONTEXT));

  DevicePathSize = GetDevicePathSize (UsbContext->Device);
  SmiHandlerUsbContext = AllocatePool (sizeof (SMI_HANDLER_PROFILE_USB_REGISTER_CONTEXT) + DevicePathSize);
  if (SmiHandlerUsbContext == NULL) {
    *SmiHandlerUsbContextSize = 0;
    return NULL;
  }
  SmiHandlerUsbContext->Type = UsbContext->Type;
  SmiHandlerUsbContext->DevicePathSize = (UINT32)DevicePathSize;
  CopyMem (SmiHandlerUsbContext + 1, UsbContext->Device, DevicePathSize);
  *SmiHandlerUsbContextSize = sizeof (SMI_HANDLER_PROFILE_USB_REGISTER_CONTEXT) + DevicePathSize;
  return SmiHandlerUsbContext;
}

/**
  Convert EFI_SMM_SW_REGISTER_CONTEXT to SMI_HANDLER_PROFILE_SW_REGISTER_CONTEXT.

  @param SwContext                    A pointer to EFI_SMM_SW_REGISTER_CONTEXT
  @param SwContextSize                The size of EFI_SMM_SW_REGISTER_CONTEXT in bytes
  @param SmiHandlerSwContextSize      The size of SMI_HANDLER_PROFILE_SW_REGISTER_CONTEXT in bytes

  @return SmiHandlerSwContext   A pointer to SMI_HANDLER_PROFILE_SW_REGISTER_CONTEXT
**/
SMI_HANDLER_PROFILE_SW_REGISTER_CONTEXT *
ConvertSmiHandlerSwContext (
  IN EFI_SMM_SW_REGISTER_CONTEXT    *SwContext,
  IN UINTN                          SwContextSize,
  OUT UINTN                         *SmiHandlerSwContextSize
  )
{
  SMI_HANDLER_PROFILE_SW_REGISTER_CONTEXT  *SmiHandlerSwContext;

  ASSERT (SwContextSize == sizeof(EFI_SMM_SW_REGISTER_CONTEXT));

  SmiHandlerSwContext = AllocatePool (sizeof (SMI_HANDLER_PROFILE_SW_REGISTER_CONTEXT));
  if (SmiHandlerSwContext == NULL) {
    *SmiHandlerSwContextSize = 0;
    return NULL;
  }
  SmiHandlerSwContext->SwSmiInputValue = SwContext->SwSmiInputValue;
  *SmiHandlerSwContextSize = sizeof (SMI_HANDLER_PROFILE_SW_REGISTER_CONTEXT);
  return SmiHandlerSwContext;
}

/**
  This function is called by SmmChildDispatcher module to report
  a new SMI handler is registered, to SmmCore.

  @param This            The protocol instance
  @param HandlerGuid     The GUID to identify the type of the handler.
                         For the SmmChildDispatch protocol, the HandlerGuid
                         must be the GUID of SmmChildDispatch protocol.
  @param Handler         The SMI handler.
  @param CallerAddress   The address of the module who registers the SMI handler.
  @param Context         The context of the SMI handler.
                         For the SmmChildDispatch protocol, the Context
                         must match the one defined for SmmChildDispatch protocol.
  @param ContextSize     The size of the context in bytes.
                         For the SmmChildDispatch protocol, the Context
                         must match the one defined for SmmChildDispatch protocol.

  @retval EFI_SUCCESS           The information is recorded.
  @retval EFI_OUT_OF_RESOURCES  There is no enough resource to record the information.
**/
EFI_STATUS
EFIAPI
SmiHandlerProfileRegisterHandler (
  IN SMI_HANDLER_PROFILE_PROTOCOL   *This,
  IN EFI_GUID                       *HandlerGuid,
  IN EFI_SMM_HANDLER_ENTRY_POINT2   Handler,
  IN PHYSICAL_ADDRESS               CallerAddress,
  IN VOID                           *Context, OPTIONAL
  IN UINTN                          ContextSize OPTIONAL
  )
{
  SMI_HANDLER  *SmiHandler;
  SMI_ENTRY    *SmiEntry;
  LIST_ENTRY   *List;

  if (((ContextSize == 0) && (Context != NULL)) ||
      ((ContextSize != 0) && (Context == NULL))) {
    return EFI_INVALID_PARAMETER;
  }

  SmiHandler = AllocateZeroPool (sizeof (SMI_HANDLER));
  if (SmiHandler == NULL) {
    return EFI_OUT_OF_RESOURCES;
  }

  SmiHandler->Signature = SMI_HANDLER_SIGNATURE;
  SmiHandler->Handler = Handler;
  SmiHandler->CallerAddr = (UINTN)CallerAddress;
  SmiHandler->Context = Context;
  SmiHandler->ContextSize = ContextSize;

  if (Context != NULL) {
    if (CompareGuid (HandlerGuid, &gEfiSmmUsbDispatch2ProtocolGuid)) {
      SmiHandler->Context = ConvertSmiHandlerUsbContext (Context, ContextSize, &SmiHandler->ContextSize);
    } else if (CompareGuid (HandlerGuid, &gEfiSmmSwDispatch2ProtocolGuid)) {
      SmiHandler->Context = ConvertSmiHandlerSwContext (Context, ContextSize, &SmiHandler->ContextSize);
    } else {
      SmiHandler->Context = AllocateCopyPool (ContextSize, Context);
    }
  }
  if (SmiHandler->Context == NULL) {
    SmiHandler->ContextSize = 0;
  }

  SmiEntry = SmmCoreFindHardwareSmiEntry (HandlerGuid, TRUE);
  if (SmiEntry == NULL) {
    if (SmiHandler->Context != NULL) {
      FreePool (SmiHandler->Context);
    }
    FreePool (SmiHandler);
    return EFI_OUT_OF_RESOURCES;
  }

  List = &SmiEntry->SmiHandlers;

  SmiHandler->SmiEntry = SmiEntry;
  InsertTailList (List, &SmiHandler->Link);

  return EFI_SUCCESS;
}

/**
  This function is called by SmmChildDispatcher module to report
  an existing SMI handler is unregistered, to SmmCore.

  @param This            The protocol instance
  @param HandlerGuid     The GUID to identify the type of the handler.
                         For the SmmChildDispatch protocol, the HandlerGuid
                         must be the GUID of SmmChildDispatch protocol.
  @param Handler         The SMI handler.
  @param Context         The context of the SMI handler.
                         If it is NOT NULL, it will be used to check what is registered.
  @param ContextSize     The size of the context in bytes.
                         If Context is NOT NULL, it will be used to check what is registered.

  @retval EFI_SUCCESS           The original record is removed.
  @retval EFI_NOT_FOUND         There is no record for the HandlerGuid and handler.
**/
EFI_STATUS
EFIAPI
SmiHandlerProfileUnregisterHandler (
  IN SMI_HANDLER_PROFILE_PROTOCOL   *This,
  IN EFI_GUID                       *HandlerGuid,
  IN EFI_SMM_HANDLER_ENTRY_POINT2   Handler,
  IN VOID                           *Context, OPTIONAL
  IN UINTN                          ContextSize OPTIONAL
  )
{
  LIST_ENTRY   *Link;
  LIST_ENTRY   *Head;
  SMI_HANDLER  *SmiHandler;
  SMI_ENTRY    *SmiEntry;
  SMI_HANDLER  *TargetSmiHandler;
  VOID         *SearchContext;
  UINTN        SearchContextSize;

  if (((ContextSize == 0) && (Context != NULL)) ||
      ((ContextSize != 0) && (Context == NULL))) {
    return EFI_INVALID_PARAMETER;
  }

  SmiEntry = SmmCoreFindHardwareSmiEntry (HandlerGuid, FALSE);
  if (SmiEntry == NULL) {
    return EFI_NOT_FOUND;
  }

  SearchContext = Context;
  SearchContextSize = ContextSize;
  if (Context != NULL) {
    if (CompareGuid (HandlerGuid, &gEfiSmmUsbDispatch2ProtocolGuid)) {
      SearchContext = ConvertSmiHandlerUsbContext (Context, ContextSize, &SearchContextSize);
    } else if (CompareGuid (HandlerGuid, &gEfiSmmSwDispatch2ProtocolGuid)) {
      SearchContext = ConvertSmiHandlerSwContext (Context, ContextSize, &SearchContextSize);
    }
  }

  TargetSmiHandler = NULL;
  Head = &SmiEntry->SmiHandlers;
  for (Link = Head->ForwardLink; Link != Head; Link = Link->ForwardLink) {
    SmiHandler = CR (Link, SMI_HANDLER, Link, SMI_HANDLER_SIGNATURE);
    if (SmiHandler->Handler == Handler) {
      if ((SearchContext == NULL) ||
          ((SearchContextSize == SmiHandler->ContextSize) && (CompareMem (SearchContext, SmiHandler->Context, SearchContextSize) == 0))) {
        TargetSmiHandler = SmiHandler;
        break;
      }
    }
  }

  if (SearchContext != NULL) {
    if (CompareGuid (HandlerGuid, &gEfiSmmUsbDispatch2ProtocolGuid)) {
      FreePool (SearchContext);
    }
  }

  if (TargetSmiHandler == NULL) {
    return EFI_NOT_FOUND;
  }
  SmiHandler = TargetSmiHandler;

  RemoveEntryList (&SmiHandler->Link);
  if (SmiHandler->Context != NULL) {
    FreePool (SmiHandler->Context);
  }
  FreePool (SmiHandler);

  if (IsListEmpty (&SmiEntry->SmiHandlers)) {
    RemoveEntryList (&SmiEntry->AllEntries);
    FreePool (SmiEntry);
  }

  return EFI_SUCCESS;
}

/**
  Initialize SmiHandler profile feature.
**/
VOID
SmmCoreInitializeSmiHandlerProfile (
  VOID
  )
{
  EFI_STATUS  Status;
  VOID        *Registration;
  EFI_HANDLE  Handle;

  if ((PcdGet8 (PcdSmiHandlerProfilePropertyMask) & 0x1) != 0) {
    InsertTailList (&mRootSmiEntryList, &mRootSmiEntry.AllEntries);

    Status = gSmst->SmmRegisterProtocolNotify (
                      &gEfiSmmReadyToLockProtocolGuid,
                      SmmReadyToLockInSmiHandlerProfile,
                      &Registration
                      );
    ASSERT_EFI_ERROR (Status);

    Handle = NULL;
    Status = gSmst->SmmInstallProtocolInterface (
                      &Handle,
                      &gSmiHandlerProfileGuid,
                      EFI_NATIVE_INTERFACE,
                      &mSmiHandlerProfile
                      );
    ASSERT_EFI_ERROR (Status);
  }
}