/** @file
  BIOS Self-Healing PEI Module.

;******************************************************************************
;* Copyright (c) 2020, Insyde Software Corporation. 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 "BiosSelfHealingPei.h"

EFI_PEI_NOTIFY_DESCRIPTOR               mSelfHealingInitDescriptor = {
  (EFI_PEI_PPI_DESCRIPTOR_NOTIFY_CALLBACK | EFI_PEI_PPI_DESCRIPTOR_TERMINATE_LIST),
#if (FixedPcdGet32 (PcdL05ChipsetName) == L05_CHIPSET_NAME_CEZANNE || \
     FixedPcdGet32 (PcdL05ChipsetName) == L05_CHIPSET_NAME_REMBRANDT)
  &gAmdPspCommonServicePpiGuid,
#else
  &gEfiPeiReadOnlyVariable2PpiGuid,
#endif
  SelfHealingInitCallback
};

EFI_PEI_NOTIFY_DESCRIPTOR               mStage2EcNotifyDescriptor = {
  (EFI_PEI_PPI_DESCRIPTOR_NOTIFY_CALLBACK | EFI_PEI_PPI_DESCRIPTOR_TERMINATE_LIST),
  &gEfiPeiMemoryDiscoveredPpiGuid,
  Stage2EcNotifyCallback
};

EFI_PEI_NOTIFY_DESCRIPTOR               mClearWdtDescriptor = {
  (EFI_PEI_PPI_DESCRIPTOR_NOTIFY_CALLBACK | EFI_PEI_PPI_DESCRIPTOR_TERMINATE_LIST),
  &gEfiEndOfPeiSignalPpiGuid,
  ClearWdtCallback
};

EFI_PEI_NOTIFY_DESCRIPTOR               mInstallFirmwareAuthHookDescriptor = {
  (EFI_PEI_PPI_DESCRIPTOR_NOTIFY_CALLBACK | EFI_PEI_PPI_DESCRIPTOR_TERMINATE_LIST),
  &gEfiPeiMemoryDiscoveredPpiGuid,
  InstallFirmwareAuthHook
};

EFI_PEI_NOTIFY_DESCRIPTOR               mFirmwareAuthHookDescriptor = {
  (EFI_PEI_PPI_DESCRIPTOR_NOTIFY_CALLBACK | EFI_PEI_PPI_DESCRIPTOR_TERMINATE_LIST),
  &gEfiPeiRecoveryModulePpiGuid,
  FirmwareAuthHookCallback
};

FIRMWARE_AUTHENTICATION_PPI             mFirmwareAuthPpi = {
  VerifyFirmware
};

EFI_PEI_PPI_DESCRIPTOR                  mFirmwareAuthPpiList = {
  (EFI_PEI_PPI_DESCRIPTOR_PPI | EFI_PEI_PPI_DESCRIPTOR_TERMINATE_LIST),
  &gFirmwareAuthenticationPpiGuid,
  &mFirmwareAuthPpi
};

UINT32                                  mIgnoreFvBase[] = {
  FixedPcdGet32 (PcdFlashNvStorageVariableDefaultsBase)
};

/**
  Check if BIOS Self-Healing is enabled or not.

  @retval TRUE                          BIOS Self-Healing is enabled.
  @retval FALSE                         BIOS Self-Healing is disabled.
**/
BOOLEAN
BiosSelfHealingEnabled (
  VOID
  )
{
  SYSTEM_CONFIGURATION                  *SetupNvData;
  UINTN                                 BufferSize;
  BOOLEAN                               BiosSelfHealingIsEnabled;

  SetupNvData              = NULL;
  BufferSize               = 0;
  BiosSelfHealingIsEnabled = FALSE;

  CommonGetVariableDataAndSize (
    L"Setup",
    &gSystemConfigurationGuid,
    &BufferSize,
    (VOID **) &SetupNvData
    );

  if (SetupNvData != NULL) {
    BiosSelfHealingIsEnabled = (SetupNvData->L05BiosSelfHealing == 1) ? TRUE : FALSE;
  }

  return BiosSelfHealingIsEnabled;
}

