/*
 * Copyright (c) Qualcomm Technologies, Inc.
 * All Rights Reserved.
 * Confidential and Proprietary - Qualcomm Technologies, Inc.
 */

#include "CAdder.h"
#include "CHash.h"
#include "CLogSinkExample_open.h"
#include "CSMCIExampleApp_open.h"
#include "CTALog.h"
#include "IAdder.h"
#include "ICallbackObjectExample.h"
#include "IHash.h"
#include "IMemRegion.h"
#include "IMemSpace.h"
#include "ISMCIExampleApp_invoke.h"
#include "ITALog_invoke.h"
#include "cdefs.h"
#include "object.h"
#include "qsee_env.h"
#include "qsee_heap.h"
#include "qsee_log.h"
#include "qsee_prng.h"
#include "stringl.h"
#include <stddef.h>
#include <stdint.h>
#include <stdio.h>

typedef struct CSMCIExampleApp {
  int refs;
} CSMCIExampleApp;

/* Call qsee_prng_getdata to get four random bytes and use them to create a uint32_t */
static uint32_t getRandomNumber(uint32_t *randNum)
{
  uint8_t prngData[4] = {0};
  uint32_t lenOut = qsee_prng_getdata(prngData, sizeof(prngData));
  *randNum = 0;

  /* We requested 4 bytes of data from qsee_prng, but it may return less. Use lenout to
   * determine how much of randNum to set. */
  if (lenOut) {
    for (uint32_t i = 0; i < lenOut; i++) {
      uint32_t tmp = prngData[i];
      *randNum |= tmp << 8 * i;
    }
  } else {
    qsee_log(QSEE_LOG_MSG_ERROR, "qsee_prng_getdata failed!");
  }

  return lenOut;
}

/* serviceExample helper function to print the IAdder object's status with qsee_log */
static void logStatus(Object adderObj)
{
  statusStruct status = {0, 0, 0, 0, 0, 0, 0};
  int32_t ret;

  ret = IAdder_getStatus(adderObj, &status);

  if (Object_isERROR(ret)) {
    qsee_log(QSEE_LOG_MSG_ERROR, "IAdder_getStatus failed with error %d!", ret);
    return;
  }
  /* Status is kept in the IAdder object's context as an IAdder_OP code; since these are
   * put into IAdder.h we can convert them to something human-readable here. */
  switch (status.lastCall) {
    case IAdder_OP_add:
      qsee_log(QSEE_LOG_MSG_DEBUG, "Last call: IAdder_add");
      break;
    case IAdder_OP_addArray:
      qsee_log(QSEE_LOG_MSG_DEBUG, "Last call: IAdder_addArray");
      break;
    case IAdder_OP_addStruct:
      qsee_log(QSEE_LOG_MSG_DEBUG, "Last call: IAdder_addStruct");
      break;
    case IAdder_OP_getStatus:
      qsee_log(QSEE_LOG_MSG_DEBUG, "Last call: IAdder_getStatus");
      break;
    default:
      qsee_log(QSEE_LOG_MSG_DEBUG, "No known call has been made to IAdder!");
  }

  /* Log if the last call to IAdder overflowed */
  if (status.resultSaturated == IAdder_OVERFLOW_TRUE) {
    qsee_log(QSEE_LOG_MSG_ERROR, "Last call to adder exceeded UINT32_MAX!");
  }
  /* Log if the IAdder total has ever overflowed */
  if (status.totalSaturated == IAdder_OVERFLOW_TRUE) {
    qsee_log(QSEE_LOG_MSG_ERROR, "IAdder total exceeded UINT32_MAX!");
  }
  qsee_log(QSEE_LOG_MSG_DEBUG, "IAdder total: %u", status.total);
}

/* Release is called to decrement the reference counter of this object. When the
 * reference count reaches 0 (i.e. everything retaining a reference to it has
 * called release), the object is freed. */
static int32_t CSMCIExampleApp_release(CSMCIExampleApp *me)
{
  if (--me->refs == 0) {
    QSEE_FREE_PTR(me);
  }
  return Object_OK;
}

