/** @file
  Api's to communicate with OP-TEE OS (Trusted OS based on ARM TrustZone) via
  secure monitor calls.

  Copyright (c) 2018, Linaro Ltd. All rights reserved.<BR>
  Copyright (c) 2021, Arm Limited. All rights reserved.<BR>

  SPDX-License-Identifier: BSD-2-Clause-Patent

**/

#include <Library/ArmMmuLib.h>
#include <Library/ArmSmcLib.h>
#include <Library/BaseMemoryLib.h>
#include <Library/BaseLib.h>
#include <Library/DebugLib.h>
#include <Library/OpteeLib.h>

#include <IndustryStandard/ArmStdSmc.h>
#include <OpteeSmc.h>
#include <Uefi.h>

STATIC OPTEE_SHARED_MEMORY_INFORMATION OpteeSharedMemoryInformation = { 0 };

/**
  Check for OP-TEE presence.
**/
BOOLEAN
EFIAPI
IsOpteePresent (
  VOID
  )
{
  ARM_SMC_ARGS ArmSmcArgs;

  ZeroMem (&ArmSmcArgs, sizeof (ARM_SMC_ARGS));
  // Send a Trusted OS Calls UID command
  ArmSmcArgs.Arg0 = ARM_SMC_ID_TOS_UID;
  ArmCallSmc (&ArmSmcArgs);

  if ((ArmSmcArgs.Arg0 == OPTEE_OS_UID0) &&
      (ArmSmcArgs.Arg1 == OPTEE_OS_UID1) &&
      (ArmSmcArgs.Arg2 == OPTEE_OS_UID2) &&
      (ArmSmcArgs.Arg3 == OPTEE_OS_UID3)) {
    return TRUE;
  } else {
    return FALSE;
  }
}

STATIC
EFI_STATUS
OpteeSharedMemoryRemap (
  VOID
  )
{
  ARM_SMC_ARGS                 ArmSmcArgs;
  EFI_PHYSICAL_ADDRESS         PhysicalAddress;
  EFI_PHYSICAL_ADDRESS         Start;
  EFI_PHYSICAL_ADDRESS         End;
  EFI_STATUS                   Status;
  UINTN                        Size;

  ZeroMem (&ArmSmcArgs, sizeof (ARM_SMC_ARGS));
  ArmSmcArgs.Arg0 = OPTEE_SMC_GET_SHARED_MEMORY_CONFIG;

  ArmCallSmc (&ArmSmcArgs);
  if (ArmSmcArgs.Arg0 != OPTEE_SMC_RETURN_OK) {
    DEBUG ((DEBUG_WARN, "OP-TEE shared memory not supported\n"));
    return EFI_UNSUPPORTED;
  }

  if (ArmSmcArgs.Arg3 != OPTEE_SMC_SHARED_MEMORY_CACHED) {
    DEBUG ((DEBUG_WARN, "OP-TEE: Only normal cached shared memory supported\n"));
    return EFI_UNSUPPORTED;
  }

  Start = (ArmSmcArgs.Arg1 + SIZE_4KB - 1) & ~(SIZE_4KB - 1);
  End = (ArmSmcArgs.Arg1 + ArmSmcArgs.Arg2) & ~(SIZE_4KB - 1);
  PhysicalAddress = Start;
  Size = End - Start;

  if (Size < SIZE_4KB) {
    DEBUG ((DEBUG_WARN, "OP-TEE shared memory too small\n"));
    return EFI_BUFFER_TOO_SMALL;
  }

  Status = ArmSetMemoryAttributes (PhysicalAddress, Size, EFI_MEMORY_WB);
  if (EFI_ERROR (Status)) {
    return Status;
  }

  OpteeSharedMemoryInformation.Base = (UINTN)PhysicalAddress;
  OpteeSharedMemoryInformation.Size = Size;

  return EFI_SUCCESS;
}

EFI_STATUS
EFIAPI
OpteeInit (
  VOID
  )
{
  EFI_STATUS      Status;

  if (!IsOpteePresent ()) {
    DEBUG ((DEBUG_WARN, "OP-TEE not present\n"));
    return EFI_UNSUPPORTED;
  }

  Status = OpteeSharedMemoryRemap ();
  if (EFI_ERROR (Status)) {
    DEBUG ((DEBUG_WARN, "OP-TEE shared memory remap failed\n"));
    return Status;
  }

  return EFI_SUCCESS;
}

