// SPDX-License-Identifier: (GPL-2.0 OR BSD-3-Clause)
/*
 * Copyright (c) 2020 MediaTek Inc.
 */

#include <linux/kernel.h>
#include <linux/module.h>
#include <linux/version.h>
#if LINUX_VERSION_CODE >= KERNEL_VERSION(5, 15, 0)
#include "tee_drv_5_15_13.h"
#else
#include <linux/tee_drv.h>
#endif

#include "mdrv_demux_tee_comm.h"
#include "mdrv_demux_debug.h"


#define TEEC_CONFIG_PAYLOAD_REF_COUNT	4

 /**
  * Encode the paramTypes according to the supplied types.
  *
  * @param p0 The first param type.
  * @param p1 The second param type.
  * @param p2 The third param type.
  * @param p3 The fourth param type.
  */
#define TEEC_PARAM_TYPES(p0, p1, p2, p3) \
	 ((p0) | ((p1) << 4) | ((p2) << 8) | ((p3) << 12))

 /**
  * Get the i_th param type from the paramType.
  *
  * @param p The paramType.
  * @param i The i-th parameter to get the type for.
  */
#define TEEC_PARAM_TYPE_GET(p, i) (((p) >> (i * 4)) & 0xF)

struct TEEC_SharedMemory_2 {
	void *buffer;
	size_t size;
	uint32_t flags;
	/*
	 * Implementation-Defined
	 */
	int id;
	size_t alloced_size;
	void *shadow_buffer;
};

struct TempMemoryReference {
	void *buffer;
	size_t size;
};

struct TempValue {
	uint32_t a;
	uint32_t b;
};

struct parameter {
	struct TempValue value;
	struct TempMemoryReference tmpref;
};

struct TEEC_Operation {
	uint32_t paramTypes;
	struct parameter params[TEEC_CONFIG_PAYLOAD_REF_COUNT];
};

struct mtk_demux_tee_context {
	struct tee_context *ctx;
	bool bInit;
} dmx_tee_ctx;


static int teec_allocatesharedmemory(struct tee_context *ctx,
		struct TEEC_SharedMemory_2 *shm)
{
	struct tee_shm *_shm = NULL;
	void *shm_va = NULL;
	size_t len;
	struct tee_context *context = ctx;

	shm->shadow_buffer = shm->buffer = NULL;
	len = shm->size;
	if (!len)
		len = 8;

	_shm = tee_shm_alloc(context, len, TEE_SHM_MAPPED | TEE_SHM_DMA_BUF);
	if (IS_ERR(_shm)) {
		DEMUX_DEBUG_ERR("[%s][%d] tee_shm_alloc fail\n", __func__, __LINE__);
		return -ENOMEM;
	}

	shm_va = tee_shm_get_va(_shm, 0);
	if (IS_ERR(shm_va)) {
		DEMUX_DEBUG_ERR("[%s][%d] tee_shm_get_va fail\n", __func__, __LINE__);
		return -ENOMEM;
	}
	shm->buffer = shm_va;
	shm->shadow_buffer = (void *)_shm;
	shm->alloced_size = len;

	return 0;
}

static void teec_releasesharedmemory(struct TEEC_SharedMemory_2 *shm)
{
	if (shm->shadow_buffer)
		tee_shm_free((struct tee_shm *)shm->shadow_buffer);
	shm->shadow_buffer = NULL;
	shm->buffer = NULL;
}

static int teec_pre_process_tmpref(struct tee_context *ctx,
		uint32_t param_type, struct TempMemoryReference *tmpref,
		struct tee_param *param, struct TEEC_SharedMemory_2 *shm)
{
	int res;

	param->attr = param_type;
	shm->size = tmpref->size;

	res = teec_allocatesharedmemory(ctx, shm);
	if (res != 0)
		return res;

	memcpy(shm->buffer, tmpref->buffer, tmpref->size);
	param->u.memref.size = tmpref->size;
	//Workaround, shm->shadow_buffer is shm object.
	param->u.memref.shm = (struct tee_shm *)shm->shadow_buffer;
	return 0;
}

static int teec_pre_process_operation(struct tee_context *ctx,
		struct TEEC_Operation *operation,
		struct tee_param *params,
		struct TEEC_SharedMemory_2 *shms)
{
	int res;
	size_t n;
	uint32_t param_type;

	memset(shms, 0, sizeof(*shms) * TEEC_CONFIG_PAYLOAD_REF_COUNT);
	if (!operation) {
		memset(params, 0, sizeof(*params) * TEEC_CONFIG_PAYLOAD_REF_COUNT);
		return 0;
	}