/* When retain is called, this ISMCIExample object's reference count is
 * incremented. This would be called when keeping a new reference to this
 * object. */
static int32_t CSMCIExampleApp_retain(CSMCIExampleApp *me)
{
  me->refs++;
  return Object_OK;
}

/* This example demonstrates how a MemObj can be used to share a buffer between a client and a
 * trusted application, whilst retaining the correct permissions. In this example we showcase
 * a buffer being used to print a message in the TA from the CA, and then update that same
 * buffer to pass a message back to the calling CA. */
static int32_t CSMCIExampleApp_sharedMemoryExample(CSMCIExampleApp *me, Object memObj)
{
  (void)me;

  /* Retrieve the TA's IMemSpace object to map the memory objet we recieved from the caller. */
  Object space = qsee_get_space();
  Object mapping = Object_NULL;
  uint32_t randNum = 0;

  uint64_t bufferAddr, bufferLen;
  uint32_t ret = Object_OK;

  /* Exit early if qsee_get_space() failed */
  if (Object_isNull(space)) {
    qsee_log(QSEE_LOG_MSG_ERROR, "qsee_get_space() failed!");
    space = Object_NULL;
    ret = Object_ERROR;
    goto sharedMemoryExample_cleanup;
  }

  ret = IMemSpace_map(space, memObj, IMemRegion_PERM_RW, &bufferAddr, &bufferLen, &mapping);

  if (Object_isERROR(ret)) {
    qsee_log(QSEE_LOG_MSG_ERROR, "Adding memory object to memspace failed with error %d!", ret);
    mapping = Object_NULL;
    goto sharedMemoryExample_cleanup;
  }

  /* Ensure that the last character of what we're about to print is \0. */
  ((char *)bufferAddr)[bufferLen - 1] = '\0';

  /* Print the message stored in the address specified by bufferAddr, by casting in to a char
   * pointer. */
  qsee_log(QSEE_LOG_MSG_DEBUG, "%s", (char *)bufferAddr);

  /* Write a new message into the buffer to demonstrate that it really is shared with the caller.
   * Insert a 'random' number into the string. */
  char const message[] = "Random number in a shared buffer from the TA:";
  uint32_t validNum = getRandomNumber(&randNum);
  if (!validNum) {
    ret = Object_ERROR;
    goto sharedMemoryExample_cleanup;
  }

  /* Print the message back to the caller to the shared buffer. */
  snprintf((char *)bufferAddr, bufferLen, "%s %u", message, randNum);

sharedMemoryExample_cleanup:
  Object_ASSIGN_NULL(mapping);
  Object_ASSIGN_NULL(space);
  return ret;
}

/* An example showing how using a TA-hosted service inside our TA works. Here we
 * open the IAdder service, hosted by our service TA. This relies on our service TA
 * already being loaded before calling this function. This function, and the IAdder service,
 * will also serve as an example of the types we can pass to functions when using SMCInvoke. */
