 /*****************************************************************************
 *
 *
 * Copyright (c) 2012 - 2015, Hefei LCFC Information Technology Co.Ltd. 
 * And/or its affiliates. All rights reserved. 
 * Hefei LCFC Information Technology Co.Ltd. PROPRIETARY/CONFIDENTIAL. 
 * Use is subject to license terms.
 * 
 *****************************************************************************/
/*
Data          Name          Version    Description
2013.12.13    dahai.zhou    v1.00      Initial release it for EDK2 projects
2013.12.16    dahai.zhou    v1.01      Fix the issue that ChgBoot.exe can't change LAN sequence
2013.12.23    dahai.zhou    v1.02      Add Lenovo Fast Boot supported
2014.01.25    dahai.zhou    v1.03      Update UsbLegacy if change boot type
2014.04.22    dahai.zhou    v1.04      Remove the UsbLegacy item since it's for BaytrailM only, and remove some project code
2014.05.21    dahai.zhou    v1.05      Update the OneTimeDisableFastBoot GUID to LFC GUID
2014.06.03    dahai.zhou    v1.06      Remove project dependent code, use library routine GetSwSmiSubFunctionNumber and SetSwSmiSubFunctionNumber instead
2015.01.05    erhu.wang     v1.07      Add debug mode secure flash method, which use a self-defined varible to control
2015.03.26    cissie.gu     v1.08      Add usb charge disable and enable function
2015.05.25    dahai.zhou    v1.09      Support tpmdis and tpmen command
2015.09.15    dahai.zhou    v1.10      Rename LOAD_WIN8_64_DEFAULT_FUNCTION to LOAD_UEFI_OS_DEFAULT_FUNCTION
                                       Support sbrestore command
2015.11.03    dahai.zhou    v1.11      Update ChgBoot load default method to use UEFI variable instead of use LVAR
                                       Add new command SB_RESTORE_FACTORY_DEFAULT_FUNCTION, sbrestore
                                       Must co-work with BIOS design guide V2.19 or newer
2016.01.03    ned.li        v1.12      Add full reset on the next reboot.
2016.02.03    rita.zhang    v1.13      Write disable Secure Boot flag when use dualboot command.
2016.05.25    rita.zhang    v1.14      Update ChgBoot driver in LfcPkg trunk according to the new spec. Add new OutPut status FUNCTION_UNSUPPORTED. 
2016.08.31    Brick.Ji      v1.15      Add /pben and /pbdis command.
2016.09.28    rita.zhang    v1.16      Add /ooen and /oodis command.
2016.09.28    Ned.li        v1.17      Add /sataraid and /sataahci command.
2018.09.29    rita.zhang    v1.18      Remove /pben and /pbdis command.
2018.12.21    wita/rita     v1.19      Add charge in battery mode code to sync with Lenovo Vantage to support AOU.
2019.01.04    rita.zhang    v1.20      Uefios&legacyos command remove delete "Setup" variable behavior.(commands sync to F9 common porting part)
2019.01.14    rita.zhang    v1.21      Add igpu/idgpu/dgpu command.
*/

#include "ChgBootSmm.h"

