/** @file
  BIOS Self-Healing DXE Driver.

;******************************************************************************
;* Copyright (c) 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.
;*
;******************************************************************************
*/

#include "BiosSelfHealingDxe.h"
//#[-start-220627-dennis0018-add]//
#ifdef LCFC_SUPPORT
#include <Protocol/LenovoVariable.h>
#endif
//#[-end-220627-dennis0018-add]//


EFI_RESET_SYSTEM                        mOriginalResetSystemPtr = NULL;
L05_BIOS_SELF_HEALING_UI_TYPE           mUiType                 = 0;
BOOLEAN                                 mTimeExpired            = FALSE;
BOOLEAN                                 mSpiBackupIsUpdated     = FALSE;

/**
  Compares two memory buffers of a given length.

  @param  DestinationBuffer             The first memory buffer.
  @param  SourceBuffer                  The second memory buffer.
  @param  Length                        Length of DestinationBuffer and SourceBuffer memory
                                        regions to compare. Must be non-zero.

  @return 0                             All Length bytes of the two buffers are identical.
  @retval Non-zero                      The first mismatched byte in SourceBuffer subtracted from the first
                                        mismatched byte in DestinationBuffer.
**/
INTN
EFIAPI
InternalCompareMem64 (
  IN  VOID                              *DestinationBuffer,
  IN  VOID                              *SourceBuffer,
  IN  UINTN                             Length
  )
{
  UINTN                                 CompareLength;

  CompareLength = Length - sizeof (INT64);
  while ((CompareLength != 0) &&
         (*(INT64*)DestinationBuffer == *(INT64*)SourceBuffer)) {
    DestinationBuffer = (VOID*)((UINTN)DestinationBuffer + sizeof (INT64));
    SourceBuffer = (VOID*)((UINTN)SourceBuffer + sizeof (INT64));
    CompareLength -= sizeof (INT64);
  }
  return (INTN)*(UINT64*)DestinationBuffer - (INTN)*(UINT64*)SourceBuffer;
}

/**
  ResetSystem dummy function.

  @param  ResetType                     UEFI defined reset type.
  @param  ResetStatus                   The status code for the reset.
  @param  DataSize                      The size of ResetData in bytes.
  @param  ResetData                     Optional element used to introduce a platform specific reset.
                                        The exact type of the reset is defined by the EFI_GUID that follows
                                        the Null-terminated Unicode string.
**/
VOID
EFIAPI
DummyResetSystem (
  IN  EFI_RESET_TYPE                    ResetType,
  IN  EFI_STATUS                        ResetStatus,
  IN  UINTN                             DataSize,
  IN  VOID                              *ResetData OPTIONAL
  )
{
  return;
}

/**
  ResetSystem hook function.

  @param  ToHook                        Point out to hook or not.
**/
VOID
HookResetSystem (
  IN  BOOLEAN                           ToHook
  )
{
  if (ToHook) {
    mOriginalResetSystemPtr = gRT->ResetSystem;
    gRT->ResetSystem = DummyResetSystem;
  } else {
    gRT->ResetSystem = mOriginalResetSystemPtr;
  }

  return;
}

/**
  Update blocks on flash.

  @param  StartAddress                  Start address to be written.
  @param  Buffer                        Data buffer want to write.
  @param  BufferSize                    The requested write size.
  @param  UseProtectTable               Need to confirm the protected area or not.
  @param  TotalFlashBlocks              Total number of flashing blocks. If non-zero, it means that multi region are currently flashing.
  @param  CurrentFlashBlocks            Number of blocks currently flashed. If non-zero, it means that multi region are currently flashing.
  @param  IsMultiRegion                 Multi region flashing.

  @retval EFI_SUCCESS                   Update BIOS successfully.
  @retval Others                        An unexpected error occurred.
**/
EFI_STATUS
FlashUpdate (
  IN  UINTN                             StartAddress,
  IN  UINTN                             Buffer,
  IN  UINTN                             BufferSize,
  IN  BOOLEAN                           UseProtectTable,
  IN  UINTN                             TotalFlashBlocks, OPTIONAL
  IN  UINTN                             CurrentFlashBlocks, OPTIONAL
  IN  BOOLEAN                           IsMultiRegion OPTIONAL
  )
{
  EFI_STATUS                            Status;
  EFI_L05_FLASH_PROTECT_REGION          *L05ProtectTable;
  UINTN                                 TableCount;
  UINTN                                 BlockSize;
  UINTN                                 RecoveryBlocks;
  UINTN                                 WriteAddress;
  UINTN                                 SrcAddress;
  UINTN                                 WriteCount;
  UINTN                                 Size;
  UINTN                                 Retry;
  UINTN                                 FlashPrecentage;
  UINTN                                 Index;
  UINT8                                 *CompareBuffer;

  Status           = EFI_SUCCESS;
  L05ProtectTable  = NULL;
  TableCount       = 0;
  BlockSize        = FLASH_SECTOR_SIZE;
  RecoveryBlocks   = (BufferSize / BlockSize);
  TotalFlashBlocks = IsMultiRegion ? TotalFlashBlocks : RecoveryBlocks;
  WriteAddress     = StartAddress;
  SrcAddress       = Buffer;
  WriteCount       = 0;
  Size             = 0;
  Retry            = 0;
  FlashPrecentage  = 0;
  CompareBuffer    = NULL;

  Status = GetFlashProtectTable (&L05ProtectTable, &TableCount, SelfHealing);

  if (EFI_ERROR (Status) || (L05ProtectTable == NULL) || (TableCount == 0)) {
    UseProtectTable = FALSE;
  }

  CompareBuffer = AllocatePool (BlockSize);

  if (CompareBuffer == NULL) {
    return EFI_OUT_OF_RESOURCES;
  }

  HookResetSystem (TRUE);

  for (WriteCount = 0; WriteCount < RecoveryBlocks; WriteCount++) {
    //
    // Confirm whether the region needs protection
    //
    if (UseProtectTable) {

      for (Index = 0; Index < TableCount; Index++) {

        if ((WriteAddress >= L05ProtectTable[Index].LinearAddress) &&
            (WriteAddress < (L05ProtectTable[Index].LinearAddress + L05ProtectTable[Index].Size))) {

          if (WriteAddress == L05ProtectTable[Index].LinearAddress) {
            DEBUG ((EFI_D_INFO, "BiosSelfHealingDxe: Skip protected area update => address: %x, size: %x\n", L05ProtectTable[Index].LinearAddress, L05ProtectTable[Index].Size));
          }
          goto Skip;
        }
      }
    }

    //
    // Update region
    //
    for (Retry = 0; Retry < FLASH_FAILURE_RETRY_COUNT; Retry++) {

      Status = FlashErase (WriteAddress, BlockSize);
      if (EFI_ERROR (Status)) {
        DEBUG ((EFI_D_ERROR, "Erase flash address 0x%x error %r\n", WriteAddress, Status));
        continue;
      }

      Size = BlockSize;
      Status = FlashProgram (
                 (UINT8 *) WriteAddress,
                 (UINT8 *) SrcAddress,
                 &Size,
                 WriteAddress
                 );
      if (EFI_ERROR (Status)) {
        DEBUG ((EFI_D_ERROR, "Program flash address 0x%x error %r\n", WriteAddress));
        continue;
      }

      ZeroMem (CompareBuffer, BlockSize);
      Status = FlashRead (
                 CompareBuffer,
                 (UINT8 *) WriteAddress,
                 BlockSize
                 );
      if (EFI_ERROR (Status)) {
        DEBUG ((EFI_D_ERROR, "Read flash address 0x%x error %r\n", WriteAddress, Status));
        continue;
      }

      if (CompareMem (CompareBuffer, (VOID *) SrcAddress, BlockSize) == 0) {
        break;
      }
    }

Skip:
    FlashPrecentage = ((WriteCount + CurrentFlashBlocks + 1) * PROCESS_ACCURACY_VALUE) / TotalFlashBlocks;
    BiosSelfHealingProgressBar (FlashPrecentage);

    WriteAddress += BlockSize;
    SrcAddress   += BlockSize;
  }

  HookResetSystem (FALSE);

  SafeFreePool ((VOID **) &CompareBuffer);
  SafeFreePool ((VOID **) &L05ProtectTable);

  return Status;
}

