// SPDX-License-Identifier: GPL-2.0-only
/*
 * Copyright (c) 2025, Linaro Limited
 */
#define pr_fmt(fmt) KBUILD_MODNAME ": " fmt

#include <linux/errno.h>
#include <linux/genalloc.h>
#include <linux/slab.h>
#include <linux/string.h>
#include <linux/tee_core.h>
#include <linux/types.h>
#include "optee_private.h"

struct optee_protmem_dyn_pool {
	struct tee_protmem_pool pool;
	struct gen_pool *gen_pool;
	struct optee *optee;
	size_t page_count;
	u32 *mem_attrs;
	u_int mem_attr_count;
	refcount_t refcount;
	u32 use_case;
	struct tee_shm *protmem;
	/* Protects when initializing and tearing down this struct */
	struct mutex mutex;
};

static struct optee_protmem_dyn_pool *
to_protmem_dyn_pool(struct tee_protmem_pool *pool)
{
	return container_of(pool, struct optee_protmem_dyn_pool, pool);
}

static int init_dyn_protmem(struct optee_protmem_dyn_pool *rp)
{
	int rc;

	rp->protmem = tee_shm_alloc_dma_mem(rp->optee->ctx, rp->page_count);
	if (IS_ERR(rp->protmem)) {
		rc = PTR_ERR(rp->protmem);
		goto err_null_protmem;
	}

	/*
	 * TODO unmap the memory range since the physical memory will
	 * become inaccesible after the lend_protmem() call.
	 *
	 * If the platform supports a hypervisor at EL2, it will unmap the
	 * intermediate physical memory for us and stop cache pre-fetch of
	 * the memory.
	 */
	rc = rp->optee->ops->lend_protmem(rp->optee, rp->protmem,
					  rp->mem_attrs,
					  rp->mem_attr_count, rp->use_case);
	if (rc)
		goto err_put_shm;
	rp->protmem->flags |= TEE_SHM_DYNAMIC;

	rp->gen_pool = gen_pool_create(PAGE_SHIFT, -1);
	if (!rp->gen_pool) {
		rc = -ENOMEM;
		goto err_reclaim;
	}

	rc = gen_pool_add(rp->gen_pool, rp->protmem->paddr,
			  rp->protmem->size, -1);
	if (rc)
		goto err_free_pool;

	refcount_set(&rp->refcount, 1);
	return 0;

err_free_pool:
	gen_pool_destroy(rp->gen_pool);
	rp->gen_pool = NULL;
err_reclaim:
	rp->optee->ops->reclaim_protmem(rp->optee, rp->protmem);
err_put_shm:
	tee_shm_put(rp->protmem);
err_null_protmem:
	rp->protmem = NULL;
	return rc;
}

static int get_dyn_protmem(struct optee_protmem_dyn_pool *rp)
{
	int rc = 0;

	if (!refcount_inc_not_zero(&rp->refcount)) {
		mutex_lock(&rp->mutex);
		if (rp->gen_pool) {
			/*
			 * Another thread has already initialized the pool
			 * before us, or the pool was just about to be torn
			 * down. Either way we only need to increase the
			 * refcount and we're done.
			 */
			refcount_inc(&rp->refcount);
		} else {
			rc = init_dyn_protmem(rp);
		}
		mutex_unlock(&rp->mutex);
	}

	return rc;
}

static void release_dyn_protmem(struct optee_protmem_dyn_pool *rp)
{
	gen_pool_destroy(rp->gen_pool);
	rp->gen_pool = NULL;

	rp->optee->ops->reclaim_protmem(rp->optee, rp->protmem);
	rp->protmem->flags &= ~TEE_SHM_DYNAMIC;

	WARN(refcount_read(&rp->protmem->refcount) != 1, "Unexpected refcount");
	tee_shm_put(rp->protmem);
	rp->protmem = NULL;
}

static void put_dyn_protmem(struct optee_protmem_dyn_pool *rp)
{
	if (refcount_dec_and_test(&rp->refcount)) {
		mutex_lock(&rp->mutex);
		if (rp->gen_pool)
			release_dyn_protmem(rp);
		mutex_unlock(&rp->mutex);
	}
}

static int protmem_pool_op_dyn_alloc(struct tee_protmem_pool *pool,
				     struct sg_table *sgt, size_t size,
				     size_t *offs)
{
	struct optee_protmem_dyn_pool *rp = to_protmem_dyn_pool(pool);
	size_t sz = ALIGN(size, PAGE_SIZE);
	phys_addr_t pa;
	int rc;

	rc = get_dyn_protmem(rp);
	if (rc)
		return rc;

	pa = gen_pool_alloc(rp->gen_pool, sz);
	if (!pa) {
		rc = -ENOMEM;
		goto err_put;
	}

	rc = sg_alloc_table(sgt, 1, GFP_KERNEL);
	if (rc)
		goto err_free;

	sg_set_page(sgt->sgl, phys_to_page(pa), size, 0);
	*offs = pa - rp->protmem->paddr;

	return 0;
err_free:
	gen_pool_free(rp->gen_pool, pa, size);
err_put:
	put_dyn_protmem(rp);

	return rc;
}

static void protmem_pool_op_dyn_free(struct tee_protmem_pool *pool,
				     struct sg_table *sgt)
{
	struct optee_protmem_dyn_pool *rp = to_protmem_dyn_pool(pool);
	struct scatterlist *sg;
	int i;

