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

#include <linux/module.h>
#include <linux/fs.h>
#include <linux/timer.h>
#include <linux/sched.h>
#include <linux/slab.h>
#include <linux/clk.h>
#include <linux/interrupt.h>
#include <linux/of.h>
#include <linux/types.h>
#include <linux/dma-mapping.h>
#include <linux/platform_device.h>
#include <linux/err.h>
#include <linux/firmware.h>
#include <asm/uaccess.h>
#include "comm_debug.h"
#include "comm_teec.h"

static DEFINE_MUTEX(ctx_mtx);
static struct TEEC_UUID uuid_list[] = {AI_TA_UUID, CPUFREQ_TA_UUID, LTP_TA_UUID};
struct comm_teec_handle teec_hnd[MAX_UUID_IDX];
static int prealloc = 1;
static struct device *dev;

static void uuid_to_octets(uint8_t d[TEE_IOCTL_UUID_LEN], const struct TEEC_UUID *s)
{
	d[0] = s->timeLow >> 24;
	d[1] = s->timeLow >> 16;
	d[2] = s->timeLow >> 8;
	d[3] = s->timeLow;
	d[4] = s->timeMid >> 8;
	d[5] = s->timeMid;
	d[6] = s->timeHiAndVersion >> 8;
	d[7] = s->timeHiAndVersion;
	memcpy(d + 8, s->clockSeqAndNode, sizeof(s->clockSeqAndNode));
}

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

static int comm_get_optee_version(struct comm_teec_handle *hnd)
{
	const struct firmware *fw;
	int ret;

	if (hnd->optee_version != -1)
		goto final;

	if (hnd->optee_version == -1) {
		char *sVer = NULL;

		sVer = kmalloc(sizeof(char) * BOOTARG_SIZE, GFP_KERNEL);
		if (sVer == NULL)
			goto final;

		ret = request_firmware_direct(&fw, "/proc/tz2_mstar/version", dev);
		if (ret) {
			hnd->optee_version = 1;
			release_firmware(fw);
			goto free_sVer;
		}

		memcpy(sVer, fw->data, BOOTARG_SIZE);
		sVer[BOOTARG_SIZE - 2] = '\n';
		sVer[BOOTARG_SIZE - 1] = '\0';

		if (strncmp(sVer, "2.4", 2) == 0)
			hnd->optee_version = 2;
		else if (strncmp(sVer, "3.2", 2) == 0)
			hnd->optee_version = 3;
		else
			hnd->optee_version = 1;

		release_firmware(fw);
free_sVer:
		kfree(sVer);
	}

final:
	return hnd->optee_version;
}

static struct comm_teec_user_ctx *comm_teec_uctx_get(struct comm_teec_handle *hnd, u32 user_id)
{
	struct comm_teec_user_ctx *uctx = NULL;
	struct list_head *ptr, *tmp;

	mutex_lock(&hnd->mtx);
	list_for_each_safe(ptr, tmp, &hnd->uctx_list) {
		uctx = list_entry(ptr, struct comm_teec_user_ctx, list);
		if (uctx->user_id == user_id)
			break;

		uctx = NULL;
	}

	if (!uctx) {
		uctx = vzalloc(sizeof(struct comm_teec_user_ctx));
		if (uctx) {
			INIT_LIST_HEAD(&uctx->list);
			uctx->user_id = user_id;
			list_add(&uctx->list, &hnd->uctx_list);
		}
	}
	mutex_unlock(&hnd->mtx);

	return uctx;
}

static void comm_teec_free_shm(struct TEEC_SharedMemory *free_shm, int force)
{
	if (force || !prealloc) {
		tee_shm_free((struct tee_shm *)free_shm->shadow_buffer);
		free_shm->id = -1;
		free_shm->shadow_buffer = NULL;
		free_shm->buffer = NULL;
		free_shm->registered_fd = -1;
		free_shm->alloced_size = 0;
	}
}