STATIC
BOOLEAN
IsOpteeSmcReturnRpc (
  UINT32 Return
  )
{
  return (Return != OPTEE_SMC_RETURN_UNKNOWN_FUNCTION) &&
         ((Return & OPTEE_SMC_RETURN_RPC_PREFIX_MASK) ==
          OPTEE_SMC_RETURN_RPC_PREFIX);
}

/**
  Does Standard SMC to OP-TEE in secure world.

  @param[in]  PhysicalArg   Physical address of message to pass to secure world

  @return                   0 on success, secure world return code otherwise

**/
STATIC
UINT32
OpteeCallWithArg (
  IN UINT64 PhysicalArg
  )
{
  ARM_SMC_ARGS ArmSmcArgs;

  ZeroMem (&ArmSmcArgs, sizeof (ARM_SMC_ARGS));
  ArmSmcArgs.Arg0 = OPTEE_SMC_CALL_WITH_ARG;
  ArmSmcArgs.Arg1 = (UINT32)(PhysicalArg >> 32);
  ArmSmcArgs.Arg2 = (UINT32)PhysicalArg;

  while (TRUE) {
    ArmCallSmc (&ArmSmcArgs);

    if (IsOpteeSmcReturnRpc (ArmSmcArgs.Arg0)) {
      switch (ArmSmcArgs.Arg0) {
      case OPTEE_SMC_RETURN_RPC_FOREIGN_INTERRUPT:
        //
        // A foreign interrupt was raised while secure world was
        // executing, since they are handled in UEFI a dummy RPC is
        // performed to let UEFI take the interrupt through the normal
        // vector.
        //
        break;

      default:
         // Do nothing in case RPC is not implemented.
        break;
      }

      ArmSmcArgs.Arg0 = OPTEE_SMC_RETURN_FROM_RPC;
    } else {
      break;
    }
  }

  return ArmSmcArgs.Arg0;
}

STATIC
VOID
EfiGuidToRfc4122Uuid (
  OUT RFC4122_UUID       *Rfc4122Uuid,
  IN EFI_GUID            *Guid
  )
{
  Rfc4122Uuid->Data1 = SwapBytes32 (Guid->Data1);
  Rfc4122Uuid->Data2 = SwapBytes16 (Guid->Data2);
  Rfc4122Uuid->Data3 = SwapBytes16 (Guid->Data3);
  CopyMem (Rfc4122Uuid->Data4, Guid->Data4, sizeof (Rfc4122Uuid->Data4));
}

EFI_STATUS
EFIAPI
OpteeOpenSession (
  IN OUT OPTEE_OPEN_SESSION_ARG      *OpenSessionArg
  )
{
  OPTEE_MESSAGE_ARG    *MessageArg;

  MessageArg = NULL;

  if (OpteeSharedMemoryInformation.Base == 0) {
    DEBUG ((DEBUG_WARN, "OP-TEE not initialized\n"));
    return EFI_NOT_STARTED;
  }

  MessageArg = (OPTEE_MESSAGE_ARG *)OpteeSharedMemoryInformation.Base;
  ZeroMem (MessageArg, sizeof (OPTEE_MESSAGE_ARG));

  MessageArg->Command = OPTEE_MESSAGE_COMMAND_OPEN_SESSION;

  //
  // Initialize and add the meta parameters needed when opening a
  // session.
  //
  MessageArg->Params[0].Attribute = OPTEE_MESSAGE_ATTRIBUTE_TYPE_VALUE_INPUT |
                                    OPTEE_MESSAGE_ATTRIBUTE_META;
  MessageArg->Params[1].Attribute = OPTEE_MESSAGE_ATTRIBUTE_TYPE_VALUE_INPUT |
                                    OPTEE_MESSAGE_ATTRIBUTE_META;
  EfiGuidToRfc4122Uuid (
    (RFC4122_UUID *)&MessageArg->Params[0].Union.Value,
    &OpenSessionArg->Uuid
    );
  ZeroMem (&MessageArg->Params[1].Union.Value, sizeof (EFI_GUID));
  MessageArg->Params[1].Union.Value.C = OPTEE_LOGIN_PUBLIC;

  MessageArg->NumParams = 2;

  if (OpteeCallWithArg ((UINTN)MessageArg) != 0) {
    MessageArg->Return = OPTEE_ERROR_COMMUNICATION;
    MessageArg->ReturnOrigin = OPTEE_ORIGIN_COMMUNICATION;
  }

  OpenSessionArg->Session = MessageArg->Session;
  OpenSessionArg->Return = MessageArg->Return;
  OpenSessionArg->ReturnOrigin = MessageArg->ReturnOrigin;

  return EFI_SUCCESS;
}