/**
  Reset system event.

  @param  Event                         Event whose notification function is being invoked.
  @param  Context                       Pointer to the notification function!|s context, which is
                                        implementation-dependent.  Context corresponds to
                                        NotifyContext in CreateEventEx().
**/
VOID
EFIAPI
ResetSystemTimerEvent (
  IN  EFI_EVENT                         Event,
  IN  VOID                              *Context
  )
{
  mTimeExpired = TRUE;

  if (mUiType == BshRecoveryDoneReset) {
    gRT->ResetSystem (EfiResetPlatformSpecific, EFI_SUCCESS, 0, NULL);
    CpuDeadLoop ();
  }
}

/**
  Trigger disable Top Swap.

  @param  None

  @retval EFI_SUCCESS                   This function execute successfully.
  @retval Orther                        An unexpected error occurred.
**/
EFI_STATUS
TriggerDisableTopSwap (
  VOID
  )
{
  EFI_STATUS                            Status;
  EFI_SMM_CONTROL2_PROTOCOL             *SmmControl;
  UINT8                                 SmiDataValue;

  Status       = EFI_SUCCESS;
  SmmControl   = NULL;
  SmiDataValue = 0;

  Status = gBS->LocateProtocol (
                  &gEfiSmmControl2ProtocolGuid,
                  NULL,
                  (VOID **) &SmmControl
                  );

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

  SmiDataValue = PcdGet8 (PcdL05TopSwapDisableSwSmi);

  if (SmiDataValue == 0xFF) {
    return EFI_INVALID_PARAMETER;
  }

  Status = SmmControl->Trigger (
                         SmmControl,
                         &SmiDataValue,
                         NULL,
                         0,
                         0
                         );

  return Status;
}