	for_each_sgtable_sg(sgt, sg, i)
		gen_pool_free(rp->gen_pool, sg_phys(sg), sg->length);
	sg_free_table(sgt);
	put_dyn_protmem(rp);
}

static int protmem_pool_op_dyn_update_shm(struct tee_protmem_pool *pool,
					  struct sg_table *sgt, size_t offs,
					  struct tee_shm *shm,
					  struct tee_shm **parent_shm)
{
	struct optee_protmem_dyn_pool *rp = to_protmem_dyn_pool(pool);

	*parent_shm = rp->protmem;

	return 0;
}

static void pool_op_dyn_destroy_pool(struct tee_protmem_pool *pool)
{
	struct optee_protmem_dyn_pool *rp = to_protmem_dyn_pool(pool);

	mutex_destroy(&rp->mutex);
	kfree(rp);
}

static struct tee_protmem_pool_ops protmem_pool_ops_dyn = {
	.alloc = protmem_pool_op_dyn_alloc,
	.free = protmem_pool_op_dyn_free,
	.update_shm = protmem_pool_op_dyn_update_shm,
	.destroy_pool = pool_op_dyn_destroy_pool,
};

static int get_protmem_config(struct optee *optee, u32 use_case,
			      size_t *min_size, u_int *pa_width,
			      u32 *mem_attrs, u_int *ma_count)
{
	struct tee_param params[2] = {
		[0] = {
			.attr = TEE_IOCTL_PARAM_ATTR_TYPE_VALUE_INOUT,
			.u.value.a = use_case,
		},
		[1] = {
			.attr = TEE_IOCTL_PARAM_ATTR_TYPE_MEMREF_OUTPUT,
		},
	};
	struct optee_shm_arg_entry *entry;
	struct tee_shm *shm_param = NULL;
	struct optee_msg_arg *msg_arg;
	struct tee_shm *shm;
	u_int offs;
	int rc;

	if (mem_attrs && *ma_count) {
		params[1].u.memref.size = *ma_count * sizeof(*mem_attrs);
		shm_param = tee_shm_alloc_priv_buf(optee->ctx,
						   params[1].u.memref.size);
		if (IS_ERR(shm_param))
			return PTR_ERR(shm_param);
		params[1].u.memref.shm = shm_param;
	}

	msg_arg = optee_get_msg_arg(optee->ctx, ARRAY_SIZE(params), &entry,
				    &shm, &offs);
	if (IS_ERR(msg_arg)) {
		rc = PTR_ERR(msg_arg);
		goto out_free_shm;
	}
	msg_arg->cmd = OPTEE_MSG_CMD_GET_PROTMEM_CONFIG;

	rc = optee->ops->to_msg_param(optee, msg_arg->params,
				      ARRAY_SIZE(params), params);
	if (rc)
		goto out_free_msg;

	rc = optee->ops->do_call_with_arg(optee->ctx, shm, offs, false);
	if (rc)
		goto out_free_msg;
	if (msg_arg->ret && msg_arg->ret != TEEC_ERROR_SHORT_BUFFER) {
		rc = -EINVAL;
		goto out_free_msg;
	}

	rc = optee->ops->from_msg_param(optee, params, ARRAY_SIZE(params),
					msg_arg->params);
	if (rc)
		goto out_free_msg;

	if (!msg_arg->ret && mem_attrs &&
	    *ma_count < params[1].u.memref.size / sizeof(*mem_attrs)) {
		rc = -EINVAL;
		goto out_free_msg;
	}

	*min_size = params[0].u.value.a;
	*pa_width = params[0].u.value.c;
	*ma_count = params[1].u.memref.size / sizeof(*mem_attrs);

	if (msg_arg->ret == TEEC_ERROR_SHORT_BUFFER) {
		rc = -ENOSPC;
		goto out_free_msg;
	}

	if (mem_attrs)
		memcpy(mem_attrs, tee_shm_get_va(shm_param, 0),
		       params[1].u.memref.size);

out_free_msg:
	optee_free_msg_arg(optee->ctx, entry, offs);
out_free_shm:
	if (shm_param)
		tee_shm_free(shm_param);
	return rc;
}

struct tee_protmem_pool *optee_protmem_alloc_dyn_pool(struct optee *optee,
						      enum tee_dma_heap_id id)
{
	struct optee_protmem_dyn_pool *rp;
	size_t min_size;
	u_int pa_width;
	int rc;

	rp = kzalloc_obj(*rp);
	if (!rp)
		return ERR_PTR(-ENOMEM);
	rp->use_case = id;

	rc = get_protmem_config(optee, id, &min_size, &pa_width, NULL,
				&rp->mem_attr_count);
	if (rc) {
		if (rc != -ENOSPC)
			goto err;
		rp->mem_attrs = kcalloc(rp->mem_attr_count,
					sizeof(*rp->mem_attrs), GFP_KERNEL);
		if (!rp->mem_attrs) {
			rc = -ENOMEM;
			goto err;
		}
		rc = get_protmem_config(optee, id, &min_size, &pa_width,
					rp->mem_attrs, &rp->mem_attr_count);
		if (rc)
			goto err_kfree_eps;
	}

	rc = optee_set_dma_mask(optee, pa_width);
	if (rc)
		goto err_kfree_eps;

	rp->pool.ops = &protmem_pool_ops_dyn;
	rp->optee = optee;
	rp->page_count = min_size / PAGE_SIZE;
	mutex_init(&rp->mutex);

	return &rp->pool;

err_kfree_eps:
	kfree(rp->mem_attrs);
err:
	kfree(rp);
	return ERR_PTR(rc);
}