/** @file

  Copyright (c) 2017, Intel Corporation. All rights reserved.<BR>
  This program and the accompanying materials
  are licensed and made available under the terms and conditions of the BSD License
  which accompanies this distribution.  The full text of the license may be found at
  http://opensource.org/licenses/bsd-license.php

  THE PROGRAM IS DISTRIBUTED UNDER THE BSD LICENSE ON AN "AS IS" BASIS,
  WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.

**/

#include "SmmTestPoint.h"

EFI_STATUS
GetAllSmmTestPointData (
  IN OUT UINTN  *DataSize,
  IN OUT VOID   *Data
  )
{
  EFI_STATUS                        Status;
  EFI_ADAPTER_INFORMATION_PROTOCOL  *Aip;
  UINTN                             NoHandles;
  EFI_HANDLE                        *Handles;
  UINTN                             HandleBufSize;
  UINTN                             Index;
  EFI_GUID                          *InfoTypesBuffer;
  UINTN                             InfoTypesBufferCount;
  UINTN                             InfoTypesIndex;
  EFI_ADAPTER_INFORMATION_PROTOCOL  *AipCandidate;
  VOID                              *InformationBlock;
  UINTN                             InformationBlockSize;
  ADAPTER_INFO_PLATFORM_TEST_POINT  *TestPoint;
  UINTN                             TotalSize;
  EFI_STATUS                        RetStatus;

  TotalSize = 0;

  Handles = NULL;
  HandleBufSize = 0;
  Status = gSmst->SmmLocateHandle (
                    ByProtocol,
                    &gEfiAdapterInformationProtocolGuid,
                    NULL,
                    &HandleBufSize,
                    Handles
                    );
  if (Status != EFI_BUFFER_TOO_SMALL) {
    RetStatus = EFI_NOT_FOUND;
    goto Done ;
  }
  Handles = AllocateZeroPool (HandleBufSize);
  if (Handles == NULL) {
    RetStatus = EFI_OUT_OF_RESOURCES;
    goto Done ;
  }
  Status = gSmst->SmmLocateHandle (
                    ByProtocol,
                    &gEfiAdapterInformationProtocolGuid,
                    NULL,
                    &HandleBufSize,
                    Handles
                    );
  if (EFI_ERROR (Status)) {
    RetStatus = Status;
    goto Done ;
  }
  NoHandles = HandleBufSize / sizeof(EFI_HANDLE);

  RetStatus = EFI_SUCCESS;
  
  TestPoint = NULL;
  Aip = NULL;
  InformationBlock = NULL;
  InformationBlockSize = 0;
  for (Index = 0; Index < NoHandles; Index++) {
    Status = gSmst->SmmHandleProtocol (
                      Handles[Index],
                      &gEfiAdapterInformationProtocolGuid,
                      (VOID **)&Aip
                      );
    if (EFI_ERROR (Status)) {
      continue;
    }

    //
    // Check AIP
    //
    Status = Aip->GetSupportedTypes (
                    Aip,
                    &InfoTypesBuffer,
                    &InfoTypesBufferCount
                    );
    if (EFI_ERROR (Status)) {
      continue;
    }

    AipCandidate = NULL;
    for (InfoTypesIndex = 0; InfoTypesIndex < InfoTypesBufferCount; InfoTypesIndex++) {
      if (CompareGuid (&InfoTypesBuffer[InfoTypesIndex], &gAdapterInfoPlatformTestPointGuid)) {
        AipCandidate = Aip;
        break;
      }
    }
    FreePool (InfoTypesBuffer);

    if (AipCandidate == NULL) {
      continue;
    }

    //
    // Check Role
    //
    Aip = AipCandidate;
    Status = Aip->GetInformation (
                    Aip,
                    &gAdapterInfoPlatformTestPointGuid,
                    &InformationBlock,
                    &InformationBlockSize
                    );
    if (EFI_ERROR (Status)) {
      continue;
    }

    if ((Data != NULL) && (TotalSize + InformationBlockSize <= *DataSize)) {
      CopyMem ((UINT8 *)Data + TotalSize, InformationBlock, InformationBlockSize);
    } else {
      RetStatus = EFI_BUFFER_TOO_SMALL;
    }
    TotalSize += InformationBlockSize;

    FreePool (InformationBlock);
  }

Done:

  *DataSize = TotalSize;

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

  return RetStatus;
}

