/** @file
  Instance of SMM/DXE BIOS Self-Healing Services Library.

;******************************************************************************
;* 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 <PiSmm.h>
#include <L05Config.h>
#include <L05BiosSelfHealingConfig.h>
#include <Library/UefiBootServicesTableLib.h>
#include <Library/SmmServicesTableLib.h>
#include <Library/BaseMemoryLib.h>
#include <Library/MemoryAllocationLib.h>
#include <Library/DebugLib.h>
#include <Protocol/SmmBase2.h>
#include <Protocol/SmmFwBlockService.h>
#include <Protocol/L05Variable.h>

BOOLEAN                                 mInSmm = FALSE;
EFI_SMM_SYSTEM_TABLE2                   *mSmst = NULL;

/**
  Wrap original FreePool call in order to decrease code length (with setting back Buffer to NULL).

  @param  Buffer                        Pointer to the allocated memory address.
**/
VOID
SafeFreePool (
  IN OUT VOID                           **Buffer
  )
{
  if (Buffer != NULL && *Buffer != NULL) {
    FreePool (*Buffer);
    *Buffer = NULL;
  }
}

/**
  Get BIOS Self-Healing Flag.

  @retval EFI_SUCCESS                   Get flag successfully.
  @retval Others                        An unexpected error occurred.
**/
EFI_STATUS
GetBiosSelfHealingFlag (
  IN  UINT8                             *BiosSelfHealingFlag
  )
{
  EFI_STATUS                            Status;
#ifdef L05_SPECIFIC_VARIABLE_SERVICE_ENABLE
  EFI_L05_VARIABLE_PROTOCOL             *L05Variable;
  UINT32                                DataLength;
#else
  EFI_L05_EEPROM_MAP_120                *EepromBuffer;
#endif

  *BiosSelfHealingFlag = 0;
  Status               = EFI_SUCCESS;

#ifdef L05_SPECIFIC_VARIABLE_SERVICE_ENABLE
  L05Variable = NULL;

  if (!mInSmm) {
    Status = gBS->LocateProtocol (
                    &gEfiL05VariableProtocolGuid,
                    NULL,
                    &L05Variable
                    );
  } else {
    Status = mSmst->SmmLocateProtocol (
                      &gEfiL05VariableProtocolGuid,
                      NULL,
                      &L05Variable
                      );
  }

  if (EFI_ERROR (Status)) {
    return Status;
  }

  DataLength = sizeof (UINT8);
  Status = L05Variable->GetVariable (
                          L05Variable,
                          &gL05BiosSelfHealingFlagGuid,
                          &DataLength,
                          BiosSelfHealingFlag
                          );

  if (EFI_ERROR (Status)) {
    return Status;
  }

#else
  EepromBuffer = (EFI_L05_EEPROM_MAP_120 *)(UINTN) PcdGet32 (PcdFlashFvEepromBase);

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

  *BiosSelfHealingFlag = EepromBuffer->BiosSelfHealingFlag[0];
#endif

  return EFI_SUCCESS;
}