/**
  Callback function to update fv trust status after verifying region.

  @param  Event                         A pointer to the Event that triggered the callback.
  @param  Handle                        Checkpoint handle.
**/
VOID
EFIAPI
CpVerifyFvCallback (
  IN  EFI_EVENT                         Event,
  IN  H2O_CP_HANDLE                     Handle
  )
{
  EFI_STATUS                            Status;
  H2O_BASE_CP_VERIFY_FV_DATA            *VerifyFvData;
  UINTN                                 Index;

  DEBUG ((DEBUG_INFO, "\nBiosSelfHealingPei: CpVerifyFvCallback is triggered\n"));

  Status = H2OCpLookup (Handle, (VOID **) &VerifyFvData, NULL);
  if (EFI_ERROR (Status)) {
    return;
  }

  for (Index = 0; Index < (sizeof (mIgnoreFvBase) / sizeof (UINT32)); Index++) {
    if (VerifyFvData->FvBase == mIgnoreFvBase[Index]) {
      VerifyFvData->Status  = H2O_BDS_TASK_UPDATE;
      VerifyFvData->Trusted = TRUE;
      DEBUG ((DEBUG_INFO, "IGNORE!!! VerifyFvData->FvBase: 0x%lx.\n", VerifyFvData->FvBase));
      break;
    }
  }
}

/**
  Callback function to perform Self-Healing function initialization.

  @param  PeiServices                   General purpose services available to every PEIM.
  @param  NotifyDescriptor              Pointer of the notificaiton data structure.
  @param  Ppi                           Pointer of PPI.

  @retval EFI_SUCCESS                   Operation is successful.
  @retval Others                        An unexpected error occurred.
**/
EFI_STATUS
SelfHealingInitCallback (
  IN  EFI_PEI_SERVICES                  **PeiServices,
  IN  EFI_PEI_NOTIFY_DESCRIPTOR         *NotifyDescriptor,
  IN  VOID                              *Ppi
  )
{
  EFI_STATUS                            Status;

  DEBUG ((DEBUG_INFO, "BiosSelfHealingPei: SelfHealingInitCallback is triggered\n"));
//[-start-220125-BAIN000092-remove]//
#ifndef LCFC_SUPPORT
#if (FixedPcdGet32 (PcdL05ChipsetName) == L05_CHIPSET_NAME_ALDERLAKE || \
     FixedPcdGet32 (PcdL05ChipsetName) == L05_CHIPSET_NAME_TIGERLAKE)
  //
  // If run crisis recovery successfully in Top Swap mode, then disable Top Swap before end of SiInit
  //
  if (PcdGetBool (PcdL05TopSwapEnable) &&
      (CheckPbbrSyncFlag () || (ReadCmos8 (EfiL05BiosSelfHealingModeSwitch) == V_EFI_L05_BIOS_SELF_HEALING_MODE_CRISIS_RECOVERY_COMPLETED))) {
    TopSwapSet (FALSE);
    ResetCold ();
  }
#endif
#endif
//[-end-220125-BAIN000092-remove]//

