/*
 * Copyright (c) 2019-2020 The Fuchsia Authors
 *
 * SPDX-License-Identifier:	BSD-3-Clause
 */

#include <common.h>
#include <tee.h>
#include <tee/ta_helper.h>
#include <tee/ta_vx.h>
#include <tee/ta_vx_helper.h>
#include <zircon_uboot/boot_args.h>

#ifdef CONFIG_TA_VX
#define VX_PERM_ATTR_HASH_SIZE 32

static struct tee_optee_ta_uuid vx_uuid = TA_VX_UUID;

static int ta_vx_read_lock_state(enum vx_lock_state *lock_state)
{
	int rc;

	struct tee_param param = { 0 };
	param.attr = TEE_PARAM_ATTR_TYPE_VALUE_OUTPUT;

	rc = ta_call(&vx_uuid, TA_VX_CMD_READ_LOCK_STATE, 1, &param);
	if (!rc) {
		*lock_state = param.u.value.a;
	}

	return rc;
}

static int ta_vx_write_lock_state(enum vx_lock_state lock_state)
{
	struct tee_param param = { 0 };
	param.attr = TEE_PARAM_ATTR_TYPE_VALUE_INPUT;
	param.u.value.a = lock_state;

	return ta_call(&vx_uuid, TA_VX_CMD_WRITE_LOCK_STATE, 1, &param);
}

int ta_vx_lock(void)
{
	int rc = zircon_clear_stored_cmdline();
	if (rc) {
		printf("Failed to clear stored boot args\n");
		return rc;
	}

	return ta_vx_write_lock_state(VX_LOCKED);
}

int ta_vx_lock_if_ephemerally_unlocked(void)
{
	int rc;
	uint32_t vars;

	rc = ta_vx_getvar_all(&vars);
	if (rc) {
		return rc;
	}

	if (!(vars & VX_VAR_EPHEMERALLY_UNLOCKED)) {
		return 0;
	}

	printf("Locking up an ephemerally unlocked device\n");
	rc = ta_vx_lock();
	if (!rc) {
		printf("ta_vx_lock() failed\n");
		return rc;
	}

	rc = ta_vx_getvar_all(&vars);
	if (rc) {
		return rc;
	}

	if (vars & VX_VAR_EPHEMERALLY_UNLOCKED) {
		printf("SHOCKED: device is still ephemerally unlocked after "
		       "successful locking\n");
		return -EINVAL;
	}

	return 0;
}

int ta_vx_unlock(void)
{
	return ta_vx_write_lock_state(VX_UNLOCKED);
}

int ta_vx_is_unlocked(bool *unlocked)
{
	enum vx_lock_state lock_state;
	if (ta_vx_read_lock_state(&lock_state)) {
		return -EIO;
	}
	*unlocked = (lock_state == VX_UNLOCKED);
	return 0;
}

int ta_vx_read_rollback_index(uint32_t slot, uint64_t *rollback_index)
{
	int rc;
	struct tee_param params[2] = { 0 };

	params[0].attr = TEE_PARAM_ATTR_TYPE_VALUE_INPUT;
	params[0].u.value.a = slot;

	params[1].attr = TEE_PARAM_ATTR_TYPE_VALUE_OUTPUT;

	rc = ta_call(&vx_uuid, TA_VX_CMD_READ_ROLLBACK_INDEX,
		     ARRAY_SIZE(params), params);
	if (!rc) {
		*rollback_index = ((uint64_t)params[1].u.value.a) << 32 |
				  params[1].u.value.b;
	}

	return rc;
}

int ta_vx_write_rollback_index(uint32_t slot, uint64_t rollback_index)
{
	struct tee_param params[2] = { 0 };

	params[0].attr = TEE_PARAM_ATTR_TYPE_VALUE_INPUT;
	params[0].u.value.a = slot;

	params[1].attr = TEE_PARAM_ATTR_TYPE_VALUE_INPUT;
	params[1].u.value.a = rollback_index >> 32;
	params[1].u.value.b = rollback_index;

	return ta_call(&vx_uuid, TA_VX_CMD_WRITE_ROLLBACK_INDEX,
		       ARRAY_SIZE(params), params);
}

