/** @file
  SPEC    : Computrace implement requirement v1.2.doc

;******************************************************************************
;* Copyright (c) 2013 - 2019, Insyde Software Corporation. All Rights Reserved.
;*
;* You may not reproduce, distribute, publish, display, perform, modify, adapt,
;* transmit, broadcast, present, recite, release, license or otherwise exploit
;* any part of this publication in any form, by any means, without the prior
;* written permission of Insyde Software Corporation.
;*
;******************************************************************************
*/

#include "L05SmmComputrace.h"

EFI_STATUS
ReadDwordRegister (
  IN  EFI_SMM_SAVE_STATE_REGISTER       RegisterNum,
  IN  UINTN                             Width,
  IN  UINTN                             CpuNum,
  OUT VOID                              *RegisterData
  )
{
  EFI_STATUS                            Status;
  
  Status = mSmmCpu->ReadSaveState (
                      mSmmCpu,
                      Width,
                      RegisterNum,
                      CpuNum,
                      RegisterData
                      );
  return Status;
}

EFI_STATUS
WriteDwordRegister (
  IN  EFI_SMM_SAVE_STATE_REGISTER       RegisterNum,
  IN  UINTN                             Width,
  IN  UINTN                             CpuNum,
  IN  VOID                              *RegisterData
  )
{
  EFI_STATUS                            Status;
  
  Status = mSmmCpu->WriteSaveState (
                      mSmmCpu,
                      Width,
                      RegisterNum,
                      CpuNum,
                      RegisterData
                      );
  return Status;
}

EFI_STATUS
UpdateL05ComputraceInfo (
  IN OUT  EFI_L05_COMPUTRACE_AREA       *L05ComputraceArea
  )
{
  if ((FdmGetNAtAddr (&gL05H2OFlashMapRegionComputraceGuid, 1) == (UINTN) L05_INVALID_VALUE) ||
      (FdmGetNAtSize (&gL05H2OFlashMapRegionComputraceGuid, 1) == (UINTN) L05_INVALID_VALUE)) {
    return EFI_UNSUPPORTED;
  }

  CopyMem (L05ComputraceArea, (VOID *) (UINTN) FdmGetNAtAddr (&gL05H2OFlashMapRegionComputraceGuid, 1), sizeof (EFI_L05_COMPUTRACE_AREA));

  return EFI_SUCCESS;
}

EFI_STATUS
SaveL05ComputraceInfo (
  IN  EFI_L05_COMPUTRACE_AREA           *L05ComputraceArea
  )
{
  EFI_SMM_FW_BLOCK_SERVICE_PROTOCOL     *SmmFwBlockService;
  EFI_STATUS                            Status;
  UINTN                                 DataLength;
  UINTN                                 ComputraceRegionBase;
  UINTN                                 ComputraceRegionSize;

  SmmFwBlockService    = NULL;
  Status               = EFI_SUCCESS;
  DataLength           = 0;
  ComputraceRegionBase = (UINTN) FdmGetNAtAddr (&gL05H2OFlashMapRegionComputraceGuid, 1);
  ComputraceRegionSize = (UINTN) FdmGetNAtSize (&gL05H2OFlashMapRegionComputraceGuid, 1);

  if ((ComputraceRegionBase  == (UINTN) L05_INVALID_VALUE) || (ComputraceRegionSize == (UINTN) L05_INVALID_VALUE)) {
    return EFI_UNSUPPORTED;
  }

  Status = gSmst->SmmLocateProtocol (&gEfiSmmFwBlockServiceProtocolGuid, NULL, &SmmFwBlockService);

  if (EFI_ERROR (Status)) {
    return Status;
  }

  DataLength = MIN (ComputraceRegionSize, sizeof (EFI_L05_COMPUTRACE_AREA));

  Status = SmmFwBlockService->EraseBlocks (SmmFwBlockService, ComputraceRegionBase, &ComputraceRegionSize);

  if (EFI_ERROR (Status)) {
    return Status;
  }

  return SmmFwBlockService->Write (SmmFwBlockService, ComputraceRegionBase, &DataLength, (UINT8 *) L05ComputraceArea);
}

