/** @file

;******************************************************************************
;* Copyright (c) 2017 - 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 <BootGuardRecoveryHook.h>
#include <Library/DebugLib.h>
#include <PiPei.h>
#include <Library/PeiServicesTablePointerLib.h>
#include <Library/HobLib.h>

EFI_PEI_NOTIFY_DESCRIPTOR
mBootGuardRecoveryNotifyList[] = {
  {
    EFI_PEI_PPI_DESCRIPTOR_NOTIFY_CALLBACK | EFI_PEI_PPI_DESCRIPTOR_TERMINATE_LIST,
    &gEfiPeiDeviceRecoveryModulePpiGuid,
    BootGuardRecoveryHookCallback
  }
};

DEVICE_RECOVERY_MODULE_HOOK_LIST        *mDeviceRecoveryModuleHookList = NULL;

/**
 To register the Notify List for Recovery with Boot Guard enabling.
 Arguments :
 FfsHeader   - Pointer of the FFS file header
 PeiServices - General purpose services available to every PEIM
 Returns :
 EFI_SUCCESS - Process complete
 Other       - Failed to register the Notify List

 @param [in]   FileHandle
 @param [in]   PeiServices      General purpose services available to every PEIM


**/
EFI_STATUS
BootGuardRecoveryHookEntry (
  IN EFI_PEI_FILE_HANDLE                   FileHandle,
  IN CONST EFI_PEI_SERVICES                **PeiServices
  )
{
  EFI_STATUS        Status;

  //
  // Shadow this PEIM to run from memory
  //
  if (!EFI_ERROR (PeiServicesRegisterForShadow (FileHandle))) {
    return EFI_SUCCESS;
  }

  Status = EFI_SUCCESS;
  if ( BootGuardPlatformLibDetermineBootState () != BootGuardBootStateLegacyBoot ) {
    Status = PeiServicesNotifyPpi (&mBootGuardRecoveryNotifyList[0]);
    ASSERT_EFI_ERROR (Status);
  }

  return Status;
}

/**
 To hook the Device Recovery Module PPIs.
 Arguments :
 PeiServices      - General purpose services available to every PEIM
 NotifyDescriptor - Pointer of the notificaiton data structure
 Ppi              - Pointer of PPI
 Returns :
 EFI_SUCCESS - Process complete
 Other       - Failed to allocate the record buffer

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


**/
EFI_STATUS
BootGuardRecoveryHookCallback (
  IN      EFI_PEI_SERVICES                 **PeiServices,
  IN      EFI_PEI_NOTIFY_DESCRIPTOR        *NotifyDescriptor,
  IN      VOID                             *Ppi
  )
{
  EFI_STATUS                                  Status;
  UINTN                                       Instance;
  EFI_PEI_DEVICE_RECOVERY_MODULE_PPI          *DeviceRecoveryModule;
  DEVICE_RECOVERY_MODULE_HOOK_LIST            **DeviceRecoveryModuleHookList;


  Instance                     = 0;
  DeviceRecoveryModule         = NULL;

  DeviceRecoveryModuleHookList = &mDeviceRecoveryModuleHookList;
  
  while ( TRUE ) {
    Status = PeiServicesLocatePpi (&gEfiPeiDeviceRecoveryModulePpiGuid, Instance, NULL, (VOID **)&DeviceRecoveryModule);
    if ( EFI_ERROR ( Status ) ) {
      if ( *DeviceRecoveryModuleHookList != NULL ) {
        *DeviceRecoveryModuleHookList = NULL;
      }

      break;
    }

    if ( *DeviceRecoveryModuleHookList == NULL ) {
      Status = PeiServicesAllocatePool (sizeof (DEVICE_RECOVERY_MODULE_HOOK_LIST), DeviceRecoveryModuleHookList);
      if ( EFI_ERROR ( Status ) || *DeviceRecoveryModuleHookList == NULL) {
        return Status;
      }

      ( *DeviceRecoveryModuleHookList )->DeviceRecoveryModule = NULL;
      ( *DeviceRecoveryModuleHookList )->LoadRecoveryCapsule  = NULL;
      ( *DeviceRecoveryModuleHookList )->NextPtr              = NULL;
    }

    if ( ( *DeviceRecoveryModuleHookList )->DeviceRecoveryModule != DeviceRecoveryModule ) {
      ( *DeviceRecoveryModuleHookList )->DeviceRecoveryModule = DeviceRecoveryModule;
      ( *DeviceRecoveryModuleHookList )->LoadRecoveryCapsule  = DeviceRecoveryModule->LoadRecoveryCapsule;

      DeviceRecoveryModule->LoadRecoveryCapsule               = BootGuardRecoveryHook;
    }

    DeviceRecoveryModuleHookList = &( ( *DeviceRecoveryModuleHookList )->NextPtr );
    Instance                     = Instance + 1;
  }

  return EFI_SUCCESS;
}