#if (FixedPcdGet32 (PcdL05ChipsetName) == L05_CHIPSET_NAME_CEZANNE || \
     FixedPcdGet32 (PcdL05ChipsetName) == L05_CHIPSET_NAME_REMBRANDT)
{
  BOOLEAN                               CmosBshFlag = FALSE;
  //
  // Sync Top Swap status to PCD
  //
  if (ReadCmos8 (EfiL05BiosSelfHealingModeSwitch) == V_EFI_L05_BIOS_SELF_HEALING_MODE_FORCE_ENTER_BSH) {
    WriteCmos8 (EfiL05BiosSelfHealingModeSwitch, V_EFI_L05_BIOS_SELF_HEALING_MODE_NORMAL);
    CmosBshFlag = TRUE;
  }
  PcdSetBoolS (PcdL05TopSwapEnable, TopSwapStatus () ? TRUE : CmosBshFlag);
  DEBUG ((DEBUG_INFO, "BiosSelfHealingPei: CmosSwitch status: %x\n", CmosBshFlag));
  DEBUG ((DEBUG_INFO, "BiosSelfHealingPei: PcdL05TopSwapEnable status: %x\n", (UINTN) PcdGetBool (PcdL05TopSwapEnable)));
}
#endif

  //
  // Check if BIOS Self-Healing function is enabled or not
  //
  PcdSetBoolS (PcdL05BiosSelfHealingEnable, BiosSelfHealingEnabled ());
  DEBUG ((DEBUG_INFO, "BiosSelfHealingPei: PcdL05BiosSelfHealingEnable status: %x\n", (UINTN) PcdGetBool (PcdL05BiosSelfHealingEnable)));

  if (!PcdGetBool (PcdL05BiosSelfHealingEnable)) {
    return EFI_SUCCESS;
  }

  //
  // [Lenovo BIOS Self-Healing Design Guidance Specification v1.9]
  //   2.1 Overview
  //     Notify EC that PEI has started (Stage 1)
  //
  OemSvcNotifyEcToPeiStart (TRUE);

  //
  // [Self-Healing EC WDT timer optimize]
  //   Split into two stages to notify EC - "PEI start (before memory training)" and "Memory ready"
  //   Before memory ready, EC will reset timer 3 times to extend WDT timer if Port 80 keep changing
  //
  Status = PeiServicesNotifyPpi (&mStage2EcNotifyDescriptor);
  if (EFI_ERROR (Status)) {
    return Status;
  }

  //
  // Register callback to clear WDT during POST
  //
  if (!PcdGetBool (PcdL05TopSwapEnable)) {
    Status = PeiServicesNotifyPpi (&mClearWdtDescriptor);
    if (EFI_ERROR (Status)) {
      return Status;
    }
  }

#if (FixedPcdGet32 (PcdL05ChipsetName) != L05_CHIPSET_NAME_ALDERLAKE)
  //
  // Register InstallFirmwareAuthHookDescriptor
  //
  Status = PeiServicesNotifyPpi (&mInstallFirmwareAuthHookDescriptor);
  if (EFI_ERROR (Status)) {
    return Status;
  }
#endif

  return EFI_SUCCESS;
}

/**
  Callback function to notify EC that memory is ready.

  @param  PeiServices                   General purpose services available to every PEIM.
  @param  NotifyDescriptor              Pointer of the notificaiton data structure.
  @param  Ppi                           Pointer of PPI.

  @retval EFI_SUCCESS                   Operation is successful.
  @retval Others                        An unexpected error occurred.
**/
EFI_STATUS
Stage2EcNotifyCallback (
  IN  EFI_PEI_SERVICES                  **PeiServices,
  IN  EFI_PEI_NOTIFY_DESCRIPTOR         *NotifyDescriptor,
  IN  VOID                              *Ppi
  )
{
  DEBUG ((DEBUG_INFO, "BiosSelfHealingPei: Stage2EcNotifyCallback is triggered\n"));

  //
  // [Self-Healing EC WDT timer optimize]
  //   Split into two stages to notify EC - "PEI start (before memory training)" and "Memory ready"
  //   Before memory ready, EC will reset timer 3 times to extend WDT timer if Port 80 keep changing
  //
  OemSvcNotifyEcToPeiStart (FALSE);

  return EFI_SUCCESS;
}

