// SPDX-License-Identifier: BSD-2-Clause
/*
 * Copyright (c) 2022 Google
 */

#include <common.h>
#include <emmc_partitions.h>
#include <tee.h>
#include <tee/optee.h>
#include <zircon_uboot/partition.h>

#include "optee_msg.h"
#include "optee_msg_supplicant.h"
#include "optee_private.h"

/*
 * Read `length` bytes from `offset` of a designated raw partition into
 * `buffer`. Buffer size is guaranteed to be >= `length`.
 */
static u32 do_read(u64 offset, u64 length, void *buffer)
{
	zircon_partition *part = zircon_get_partition(MMC_SECURE_STORAGE_NAME);
	if (!part) {
		return TEE_ERROR_ITEM_NOT_FOUND;
	}

	if (!part->read) {
		printf("Optee raw IO: no read implementation\n");
		zircon_free_partition(part);
		return TEE_ERROR_NOT_SUPPORTED;
	}

	// Boundary check will be performed in read(). The API supports unaligned
	// read.
	int res = part->read(part, offset, buffer, length);
	if (res) {
		printf("Optee raw IO: read failed. %d\n", res);
		zircon_free_partition(part);
		return TEE_ERROR_GENERIC;
	}

	zircon_free_partition(part);
	return TEE_SUCCESS;
}

/*
 * Write `length` bytes from `buffer` into a designated raw partition at
 * `offset`. Buffer size is guaranteed to be >= `length`.
 */
static u32 do_write(u64 offset, u64 length, const void *buffer)
{
	zircon_partition *part = zircon_get_partition(MMC_SECURE_STORAGE_NAME);
	if (!part) {
		return TEE_ERROR_ITEM_NOT_FOUND;
	}

	if (!part->write) {
		printf("Optee raw IO: no write implementation\n");
		zircon_free_partition(part);
		return TEE_ERROR_NOT_SUPPORTED;
	}

	// Boundary check will be performed in write(). The API supports unaligned
	// write.
	int res = part->write(part, offset, buffer, length);
	if (res) {
		printf("Optee raw IO: write failed. %d\n", res);
		zircon_free_partition(part);
		return TEE_ERROR_GENERIC;
	}

	zircon_free_partition(part);
	return TEE_SUCCESS;
}

void optee_suppl_cmd_raw_io(struct optee_msg_arg *arg)
{
	u64 offset;
	u64 length;
	void *buffer;
	u64 buffer_size;

	/* The supplicant is part of the REE<>TEE communications stack. */
	arg->ret_origin = TEE_ORIGIN_COMMS;

	if (!(arg->num_params == 2 &&
	      arg->params[0].attr == OPTEE_MSG_ATTR_TYPE_VALUE_INPUT)) {
		arg->ret = TEE_ERROR_BAD_PARAMETERS;
		return;
	}

	offset = arg->params[0].u.value.b;
	length = arg->params[0].u.value.c;
	buffer = (u8 *)arg->params[1].u.tmem.buf_ptr;
	buffer_size = arg->params[1].u.tmem.size;
	if (!(buffer && buffer_size > 0 && buffer_size >= length)) {
		arg->ret = TEE_ERROR_BAD_PARAMETERS;
		return;
	}

	switch (arg->params[0].u.value.a) {
	case OPTEE_MRC_RAW_READ:
		if (arg->params[1].attr != OPTEE_MSG_ATTR_TYPE_TMEM_OUTPUT) {
			arg->ret = TEE_ERROR_BAD_PARAMETERS;
			return;
		}

		arg->ret = do_read(offset, length, buffer);
		break;
	case OPTEE_MRC_RAW_WRITE:
		if (arg->params[1].attr != OPTEE_MSG_ATTR_TYPE_TMEM_INPUT) {
			arg->ret = TEE_ERROR_BAD_PARAMETERS;
			return;
		}

		arg->ret = do_write(offset, length, buffer);
		break;
	default:
		arg->ret = TEE_ERROR_BAD_PARAMETERS;
		break;
	}

	return;
}

#if defined(DEV_BUILD_CONFIG)
int optee_suppl_cmd_raw_io_wrapper(enum optee_suppl_cmd_raw_io_op op,
				   void *buffer, size_t offset, size_t length)
{
	struct optee_suppl_cmd_raw_io_test_param {
		struct optee_msg_arg arg;
		struct optee_msg_param params[2];
	} test_arg = {
		.arg = {
			.num_params = 2,
		},
		.params[0] = {
			.attr = OPTEE_MSG_ATTR_TYPE_VALUE_INPUT,
			.u.value.b = offset,
			.u.value.c = length,
		},
		.params[1] = {
			.u.tmem.buf_ptr = (u64)buffer,
			.u.tmem.size = length,
		},
	};
	switch (op) {
	case optee_suppl_cmd_raw_io_read:
		test_arg.arg.params[1].attr = OPTEE_MSG_ATTR_TYPE_TMEM_OUTPUT;
		test_arg.arg.params[0].u.value.a = OPTEE_MRC_RAW_READ;
		break;
	case optee_suppl_cmd_raw_io_write:
		test_arg.arg.params[1].attr = OPTEE_MSG_ATTR_TYPE_TMEM_INPUT;
		test_arg.arg.params[0].u.value.a = OPTEE_MRC_RAW_WRITE;
		break;
	default:
		return -1;
	}

	optee_suppl_cmd_raw_io(&test_arg.arg);
	return test_arg.arg.ret == TEE_SUCCESS ? 0 : -1;
}
#endif