/** @file
  Lite version of SeamlessRecoveryLib for PEIMs before memory initialization

;******************************************************************************
;* Copyright (c) 2020 - 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/PeiServicesLib.h>
#include <Library/HobLib.h>
#include <Library/BaseMemoryLib.h>
#include <Library/FlashRegionLib.h>
#include <Library/VariableLib.h>
#include <Library/SeamlessRecoveryLib.h>
#include <Library/SpiAccessLib.h>
#include <Guid/SysFwUpdateProgress.h>
#include <Library/PciSegmentLib.h>
#include <Library/PchPciBdfLib.h>
#include <Register/SpiRegs.h>
/**
  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) FdmGetNAtAddr (&gH2OFlashMapRegionFtwBackupGuid ,1);
  if (RecoverySig != NULL && CompareGuid( PcdGetPtr(PcdSeamlessRecoverySignature), RecoverySig) ) {
    return TRUE;
  }

  return FALSE;
}

/**
  Set firmware updating in process signature, not supported if PEI memory is not initialized

  @param  None

  @returns None

**/
VOID
EFIAPI
SetFirmwareUpdatingFlag (
  BOOLEAN FirmwareUpdatingFlag
  )
{
  //
  // not supported
  //
  return;
}

/**
  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
  )
{
  //
  // not supported
  //
  return EFI_UNSUPPORTED;
}

/**
  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_HOB_GUID_TYPE                 *GuidHob;

  GuidHob = GetFirstGuidHob (&gH2OSeamlessRecoveryGuid);
  if (GuidHob == NULL) {
    return EFI_NOT_FOUND;
  }

  *Progress = *(UINT32 *)GET_GUID_HOB_DATA (GuidHob);
  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;
  UINT32                            Data32;

  Status = GetFirmwareUpdateProgress (&Data32);
  if (EFI_ERROR (Status)) {
    DEBUG ((DEBUG_ERROR, "Failed to get Progress hob -%r\n"));
    return FALSE;
  }

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

/**
  Clear firmware updating progress.

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

**/
EFI_STATUS
EFIAPI
ClearFirmwareUpdateProgress (
  VOID
  )
{
  //
  // not supported
  //
  return EFI_UNSUPPORTED;
}

/**
  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
PeiIsTopSwapEnabled (
  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 (PeiIsTopSwapEnabled()){
      return TRUE;
    } else if (TopSwapStatus) {
      return TRUE;
    }
    return FALSE;
  }

  return (Data32 == FlashResiliency);
}

/**
  Create hobs for H2O progress.

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

**/
STATIC
EFI_STATUS
CreateH2OSeamlessRecoveryHob (
  VOID
  )
{
  EFI_STATUS                        Status;
  UINT32                            Data32;
  UINTN                             DataSize;
  EFI_HOB_GUID_TYPE                 *GuidHob;

  Status = EFI_SUCCESS;

  GuidHob = GetFirstGuidHob (&gH2OSeamlessRecoveryGuid);
  if (GuidHob == NULL) {
    DataSize = sizeof (UINT32);
    Status = CommonGetVariable (
               UPDATE_PROGRESS_NAME,
               &gH2OSeamlessRecoveryGuid,
               &DataSize,
               &Data32
               );
    if (!EFI_ERROR (Status)) {
      BuildGuidDataHob (&gH2OSeamlessRecoveryGuid, &Data32, sizeof (UINT32));
    }
  }

  return Status;
}

/**
  Create hobs for digest.

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

**/
STATIC
EFI_STATUS
CreateH2OSeamlessRecoveryDigestHob (
  VOID
  )
{
  EFI_STATUS                        Status;
  EFI_HOB_GUID_TYPE                 *GuidHob;
  VOID                              *Sha256DigestInFv;
  EFI_PEI_FILE_HANDLE               DigestFileHandle;

  Sha256DigestInFv   = NULL;
  DigestFileHandle   = NULL;

  GuidHob = GetFirstGuidHob (&gH2OSeamlessRecoveryDigestGuid);
  if (GuidHob == NULL) {

    Status = PeiServicesFfsFindFileByName (&gSbbDigestGuid, (EFI_PEI_FV_HANDLE) FixedPcdGet32 (PcdFlashFvRecoveryBase), &DigestFileHandle);
    if (!EFI_ERROR (Status)) {
      //
      // Search RAW section.
      //
      Status = PeiServicesFfsFindSectionData (EFI_SECTION_RAW, DigestFileHandle, &Sha256DigestInFv);
      if (!EFI_ERROR (Status)) {
        DEBUG ((DEBUG_INFO, "Find Obb digest on ROM:\n"));
        BuildGuidDataHob (&gH2OSeamlessRecoveryDigestGuid, Sha256DigestInFv, SHA256_DIGEST_SIZE);
      }
    }
  }

  return EFI_SUCCESS;
}