/**
  Set BIOS Self-Healing Flag.

  @retval EFI_SUCCESS                   Set flag successfully.
  @retval Others                        An unexpected error occurred.
**/
EFI_STATUS
SetBiosSelfHealingFlag (
  IN  UINT8                             BiosSelfHealingFlag
  )
{
  EFI_STATUS                            Status;
#ifdef L05_SPECIFIC_VARIABLE_SERVICE_ENABLE
  EFI_L05_VARIABLE_PROTOCOL             *L05Variable;
  UINT32                                DataLength;
#else
  EFI_SMM_FW_BLOCK_SERVICE_PROTOCOL     *SmmFwBlockService;
  UINTN                                 EepromBase;
  UINTN                                 EepromSize;
  VOID                                  *Buffer;
  UINTN                                 BufferSize;
#endif

  Status = EFI_SUCCESS;

#ifdef L05_SPECIFIC_VARIABLE_SERVICE_ENABLE
  L05Variable = NULL;

  if (!mInSmm) {
    Status = gBS->LocateProtocol (
                    &gEfiL05VariableProtocolGuid,
                    NULL,
                    &L05Variable
                    );
  } else {
    Status = mSmst->SmmLocateProtocol (
                      &gEfiL05VariableProtocolGuid,
                      NULL,
                      &L05Variable
                      );
  }

  if (EFI_ERROR (Status)) {
    return Status;
  }

  DataLength = sizeof (UINT8);
  Status = L05Variable->SetVariable (
                          L05Variable,
                          &gL05BiosSelfHealingFlagGuid,
                          DataLength,
                          &BiosSelfHealingFlag
                          );

  if (EFI_ERROR (Status)) {
    return Status;
  }

#else
  SmmFwBlockService = NULL;
  EepromBase        = 0;
  EepromSize        = 0;
  Buffer            = NULL;
  BufferSize        = 0;

  if (!mInSmm) {
    Status = gBS->LocateProtocol (
                    &gEfiSmmFwBlockServiceProtocolGuid,
                    NULL,
                    &SmmFwBlockService
                    );
  } else {
    Status = mSmst->SmmLocateProtocol (
                      &gEfiSmmFwBlockServiceProtocolGuid,
                      NULL,
                      &SmmFwBlockService
                      );
  }

  if (EFI_ERROR (Status)) {
    return Status;
  }

  EepromBase = (UINTN) PcdGet32 (PcdFlashFvEepromBase);
  EepromSize = (UINTN) PcdGet32 (PcdFlashFvEepromSize);

  if ((EepromBase == 0) || (EepromSize == 0)) {
    return EFI_NOT_FOUND;
  }

  Buffer = AllocatePages (EFI_SIZE_TO_PAGES (EepromSize));

  if (Buffer == NULL) {
    return EFI_OUT_OF_RESOURCES;
  }

  BufferSize = EFI_PAGES_TO_SIZE (EFI_SIZE_TO_PAGES (EepromSize));
  ZeroMem (Buffer, BufferSize);
  CopyMem (Buffer, (VOID *) EepromBase, BufferSize);

  Status = SmmFwBlockService->EraseBlocks (
                                SmmFwBlockService,
                                EepromBase,
                                &BufferSize
                                );

  if (EFI_ERROR (Status)) {
    return Status;
  }

  ((EFI_L05_EEPROM_MAP_120 *)Buffer)->BiosSelfHealingFlag[0] = BiosSelfHealingFlag;

  Status = SmmFwBlockService->Write (
                                SmmFwBlockService,
                                EepromBase,
                                &BufferSize,
                                Buffer
                                );

  if (EFI_ERROR (Status)) {
    return Status;
  }

  FreePages (Buffer, EFI_SIZE_TO_PAGES (EepromSize));
#endif

  return EFI_SUCCESS;
}

/**
 Constructor which locates the relevant protocols.

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

 @retval EFI_SUCCESS                    The operation completed successfully.
 @retval Others                         An unexpected error occurred.
*/
EFI_STATUS
EFIAPI
SmmDxeBiosSelfHealingLibConstructor (
  IN  EFI_HANDLE                        ImageHandle,
  IN  EFI_SYSTEM_TABLE                  *SystemTable
  )
{
  EFI_STATUS                            Status;
  EFI_SMM_BASE2_PROTOCOL                *SmmBase2;

  Status   = EFI_SUCCESS;
  SmmBase2 = NULL;

  //
  // SMM check
  //
  Status = gBS->LocateProtocol (
                  &gEfiSmmBase2ProtocolGuid,
                  NULL,
                  &SmmBase2
                  );

  if (!EFI_ERROR (Status)) {
    Status = SmmBase2->InSmm (SmmBase2, &mInSmm);
  } else {
    mInSmm = FALSE;
  }

  if (!mInSmm) {
    return EFI_SUCCESS;
  }

  //
  // We're now in SMM
  // Get SMM System Table
  //
  Status = SmmBase2->GetSmstLocation (SmmBase2, &mSmst);

  if (EFI_ERROR (Status)) {
    return Status;
  }

  return EFI_SUCCESS;
}