/** @file
   Seamless Recovery support library for Windows UEFI Firmware Update Platform

;******************************************************************************
;* Copyright (c) 2013 - 2021, 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 <Uefi.h>
#include <Library/PcdLib.h>
#include <Library/DebugLib.h>
#include <Library/BaseMemoryLib.h>
#include <Library/MemoryAllocationLib.h>
#include <Library/FdSupportLib.h>
#include <Library/FlashRegionLib.h>
#include <Library/UefiBootServicesTableLib.h>
#include <Library/VariableLib.h>
#include <Library/SeamlessRecoveryLib.h>
#include <Protocol/SmmFwBlockService.h>
#include <Protocol/SwapAddressRange.h>
#include <Guid/SysFwUpdateProgress.h>
#include <Library/PciSegmentLib.h>
#include <Library/PchPciBdfLib.h>
#include <Register/SpiRegs.h>
#include <Library/HobLib.h>

STATIC  UINT64       mFtwBackupAddress;

/**
  Trunk SysFwUpdate & H2O progress.

  @retval EFI_SUCCESS         Operation completed successfully.
  @retval Others              Operation was unsuccessful.

**/
EFI_STATUS
SysFwUpdateVariableThunk (
  VOID
  );

/**
  Check whether it is in firmware failure recovery mode

  @param  None

  @retval TRUE      The system is in firmware failure recovery mode
          FALSE     The system is not in firmware failure recovery mode

**/
BOOLEAN
EFIAPI
IsFirmwareFailureRecovery (
  VOID
  )
{
  EFI_GUID *RecoverySig;

  RecoverySig = (EFI_GUID *)(UINTN)mFtwBackupAddress;
  if (RecoverySig != NULL && CompareGuid (PcdGetPtr (PcdSeamlessRecoverySignature), RecoverySig)) {
    return TRUE;
  }

  return FALSE;
}

/**
  Set firmware updating in process signature

  @param  None

  @returns None

**/
VOID
EFIAPI
SetFirmwareUpdatingFlag (
  BOOLEAN FirmwareUpdatingFlag
  )
{
  EFI_STATUS  Status;
  UINTN       Size;
  EFI_SMM_FW_BLOCK_SERVICE_PROTOCOL     *SmmFwbService;
  UINTN                                 BlocksSize;

  if ((!FirmwareUpdatingFlag) && (!IsFirmwareFailureRecovery())) {
    return;
  }

  BlocksSize = GetFlashBlockSize ();

  Status = gBS->LocateProtocol (
                  &gEfiSmmFwBlockServiceProtocolGuid,
                  NULL,
                  (VOID **)&SmmFwbService
                  );
  if (EFI_ERROR (Status)) {
    Status = FlashErase ((UINTN)mFtwBackupAddress, BlocksSize);
    if (EFI_ERROR(Status)) {
      ASSERT_EFI_ERROR(Status);
      return;
    }
    if (FirmwareUpdatingFlag) {
      Size = sizeof(EFI_GUID);
      FlashProgram (
        (UINT8 *)(UINTN)mFtwBackupAddress,
        PcdGetPtr(PcdSeamlessRecoverySignature),
        &Size,
        (UINTN)mFtwBackupAddress
        );
    }
  } else {

    Status = SmmFwbService->EraseBlocks(
                              SmmFwbService,
                              (UINTN)mFtwBackupAddress,
                              &BlocksSize
                              );
    if (EFI_ERROR (Status)) {
      ASSERT_EFI_ERROR (Status);
      return;
    }
    if (FirmwareUpdatingFlag) {
      Size = sizeof (EFI_GUID);
      Status = SmmFwbService->Write (
                                SmmFwbService,
                                (UINTN)mFtwBackupAddress,
                                (UINTN *)&Size,
                                (UINT8 *)PcdGetPtr (PcdSeamlessRecoverySignature)
                                );
    }
  }
}

/**
  Set firmware updating progress.

  @param[in]  Progress         Update progress

  @retval EFI_SUCCESS         Operation completed successfully.
  @retval Others              Operation was unsuccessful.

**/
EFI_STATUS
EFIAPI
SetFirmwareUpdateProgress (
  IN  UINT32    Progress
  )
{
  return SetVariableToSensitiveVariable (
           UPDATE_PROGRESS_NAME,
           &gH2OSeamlessRecoveryGuid,
           EFI_VARIABLE_BOOTSERVICE_ACCESS | EFI_VARIABLE_NON_VOLATILE,
           sizeof (UINT32),
           &Progress
           );
}