EFI_STATUS
EFIAPI
OpteeCloseSession (
  IN UINT32                  Session
  )
{
  OPTEE_MESSAGE_ARG    *MessageArg;

  MessageArg = NULL;

  if (OpteeSharedMemoryInformation.Base == 0) {
    DEBUG ((DEBUG_WARN, "OP-TEE not initialized\n"));
    return EFI_NOT_STARTED;
  }

  MessageArg = (OPTEE_MESSAGE_ARG *)OpteeSharedMemoryInformation.Base;
  ZeroMem (MessageArg, sizeof (OPTEE_MESSAGE_ARG));

  MessageArg->Command = OPTEE_MESSAGE_COMMAND_CLOSE_SESSION;
  MessageArg->Session = Session;

  OpteeCallWithArg ((UINTN)MessageArg);

  return EFI_SUCCESS;
}

STATIC
EFI_STATUS
OpteeToMessageParam (
  OUT OPTEE_MESSAGE_PARAM    *MessageParams,
  IN UINT32                  NumParams,
  IN OPTEE_MESSAGE_PARAM     *InParams
  )
{
  UINT32                  Idx;
  UINTN                   ParamSharedMemoryAddress;
  UINTN                   SharedMemorySize;
  UINTN                   Size;

  Size = (sizeof (OPTEE_MESSAGE_ARG) + sizeof (UINT64) - 1) &
          ~(sizeof (UINT64) - 1);
  ParamSharedMemoryAddress = OpteeSharedMemoryInformation.Base + Size;
  SharedMemorySize = OpteeSharedMemoryInformation.Size - Size;

  for (Idx = 0; Idx < NumParams; Idx++) {
    CONST OPTEE_MESSAGE_PARAM    *InParam;
    OPTEE_MESSAGE_PARAM          *MessageParam;
    UINT32                       Attribute;

    InParam = InParams + Idx;
    MessageParam = MessageParams + Idx;
    Attribute = InParam->Attribute & OPTEE_MESSAGE_ATTRIBUTE_TYPE_MASK;

    switch (Attribute) {
    case OPTEE_MESSAGE_ATTRIBUTE_TYPE_NONE:
      MessageParam->Attribute = OPTEE_MESSAGE_ATTRIBUTE_TYPE_NONE;
      ZeroMem (&MessageParam->Union, sizeof (MessageParam->Union));
      break;

    case OPTEE_MESSAGE_ATTRIBUTE_TYPE_VALUE_INPUT:
    case OPTEE_MESSAGE_ATTRIBUTE_TYPE_VALUE_OUTPUT:
    case OPTEE_MESSAGE_ATTRIBUTE_TYPE_VALUE_INOUT:
      MessageParam->Attribute = Attribute;
      MessageParam->Union.Value.A = InParam->Union.Value.A;
      MessageParam->Union.Value.B = InParam->Union.Value.B;
      MessageParam->Union.Value.C = InParam->Union.Value.C;
      break;

    case OPTEE_MESSAGE_ATTRIBUTE_TYPE_MEMORY_INPUT:
    case OPTEE_MESSAGE_ATTRIBUTE_TYPE_MEMORY_OUTPUT:
    case OPTEE_MESSAGE_ATTRIBUTE_TYPE_MEMORY_INOUT:
      MessageParam->Attribute = Attribute;

      if (InParam->Union.Memory.Size > SharedMemorySize) {
        return EFI_OUT_OF_RESOURCES;
      }

      CopyMem (
        (VOID *)ParamSharedMemoryAddress,
        (VOID *)(UINTN)InParam->Union.Memory.BufferAddress,
        InParam->Union.Memory.Size
        );
      MessageParam->Union.Memory.BufferAddress = (UINT64)ParamSharedMemoryAddress;
      MessageParam->Union.Memory.Size = InParam->Union.Memory.Size;

      Size = (InParam->Union.Memory.Size + sizeof (UINT64) - 1) &
              ~(sizeof (UINT64) - 1);
      ParamSharedMemoryAddress += Size;
      SharedMemorySize -= Size;
      break;

    default:
      return EFI_INVALID_PARAMETER;
    }
  }

  return EFI_SUCCESS;
}