EFI_STATUS
LoadL05ComputraceInfo (
  IN OUT  EFI_L05_COMPUTRACE_AREA       *L05ComputraceArea
  )
{
  EFI_STATUS                            Status;
  UINT32                                ComputraceRegionBase;
  UINT32                                ComputraceRegionSize;

  Status               = EFI_SUCCESS;
  ComputraceRegionBase = (UINT32) FdmGetNAtAddr (&gL05H2OFlashMapRegionComputraceGuid, 1);
  ComputraceRegionSize = (UINT32) FdmGetNAtSize (&gL05H2OFlashMapRegionComputraceGuid, 1);

  if ((ComputraceRegionBase == (UINTN) L05_INVALID_VALUE) || (ComputraceRegionSize == (UINTN) L05_INVALID_VALUE)) {
    return EFI_UNSUPPORTED;
  }

  CopyMem (L05ComputraceArea, (VOID *) (UINTN) ComputraceRegionBase, sizeof (EFI_L05_COMPUTRACE_AREA));

  return Status;
}

EFI_STATUS
SetComputraceDisableKey (
  IN     UINT32                         *EfiL05ComputraceDisableKey
  )
{
  EFI_STATUS                            Status;
  EFI_L05_COMPUTRACE_AREA               L05ComputraceArea;

  Status = EFI_SUCCESS;
  
  Status = LoadL05ComputraceInfo (&L05ComputraceArea);

  if (EFI_ERROR (Status)) {
    return Status;
  }

  L05ComputraceArea.DisableKey = *EfiL05ComputraceDisableKey ^ L05ComputraceArea.Mask;

  return SaveL05ComputraceInfo (&L05ComputraceArea);
}

EFI_STATUS
GetComputraceDisableKey (
  IN OUT UINT32                         *EfiL05ComputraceDisableKey
  )
{
  EFI_STATUS                            Status;
  EFI_L05_COMPUTRACE_AREA               L05ComputraceArea;

  Status = EFI_SUCCESS;

  Status = LoadL05ComputraceInfo (&L05ComputraceArea);

  if (EFI_ERROR (Status)) {
    return Status;
  }

  *EfiL05ComputraceDisableKey = L05ComputraceArea.DisableKey ^ L05ComputraceArea.Mask;

  return Status;
}

