/** @file

;******************************************************************************
;* Copyright (c) 2012 - 2017, 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 <Library/PeiServicesLib.h>
#include <Library/HobLib.h>
#include <Library/BaseMemoryLib.h>
#include <Library/VariableLib.h>

#include <Ppi/MemoryDiscovered.h>
#include <Ppi/MasterBootMode.h>
#include <SecureFlash.h>   //For SecureFlash feature

#define  MEMORY_SIZE_FOR_RESERVED          PcdGet32 (PcdReservedMemorySizeForSecureFlash)
#define  MEMORY_SIZE_4KB                   0x1000

EFI_STATUS
EFIAPI
SecureFlashPeiCallback (
  IN EFI_PEI_SERVICES           **PeiServices,
  IN EFI_PEI_NOTIFY_DESCRIPTOR  *NotifyDescriptor,
  IN VOID                       *Ppi
  );

EFI_PEI_NOTIFY_DESCRIPTOR  mMemoryReservedNotify = {
  (EFI_PEI_PPI_DESCRIPTOR_NOTIFY_CALLBACK | EFI_PEI_PPI_DESCRIPTOR_TERMINATE_LIST),
  &gEfiPeiMemoryDiscoveredPpiGuid,
  SecureFlashPeiCallback
  };

EFI_PEI_NOTIFY_DESCRIPTOR  mPeiMasterBootModeNotify = {
  (EFI_PEI_PPI_DESCRIPTOR_NOTIFY_CALLBACK | EFI_PEI_PPI_DESCRIPTOR_TERMINATE_LIST),
  &gEfiPeiMasterBootModePpiGuid,
  SecureFlashPeiCallback
  };

EFI_STATUS
EFIAPI
SecureFlashPeiMain (
  IN EFI_PEI_FILE_HANDLE             FileHandle,
  IN CONST EFI_PEI_SERVICES          **PeiServices
  )
{
  return PeiServicesNotifyPpi (&mMemoryReservedNotify);
}

/**
 (1)Set memory type to "EFI_RESOURCE_MEMORY_RESERVED" so it won't be destroy during DXE phase.
 (2)Save image information to Hob.


 @retval None

**/
EFI_STATUS
EFIAPI
SecureFlashPeiCallback (
  IN EFI_PEI_SERVICES           **PeiServices,
  IN EFI_PEI_NOTIFY_DESCRIPTOR  *NotifyDescriptor,
  IN VOID                       *Ppi
  )
{
  EFI_STATUS                      Status;
  EFI_STATUS                      VariableStatus;
  UINTN                           Size;
  UINTN                           VariableSize;
  IMAGE_INFO                      ImageInfo;
  EFI_PHYSICAL_ADDRESS            ReservedMemAddress;
  EFI_BOOT_MODE                   BootMode;

  Status = PeiServicesLocatePpi (
             &gEfiPeiMasterBootModePpiGuid,
             0,
             NULL,
             NULL
             );
  if (EFI_ERROR (Status)) {
    return PeiServicesNotifyPpi (&mPeiMasterBootModeNotify);
  }

  BootMode = GetBootModeHob ();
  if (BootMode == BOOT_IN_RECOVERY_MODE) {
    return EFI_SUCCESS;
  }

  if (PcdGet32 (PcdReservedMemorySizeForSecureFlash) == 0) {
    return EFI_SUCCESS;
  }

  Size = MEMORY_SIZE_FOR_RESERVED / MEMORY_SIZE_4KB;
  if ((MEMORY_SIZE_FOR_RESERVED % MEMORY_SIZE_4KB) != 0) {
    Size++;
  }

  VariableSize = sizeof (ImageInfo);
  VariableStatus = CommonGetVariable (
                     SECURE_FLASH_INFORMATION_NAME,
                     &gSecureFlashInfoGuid,
                     &VariableSize,
                     &ImageInfo
                     );
  if (!(EFI_ERROR(VariableStatus)) && (ImageInfo.ImageSize > MEMORY_SIZE_FOR_RESERVED)) {
    return EFI_OUT_OF_RESOURCES;
  }

  //
  // Allocate a reserved memory for Secure Flash
  //
  Status = (*PeiServices)->AllocatePages (
                             (CONST EFI_PEI_SERVICES **)PeiServices,
                             EfiACPIMemoryNVS,
                             Size,
                             &ReservedMemAddress
                             );
  if (EFI_ERROR (Status)) {
    return Status;
  }

  if (!(EFI_ERROR(VariableStatus)) && (ImageInfo.FlashMode)) {
    CopyMem ((VOID*)(UINTN)ReservedMemAddress, (VOID*)(UINTN)ImageInfo.ImageAddress, ImageInfo.ImageSize);
  } else {
    ImageInfo.ImageSize = MEMORY_SIZE_FOR_RESERVED;
  }
  ImageInfo.ImageAddress = (UINTN)ReservedMemAddress;
  //
  // Save the data of ImageAddress and ImageSize to Hob.
  // Because it cannot SetVariable in PEI phase.
  //
  BuildCvHob (ImageInfo.ImageAddress, ImageInfo.ImageSize);

  return EFI_SUCCESS;
}