STATIC
EFI_STATUS
OpteeFromMessageParam (
  OUT OPTEE_MESSAGE_PARAM    *OutParams,
  IN UINT32                  NumParams,
  IN OPTEE_MESSAGE_PARAM     *MessageParams
  )
{
  UINT32                 Idx;

  for (Idx = 0; Idx < NumParams; Idx++) {
    OPTEE_MESSAGE_PARAM          *OutParam;
    CONST OPTEE_MESSAGE_PARAM    *MessageParam;
    UINT32                   Attribute;

    OutParam = OutParams + Idx;
    MessageParam = MessageParams + Idx;
    Attribute = MessageParam->Attribute & OPTEE_MESSAGE_ATTRIBUTE_TYPE_MASK;

    switch (Attribute) {
    case OPTEE_MESSAGE_ATTRIBUTE_TYPE_NONE:
      OutParam->Attribute = OPTEE_MESSAGE_ATTRIBUTE_TYPE_NONE;
      ZeroMem (&OutParam->Union, sizeof (OutParam->Union));
      break;

    case OPTEE_MESSAGE_ATTRIBUTE_TYPE_VALUE_INPUT:
    case OPTEE_MESSAGE_ATTRIBUTE_TYPE_VALUE_OUTPUT:
    case OPTEE_MESSAGE_ATTRIBUTE_TYPE_VALUE_INOUT:
      OutParam->Attribute = Attribute;
      OutParam->Union.Value.A = MessageParam->Union.Value.A;
      OutParam->Union.Value.B = MessageParam->Union.Value.B;
      OutParam->Union.Value.C = MessageParam->Union.Value.C;
      break;

    case OPTEE_MESSAGE_ATTRIBUTE_TYPE_MEMORY_INPUT:
    case OPTEE_MESSAGE_ATTRIBUTE_TYPE_MEMORY_OUTPUT:
    case OPTEE_MESSAGE_ATTRIBUTE_TYPE_MEMORY_INOUT:
      OutParam->Attribute = Attribute;

      if (MessageParam->Union.Memory.Size > OutParam->Union.Memory.Size) {
        return EFI_BAD_BUFFER_SIZE;
      }

      CopyMem (
        (VOID *)(UINTN)OutParam->Union.Memory.BufferAddress,
        (VOID *)(UINTN)MessageParam->Union.Memory.BufferAddress,
        MessageParam->Union.Memory.Size
        );
      OutParam->Union.Memory.Size = MessageParam->Union.Memory.Size;
      break;

    default:
      return EFI_INVALID_PARAMETER;
    }
  }

  return EFI_SUCCESS;
}

EFI_STATUS
EFIAPI
OpteeInvokeFunction (
  IN OUT OPTEE_INVOKE_FUNCTION_ARG       *InvokeFunctionArg
  )
{
  EFI_STATUS       Status;
  OPTEE_MESSAGE_ARG    *MessageArg;

  MessageArg = NULL;

  if (OpteeSharedMemoryInformation.Base == 0) {
    DEBUG ((DEBUG_WARN, "OP-TEE not initialized\n"));
    return EFI_NOT_STARTED;
  }

  MessageArg = (OPTEE_MESSAGE_ARG *)OpteeSharedMemoryInformation.Base;
  ZeroMem (MessageArg, sizeof (OPTEE_MESSAGE_ARG));

  MessageArg->Command = OPTEE_MESSAGE_COMMAND_INVOKE_FUNCTION;
  MessageArg->Function = InvokeFunctionArg->Function;
  MessageArg->Session = InvokeFunctionArg->Session;

  Status = OpteeToMessageParam (
             MessageArg->Params,
             OPTEE_MAX_CALL_PARAMS,
             InvokeFunctionArg->Params
             );
  if (Status) {
    return Status;
  }

  MessageArg->NumParams = OPTEE_MAX_CALL_PARAMS;

  if (OpteeCallWithArg ((UINTN)MessageArg) != 0) {
    MessageArg->Return = OPTEE_ERROR_COMMUNICATION;
    MessageArg->ReturnOrigin = OPTEE_ORIGIN_COMMUNICATION;
  }

  if (OpteeFromMessageParam (
        InvokeFunctionArg->Params,
        OPTEE_MAX_CALL_PARAMS,
        MessageArg->Params
        ) != 0) {
    MessageArg->Return = OPTEE_ERROR_COMMUNICATION;
    MessageArg->ReturnOrigin = OPTEE_ORIGIN_COMMUNICATION;
  }

  InvokeFunctionArg->Return = MessageArg->Return;
  InvokeFunctionArg->ReturnOrigin = MessageArg->ReturnOrigin;

  return EFI_SUCCESS;
}