/** @file
  DXE capsule report related function.

  Copyright (c) 2016 - 2021, Intel Corporation. All rights reserved.<BR>
  SPDX-License-Identifier: BSD-2-Clause-Patent

**/

#include <PiDxe.h>
#include <Protocol/FirmwareManagement.h>
#include <Guid/CapsuleReport.h>
#include <Guid/FmpCapsule.h>
#include <Guid/CapsuleVendor.h>

#include <Library/BaseLib.h>
#include <Library/DebugLib.h>
#include <Library/BaseMemoryLib.h>
#include <Library/UefiBootServicesTableLib.h>
#include <Library/UefiRuntimeServicesTableLib.h>
#include <Library/MemoryAllocationLib.h>
#include <Library/UefiLib.h>
#include <Library/PcdLib.h>
#include <Library/HobLib.h>
#include <Library/PrintLib.h>
#include <Library/ReportStatusCodeLib.h>
#include <Library/DevicePathLib.h>
#include <Library/CapsuleLib.h>
#include <Library/VariablePolicyHelperLib.h>

#include <IndustryStandard/WindowsUxCapsule.h>

/**
  This routine is called to clear CapsuleOnDisk Relocation Info variable.
  Total Capsule On Disk length is recorded in this variable

  @retval EFI_SUCCESS   Capsule On Disk flags are cleared

**/
EFI_STATUS
CoDClearCapsuleRelocationInfo (
  VOID
  );

/**
  Get current capsule last variable index.

  @return Current capsule last variable index.
  @retval -1  No current capsule last variable.
**/
INTN
GetCurrentCapsuleLastIndex (
  VOID
  )
{
  UINTN       Size;
  CHAR16      CapsuleLastStr[sizeof ("Capsule####")];
  EFI_STATUS  Status;
  UINT16      CurrentIndex;

  Size   = sizeof (L"Capsule####") - sizeof (CHAR16); // no zero terminator
  Status = gRT->GetVariable (
                  L"CapsuleLast",
                  &gEfiCapsuleReportGuid,
                  NULL,
                  &Size,
                  CapsuleLastStr
                  );
  if (EFI_ERROR (Status)) {
    return -1;
  }

  CurrentIndex = (UINT16)StrHexToUintn (&CapsuleLastStr[sizeof ("Capsule") - 1]);
  return CurrentIndex;
}

/**
  Get a new capsule status variable index.

  @return A new capsule status variable index.
  @retval 0  No new capsule status variable index. Rolling over.
**/
INTN
GetNewCapsuleResultIndex (
  VOID
  )
{
  INTN  CurrentIndex;

  CurrentIndex = GetCurrentCapsuleLastIndex ();
  if (CurrentIndex >= PcdGet16 (PcdCapsuleMax)) {
    DEBUG ((DEBUG_INFO, "  CapsuleResult variable Rolling Over!\n"));
    return 0;
  }

  return CurrentIndex + 1;
}

/**
  Lock Variable by variable policy.

  @param[in] VariableGuid         The Guid of the variable to be locked
  @param[in] VariableName         The name of the variable to be locked
  @param[in] VariablePolicy       The pointer of variable lock policy
**/
VOID
LockVariable (
  IN CONST  EFI_GUID                 VariableGuid,
  IN CHAR16                          *VariableName,
  IN EDKII_VARIABLE_POLICY_PROTOCOL  *VariablePolicy
  )
{
  EFI_STATUS  Status;

  // Set the policies to protect the target variables
  Status = RegisterBasicVariablePolicy (
             VariablePolicy,
             &VariableGuid,
             VariableName,
             VARIABLE_POLICY_NO_MIN_SIZE,
             VARIABLE_POLICY_NO_MAX_SIZE,
             VARIABLE_POLICY_NO_MUST_ATTR,
             VARIABLE_POLICY_NO_CANT_ATTR,
             VARIABLE_POLICY_TYPE_LOCK_NOW
             );
  if (EFI_ERROR (Status)) {
    DEBUG ((
      DEBUG_ERROR,
      "DxeCapsuleLibFmp: Failed to lock variable %g %s.  Status = %r\n",
      &VariableGuid,
      VariableName,
      Status
      ));
    ASSERT_EFI_ERROR (Status);
  }
}