int ta_vx_read_persistent_value(const char *name, void *buf, size_t buf_len,
				size_t *bytes_read)
{
	int rc;
	struct ta_context context = { 0 };
	struct tee_param params[2] = { 0 };

	struct tee_shm *shm_key = NULL;
	struct tee_shm *shm_buf = NULL;

	size_t key_len = strlen(name) + 1;

	/* Connect to TA. */
	rc = ta_open(&vx_uuid, &context);
	if (rc)
		return rc;

	/* Pass in @key via shared memory. */
	rc = tee_shm_alloc(context.tee, key_len, TEE_SHM_ALLOC, &shm_key);
	if (rc)
		goto out;

	memcpy(shm_key->addr, name, key_len);
	params[0].attr = TEE_PARAM_ATTR_TYPE_MEMREF_INPUT;
	params[0].u.memref.shm = shm_key;
	params[0].u.memref.size = key_len;

	/* Pass in buf via shared memory. */
	rc = tee_shm_alloc(context.tee, buf_len, TEE_SHM_ALLOC, &shm_buf);
	if (rc)
		goto out;

	params[1].attr = TEE_PARAM_ATTR_TYPE_MEMREF_INOUT;
	params[1].u.memref.shm = shm_buf;
	params[1].u.memref.size = buf_len;

	/* Invoke the TA. */
	rc = ta_invoke(&context, TA_VX_CMD_READ_PERSIST_VALUE,
		       ARRAY_SIZE(params), params);
	if (!rc) {
		memcpy(buf, (void *)(params[1].u.memref.shm->addr),
		       params[1].u.memref.size);
		if (bytes_read) {
			*bytes_read = params[1].u.memref.size;
		}
	}

out:
	tee_shm_free(shm_key);
	tee_shm_free(shm_buf);

	ta_close(&context);
	return rc;
}

int ta_vx_write_persistent_value(const char *name, const void *buf,
				 size_t buf_len)
{
	int rc;
	struct ta_context context = { 0 };
	struct tee_param params[2] = { 0 };

	struct tee_shm *shm_key = NULL;
	struct tee_shm *shm_buf = NULL;

	size_t key_len = strlen(name) + 1;

	/* Connect to TA. */
	rc = ta_open(&vx_uuid, &context);
	if (rc)
		return rc;

	/* Pass in @key via shared memory. */
	rc = tee_shm_alloc(context.tee, key_len, TEE_SHM_ALLOC, &shm_key);
	if (rc)
		goto out;

	memcpy(shm_key->addr, name, key_len);
	params[0].attr = TEE_PARAM_ATTR_TYPE_MEMREF_INPUT;
	params[0].u.memref.shm = shm_key;
	params[0].u.memref.size = key_len;

	/* Pass in buf via shared memory. */
	rc = tee_shm_alloc(context.tee, buf_len, TEE_SHM_ALLOC, &shm_buf);
	if (rc)
		goto out;

	memcpy(shm_buf->addr, buf, buf_len);
	params[1].attr = TEE_PARAM_ATTR_TYPE_MEMREF_INPUT;
	params[1].u.memref.shm = shm_buf;
	params[1].u.memref.size = buf_len;

	/* Invoke the TA. */
	rc = ta_invoke(&context, TA_VX_CMD_WRITE_PERSIST_VALUE,
		       ARRAY_SIZE(params), params);

out:
	tee_shm_free(shm_key);
	tee_shm_free(shm_buf);

	ta_close(&context);
	return rc;
}

int ta_vx_delete_persistent_value(const char *name)
{
	int rc;
	struct ta_context context = { 0 };
	struct tee_param params[1] = { 0 };

	struct tee_shm *shm_key = NULL;
	size_t key_len = strlen(name) + 1;

	/* Connect to TA. */
	rc = ta_open(&vx_uuid, &context);
	if (rc)
		return rc;

	/* Pass in @key via shared memory. */
	rc = tee_shm_alloc(context.tee, key_len, TEE_SHM_ALLOC, &shm_key);
	if (rc)
		goto out;
	memcpy(shm_key->addr, name, key_len);
	params[0].attr = TEE_PARAM_ATTR_TYPE_MEMREF_INPUT;
	params[0].u.memref.shm = shm_key;
	params[0].u.memref.size = key_len;

	/* Invoke the TA. */
	rc = ta_invoke(&context, TA_VX_CMD_DELETE_PERSIST_VALUE,
		       ARRAY_SIZE(params), params);

out:
	tee_shm_free(shm_key);
	ta_close(&context);
	return rc;
}