EFI_STATUS
CheckLcfcSubFunction (
  VOID
)
{
  EFI_STATUS                            Status;
  UINT8                                 SmiData;

  GetSwSmiSubFunctionNumber (&SmiData);  // Read SMI ID

  switch (SmiData) {
    case FUNCTION_ENABLE_HOT_KEY_MODE:
    case FUNCTION_DISABLE_HOT_KEY_MODE:
      {
        UINTN                                 DataSize;
        EFI_SMM_VARIABLE_PROTOCOL             *SmmVariable;
        SYSTEM_CONFIGURATION                  SetupVariable;
        EFI_GUID                              SystemConfigurationGuid  = SYSTEM_CONFIGURATION_GUID;

        Status = gSmst->SmmLocateProtocol (&gEfiSmmVariableProtocolGuid, NULL, &SmmVariable);
        ASSERT_EFI_ERROR (Status);

        DataSize = sizeof (SYSTEM_CONFIGURATION);
        Status = SmmVariable->SmmGetVariable (
                                L"Setup",
                                &SystemConfigurationGuid,
                                NULL,
                                &DataSize,
                                &SetupVariable
                                );
        ASSERT_EFI_ERROR (Status);

        SetupVariable.L05HotKeyMode = (SmiData == FUNCTION_ENABLE_HOT_KEY_MODE)?1:0;
        Status = SmmVariable->SmmSetVariable (
                                L"Setup",
                                &SystemConfigurationGuid,
                                EFI_VARIABLE_BOOTSERVICE_ACCESS | EFI_VARIABLE_RUNTIME_ACCESS | EFI_VARIABLE_NON_VOLATILE,
                                DataSize,
                                &SetupVariable
                                );
        ASSERT_EFI_ERROR (Status);
        Status = EFI_SUCCESS;
        break;
      }
    
    case FUNCTION_ENABLE_USB_CHARGE_MODE:
    case FUNCTION_DISABLE_USB_CHARGE_MODE:
      {
        UINTN                                 DataSize;
        EFI_SMM_VARIABLE_PROTOCOL             *SmmVariable;
        SYSTEM_CONFIGURATION                  SetupVariable;
        EFI_GUID                              SystemConfigurationGuid  = SYSTEM_CONFIGURATION_GUID;

        Status = gSmst->SmmLocateProtocol (&gEfiSmmVariableProtocolGuid, NULL, &SmmVariable);
        ASSERT_EFI_ERROR (Status);

        DataSize = sizeof (SYSTEM_CONFIGURATION);
        Status = SmmVariable->SmmGetVariable (
                                L"Setup",
                                &SystemConfigurationGuid,
                                NULL,
                                &DataSize,
                                &SetupVariable
                                );
        ASSERT_EFI_ERROR (Status);

        SetupVariable.L05AlwaysOnUSB = (SmiData == FUNCTION_ENABLE_USB_CHARGE_MODE)?1:0;
        Status = SmmVariable->SmmSetVariable (
                                L"Setup",
                                &SystemConfigurationGuid,
                                EFI_VARIABLE_BOOTSERVICE_ACCESS | EFI_VARIABLE_RUNTIME_ACCESS | EFI_VARIABLE_NON_VOLATILE,
                                DataSize,
                                &SetupVariable
                                );
        ASSERT_EFI_ERROR (Status);
        Status = EFI_SUCCESS;
        break;
      }

    case FUNCTION_ENABLE_CHARGE_IN_BATTERY_MODE:
    case FUNCTION_DISABLE_CHARGE_IN_BATTERY_MODE:
      {
        UINTN                                 DataSize;
        EFI_SMM_VARIABLE_PROTOCOL             *SmmVariable;
        SYSTEM_CONFIGURATION                  SetupVariable;
        EFI_GUID                              SystemConfigurationGuid  = SYSTEM_CONFIGURATION_GUID;

        Status = gSmst->SmmLocateProtocol (&gEfiSmmVariableProtocolGuid, NULL, &SmmVariable);
        ASSERT_EFI_ERROR (Status);

        DataSize = sizeof (SYSTEM_CONFIGURATION);
        Status = SmmVariable->SmmGetVariable (
                                L"Setup",
                                &SystemConfigurationGuid,
                                NULL,
                                &DataSize,
                                &SetupVariable
                                );
        ASSERT_EFI_ERROR (Status);

        SetupVariable.L05ChargeInBatteryMode = (SmiData == FUNCTION_ENABLE_CHARGE_IN_BATTERY_MODE)?1:0;
        Status = SmmVariable->SmmSetVariable (
                                L"Setup",
                                &SystemConfigurationGuid,
                                EFI_VARIABLE_BOOTSERVICE_ACCESS | EFI_VARIABLE_RUNTIME_ACCESS | EFI_VARIABLE_NON_VOLATILE,
                                DataSize,
                                &SetupVariable
                                );
        ASSERT_EFI_ERROR (Status);
        Status = EFI_SUCCESS;
        break;
      }
		case SET_GFX_IGPU_ONLY_FUNCTION:
    case SET_GFX_IGPU_DGPU_FUNCTION:
		case SET_GFX_DGPU_ONLY_FUNCTION:
      {
        UINTN                                 DataSize;
        EFI_SMM_VARIABLE_PROTOCOL             *SmmVariable;
        SA_SETUP                              SaSetup; 
        
        Status = gSmst->SmmLocateProtocol (&gEfiSmmVariableProtocolGuid, NULL, &SmmVariable);
        ASSERT_EFI_ERROR (Status);
        
        DataSize = sizeof (SA_SETUP);
        Status = SmmVariable->SmmGetVariable (
                                L"SaSetup",
                                &gSaSetupVariableGuid,
                                NULL,
                                &DataSize,
                                &SaSetup
                                );
        ASSERT_EFI_ERROR (Status);
        //SET_GFX_IGPU_ONLY_FUNCTION set PrimaryDisplay = 0;
        //SET_GFX_IGPU_DGPU_FUNCTION set PrimaryDisplay = 4;
        //SET_GFX_DGPU_ONLY_FUNCTION set PrimaryDisplay = 1;
        SaSetup.PrimaryDisplay = (SmiData == SET_GFX_IGPU_ONLY_FUNCTION)?0:((SmiData == SET_GFX_DGPU_ONLY_FUNCTION)?1:4); 
        Status = SmmVariable->SmmSetVariable (
                                L"SaSetup",
                                &gSaSetupVariableGuid,
                                EFI_VARIABLE_BOOTSERVICE_ACCESS | EFI_VARIABLE_RUNTIME_ACCESS | EFI_VARIABLE_NON_VOLATILE,
                                DataSize,
                                &SaSetup
                                );
        ASSERT_EFI_ERROR (Status);
        Status = EFI_SUCCESS;
        break;
      }
    default:
      Status = EFI_UNSUPPORTED;
      break;
  }

  SmiData = 0;
  SetSwSmiSubFunctionNumber (SmiData);  // Clear this SMI data

  return Status;
}

