/** @file

;******************************************************************************
;* Copyright (c) 2016 - 2020, Insyde Software Corp. 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.
;*
;******************************************************************************
;
; Abstract:
;    Create a SW SMI for support ASF chnaged secure boot mode command
;
*/

#include <Library/BaseLib.h>
#include <Library/BaseMemoryLib.h>
#include <Library/MemoryAllocationLib.h>
#include <Library/DebugLib.h>
#include <Library/PcdLib.h>
#include <Library/VariableLib.h>

#include <Library/UefiLib.h>
#include <Library/SmmServicesTableLib.h>
#include <Library/UefiRuntimeServicesTableLib.h>
#include <Library/UefiBootServicesTableLib.h>

#include <Protocol/SmmBase2.h>
#include <Protocol/SmmSwDispatch2.h>
#include <Protocol/SmmCpu.h>
#include <Protocol/SmmReadyToBoot.h>

#include <Guid/AdmiSecureBoot.h>
#include <Guid/DebugMask.h>
#include <Guid/EventGroup.h>
#include <Guid/AsfSecureBootVariable.h>

#include <ChipsetSmiTable.h>

#define H2O_SIGNATURE        SIGNATURE_32 ('O', '2', 'H', '$')

EFI_SMM_CPU_PROTOCOL         *mSmmCpu;
BOOLEAN                      mAtRuntime              = FALSE;


EFI_STATUS
EFIAPI
AsfSmmReadyToBootNotify (
  IN     CONST EFI_GUID                *Protocol,
  IN     VOID                          *Interface,
  IN     EFI_HANDLE                    Handle
  )
{
  mAtRuntime = TRUE;
  return EFI_SUCCESS;
}

STATIC
BOOLEAN
AtRuntime (
  VOID
  )
{
  return mAtRuntime;
}

/**
 This notification function is called when an SMM Mode is invoked through SMI.
 This may happen during RT, so it must be RT safe.

 @param[in]         DispatchHandle      EFI Handle
 @param[in]         Context             Pointer to the EFI_SMM_SW_REGISTER_CONTEXT

 @retval            EFI_SUCCESS         If this function did handle its specific event.
 @retval            other               This function did not need to handle its events.
*/
EFI_STATUS
EFIAPI
AsfSecureBootChangedCallback (
  IN EFI_HANDLE                        DispatchHandle,
  IN CONST EFI_SMM_SW_REGISTER_CONTEXT *Context,
  IN OUT EFI_SMM_SW_CONTEXT            *SwContext,
  IN OUT UINTN                         *CommBufferSize
  )
{
  EFI_STATUS                    Status;
  UINTN                         Index;
  UINT32                        RegisterValue;
  UINTN                         VarSize;
  UINT8                         AsfData;
  UINT8                         SecureBootEnable;


  Status  = EFI_SUCCESS;
  VarSize = 0;
  SecureBootEnable = 0;
  Index = 0;

  if(!AtRuntime()) {
    //
    // Get the pointer to saved CPU regs.
    //
    Status = mSmmCpu->ReadSaveState (
                        mSmmCpu,
                        sizeof (UINT32),
                        EFI_SMM_SAVE_STATE_REGISTER_RBX,
                        SwContext->SwSmiCpuIndex,
                        &RegisterValue
                        );
    for (Index = 0; Index < gSmst->NumberOfCpus; Index++) {
      Status = mSmmCpu->ReadSaveState ( mSmmCpu,
                                        sizeof (UINT32),
                                        EFI_SMM_SAVE_STATE_REGISTER_RBX,
                                        Index,
                                        &RegisterValue
                                        );
      if ((Status == EFI_SUCCESS) && (RegisterValue == H2O_SIGNATURE)) {
        break;
      }
    }

    if (Index == gSmst->NumberOfCpus) {
      //
      // Error out due to CPU not found
      //
      goto Done;
    }


    //
    // Set ASBN
    //
    VarSize = sizeof (UINT8);

    Status = mSmmCpu->ReadSaveState (
                        mSmmCpu,
                        sizeof (UINT32),
                        EFI_SMM_SAVE_STATE_REGISTER_RAX,
                        SwContext->SwSmiCpuIndex,
                        &RegisterValue
                        );
    if (EFI_ERROR (Status)) {
      goto Done;
    }

    AsfData = (UINT8) (RegisterValue >> 8);


    Status = CommonSetVariable (
               ASF_SECURE_BOOT_VARIABLE_NAME,
               &gAsfSecureBootVariableGuid,
               EFI_VARIABLE_NON_VOLATILE | EFI_VARIABLE_BOOTSERVICE_ACCESS | EFI_VARIABLE_RUNTIME_ACCESS,
               sizeof(UINT8),
               &AsfData
               );
    if (EFI_ERROR (Status)) {
      goto Done;
    }




      SecureBootEnable = 0x01;
      Status = CommonSetVariable (
                EFI_ADMINISTER_SECURE_BOOT_NAME,
                &gEfiGenericVariableGuid,
                EFI_VARIABLE_NON_VOLATILE | EFI_VARIABLE_BOOTSERVICE_ACCESS | EFI_VARIABLE_RUNTIME_ACCESS,
                1,
                &SecureBootEnable
                );

  }

Done:
    RegisterValue = (UINT32)Status;
    Status = mSmmCpu->ReadSaveState ( mSmmCpu,
                                      sizeof (UINT32),
                                      EFI_SMM_SAVE_STATE_REGISTER_RAX,
                                      SwContext->SwSmiCpuIndex,
                                      &RegisterValue
                                      );
  return EFI_SUCCESS;
}

/**
 Initializes the SMM Platfrom Driver

 @param[in]         ImageHandle         Pointer to the loaded image protocol for this driver
 @param[in]         SystemTable         Pointer to the EFI System Table

 @retval            EFI_SUCCESS
 @retval            Assert, otherwise.
*/
EFI_STATUS
EFIAPI
AsfSecureBootSmmEntrypoint (
  IN EFI_HANDLE        ImageHandle,
  IN EFI_SYSTEM_TABLE  *SystemTable
  )
{
  EFI_STATUS                                Status;
  EFI_SMM_SW_DISPATCH2_PROTOCOL             *SwDispatch;
  EFI_SMM_SW_REGISTER_CONTEXT               SwContext;
  EFI_HANDLE                                SwHandle;
  VOID                                      *Registration;
  //
  // Get SMM CPU protocol
  //
  Status = gSmst->SmmLocateProtocol (
                    &gEfiSmmCpuProtocolGuid,
                    NULL,
                    (VOID **)&mSmmCpu
                    );
  ASSERT_EFI_ERROR (Status);

  //
  //  Get the Sw dispatch protocol
  //
  Status = gSmst->SmmLocateProtocol (
                    &gEfiSmmSwDispatch2ProtocolGuid,
                    NULL,
                    (VOID**)&SwDispatch
                    );
  ASSERT_EFI_ERROR (Status);

  //
  // Register SMI
  //
  SwContext.SwSmiInputValue = ASF_SECURE_BOOT_SMI;
  Status = SwDispatch->Register (SwDispatch, AsfSecureBootChangedCallback, &SwContext, &SwHandle);
  ASSERT_EFI_ERROR (Status);

  Status = gSmst->SmmRegisterProtocolNotify (
                  &gEdkiiSmmReadyToBootProtocolGuid,
                  AsfSmmReadyToBootNotify,
                  &Registration
                  );

  ASSERT_EFI_ERROR (Status);

  return EFI_SUCCESS;
}