/**
  Write a new capsule status variable.

  @param[in] CapsuleResult      The capsule status variable
  @param[in] CapsuleResultSize  The size of the capsule stauts variable in bytes

  @retval EFI_SUCCESS          The capsule status variable is recorded.
  @retval EFI_OUT_OF_RESOURCES No resource to record the capsule status variable.
**/
EFI_STATUS
WriteNewCapsuleResultVariable (
  IN VOID   *CapsuleResult,
  IN UINTN  CapsuleResultSize
  )
{
  INTN        CapsuleResultIndex;
  CHAR16      CapsuleResultStr[sizeof ("Capsule####")];
  UINTN       Size;
  EFI_STATUS  Status;

  CapsuleResultIndex = GetNewCapsuleResultIndex ();
  DEBUG ((DEBUG_INFO, "New CapsuleResultIndex - 0x%x\n", CapsuleResultIndex));

  UnicodeSPrint (
    CapsuleResultStr,
    sizeof (CapsuleResultStr),
    L"Capsule%04x",
    CapsuleResultIndex
    );

  Status = gRT->SetVariable (
                  CapsuleResultStr,
                  &gEfiCapsuleReportGuid,
                  EFI_VARIABLE_NON_VOLATILE | EFI_VARIABLE_BOOTSERVICE_ACCESS | EFI_VARIABLE_RUNTIME_ACCESS,
                  CapsuleResultSize,
                  CapsuleResult
                  );
  if (!EFI_ERROR (Status)) {
    Size = sizeof (L"Capsule####") - sizeof (CHAR16); // no zero terminator
    DEBUG ((DEBUG_INFO, "Set CapsuleLast - %s\n", CapsuleResultStr));
    Status = gRT->SetVariable (
                    L"CapsuleLast",
                    &gEfiCapsuleReportGuid,
                    EFI_VARIABLE_NON_VOLATILE | EFI_VARIABLE_BOOTSERVICE_ACCESS | EFI_VARIABLE_RUNTIME_ACCESS,
                    Size,
                    CapsuleResultStr
                    );
  }

  return Status;
}

/**
  Record capsule status variable and to local cache.

  @param[in] CapsuleHeader  The capsule image header
  @param[in] CapsuleStatus  The capsule process stauts

  @retval EFI_SUCCESS          The capsule status variable is recorded.
  @retval EFI_OUT_OF_RESOURCES No resource to record the capsule status variable.
**/
EFI_STATUS
RecordCapsuleStatusVariable (
  IN EFI_CAPSULE_HEADER  *CapsuleHeader,
  IN EFI_STATUS          CapsuleStatus
  )
{
  EFI_CAPSULE_RESULT_VARIABLE_HEADER  CapsuleResultVariable;
  EFI_STATUS                          Status;

  CapsuleResultVariable.VariableTotalSize = sizeof (CapsuleResultVariable);
  CapsuleResultVariable.Reserved          = 0;
  CopyGuid (&CapsuleResultVariable.CapsuleGuid, &CapsuleHeader->CapsuleGuid);
  ZeroMem (&CapsuleResultVariable.CapsuleProcessed, sizeof (CapsuleResultVariable.CapsuleProcessed));
  gRT->GetTime (&CapsuleResultVariable.CapsuleProcessed, NULL);
  CapsuleResultVariable.CapsuleStatus = CapsuleStatus;

  Status = EFI_SUCCESS;
  if ((CapsuleHeader->Flags & CAPSULE_FLAGS_PERSIST_ACROSS_RESET) != 0) {
    Status = WriteNewCapsuleResultVariable (&CapsuleResultVariable, sizeof (CapsuleResultVariable));
  }

  return Status;
}