	for (n = 0; n < TEEC_CONFIG_PAYLOAD_REF_COUNT; n++) {
		param_type = TEEC_PARAM_TYPE_GET(operation->paramTypes, n);
		switch (param_type) {
		case TEE_IOCTL_PARAM_ATTR_TYPE_NONE:
			params[n].attr = param_type;
			break;
		case TEE_IOCTL_PARAM_ATTR_TYPE_VALUE_INPUT:
		case TEE_IOCTL_PARAM_ATTR_TYPE_VALUE_OUTPUT:
		case TEE_IOCTL_PARAM_ATTR_TYPE_VALUE_INOUT:
			params[n].attr = param_type;
			params[n].u.value.a = operation->params[n].value.a;
			params[n].u.value.b = operation->params[n].value.b;
			break;
		case TEE_IOCTL_PARAM_ATTR_TYPE_MEMREF_INPUT:
		case TEE_IOCTL_PARAM_ATTR_TYPE_MEMREF_OUTPUT:
		case TEE_IOCTL_PARAM_ATTR_TYPE_MEMREF_INOUT:
			res = teec_pre_process_tmpref(ctx, param_type,
					&operation->params[n].tmpref,
					params + n, shms + n);
			if (res != 0)
				return res;
			break;
		default:
			DEMUX_DEBUG_ERR("[%s][%d] param_type not support!\n",
				__func__, __LINE__);
			return -EINVAL;
		}
	}

	return 0;
}

static void teec_post_process_tmpref(uint32_t param_type,
		struct TempMemoryReference *tmpref,
		struct tee_param *param,
		struct TEEC_SharedMemory_2 *shm)
{
	if (param_type != TEE_IOCTL_PARAM_ATTR_TYPE_MEMREF_INPUT) {
		if (param->u.memref.size <= tmpref->size && tmpref->buffer)
			memcpy(tmpref->buffer, shm->buffer, param->u.memref.size);

		tmpref->size = param->u.memref.size;
	}
}

static void teec_post_process_operation(struct TEEC_Operation *operation,
		struct tee_param *params, struct TEEC_SharedMemory_2 *shms)
{
	size_t n;

	if (!operation)
		return;

	for (n = 0; n < TEEC_CONFIG_PAYLOAD_REF_COUNT; n++) {
		uint32_t param_type;

		param_type = TEEC_PARAM_TYPE_GET(operation->paramTypes, n);
		switch (param_type) {
		case TEE_IOCTL_PARAM_ATTR_TYPE_VALUE_INPUT:
			break;
		case TEE_IOCTL_PARAM_ATTR_TYPE_VALUE_OUTPUT:
		case TEE_IOCTL_PARAM_ATTR_TYPE_VALUE_INOUT:
			operation->params[n].value.a = params[n].u.value.a;
			operation->params[n].value.b = params[n].u.value.b;
			break;
		case TEE_IOCTL_PARAM_ATTR_TYPE_MEMREF_INPUT:
		case TEE_IOCTL_PARAM_ATTR_TYPE_MEMREF_OUTPUT:
		case TEE_IOCTL_PARAM_ATTR_TYPE_MEMREF_INOUT:
			teec_post_process_tmpref(param_type,
				&operation->params[n].tmpref,
				params + n, shms + n);
			break;
		default:
			break;
		}
	}
}

static void teec_free_temp_refs(struct TEEC_Operation *operation,
		struct TEEC_SharedMemory_2 *shms)
{
	size_t n;

	if (!operation)
		return;

	for (n = 0; n < TEEC_CONFIG_PAYLOAD_REF_COUNT; n++) {
		switch (TEEC_PARAM_TYPE_GET(operation->paramTypes, n)) {
		case TEE_IOCTL_PARAM_ATTR_TYPE_MEMREF_INPUT:
		case TEE_IOCTL_PARAM_ATTR_TYPE_MEMREF_OUTPUT:
		case TEE_IOCTL_PARAM_ATTR_TYPE_MEMREF_INOUT:
			teec_releasesharedmemory(shms + n);
			break;
		default:
			break;
		}
	}
}

static int _optee_match(struct tee_ioctl_version_data *data, const void *vers)
{
	return 1;
}

