//*****************************************************************************
 //
 //
 // Copyright (c) 2012 - 2020, 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
2014.08.12    dahai.zhou    v1.10      Reconstitution it
2014.08.13    dahai.zhou    v1.11      Fix the issue that find enough space wrong
2014.08.19    dahai.zhou    v1.12      Support AMD platform
                                       Fix spell error
2014.09.12    feng.gu       v1.20      Record set,lock,unlock LenovoVariable actions into BIOS ROM
                                       Support Lvar Tool dump recorded debug datas into a binary file
2014.09.12    feng.gu       v1.21      Fix the issue that compare decoded debug datas in ram with encoded debug datas in BIOS ROM
2015.01.15    feng.gu       v1.22      Use EncodeOrDecodeBuffer and FlashToBiosRom functions uniformly.
2016.12.28    feng.gu       v1.24      Fix the bug that debug region will overflow
2020.03.02    dahai.zhou    v1.25      Fix the hang 0x15 issue from TGL platform, variable attibute mismatch
2020.03.02    feng.gu       v1.26      Modify lock,unlock and delete action log data and add SMI data port for DXE to SMM. 
*/

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

#include <Library/SmmServicesTableLib.h>
#include <Protocol/SmmSwDispatch2.h>
#include <Protocol/SmmBase2.h>
#include <Protocol/SmmCpu.h>
#include <Protocol/SmmVariable.h>
#include <Protocol/VariableWrite.h>
#include <Library/IoLib.h>

#include <Protocol/LenovoVariable.h>
#include <Library/LfcFlashDeviceLib.h>
#include <Library/LfcSwSmiCpuRegisterLib.h>
#include <Guid/LfcVariableGuid.h>
#include <Lfc.h>
#include <Library/OemSvcLfcGetLvarAddress.h>
#include <Library/SmmMemLib.h>
//[-start-220210-Dongxu0040-Modify]//
#include <Library/LfcEcLib.h>
//[-end-220210-Dongxu0040-Modify]//
UINT32 Sub1BaseInRom = 0;
UINT32 Sub1SizeInRom = 0;
UINT32 Sub2BaseInRom = 0;
UINT32 Sub2SizeInRom = 0;
UINT32 DebugBaseInRom = 0;
UINT32 DebugSizeInRom = 0;

EFI_PHYSICAL_ADDRESS                  mApplicationBufferAddress;

EFI_STATUS
FlashToBiosRom (
  IN     UINTN  WriteFlashAddress,
  IN     UINTN  WriteFlashSize,
  IN     UINT8  *Buffer
  )
{ 
  EFI_STATUS                              Status = EFI_SUCCESS;
  Status = LfcFdLibBlockErase (WriteFlashAddress, WriteFlashSize);
  ASSERT_EFI_ERROR (Status);
  Status = LfcFdLibWrite (WriteFlashAddress, &WriteFlashSize, Buffer);
  ASSERT_EFI_ERROR (Status);
  return Status;
}

EFI_STATUS
EncodeOrDecodeBuffer (
  IN     UINT8  *Buffer,
  IN     UINTN  Size,
  IN     UINTN  Header,
  IN     UINT8  RamdomNum
  )
{
  UINT8                                 *BufferPointer;
  UINTN                                 Index;
  BufferPointer = Buffer + Header;
  for (Index = 0; Index < Size - Header; Index++) {
    *(BufferPointer + Index) = (*(BufferPointer + Index))^RamdomNum;
  }
  return EFI_SUCCESS;
}

// Add and save Data about SetVariable, LockVariable, UnlockVariable action to BIOS ROM
EFI_STATUS
LenovoVariableSmmAddDebugData (
  IN     LENOVO_VARIABLE_PRIVATE_DATA  *PrivateData, 
  IN     EFI_GUID                      *VariableGuidName,
  IN     UINT8                         ActionItem,
  IN     UINT32                        DataSize
  )
{   
  LENOVO_DEBUG_RIGION_HEADER              *DebugAddress;
  LENOVO_DEBUG_DATA                       *Next;
  EFI_STATUS                              Status = EFI_SUCCESS;
  UINTN                                   WriteFlashAddress = 0;
  UINTN                                   WriteFlashSize = 0;

  DebugAddress = (LENOVO_DEBUG_RIGION_HEADER *)(PrivateData->DebugBuffer);
  Next = (LENOVO_DEBUG_DATA *)(PrivateData->DebugBuffer + DebugAddress->NextUsable);

  // Save action time
  IoWrite8 (0x70, 0x09);
  *((UINT16 *)(Next->Time)) = 0x2000 + IoRead8 (0x71); // Year
  IoWrite8 (0x70, 0x08);
  *(Next->Time + 2) = IoRead8 (0x71); // Month
  IoWrite8 (0x70, 0x07);
  *(Next->Time + 3) = IoRead8 (0x71); // Date
  IoWrite8 (0x70, 0x4);
  *(Next->Time + 4) = IoRead8 (0x71); // Hour
  IoWrite8 (0x70, 0x02);
  *(Next->Time + 5) = IoRead8 (0x71); // Minite
  IoWrite8 (0x70, 0x00);
  *(Next->Time + 6) = IoRead8 (0x71); // Second

  // Save action type:set:2,lock:6,unlock:7 
  Next->Action = ActionItem;
  Next->Guid = *VariableGuidName;
  Next->DataLength = DataSize;
  SetMem (Next->Reserved, sizeof (Next->Reserved), 0);
  DebugAddress->NextUsable += sizeof (LENOVO_DEBUG_DATA);
  if (DebugAddress->NextUsable >= DebugSizeInRom ) {
    DebugAddress->NextUsable = sizeof (LENOVO_DEBUG_RIGION_HEADER);
  }

  // Encode
  EncodeOrDecodeBuffer (PrivateData->DebugBuffer,
                       DebugSizeInRom,
                       sizeof (LENOVO_DEBUG_RIGION_HEADER),
                       ((LENOVO_VARIABLE_RIGION_HEADER*)PrivateData->VariableBuffer)->RandomNumber);

  // Write Debug data to flash
  WriteFlashAddress = DebugBaseInRom;
  WriteFlashSize = DebugSizeInRom;
  FlashToBiosRom (WriteFlashAddress, WriteFlashSize, PrivateData->DebugBuffer);

  // Decode
  EncodeOrDecodeBuffer (PrivateData->DebugBuffer,
                       DebugSizeInRom,
                       sizeof (LENOVO_DEBUG_RIGION_HEADER),
                       ((LENOVO_VARIABLE_RIGION_HEADER*)PrivateData->VariableBuffer)->RandomNumber);

  return Status;
}