int ta_vx_cprng_draw(void *buf, size_t buf_len)
{
	int rc;
	struct ta_context context = { 0 };
	struct tee_param params[1] = { 0 };
	struct tee_shm *shm_buf = NULL;

	if (!(buf && buf_len))
		return -EINVAL;

	rc = ta_open(&vx_uuid, &context);
	if (rc)
		return rc;

	rc = tee_shm_alloc(context.tee, buf_len, TEE_SHM_ALLOC, &shm_buf);
	if (rc)
		goto out;

	params[0].attr = TEE_PARAM_ATTR_TYPE_MEMREF_INOUT;
	params[0].u.memref.shm = shm_buf;
	params[0].u.memref.size = buf_len;

	rc = ta_invoke(&context, TA_VX_CMD_CPRNG_DRAW, ARRAY_SIZE(params),
		       params);
	if (!rc) {
		memcpy(buf, (void *)(params[0].u.memref.shm->addr),
		       params[0].u.memref.size);
	}

out:
	tee_shm_free(shm_buf);

	ta_close(&context);
	return rc;
}

void ta_vx_exit_bootloader_or_panic(void)
{
	uint32_t vars;

	if (ta_call(&vx_uuid, TA_VX_CMD_EXIT_BOOTLOADER, 0, NULL)) {
		panic("VX finalization failed!");
	}

	/* Make sure we are indeed beyond the bootloader security domain. */
	if (ta_vx_getvar_all(&vars) || !(vars & VX_VAR_OUT_OF_BOOTLOADER)) {
		panic("VX finalization verification failed.");
	}

	return;
}

int ta_vx_read_perm_attr(void *buf, size_t buf_len, size_t *bytes_read)
{
	int rc;
	struct ta_context context = { 0 };
	struct tee_param params[1] = { 0 };

	struct tee_shm *shm_buf = NULL;

	if (!(buf && buf_len)) {
		return -EINVAL;
	}

	/* Connect to TA. */
	rc = ta_open(&vx_uuid, &context);
	if (rc) {
		return rc;
	}

	/* Pass in buf via shared memory. */
	rc = tee_shm_alloc(context.tee, buf_len, TEE_SHM_ALLOC, &shm_buf);
	if (rc) {
		ta_close(&context);
		return rc;
	}

	params[0].attr = TEE_PARAM_ATTR_TYPE_MEMREF_INOUT;
	params[0].u.memref.shm = shm_buf;
	params[0].u.memref.size = buf_len;

	/* Invoke the TA. */
	rc = ta_invoke(&context, TA_VX_CMD_READ_PERM_ATTR, ARRAY_SIZE(params),
		       params);
	if (!rc) {
		memcpy(buf, (void *)(params[0].u.memref.shm->addr),
		       params[0].u.memref.size);
		if (bytes_read) {
			*bytes_read = params[0].u.memref.size;
		}
	}

	tee_shm_free(shm_buf);

	ta_close(&context);
	return rc;
}

int ta_vx_write_perm_attr(const void *buf, size_t buf_len)
{
	int rc;
	struct ta_context context = { 0 };
	struct tee_param params[1] = { 0 };

	struct tee_shm *shm_buf = NULL;

	if (!(buf && buf_len))
		return -EINVAL;

	/* Connect to TA. */
	rc = ta_open(&vx_uuid, &context);
	if (rc) {
		return rc;
	}

	/* Pass in buf via shared memory. */
	rc = tee_shm_alloc(context.tee, buf_len, TEE_SHM_ALLOC, &shm_buf);
	if (rc) {
		ta_close(&context);
		return rc;
	}

	memcpy(shm_buf->addr, buf, buf_len);
	params[0].attr = TEE_PARAM_ATTR_TYPE_MEMREF_INPUT;
	params[0].u.memref.shm = shm_buf;
	params[0].u.memref.size = buf_len;

	/* Invoke the TA. */
	rc = ta_invoke(&context, TA_VX_CMD_WRITE_PERM_ATTR, ARRAY_SIZE(params),
		       params);

	tee_shm_free(shm_buf);

	ta_close(&context);
	return rc;
}