int mdrv_demux_tee_open_ctx(void)
{
	struct tee_ioctl_version_data vers = {
		.impl_id = TEE_IMPL_ID_OPTEE,
		.impl_caps = TEE_OPTEE_CAP_TZ,
		.gen_caps = TEE_GEN_CAP_GP,
	};

	if (dmx_tee_ctx.bInit) {
		DEMUX_DEBUG_ERR("[%s][%d] optee context already opened\n", __func__, __LINE__);
		return 0;
	}

	dmx_tee_ctx.ctx = tee_client_open_context(NULL,
	_optee_match, NULL, &vers);

	if (IS_ERR(dmx_tee_ctx.ctx)) {
		DEMUX_DEBUG_ERR("[%s][%d] failed to open context\n", __func__, __LINE__);
		return -EFAULT;
	}

	dmx_tee_ctx.bInit = true;
	return 0;
}

int mdrv_demux_tee_close_ctx(void)
{
	if (!dmx_tee_ctx.bInit) {
		DEMUX_DEBUG_ERR("[%s][%d] optee context already closed\n", __func__, __LINE__);
		return 0;
	}

	tee_client_close_context(dmx_tee_ctx.ctx);
	dmx_tee_ctx.bInit = false;
	return 0;
}

int mdrv_demux_tee_open_session(u32 *session_id, uuid_t uuid)
{
	struct tee_ioctl_open_session_arg arg = { 0 };
	struct tee_param param = { 0 };
	int rc;

	if (!session_id) {
		DEMUX_DEBUG_ERR("[%s][%d] invalid parameter\n", __func__, __LINE__);
		return -EINVAL;
	}

	memcpy(arg.uuid, &uuid, TEE_IOCTL_UUID_LEN);
	arg.clnt_login = TEE_IOCTL_LOGIN_PUBLIC;

	rc = tee_client_open_session(dmx_tee_ctx.ctx, &arg, (struct tee_param *)&param);

	if (rc == 0) {
		*session_id = arg.session;
		DEMUX_DEBUG_ERR("[%s][%d] session id = %u\n",
						__func__, __LINE__, *session_id);
	} else {
		DEMUX_DEBUG_ERR("[%s][%d] tee_client_open_session() fail,rc=0x%x\n",
						__func__, __LINE__, rc);
		return -EINVAL;
	}

	return 0;
}

int mdrv_demux_tee_close_session(u32 session_id)
{
	return tee_client_close_session(dmx_tee_ctx.ctx, session_id);
}

int mdrv_demux_tee_invoke_cmd(u32 session_id, u32 cmd, void *param, size_t size)
{
	struct TEEC_Operation op = { 0 };
	struct tee_ioctl_invoke_arg arg = { 0 };
	struct tee_param params[TEEC_CONFIG_PAYLOAD_REF_COUNT] = { 0 };
	struct TEEC_SharedMemory_2 shm[TEEC_CONFIG_PAYLOAD_REF_COUNT] = { 0 };
	int res, rc;
	uint32_t eorig = 0;

	if (size) {
		op.params[0].tmpref.buffer = param;
		op.params[0].tmpref.size = size;
		op.paramTypes = TEEC_PARAM_TYPES(TEE_IOCTL_PARAM_ATTR_TYPE_MEMREF_INOUT,
							TEE_IOCTL_PARAM_ATTR_TYPE_NONE,
							TEE_IOCTL_PARAM_ATTR_TYPE_NONE,
							TEE_IOCTL_PARAM_ATTR_TYPE_NONE);
	} else {
		op.paramTypes = TEEC_PARAM_TYPES(TEE_IOCTL_PARAM_ATTR_TYPE_NONE,
							TEE_IOCTL_PARAM_ATTR_TYPE_NONE,
							TEE_IOCTL_PARAM_ATTR_TYPE_NONE,
							TEE_IOCTL_PARAM_ATTR_TYPE_NONE);
	}

	arg.num_params = TEEC_CONFIG_PAYLOAD_REF_COUNT;
	arg.session = session_id;
	arg.func = cmd;

	res = teec_pre_process_operation(dmx_tee_ctx.ctx, &op, params, shm);
	if (res != 0) {
		eorig = 1;
		goto out_free_temp_refs;
	}

	rc = tee_client_invoke_func(dmx_tee_ctx.ctx, &arg, params);
	if (rc) {
		eorig = 2;
		res = rc;
		goto out_free_temp_refs;
	}

	res = arg.ret;
	eorig = arg.ret_origin;
	teec_post_process_operation(&op, params, shm);

out_free_temp_refs:
	teec_free_temp_refs(&op, shm);
	return res;
}