/**
  Clear WDT callback.

  @param  PeiServices                   General purpose services available to every PEIM.
  @param  NotifyDescriptor              Pointer of the notificaiton data structure.
  @param  Ppi                           Pointer of PPI.

  @retval EFI_SUCCESS                   Operation is successful.
  @retval Others                        An unexpected error occurred.
**/
EFI_STATUS
ClearWdtCallback (
  IN  EFI_PEI_SERVICES                  **PeiServices,
  IN  EFI_PEI_NOTIFY_DESCRIPTOR         *NotifyDescriptor,
  IN  VOID                              *Ppi
  )
{
  DEBUG ((DEBUG_INFO, "BiosSelfHealingPei: ClearWdtCallback is triggered\n"));

  //
  // [Lenovo BIOS Self-Healing Design Guidance Specification v1.9]
  //   2.2 Detect
  //     When BIOS flash happened, before erase/write BIOS should notify EC to set flag (WdtFlag should be
  //     keeped even EC power lost, such as stored in the EC EEPROM) to enable WDT, the WdtFlag should be
  //     cleared when EC received BIOS "cancel WDT" command.
  //     On every boot, EC check the WdtFlag flag to decide if need to start a 15s WDT,
  //     and stop the WDT if BIOS send "stop/cancel WDT" command.
  //
  OemSvcNotifyEcToClearWdt ();

  return EFI_SUCCESS;
}

/**
  Install FirmwareAuthHookDescriptor.

  @param  PeiServices                   General purpose services available to every PEIM.
  @param  NotifyDescriptor              Pointer of the notificaiton data structure.
  @param  Ppi                           Pointer of PPI.

  @retval EFI_SUCCESS                   Operation is successful.
  @retval Others                        An unexpected error occurred.
**/
EFI_STATUS
InstallFirmwareAuthHook (
  IN  EFI_PEI_SERVICES                  **PeiServices,
  IN  EFI_PEI_NOTIFY_DESCRIPTOR         *NotifyDescriptor,
  IN  VOID                              *Ppi
  )
{
  EFI_STATUS                            Status;

  //
  // If system in Top Swap mode and BIOS recovery hotkey is not pressed,
  // register callback to hook FirmwareAuthenticationPpi
  //
  if (PcdGetBool (PcdL05TopSwapEnable) && !PcdGetBool (PcdL05BiosRecoveryHotkeyFlag)) {
    Status = PeiServicesNotifyPpi (&mFirmwareAuthHookDescriptor);
    if (EFI_ERROR (Status)) {
      return Status;
    }
  }

  return EFI_SUCCESS;
}

/**
  Callback function to hook FirmwareAuthenticationPpi.

  @param  PeiServices                   General purpose services available to every PEIM.
  @param  NotifyDescriptor              Pointer of the notificaiton data structure.
  @param  Ppi                           Pointer of PPI.

  @retval EFI_SUCCESS                   Operation is successful.
  @retval Others                        An unexpected error occurred.
**/
EFI_STATUS
FirmwareAuthHookCallback (
  IN  EFI_PEI_SERVICES                  **PeiServices,
  IN  EFI_PEI_NOTIFY_DESCRIPTOR         *NotifyDescriptor,
  IN  VOID                              *Ppi
  )
{
  EFI_STATUS                            Status;
  EFI_PEI_PPI_DESCRIPTOR                *OldFirmwareAuthPpiDescriptor;
  FIRMWARE_AUTHENTICATION_PPI           *OldFirmwareAuthPpi;

  DEBUG ((DEBUG_INFO, "BiosSelfHealingPei: FirmwareAuthHookCallback is triggered\n"));

  Status = (*PeiServices)->LocatePpi (
                             PeiServices,
                             &gFirmwareAuthenticationPpiGuid,
                             0,
                             &OldFirmwareAuthPpiDescriptor,
                             (VOID **)&OldFirmwareAuthPpi
                             );

  if (Status == EFI_SUCCESS) {
    Status = (*PeiServices)->ReInstallPpi (
                               PeiServices,
                               OldFirmwareAuthPpiDescriptor,
                               &mFirmwareAuthPpiList
                               );
  } else {
    Status = (*PeiServices)->InstallPpi (
                               PeiServices,
                               &mFirmwareAuthPpiList
                               );
  }

  return EFI_SUCCESS;
}