EFI_STATUS
LenovoVariableInitRegionAndBuffer (
  IN     LENOVO_VARIABLE_PRIVATE_DATA   *PrivateData
  )
{
  UINTN                                 WriteFlashAddress;
  UINTN                                 WriteFlashSize;
  UINT8                                 *BufferAddress; 

  BufferAddress = (UINT8*)PrivateData->VariableBuffer;
  SetMem (BufferAddress, Sub1SizeInRom, 0);

  *(BufferAddress) = 'L';
  *(BufferAddress + 1) = 'E';
  *(BufferAddress + 2) = 'N';
  *(BufferAddress + 3) = 'V';

  WriteFlashAddress = Sub1BaseInRom;
  WriteFlashSize = Sub1SizeInRom;

  FlashToBiosRom (WriteFlashAddress, WriteFlashSize, BufferAddress);
  
  WriteFlashAddress = Sub2BaseInRom;
  WriteFlashSize = Sub2SizeInRom;

  FlashToBiosRom (WriteFlashAddress, WriteFlashSize, BufferAddress);
  
  return EFI_SUCCESS;
}

VOID
LenovoVariableSearchVar (
  IN     LENOVO_VARIABLE_PRIVATE_DATA   *PrivateData,
  IN     EFI_GUID                       *VariableGuidName,
  OUT    UINT8                          **VariableAddress
  )
/*
  Search VariableGuidName from PrivateData->VariableBuffer
  If found, VariableSizeAddress point to the variable address
  If not found, VariableSizeAddress = NULL
*/
{
  UINT8                                 *BufferPointer;
  UINTN                                 Index;

  BufferPointer = PrivateData->VariableBuffer + sizeof (LENOVO_VARIABLE_RIGION_HEADER);
  for (Index = 0; Index < ((LENOVO_VARIABLE_RIGION_HEADER*)PrivateData->VariableBuffer)->VariableCount; Index++) {
    if (CompareGuid (VariableGuidName, &(((LENOVO_VARIABLE_VARIABLE *)BufferPointer)->VariableGuidName))) {
      *VariableAddress = BufferPointer;
      return;
    } else {
      BufferPointer = BufferPointer + ((LENOVO_VARIABLE_VARIABLE *)BufferPointer)->VariableDataSize;
      BufferPointer = BufferPointer + sizeof (LENOVO_VARIABLE_VARIABLE) - 1;
    }
  }

  *VariableAddress = NULL;
  return;
}



EFI_STATUS
LenovoVariableWriteBufferToFlash (
  IN     LENOVO_VARIABLE_PRIVATE_DATA   *PrivateData
  )
{
  EFI_STATUS                            Status = EFI_SUCCESS;
  UINTN                                 WriteFlashAddress = 0;
  UINTN                                 WriteFlashSize;

  switch (PrivateData->Region) {

    case LENOVO_VARIABLE_NO_DATA:
      WriteFlashAddress = Sub1BaseInRom;
      break;

    case LENOVO_VARIABLE_USING_REGION_1:
      WriteFlashAddress = Sub2BaseInRom;
      break;

    case LENOVO_VARIABLE_USING_REGION_2:
      WriteFlashAddress = Sub1BaseInRom;
      break;

    default:
      break;
  }

  ((LENOVO_VARIABLE_RIGION_HEADER*)(PrivateData->VariableBuffer))->PriorityFlag++;

  // Encode
  EncodeOrDecodeBuffer (PrivateData->VariableBuffer,
                       Sub1SizeInRom,
                       sizeof (LENOVO_VARIABLE_RIGION_HEADER),
                       ((LENOVO_VARIABLE_RIGION_HEADER*)PrivateData->VariableBuffer)->RandomNumber);

  // Update checksum
  {
    UINT8                                 *BufferPointer;
    UINTN                                 Index;
    UINT16                                Checksum = 0;

    BufferPointer = PrivateData->VariableBuffer + sizeof (LENOVO_VARIABLE_RIGION_HEADER);
    for (Index = 0; Index < Sub1SizeInRom - sizeof (LENOVO_VARIABLE_RIGION_HEADER); Index++) {
      Checksum += *(BufferPointer + Index);
    }
    ((LENOVO_VARIABLE_RIGION_HEADER*)PrivateData->VariableBuffer)->Checksum = Checksum;
  }

  WriteFlashSize = Sub1SizeInRom;

  FlashToBiosRom (WriteFlashAddress, WriteFlashSize, PrivateData->VariableBuffer);

  // Update region
  if (!EFI_ERROR(Status)) {
    PrivateData->Region = (PrivateData->Region == LENOVO_VARIABLE_USING_REGION_1)?LENOVO_VARIABLE_USING_REGION_2:LENOVO_VARIABLE_USING_REGION_1;
  }

  // Decode
  EncodeOrDecodeBuffer (PrivateData->VariableBuffer,
                       Sub1SizeInRom,
                       sizeof (LENOVO_VARIABLE_RIGION_HEADER),
                       ((LENOVO_VARIABLE_RIGION_HEADER*)PrivateData->VariableBuffer)->RandomNumber);

  return Status;
}