int ta_vx_get_rpmb_status(uint32_t *out_status, uint32_t *out_write_count)
{
	int rc;
	struct tee_param params[1] = { 0 };

	params[0].attr = TEE_PARAM_ATTR_TYPE_VALUE_OUTPUT;
	rc = ta_call(&vx_uuid, TA_VX_CMD_GET_RPMB_STATUS, ARRAY_SIZE(params),
		     params);
	if (rc) {
		return rc;
	}

	if (out_status) {
		*out_status = params[0].u.value.a;
	}
	if (out_write_count) {
		*out_write_count = params[0].u.value.b;
	}

	return 0;
}

int ta_vx_provision_rpmb(void)
{
	return ta_call(&vx_uuid, TA_VX_CMD_PROVISION_RPMB, 0, NULL);
}

int ta_vx_reroute_rpmb_till_reboot(void)
{
	return ta_call(&vx_uuid, TA_VX_CMD_REROUTE_RPMB_TILL_REBOOT, 0, NULL);
}

int ta_vx_reroute_rpmb_to_software(void)
{
	return ta_call(&vx_uuid, TA_VX_CMD_REROUTE_RPMB_TO_SOFTWARE, 0, NULL);
}

int ta_vx_reroute_rpmb_to_hardware(void)
{
	return ta_call(&vx_uuid, TA_VX_CMD_REROUTE_RPMB_TO_HARDWARE, 0, NULL);
}

int ta_vx_delete_perm_attr(void)
{
	return ta_call(&vx_uuid, TA_VX_CMD_DELETE_PERM_ATTR, 0, NULL);
}

int ta_vx_read_perm_attr_hash(void *buf, size_t buf_len)
{
	int rc;
	struct ta_context context = { 0 };
	struct tee_param params[1] = { 0 };

	struct tee_shm *shm_buf = NULL;

	if (!buf || (buf_len < VX_PERM_ATTR_HASH_SIZE))
		return -EINVAL;

	/* Connect to TA. */
	rc = ta_open(&vx_uuid, &context);
	if (rc) {
		return rc;
	}

	/* Pass in buf via shared memory. */
	rc = tee_shm_alloc(context.tee, VX_PERM_ATTR_HASH_SIZE, TEE_SHM_ALLOC,
			   &shm_buf);
	if (rc) {
		ta_close(&context);
		return rc;
	}

	params[0].attr = TEE_PARAM_ATTR_TYPE_MEMREF_INOUT;
	params[0].u.memref.shm = shm_buf;
	params[0].u.memref.size = VX_PERM_ATTR_HASH_SIZE;

	/* Invoke the TA. */
	rc = ta_invoke(&context, TA_VX_CMD_READ_PERM_ATTR_HASH,
		       ARRAY_SIZE(params), params);
	if (!rc) {
		memcpy(buf, (void *)(params[0].u.memref.shm->addr),
		       params[0].u.memref.size);
	}

	tee_shm_free(shm_buf);

	ta_close(&context);
	return rc;
}

int ta_vx_get_perm_attr_status(uint32_t *status)
{
	int rc;

	struct tee_param param = { 0 };
	param.attr = TEE_PARAM_ATTR_TYPE_VALUE_OUTPUT;

	rc = ta_call(&vx_uuid, TA_VX_CMD_GET_PERM_ATTR_STATUS, 1, &param);
	if (!rc) {
		if (status) {
			*status = param.u.value.a;
		}
	}

	return rc;
}

int ta_vx_lock_perm_attr(void)
{
	return ta_call(&vx_uuid, TA_VX_CMD_LOCK_PERM_ATTR, 0, NULL);
}

int ta_vx_provision_usb_hash(void)
{
	return ta_call(&vx_uuid, TA_VX_CMD_PROVISION_USBBOOT_PWD_HASH, 0, NULL);
}

int ta_vx_get_usbboot_status(uint32_t *status)
{
	int rc;

	struct tee_param param = { 0 };
	param.attr = TEE_PARAM_ATTR_TYPE_VALUE_OUTPUT;

	rc = ta_call(&vx_uuid, TA_VX_CMD_GET_USBBOOT_STATUS, 1, &param);
	if (!rc) {
		if (status) {
			*status = param.u.value.a;
		}
	}

	return rc;
}

int ta_vx_getvar_all(uint32_t *out_vars)
{
	int rc;
	struct tee_param params[1] = { 0 };

	if (!out_vars) {
		return -EINVAL;
	}

	params[0].attr = TEE_PARAM_ATTR_TYPE_VALUE_OUTPUT;
	rc = ta_call(&vx_uuid, TA_VX_CMD_GETVAR_ALL, ARRAY_SIZE(params),
		     params);
	if (rc) {
		return rc;
	}

	*out_vars = params[0].u.value.a;
	return 0;
}