EFI_STATUS
EFIAPI
L05SmmComputraceGetSetStateCallback (
  IN  UINTN                             CpuNum
  )
{
  EFI_STATUS                            Status;
  UINT32                                Eax;
  UINT32                                Ebx;
  EFI_L05_COMPUTRACE_AREA               L05ComputraceArea;
  UINT32                                DisableKey;
  UINT32                                RegisterBuf;
  UINT32                                Index;
  static UINT32                         TrialCount = OEM_L05_COMPUTRACE_MAX_RETRY;
  
  DisableKey  = 0;
  Index       = 0;
  Eax         = 0;
  Ebx         = 0;
  RegisterBuf = OEM_L05_COMPUTRACE_SET_STATE_NOT_SUPPORTED;

  //
  // Check function signature is valid or not
  //
  ReadDwordRegister (EFI_SMM_SAVE_STATE_REGISTER_RAX, sizeof (UINT32), CpuNum, &Eax);

  if ((Eax != OEM_L05_COMPUTRACE_SET_STATE_ENABLE_FUNCTION)  &&
      (Eax != OEM_L05_COMPUTRACE_SET_STATE_DISABLE_FUNCTION) &&
      (Eax != OEM_L05_COMPUTRACE_GET_STATE_FUNCTION)) {
    return EFI_UNSUPPORTED;
  }

  Status = UpdateL05ComputraceInfo (&L05ComputraceArea);

  if (EFI_ERROR (Status)) {
    RegisterBuf = OEM_L05_COMPUTRACE_SET_STATE_NOT_SUPPORTED;
    WriteDwordRegister (EFI_SMM_SAVE_STATE_REGISTER_RAX, sizeof (UINT32), CpuNum, &RegisterBuf);
    return EFI_SUCCESS;
  }

  ReadDwordRegister (EFI_SMM_SAVE_STATE_REGISTER_RAX, sizeof (UINT32), CpuNum, &Eax);

  switch (Eax) {

  case OEM_L05_COMPUTRACE_GET_STATE_FUNCTION:

    switch (L05ComputraceArea.ComputraceState) {

    case EFI_L05_COMPUTRACE_ENABLING:
    case EFI_L05_COMPUTRACE_ENABLED:
      RegisterBuf = OEM_L05_COMPUTRACE_GET_STATE_ENABLED;
      WriteDwordRegister (EFI_SMM_SAVE_STATE_REGISTER_RAX, sizeof (UINT32), CpuNum, &RegisterBuf);
      break;

    case EFI_L05_COMPUTRACE_DISABLED:
      RegisterBuf = OEM_L05_COMPUTRACE_GET_STATE_DISABLED;
      WriteDwordRegister (EFI_SMM_SAVE_STATE_REGISTER_RAX, sizeof (UINT32), CpuNum, &RegisterBuf);
      break;

    default:
      RegisterBuf = OEM_L05_COMPUTRACE_SET_STATE_NOT_SUPPORTED;
      break;
    }

    break;

  case OEM_L05_COMPUTRACE_SET_STATE_ENABLE_FUNCTION:
    switch (L05ComputraceArea.ComputraceState) {

    case EFI_L05_COMPUTRACE_ENABLING:
    case EFI_L05_COMPUTRACE_ENABLED:
      RegisterBuf = OEM_L05_COMPUTRACE_SET_STATE_ALREADY_ENABLED;
      WriteDwordRegister (EFI_SMM_SAVE_STATE_REGISTER_RAX, sizeof (UINT32), CpuNum, &RegisterBuf);
      break;

    case EFI_L05_COMPUTRACE_DISABLED:
      L05ComputraceArea.Mask = L05_COMPUTRACE_DEFAULT_MASK;
      L05ComputraceArea.ComputraceState = EFI_L05_COMPUTRACE_ENABLING;
      Status = SaveL05ComputraceInfo (&L05ComputraceArea);

      if (EFI_ERROR (Status)) {
        RegisterBuf = OEM_L05_COMPUTRACE_SET_STATE_NOT_SUPPORTED;
        WriteDwordRegister (EFI_SMM_SAVE_STATE_REGISTER_RAX, sizeof (UINT32), CpuNum, &RegisterBuf);
        return EFI_SUCCESS;
      }

      ReadDwordRegister (EFI_SMM_SAVE_STATE_REGISTER_RBX, sizeof (UINT32), CpuNum, &DisableKey);
      Status = SetComputraceDisableKey (&DisableKey);

      if (EFI_ERROR (Status)) {
        RegisterBuf = OEM_L05_COMPUTRACE_SET_STATE_NOT_SUPPORTED;
        WriteDwordRegister (EFI_SMM_SAVE_STATE_REGISTER_RAX, sizeof (UINT32), CpuNum, &RegisterBuf);
        return EFI_SUCCESS;
      }

      RegisterBuf = OEM_L05_COMPUTRACE_SET_STATE_SUCCESS;
      WriteDwordRegister (EFI_SMM_SAVE_STATE_REGISTER_RAX, sizeof (UINT32), CpuNum, &RegisterBuf);
      break;

    default:
      RegisterBuf = OEM_L05_COMPUTRACE_SET_STATE_NOT_SUPPORTED;
      WriteDwordRegister (EFI_SMM_SAVE_STATE_REGISTER_RAX, sizeof (UINT32), CpuNum, &RegisterBuf);
      break;
    }

    break;

  case OEM_L05_COMPUTRACE_SET_STATE_DISABLE_FUNCTION:

    //
    //  By L05 spec, Disable Computrace must has trial count to prevent brute-force attack.
    //
    if (TrialCount == 0) {
      RegisterBuf = OEM_L05_COMPUTRACE_SET_STATE_NOT_SUPPORTED;
      WriteDwordRegister (EFI_SMM_SAVE_STATE_REGISTER_RAX, sizeof (UINT32), CpuNum, &RegisterBuf);
      return EFI_SUCCESS;
    }

    switch (L05ComputraceArea.ComputraceState) {

    case EFI_L05_COMPUTRACE_ENABLING:
    case EFI_L05_COMPUTRACE_ENABLED:
      Status = GetComputraceDisableKey (&DisableKey);

      if (EFI_ERROR (Status)) {
        RegisterBuf = OEM_L05_COMPUTRACE_SET_STATE_NOT_SUPPORTED;
        WriteDwordRegister (EFI_SMM_SAVE_STATE_REGISTER_RAX, sizeof (UINT32), CpuNum, &RegisterBuf);
        return EFI_SUCCESS;
      }


      ReadDwordRegister (EFI_SMM_SAVE_STATE_REGISTER_RBX, sizeof (UINT32), CpuNum, &Ebx);

      if (DisableKey != Ebx) {
        RegisterBuf = OEM_L05_COMPUTRACE_SET_STATE_KEY_NOT_MATCH;
        WriteDwordRegister (EFI_SMM_SAVE_STATE_REGISTER_RAX, sizeof (UINT32), CpuNum, &RegisterBuf);
        TrialCount--;
        return EFI_SUCCESS;
      }

      //
      //  Clear All Computrace Area when disabling.
      //
      for (Index = 0; Index < sizeof (EFI_L05_COMPUTRACE_AREA); Index++) {
        ((UINT8 *) &L05ComputraceArea)[Index] = 0xFF;
      }

      L05ComputraceArea.ComputraceState = EFI_L05_COMPUTRACE_DISABLED;
      Status = SaveL05ComputraceInfo (&L05ComputraceArea);

      if (EFI_ERROR (Status)) {
        RegisterBuf = OEM_L05_COMPUTRACE_SET_STATE_NOT_SUPPORTED;
        WriteDwordRegister (EFI_SMM_SAVE_STATE_REGISTER_RAX, sizeof (UINT32), CpuNum, &RegisterBuf);
        return EFI_SUCCESS;
      }

      RegisterBuf = OEM_L05_COMPUTRACE_SET_STATE_SUCCESS;
      WriteDwordRegister (EFI_SMM_SAVE_STATE_REGISTER_RAX, sizeof (UINT32), CpuNum, &RegisterBuf);
      break;

    case EFI_L05_COMPUTRACE_DISABLED:
      RegisterBuf = OEM_L05_COMPUTRACE_SET_STATE_ALREADY_DISABLED;
      WriteDwordRegister (EFI_SMM_SAVE_STATE_REGISTER_RAX, sizeof (UINT32), CpuNum, &RegisterBuf);
      break;

    default:
      RegisterBuf = OEM_L05_COMPUTRACE_SET_STATE_NOT_SUPPORTED;
      WriteDwordRegister (EFI_SMM_SAVE_STATE_REGISTER_RAX, sizeof (UINT32), CpuNum, &RegisterBuf);
      break;
    }

    break;


  default:
    RegisterBuf = OEM_L05_COMPUTRACE_SET_STATE_NOT_SUPPORTED;
    WriteDwordRegister (EFI_SMM_SAVE_STATE_REGISTER_RAX, sizeof (UINT32), CpuNum, &RegisterBuf);
    break;
  }

  return EFI_SUCCESS;
}