/**
  Get SysFw update progress.

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

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


  //
  // Only flash interrupted case will exist progress variable.
  //
  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;
}

/**
  Get sysfw update digest.

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

**/
STATIC
EFI_STATUS
CreateSysFwUpdateDigestHob (
  VOID
  )
{
  EFI_STATUS                        Status;
  UINTN                             DataSize;
  UINT8                             Sha256Variable[SHA256_DIGEST_SIZE];
  EFI_HOB_GUID_TYPE                 *GuidHob;

  GuidHob = GetFirstGuidHob (&gSysFwUpdateDigiestGuid);
  if (GuidHob != NULL) {
    //
    // Already exists
    //
    return EFI_SUCCESS;
  }

  DataSize = SHA256_DIGEST_SIZE;
  ZeroMem (Sha256Variable, SHA256_DIGEST_SIZE);
  Status = CommonGetVariable (
             SYSBIOS_NEW_CAPSULE_DIGEST_VARIABLE_NAMEN,
             &gSysFwUpdateDigiestGuid,
             &DataSize,
             &Sha256Variable
             );
  if (EFI_ERROR (Status)) {
    return Status;
  }

  BuildGuidDataHob (&gSysFwUpdateDigiestGuid, Sha256Variable, DataSize);

  return EFI_SUCCESS;
}

/**
  Trunk SysFwUpdate & H2O progress.

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

**/
STATIC
EFI_STATUS
SysFwUpdateProgressHobThunk (
  VOID
  )
{
  EFI_STATUS                        Status;
  SYSTEM_FIRMWARE_UPDATE_PROGRESS   UpdatePorgress;
  UINT32                            H2OProgress;
  EFI_HOB_GUID_TYPE                 *GuidHob;

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

  GuidHob = GetFirstGuidHob (&gSysFwUpdateProgressGuid);
  if (GuidHob == NULL) { 
    BuildGuidDataHob (&gSysFwUpdateProgressGuid, &UpdatePorgress, sizeof (SYSTEM_FIRMWARE_UPDATE_PROGRESS));
  }

  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 = FlashPbb;
  } else if (UpdatePorgress.Component == UpdatingBtGAcm) {
    switch (UpdatePorgress.Progress) {
    case BiosIbbR:
      H2OProgress = FlashPbb;
      break;    
    default:
      H2OProgress = FlashMax;
      break;
     }
  }

  GuidHob = GetFirstGuidHob (&gH2OSeamlessRecoveryGuid);
  if (GuidHob == NULL) { 
    BuildGuidDataHob (&gH2OSeamlessRecoveryGuid, &H2OProgress, sizeof (UINT32));
  }

  CreateSysFwUpdateDigestHob ();

  return EFI_SUCCESS;
}