#ifdef CONFIG_TA_VX_TESTS
int ta_vx_run_tests(const char *name)
{
	int rc;
	struct ta_context context = { 0 };
	struct tee_param params[1] = { 0 };
	size_t name_len = strlen(name) + 1;
	struct tee_shm *shm_buf = NULL;

	rc = ta_open(&vx_uuid, &context);
	if (rc)
		return rc;

	rc = tee_shm_alloc(context.tee, name_len, TEE_SHM_ALLOC, &shm_buf);
	if (rc)
		goto out;

	memcpy(shm_buf->addr, name, name_len);
	params[0].attr = TEE_PARAM_ATTR_TYPE_MEMREF_INPUT;
	params[0].u.memref.shm = shm_buf;
	params[0].u.memref.size = name_len;

	rc = ta_invoke(&context, TA_VX_CMD_RUN_TESTS, ARRAY_SIZE(params),
		       params);

out:
	tee_shm_free(shm_buf);

	ta_close(&context);
	return rc;
}
#endif // CONFIG_TA_VX_TESTS

#else // CONFIG_TA_VX

static int ta_vx_fail(const char *func)
{
	printf("%s: failure due to disabled VX TA\n", func);
	return -1;
}

int ta_vx_lock()
{
	return ta_vx_fail(__func__);
}

int ta_vx_lock_if_ephemerally_unlocked()
{
	return ta_vx_fail(__func__);
}

int ta_vx_unlock()
{
	return ta_vx_fail(__func__);
}

int ta_vx_is_unlocked(bool *unlocked)
{
	return ta_vx_fail(__func__);
}

int ta_vx_read_rollback_index(uint32_t slot, uint64_t *rollback_index)
{
	return ta_vx_fail(__func__);
}

int ta_vx_write_rollback_index(uint32_t slot, uint64_t rollback_index)
{
	return ta_vx_fail(__func__);
}

int ta_vx_read_persistent_value(const char *name, void *buf, size_t buf_len,
				size_t *bytes_read)
{
	return ta_vx_fail(__func__);
}

int ta_vx_write_persistent_value(const char *name, const void *val_buf,
				 size_t val_buf_len)
{
	return ta_vx_fail(__func__);
}

int ta_vx_delete_persistent_value(const char *name)
{
	return ta_vx_fail(__func__);
}

int ta_vx_cprng_draw(void *buf, size_t buf_len)
{
	return ta_vx_fail(__func__);
}

void ta_vx_exit_bootloader_or_panic()
{
	/* Unlike other VX stubs, this one succeeds by default as a no-op if
	   the VX TA is disabled. */
}

int ta_vx_read_perm_attr(void *buf, size_t buf_len, size_t *bytes_read)
{
	return ta_vx_fail(__func__);
}

int ta_vx_write_perm_attr(const void *val_buf, size_t val_buf_len)
{
	return ta_vx_fail(__func__);
}

int ta_vx_read_perm_attr_hash(void *buf, size_t buf_len)
{
	return ta_vx_fail(__func__);
}

int ta_vx_delete_perm_attr(void)
{
	return ta_vx_fail(__func__);
}

int ta_vx_get_perm_attr_status(uint32_t *status)
{
	return ta_vx_fail(__func__);
}

int ta_vx_lock_perm_attr(void)
{
	return ta_vx_fail(__func__);
}

int ta_vx_get_rpmb_status(uint32_t *out_status, uint32_t *out_write_count)
{
	return ta_vx_fail(__func__);
}

int ta_vx_provision_rpmb(void)
{
	return ta_vx_fail(__func__);
}

int ta_vx_reroute_rpmb_till_reboot(void)
{
	return ta_vx_fail(__func__);
}

int ta_vx_reroute_rpmb_to_software(void)
{
	return ta_vx_fail(__func__);
}

int ta_vx_reroute_rpmb_to_hardware(void)
{
	return ta_vx_fail(__func__);
}

int ta_vx_provision_usb_hash(void)
{
	return ta_vx_fail(__func__);
}

int ta_vx_get_usbboot_status(uint32_t *status)
{
	return ta_vx_fail(__func__);
}

int ta_vx_getvar_all(uint32_t *out_vars)
{
	return ta_vx_fail(__func__);
}

#endif // CONFIG_TA_VX