/**
  Record FMP capsule status variable and to local cache.

  @param[in] CapsuleHeader  The capsule image header
  @param[in] CapsuleStatus  The capsule process stauts
  @param[in] PayloadIndex   FMP payload index
  @param[in] ImageHeader    FMP image header
  @param[in] FmpDevicePath  DevicePath associated with the FMP producer
  @param[in] CapFileName    Capsule file name

  @retval EFI_SUCCESS          The capsule status variable is recorded.
  @retval EFI_OUT_OF_RESOURCES No resource to record the capsule status variable.
**/
EFI_STATUS
RecordFmpCapsuleStatusVariable (
  IN EFI_CAPSULE_HEADER                            *CapsuleHeader,
  IN EFI_STATUS                                    CapsuleStatus,
  IN UINTN                                         PayloadIndex,
  IN EFI_FIRMWARE_MANAGEMENT_CAPSULE_IMAGE_HEADER  *ImageHeader,
  IN EFI_DEVICE_PATH_PROTOCOL                      *FmpDevicePath  OPTIONAL,
  IN CHAR16                                        *CapFileName    OPTIONAL
  )
{
  EFI_CAPSULE_RESULT_VARIABLE_HEADER  *CapsuleResultVariableHeader;
  EFI_CAPSULE_RESULT_VARIABLE_FMP     *CapsuleResultVariableFmp;
  EFI_STATUS                          Status;
  UINT8                               *CapsuleResultVariable;
  UINTN                               CapsuleResultVariableSize;
  CHAR16                              *DevicePathStr;
  UINTN                               DevicePathStrSize;
  UINTN                               CapFileNameSize;

  DevicePathStr   = NULL;
  CapFileNameSize = sizeof (CHAR16);

  if (FmpDevicePath != NULL) {
    DevicePathStr = ConvertDevicePathToText (FmpDevicePath, FALSE, FALSE);
  }

  if (DevicePathStr != NULL) {
    DevicePathStrSize = StrSize (DevicePathStr);
  } else {
    DevicePathStrSize = sizeof (CHAR16);
  }

  if (CapFileName != NULL) {
    CapFileNameSize = StrSize (CapFileName);
  }

  //
  // Allocate room for CapsuleFileName.
  //
  CapsuleResultVariableSize = sizeof (EFI_CAPSULE_RESULT_VARIABLE_HEADER) + sizeof (EFI_CAPSULE_RESULT_VARIABLE_FMP) + CapFileNameSize + DevicePathStrSize;

  CapsuleResultVariable = AllocateZeroPool (CapsuleResultVariableSize);
  if (CapsuleResultVariable == NULL) {
    return EFI_OUT_OF_RESOURCES;
  }

  CapsuleResultVariableHeader                    = (VOID *)CapsuleResultVariable;
  CapsuleResultVariableHeader->VariableTotalSize = (UINT32)CapsuleResultVariableSize;
  CapsuleResultVariableHeader->Reserved          = 0;
  CopyGuid (&CapsuleResultVariableHeader->CapsuleGuid, &CapsuleHeader->CapsuleGuid);
  ZeroMem (&CapsuleResultVariableHeader->CapsuleProcessed, sizeof (CapsuleResultVariableHeader->CapsuleProcessed));
  gRT->GetTime (&CapsuleResultVariableHeader->CapsuleProcessed, NULL);
  CapsuleResultVariableHeader->CapsuleStatus = CapsuleStatus;

  CapsuleResultVariableFmp                   = (VOID *)(CapsuleResultVariable + sizeof (EFI_CAPSULE_RESULT_VARIABLE_HEADER));
  CapsuleResultVariableFmp->Version          = 0x1;
  CapsuleResultVariableFmp->PayloadIndex     = (UINT8)PayloadIndex;
  CapsuleResultVariableFmp->UpdateImageIndex = ImageHeader->UpdateImageIndex;
  CopyGuid (&CapsuleResultVariableFmp->UpdateImageTypeId, &ImageHeader->UpdateImageTypeId);

  if (CapFileName != NULL) {
    CopyMem ((UINT8 *)CapsuleResultVariableFmp + sizeof (EFI_CAPSULE_RESULT_VARIABLE_FMP), CapFileName, CapFileNameSize);
  }

  if (DevicePathStr != NULL) {
    CopyMem ((UINT8 *)CapsuleResultVariableFmp + sizeof (EFI_CAPSULE_RESULT_VARIABLE_FMP) + CapFileNameSize, DevicePathStr, DevicePathStrSize);
    FreePool (DevicePathStr);
    DevicePathStr = NULL;
  }

  Status = EFI_SUCCESS;
  if ((CapsuleHeader->Flags & CAPSULE_FLAGS_PERSIST_ACROSS_RESET) != 0) {
    Status = WriteNewCapsuleResultVariable (CapsuleResultVariable, CapsuleResultVariableSize);
  }

  FreePool (CapsuleResultVariable);
  return Status;
}