EFI_STATUS
L05SmmComputraceEntryPoint (
  IN EFI_HANDLE                         ImageHandle,
  IN EFI_SYSTEM_TABLE                   *SystemTable
  )
{
  EFI_STATUS                            Status;
  BOOLEAN                               ComputraceSupported;
  EFI_L05_SMM_SW_SMI_INTERFACE_PROTOCOL *L05SwSmiPtr;

  Status              = EFI_SUCCESS;
  ComputraceSupported = TRUE;
  L05SwSmiPtr         = NULL;

  //
  // Skip Computrace init if platform doesn't support.
  //
  Status = OemSvcCheckComputraceSupportStatus (&ComputraceSupported);

  if (Status == EFI_MEDIA_CHANGED && !ComputraceSupported) {
    return EFI_SUCCESS;
  }

  //
  // Locate Smm Cpu protocol for Cpu save state manipulation
  //
  Status = gSmst->SmmLocateProtocol (
                    &gEfiSmmCpuProtocolGuid,
                    NULL,
                    &mSmmCpu
                    );

  if (EFI_ERROR (Status)) {
    return Status;
  }

  //
  // Register Software SMI function with L05SwSmi Driver
  //
  Status = gSmst->SmmLocateProtocol (
                        &gEfiL05SmmSwSmiInterfaceProtocolGuid,
                        NULL,
                        &L05SwSmiPtr
                        );

  if (EFI_ERROR (Status)) {
    return Status;
  }

  Status = L05SwSmiPtr->RegisterCallbackFunction (
                          L05SwSmiPtr,
                          L05_SECURITY_SW_SMI,
                          FeatureCallbackType,
                          L05SmmComputraceGetSetStateCallback
                          );

  if (EFI_ERROR (Status)) {
    return Status;
  }

  return Status;
}