static int comm_teec_alloc_shm(struct TEEC_Context *ctx, struct TEEC_SharedMemory *shm)
{
	struct tee_shm *_shm = NULL;
	void *shm_va = NULL;
	size_t len = shm->size;
	struct tee_context *context = (struct tee_context *) ctx->fd;

	len = (shm->size) ? shm->size : 8;

	if (shm->shadow_buffer)
		comm_teec_free_shm(shm, 1);

	shm->shadow_buffer = shm->buffer = NULL;

#if LINUX_VERSION_CODE < KERNEL_VERSION(6, 6, 0)
	_shm = tee_shm_alloc(context, len, TEE_SHM_MAPPED | TEE_SHM_DMA_BUF);
#else
	_shm = tee_shm_alloc_kernel_buf(context, len);
#endif
	if (IS_ERR_OR_NULL(_shm)) {
		comm_err("[%s] tee_shm_alloc fail\n", __func__);
		return -ENOMEM;
	}

	shm_va = tee_shm_get_va(_shm, 0);
	if (IS_ERR_OR_NULL(shm_va)) {
		comm_err("[%s] tee_shm_get_va fail\n", __func__);
		return -ENOMEM;
	}

	shm->buffer = shm_va;
	shm->shadow_buffer = (void *)_shm;
	shm->alloced_size = len;
	shm->registered_fd = -1;

	return 0;
}

static void comm_teec_post_proc_tmpref(uint32_t param_type,
		struct TEEC_TempMemoryReference *tmpref,
		struct tee_param *param,
		struct TEEC_SharedMemory *shm)
{
	if (param_type != TEEC_MEMREF_TEMP_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;
	}

	comm_teec_free_shm(shm, 0);
}

static void comm_teec_post_proc_whole(struct TEEC_RegisteredMemoryReference *memref,
		struct tee_param *param)
{
	struct TEEC_SharedMemory *shm = memref->parent;