/**
 SHA256 hash calculation.

 @param  Message                        The message data to be calculated.
 @param  MessageSize                    The size in byte of the message data.
 @param  Digest                         The caclulated HASH digest.

 @retval EFI_SUCCESS                    The HASH value is calculated.
 @retval EFI_SECURITY_VIOLATION         Failed to calculate the hash.
**/
EFI_STATUS
CalculateSha256Hash (
  IN  UINT8                             *Message,
  IN  UINTN                             MessageSize,
  OUT UINT8                             *Digest
  )
{
  EFI_STATUS                            Status;
  VOID                                  *HashCtx;
  UINTN                                 CtxSize;

  SetMem (Digest, SHA256_DIGEST_SIZE, 0);
  CtxSize = Sha256GetContextSize ();
  HashCtx = NULL;
  HashCtx = AllocatePool (CtxSize);
  if (HashCtx == NULL) {
    return EFI_OUT_OF_RESOURCES;
  }

  if (!Sha256Init (HashCtx)) {
    Status = EFI_SECURITY_VIOLATION;
    goto Done;
  }

  if(!Sha256Update (HashCtx, Message, MessageSize)) {
    Status = EFI_SECURITY_VIOLATION;
    goto Done;
  }

  if(!Sha256Final (HashCtx, Digest)) {
    Status = EFI_SECURITY_VIOLATION;
  } else {
    Status = EFI_SUCCESS;
  }

Done:
  FreePool (HashCtx);
  return Status;
}

/**
  This function looks up the FV hash in the FDM.

  @param  Fdm                           Point to the FDM header.
  @param  Guid                          Specify the GUID of an FV.

  @retval Non-null                      Success.
  @retval Null                          Not found.
**/
VOID *
LookUpFvHash (
  IN  UINT8                             *Fdm,
  IN  EFI_GUID                          *Guid
  )
{
  STATIC CHAR8                          FdmSignature[] = {'H', 'F', 'D', 'M'};
  UINT32                                EntrySize;
  UINT8                                 *EndPtr;
  H2O_FLASH_DEVICE_MAP_ENTRY            *Entry;

  if ((Fdm == NULL) || CompareMem (Fdm, FdmSignature, sizeof (FdmSignature))) {
    return NULL;
  }

  EntrySize  = ((H2O_FLASH_DEVICE_MAP_HEADER *) Fdm)->EntrySize;
  EndPtr     = Fdm + ((H2O_FLASH_DEVICE_MAP_HEADER *) Fdm)->Size;

  for (Entry = (H2O_FLASH_DEVICE_MAP_ENTRY *)(Fdm + ((H2O_FLASH_DEVICE_MAP_HEADER *) Fdm)->Offset);
       (UINT8 *)Entry < EndPtr;
       Entry = (H2O_FLASH_DEVICE_MAP_ENTRY *)(((UINT8 *) Entry) + EntrySize)) {
    if (CompareGuid ((EFI_GUID *) &Entry->RegionId, Guid)) {
      return ((UINT8 *) Entry + sizeof (H2O_FLASH_DEVICE_MAP_ENTRY));
    }
  }

  return NULL;
}