static int32_t CSMCIExampleApp_serviceExample(CSMCIExampleApp *me)
{
  (void) me;

  Object adderObj = Object_NULL;  // Object reference to IAdder interface
  uint32_t additionResult = 0;
  const uint32_t additionVar1 = 1;  // additionVar1 and additionVar2 are used in IAdder_add() call.
  const uint32_t additionVar2 = 2;
  const uint32_t additionArray[] = {1, 2, 3, 4};  // used in IAdder_addArray
  size_t additionArrayLen = C_LENGTHOF(additionArray);
  additionStruct additionStruct;  // used in IAdder_addStruct

  additionStruct.int1 = 5;
  additionStruct.int2 = 6;

  /* qsee_open() is the function used to obtain a service interface object based on its
   * UID. The IAdder object is hosted in the smci_example_svc TA. The QTEE kernel therefore calls
   * its tz_module_open() implementation. */
  int32_t ret = qsee_open(CAdder_UID, &adderObj);

  if (Object_isERROR(ret)) {
    qsee_log(QSEE_LOG_MSG_ERROR, "Opening IAdder interface failed with error: %d!", ret);
    adderObj = Object_NULL;
    goto serviceExample_cleanup;
  }

  ret = IAdder_add(adderObj, additionVar1, additionVar2, &additionResult);
  if (Object_isERROR(ret)) {
    qsee_log(QSEE_LOG_MSG_ERROR, "IAdder_add failed with error: %d!", ret);
    goto serviceExample_cleanup;
  }
  qsee_log(QSEE_LOG_MSG_DEBUG, "IAdder_add result: %d", additionResult);
  logStatus(adderObj);

  ret = IAdder_addArray(adderObj, additionArray, additionArrayLen, &additionResult);
  if (Object_isERROR(ret)) {
    qsee_log(QSEE_LOG_MSG_ERROR, "Iadder_addArray failed with error :%d!", ret);
    goto serviceExample_cleanup;
  }
  qsee_log(QSEE_LOG_MSG_DEBUG, "IAdder_addArray result: %d", additionResult);
  logStatus(adderObj);

  ret = IAdder_addStruct(adderObj, &additionStruct, &additionResult);
  if (Object_isERROR(ret)) {
    qsee_log(QSEE_LOG_MSG_ERROR, "IAdder_addStruct failed with error: %d!", ret);
    goto serviceExample_cleanup;
  }
  qsee_log(QSEE_LOG_MSG_DEBUG, "IAdder_addStruct result: %d", additionResult);
  logStatus(adderObj);

serviceExample_cleanup:
  Object_ASSIGN_NULL(adderObj);
  return ret;
}

/* This function will open the IHash service using qsee_open and compute the
 * hash of an input buffer from the CA. It's worth noting that this example
 * is similar to the run_invoke_hash example in the QSEECom example app. The
 * important difference is, using pure SMCInvoke, we don't have to worry about
 * managing req/rsp structures and can pass our arguments directly. */
static int32_t CSMCIExampleApp_computeHash(CSMCIExampleApp *me,
                                           const void *input_ptr,
                                           size_t input_len,
                                           uint32_t algorithm_val,
                                           void *hash_ptr,
                                           size_t hash_len,
                                           size_t *hash_len_out)
{
  (void) me;

  int32_t ret = Object_ERROR;
  Object hashObj = Object_NULL;  // Will be populated by an IHash object.

  if (input_len == 0) {  //
    qsee_log(QSEE_LOG_MSG_ERROR, "Input string was too short!");
    return ISMCIExampleApp_ERROR_INPUT_BUFFER_TOO_SMALL;
  }

  if (hash_len < ISMCIExampleApp_HASH_SHA256_SIZE) {
    qsee_log(QSEE_LOG_MSG_ERROR, "Outbut buffer is too small!");
    return Object_ERROR_SIZE_OUT;
  }

  /* Open the hash interface based on the input algorithm. Example only has SHA256, but
   * could easily be extended to SHA1, SHA512, etc. */

  switch (algorithm_val) {
    case ISMCIExampleApp_HASH_SHA256:
      ret = qsee_open(CHashSHA256_UID, &hashObj);
      break;

    default:
      ret = Object_ERROR;
      qsee_log(QSEE_LOG_MSG_ERROR, "Input algorithm not implemented.");
      break;
  }

  if (Object_isERROR(ret)) {
    hashObj = Object_NULL;
    goto hashExample_cleanup;
  }

  ret = IHash_update(hashObj, input_ptr, input_len);
  if (Object_isERROR(ret)) {
    qsee_log(QSEE_LOG_MSG_ERROR, "IHash_update failed with error: %d!", ret);
    goto hashExample_cleanup;
  }

  ret = IHash_final(hashObj, hash_ptr, hash_len, hash_len_out);
  if (Object_isERROR(ret)) {
    qsee_log(QSEE_LOG_MSG_ERROR, "IHash_final failed with error: %d!", ret);
  }

hashExample_cleanup:
  Object_ASSIGN_NULL(hashObj);
  return ret;
}