/**
  Initialize CapsuleMax variables.

  @param[in] VariablePolicy       The pointer of variable lock policy
**/
VOID
InitCapsuleMaxVariable (
  EDKII_VARIABLE_POLICY_PROTOCOL  *VariablePolicy
  )
{
  EFI_STATUS  Status;
  UINTN       Size;
  CHAR16      CapsuleMaxStr[sizeof ("Capsule####")];

  UnicodeSPrint (
    CapsuleMaxStr,
    sizeof (CapsuleMaxStr),
    L"Capsule%04x",
    PcdGet16 (PcdCapsuleMax)
    );

  Size   = sizeof (L"Capsule####") - sizeof (CHAR16); // no zero terminator
  Status = gRT->SetVariable (
                  L"CapsuleMax",
                  &gEfiCapsuleReportGuid,
                  EFI_VARIABLE_BOOTSERVICE_ACCESS | EFI_VARIABLE_RUNTIME_ACCESS,
                  Size,
                  CapsuleMaxStr
                  );
  if (!EFI_ERROR (Status)) {
    // Lock it per UEFI spec.
    LockVariable (gEfiCapsuleReportGuid, L"CapsuleMax", VariablePolicy);
  }
}

/**
  Initialize CapsuleLast variables.

  @param[in] VariablePolicy       The pointer of variable lock policy
**/
VOID
InitCapsuleLastVariable (
  EDKII_VARIABLE_POLICY_PROTOCOL  *VariablePolicy
  )
{
  EFI_STATUS     Status;
  EFI_BOOT_MODE  BootMode;
  VOID           *CapsuleResult;
  UINTN          Size;
  CHAR16         CapsuleLastStr[sizeof ("Capsule####")];

  BootMode = GetBootModeHob ();
  if (BootMode == BOOT_ON_FLASH_UPDATE) {
    Status = gRT->SetVariable (
                    L"CapsuleLast",
                    &gEfiCapsuleReportGuid,
                    EFI_VARIABLE_NON_VOLATILE | EFI_VARIABLE_BOOTSERVICE_ACCESS | EFI_VARIABLE_RUNTIME_ACCESS,
                    0,
                    NULL
                    );
    // Do not lock it because it will be updated later.
  } else {
    //
    // Check if OS/APP cleared L"Capsule####"
    //
    ZeroMem (CapsuleLastStr, sizeof (CapsuleLastStr));
    Size   = sizeof (L"Capsule####") - sizeof (CHAR16); // no zero terminator
    Status = gRT->GetVariable (
                    L"CapsuleLast",
                    &gEfiCapsuleReportGuid,
                    NULL,
                    &Size,
                    CapsuleLastStr
                    );
    if (!EFI_ERROR (Status)) {
      //
      // L"CapsuleLast" is got, check if data is there.
      //
      Status = GetVariable2 (
                 CapsuleLastStr,
                 &gEfiCapsuleReportGuid,
                 (VOID **)&CapsuleResult,
                 NULL
                 );
      if (EFI_ERROR (Status)) {
        //
        // If no data, delete L"CapsuleLast"
        //
        Status = gRT->SetVariable (
                        L"CapsuleLast",
                        &gEfiCapsuleReportGuid,
                        EFI_VARIABLE_NON_VOLATILE | EFI_VARIABLE_BOOTSERVICE_ACCESS | EFI_VARIABLE_RUNTIME_ACCESS,
                        0,
                        NULL
                        );
      } else {
        if (CapsuleResult != NULL) {
          FreePool (CapsuleResult);
        }
      }
    }

    // Lock it in normal boot path per UEFI spec.
    LockVariable (gEfiCapsuleReportGuid, L"CapsuleLast", VariablePolicy);
  }
}