/**
  SMM test point SMI handler to get info.

  @param SmiHandlerTestPointParameterGetInfo The parameter of SMM test point SMI handler get info.

**/
VOID
SmmTestPointSmiHandlerGetInfo (
  IN SMI_HANDLER_TEST_POINT_PARAMETER_GET_INFO   *SmiHandlerTestPointParameterGetInfo
  )
{
  UINTN       DataSize;
  EFI_STATUS  Status;
  
  DataSize = 0;
  Status = GetAllSmmTestPointData (&DataSize, NULL);
  if (Status == EFI_BUFFER_TOO_SMALL) {
    SmiHandlerTestPointParameterGetInfo->DataSize = DataSize;
    SmiHandlerTestPointParameterGetInfo->Header.ReturnStatus = 0;
  } else {
    SmiHandlerTestPointParameterGetInfo->DataSize = 0;
    SmiHandlerTestPointParameterGetInfo->Header.ReturnStatus = (UINT64)(INT64)(INTN)EFI_NOT_FOUND;
  }
}

/**
  Copy SMM Test Point data.

  @param DataBuffer  The buffer to hold SMM Test Point 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
SmiHandlerTestPointCopyData (
  IN VOID       *InputData,
  IN UINTN      InputDataSize,
  OUT VOID      *DataBuffer,
  IN OUT UINT64 *DataSize,
  IN OUT UINT64 *DataOffset
  )
{
  if (*DataOffset >= InputDataSize) {
    *DataOffset = InputDataSize;
    return;
  }
  if (InputDataSize - *DataOffset < *DataSize) {
    *DataSize = InputDataSize - *DataOffset;
  }

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

/**
  SMM test point SMI handler to get data by offset.

  @param SmiHandlerTestPointParameterGetDataByOffset   The parameter of SMM test point SMI handler get data by offset.

**/
VOID
SmmTestPointSmiHandlerGetDataByOffset (
  IN SMI_HANDLER_TEST_POINT_PARAMETER_GET_DATA_BY_OFFSET     *SmiHandlerTestPointParameterGetDataByOffset
  )
{
  SMI_HANDLER_TEST_POINT_PARAMETER_GET_DATA_BY_OFFSET    SmiHandlerTestPointGetDataByOffset;
  VOID                                                   *Data;
  UINTN                                                  DataSize;
  EFI_STATUS                                             Status;

  Data = NULL;

  CopyMem (
    &SmiHandlerTestPointGetDataByOffset,
    SmiHandlerTestPointParameterGetDataByOffset,
    sizeof(SmiHandlerTestPointGetDataByOffset)
    );

  //
  // Sanity check
  //
  if (!SmmIsBufferOutsideSmmValid((UINTN)SmiHandlerTestPointGetDataByOffset.DataBuffer, (UINTN)SmiHandlerTestPointGetDataByOffset.DataSize)) {
    DEBUG((DEBUG_ERROR, "SmmTestPointSmiHandlerGetDataByOffset: SmmTestPoint get data in SMRAM or overflow!\n"));
    SmiHandlerTestPointParameterGetDataByOffset->Header.ReturnStatus = (UINT64)(INT64)(INTN)EFI_ACCESS_DENIED;
    goto Done;
  }
  
  DataSize = 0;
  Status = GetAllSmmTestPointData (&DataSize, NULL);
  if (Status != EFI_BUFFER_TOO_SMALL) {
    SmiHandlerTestPointParameterGetDataByOffset->Header.ReturnStatus = (UINT64)(INT64)(INTN)EFI_NOT_FOUND;
    goto Done;
  }
  Data = AllocatePool (DataSize);
  if (Data == NULL) {
    SmiHandlerTestPointParameterGetDataByOffset->Header.ReturnStatus = (UINT64)(INT64)(INTN)EFI_OUT_OF_RESOURCES;
    goto Done;
  }
  Status = GetAllSmmTestPointData (&DataSize, Data);
  if (EFI_ERROR(Status)) {
    SmiHandlerTestPointParameterGetDataByOffset->Header.ReturnStatus = (UINT64)(INT64)(INTN)Status;
    goto Done;
  }
  
  SmiHandlerTestPointCopyData (
    Data,
    DataSize,
    (VOID *)(UINTN)SmiHandlerTestPointGetDataByOffset.DataBuffer,
    &SmiHandlerTestPointGetDataByOffset.DataSize,
    &SmiHandlerTestPointGetDataByOffset.DataOffset
    );

  CopyMem (
    SmiHandlerTestPointParameterGetDataByOffset,
    &SmiHandlerTestPointGetDataByOffset,
    sizeof(SmiHandlerTestPointGetDataByOffset)
    );

  SmiHandlerTestPointParameterGetDataByOffset->Header.ReturnStatus = 0;

Done:
  if (Data != NULL) {
    FreePool (Data);
  }
}

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

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

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

  TempCommBufferSize = *CommBufferSize;

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

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

  SmiHandlerTestPointParameterHeader = (SMI_HANDLER_TEST_POINT_PARAMETER_HEADER *)((UINTN)CommBuffer);
  SmiHandlerTestPointParameterHeader->ReturnStatus = (UINT64)-1;

  switch (SmiHandlerTestPointParameterHeader->Command) {
  case SMI_HANDLER_TEST_POINT_COMMAND_GET_INFO:
    DEBUG((DEBUG_ERROR, "SmiHandlerTestPointHandlerGetInfo\n"));
    if (TempCommBufferSize != sizeof(SMI_HANDLER_TEST_POINT_PARAMETER_GET_INFO)) {
      DEBUG((DEBUG_ERROR, "SmmTestPointSmiHandler: SMM communication buffer size invalid!\n"));
      return EFI_SUCCESS;
    }
    SmmTestPointSmiHandlerGetInfo((SMI_HANDLER_TEST_POINT_PARAMETER_GET_INFO *)(UINTN)CommBuffer);
    break;
  case SMI_HANDLER_TEST_POINT_COMMAND_GET_DATA_BY_OFFSET:
    DEBUG((DEBUG_ERROR, "SmiHandlerTestPointHandlerGetDataByOffset\n"));
    if (TempCommBufferSize != sizeof(SMI_HANDLER_TEST_POINT_PARAMETER_GET_DATA_BY_OFFSET)) {
      DEBUG((DEBUG_ERROR, "SmmTestPointSmiHandler: SMM communication buffer size invalid!\n"));
      return EFI_SUCCESS;
    }
    SmmTestPointSmiHandlerGetDataByOffset((SMI_HANDLER_TEST_POINT_PARAMETER_GET_DATA_BY_OFFSET *)(UINTN)CommBuffer);
    break;
  default:
    break;
  }

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

  return EFI_SUCCESS;
}

/**
  Register SMM TestPoint handler.
**/
VOID
RegisterSmmTestPointSmiHandler (
  VOID
  )
{
  EFI_HANDLE                       DispatchHandle;
  EFI_STATUS                       Status;
  STATIC BOOLEAN                   Registered = FALSE;

  if (Registered) {
    return ;
  }
  
  Status = gSmst->SmiHandlerRegister (
                    SmmTestPointSmiHandler,
                    &gAdapterInfoPlatformTestPointGuid,
                    &DispatchHandle
                    );
  ASSERT_EFI_ERROR (Status);
  Registered = TRUE;
}