/**
  Get firmware updating progress.

  @param[out]  Progress         Update progress

  @retval EFI_SUCCESS         Operation completed successfully.
  @retval Others              Operation was unsuccessful.

**/
EFI_STATUS
EFIAPI
GetFirmwareUpdateProgress (
  OUT  UINT32    *Progress
  )
{
  EFI_STATUS                        Status;
  UINTN                             DataSize;
  UINT32                            Data32;

  SysFwUpdateVariableThunk ();

  DataSize = sizeof (UINT32);
  Status = CommonGetVariable (
             UPDATE_PROGRESS_NAME,
             &gH2OSeamlessRecoveryGuid,
             &DataSize,
             &Data32
             );
  if (EFI_ERROR (Status)) {
    DEBUG ((DEBUG_ERROR, "Failed to get Progress variable -%r\n"));
    return Status;
  }

  *Progress = Data32;
  return EFI_SUCCESS;
}

/**
  Check whether there was a firmware update interrupted.

  @retval TRUE      Fimware update was interrupted.
          FALSE     No fimware update was interrupted.

**/
BOOLEAN
EFIAPI
IsFirmwareUpdateInterrupted (
  VOID
  )
{
  EFI_STATUS                        Status;
  UINTN                             DataSize;
  UINT32                            Data32;

  DataSize = sizeof (UINT32);
  Status = CommonGetVariable (
             UPDATE_PROGRESS_NAME,
             &gH2OSeamlessRecoveryGuid,
             &DataSize,
             &Data32
             );
  if (EFI_ERROR (Status)) {
    DEBUG ((DEBUG_ERROR, "Failed to get Progress variable -%r\n"));
    return FALSE;
  }

  return ((Data32 > FlashStart) && (Data32 < FlashResiliency));
}

/**
  Clear firmware updating progress.

  @retval EFI_SUCCESS         Operation completed successfully.

**/
EFI_STATUS
EFIAPI
ClearFirmwareUpdateProgress (
  VOID
  )
{

  SetVariableToSensitiveVariable (
    UPDATE_PROGRESS_NAME,
    &gH2OSeamlessRecoveryGuid,
    0,
    0,
    NULL
    );

  CommonSetVariable (
    SYSFW_UPDATE_PROGRESS_VARIABLE_NAME,
    &gSysFwUpdateProgressGuid,
    0,
    0,
    NULL
    );

  CommonSetVariable (
    L"FmpCapsuleInfo",
    &gFmpCapsuleInfoGuid,
    0,
    0,
    NULL
    );

  return EFI_SUCCESS;
}

/**
  Check whether firmware update resiliency is supported.

  @retval TRUE      Firmware update resiliecny is supported.
          FALSE     Firmware update resiliecny is not supported.

**/
BOOLEAN
EFIAPI
IsFirmwareUpdateResiliencySupported (
  VOID
  )
{
  return FeaturePcdGet (PcdH2OBiosUpdateFaultToleranceResiliencyEnabled);
}

BOOLEAN
DxeIsTopSwapEnabled (
  VOID
  )
{
  return !!(PciSegmentRead32 (SpiPciCfgBase () + R_SPI_CFG_BC) & B_SPI_CFG_BC_TSS);
}
/**
  Check whether firmware update resiliency is enabled.

  @retval TRUE      Firmware update resiliency is enabled.
          FALSE     Firmware update resiliency is not enabled.

**/
BOOLEAN
EFIAPI
IsFirmwareUpdateResiliencyEnabled (
  VOID
  )
{
  EFI_STATUS            Status;
  UINT32                Data32;
  EFI_HOB_GUID_TYPE     *GuidHob;
  BOOLEAN               TopSwapStatus;

  TopSwapStatus = FALSE;

  if (!IsFirmwareUpdateResiliencySupported ()) {
    return FALSE;
  }

  GuidHob = GetFirstGuidHob (&gChasmfallsTopSwapStatusGuid);
  if (GuidHob != NULL) {
    TopSwapStatus = *((BOOLEAN *) GET_GUID_HOB_DATA (GuidHob));
  }

  Status = GetFirmwareUpdateProgress (&Data32);
  if (EFI_ERROR (Status)) {
    if (DxeIsTopSwapEnabled()){
      return TRUE;
    } else if (TopSwapStatus) {
      return TRUE;
    }
    return FALSE;
  }

  return (Data32 == FlashResiliency);
}

