/** @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>

  This program and the accompanying materials
  are licensed and made available under the terms and conditions of the BSD License
  which accompanies this distribution.  The full text of the license may be found at
  http://opensource.org/licenses/bsd-license.php

  THE PROGRAM IS DISTRIBUTED UNDER THE BSD LICENSE ON AN "AS IS" BASIS,
  WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.

**/

#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;
}

/**
  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 (ArmSmcArgs.Arg0 == 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.
      //
      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)) {
    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)) {
    MessageArg->Return = OPTEE_ERROR_COMMUNICATION;
    MessageArg->ReturnOrigin = OPTEE_ORIGIN_COMMUNICATION;
  }

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

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

  return EFI_SUCCESS;
}