/**
  Set BIOS Self-Healing Feature to enable or disable.

  @param  EnableFlag                    Enable or Disable BIOS Self-Healing.

  @retval EFI_SUCCESS                   Update BIOS successfully.
  @retval Others                        An unexpected error occurred.
**/
EFI_STATUS
EFIAPI
BiosSelfHealingFunctionSwitch (
  IN  BOOLEAN                           EnableFlag
  )
{
  EFI_STATUS                            Status;
  UINT8                                 BiosSelfHealingFlag;
  UINTN                                 NumberFileSystemHandles;
  EFI_HANDLE                            *FileSystemHandles;
  UINTN                                 Index;
  EFI_SIMPLE_FILE_SYSTEM_PROTOCOL       *Esp;
  EFI_SIMPLE_FILE_SYSTEM_PROTOCOL       *Volume;
  EFI_FILE_HANDLE                       RootFs;
  EFI_FILE_HANDLE                       FileHandle;
  CHAR16                                *RecoveryFileFullPath;
  UINTN                                 RecoveryFileFullPathSize;

  Status                  = EFI_SUCCESS;
  BiosSelfHealingFlag     = 0;
  NumberFileSystemHandles = 0;
  FileSystemHandles       = NULL;
  Index                   = 0;
  Esp                     = NULL;
  Volume                  = NULL;
  RootFs                  = NULL;
  FileHandle              = NULL;
  RecoveryFileFullPath    = NULL;
  RecoveryFileFullPathSize = 0; 

  //
  // [Lenovo BIOS Self-Healing Design Guidance Specification v1.9]
  //   2.6 BIOS Setup Option
  //     When "BIOS Self Healing" switch is enabled or disabled,
  //     BIOS should notify EC to enable or disable WDT permanently.
  //
  Status = OemSvcNotifyEcToSetWdtFunction (EnableFlag);

#if (FixedPcdGet32 (PcdL05ChipsetName) == L05_CHIPSET_NAME_ALDERLAKE)
  return EFI_SUCCESS;
#else
  //
  // If switch the function to enabled, just return here
  //
  if (EnableFlag) {
    return EFI_SUCCESS;
  }

  //
  // If switch the function to disabled, clear BackupBlockSynced Flag
  //
  Status = GetBiosSelfHealingFlag (&BiosSelfHealingFlag);

  if (!EFI_ERROR (Status)) {
    BiosSelfHealingFlag &= ~L05_BIOS_SELF_HEALING_FLAG_BACKUP_BLOCK_SYNCED;
    Status = SetBiosSelfHealingFlag (BiosSelfHealingFlag);
  }

  //
  // Find storage backup file & delete it
  //
  RecoveryFileFullPathSize = StrSize ((CHAR16 *) PcdGetPtr (PcdL05SelfRecoveryFolder)) +
                             StrSize ((CHAR16 *) PcdGetPtr (PcdL05SelfRecoveryFile)) +
                             sizeof (CHAR16);

  RecoveryFileFullPath = AllocateZeroPool (RecoveryFileFullPathSize);

  if (RecoveryFileFullPath == NULL) {
    Status = EFI_OUT_OF_RESOURCES;
    goto FunctionSwitchError;
  }

  StrCpyS (RecoveryFileFullPath, RecoveryFileFullPathSize / sizeof (CHAR16), (CHAR16 *) PcdGetPtr (PcdL05SelfRecoveryFolder));
  StrCatS (RecoveryFileFullPath, RecoveryFileFullPathSize / sizeof (CHAR16), L"\\\\");
  StrCatS (RecoveryFileFullPath, RecoveryFileFullPathSize / sizeof (CHAR16), (CHAR16 *) PcdGetPtr (PcdL05SelfRecoveryFile));

  Status = gBS->LocateHandleBuffer (
                  ByProtocol,
                  &gEfiSimpleFileSystemProtocolGuid,
                  NULL,
                  &NumberFileSystemHandles,
                  &FileSystemHandles
                  );

  if (EFI_ERROR (Status)) {
    goto FunctionSwitchError;
  }

  for (Index = 0; Index < NumberFileSystemHandles; Index++) {

    Status = gBS->HandleProtocol (
                    FileSystemHandles[Index],
                    &gEfiPartTypeSystemPartGuid,
                    (VOID **) &Esp
                    );

    if (!EFI_ERROR (Status)) {

      Status = gBS->HandleProtocol (
                      FileSystemHandles[Index],
                      &gEfiSimpleFileSystemProtocolGuid,
                      (VOID *) &Volume
                      );

      if (EFI_ERROR (Status)) {
        continue;
      }

      Status = Volume->OpenVolume (
                         Volume,
                         &RootFs
                         );

      if (EFI_ERROR (Status)) {
        continue;
      }

      Status = RootFs->Open (
                         RootFs,
                         &FileHandle,
                         RecoveryFileFullPath,
                         EFI_FILE_MODE_READ | EFI_FILE_MODE_WRITE,
                         0
                         );

      if (!EFI_ERROR (Status)) {
        //
        // Delete file
        //
        FileHandle->Delete (FileHandle);
        RootFs->Close (RootFs);
        break;
      }

      RootFs->Close (RootFs);
    }
  }

FunctionSwitchError:
  SafeFreePool ((VOID **) &RecoveryFileFullPath);

  return Status;
#endif
}

/**
  Display prompt message.

  @param  UiType                        Corresponding to the UI type.
  @param  DisableVendorUi               A flag indicates whether to mask vendor user interfaces.
  @param  DisplayLogo                   A flag indicates whether to display logo.

  @retval EFI_SUCCESS                   The operation completed successfully.
**/
EFI_STATUS
EFIAPI
BiosSelfHealingDisplayMessage (
  IN  L05_BIOS_SELF_HEALING_UI_TYPE     UiType,
  IN  BOOLEAN                           DisableVendorUi,
  IN  BOOLEAN                           DisplayLogo
  )
{
  EFI_STATUS                            Status;
  EFI_EVENT                             TimerEvent;
  EFI_SIMPLE_TEXT_INPUT_EX_PROTOCOL     *SimpleTextInEx;
  UINTN                                 EventIndex;
  EFI_KEY_DATA                          KeyData;

  Status         = EFI_SUCCESS;
  TimerEvent     = NULL;
  SimpleTextInEx = NULL;
  EventIndex     = 0;

  if ((UiType == BshRecoveryDone) || (UiType == BshRecoveryDoneReset)) {
    mUiType = UiType;
    SetResolutionToTarget (FALSE);
    PrintAt (0, 0, L"%s", SELF_HEALING_RECOVERY_DONE_STRING);

    Status = gBS->CreateEvent (
                    EVT_TIMER | EVT_NOTIFY_SIGNAL,
                    TPL_NOTIFY,
                    ResetSystemTimerEvent,
                    NULL,
                    &TimerEvent
                    );

    if (!EFI_ERROR (Status)) {
      gBS->SetTimer (
             TimerEvent,
             TimerPeriodic,
             EFI_TIMER_PERIOD_SECONDS (5)
             );
    }

    Status = gBS->HandleProtocol (gST->ConsoleInHandle, &gEfiSimpleTextInputExProtocolGuid, (VOID **) &SimpleTextInEx);

    if (!EFI_ERROR (Status)) {
      do {
        ZeroMem (&KeyData, sizeof (EFI_KEY_DATA));
        gBS->WaitForEvent (1, &(SimpleTextInEx->WaitForKeyEx), &EventIndex);
        SimpleTextInEx->ReadKeyStrokeEx (SimpleTextInEx, &KeyData);
      } while ((KeyData.Key.ScanCode != SCAN_ESC) && !mTimeExpired);
    }

    if (UiType == BshRecoveryDoneReset) {
      gRT->ResetSystem (EfiResetPlatformSpecific, EFI_SUCCESS, 0, NULL);
      CpuDeadLoop ();
    }

    return EFI_SUCCESS;
  }

  if (DisplayLogo) {
    DisplayBootLogo ();
  }

  DisplayMessage (UiType, DisableVendorUi);

  return EFI_SUCCESS;
}