/**
  Get SysFw update progress.

  @retval EFI_SUCCESS         Operation completed successfully.
  @retval Others              Operation was unsuccessful.

**/
EFI_STATUS
GetSysFwUpdateProgress (
  OUT  SYSTEM_FIRMWARE_UPDATE_PROGRESS   *UpdatePorgress
  )
{
  EFI_STATUS                        Status;
  UINTN                             DataSize;
  SYSTEM_FIRMWARE_UPDATE_PROGRESS   TempUpdatePorgress;

  DataSize = sizeof (SYSTEM_FIRMWARE_UPDATE_PROGRESS);
  Status = CommonGetVariable (
             SYSFW_UPDATE_PROGRESS_VARIABLE_NAME,
             &gSysFwUpdateProgressGuid,
             &DataSize,
             &TempUpdatePorgress
             );
  if (EFI_ERROR (Status)) {
    return Status;
  }

  CopyMem (UpdatePorgress, &TempUpdatePorgress, sizeof (SYSTEM_FIRMWARE_UPDATE_PROGRESS));

  return EFI_SUCCESS;
}

/**
  Trunk SysFwUpdate & H2O progress.

  @retval EFI_SUCCESS         Operation completed successfully.
  @retval Others              Operation was unsuccessful.

**/
EFI_STATUS
SysFwUpdateVariableThunk (
  VOID
  )
{
  EFI_STATUS                        Status;
  SYSTEM_FIRMWARE_UPDATE_PROGRESS   UpdatePorgress;
  UINT32                            H2OProgress;

  Status = GetSysFwUpdateProgress (&UpdatePorgress);
  if (EFI_ERROR (Status)) {
    return Status;
  }

  if ((UpdatePorgress.Component != UpdatingBios) &&
     (UpdatePorgress.Component != UpdatingResiliency) &&
     (UpdatePorgress.Component != UpdatinguCode) &&
     (UpdatePorgress.Component != UpdatingBtGAcm)) {
    return EFI_UNSUPPORTED;
  }

  if (UpdatePorgress.Component == UpdatingBios) {
    if (IsFirmwareUpdateResiliencySupported ()) {
      switch (UpdatePorgress.Progress) {
      case BiosIbbR:
        H2OProgress = FlashPbb;
        break;
      case BiosIbb:
        H2OProgress = FlashSbb;
        break;
      default:
        H2OProgress = FlashMax;
        break;
      }
    } else {
      switch (UpdatePorgress.Progress) {
      case BiosIbbR:
        H2OProgress = FlashPbbR;
        break;
      case BiosIbb:
        H2OProgress = FlashPbb;
        break;
      case BiosObb:
        H2OProgress = FlashSbb;
        break;
      default:
        H2OProgress = FlashMax;
        break;
      }
    }
  } else if (UpdatePorgress.Component == UpdatingResiliency) {
    H2OProgress = FlashResiliency;
  } else if (UpdatePorgress.Component == UpdatinguCode) {
    H2OProgress = UpdatePorgress.Progress;
  } else if (UpdatePorgress.Component == UpdatingBtGAcm) {
    switch (UpdatePorgress.Progress) {
    case BiosIbbR:
      H2OProgress = FlashPbb;
      break;    
    default:
      H2OProgress = FlashMax;
      break;
     }     
  }
  Status = SetFirmwareUpdateProgress (H2OProgress);
  ASSERT_EFI_ERROR (Status);

  return EFI_SUCCESS;
}

/**
  The constructor function stores the FDM information to prevent from the
  FDM location is changed after BIOS is updated.

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

  @retval EFI_SUCCESS     The constructor always returns EFI_SUCCESS.
**/
EFI_STATUS
EFIAPI
SeamlessRecoveryLibConstructor (
  IN      EFI_HANDLE                ImageHandle,
  IN      EFI_SYSTEM_TABLE          *SystemTable
  )
{

  mFtwBackupAddress = FdmGetNAtAddr (&gH2OFlashMapRegionFtwBackupGuid ,1);

  SysFwUpdateVariableThunk ();

  return EFI_SUCCESS;
}