EFI_STATUS
LenovoVariableSmmGetVariable (
  IN     LENOVO_VARIABLE_PROTOCOL       *This,
  IN     EFI_GUID                       *VariableGuidName,
  IN OUT UINT32                         *DataSize,
  OUT    VOID                           *Data
  )
{
  LENOVO_VARIABLE_PRIVATE_DATA          *PrivateData;
  UINT32                                VariableSize;
  UINT8                                 *VariableAddress;

  PrivateData = INSTANCE_FROM_LENV_THIS (This);

  if (VariableGuidName == NULL || DataSize == NULL) {
    return EFI_INVALID_PARAMETER;
  }

  if (((LENOVO_VARIABLE_RIGION_HEADER*)PrivateData->VariableBuffer)->VariableCount == 0) {
    return EFI_NOT_FOUND;
  }

  LenovoVariableSearchVar (PrivateData, VariableGuidName, &VariableAddress);

  if (VariableAddress != NULL) {
    VariableSize = ((LENOVO_VARIABLE_VARIABLE *)VariableAddress)->VariableDataSize;
    if (*DataSize < VariableSize) {
      *DataSize = VariableSize;
      return EFI_BUFFER_TOO_SMALL;
    }
    if (Data == NULL) {
      return EFI_INVALID_PARAMETER;
    }
    *DataSize = VariableSize;
    CopyMem (Data, ((LENOVO_VARIABLE_VARIABLE *)VariableAddress)->VariableData, VariableSize);
    return EFI_SUCCESS;
  }

  return EFI_NOT_FOUND;
}

EFI_STATUS
LenovoVariableSmmDeleteVariable (
  IN     LENOVO_VARIABLE_PROTOCOL       *This,
  IN     EFI_GUID                       *VariableGuidName
  )
{
  EFI_STATUS                            Status;
  LENOVO_VARIABLE_PRIVATE_DATA          *PrivateData;
  UINT8                                 *VariableAddress;
  UINT8                                 *ShiftBegin;
  UINTN                                 ShiftCount;

  PrivateData = INSTANCE_FROM_LENV_THIS (This);

  if (VariableGuidName == NULL) {
    return EFI_INVALID_PARAMETER;
  }

  if (((LENOVO_VARIABLE_RIGION_HEADER*)PrivateData->VariableBuffer)->LockFlag == LENOVO_VARIABLE_REGION_LOCKED) {
    return EFI_ACCESS_DENIED;
  }

  LenovoVariableSearchVar (PrivateData, VariableGuidName, &VariableAddress);
  if (VariableAddress == NULL) {
    return EFI_NOT_FOUND;
  }

  if ((((LENOVO_VARIABLE_VARIABLE *)VariableAddress)->VariableAttribute & LENOVO_VARIABLE_ATTRIBUTE_LOCKED) == LENOVO_VARIABLE_ATTRIBUTE_LOCKED) {
    return EFI_WRITE_PROTECTED;
  }

  ShiftBegin = VariableAddress + ((LENOVO_VARIABLE_VARIABLE *)VariableAddress)->VariableDataSize;
  ShiftBegin = ShiftBegin + sizeof (LENOVO_VARIABLE_VARIABLE) - 1;

  ShiftCount = 0 + ((LENOVO_VARIABLE_VARIABLE *)VariableAddress)->VariableDataSize;
  ShiftCount = ShiftCount + sizeof (LENOVO_VARIABLE_VARIABLE) - 1;

  // Begin shift
  while (ShiftBegin <= PrivateData->LastByte) {
    *(ShiftBegin - ShiftCount) = *ShiftBegin;
    ShiftBegin++;
  }

  SetMem (ShiftBegin - ShiftCount, ShiftCount, 0);

  ((LENOVO_VARIABLE_RIGION_HEADER*)PrivateData->VariableBuffer)->VariableCount--;
  PrivateData->LastByte -= ShiftCount;

  Status = LenovoVariableWriteBufferToFlash (PrivateData);

  return Status;
}