/**
  Initialize capsule update variables.
**/
VOID
InitCapsuleUpdateVariable (
  VOID
  )
{
  EFI_STATUS  Status;
  UINTN       Index;
  CHAR16      CapsuleVarName[30];
  CHAR16      *TempVarName;

  //
  // Clear all the capsule variables CapsuleUpdateData, CapsuleUpdateData1, CapsuleUpdateData2...
  // as early as possible which will avoid the next time boot after the capsule update
  // will still into the capsule loop
  //
  StrCpyS (CapsuleVarName, sizeof (CapsuleVarName)/sizeof (CapsuleVarName[0]), EFI_CAPSULE_VARIABLE_NAME);
  TempVarName = CapsuleVarName + StrLen (CapsuleVarName);
  Index       = 0;
  while (TRUE) {
    if (Index > 0) {
      UnicodeValueToStringS (
        TempVarName,
        sizeof (CapsuleVarName) - ((UINTN)TempVarName - (UINTN)CapsuleVarName),
        0,
        Index,
        0
        );
    }

    Status = gRT->SetVariable (
                    CapsuleVarName,
                    &gEfiCapsuleVendorGuid,
                    EFI_VARIABLE_NON_VOLATILE | EFI_VARIABLE_RUNTIME_ACCESS | EFI_VARIABLE_BOOTSERVICE_ACCESS,
                    0,
                    (VOID *)NULL
                    );
    if (EFI_ERROR (Status)) {
      //
      // There is no capsule variables, quit
      //
      break;
    }

    Index++;
  }
}

/**
  Initialize capsule relocation info variable.

  @param[in] VariablePolicy       The pointer of variable lock policy
**/
VOID
InitCapsuleRelocationInfo (
  EDKII_VARIABLE_POLICY_PROTOCOL  *VariablePolicy
  )
{
  CoDClearCapsuleRelocationInfo ();

  //
  // Unlock Capsule On Disk relocation Info variable only when Capsule On Disk flag is enabled
  //
  if (!CoDCheckCapsuleOnDiskFlag ()) {
    LockVariable (gEfiCapsuleVendorGuid, COD_RELOCATION_INFO_VAR_NAME, VariablePolicy);
  }
}

/**
  Initialize capsule related variables.
**/
VOID
InitCapsuleVariable (
  VOID
  )
{
  EFI_STATUS                      Status;
  EDKII_VARIABLE_POLICY_PROTOCOL  *VariablePolicy;

  // Locate the VariablePolicy protocol
  Status = gBS->LocateProtocol (&gEdkiiVariablePolicyProtocolGuid, NULL, (VOID **)&VariablePolicy);
  if (EFI_ERROR (Status)) {
    DEBUG ((DEBUG_ERROR, "DxeCapsuleReportLib %a - Could not locate VariablePolicy protocol! %r\n", __FUNCTION__, Status));
    ASSERT_EFI_ERROR (Status);
  }

  InitCapsuleUpdateVariable ();
  InitCapsuleMaxVariable (VariablePolicy);
  InitCapsuleLastVariable (VariablePolicy);
  InitCapsuleRelocationInfo (VariablePolicy);

  //
  // No need to clear L"Capsule####", because OS/APP should refer L"CapsuleLast"
  // to check status and delete them.
  //
}