	if (shm->flags & TEEC_MEM_OUTPUT) {

		/*
		 * We're using a shadow buffer in this reference, copy back
		 * the shadow buffer into the real buffer now that we've
		 * returned from secure world.
		 */
		if (shm->shadow_buffer && param->u.memref.size <= shm->size)
			memcpy(shm->buffer, shm->shadow_buffer, param->u.memref.size);

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

static void comm_teec_post_proc_partial(uint32_t param_type,
		struct TEEC_RegisteredMemoryReference *memref,
		struct tee_param *param)
{
	if (param_type != TEEC_MEMREF_PARTIAL_INPUT) {
		struct TEEC_SharedMemory *shm = memref->parent;

		/*
		 * We're using a shadow buffer in this reference, copy back
		 * the shadow buffer into the real buffer now that we've
		 * returned from secure world.
		 */
		if (shm->shadow_buffer && param->u.memref.size <= memref->size)
			memcpy((char *)shm->buffer + memref->offset,
					(char *)shm->shadow_buffer + memref->offset,
					param->u.memref.size);

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

static void comm_teec_post_proc_operation(struct TEEC_Operation *operation,
		struct tee_param *params,
		struct TEEC_SharedMemory *shms)
{
	size_t n = 0;
	uint32_t param_type = 0;

	if (!operation)
		return;

	for (n = 0; n < TEEC_CONFIG_PAYLOAD_REF_COUNT; n++) {
		param_type = TEEC_PARAM_TYPE_GET(operation->paramTypes, n);

		switch (param_type) {
		case TEEC_NONE:
		case TEEC_VALUE_INPUT:
			break;
		case TEEC_VALUE_OUTPUT:
		case TEEC_VALUE_INOUT:
			operation->params[n].value.a = params[n].u.value.a;
			operation->params[n].value.b = params[n].u.value.b;
			break;
		case TEEC_MEMREF_TEMP_INPUT:
		case TEEC_MEMREF_TEMP_OUTPUT:
		case TEEC_MEMREF_TEMP_INOUT:
			comm_teec_post_proc_tmpref(param_type,
					&operation->params[n].tmpref, params + n, shms + n);
			break;
		case TEEC_MEMREF_WHOLE:
			comm_teec_post_proc_whole(&operation->params[n].memref,
					params + n);
			break;
		case TEEC_MEMREF_PARTIAL_INPUT:
		case TEEC_MEMREF_PARTIAL_OUTPUT:
		case TEEC_MEMREF_PARTIAL_INOUT:
			comm_teec_post_proc_partial(param_type,
					&operation->params[n].memref, params + n);
			break;
		default:
			break;
		}
	}
}

static int comm_teec_pre_proc_tmpref(struct TEEC_Context *ctx,
		uint32_t param_type, struct TEEC_TempMemoryReference *tmpref,
		struct tee_param *param, struct TEEC_SharedMemory *shm)
{
	switch (param_type) {
	case TEEC_MEMREF_TEMP_INPUT:
		param->attr = TEE_IOCTL_PARAM_ATTR_TYPE_MEMREF_INPUT;
		shm->flags = TEEC_MEM_INPUT;
		break;
	case TEEC_MEMREF_TEMP_OUTPUT:
		param->attr = TEE_IOCTL_PARAM_ATTR_TYPE_MEMREF_OUTPUT;
		shm->flags = TEEC_MEM_OUTPUT;
		break;
	case TEEC_MEMREF_TEMP_INOUT:
		param->attr = TEE_IOCTL_PARAM_ATTR_TYPE_MEMREF_INOUT;
		shm->flags = TEEC_MEM_INPUT | TEEC_MEM_OUTPUT;
		break;
	default:
		return -EINVAL;
	}

	if (!shm->buffer || shm->alloced_size < tmpref->size) {
		shm->size = tmpref->size;

		if (comm_teec_alloc_shm(ctx, shm))
			return -EPERM;
	}

	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 comm_teec_pre_proc_whole(
		struct TEEC_RegisteredMemoryReference *memref,
		struct tee_param *param)
{
	const uint32_t inout = TEEC_MEM_INPUT | TEEC_MEM_OUTPUT;
	uint32_t flags = memref->parent->flags & inout;
	struct TEEC_SharedMemory *shm = NULL;

	if (flags == inout)
		param->attr = TEE_IOCTL_PARAM_ATTR_TYPE_MEMREF_INOUT;
	else if (flags & TEEC_MEM_INPUT)
		param->attr = TEE_IOCTL_PARAM_ATTR_TYPE_MEMREF_INPUT;
	else if (flags & TEEC_MEM_OUTPUT)
		param->attr = TEE_IOCTL_PARAM_ATTR_TYPE_MEMREF_OUTPUT;
	else
		return -EINVAL;

	shm = memref->parent;
	/*
	 * We're using a shadow buffer in this reference, copy the real buffer
	 * into the shadow buffer if needed. We'll copy it back once we've
	 * returned from the call to secure world.
	 */
	if (shm->shadow_buffer && (flags & TEEC_MEM_INPUT))
		memcpy(shm->shadow_buffer, shm->buffer, shm->size);

	param->u.memref.shm = shm->shadow_buffer;
	param->u.memref.size = shm->size;

	return 0;
}

static int comm_teec_pre_proc_partial(uint32_t param_type,
		struct TEEC_RegisteredMemoryReference *memref,
		struct tee_param *param)
{
	uint32_t req_shm_flags = 0;
	struct TEEC_SharedMemory *shm = NULL;

	switch (param_type) {
	case TEEC_MEMREF_PARTIAL_INPUT:
		req_shm_flags = TEEC_MEM_INPUT;
		param->attr = TEE_IOCTL_PARAM_ATTR_TYPE_MEMREF_INPUT;
		break;
	case TEEC_MEMREF_PARTIAL_OUTPUT:
		req_shm_flags = TEEC_MEM_OUTPUT;
		param->attr = TEE_IOCTL_PARAM_ATTR_TYPE_MEMREF_OUTPUT;
		break;
	case TEEC_MEMREF_PARTIAL_INOUT:
		req_shm_flags = TEEC_MEM_OUTPUT | TEEC_MEM_INPUT;
		param->attr = TEE_IOCTL_PARAM_ATTR_TYPE_MEMREF_INOUT;
		break;
	default:
		return -EINVAL;
	}

	shm = memref->parent;

	if ((shm->flags & req_shm_flags) != req_shm_flags)
		return -EPERM;

	/*
	 * We're using a shadow buffer in this reference, copy the real buffer
	 * into the shadow buffer if needed. We'll copy it back once we've
	 * returned from the call to secure world.
	 */
	if (shm->shadow_buffer && param_type != TEEC_MEMREF_PARTIAL_OUTPUT)
		memcpy((char *)shm->shadow_buffer + memref->offset,
				(char *)shm->buffer + memref->offset, memref->size);

	param->u.memref.shm = shm->shadow_buffer;
	param->u.memref.shm_offs = memref->offset;
	param->u.memref.size = memref->size;

	return 0;
}

static int comm_teec_pre_proc_operation(
		struct TEEC_Context *ctx,
		struct TEEC_Operation *operation,
		struct tee_param *params,
		struct TEEC_SharedMemory *shms)
{
	int ret = 0;
	size_t n = 0;
	uint32_t param_type = 0;

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

	for (n = 0; n < TEEC_CONFIG_PAYLOAD_REF_COUNT; n++) {
		param_type = TEEC_PARAM_TYPE_GET(operation->paramTypes, n);

		switch (param_type) {
		case TEEC_NONE:
			params[n].attr = param_type;
			break;
		case TEEC_VALUE_INPUT:
		case TEEC_VALUE_OUTPUT:
		case TEEC_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 TEEC_MEMREF_TEMP_INPUT:
		case TEEC_MEMREF_TEMP_OUTPUT:
		case TEEC_MEMREF_TEMP_INOUT:
			ret = comm_teec_pre_proc_tmpref(ctx, param_type,
				&(operation->params[n].tmpref), params + n, shms + n);
			break;
		case TEEC_MEMREF_WHOLE:
			ret = comm_teec_pre_proc_whole(
					&operation->params[n].memref, params + n);
			break;
		case TEEC_MEMREF_PARTIAL_INPUT:
		case TEEC_MEMREF_PARTIAL_OUTPUT:
		case TEEC_MEMREF_PARTIAL_INOUT:
			ret = comm_teec_pre_proc_partial(param_type,
					&operation->params[n].memref, params + n);
			break;
		default:
			ret = -EINVAL;
			comm_drv_debug("invalid param_type : %x\n", param_type);
			break;
		}
	}

	return ret;
}

static int comm_teec_invoke_cmd(struct comm_teec_handle *hnd,
		struct TEEC_Session *session, u32 user_id, u32 cmd_id,
		struct TEEC_Operation *operation, u32 *error_origin)
{
	int ret = -EPERM;
	uint32_t err_ori = 0;
	uint64_t buf[(sizeof(struct tee_ioctl_invoke_arg) +
			TEEC_CONFIG_PAYLOAD_REF_COUNT *
			sizeof(struct tee_param)) / sizeof(uint64_t)] = { 0 };
	struct tee_ioctl_invoke_arg *arg = NULL;
	struct tee_param *params = NULL;
	//struct TEEC_SharedMemory *shm = hnd->shm;
	struct comm_teec_user_ctx *uctx;
	struct TEEC_SharedMemory *shm;

	uctx = comm_teec_uctx_get(hnd, user_id);
	if (!uctx)
		return -ENOMEM;

	shm = uctx->shm;
	arg = (struct tee_ioctl_invoke_arg *)buf;
	arg->num_params = TEEC_CONFIG_PAYLOAD_REF_COUNT;
	params = (struct tee_param *)(arg + 1);
	arg->session = session->session_id;
	arg->func = cmd_id;

	if (operation)
		operation->session = session;

	ret = comm_teec_pre_proc_operation(&session->ctx, operation, params, shm);
	if (ret) {
		err_ori = TEEC_ORIGIN_API;
		goto post_proc;
	}

	ret = tee_client_invoke_func((struct tee_context *)session->ctx.fd, arg, params);
	if (ret) {
		err_ori = TEEC_ORIGIN_COMMS;
		goto post_proc;
	}

	ret = arg->ret;
	err_ori = arg->ret_origin;
post_proc:
	comm_teec_post_proc_operation(operation, params, shm);

	if (error_origin)
		*error_origin = err_ori;

	return ret;
}

static void comm_teec_close_session(struct TEEC_Session *session)
{
	if (!session)
		return;
	if (!session->ctx.fd)
		return;

	tee_client_close_session((struct tee_context *)session->ctx.fd,
			session->session_id);
}

static int comm_teec_open_session(struct TEEC_Context *context,
			struct TEEC_Session *session,
			const struct TEEC_UUID *destination,
			struct TEEC_Operation *operation,
			u32 *error_origin)
{
	struct tee_ioctl_open_session_arg arg;
	struct tee_param param;
	int rc;

	if ((!session) || (!context))
		return -EPERM;

	memset(&arg, 0, sizeof(struct tee_ioctl_open_session_arg));
	memset(&param, 0, sizeof(struct tee_param));
	uuid_to_octets(arg.uuid, destination);
	arg.clnt_login = TEE_IOCTL_LOGIN_PUBLIC;
	arg.num_params = 0;
	if (operation != NULL)
		param.attr = operation->paramTypes;

	rc = tee_client_open_session(context->ctx, &arg, (struct tee_param *) &param);
	if ((rc == 0) && (arg.ret == 0)) {
		session->session_id = arg.session;
		session->ctx.fd = (uintptr_t)context->ctx;
		if (operation != NULL) {
			operation->params[0].value.a = param.u.value.a;
			operation->params[0].value.b = param.u.value.b;
		}
		return 0;
	}

	if (error_origin != NULL)
		*error_origin = arg.ret_origin;

	return -EPERM;
}


static void comm_teec_exit_context(struct TEEC_Context *context)
{
	if (!context || !context->ctx)
		return;

	tee_client_close_context(context->ctx);
}

static int comm_teec_init_context(struct comm_teec_handle *hnd,
		struct TEEC_Context *context)
{
	if (comm_get_optee_version(hnd) == 3) {
		struct tee_ioctl_version_data vers = {
			.impl_id = TEE_OPTEE_CAP_TZ,
			.impl_caps = TEE_IMPL_ID_OPTEE,
			.gen_caps = TEE_GEN_CAP_GP,
		};

		context->ctx = tee_client_open_context(NULL, _optee_match, NULL, &vers);

		if (IS_ERR_OR_NULL(context->ctx)) {
			comm_err("open context is fail\n");
			return -EPERM;
		}
		return 0;
	}

	return -EPERM;
}

static struct comm_teec_handle *comm_teec_open_context_internal(u32 uuid_idx, u32 invoke)
{
	struct comm_teec_handle *hnd;
	struct TEEC_Context *context;
	struct TEEC_Session *session;
	u32 error = 0;
	int ret = 0;

	if (uuid_idx >= MAX_UUID_IDX)
		return NULL;

	hnd = &teec_hnd[uuid_idx];

	mutex_lock(&ctx_mtx);
	if (kref_read(&hnd->ref) == 0) {
		hnd->optee_version = 3;
		context = &hnd->stContext;
		session = &hnd->stSession;

		ret = comm_teec_init_context(hnd, context);
		if (ret) {
			comm_err("%s(): init context fail\n", __func__);
			comm_teec_exit_context(context);
			goto done;
		}

		ret = comm_teec_open_session(context, session,
				&uuid_list[uuid_idx], NULL, &error);
		if (ret) {
			comm_err("%s(): open session fail, with error(%u)\n",
					__func__, error);
			comm_teec_close_session(session);
			comm_teec_exit_context(context);
			goto done;
		}

		mutex_init(&hnd->mtx);
		kref_init(&hnd->ref);
		INIT_LIST_HEAD(&hnd->uctx_list);
	}

	if (invoke)
		kref_get(&hnd->ref);
done:
	mutex_unlock(&ctx_mtx);

	return (ret) ? NULL : hnd;
}

static void comm_teec_release_context(struct kref *kref)
{
	struct comm_teec_handle *hnd =
			container_of(kref, struct comm_teec_handle, ref);
	struct comm_teec_user_ctx *uctx, *tmp;
	struct TEEC_Context *context;
	struct TEEC_Session *session;
	int i;

	mutex_lock(&hnd->mtx);
	list_for_each_entry_safe(uctx, tmp, &hnd->uctx_list, list) {
		list_del(&uctx->list);

		for (i = 0; i < TEEC_CONFIG_PAYLOAD_REF_COUNT; i++) {
			if (uctx->shm[i].shadow_buffer)
				comm_teec_free_shm(&uctx->shm[i], 1);
		}
		vfree(uctx);
	}
	mutex_unlock(&hnd->mtx);

	context = &hnd->stContext;
	session = &hnd->stSession;

	comm_teec_close_session(session);
	comm_teec_exit_context(context);

	mutex_unlock(&ctx_mtx);
}

void comm_teec_close_context_internal(struct comm_teec_handle *hnd)
{
	kref_put_mutex(&hnd->ref, comm_teec_release_context, &ctx_mtx);
}

void comm_teec_close_context(void *hnd)
{
	if (!hnd)
		return;

	comm_teec_close_context_internal((struct comm_teec_handle *)hnd);
}

void *comm_teec_open_context(u32 uuid_idx)
{
	return comm_teec_open_context_internal(uuid_idx, 0);
}

int comm_teec_send_command(u32 uuid_idx, u32 user_id,
	u32 cmd_id, struct TEEC_Operation *op, u32 *error)
{
	int ret = 0;
	struct comm_teec_handle *hnd;

	if (user_id >= COMM_TEEC_UCTX_MAX)
		return -EFAULT;

	hnd = comm_teec_open_context_internal(uuid_idx, 1);
	if (!hnd) {
		comm_err("[%s] open context fail\n", __func__);
		return -EFAULT;
	}

	ret = comm_teec_invoke_cmd(hnd, &hnd->stSession, user_id, cmd_id, op, error);
	if (ret) {
		comm_err("[%s] invoke fail, error(%u)\n", __func__, *error);
		ret = -EPERM;
	}

	comm_teec_close_context_internal(hnd);

	return ret;
}

int comm_teec_init(struct platform_device *pdev)
{
	dev = &pdev->dev;

	return 0;
}