EFI_STATUS
LenovoVariableSmmSetVariable (
  IN     LENOVO_VARIABLE_PROTOCOL       *This,
  IN     EFI_GUID                       *VariableGuidName,
  IN     UINT32                         DataSize,
  IN     VOID                           *Data
  )
{
  EFI_STATUS                            Status;
  LENOVO_VARIABLE_PRIVATE_DATA          *PrivateData;

  PrivateData = INSTANCE_FROM_LENV_THIS (This);

  if (VariableGuidName == NULL) {
    return EFI_INVALID_PARAMETER;
  }

  if (((LENOVO_VARIABLE_RIGION_HEADER*)PrivateData->VariableBuffer)->LockFlag == LENOVO_VARIABLE_REGION_LOCKED) {
    return EFI_ACCESS_DENIED;
  }

  // If DataSize = 0, then delete variable, sync with EFI variable
  if (DataSize == 0) {
    Status = LenovoVariableSmmDeleteVariable (This, VariableGuidName);
    LenovoVariableSmmAddDebugData (PrivateData, VariableGuidName, SET_VARIABLE_FUNCTION, DataSize);

    return Status;
  }

  if (Data == NULL) {
    return EFI_INVALID_PARAMETER;
  }

  // If variable exist, need delete it first
  Status = LenovoVariableSmmDeleteVariable (This, VariableGuidName);
  if (Status != EFI_NOT_FOUND && Status != EFI_SUCCESS) {
    return Status;
  }

  // Make sure we have enough space
  // Last + Need - Base > Size, then error
  if ((UINTN)(PrivateData->LastByte) + sizeof (LENOVO_VARIABLE_VARIABLE) + DataSize - (UINTN)(PrivateData->VariableBuffer) > \
      Sub2SizeInRom) {
    return EFI_OUT_OF_RESOURCES;
  }

  // Add a variable
  {
    LENOVO_VARIABLE_VARIABLE  Variable;
    CopyMem (&(Variable.VariableGuidName), VariableGuidName, sizeof (*VariableGuidName));
    Variable.VariableDataSize = DataSize;
    Variable.VariableAttribute = 0;

    CopyMem ((PrivateData->LastByte) + 1, &Variable, sizeof (Variable) - 1);
    CopyMem ((PrivateData->LastByte) + 1 + sizeof (Variable) - 1, Data, DataSize);
    PrivateData->LastByte = PrivateData->LastByte + sizeof (Variable) - 1 + DataSize;
  }
  ((LENOVO_VARIABLE_RIGION_HEADER*)PrivateData->VariableBuffer)->VariableCount++;

  Status = LenovoVariableWriteBufferToFlash (PrivateData);
  LenovoVariableSmmAddDebugData (PrivateData, VariableGuidName, SET_VARIABLE_FUNCTION, DataSize);

  return Status;
}

EFI_STATUS
LenovoVariableSmmLockVariable (
  IN     LENOVO_VARIABLE_PROTOCOL       *This,
  IN     EFI_GUID                       *VariableGuidName
  )
{
  LENOVO_VARIABLE_PRIVATE_DATA          *PrivateData;
  UINT8                                 *VariableAddress;

  PrivateData = INSTANCE_FROM_LENV_THIS (This);

  if (((LENOVO_VARIABLE_RIGION_HEADER*)PrivateData->VariableBuffer)->LockFlag == LENOVO_VARIABLE_REGION_LOCKED) {
    return EFI_ACCESS_DENIED;
  }

  if ( ((LENOVO_VARIABLE_RIGION_HEADER*)PrivateData->VariableBuffer)->VariableCount == 0) {
    return EFI_NOT_FOUND;
  }

  LenovoVariableSearchVar (PrivateData, VariableGuidName, &VariableAddress);
  if (VariableAddress != NULL) {
    ((LENOVO_VARIABLE_VARIABLE *)VariableAddress)->VariableAttribute |= LENOVO_VARIABLE_ATTRIBUTE_LOCKED;
    LenovoVariableSmmAddDebugData (PrivateData, VariableGuidName, LOCK_VARIABLE_FUNCTION, 0);
    return LenovoVariableWriteBufferToFlash (PrivateData);
  } else {
    return EFI_NOT_FOUND;
  }
}

EFI_STATUS
LenovoVariableSmmUnlockVariable (
  IN     LENOVO_VARIABLE_PROTOCOL       *This,
  IN     EFI_GUID                       *VariableGuidName,
  IN     CHAR16                         *Password
  )
{
  LENOVO_VARIABLE_PRIVATE_DATA          *PrivateData;
  UINT8                                 *VariableAddress;
  CHAR16                                *LenovoVariablePassword = LENOVO_VARIABLE_UNLOCK_PASSWORD;
  static UINTN                          FailCount = 0;

  PrivateData = INSTANCE_FROM_LENV_THIS (This);

  if (((LENOVO_VARIABLE_RIGION_HEADER*)PrivateData->VariableBuffer)->LockFlag == LENOVO_VARIABLE_REGION_LOCKED) {
    return EFI_ACCESS_DENIED;
  }

  if (((LENOVO_VARIABLE_RIGION_HEADER*)PrivateData->VariableBuffer)->VariableCount == 0) {
    return EFI_NOT_FOUND;
  }

  LenovoVariableSearchVar (PrivateData, VariableGuidName, &VariableAddress);
  if (VariableAddress != NULL) {
    // Verify Password
    if (FailCount == LENOVO_VARIABLE_UNLOCK_FAIL_MAX) {
      return EFI_SECURITY_VIOLATION;
    }
    if (StrCmp (Password, LenovoVariablePassword) != 0) {
      FailCount++;
      return EFI_SECURITY_VIOLATION;
    }
    ((LENOVO_VARIABLE_VARIABLE *)VariableAddress)->VariableAttribute &= ~LENOVO_VARIABLE_ATTRIBUTE_LOCKED;
    LenovoVariableSmmAddDebugData (PrivateData, VariableGuidName, UNLOCK_VARIABLE_FUNCTION, (UINT32)StrLen (Password));
    return LenovoVariableWriteBufferToFlash (PrivateData);
  } else {
    return EFI_NOT_FOUND;
  }
}

EFI_STATUS
LenovoLockVariableRegion (
  IN     LENOVO_VARIABLE_PROTOCOL       *This
  )
{
  return EFI_UNSUPPORTED;
}

EFI_STATUS
LenovoUnlockVariableRegion (
  IN     LENOVO_VARIABLE_PROTOCOL       *This,
  IN     CHAR16                         *Password
  )
{
  return EFI_UNSUPPORTED;
}