/**
  Trunk SysFwUpdate & H2O progress.

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

**/
STATIC
EFI_STATUS
FmpCapsuleInfoThunk (
  VOID
  )
{
  EFI_STATUS                        Status;
  EFI_HOB_GUID_TYPE                 *GuidHob;
  UINT32                            H2OProgress;
  UINTN                             DataSize;
  SYSTEM_FIRMWARE_COMPONENT         Component;
  UINT64                            OsIndications;
  SYSTEM_FIRMWARE_UPDATE_PROGRESS   UpdatePorgress;
  UINT32                            Data32;
  
  if (PcdGet8(PcdChasmFallsSupport) == 1){
    GuidHob = GetFirstGuidHob (&gChasmfallsCrisisRecoveryGuid);
    if (GuidHob != NULL) { 
      H2OProgress = FlashStart;
      UpdatePorgress.Component = UpdatingBios;
      UpdatePorgress.Progress = 0;
      goto BuildHob;
    }
  }
  //
  // Capsule put in HDD.
  // Insyde capsule will set progress variable to FlashStart in DxeCapsuleLib.
  //
  DataSize = sizeof (SYSTEM_FIRMWARE_COMPONENT);
  Status = CommonGetVariable (
             L"FmpCapsuleInfo",
             &gFmpCapsuleInfoGuid,
             &DataSize,
             &Component
             );
  if (EFI_ERROR (Status)) {
    return Status;
  }

  //
  // If it is capsule update?
  //
  DataSize = sizeof (UINT64);
  Status = CommonGetVariable (
             L"OsIndications",
             &gEfiGlobalVariableGuid,
             &DataSize,
             &OsIndications
             );
  if (EFI_ERROR (Status)) {
    return Status;
  }

  if ((OsIndications & EFI_OS_INDICATIONS_FILE_CAPSULE_DELIVERY_SUPPORTED) == 0) {
    return EFI_UNSUPPORTED;
  }

  if (Component == UpdatingBios) {
    H2OProgress = FlashStart;
    
    UpdatePorgress.Component = UpdatingBios;
    UpdatePorgress.Progress = 0;
  } else if (Component == UpdatinguCode) {
    H2OProgress = FlashStart;

    UpdatePorgress.Component = UpdatinguCode;
    UpdatePorgress.Progress = 0;
  } else if (Component == UpdatingBtGAcm) {
    H2OProgress = FlashStart;

    UpdatePorgress.Component = UpdatingBtGAcm;
    UpdatePorgress.Progress = 0;
  } else if (Component == UpdatingTypeMax){
    // Monolithic update. Sync bios's progress.
    DataSize = sizeof (UINT32);
    Status = CommonGetVariable (
               UPDATE_PROGRESS_NAME,
               &gH2OSeamlessRecoveryGuid,
               &DataSize,
               &Data32
               );
    if (!EFI_ERROR (Status)) {
      BuildGuidDataHob (&gH2OSeamlessRecoveryGuid, &Data32, sizeof (UINT32));
      return EFI_SUCCESS;
    } else {
      // First time do capsule update.
      H2OProgress = FlashStart;

      UpdatePorgress.Component = UpdatingBios;
      UpdatePorgress.Progress = 0;
    } 
  } else {
    //
    // Not support device firmware case. 
    //
    return EFI_UNSUPPORTED;
  }

BuildHob:
  //
  // Create insyde progress hob to disable bioslock.
  //
  GuidHob = GetFirstGuidHob (&gH2OSeamlessRecoveryGuid);
  if (GuidHob == NULL) { 
    BuildGuidDataHob (&gH2OSeamlessRecoveryGuid, &H2OProgress, sizeof (UINT32));
  }

  return EFI_SUCCESS;
}

/**
  The constructor function of PeiSeamlessRecoveryLib.

  @param  FileHandle   The handle of FFS header the loaded driver.
  @param  PeiServices  The pointer to the PEI services.

  @retval EFI_SUCCESS           The constructor executes successfully.
  @retval EFI_OUT_OF_RESOURCES  There is no enough resource for the constructor.

**/
EFI_STATUS
EFIAPI
SeamlessRecoveryLibConstructor (
  IN EFI_PEI_FILE_HANDLE        FileHandle,
  IN CONST EFI_PEI_SERVICES     **PeiServices
  )
{
  EFI_STATUS                        Status;
  // 
  // First sync intel progress variable to insyde progress variable. 
  // There are two case.
  //
  //
  // (2) Flash interrupted case,
  //     Sync progress varialbe for recovery to load old/new dxe and disable bios lock. 
  //
  Status = SysFwUpdateProgressHobThunk ();
  if (EFI_ERROR (Status)) {
    //
    // (1) First time to do capsule update.
    //     Build progress hob to disable bios lock. 
    //
    Status = FmpCapsuleInfoThunk ();
  }

  //
  // No intel capsule. 
  // Only build insyde progress hob to disable bioslock or recovery to load old dxe or new dxe.
  //
  if (EFI_ERROR (Status)) {
    CreateH2OSeamlessRecoveryHob ();
  }

  //
  // Create sbb digest hob for case(2) to load old obb.
  // Both Intel and Insyde used.
  //
  CreateH2OSeamlessRecoveryDigestHob ();

  return EFI_SUCCESS;
}