/**
  Display update progress bar.

  @param  Completion                    A value between 1 and 100 indicating the current completion progress of the update.

  @retval EFI_SUCCESS                   The operation completed successfully.
**/
EFI_STATUS
EFIAPI
BiosSelfHealingProgressBar (
  IN  UINTN                             Completion
  )
{
  DrawProgressBar (Completion);

  return EFI_SUCCESS;
}

#if (FixedPcdGet32 (PcdL05ChipsetName) != L05_CHIPSET_NAME_ALDERLAKE)
/**
  BIOS recovery callback.

  @param  Event                         Event whose notification function is being invoked.
  @param  Context                       Pointer to the notification function's context.
**/
VOID
EFIAPI
BiosRecoveryCallback (
  IN  EFI_EVENT                         Event,
  IN  VOID                              *Context
  )
{
  EFI_STATUS                            Status;
  VOID                                  *EndOfBdsBootSelection;
  VOID                                  *HobList;
  EFI_PEI_HOB_POINTERS                  RecoveryHob;
  UINTN                                 Buffer;
  UINTN                                 BufferSize;
  UINTN                                 FlashAddress;

  Status                = EFI_SUCCESS;
  EndOfBdsBootSelection = NULL;
  HobList               = NULL;
  Buffer                = 0;
  BufferSize            = 0;
  FlashAddress          = 0;

  Status = gBS->LocateProtocol (
                  &gEndOfBdsBootSelectionProtocolGuid,
                  NULL,
                  (VOID **) &EndOfBdsBootSelection
                  );

  if (EFI_ERROR (Status)) {
    return;
  }

  gBS->CloseEvent (Event);

  //
  // Get recovery file's hob
  //
  HobList = GetHobList ();
  RecoveryHob.Raw = HobList;
  RecoveryHob.Raw = GetNextHob (EFI_HOB_TYPE_MEMORY_ALLOCATION, RecoveryHob.Raw);

  while (RecoveryHob.Header->HobType == EFI_HOB_TYPE_MEMORY_ALLOCATION &&
         !CompareGuid (&RecoveryHob.MemoryAllocationModule->MemoryAllocationHeader.Name,
         &gEfiRecoveryFileAddressGuid)) {
    RecoveryHob.Raw = GET_NEXT_HOB (RecoveryHob);
    RecoveryHob.Raw = GetNextHob (EFI_HOB_TYPE_MEMORY_ALLOCATION, RecoveryHob.Raw);
  }

  Buffer     = (UINTN) RecoveryHob.MemoryAllocationModule->MemoryAllocationHeader.MemoryBaseAddress;
  BufferSize = (UINTN) RecoveryHob.MemoryAllocationModule->MemoryAllocationHeader.MemoryLength;

  //
  // Recovery BIOS
  //
  BiosSelfHealingDisplayMessage (BshRecovery, FALSE, TRUE);

  FlashAddress = (UINTN) PcdGet32 (PcdFlashAreaBaseAddress);
  Status = FlashUpdate (
             FlashAddress,
             Buffer,
             BufferSize,
             TRUE,
             0,
             0,
             FALSE
             );
  DEBUG ((EFI_D_INFO, "BiosSelfHealingDxe: BIOS recovery status: %r\n", Status));

  //
  // [Lenovo BIOS Self-Healing Design Guidance Specification v1.9]
  //   2.2 Detect
  //     On Top-Swap enabled boot,
  //     BIOS code will notify EC to disable Top-Swap pin after recovery successfully.
  //
  OemSvcNotifyEcToDisableTopSwap ();

#ifdef L05_BIOS_SELF_HEALING_TEST_ON_CRB
  TriggerDisableTopSwap ();
#endif

  BiosSelfHealingDisplayMessage (BshRecoveryDoneReset, FALSE, FALSE);

  return;
}