EFI_STATUS
EFIAPI
LenovoVariableDxeToSmm (
  VOID
  )
{
  EFI_STATUS                            Status;
  SMAPI_DATA                            *DxeSmmDataRt = NULL;
  UINTN                                 DataSize = sizeof (SMAPI_DATA*);
  LENOVO_VARIABLE_PROTOCOL              *LenovoVariableInstance;
  EFI_SMM_VARIABLE_PROTOCOL             *SmmVariable;
  UINT32                                Attributes;

  Status = gSmst->SmmLocateProtocol (&gLenovoVariableProtocolGuid, NULL, &LenovoVariableInstance);
  ASSERT_EFI_ERROR (Status);

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

  Status = SmmVariable->SmmGetVariable (
                          LENOVO_VARIABLE_DXE_SMM_DATA_VARIABLE_NAME,
                          &gLfcVariableGuid,
                          &Attributes,
                          &DataSize,
                          &DxeSmmDataRt
                          );

  if (!SmmIsBufferOutsideSmmValid ((UINTN) (DxeSmmDataRt),(UINTN)(DataSize))) {
    DEBUG ((DEBUG_ERROR, "LenovoVariableDxeToSmm: Runtime non-volatile cache buffer in SMRAM or overflow!\n"));
    return  EFI_UNSUPPORTED;
  }
  if (EFI_ERROR (Status) || (Attributes != EFI_VARIABLE_BOOTSERVICE_ACCESS)) {
    return EFI_SUCCESS;
  }

  switch (DxeSmmDataRt->FunctionCode) {
    case GET_VARIABLE_FUNCTION:
      *((EFI_STATUS *)(DxeSmmDataRt->Name)) = LenovoVariableInstance->GetVariable (
                                                           LenovoVariableInstance,
                                                           &(DxeSmmDataRt->VariableGuidName),
                                                           &(DxeSmmDataRt->DataLength),
                                                           DxeSmmDataRt->Data
                                                           );
      break;

    case SET_VARIABLE_FUNCTION:
       *((EFI_STATUS *)(DxeSmmDataRt->Name))  = LenovoVariableInstance->SetVariable (
                                                           LenovoVariableInstance,
                                                           &(DxeSmmDataRt->VariableGuidName),
                                                           DxeSmmDataRt->DataLength,
                                                           DxeSmmDataRt->Data
                                                           );
      break;

    case LOCK_VARIABLE_FUNCTION:
       *((EFI_STATUS *)(DxeSmmDataRt->Name))  = LenovoVariableInstance->LockVariable (
                                                           LenovoVariableInstance,
                                                           &(DxeSmmDataRt->VariableGuidName)
                                                           );
      break;

    case UNLOCK_VARIABLE_FUNCTION:
       *((EFI_STATUS *)(DxeSmmDataRt->Name))  = LenovoVariableInstance->UnlockVariable (
                                                           LenovoVariableInstance,
                                                           &(DxeSmmDataRt->VariableGuidName),
                                                           (CHAR16*)(DxeSmmDataRt->Data)
                                                           );
      break;

    default:
      break;
  }

  // Delete variable after use
  Status = SmmVariable->SmmSetVariable (
                          LENOVO_VARIABLE_DXE_SMM_DATA_VARIABLE_NAME,
                          &gLfcVariableGuid,
                          EFI_VARIABLE_BOOTSERVICE_ACCESS,
                          0,
                          NULL
                          );

  return EFI_SUCCESS;
}

EFI_STATUS
EFIAPI
LenovoVariableSmiCallback (
  IN EFI_HANDLE     DispatchHandle,
  IN CONST VOID     *Context         OPTIONAL,
  IN OUT VOID       *CommBuffer      OPTIONAL,
  IN OUT UINTN      *CommBufferSize  OPTIONAL
  )