EFI_STATUS
EFIAPI
ChgBootSmiCallback (
  IN EFI_HANDLE     DispatchHandle,
  IN CONST VOID     *Context         OPTIONAL,
  IN OUT VOID       *CommBuffer      OPTIONAL,
  IN OUT UINTN      *CommBufferSize  OPTIONAL
  )
/*
SW SMI handler.
Input:
AX = 0x53ca
BX = Function number
ECX = Physics address low 32bit of BOOT_DEVICE_DATA
EDI = Physics address high 32bit of BOOT_DEVICE_DATA

Output:
AL = Return status
*/
{
  EFI_STATUS                            Status;
  UINTN                                 CpuIndex;
  UINT32                                TempValue;
//  UINT8                                 DefaultBootTypeOrderSequence[] = {L05_BOOT_ORDER_SEQUENCE};  //From L05Project.h
//  BOOT_DEVICE_DATA                      *SmiData = NULL;
//  UINTN                                 SmiDataAddressLow32, SmiDataAddressHigh32;
  EFI_SMM_VARIABLE_PROTOCOL             *SmmVariable;
  SYSTEM_CONFIGURATION                  SetupVariable;
  EFI_GUID                              SystemConfigurationGuid  = SYSTEM_CONFIGURATION_GUID;
  UINTN                                 DataSize;
  UINT8                                 IsNeedUpdateSetup;

  if (CheckLcfcSubFunction () == EFI_SUCCESS) {
    return EFI_SUCCESS;
  }

  Status = IdentifyCpuIndexByEax (BOOT_SEQUENCE_UPDATE_FLAG|BOOT_SEQUENCE_UPDATE_SWSMI, 0xffff, &CpuIndex);
  if (EFI_ERROR (Status)) {
    return EFI_SUCCESS;
  }

//  Status = GetDwordRegisterByCpuIndex (EFI_SMM_SAVE_STATE_REGISTER_RCX, CpuIndex, sizeof (UINT32), &TempValue);
//  ASSERT_EFI_ERROR (Status);
//  SmiDataAddressLow32 = (UINTN)TempValue;
//
//  Status = GetDwordRegisterByCpuIndex (EFI_SMM_SAVE_STATE_REGISTER_RDI, CpuIndex, sizeof (UINT32), &TempValue);
//  ASSERT_EFI_ERROR (Status);
//  SmiDataAddressHigh32 = (UINTN)(((UINT64)TempValue)<<32);
//
//  SmiData = (SMAPI_DATA *)(SmiDataAddressLow32 + SmiDataAddressHigh32);

  Status = gSmst->SmmLocateProtocol (&gEfiSmmVariableProtocolGuid, NULL, &SmmVariable);
  ASSERT_EFI_ERROR (Status);

  Status = GetDwordRegisterByCpuIndex (EFI_SMM_SAVE_STATE_REGISTER_RBX, CpuIndex, sizeof (UINT32), &TempValue);
  ASSERT_EFI_ERROR (Status);

  DataSize = sizeof (SYSTEM_CONFIGURATION);
  Status = SmmVariable->SmmGetVariable (
                          L"Setup",
                          &SystemConfigurationGuid,
                          NULL,
                          &DataSize,
                          &SetupVariable
                          );

  IsNeedUpdateSetup = TRUE;

  switch(TempValue) {
    case SET_DEFAULT_FUNCTION:
    case SET_SECURE_BOOT_ENABLE_FUNCTION:
    case LOAD_UEFI_OS_DEFAULT_FUNCTION:
    case LOAD_LEGACY_OS_DEFAULT_FUNCTION:
    case LCFC_SECURE_FLASH_ENABLE_FUNCTION:
    case SB_RESTORE_FACTORY_DEFAULT_FUNCTION:
    case SET_FULLRESET_FUNCTION:
    case SET_SECURE_BOOT_DISABLE_FUNCTION:
    case SET_PXE_FIRST_FUNCTION:
    case SET_HDD_FIRST_FUNCTION:
    case SET_USB_FIRST_FUNCTION:
    case SET_EFI_PXE_ONE_TIME_FUNCTION:
      IsNeedUpdateSetup = FALSE; 
      break;  
    default:
      if(EFI_ERROR(Status)) {
        IsNeedUpdateSetup = FALSE;
        TempValue = SETUP_NOT_FOUND;
        SetDwordRegisterByCpuIndex (EFI_SMM_SAVE_STATE_REGISTER_RAX, CpuIndex, sizeof (UINT32), &TempValue);
      }
      break;
  }
  
  if(TempValue == SETUP_NOT_FOUND) {
    return EFI_SUCCESS;
  }

  switch ((UINT16)TempValue) {
//    case GET_CURRENT_SEQ_FUNCTION:
//      CopyMem (SmiData->BootTypeOrderSequence, SetupVariable.BootTypeOrder, MAX_BOOT_ORDER_NUMBER);
//      return;
//
    case SET_DEFAULT_FUNCTION:
      {
        UINT8                    ChgBootDefaultTypeData = 0;
        UINTN                    ChgBootDefaultTypeDataSize = sizeof (ChgBootDefaultTypeData);
        SmmVariable->SmmSetVariable (
                       L"ChgBootBootOrderSetDefault",
                       &gLfcVariableGuid,
                       EFI_VARIABLE_BOOTSERVICE_ACCESS | EFI_VARIABLE_RUNTIME_ACCESS | EFI_VARIABLE_NON_VOLATILE,
                       ChgBootDefaultTypeDataSize,
                       &ChgBootDefaultTypeData
                       );
      }
      break;

    case SET_UEFI_FIRST_FUNCTION:
      SetupVariable.BootNormalPriority = 0;
      break;

    case SET_LEGACY_FIRST_FUNCTION:
      SetupVariable.BootNormalPriority = 1;
      break;

    case SET_UEFI_ONLY_FUNCTION:
      //SetupVariable.BootType = 2;  //UEFI only
      {
        UINT8  BootTypeVariable = 2;
        Status = SmmVariable->SmmSetVariable (
                                L"BootType",
                                &gSystemConfigurationGuid,
                                EFI_VARIABLE_BOOTSERVICE_ACCESS | EFI_VARIABLE_RUNTIME_ACCESS | EFI_VARIABLE_NON_VOLATILE,
                                sizeof (BootTypeVariable),
                                &BootTypeVariable
                                );
      }
//    SetupVariable.NetworkProtocol  = 2;  // Special case for LAN boot
      SetupVariable.BootType = 2;  // UEFI only
      break;

    case SET_DUAL_BOOT_FUNCTION:
      //SetupVariable.BootType = 0;  //Dual boot mode
      {
        UINT8  BootTypeVariable = 0;
        Status = SmmVariable->SmmSetVariable (
                                L"BootType",
                                &gSystemConfigurationGuid,
                                EFI_VARIABLE_BOOTSERVICE_ACCESS | EFI_VARIABLE_RUNTIME_ACCESS | EFI_VARIABLE_NON_VOLATILE,
                                sizeof (BootTypeVariable),
                                &BootTypeVariable
                                );
      }
//    SetupVariable.NetworkProtocol  = 3;  // Special case for LAN boot
      SetupVariable.BootType = 0; //Dual boot mode
      SetupVariable.L05SecureBoot    = 0;
      {
        UINT8                    TempVariable = 0;
        UINT32                   TempDataLength = sizeof (UINT8);
        SmmVariable->SmmSetVariable (
                       L"ChgBootChangeLegacy",
                       &gLfcVariableGuid,
                       EFI_VARIABLE_BOOTSERVICE_ACCESS | EFI_VARIABLE_RUNTIME_ACCESS | EFI_VARIABLE_NON_VOLATILE,
                       TempDataLength,
                       &TempVariable
                       );
      }
      break;

    case SET_SECURE_BOOT_ENABLE_FUNCTION:
      // TODO: reset factory default
      {
        UINT8                    TempVariable = 0;
        UINT32                   TempDataLength = sizeof (UINT8);
        SmmVariable->SmmSetVariable (
                       L"ChgBootSecureBootEnable",
                       &gLfcVariableGuid,
                       EFI_VARIABLE_BOOTSERVICE_ACCESS | EFI_VARIABLE_RUNTIME_ACCESS | EFI_VARIABLE_NON_VOLATILE,
                       TempDataLength,
                       &TempVariable
                       );
      }
      break;

    case SET_FAST_BOOT_ENABLE_FUNCTION:
      // Enable this code if Lenovo Fast Boot implemented
      // SetupVariable.QuickBoot = 1;
      SetupVariable.L05FastBoot = 1;
      break;

    case SET_FAST_BOOT_DISABLE_FUNCTION:
      // Enable this code if Lenovo Fast Boot implemented
      // SetupVariable.QuickBoot = 0;
      SetupVariable.L05FastBoot = 0;
      break;

    case SET_PXE_BOOT_ENABLE_FUNCTION:
      SetupVariable.PxeBootToLan = 1;
      break;

    case SET_PXE_BOOT_DISABLE_FUNCTION:
      SetupVariable.PxeBootToLan = 0;
      break;

    case SET_BACK_FLASH_ENABLE_FUNCTION:
      SetupVariable.L05BiosBackFlash = 1;
      break;

    case SET_BACK_FLASH_DISABLE_FUNCTION:
      SetupVariable.L05BiosBackFlash = 0;
      break;

    case SET_TPM_ENABLE_FUNCTION:
      #if (FixedPcdGet8 (PcdPlatformIntelOrAmd) == 1)
        SetupVariable.PTTEnable = 1;
      #else
        SetupVariable.EnableTPM = 2;
      #endif
      break;

    case SET_TPM_DISABLE_FUNCTION:
      #if (FixedPcdGet8 (PcdPlatformIntelOrAmd) == 1)
        SetupVariable.PTTEnable = 0;
      #else
        SetupVariable.EnableTPM = 0;
      #endif
      break;

//    case SET_POWER_BEEP_ENABLE_FUNCTION:
//      SetupVariable.L05PowerBeepFunction = 1;
//      break;

//    case SET_POWER_BEEP_DISABLE_FUNCTION:
//      SetupVariable.L05PowerBeepFunction = 0;
//      break;

    case SET_OS_OPTINIZED_ENABLE_FUNCTION:
      SetupVariable.L05DefaultType = 0;
      break;

    case SET_OS_OPTINIZED_DISABLE_FUNCTION:
      SetupVariable.L05DefaultType = 1;
      break;
// LCFCTODO: You must follow your project setting.
// In Big Core, it should be PCH_SETUP.SataInterfaceMode.
// In Small Core, it should be SetupVariable.SataInterfaceMode.
// case SET_SATA_RAID_FUNCTION:
//   SetupVariable.SataInterfaceMode = 1;
//   break;
// case SET_SATA_AHCI_FUNCTION:
//   SetupVariable.SataInterfaceMode = 0;		
//   break;
    case LOAD_UEFI_OS_DEFAULT_FUNCTION:
    case LOAD_LEGACY_OS_DEFAULT_FUNCTION:
      {
        UINT8                    TempVariable = 0;
        UINT32                   TempDataLength = sizeof (UINT8);

        if (TempValue == LOAD_UEFI_OS_DEFAULT_FUNCTION) {
          SmmVariable->SmmSetVariable (
                         L"ChgBootUefiLoadDefault",
                         &gLfcVariableGuid,
                         EFI_VARIABLE_BOOTSERVICE_ACCESS | EFI_VARIABLE_RUNTIME_ACCESS | EFI_VARIABLE_NON_VOLATILE,
                         TempDataLength,
                         &TempVariable
                         );
        } else {
          SmmVariable->SmmSetVariable (
                         L"ChgBootLegacyLoadDefault",
                         &gLfcVariableGuid,
                         EFI_VARIABLE_BOOTSERVICE_ACCESS | EFI_VARIABLE_RUNTIME_ACCESS | EFI_VARIABLE_NON_VOLATILE,
                         TempDataLength,
                         &TempVariable
                         );
        }
        ASSERT_EFI_ERROR (Status);

        //Delete it to make it load default at next boot
        //DataSize = 0;
        //Status = SmmVariable->SmmSetVariable (
        //                        L"Setup",
        //                        &SystemConfigurationGuid,
        //                        EFI_VARIABLE_BOOTSERVICE_ACCESS | EFI_VARIABLE_RUNTIME_ACCESS | EFI_VARIABLE_NON_VOLATILE,
        //                        DataSize,
        //                        &SetupVariable
        //                        );
        //if (Status == EFI_NOT_FOUND) {
        //  //Already deleted, so skip next code
        //} else {
        //  ASSERT_EFI_ERROR (Status);
        //}
      }
      break;

//    case SET_SEQ_FUNCTION:
//      CopyMem (SetupVariable.BootTypeOrder, SmiData->BootTypeOrderSequence, MAX_BOOT_ORDER_NUMBER);
//      //Clear BootOrder variable to make it re-enumerate all boot device when next boot
//      Status = SmmVariable->SmmSetVariable (
//                              L"BootOrder",
//                              &gEfiGlobalVariableGuid,
//                              EFI_VARIABLE_BOOTSERVICE_ACCESS | EFI_VARIABLE_RUNTIME_ACCESS | EFI_VARIABLE_NON_VOLATILE,
//                              0,
//                              NULL
//                              );
//      break;

    case LCFC_SECURE_FLASH_ENABLE_FUNCTION:
      {
        BOOLEAN     LfcSecureFlashCer = TRUE;
        // It is a BOOLEAN type flag
        // It exist = use ODM itself secure flash ceritfication
        // It doesn't exist = use formal secure flash certification
        // No matter its value
        SmmVariable->SmmSetVariable (
                       L"@Rm",
                       &gLfcVariableGuid,
                       EFI_VARIABLE_BOOTSERVICE_ACCESS | EFI_VARIABLE_RUNTIME_ACCESS | EFI_VARIABLE_NON_VOLATILE,
                       sizeof (LfcSecureFlashCer),
                       &LfcSecureFlashCer
                       );
      }
      break;

    case SB_RESTORE_FACTORY_DEFAULT_FUNCTION:
      {
        UINT8                    TempVariable = 0;
        UINT32                   TempDataLength = sizeof (UINT8);
        SmmVariable->SmmSetVariable (
                       L"ChgBootRestoreFactory",
                       &gLfcVariableGuid,
                       EFI_VARIABLE_BOOTSERVICE_ACCESS | EFI_VARIABLE_RUNTIME_ACCESS | EFI_VARIABLE_NON_VOLATILE,
                       TempDataLength,
                       &TempVariable
                       );
         ASSERT_EFI_ERROR (Status);
      }
      break;

    case SET_FULLRESET_FUNCTION:
      {
        UINT8                    ChgBootDefaultTypeData = 0;
        UINTN                    ChgBootDefaultTypeDataSize = sizeof (ChgBootDefaultTypeData);
        SmmVariable->SmmSetVariable (
                       L"ChgBootFullReset",
                       &gLfcVariableGuid,
                       EFI_VARIABLE_BOOTSERVICE_ACCESS | EFI_VARIABLE_RUNTIME_ACCESS | EFI_VARIABLE_NON_VOLATILE,
                       ChgBootDefaultTypeDataSize,
                       &ChgBootDefaultTypeData
                       );
      }
      break;

    case SET_SECURE_BOOT_DISABLE_FUNCTION:
      {
        UINT8                    TempVariable = 0;
        UINT32                   TempDataLength = sizeof (UINT8);
        SmmVariable->SmmSetVariable (
                       L"ChgBootSecureBootDisable",
                       &gLfcVariableGuid,
                       EFI_VARIABLE_BOOTSERVICE_ACCESS | EFI_VARIABLE_RUNTIME_ACCESS | EFI_VARIABLE_NON_VOLATILE,
                       TempDataLength,
                       &TempVariable
                       );
      }
      break;
      
    case SET_PXE_FIRST_FUNCTION:
      {
        UINT8                    TempVariable = 0;
        UINT32                   TempDataLength = sizeof (UINT8);
        SmmVariable->SmmSetVariable (
                       L"ChgBootSetPxeToFirst",
                       &gLfcVariableGuid,
                       EFI_VARIABLE_BOOTSERVICE_ACCESS | EFI_VARIABLE_RUNTIME_ACCESS | EFI_VARIABLE_NON_VOLATILE,
                       TempDataLength,
                       &TempVariable
                       );
      }      
      break;

    case SET_HDD_FIRST_FUNCTION:
      
      break;
      
    case SET_USB_FIRST_FUNCTION:
      
      break;

    case SET_EFI_PXE_ONE_TIME_FUNCTION:
      {
        UINT8                    TempVariable = 0;
        UINT32                   TempDataLength = sizeof (UINT8);
        SmmVariable->SmmSetVariable (
                       L"ChgBootSetEfiPxeOneTime",
                       &gLfcVariableGuid,
                       EFI_VARIABLE_BOOTSERVICE_ACCESS | EFI_VARIABLE_RUNTIME_ACCESS | EFI_VARIABLE_NON_VOLATILE,
                       TempDataLength,
                       &TempVariable
                       );
      }  
      break;
   case GET_OPTION_FUNCTION:
      {
        UINT16 Data=0x00;
        UINT8  Data1 =0x00;
        UINT32 Data2 =0x00;
        DataSize = sizeof (SYSTEM_CONFIGURATION);
        Status = SmmVariable->SmmGetVariable (
                                L"Setup",
                                &SystemConfigurationGuid,
                                NULL,
                                &DataSize,
                                &SetupVariable
                                );

        Data =(UINT16)(TempValue>>16);       
        Data1 = (((UINT8*)(UINTN)(&SetupVariable))[Data]);
        Data2=(UINT32)Data1;
        SetDwordRegisterByCpuIndex (EFI_SMM_SAVE_STATE_REGISTER_RBX, CpuIndex, sizeof (UINT32), &Data2);
      }  
       
      break;
		case SET_GFX_IGPU_ONLY_FUNCTION:
    case SET_GFX_IGPU_DGPU_FUNCTION:
		case SET_GFX_DGPU_ONLY_FUNCTION:
      {
        UINTN                                 DataSize;
        EFI_SMM_VARIABLE_PROTOCOL             *SmmVariable;
        SA_SETUP                              SaSetup; 
        
        Status = gSmst->SmmLocateProtocol (&gEfiSmmVariableProtocolGuid, NULL, &SmmVariable);
        ASSERT_EFI_ERROR (Status);
        
        DataSize = sizeof (SA_SETUP);
        Status = SmmVariable->SmmGetVariable (
                                L"SaSetup",
                                &gSaSetupVariableGuid,
                                NULL,
                                &DataSize,
                                &SaSetup
                                );
        ASSERT_EFI_ERROR (Status);
        //SET_GFX_IGPU_ONLY_FUNCTION set PrimaryDisplay = 0;
        //SET_GFX_IGPU_DGPU_FUNCTION set PrimaryDisplay = 4;
        //SET_GFX_DGPU_ONLY_FUNCTION set PrimaryDisplay = 1;
        SaSetup.PrimaryDisplay = ((UINT16)TempValue == SET_GFX_IGPU_ONLY_FUNCTION)?0:(((UINT16)TempValue == SET_GFX_DGPU_ONLY_FUNCTION)?1:4); 
        Status = SmmVariable->SmmSetVariable (
                                L"SaSetup",
                                &gSaSetupVariableGuid,
                                EFI_VARIABLE_BOOTSERVICE_ACCESS | EFI_VARIABLE_RUNTIME_ACCESS | EFI_VARIABLE_NON_VOLATILE,
                                DataSize,
                                &SaSetup
                                );
        ASSERT_EFI_ERROR (Status);
        Status = EFI_SUCCESS;
        break;
      }
    default:
      TempValue = FUNCTION_UNSUPPORTED;
      SetDwordRegisterByCpuIndex (EFI_SMM_SAVE_STATE_REGISTER_RAX, CpuIndex, sizeof (UINT32), &TempValue);
      return EFI_SUCCESS;
  }
  
  if(IsNeedUpdateSetup) {
    DataSize = sizeof (SYSTEM_CONFIGURATION);
    Status = SmmVariable->SmmSetVariable (
                            L"Setup",
                            &SystemConfigurationGuid,
                            EFI_VARIABLE_BOOTSERVICE_ACCESS | EFI_VARIABLE_RUNTIME_ACCESS | EFI_VARIABLE_NON_VOLATILE,
                            DataSize,
                            &SetupVariable
                            );
    ASSERT_EFI_ERROR (Status);
  }

  // Need disable next time smart boot
  // Enable this code if Lenovo Fast Boot implemented
  {
    BOOLEAN                             DisableFastBootFlag = 1;
    // It is a BOOLEAN type flag
    // It exist = disable Fast Boot in next boot
    // It doesn't exist = don't
    // No matter its value
    SmmVariable->SmmSetVariable (L"OneTimeDisableFastBoot", &gLfcVariableGuid, EFI_VARIABLE_BOOTSERVICE_ACCESS | EFI_VARIABLE_RUNTIME_ACCESS | EFI_VARIABLE_NON_VOLATILE, sizeof (DisableFastBootFlag), &DisableFastBootFlag);

    //Set CMOS flag also to double confirm disable Fast Boot
    IoWrite8 (LFC_CMOS_INDEX, LFC_FAST_BOOT_PREVIOUS_BOOT_SUCCESS_INDEX);
    IoWrite8 (LFC_CMOS_DATA, FAST_BOOT_BOOT_FAIL);
  }

  TempValue = SUCCESS;
  SetDwordRegisterByCpuIndex (EFI_SMM_SAVE_STATE_REGISTER_RAX, CpuIndex, sizeof (UINT32), &TempValue);

  return EFI_SUCCESS;
}

EFI_STATUS
ChgBootSmmEntry (
  IN EFI_HANDLE           ImageHandle,
  IN EFI_SYSTEM_TABLE     *SystemTable
  )
{
  if (!InSmm ()) {
    return EFI_UNSUPPORTED;
  }

  //Register Software SMI
  {
    EFI_SMM_SW_DISPATCH2_PROTOCOL         *SwDispatch;
    EFI_SMM_SW_REGISTER_CONTEXT           SwContext;
    EFI_HANDLE                            SwHandle;
    EFI_STATUS                            Status;

    Status = gSmst->SmmLocateProtocol (
                      &gEfiSmmSwDispatch2ProtocolGuid,
                      NULL,
                      &SwDispatch
                      );
    ASSERT_EFI_ERROR (Status);

    SwContext.SwSmiInputValue = BOOT_SEQUENCE_UPDATE_SWSMI;
    Status = SwDispatch->Register (
                           SwDispatch,
                           ChgBootSmiCallback,
                           &SwContext,
                           &SwHandle
                           );
    ASSERT_EFI_ERROR (Status);
  }

  return EFI_SUCCESS;
}