/**
 Hook to add the checking of the Boot Guard Revocation Values.
 Arguments :
 PeiServices     - General purpose services available to every PEIM
 This            - Device Recovery Module PPI
 CapsuleInstance - Target instance of the capsule
 Buffer          - Device Recovery Module PPI
 Returns :
 EFI_SUCCESS            - Process complete and the Revocation Values are valid
 EFI_NOT_FOUND          - The record is missing
 EFI_SECURITY_VIOLATION - Invalid Revocation Values
 Other                  - Failed to get required data

 @param [in, out] PeiServices   General purpose services available to every PEIM
 @param [in]   This             Device Recovery Module PPI
 @param [in]   CapsuleInstance  Target instance of the capsule
 @param [out]  Buffer           Device Recovery Module PPI


**/
EFI_STATUS
EFIAPI
BootGuardRecoveryHook (
  IN OUT  EFI_PEI_SERVICES                            **PeiServices,
  IN      EFI_PEI_DEVICE_RECOVERY_MODULE_PPI          *This,
  IN      UINTN                                       CapsuleInstance,
     OUT  VOID                                        *Buffer
  )
{
  EFI_STATUS                                  Status;
  UINTN                                       Instance;
  EFI_PEI_DEVICE_RECOVERY_MODULE_PPI          *DeviceRecoveryModule;
  DEVICE_RECOVERY_MODULE_HOOK_LIST            *DeviceRecoveryModuleHookList;
  UINTN                                       ImageLength;
  EFI_GUID                                    CapsuleType;
  UINT8                                       *Image;
  REVOCATION_VALUE                            Revocation;
  REVOCATION_VALUE                            CurrentRevocation;
  REVOCATION_VALUE                            *ArbSvnInfoHobData;
  EFI_HOB_GUID_TYPE                           *GuidHobPtr;

  Instance                     = 0;
  DeviceRecoveryModule         = NULL;
  DeviceRecoveryModuleHookList = NULL;
  ImageLength                  = 0;
  Image                        = NULL;
  Revocation.ACMSVN            = 0;
  Revocation.BPMSVN            = 0;
  Revocation.KMSVN             = 0;
  Revocation.KMID              = 0;
  Revocation.BootGuardStatus   = BootGuardBootStateLegacyBoot;

  CurrentRevocation.ACMSVN            = 0;
  CurrentRevocation.BPMSVN            = 0;
  CurrentRevocation.KMSVN             = 0;
  CurrentRevocation.KMID              = 0;
  CurrentRevocation.BootGuardStatus   = BootGuardBootStateLegacyBoot;
  
  GuidHobPtr                   = NULL;

  DeviceRecoveryModuleHookList = mDeviceRecoveryModuleHookList;
  while (DeviceRecoveryModuleHookList != NULL) {
    if (DeviceRecoveryModuleHookList->DeviceRecoveryModule == This) {
      break;
    }
    DeviceRecoveryModuleHookList = DeviceRecoveryModuleHookList->NextPtr;
  }

  if (DeviceRecoveryModuleHookList == NULL) {
    return EFI_NOT_FOUND;
  }

  Status = EFI_SUCCESS;
  if (!EFI_ERROR (Status)) {
    //
    // Try to restore SVN related data from HOB.
    //
    GuidHobPtr = GetFirstGuidHob (&gArbSvnInfoHobGuid);
    if (GuidHobPtr != NULL) {
      ArbSvnInfoHobData = (REVOCATION_VALUE *)GET_GUID_HOB_DATA (GuidHobPtr);
      CurrentRevocation.ACMSVN = ArbSvnInfoHobData->ACMSVN;
      CurrentRevocation.BPMSVN = ArbSvnInfoHobData->BPMSVN;
      CurrentRevocation.KMSVN  = ArbSvnInfoHobData->KMSVN;
      CurrentRevocation.KMID   = ArbSvnInfoHobData->KMID;
      Status                   = EFI_SUCCESS;
      CurrentRevocation.BootGuardStatus = ArbSvnInfoHobData->BootGuardStatus;
    } else {
      //
      //  Get ArbSvnInfo by Dex HECI command and by pass by PEI HOB
      //  This Hob should always exist
      //  If ArbSvnInfo Hob not exist, Btg SVN can't be verified
      //
      DEBUG ((DEBUG_ERROR, "Failed to get gArbSvnInfoHobGuid HOB\n"));
      return EFI_SECURITY_VIOLATION;
    }
  }

  //
  // It's not Bootguard SKU, pass Bootguard check
  //
  if (CurrentRevocation.BootGuardStatus == BootGuardBootStateLegacyBoot) {
    return EFI_SUCCESS;
  }

  if (!EFI_ERROR (Status)) {
    Status = This->GetRecoveryCapsuleInfo ( PeiServices, This, CapsuleInstance, &ImageLength, &CapsuleType );
  }
  if (!EFI_ERROR (Status)) {
    Status = (DeviceRecoveryModuleHookList->LoadRecoveryCapsule)(PeiServices, This, CapsuleInstance, Buffer);
  }
  if (!EFI_ERROR (Status)) {
    Image  = Buffer;
    Status = ExtractBIOSFromCapsule (&Image, &ImageLength);
  }

  if (!EFI_ERROR (Status)) {
    Status = BootGuardPlatformLibGetRevocationValues (Image, ImageLength, &Revocation, &CurrentRevocation.KMID);
  }


  if ( !EFI_ERROR ( Status ) ) {
    if ((Revocation.ACMSVN < CurrentRevocation.ACMSVN) || (Revocation.BPMSVN < CurrentRevocation.BPMSVN) || (Revocation.KMSVN < CurrentRevocation.KMSVN)) {
      Status = EFI_SECURITY_VIOLATION;
    }
  }

  return Status;
}