/*
SW SMI handler.
Input:
AX = 5380h
BX = Function number
ECX = Physical memory address (SMAPI_DATA)
EDI = Physics address high 32bit of BOOT_DEVICE_DATA

Output:
AH - C Return status
*/
{
  EFI_STATUS                            Status;
  UINTN                                 CpuIndex;
  UINT32                                TempValue;
  SMAPI_DATA                            *SmiData = NULL;
  UINTN                                 SmiDataAddressLow32, SmiDataAddressHigh32;
  LENOVO_VARIABLE_PROTOCOL              *LenovoVariableInstance;
  LENOVO_VARIABLE_PRIVATE_DATA          *PrivateData;
  UINT8                                 *Head;
  CHAR16                                *LenovoVariablePassword = LENOVO_VARIABLE_DUMP_PASSWORD;
  static UINTN                          FailCount = 0;
  UINT8                                 SmiFunctionNum;
  //[-start-220210-Dongxu0040-Modify]//
   UINT8                                ValidTool = 0;  
   //[-end-220210-Dongxu0040-Modify]//
  Status = IdentifyCpuIndexByEax (LENOVO_VARIABLE_FLAG|LENOVO_VARIABLE_CALLBACK, 0xffff, &CpuIndex);

  if (EFI_ERROR (Status)) {

    GetSwSmiSubFunctionNumber (&SmiFunctionNum);
    
    if (SmiFunctionNum == LENOVO_VARIABLE_DXE_TO_SMM_FLAG && IdentifyCpuIndexByEax (LENOVO_VARIABLE_CALLBACK, 0xff, &CpuIndex) == EFI_SUCCESS ){
         //Clear SMI data port
         SetSwSmiSubFunctionNumber (0);
        LenovoVariableDxeToSmm ();

        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 (&gLenovoVariableProtocolGuid, NULL, &LenovoVariableInstance);
  ASSERT_EFI_ERROR (Status);

  Status = GetDwordRegisterByCpuIndex (EFI_SMM_SAVE_STATE_REGISTER_RBX, CpuIndex, sizeof (UINT16), &TempValue);
  ASSERT_EFI_ERROR (Status);
  if (((UINT16)TempValue!=GET_BUFFER_ADDRESS_FUNCTION )&&((UINT16)TempValue!=DUMP_REGION_FUNCTION)){
    if (!SmmIsBufferOutsideSmmValid (
       (UINTN) (SmiDataAddressLow32 + SmiDataAddressHigh32),
       (UINTN)(sizeof(SMAPI_DATA) + SmiData->DataLength -1))) {
         DEBUG ((DEBUG_ERROR, "LenovoVariableSmiCallback: Runtime non-volatile cache buffer in SMRAM or overflow!\n"));
         return  EFI_UNSUPPORTED;
    }
  }
  //[-start-220210-Dongxu0040-Modify]//
  Status = LfcCheckValidTool(&ValidTool);
  if(ValidTool != 0X55){
  return EFI_UNSUPPORTED;
  }
//[-end-220210-Dongxu0040-Modify]//
  switch ((UINT16)TempValue) {
    case GET_VARIABLE_FUNCTION:
      Status = LenovoVariableInstance->GetVariable (
                                         LenovoVariableInstance,
                                         &(SmiData->VariableGuidName),
                                         &(SmiData->DataLength),
                                         SmiData->Data
                                         );
      if (!EFI_ERROR(Status)) {
        TempValue = LVAR_SUCCESS<<8;
      } else if (Status == EFI_BUFFER_TOO_SMALL) {
        TempValue = LVAR_BUFFER_TOO_SMALL<<8;
      } else if (Status == EFI_NOT_FOUND) {
        TempValue = LVAR_NOT_FOUND<<8;
      } else {
        TempValue = LVAR_OTHER_ERROR<<8;
      }
      break;

    case SET_VARIABLE_FUNCTION:
      Status = LenovoVariableInstance->SetVariable (
                                         LenovoVariableInstance,
                                         &(SmiData->VariableGuidName),
                                         SmiData->DataLength,
                                         SmiData->Data
                                         );
      if (!EFI_ERROR(Status)) {
        TempValue = LVAR_SUCCESS<<8;
      } else if (Status == EFI_ACCESS_DENIED) {
        TempValue = LVAR_REGIONE_LOCKED<<8;
      } else if (Status == EFI_NOT_FOUND) {
        TempValue = LVAR_NOT_FOUND<<8;
      } else if (Status == EFI_WRITE_PROTECTED) {
        TempValue = LVAR_VARIABLE_LOCKED<<8;
      } else {
        TempValue = LVAR_OTHER_ERROR<<8;
      }
      break;

    // Keep this for old version tool
    case DELETE_VARIABLE_FUNCTION:
      Status = LenovoVariableInstance->SetVariable (
                                         LenovoVariableInstance,
                                         &(SmiData->VariableGuidName),
                                         0,
                                         SmiData->Data
                                         );
      if (!EFI_ERROR(Status)) {
        TempValue = LVAR_SUCCESS<<8;
      } else if (Status == EFI_ACCESS_DENIED) {
        TempValue = LVAR_REGIONE_LOCKED<<8;
      } else if (Status == EFI_NOT_FOUND) {
        TempValue = LVAR_NOT_FOUND<<8;
      } else if (Status == EFI_WRITE_PROTECTED) {
        TempValue = LVAR_VARIABLE_LOCKED<<8;
      } else {
        TempValue = LVAR_OTHER_ERROR<<8;
      }
      break;

    case LOCK_VARIABLE_FUNCTION:
      Status = LenovoVariableInstance->LockVariable (
                                         LenovoVariableInstance,
                                         &(SmiData->VariableGuidName)
                                         );
      if (!EFI_ERROR(Status)) {
        TempValue = LVAR_SUCCESS<<8;
      } else if (Status == EFI_ACCESS_DENIED) {
        TempValue = LVAR_REGIONE_LOCKED<<8;
      } else if (Status == EFI_NOT_FOUND) {
        TempValue = LVAR_NOT_FOUND<<8;
      } else {
        TempValue = LVAR_OTHER_ERROR<<8;
      }
      break;

    case UNLOCK_VARIABLE_FUNCTION:
      Status = LenovoVariableInstance->UnlockVariable (
                                         LenovoVariableInstance,
                                         &(SmiData->VariableGuidName),
                                         (CHAR16 *)SmiData->Data
                                         );
      if (!EFI_ERROR(Status)) {
        TempValue = LVAR_SUCCESS<<8;
      } else if (Status == EFI_SECURITY_VIOLATION) {
        TempValue = LVAR_PASSWORD_WRONG<<8;
      } else if (Status == EFI_ACCESS_DENIED) {
        TempValue = LVAR_REGIONE_LOCKED<<8;
      } else if (Status == EFI_NOT_FOUND) {
        TempValue = LVAR_NOT_FOUND<<8;
      } else {
        TempValue = LVAR_OTHER_ERROR<<8;
      }
      break;

    case DUMP_REGION_FUNCTION:
      // If SmiData has enough space about 16KB to save data 
      if (SmiData->DataLength == 0x4000) {
        Head = SmiData->Data; 
        PrivateData = INSTANCE_FROM_LENV_THIS (LenovoVariableInstance);

        // Check password in SmiData->Data transmited from input of tool 
        if (FailCount == LENOVO_VARIABLE_UNLOCK_FAIL_MAX) {
          TempValue = LVAR_PASSWORD_WRONG<<8;
		      SetDwordRegisterByCpuIndex (EFI_SMM_SAVE_STATE_REGISTER_RAX, CpuIndex, sizeof (UINT32), &TempValue);
		      return EFI_SUCCESS;
        }
        if (StrCmp ((CHAR16 *)Head, LenovoVariablePassword) != 0) {
          FailCount++;
          TempValue = LVAR_PASSWORD_WRONG<<8;
		      SetDwordRegisterByCpuIndex (EFI_SMM_SAVE_STATE_REGISTER_RAX, CpuIndex, sizeof (UINT32), &TempValue);
		      return EFI_SUCCESS;
        }

        // Copy LenovoVariable region1 and 2 and bedug region to BIOS ROM
        CopyMem (Head, (VOID *)(UINTN)Sub1BaseInRom, Sub1SizeInRom);
        EncodeOrDecodeBuffer (Head,
                             Sub1SizeInRom,
                             sizeof (LENOVO_VARIABLE_RIGION_HEADER),
                             ((LENOVO_VARIABLE_RIGION_HEADER*)PrivateData->VariableBuffer)->RandomNumber);
        Head += Sub1SizeInRom;
        CopyMem (Head, (VOID *)(UINTN)Sub2BaseInRom, Sub2SizeInRom);
        EncodeOrDecodeBuffer (Head,
                             Sub2SizeInRom,
                             sizeof (LENOVO_VARIABLE_RIGION_HEADER),
                             ((LENOVO_VARIABLE_RIGION_HEADER*)PrivateData->VariableBuffer)->RandomNumber);
        Head += Sub2SizeInRom;
        CopyMem (Head, (VOID *)(UINTN)DebugBaseInRom, DebugSizeInRom);
        EncodeOrDecodeBuffer (Head,
                             DebugSizeInRom,
                             sizeof (LENOVO_DEBUG_RIGION_HEADER),
                             ((LENOVO_VARIABLE_RIGION_HEADER*)PrivateData->VariableBuffer)->RandomNumber);

      } else {
        TempValue = LVAR_OTHER_ERROR<<8;
      }
      break;

    case GET_BUFFER_ADDRESS_FUNCTION:
      TempValue = (UINT32)mApplicationBufferAddress;
      Status = SetDwordRegisterByCpuIndex (EFI_SMM_SAVE_STATE_REGISTER_RCX, CpuIndex, sizeof (UINT32), &TempValue);  // Low 32bit
      TempValue = (UINT32)(mApplicationBufferAddress>>32);
      Status = SetDwordRegisterByCpuIndex (EFI_SMM_SAVE_STATE_REGISTER_RDI, CpuIndex, sizeof (UINT32), &TempValue);  // High 32bit
      TempValue = LVAR_SUCCESS<<8;
	  break;

    default:
      TempValue = LVAR_FUNCTION_NUMBER_WRONG<<8;
      break;
  }
  SetDwordRegisterByCpuIndex (EFI_SMM_SAVE_STATE_REGISTER_RAX, CpuIndex, sizeof (UINT32), &TempValue);

  return EFI_SUCCESS;
}

EFI_STATUS
EFIAPI
LenovoVariableSmmEntry (
  IN EFI_HANDLE           ImageHandle,
  IN EFI_SYSTEM_TABLE     *SystemTable
  )
/*
This driver will decode the LENOVO_VARIABLE_REGION_SUBREGION_1 or 2 and copy it to memory
Which sub region is depend on the PriorityFlag field
And install a protocol, later read/write will be in memory
*/
{
  EFI_STATUS                            Status;
  LENOVO_VARIABLE_PRIVATE_DATA          *PrivateData;
  LENOVO_VARIABLE_PRIVATE_DATA          *PrivateDataSmm;
  EFI_PHYSICAL_ADDRESS                  BufferAddressSmm;
  EFI_PHYSICAL_ADDRESS                  DebugAddressSmm;
  LENOVO_VARIABLE_PROTOCOL              *LenovoVariable = NULL;

  if (!InSmm ()) {
    return EFI_UNSUPPORTED;
  }

    Status = OemSvcLfcDxeGetLvarAddress (&Sub1BaseInRom, &Sub1SizeInRom, &Sub2BaseInRom, &Sub2SizeInRom, &DebugBaseInRom, &DebugSizeInRom);
    if (Sub1BaseInRom == 0 || Sub1SizeInRom == 0 || Sub2BaseInRom == 0 || Sub2SizeInRom == 0 || DebugBaseInRom == 0 || DebugSizeInRom == 0) {
      LfcLibLogError (LFC_LOG_LVAR_SMM_ROM_ERROR);
      return EFI_SUCCESS;
    }

  // Prepare buffer for OS application
  mApplicationBufferAddress = 0xFFFFFFFF;
  gBS->AllocatePages (AllocateMaxAddress, EfiReservedMemoryType, APP_BUFFER_SIZE_IN_PAGE, &mApplicationBufferAddress);

  Status = gBS->LocateProtocol (&gLenovoVariableProtocolGuid, NULL, &LenovoVariable);
  if (EFI_ERROR (Status)) {
    return EFI_SUCCESS;
  }
  PrivateData = INSTANCE_FROM_LENV_THIS (LenovoVariable);

  // Prepare private data
  gSmst->SmmAllocatePool (EfiRuntimeServicesData, sizeof (LENOVO_VARIABLE_PRIVATE_DATA), &PrivateDataSmm);
  SetMem (PrivateDataSmm, sizeof (LENOVO_VARIABLE_PRIVATE_DATA), 0);

  // Prepare variable buffer and verify
  BufferAddressSmm = 0xFFFFFFFF;
  gSmst->SmmAllocatePages (AllocateMaxAddress, EfiRuntimeServicesData, Sub1SizeInRom / EFI_PAGE_SIZE, &BufferAddressSmm);
  CopyMem ((VOID *)(UINTN)BufferAddressSmm, PrivateData->VariableBuffer, Sub1SizeInRom);
  PrivateDataSmm->Signature = LENV_SIGNATURE;
  PrivateDataSmm->Region = PrivateData->Region;
  PrivateDataSmm->VariableBuffer = (UINT8 *)(UINTN)BufferAddressSmm;
  PrivateDataSmm->LastByte = PrivateDataSmm->VariableBuffer + (PrivateData->LastByte - PrivateData->VariableBuffer);

  // Encode
  EncodeOrDecodeBuffer (PrivateDataSmm->VariableBuffer,
                       Sub1SizeInRom,
                       sizeof (LENOVO_VARIABLE_RIGION_HEADER),
                       ((LENOVO_VARIABLE_RIGION_HEADER*)PrivateDataSmm->VariableBuffer)->RandomNumber);

  if ((CompareMem (PrivateDataSmm->VariableBuffer, (VOID *)(UINTN)Sub1BaseInRom, Sub1SizeInRom) != 0) && \
      (CompareMem (PrivateDataSmm->VariableBuffer, (VOID *)(UINTN)Sub2BaseInRom, Sub2SizeInRom) != 0)) {
    // Decode
    EncodeOrDecodeBuffer (PrivateDataSmm->VariableBuffer,
                         Sub1SizeInRom,
                         sizeof (LENOVO_VARIABLE_RIGION_HEADER),
                         ((LENOVO_VARIABLE_RIGION_HEADER*)PrivateDataSmm->VariableBuffer)->RandomNumber);
    LenovoVariableWriteBufferToFlash (PrivateDataSmm);
  } else {
    // EncodeOrDecode (PrivateDataSmm);
    EncodeOrDecodeBuffer (PrivateDataSmm->VariableBuffer,
                         Sub1SizeInRom,
                         sizeof (LENOVO_VARIABLE_RIGION_HEADER),
                         ((LENOVO_VARIABLE_RIGION_HEADER*)PrivateDataSmm->VariableBuffer)->RandomNumber);
  }
  // Copy debug data from Dxe phase to SMM RAM
  DebugAddressSmm = 0xFFFFFFFF;
  gSmst->SmmAllocatePages (AllocateMaxAddress, EfiRuntimeServicesData, DebugSizeInRom / EFI_PAGE_SIZE, &DebugAddressSmm);
  CopyMem ((VOID *)(UINTN)DebugAddressSmm, PrivateData->DebugBuffer, DebugSizeInRom);
  PrivateDataSmm->DebugBuffer = (UINT8 *)(UINTN)DebugAddressSmm;
  
  // Encode
  EncodeOrDecodeBuffer (PrivateDataSmm->DebugBuffer,
                        DebugSizeInRom,
                        sizeof (LENOVO_DEBUG_RIGION_HEADER),
                        ((LENOVO_VARIABLE_RIGION_HEADER*)PrivateDataSmm->VariableBuffer)->RandomNumber);
							
  // Flash debug data to BIOS ROM when SMM RAM is different from BIOS ROM
  if (CompareMem (PrivateDataSmm->DebugBuffer, (VOID *)(UINTN)DebugBaseInRom, DebugSizeInRom) != 0) {
    FlashToBiosRom (DebugBaseInRom, DebugSizeInRom, PrivateDataSmm->DebugBuffer);   
  }
  // Decode
  EncodeOrDecodeBuffer (PrivateDataSmm->DebugBuffer,
                       DebugSizeInRom,
                       sizeof (LENOVO_DEBUG_RIGION_HEADER),
                       ((LENOVO_VARIABLE_RIGION_HEADER*)PrivateDataSmm->VariableBuffer)->RandomNumber);
  // Install protocol
  {      
    EFI_HANDLE                            LenovoVariableHandle = NULL;

    PrivateDataSmm->LenovoVariableInstance.GetVariable          = LenovoVariableSmmGetVariable;
    PrivateDataSmm->LenovoVariableInstance.SetVariable          = LenovoVariableSmmSetVariable;
    PrivateDataSmm->LenovoVariableInstance.LockVariable         = LenovoVariableSmmLockVariable;
    PrivateDataSmm->LenovoVariableInstance.UnlockVariable       = LenovoVariableSmmUnlockVariable;
    Status = gSmst->SmmInstallProtocolInterface (
                      &LenovoVariableHandle,
                      &gLenovoVariableProtocolGuid,
                      EFI_NATIVE_INTERFACE,
                      &(PrivateDataSmm->LenovoVariableInstance)
                      );
    ASSERT_EFI_ERROR (Status);
  }

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

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

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

  // Tell DXE protocol that SMI is ready
  {
    UINT8                                 Ready = 1;

    Status = gRT->SetVariable (
                    LENOVO_VARIABLE_SMI_READY_FLAG_VARIABLE_NAME,
                    &gLfcVariableGuid,
                    EFI_VARIABLE_BOOTSERVICE_ACCESS | EFI_VARIABLE_RUNTIME_ACCESS,
                    1,
                    &Ready
                    );
  }

  return EFI_SUCCESS;
}