/**
  Update EEPROM/CMOS Flag callback.

  @param  Event                         Event whose notification function is being invoked.
  @param  Handle                        Checkpoint handle.
**/
VOID
EFIAPI
UpdateFlagCallback (
  IN  EFI_EVENT                         Event,
  IN  H2O_CP_HANDLE                     Handle
  )
{
  EFI_STATUS                            Status;
  UINT8                                 BiosSelfHealingFlag;

  Status              = EFI_SUCCESS;
  BiosSelfHealingFlag = 0;

  H2OCpUnregisterHandler (Handle);

  //
  // Confirm whether the previous boot state is PEI Crisis Recovery
  //
//[-start-220125-BAIN000092-modify]//
#ifdef LCFC_SUPPORT
  if (ReadCmos8 (EfiL05BiosSelfHealingModeSwitch) == V_EFI_L05_BIOS_SELF_HEALING_MODE_CRISIS_RECOVERY) {
#else
  if (ReadCmos8 (EfiL05BiosSelfHealingModeSwitch) == V_EFI_L05_BIOS_SELF_HEALING_MODE_CRISIS_RECOVERY_COMPLETED) {
#endif
//[-end-220125-BAIN000092-modify]//
    Status = GetBiosSelfHealingFlag (&BiosSelfHealingFlag);

    if (!EFI_ERROR (Status)) {
      //
      // Clear BackupBlockSynced Flag
      //
      BiosSelfHealingFlag &= ~L05_BIOS_SELF_HEALING_FLAG_BACKUP_BLOCK_SYNCED;
      Status = SetBiosSelfHealingFlag (BiosSelfHealingFlag);
    }
  }

  //
  // Reset CMOS Flag
  //
  WriteCmos8 (EfiL05BiosSelfHealingModeSwitch, V_EFI_L05_BIOS_SELF_HEALING_MODE_NORMAL);

  return;
}

/**
  Update SPI backup callback.

  @param  Event                         Event whose notification function is being invoked.
  @param  Handle                        Checkpoint handle.
**/
VOID
EFIAPI
UpdateSpiBackupCallback (
  IN  EFI_EVENT                         Event,
  IN  H2O_CP_HANDLE                     Handle
  )
{
  EFI_STATUS                            Status;
  UINT8                                 BiosSelfHealingFlag;
  UINTN                                 Index;
  UINTN                                 FvCount;
  UINT32                                *PrimaryFvBase;
  UINT32                                *BackupFvBase;
  UINT32                                *BackupFvSize;
  VOID                                  *Source;
  VOID                                  *Destination;
  UINTN                                 TotalFlashBlocks;
  UINTN                                 CurrentFlashBlocks;
  BOOLEAN                               IsMultiRegion;
  BOOLEAN                               *UpdateFlag;

  Status              = EFI_SUCCESS;
  BiosSelfHealingFlag = 0;
  FvCount             = 0;
  PrimaryFvBase       = NULL;
  BackupFvBase        = NULL;
  BackupFvSize        = NULL;
  Source              = NULL;
  Destination         = NULL;
  TotalFlashBlocks    = 0;
  CurrentFlashBlocks  = 0;
  IsMultiRegion       = FALSE;
  UpdateFlag          = NULL;

  H2OCpUnregisterHandler (Handle);

  //
  // [Lenovo BIOS Self-Healing Design Guidance Specification v1.9]
  //   2.1 Overview
  //     The self healing (EC boot flow and BIOS top swap) flow :
  //       Only confirm the SPI Backup boot block in the case of the first boot after BIOS flash/recovery or
  //       the function is from disable to enable.
  //
  Status = GetBiosSelfHealingFlag (&BiosSelfHealingFlag);

  if (!EFI_ERROR (Status) &&
        ((BiosSelfHealingFlag & L05_BIOS_SELF_HEALING_FLAG_BACKUP_BLOCK_SYNCED) == L05_BIOS_SELF_HEALING_FLAG_BACKUP_BLOCK_SYNCED)) {
    //
    // Backup boot block has been confirmed, so just return here
    //
    DEBUG ((EFI_D_INFO, "BiosSelfHealingDxe: Backup boot block has been confirmed. (BackupBlockSyncFlag is setted)\n"));
    return;
  }

  //
  // Check backup region and calculate the total number of flash blocks
  //
  DEBUG ((EFI_D_INFO, "BiosSelfHealingDxe: Update Backup IBB - Start\n"));

  BiosSelfHealingDisplayMessage (BshBackup, FALSE, TRUE);

  FvCount = RetrieveIbbFvInfo (&PrimaryFvBase, &BackupFvBase, &BackupFvSize);
  UpdateFlag = AllocateZeroPool (FvCount);
  if (FvCount == 0 || UpdateFlag == NULL) {
    goto ExitSpiBackup;
  }

  for (Index = 0; Index < FvCount; Index++) {

    Source      = AllocatePool (BackupFvSize[Index]);
    Destination = AllocatePool (BackupFvSize[Index]);
    if (Source == NULL || Destination == NULL) {
      goto ExitSpiBackup;
    }

    Status = FlashRead (
               Source,
               (UINT8 *)(UINTN) PrimaryFvBase[Index],
               BackupFvSize[Index]
               );
    if (EFI_ERROR (Status)) {
      DEBUG ((EFI_D_ERROR, "Read flash address 0x%x error %r\n", PrimaryFvBase[Index], Status));
      goto ExitSpiBackup;
    }

    Status = FlashRead (
               Destination,
               (UINT8 *)(UINTN) BackupFvBase[Index],
               BackupFvSize[Index]
               );
    if (EFI_ERROR (Status)) {
      DEBUG ((EFI_D_ERROR, "Read flash address 0x%x error %r\n", BackupFvBase[Index], Status));
      goto ExitSpiBackup;
    }

    PatchDataToBuffer (BackupFvBase[Index], Source);

    if (InternalCompareMem64 (Destination, Source, BackupFvSize[Index]) != 0) {
      if (!IsMultiRegion && TotalFlashBlocks > 0) {
        IsMultiRegion = TRUE;
      }
      TotalFlashBlocks += (BackupFvSize[Index] / FLASH_SECTOR_SIZE);
      UpdateFlag[Index] = TRUE;
    }

    SafeFreePool ((VOID **) &Source);
    SafeFreePool ((VOID **) &Destination);
  }

  //
  // Update Backup IBB
  //
  for (Index = 0; Index < FvCount; Index++) {
    DEBUG ((EFI_D_INFO, "Update %x, size %x, Status = ", BackupFvBase[Index], BackupFvSize[Index]));

    if (!UpdateFlag[Index]) {
      DEBUG ((EFI_D_INFO, "%x is valid backup fv, no need to update\n", BackupFvBase[Index]));
      continue;
    }

    Status      = EFI_SUCCESS;
    Source      = AllocatePool (BackupFvSize[Index]);
    Destination = AllocatePool (BackupFvSize[Index]);
    if (Source == NULL || Destination == NULL) {
      goto ExitSpiBackup;
    }

    Status = FlashRead (
               Source,
               (UINT8 *)(UINTN) PrimaryFvBase[Index],
               BackupFvSize[Index]
               );
    if (EFI_ERROR (Status)) {
      DEBUG ((EFI_D_ERROR, "Read flash address 0x%x error %r\n", PrimaryFvBase[Index], Status));
      goto ExitSpiBackup;
    }

    Status = FlashRead (
               Destination,
               (UINT8 *)(UINTN) BackupFvBase[Index],
               BackupFvSize[Index]
               );
    if (EFI_ERROR (Status)) {
      DEBUG ((EFI_D_ERROR, "Read flash address 0x%x error %r\n", BackupFvBase[Index], Status));
      goto ExitSpiBackup;
    }

    PatchDataToBuffer (BackupFvBase[Index], Source);

    Status = FlashUpdate (
               BackupFvBase[Index],
               (UINTN) Source,
               BackupFvSize[Index],
               FALSE,
               (IsMultiRegion ? TotalFlashBlocks : 0),
               (IsMultiRegion ? CurrentFlashBlocks : 0),
               IsMultiRegion
               );
    DEBUG ((EFI_D_INFO, "%r\n", Status));
    CurrentFlashBlocks += (BackupFvSize[Index] / FLASH_SECTOR_SIZE);

    SafeFreePool ((VOID **) &Source);
    SafeFreePool ((VOID **) &Destination);

    if (EFI_ERROR (Status)) {
      goto ExitSpiBackup;
    }
  }

  //
  // Set BackupBlockSynced Flag
  //
  BiosSelfHealingFlag |= L05_BIOS_SELF_HEALING_FLAG_BACKUP_BLOCK_SYNCED;
  Status = SetBiosSelfHealingFlag (BiosSelfHealingFlag);

  if (TotalFlashBlocks > 0) {
    mSpiBackupIsUpdated = TRUE;
  }

  BiosSelfHealingProgressBar (100);

ExitSpiBackup:
  SafeFreePool ((VOID **) &UpdateFlag);
  SafeFreePool ((VOID **) &Source);
  SafeFreePool ((VOID **) &Destination);

  DEBUG ((EFI_D_INFO, "BiosSelfHealingDxe: Update Backup IBB - End\n"));

  return;
}

/**
  Update Storage backup callback.

  @param  Event                         Event whose notification function is being invoked.
  @param  Context                       Pointer to the notification function's context.
**/
VOID
EFIAPI
UpdateStorageBackupCallback (
  IN  EFI_EVENT                         Event,
  IN  VOID                              *Context
  )
{
  EFI_STATUS                            Status;
  VOID                                  *EndOfBdsBootSelection;
  UINTN                                 NumberFileSystemHandles;
  EFI_HANDLE                            *FileSystemHandles;
  UINTN                                 Index;
  EFI_SIMPLE_FILE_SYSTEM_PROTOCOL       *Esp;
  EFI_SIMPLE_FILE_SYSTEM_PROTOCOL       *Volume;
  EFI_FILE_HANDLE                       RootFs;
  EFI_FILE_HANDLE                       FileHandle;
  CHAR16                                *RecoveryFolder;
  CHAR16                                *RecoveryFolderTemp;
  CHAR16                                *Walker;
  CHAR16                                *RecoveryFileFullPath;
  UINTN                                 FlashAreaBaseAddress;
  UINTN                                 FlashAreaSize;
  UINTN                                 BufferSize;
  EFI_FILE_INFO                         *FileInfo;
  UINT8                                 *Buffer;
  UINT64                                BvdtAddr;
  CHAR8                                 *ImageBvdtPtr;
  CHAR8                                 *OnboardBvdtStr;
  UINTN                                 RecoveryFileFullPathSize;

  Status                  = EFI_SUCCESS;
  EndOfBdsBootSelection   = NULL;
  NumberFileSystemHandles = 0;
  FileSystemHandles       = NULL;
  Index                   = 0;
  Esp                     = NULL;
  Volume                  = NULL;
  RootFs                  = NULL;
  FileHandle              = NULL;
  RecoveryFolder          = NULL;
  RecoveryFolderTemp      = NULL;
  Walker                  = NULL;
  RecoveryFileFullPath    = NULL;
  FlashAreaBaseAddress    = (UINTN) PcdGet32 (PcdFlashAreaBaseAddress);
  FlashAreaSize           = (UINTN) PcdGet32 (PcdFlashAreaSize);
  BufferSize              = 0;
  FileInfo                = 0;
  Buffer                  = NULL;
  BvdtAddr                = 0;
  ImageBvdtPtr            = NULL;
  OnboardBvdtStr          = NULL;
  RecoveryFileFullPathSize = 0;

  Status = gBS->LocateProtocol (
                  &gEndOfBdsBootSelectionProtocolGuid,
                  NULL,
                  (VOID **) &EndOfBdsBootSelection
                  );

  if (EFI_ERROR (Status)) {
    return;
  }

  gBS->CloseEvent (Event);

  BufferSize = StrSize ((CHAR16 *) PcdGetPtr (PcdL05SelfRecoveryFolder)) + sizeof (CHAR16) * 2;
  RecoveryFolder = AllocateZeroPool (BufferSize);
  if (RecoveryFolder == NULL) {
    Status = EFI_OUT_OF_RESOURCES;
    goto ExitStorageBackup;
  }

  RecoveryFolderTemp = AllocateZeroPool (BufferSize);
  if (RecoveryFolderTemp == NULL) {
    Status = EFI_OUT_OF_RESOURCES;
    goto ExitStorageBackup;
  }

  RecoveryFileFullPathSize = BufferSize + StrSize ((CHAR16 *) PcdGetPtr (PcdL05SelfRecoveryFile));
  RecoveryFileFullPath = AllocateZeroPool (RecoveryFileFullPathSize);
  if (RecoveryFileFullPath == NULL) {
    Status = EFI_OUT_OF_RESOURCES;
    goto ExitStorageBackup;
  }

  StrCpyS (RecoveryFolder, BufferSize / sizeof (CHAR16), (CHAR16 *) PcdGetPtr (PcdL05SelfRecoveryFolder));
  StrCatS (RecoveryFolder, BufferSize / sizeof (CHAR16), L"\\\\");

  StrCpyS (RecoveryFileFullPath, RecoveryFileFullPathSize / sizeof (CHAR16), RecoveryFolder);
  StrCatS (RecoveryFileFullPath, RecoveryFileFullPathSize / sizeof (CHAR16), (CHAR16 *) PcdGetPtr (PcdL05SelfRecoveryFile));

  Status = gBS->LocateHandleBuffer (
                  ByProtocol,
                  &gEfiSimpleFileSystemProtocolGuid,
                  NULL,
                  &NumberFileSystemHandles,
                  &FileSystemHandles
                  );

  if (EFI_ERROR (Status)) {
    goto ExitStorageBackup;
  }

  //
  // Scan ESP partition
  //
  for (Index = 0; Index < NumberFileSystemHandles; Index++) {

    Status = gBS->HandleProtocol (
                    FileSystemHandles[Index],
                    &gEfiPartTypeSystemPartGuid,
                    (VOID **) &Esp
                    );
    if (EFI_ERROR (Status)) {
      continue;
    }

    Status = gBS->HandleProtocol (
                    FileSystemHandles[Index],
                    &gEfiSimpleFileSystemProtocolGuid,
                    (VOID *) &Volume
                    );
    if (EFI_ERROR (Status)) {
      continue;
    }

    Status = Volume->OpenVolume (
                       Volume,
                       &RootFs
                       );
    if (EFI_ERROR (Status)) {
      continue;
    }

    //
    // 1. Try to open folder structure
    //    If one of the folders does not exist, a new folder will be created automatically
    //
    StrCpyS (RecoveryFolderTemp, BufferSize / sizeof (CHAR16), RecoveryFolder);
    Walker = RecoveryFolderTemp;

    while (1) {

      Walker = StrStr (Walker, L"\\\\");
      if (Walker == NULL) {
        break;
      }
      ZeroMem ((VOID *) Walker, sizeof (CHAR16) * 2);

      Status = RootFs->Open (
                         RootFs,
                         &FileHandle,
                         RecoveryFolderTemp,
                         EFI_FILE_MODE_READ,
                         EFI_FILE_DIRECTORY
                         );
  
      if (EFI_ERROR (Status)) {
        Status = RootFs->Open (
                           RootFs,
                           &FileHandle,
                           RecoveryFolderTemp,
                           EFI_FILE_MODE_CREATE | EFI_FILE_MODE_READ | EFI_FILE_MODE_WRITE,
                           EFI_FILE_ARCHIVE | EFI_FILE_DIRECTORY
                           );
  
        if (EFI_ERROR (Status)) {
          RootFs->Close (RootFs);
          continue;
        }
      }
  
      FileHandle->Close (FileHandle);

      StrCpyS (RecoveryFolderTemp, BufferSize / sizeof (CHAR16), RecoveryFolder);
      Walker += 2;
    }

    //
    // 2. Try to get file info
    //
    Status = RootFs->Open (
                       RootFs,
                       &FileHandle,
                       RecoveryFileFullPath,
                       EFI_FILE_MODE_READ | EFI_FILE_MODE_WRITE,
                       0
                       );

    if (EFI_ERROR (Status)) {
      //
      // File does not exist, so just create it directly
      //
      goto UpdateFile;
    }

    BufferSize = 0;
    Status = FileHandle->GetInfo (
                           FileHandle,
                           &gEfiFileInfoGuid,
                           &BufferSize,
                           FileInfo
                           );
      
    if (Status == EFI_BUFFER_TOO_SMALL) {
      FileInfo = AllocateZeroPool (BufferSize);
      if (FileInfo == NULL) {
        Status = EFI_OUT_OF_RESOURCES;
        goto CloseFileAndContinue;
      }
      Status = FileHandle->GetInfo (
                             FileHandle,
                             &gEfiFileInfoGuid,
                             &BufferSize,
                             FileInfo
                             );
    }

    if (EFI_ERROR (Status) || FileInfo->FileSize == 0) {
      Status = EFI_SUCCESS;
      goto UpdateFile;
    }

    //
    // 3. Check storage backup file version
    //
    BufferSize = FileInfo->FileSize;
    Buffer = AllocateZeroPool (BufferSize);
    if (Buffer == NULL) {
      Status = EFI_OUT_OF_RESOURCES;
      goto CloseFileAndContinue;
    }

    Status = FileHandle->Read (
                           FileHandle,
                           &BufferSize,
                           Buffer
                           );

    if (EFI_ERROR (Status)) {
      Status = EFI_SUCCESS;
      goto UpdateFile;
    }

    BvdtAddr = FdmGetNAtAddr (&gH2OFlashMapRegionBvdtGuid, 1);
    ImageBvdtPtr = (CHAR8 *) ((UINTN) Buffer + ((UINTN) BvdtAddr - FlashAreaBaseAddress) + BIOS_VERSION_OFFSET);
    OnboardBvdtStr = (CHAR8 *) (UINTN) (BvdtAddr + BIOS_VERSION_OFFSET);

    if (AsciiStrCmp (OnboardBvdtStr, ImageBvdtPtr) == 0) {
      //
      // Storage backup file is the same as onboard BIOS, so skip updates
      //
      Status = EFI_ABORTED;
      goto CloseFileAndContinue;
    }

    SafeFreePool ((VOID **)  &FileInfo);
    SafeFreePool ((VOID **)  &Buffer);

UpdateFile:
    //
    // 4. Prepare backup file
    //
    if (!EFI_ERROR (Status)) {
      FileHandle->Delete (FileHandle);
    }

    Status = RootFs->Open (
                       RootFs,
                       &FileHandle,
                       RecoveryFileFullPath,
                       EFI_FILE_MODE_CREATE | EFI_FILE_MODE_READ | EFI_FILE_MODE_WRITE,
                       0
                       );

    if (EFI_ERROR (Status)) {
      RootFs->Close (RootFs);
      continue;
    }

    BufferSize = FlashAreaSize;
    Buffer = AllocateZeroPool (BufferSize);
    if (Buffer == NULL) {
      Status = EFI_OUT_OF_RESOURCES;
      goto CloseFileAndContinue;
    }

    //
    // 5. Write file
    //
    if (!mSpiBackupIsUpdated) {
      BiosSelfHealingDisplayMessage (BshBackup, FALSE, TRUE);
      BiosSelfHealingProgressBar (0);
    }

    Status = FlashRead (
               Buffer,
               (UINT8 *) FlashAreaBaseAddress,
               FlashAreaSize
               );

    if (EFI_ERROR (Status)) {
      DEBUG ((EFI_D_ERROR, "Read flash address 0x%x error %r\n", FlashAreaBaseAddress, Status));
      goto CloseFileAndContinue;
    }

    Status = FileHandle->Write (
                           FileHandle,
                           &BufferSize,
                           Buffer
                           );

    if (!mSpiBackupIsUpdated) {
      BiosSelfHealingProgressBar (100);
      gBS->Stall (EFI_STALL_SECONDS (3));
    }

CloseFileAndContinue:
    DEBUG ((EFI_D_INFO, "BiosSelfHealingDxe: Update Storage backup file status: %r\n", Status));

    SafeFreePool ((VOID **)  &FileInfo);
    SafeFreePool ((VOID **)  &Buffer);

    FileHandle->Close (FileHandle);
    RootFs->Close (RootFs);
  }

ExitStorageBackup:
  SafeFreePool ((VOID **) &RecoveryFolder);
  SafeFreePool ((VOID **) &RecoveryFolderTemp);
  SafeFreePool ((VOID **) &RecoveryFileFullPath);

  return;
}

/**
  Install notification function.

  @retval EFI_SUCCESS                   The operation completed successfully.
  @retval Others                        An unexpected error occurred.
**/
EFI_STATUS
InstallNotificationFunction (
  VOID
  )
{
  EFI_STATUS                            Status;
  EFI_BOOT_MODE                         BootMode;
  VOID                                  *Registration;
  H2O_CP_HANDLE                         CpHandle;

  Status       = EFI_SUCCESS;
  BootMode     = GetBootModeHob ();
  Registration = NULL;
  CpHandle     = NULL;

  //
  // Check if BIOS Self-Healing function is enabled or not
  //
  if (!PcdGetBool (PcdL05BiosSelfHealingEnable)) {
    return EFI_SUCCESS;
  }

  //
  // Register BIOS recovery callback
  //
  Registration = NULL;

  if ((BootMode == BOOT_IN_RECOVERY_MODE) && PcdGetBool (PcdL05TopSwapEnable)) {
    EfiCreateProtocolNotifyEvent (
      &gEndOfBdsBootSelectionProtocolGuid,
      TPL_CALLBACK,
      BiosRecoveryCallback,
      NULL,
      (VOID **) &Registration
      );
  }

  //
  // Register EEPROM/CMOS Flag update callback on H2O_CP_MEDIUM_HIGH of gH2OBdsCpDxeSmmReadyToLockBeforeGuid to
  // ensure it runs before UpdateSpiBackupCallback
  //
  CpHandle = NULL;

  Status = H2OCpRegisterHandler (
             &gH2OBdsCpDxeSmmReadyToLockBeforeGuid,
             UpdateFlagCallback,
             H2O_CP_MEDIUM_HIGH,
             &CpHandle
             );

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

  //
  // Register SPI Backup (IBB_R) update callback on H2O_CP_MEDIUM of gH2OBdsCpDxeSmmReadyToLockBeforeGuid to
  // ensure it runs before BiosProtectionStage1
  //
  CpHandle = NULL;

  Status = H2OCpRegisterHandler (
             &gH2OBdsCpDxeSmmReadyToLockBeforeGuid,
             UpdateSpiBackupCallback,
             H2O_CP_MEDIUM,
             &CpHandle
             );

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

  //
  // Register Storage Backup (OBB_R) update callback
  //
  Registration = NULL;

  EfiCreateProtocolNotifyEvent (
    &gEndOfBdsBootSelectionProtocolGuid,
    TPL_CALLBACK - 1,
    UpdateStorageBackupCallback,
    NULL,
    (VOID **) &Registration
    );

  return Status;
}
#endif

/**
  BIOS Self-Healing DXE Entry.

  @param  ImageHandle                   The firmware allocated handle for the UEFI image.
  @param  SystemTable                   A pointer to the EFI System Table.

  @retval EFI_SUCCESS                   The operation completed successfully.
  @retval Others                        An unexpected error occurred.
**/
EFI_STATUS
BiosSelfHealingDxeEntryPoint (
  IN  EFI_HANDLE                        ImageHandle,
  IN  EFI_SYSTEM_TABLE                  *SystemTable
  )
{
//#[-start-220627-dennis0018-modify]//
#ifdef LCFC_SUPPORT
  LENOVO_VARIABLE_PROTOCOL                *LenovoVariable   = NULL;
  EFI_GUID                                MtmNumberGuid     = LVAR_MTM_NUMBER_GUID;
  UINT32                                  MtmDataSize       = 0x40;
  UINT8                                   MtmBufferPtr[0x40];
#endif

  EFI_STATUS                            Status;
  EFI_L05_BIOS_SELF_HEALING_PROTOCOL    *L05BiosSelfHealingPtr;

  Status                = EFI_SUCCESS;
  L05BiosSelfHealingPtr = NULL;
#ifdef LCFC_SUPPORT
 Status = gBS->LocateProtocol (
                  &gLenovoVariableProtocolGuid, 
                  NULL, 
                  &LenovoVariable
                  );

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

  Status = LenovoVariable->GetVariable (
                             LenovoVariable, 
                             &MtmNumberGuid, 
                             &MtmDataSize, 
                             MtmBufferPtr
                             );

  if (EFI_ERROR (Status)) {
    return EFI_UNSUPPORTED;
  }
#endif
//#[-end-220627-dennis0018-modify]//

  //
  // Install Protocol
  //
  L05BiosSelfHealingPtr = AllocateZeroPool (sizeof (EFI_L05_BIOS_SELF_HEALING_PROTOCOL));

  if (L05BiosSelfHealingPtr == NULL) {
    return EFI_OUT_OF_RESOURCES;
  }

  L05BiosSelfHealingPtr->FunctionSwitch = BiosSelfHealingFunctionSwitch;
  L05BiosSelfHealingPtr->DisplayMessage = BiosSelfHealingDisplayMessage;
  L05BiosSelfHealingPtr->ProgressBar    = BiosSelfHealingProgressBar;

  Status = gBS->InstallProtocolInterface (
                  &ImageHandle,
                  &gEfiL05BiosSelfHealingProtocolGuid,
                  EFI_NATIVE_INTERFACE,
                  L05BiosSelfHealingPtr
                  );

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

#if (FixedPcdGet32 (PcdL05ChipsetName) != L05_CHIPSET_NAME_ALDERLAKE)
  Status = InstallNotificationFunction ();

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

  return Status;
}