/* Demonstrates a callback object invocation. */
static int32_t CSMCIExampleApp_callbackObjectExample(CSMCIExampleApp *me, Object callbackObj)
{
  int32_t ret;

  (void)me;

  static const char msg[] = "Callback object invocation.";

  if (Object_isNull(callbackObj)) {
    qsee_log(QSEE_LOG_MSG_ERROR, "ISMCIExampleApp_callbackObjectExample received NULL object!");
    return Object_ERROR;
  }

  ret = ICallbackObjectExample_print(callbackObj, msg, sizeof(msg));
  if (Object_isERROR(ret)) {
    qsee_log(QSEE_LOG_MSG_ERROR, "ICallbackObjectExample_print failed with error: %d!", ret);
    return Object_ERROR;
  }

  return Object_OK;
}

static int32_t CSMCIExampleApp_logSinkExample(CSMCIExampleApp *me)
{
  int32_t ret;
  Object taLog = Object_NULL;
  Object logSink = Object_NULL;

  (void)me;

  /* Open the TA Logging service. This gives the capability to the TA to
   * register a custom LogSink.
   */
  ret = qsee_open(CTALog_UID, &taLog);
  if (Object_isERROR(ret)) {
    qsee_log(QSEE_LOG_MSG_ERROR, "CTALog_UID open failed: %d", ret);
    return Object_ERROR;
  }

  /* Open the local TA LogSink implementation. This LogSink is implemented by
   * the TA.
   */
  ret = CLogSink_open(&logSink);
  if (Object_isERROR(ret)) {
    qsee_log(QSEE_LOG_MSG_ERROR, "CLogSink_open failed: %d", ret);
    goto bail;
  }

  /* Without a handler, the qsee_log() behaves as usual. */
  qsee_log(QSEE_LOG_MSG_DEBUG, "Without handler.");

  /* We now register the local TA LogSink implementation. */
  ret = ITALog_registerSink(taLog, logSink);
  if (Object_isERROR(ret)) {
    qsee_log(QSEE_LOG_MSG_ERROR, "ITALog_registerSink w/ sink failed: %d", ret);
    goto bail;
  }

  /* The following 2 messages will go through the common qsee_log() _and_ our
   * LogSink. As an implementation example, logSink stores the last message.
   */
  qsee_log(QSEE_LOG_MSG_DEBUG, "With handler 1.");
  qsee_log(QSEE_LOG_MSG_DEBUG, "With handler 2.");

  /* To deregister the sink we pass `Object_NULL` to registerSink(). */
  ret = ITALog_registerSink(taLog, Object_NULL);
  if (Object_isERROR(ret)) {
    qsee_log(QSEE_LOG_MSG_ERROR, "ITALog_registerSink w/ NULL failed: %d", ret);
    goto bail;
  }

  /* One final message that is not passing through our sink. */
  qsee_log(QSEE_LOG_MSG_DEBUG, "Without handler again.");

bail:
  /* Releasing the logSink object should print the last message the sink had. */
  Object_ASSIGN_NULL(logSink);
  Object_ASSIGN_NULL(taLog);

  return ret;
}

/* Function declaration for the invoke function of the ISMCIExample interface.
 */
static ISMCIExampleApp_DEFINE_INVOKE(CSMCIExampleApp_invoke, CSMCIExampleApp_,
                                     CSMCIExampleApp *);

int32_t CSMCIExampleApp_open(Object *objOut)
{
  CSMCIExampleApp *me = QSEE_ZALLOC_TYPE(CSMCIExampleApp);
  if (!me) {
    qsee_log(QSEE_LOG_MSG_ERROR, "Memory allocation for CSMCIExampleApp failed!");
    return Object_ERROR_MEM;
  }

  me->refs = 1;
  *objOut = (Object){CSMCIExampleApp_invoke, me};
  return Object_OK;
}
