/** @file
 Implement the Chipset Servcie IHISI FMTS subfunction for this driver.

***************************************************************************
* Copyright (c) 2019 - 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 "Fmts.h"
#include "MeCapsuleSmm.h"
#include <SecureFlash.h>
#include <Guid/ZeroGuid.h>

#define  FMTSWriteUseIntelLib       0x32

IHISI_SMI_SUB_FUNCTION
FMTS_FUNCTION_TABLE[] = {
  { FMTSWriteUseIntelLib,        FmtsUseIntelLib,          IhisiNormalPriority}
};

EFI_BOOT_SERVICES           *BS      = NULL;
EFI_PHYSICAL_ADDRESS        *mCollectImageBuffer = NULL;
UINTN                       mCollectImageSize = 0;
BOOLEAN                     mIsMeUpdate = FALSE;
//
//  MSFT compiler uses this variable to inform itself whether floating
//  point support is loaded.
//
static int                         _fltused;

VOID
SetFlag(
  VOID
  )
{
  UINTN                          Size;
  SYSTEM_FIRMWARE_COMPONENT      Component;
  EFI_STATUS                     Status;

  Size = sizeof(SYSTEM_FIRMWARE_COMPONENT);
  Status = CommonGetVariable (
             FW_UPDATEINFO_VARIABLE_NAME,
             &gSysFwUpdateProgressGuid,
             &Size,
             &Component
             );

  if (!EFI_ERROR (Status)) {
    if(Component == UpdatingTypeMax){
      Component = UpdatingBios;
      Status = SetVariableToSensitiveVariable (
                 FW_UPDATEINFO_VARIABLE_NAME,
                 &gSysFwUpdateProgressGuid,
                 EFI_VARIABLE_BOOTSERVICE_ACCESS | EFI_VARIABLE_NON_VOLATILE,
                 sizeof(SYSTEM_FIRMWARE_COMPONENT),
                 &Component
                 );
    } else {
      Status = SetVariableToSensitiveVariable (
                 FW_UPDATEINFO_VARIABLE_NAME,
                 &gSysFwUpdateProgressGuid,
                 0,
                 0,
                 NULL
                 );
    }
    DEBUG ((DEBUG_INFO, "Set CapsuleInfo to UpdatingMe (%r)\n", Status));
  }
}


VOID
DisplaySendStatus (
  UINT32   BytesSent,
  UINT32   BytestobeSent
  )
{
//
//  Un-support print any message on flash mode
//
//  float   value = (BytesSent / BytestobeSent) * 100;
//  UINT32  pValue = (UINT32)value;
//
//  if (pValue != 100) {
//      Print (L"Sending the update image to FW for verification:  [ %d%% ]\r", pValue);
//  } else {
//      Print (L"Sending the update image to FW for verification:  [ COMPLETE ] \n");
//  }
  return;
}

EFI_STATUS
EFIAPI
FmtsUseIntelLib(
  VOID
  )
{
  EFI_STATUS                     Status;
  VOID                           *WriteDataBuffer;
  _UUID                          OemId;
//  CHAR8                          UsePassword[9];
  FMTS_INPUT_DATA_STRUCTURE      *InputData;
  FMTS_OUTPUT_DATA_STRUCTURE     *OutputData;
  UINT8                          FunctionType;
  UINT32                         IntelNativeMESize;
  UINT32                         PercentWritten;
  UINTN                          WriteDataBufferSize;

  //
  // Local Variable Initialize
  //
  Status            = EFI_SUCCESS;
  WriteDataBuffer   = NULL;
  InputData         = NULL;
  OutputData        = NULL;
  FunctionType      = 0x00;
  IntelNativeMESize = 0x00;
  PercentWritten    = 0x00;
  ZeroMem (&OemId, sizeof (OemId));

  InputData  = (FMTS_INPUT_DATA_STRUCTURE *)(UINTN) mH2OIhisi->ReadCpuReg32 (EFI_SMM_SAVE_STATE_REGISTER_RCX);
  WriteDataBuffer = (VOID *)(UINTN)mH2OIhisi->ReadCpuReg32 (EFI_SMM_SAVE_STATE_REGISTER_RSI);
  WriteDataBufferSize = (UINTN)mH2OIhisi->ReadCpuReg32 (EFI_SMM_SAVE_STATE_REGISTER_RDI);

  if (!mH2OIhisi->BufferInCmdBuffer ((VOID *) InputData, sizeof(FMTS_INPUT_DATA_STRUCTURE)) || mH2OIhisi->BufferOverlapSmram  ((VOID *) WriteDataBuffer, WriteDataBufferSize)) {
    return EFI_UNSUPPORTED;
  }

  FunctionType = InputData->FunctionType;
  OemId = InputData->OemId;
  IntelNativeMESize = InputData->DataSize;
  //
  // Check pre-reserved memory is enough for the ME firmware.
  //
  if (mCollectImageBuffer == 0 || IntelNativeMESize > PcdGet32 (PcdMeFirmwareUpdateReservedMemorySize)) {
    return EFI_BUFFER_TOO_SMALL;
  }

  if (FeaturePcdGet(PcdH2OIhisiCmdBufferSupported)) {
    if (WriteDataBufferSize != 0) {
      //
      // When EDI is non-zero, the data buffer only contains a part of image. BIOS need to collect each part
      // of image until the total size match the value in ECX Data Size field.
      //
      if (mCollectImageSize != IntelNativeMESize) {
        CopyMem((VOID *)(((UINT8*)(UINTN)mCollectImageBuffer) + mCollectImageSize), WriteDataBuffer, WriteDataBufferSize);
        mCollectImageSize += WriteDataBufferSize;
      }
    }
  }

  OutputData  = (FMTS_OUTPUT_DATA_STRUCTURE *)(UINTN) mH2OIhisi->ReadCpuReg32 (EFI_SMM_SAVE_STATE_REGISTER_RCX);
  OutputData->Signature = FMTS_OUTPUT_SIGNATURE;
  OutputData->StructureSize = sizeof (FMTS_OUTPUT_DATA_STRUCTURE);
  OutputData->Percentage = 0x0;
  OutputData->ReturnStatusCode = (UINT32)EFI_SUCCESS;

  if (FeaturePcdGet(PcdH2OIhisiCmdBufferSupported)) {

    switch (FunctionType) {
      case TYPE_ME:
        if (mCollectImageSize != IntelNativeMESize) {
          return EFI_SUCCESS;
        }
        mIsMeUpdate = TRUE;
        Status = FmtsUseIntelLibWrite((VOID *)mCollectImageBuffer,(UINT32) IntelNativeMESize, OemId);
        mCollectImageSize = 0;
        OutputData->ReturnStatusCode = (UINT32)Status;  // Status
        if (Status != EFI_SUCCESS) {
          return Status;
        }
        break;
      case TYPE_PERCENTAGE:
        Status = FmtsUseIntelLibWriteUpdateStages (&PercentWritten);
        OutputData->ReturnStatusCode = (UINT32)Status;  // Status
        OutputData->Percentage = (UINT8)PercentWritten; // percentage
        if (PcdGetBool (PcdMeResiliencyEnable)) {
          if(mIsMeUpdate == TRUE && PercentWritten == 100) {
            SetFlag ();
          }
        }
        break;
      case TYPE_PASSWARD:
          break;
      case TYPE_ISH:
      case TYPE_IOM:
        case TYPE_MGPHY:
      case TYPE_TBT:
        if (mCollectImageSize != IntelNativeMESize) {
          return EFI_SUCCESS;
        }
        Status = FmtsUseIntelLibPartialWrite ((VOID *) mCollectImageBuffer, (UINT32) IntelNativeMESize, OemId, FunctionType);
        mCollectImageSize = 0;
        OutputData->ReturnStatusCode = (UINT32) Status;
        if (Status != EFI_SUCCESS) {
          return Status;
        }

        OutputData->Percentage = 100;
        break;
      case TYPE_PDT:
        if (mCollectImageSize != IntelNativeMESize) {
          return EFI_SUCCESS;
        }
        Status = FmtsUseIntelLibSendPdt ((VOID *) mCollectImageBuffer, (UINT32) IntelNativeMESize);
        mCollectImageSize = 0;
        OutputData->ReturnStatusCode = (UINT32) Status;  // Status
        if (Status != EFI_SUCCESS) {
          return Status;
        }
        break;
      default:
        break;
    }

  } else {

    switch (FunctionType) {
      case TYPE_ME:
        WriteDataBuffer = (VOID *)(UINTN) mH2OIhisi->ReadCpuReg32(EFI_SMM_SAVE_STATE_REGISTER_RSI);
        // OutputData->Percentage = 0x0;  // percentage
        Status = FmtsUseIntelLibWrite(WriteDataBuffer,(UINT32) IntelNativeMESize, OemId);
        OutputData->ReturnStatusCode = (UINT32)Status;  // Status
        if (Status != EFI_SUCCESS) {
          return Status;
        }
        break;
      case TYPE_PERCENTAGE:
        Status = FmtsUseIntelLibWriteUpdateStages (&PercentWritten);
        OutputData->ReturnStatusCode = (UINT32)Status;  // Status
        OutputData->Percentage = (UINT8)PercentWritten;  // percentage
        break;
      case TYPE_PASSWARD:
        break;
      case TYPE_ISH:
      case TYPE_IOM:
        case TYPE_MGPHY:
      case TYPE_TBT:
        WriteDataBuffer = (VOID *)(UINTN) mH2OIhisi->ReadCpuReg32(EFI_SMM_SAVE_STATE_REGISTER_RSI);
        Status = FmtsUseIntelLibPartialWrite (WriteDataBuffer, (UINT32) IntelNativeMESize, OemId, FunctionType);
        OutputData->ReturnStatusCode = (UINT32) Status;
        if (Status != EFI_SUCCESS) {
          // OutputData->Percentage = 0x00;
          return Status;
        }
        OutputData->Percentage = 100;
        break;
      case TYPE_PDT:
        WriteDataBuffer = (VOID *)(UINTN) mH2OIhisi->ReadCpuReg32(EFI_SMM_SAVE_STATE_REGISTER_RSI);
        //OutputData->Percentage = 0x0;  // percentage
        Status = FmtsUseIntelLibSendPdt (WriteDataBuffer, (UINT32) IntelNativeMESize);
        OutputData->ReturnStatusCode = (UINT32)Status;  // Status
        if (Status != EFI_SUCCESS) {
          return Status;
        }
        break;
      default:
        break;
    }
  }

  return EFI_SUCCESS;
}

EFI_STATUS
EFIAPI
FmtsUseIntelLibInit (
  VOID
  )
{
  EFI_STATUS            Status;
  EFI_STATUS            StatusVar;
  EFI_BOOT_MODE         BootMode;
  UINTN                 Size;
  UINT64                OsIndications;
  BOOLEAN               ReservedMemory;
  IMAGE_INFO            ImageInfo;

  ReservedMemory = FALSE;

  Status = mH2OIhisi->RegisterCommand (FMTSWriteUseIntelLib, FmtsUseIntelLib, IhisiNormalPriority);
  ASSERT_EFI_ERROR (Status);

  if (!EFI_ERROR (Status) && (mCollectImageBuffer == 0)){
    BootMode = GetBootModeHob ();
    //
    // Check Boot mode and Recovery status
    //
    if (BootMode == BOOT_ON_FLASH_UPDATE || BootMode == BOOT_IN_RECOVERY_MODE) {
      ReservedMemory = TRUE;
    } else {
      Size = sizeof (IMAGE_INFO);
      StatusVar = gRT->GetVariable (
                         SECURE_FLASH_INFORMATION_NAME,
                         &gSecureFlashInfoGuid,
                         NULL,
                         &Size,
                         &ImageInfo
                         );

      if ((StatusVar == EFI_SUCCESS) && (ImageInfo.FlashMode)) {
        ReservedMemory = TRUE;
      }
    }
    //
    // Check Capsule enable or not.
    //
    Size          = sizeof (UINT64);
    OsIndications = 0;
    StatusVar = gRT->GetVariable (
                       L"OsIndications",
                       &gEfiGlobalVariableGuid,
                       NULL,
                       &Size,
                       &OsIndications
                       );

    if (!EFI_ERROR (StatusVar) && (OsIndications & EFI_OS_INDICATIONS_FILE_CAPSULE_DELIVERY_SUPPORTED)) {
      ReservedMemory = TRUE;
    }
    //
    // Check whether ME is in recovery mode
    //
    if (IsMeStateRecovery ()) {
      if (!IsMaxRecoveryTry (FALSE)) {
        ReservedMemory = TRUE;
        DEBUG ((DEBUG_INFO, "ME current state is in recovery mode and not exceed the max recovery try\n"));
      }
    }

    if (ReservedMemory) {
      Size = PcdGet32 (PcdMeFirmwareUpdateReservedMemorySize);
      Status = gBS->AllocatePool (
                      EfiReservedMemoryType,
                      Size,
                      (VOID **) (UINTN) (&mCollectImageBuffer)
                      );
      if (!EFI_ERROR (Status)){
        ZeroMem ((VOID *) mCollectImageBuffer, Size);
      }
    }
  }

  return Status;
}

EFI_STATUS
EFIAPI
FmtsUseIntelLibWrite(
  VOID                          *MeDataBuffer,
  UINT32                         MeDataSize,
  _UUID                          OemId
//  CHAR8                          *UserPassword
  )
{
  UINT32                         Status;
  UINT32                         FwUpdateApiStatus;
  UINT16                         EnabledState;
  UINT32                         FlashFwType;
  UINT32                         UpdateFwType;
  UINT32                         FlashPchSku;
  UINT32                         UpdatePchSku;
  UINT32                         FlashProductSegment;
  UINT32                         UpdateProductSegment;
  _UUID                          Uuid;

  EnabledState    = 0;
  FlashFwType = FWU_FW_TYPE_INVALID;
  UpdateFwType = FWU_FW_TYPE_INVALID;
  FlashPchSku = FWU_PCH_SKU_INVALID;
  UpdatePchSku = FWU_PCH_SKU_INVALID;
  FlashProductSegment = 0;
  UpdateProductSegment = 0;
  BS = gBS;

  FwUpdateApiStatus = FwuEnabledState (&EnabledState);
  if (FwUpdateApiStatus != EFI_SUCCESS) {
    DEBUG ((DEBUG_ERROR, "UpdateMeByHeci: FwuEnabledState failed: %r.\n", FwUpdateApiStatus));
    return EFI_DEVICE_ERROR;
  }
  DEBUG ((DEBUG_INFO, "UpdateMeByHeci: FwuEnabledState is 0x%x.\n", EnabledState));

  switch (EnabledState) {
    case FW_UPDATE_DISABLED:
      DEBUG ((DEBUG_ERROR, "UpdateMeByHeci: FWUpdate is disabled.\n"));
      return EFI_DEVICE_ERROR;

    case FW_UPDATE_ENABLED:
      break;

    default:
      break;
  }

  Status = FwuCheckCompatibilityFromBuffer((UINT8*) MeDataBuffer, (UINT32) MeDataSize);
  if (Status != EFI_SUCCESS) {
    DEBUG ((DEBUG_ERROR, "UpdateMeByHeci: FwuCheckCompatibilityFromBuffer failed: %r.\n", Status));
    return Status;
  }
  //
  // FW Type
  //
  Status = FwuFwType (&FlashFwType);
  if (Status != EFI_SUCCESS) {
    DEBUG ((DEBUG_ERROR, "UpdateMeByHeci: FwuFwType failed: %r.\n", Status));
    return Status;
  }
  Status = FwuFwTypeFromBuffer ((UINT8*) MeDataBuffer, (UINT32) MeDataSize, &UpdateFwType);
  if (Status != EFI_SUCCESS) {
    DEBUG ((DEBUG_ERROR, "UpdateMeByHeci: FwuFwTypeFromBuffer failed: %r.\n", Status));
    return Status;
  }
  if (FlashFwType != UpdateFwType) {
    DEBUG ((DEBUG_INFO, "UpdateMeByHeci: FlashFwType: 0x%x, UpdateFwType: 0x%x.\n", FlashFwType, UpdateFwType));
    return EFI_INCOMPATIBLE_VERSION;
  }
  //
  // PCH SKU
  //
  Status = FwuPchSku (&FlashPchSku);
  if (Status != EFI_SUCCESS) {
    DEBUG ((DEBUG_ERROR, "UpdateMeByHeci: FwuPchSku failed: %r.\n", Status));
    return Status;
  }
  Status = FwuPchSkuFromBuffer ((UINT8*) MeDataBuffer, (UINT32) MeDataSize, &UpdatePchSku);
  if (Status != EFI_SUCCESS) {
    DEBUG ((DEBUG_ERROR, "UpdateMeByHeci: FwuPchSkuFromBuffer failed: %r.\n", Status));
    return Status;
  }
  if (FlashPchSku != UpdatePchSku) {
    DEBUG ((DEBUG_INFO, "UpdateMeByHeci: FlashPchSku: 0x%x, UpdatePchSku: 0x%x.\n", FlashPchSku, UpdatePchSku));
    return EFI_INCOMPATIBLE_VERSION;
  }
  //
  // Product Segment
  //
  Status = FwuProductSegmentFromFlash(&FlashProductSegment);
  if (Status != EFI_SUCCESS) {
    DEBUG ((DEBUG_ERROR, "UpdateMeByHeci: FwuProductSegmentFromFlash failed: %r.\n", Status));
    return Status;
  }
  Status = FwuProductSegmentFromBuffer ((UINT8*) MeDataBuffer, (UINT32) MeDataSize, &UpdateProductSegment);
  if (Status != EFI_SUCCESS) {
    DEBUG ((DEBUG_ERROR, "UpdateMeByHeci: FwuProductSegmentFromBuffer failed: %r.\n", Status));
    return Status;
  }
  if (FlashProductSegment != UpdateProductSegment) {
    DEBUG ((DEBUG_INFO, "UpdateMeByHeci: FlashProductSegment: 0x%x, UpdateProductSegment: 0x%x.\n", FlashProductSegment, UpdateProductSegment));
    return EFI_INCOMPATIBLE_VERSION;
  }

  DEBUG ((DEBUG_INFO, "UpdateMeByHeci: Executing Full FW Update.\n"));

  Status = FwuOemId (&Uuid);
  DEBUG ((DEBUG_INFO, "Firmware OEM ID: %g [%r]\r\n", &Uuid, Status));
  if (!EFI_ERROR (Status) && !CompareGuid ((EFI_GUID *) &Uuid, &gZeroGuid)) {
    Status = FwuFullUpdateFromBuffer ((UINT8*)MeDataBuffer, (UINT32) MeDataSize, &Uuid, &DisplaySendStatus);
  } else {
    Status = FwuFullUpdateFromBuffer ((UINT8*)MeDataBuffer, (UINT32) MeDataSize, NULL, &DisplaySendStatus);
  }
  
  if (Status != EFI_SUCCESS) {
    DEBUG ((DEBUG_ERROR, "UpdateMeByHeci: FwuFullUpdateFromBuffer failed: %r.\n", Status));
    return Status;
  }
  return EFI_SUCCESS;
}

EFI_STATUS
EFIAPI
FmtsUseIntelLibPartialWrite(
  VOID                          *MeDataBuffer,
  UINT32                         MeDataSize,
  _UUID                          OemId,
  UINT8                          FunctionType
  )
{
  UINT32                         Status;
  UINT32                         FwUpdateApiStatus;
  UINT16                         EnabledState;
  UINT16                         BufferMajor;
  UINT16                         BufferMinor;
  UINT16                         BufferHotfix;
  UINT16                         BufferBuild;

  EnabledState  = 0;
  BufferMajor   = 0;
  BufferMinor   = 0;
  BufferHotfix  = 0;
  BufferBuild   = 0;
  BS            = gBS;
  Status        = 0;

  FwUpdateApiStatus = FwuEnabledState (&EnabledState);
  if (FwUpdateApiStatus != EFI_SUCCESS) {
    DEBUG ((DEBUG_ERROR, "UpdateMeByHeci: FwuEnabledState failed: %r.\n", FwUpdateApiStatus));
    return EFI_DEVICE_ERROR;
  }
  DEBUG ((DEBUG_INFO, "UpdateMeByHeci: FwuEnabledState is 0x%x.\n", EnabledState));

  switch (EnabledState) {
    case FW_UPDATE_DISABLED:
      DEBUG ((DEBUG_ERROR, "UpdateMeByHeci: FWUpdate is disabled.\n"));
      return EFI_DEVICE_ERROR;

    case FW_UPDATE_ENABLED:
      break;
    default:
      break;
  }

  switch (FunctionType) {
    case TYPE_ISH:
      Status = FwuPartialUpdateFromBuffer((UINT8 *)MeDataBuffer, MeDataSize, FPT_PARTITION_NAME_ISHC, &DisplaySendStatus);
      break;
    case TYPE_IOM:
      Status = FwuPartialUpdateFromBuffer((UINT8 *)MeDataBuffer, MeDataSize, FPT_PARTITION_NAME_IOMP, &DisplaySendStatus);
      break;
    case TYPE_MGPHY:
      Status = FwuPartialUpdateFromBuffer((UINT8 *)MeDataBuffer, MeDataSize, FPT_PARTITION_NAME_NPHY, &DisplaySendStatus);
      break;
    case TYPE_TBT:
      Status = FwuPartialUpdateFromBuffer((UINT8 *)MeDataBuffer, MeDataSize, FPT_PARTITION_NAME_TBTP, &DisplaySendStatus);
      break;
    default:
      break;
  }
  if (Status != EFI_SUCCESS) {
    return Status;
  }
  return EFI_SUCCESS;
}

EFI_STATUS
EFIAPI
FmtsUseIntelLibSendPdt(
  VOID                           *MeDataBuffer,
  UINT32                         MeDataSize
  )
{
  EFI_STATUS                     Status;
  UINT32                         IshVersion;
  UINT32                         PdtVersion;
  UINTN                          Size;

  BS = gBS;

  Status = FwuSetIshConfig((UINT8*)MeDataBuffer, MeDataSize);
  if (EFI_ERROR (Status)) {
    return Status;
  }
  DEBUG ((DEBUG_ERROR | DEBUG_INFO, "Sending Image for Executing PDT Update Succeeded.\n"));

  Size = sizeof (IshVersion);
  Status = gRT->GetVariable (
                  L"IshFwVersion",
                  &gIshUpdVerGuid,
                  NULL,
                  &Size,
                  &IshVersion
                  );
  if (!EFI_ERROR (Status)) {
    Size = sizeof (PdtVersion);
    Status = gRT->GetVariable (
                    L"PdtUpdVersion",
                    &gPdtUpdCountGuid,
                    NULL,
                    &Size,
                    &PdtVersion
                    );

    if (IshVersion < PdtVersion && IshVersion == (PdtVersion >> 16)) {
      PdtVersion++;
    }

    if (IshVersion > (PdtVersion >> 16)) {
      PdtVersion = (IshVersion << 16);
    }

    Status = gRT->SetVariable (
                    L"PdtUpdVersion",
                    &gPdtUpdCountGuid,
                    EFI_VARIABLE_BOOTSERVICE_ACCESS | EFI_VARIABLE_RUNTIME_ACCESS | EFI_VARIABLE_NON_VOLATILE,
                    Size,
                    &PdtVersion
                    );
  }

  return EFI_SUCCESS;
}

EFI_STATUS
EFIAPI
FmtsUseIntelLibWriteUpdateStages(
  UINT32   *PercentWritten
  )
{
  UINT32                         Status;
  BOOLEAN                        InProgress;
  UINT32                         FwUpdateStatus;
  UINT32                         NeededResetType;

  InProgress      = FALSE;
  FwUpdateStatus  = 0;
  NeededResetType = MFT_PART_INFO_EXT_UPDATE_ACTION_NONE;

  Status = FwuCheckUpdateProgress(&InProgress, PercentWritten, &FwUpdateStatus, &NeededResetType);

  return Status;
}