/**
 Firmware verification with SHA256 hash.

 @param  FirmwareFileData               Firmware file data buffer.
 @param  FirmwareFileSize               The firmware file size including signature.

 @retval EFI_SUCCESS                    The firmware verification is successful.
 @retval EFI_OUT_OF_RESOURCES           Out of resources.
 @retval EFI_SECURITY_VIOLATION         Failed to verify the firmware.
**/
EFI_STATUS
VerifyFirmware (
  IN  UINT8                             *FirmwareFileData,
  IN  UINTN                             FirmwareFileSize
  )
{
  EFI_STATUS                            Status;
  UINTN                                 BiosImageOffset;
  UINT8                                 *FvData;
  UINTN                                 FvSize;
  UINTN                                 HashSize;
  UINT8                                 *OnboardFvHash;
  UINT8                                 *Digest;

  Status        = EFI_SUCCESS;
  OnboardFvHash = NULL;

  //
  // Initial setting for Image data and hash information
  //
  BiosImageOffset = 0;
  HashSize        = SHA256_DIGEST_SIZE;

  //
  // Look up FV hash in onboard FDM
  //
  OnboardFvHash = LookUpFvHash ((UINT8 *)(UINTN) PcdGet64 (PcdH2OFlashDeviceMapStart), &gH2OFlashMapRegionDxeFvGuid);
  if (OnboardFvHash == NULL) {
    return EFI_SECURITY_VIOLATION;
  }

  //
  // Calculate FV hash of Image
  //
  FvData = FirmwareFileData + (BiosImageOffset + (UINTN) (PcdGet32 (PcdFlashFvMainBase) - PcdGet32 (PcdFlashAreaBaseAddress)));
  FvSize = (UINTN) PcdGet32 (PcdFlashFvMainSize);

  Digest = AllocatePool (HashSize);
  if (Digest == NULL){
    return EFI_SECURITY_VIOLATION;
  }

  Status = CalculateSha256Hash (FvData, FvSize, Digest);
  if (EFI_ERROR (Status)) {
    FreePool (Digest);
    return EFI_SECURITY_VIOLATION;
  }

  if (CompareMem (OnboardFvHash, Digest, HashSize)) {
    Status = EFI_SECURITY_VIOLATION;
  } else {
    Status = EFI_SUCCESS;
  }

  DEBUG ((DEBUG_INFO, "BiosSelfHealingPei: VerifyFirmware status: %r\n", Status));

  FreePool (Digest);
  return Status;
}

/**
  BIOS Self-Healing PEI Entry.

  @param  FfsHeader                     Points to the FFS file header to be checked.
  @param  PeiServices                   General purpose services available to every PEIM.

  @retval EFI_SUCCESS                   The operation completed successfully.
  @retval Others                        An unexpected error occurred.
**/
EFI_STATUS
EFIAPI
BiosSelfHealingPeiEntryPoint (
  IN  EFI_PEI_FILE_HANDLE               FfsHeader,
  IN  CONST EFI_PEI_SERVICES            **PeiServices
  )
{
  EFI_STATUS                            Status;
  H2O_CP_HANDLE                         CpHandle;

  Status   = EFI_SUCCESS;
  CpHandle = NULL;

#if (FixedPcdGet32 (PcdL05ChipsetName) == L05_CHIPSET_NAME_TIGERLAKE  || \
     FixedPcdGet32 (PcdL05ChipsetName) == L05_CHIPSET_NAME_ALDERLAKE)
  //
  // Sync Top Swap status to PCD
  //
  PcdSetBoolS (PcdL05TopSwapEnable, TopSwapStatus ());
  DEBUG ((DEBUG_INFO, "BiosSelfHealingPei: PcdL05TopSwapEnable status: %x\n", (UINTN) PcdGetBool (PcdL05TopSwapEnable)));
#endif

  //
  // Register checkpoint callback to update fv trust status after verifying region
  // (only execute in Top Swap mode)
  //
  if (PcdGetBool (PcdL05TopSwapEnable)) {
    if (FeaturePcdGet (PcdH2OBaseCpVerifyFvSupported)) {
      Status = H2OCpRegisterHandler (
                 &gH2OBaseCpVerifyFvGuid,
                 CpVerifyFvCallback,
                 H2O_CP_MEDIUM,
                 &CpHandle
                 );
      if (EFI_ERROR (Status)) {
        return Status;
      }
    }
  }

  //
  // Register callback to perform Self-Healing function initialization
  //
  Status = PeiServicesNotifyPpi (&mSelfHealingInitDescriptor);
  if (EFI_ERROR (Status)) {
    return Status;
  }

  return Status;
}