/**
 To extract the BIOS image from the capsule.
 Arguments :
 Image       - Pointer of the capsule, or of the extracted BIOS image
 ImageLength - Length of the capsule, or of the extracted BIOS image
 Returns :
 EFI_SUCCESS           - BIOS is extracted
 EFI_INVALID_PARAMETER - Invalid parameter inputted
 EFI_NOT_FOUND         - BIOS is not found

 @param [in, out] Image         Pointer of the capsule, or of the extracted BIOS image
 @param [in, out] ImageLength   Length of the capsule, or of the extracted BIOS image


**/
EFI_STATUS
ExtractBIOSFromCapsule (
  IN OUT  UINT8        **Image,
  IN OUT  UINTN        *ImageLength
  )
{
  EFI_STATUS                        Status;
  UINT8                             *Buffer;
  UINTN                             BufferLength;
//#ifdef SECURE_FLASH_SUPPORT
//BOOLEAN                           InRecoveryOnlyPEI;
  EFI_BOOT_MODE                     BootMode;
  CONST EFI_PEI_SERVICES            **PeiServices;
  EFI_PEI_PPI_DESCRIPTOR            *Descriptor;
  VOID                              *PPI;
  UINTN                             Index;
  ISFLASH_DATA_REGION_HEADER        *DataRegion;
//#endif
  UINTN                             FDBAR;
  UINT32                            *Ptr32;
  UINTN                             FRBA;
  UINTN                             RegionBase;
  UINTN                             RegionLimit;
  EFI_CAPSULE_HEADER                *CapsuleBuffer;
  UINT8                             *CurrentPtr;

  Buffer            = NULL;
  BufferLength      = 0;
//#ifdef SECURE_FLASH_SUPPORT
//InRecoveryOnlyPEI = FALSE;
  Descriptor        = NULL;
  PPI               = NULL;
  Index             = 0;
  DataRegion        = NULL;
//#endif
  FDBAR             = 0;
  Ptr32             = NULL;
  FRBA              = 0;
  RegionBase        = 0;
  RegionLimit       = 0;
  CapsuleBuffer     = NULL;
  CurrentPtr        = NULL;

  if ( ( Image == NULL ) || ( ImageLength == NULL ) ) {
    return EFI_INVALID_PARAMETER;
  }

  Status       = EFI_SUCCESS;
  Buffer       = *Image;
  BufferLength = *ImageLength;

  //
  // Get PeiService pointer
  //
  PeiServices = GetPeiServicesTablePointer ();

  if (FeaturePcdGet (PcdSecureFlashSupported)) {
    //InRecoveryOnlyPEI = EFI_ERROR ( PeiServicesLocatePpi ( &gEmuPeiPpiGuid, 0, &Descriptor, (VOID **)&PPI ) ) ? FALSE : TRUE;
    Status = (*PeiServices)->GetBootMode (PeiServices, &BootMode);
    //if ( ( !EFI_ERROR ( Status ) ) && ( !InRecoveryOnlyPEI ) ) {
    if ((!EFI_ERROR (Status)) && (BootMode == BOOT_IN_RECOVERY_MODE)) {
      //
      // Extract the whole binary
      //
      if (GetFirstGuidHob (&gSysFwUpdateProgressGuid) != NULL) { 

        CapsuleBuffer = (EFI_CAPSULE_HEADER*)Buffer;

        CurrentPtr = (UINT8 *)((UINTN)CapsuleBuffer + CapsuleBuffer->CapsuleImageSize - PcdGet32(PcdFlashAreaSize));
   
        if (CompareMem(CurrentPtr - sizeof(EFI_FFS_FILE_HEADER), &gCapsuleBiosImageFileGuid, sizeof(EFI_GUID)) == 0) {
          Buffer = CurrentPtr;
          BufferLength = PcdGet32(PcdFlashAreaSize);
        }
   
    
      } else {
        if ( ( Buffer[0] == ( UINT8 )( 'M' ) ) && ( Buffer[1] == ( UINT8 )( 'Z' ) ) ) {
          //
          // PE32
          //
          for ( Index = 0 ; Index < ( BufferLength - ISFLASH_IMAGE_SIGNATURE_SIZE ) ; Index = Index + 1 ) {
            if ( CompareMem ( ( UINT8 * )( UINTN )( ( UINTN )Buffer + Index ), ISFLASH_IMAGE_SIGNATURE, ISFLASH_IMAGE_SIGNATURE_SIZE ) == 0 ) {
              break;
            }
          }
          if ( Index >= ( BufferLength - ISFLASH_IMAGE_SIGNATURE_SIZE ) ) {
            Status = EFI_NOT_FOUND;
          }

          if ( !EFI_ERROR ( Status ) ) {
            DataRegion = ( ISFLASH_DATA_REGION_HEADER * )( UINTN )( ( UINTN )Buffer + Index );
            if ( ( Index + sizeof ( ISFLASH_DATA_REGION_HEADER ) + DataRegion->DataSize ) > BufferLength ) {
              Status = EFI_NOT_FOUND;
            }
          }

          if ( !EFI_ERROR ( Status ) ) {
            Buffer       = ( UINT8 * )( UINTN )( ( UINTN )( DataRegion ) + sizeof ( ISFLASH_DATA_REGION_HEADER ) );
            BufferLength = DataRegion->DataSize;
          }
        } else {
          BufferLength = BufferLength - SECURE_FLASH_SIGNATURE_SIZE;
        }
      }
    }
  }  
  

  if ( !EFI_ERROR ( Status ) ) {
    //
    // Extract the BIOS
    //
    FDBAR = ( UINTN )Buffer;
    Ptr32 = ( UINT32 * )( UINTN )( FDBAR + R_DESCRIPTOR_FDBAR_FLVALSIG );
    if ( *Ptr32 == V_DESCRIPTOR_FDBAR_FLVALSIG ) {
      Ptr32       = ( UINT32 * )( UINTN )( FDBAR + R_DESCRIPTOR_FDBAR_FLMAP0 );
      FRBA        = FDBAR + ( ( ( *Ptr32 & B_DESCRIPTOR_FDBAR_FLMAP0_FRBA ) >> N_DESCRIPTOR_FDBAR_FLMAP0_FRBA ) << 4 );

      Ptr32       = ( UINT32 * )( UINTN )( FRBA + R_DESCRIPTOR_FRBA_FLASH_REGION_BIOS );
      RegionBase  = ( *Ptr32 & B_DESCRIPTOR_FRBA_FLREG1_REGION_BASE ) >> N_DESCRIPTOR_FRBA_FLREG1_REGION_BASE;
      RegionLimit = ( *Ptr32 & B_DESCRIPTOR_FRBA_FLREG1_REGION_LIMIT ) >> N_DESCRIPTOR_FRBA_FLREG1_REGION_LIMIT;
      if ( RegionBase > RegionLimit ) {
        Status = EFI_NOT_FOUND;
      }

      if ( !EFI_ERROR ( Status ) ) {
        Buffer       = ( UINT8 * )( UINTN )( ( UINTN )Buffer + ( RegionBase << 12 ) );
        BufferLength = ( RegionLimit - RegionBase + 1 ) << 12;
      }
    }
  }

  *Image       = Buffer;
  *ImageLength = BufferLength;

  return Status;
}