// SPDX-License-Identifier: BSD-2-Clause
/*
 * SPDX-FileCopyrightText: Copyright (c) 2021-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 */

#include <config.h>
#include <crypto/crypto.h>
#include <drivers/tegra/tegra_fuse.h>
#include <drivers/tegra/tegra_se_aes.h>
#include <drivers/tegra/tegra_se_kdf.h>
#include <drivers/tegra/tegra_se_rng.h>
#include <initcall.h>
#include <pta_jetson_user_key.h>
#include <kernel/panic.h>
#include <kernel/boot.h>
#include <kernel/dt.h>
#include <kernel/pseudo_ta.h>
#include <libfdt.h>
#include <stdio.h>
#include <string.h>
#include <trace.h>
#include <sys/queue.h>

#include <mm/core_memprot.h>
#include "jetson_decrypt_cpubl_payload.h"
#include "jetson_user_key_pta_compute_cmac.h"
#include "jetson_user_key_pta_ftpm.h"
#include "jetson_ekb_extraction.h"
#include "jetson_ekb_user_key_management.h"
#include "jetson_nist_sp_800_108_kdf.h"

#define JETSON_USER_KEY_PTA_NAME	"jetson_user_key_pta.ta"

#if defined(PLATFORM_FLAVOR_t234) || defined(PLATFORM_FLAVOR_t264)
#define EKB_MEMORY_AREA_TYPE		MEM_AREA_RAM_SEC
#endif

static vaddr_t ekb_base_addr;

/*
 * Jetson service flag:
 *    LUKS_SRV_FLAG: Turn on/off for CA to query passphrase
 *    PTA_SRV_FLAG: Indicate the TA is available or not
 */
static uint32_t jetson_service_flag = LUKS_SRV_FLAG |
				      PTA_SRV_FLAG;

static TEE_Result jetson_user_key_pta_get_user_key(uint8_t **key,
						   uint32_t *key_len,
						   uint32_t key_type)
{
	TEE_Result rc = TEE_SUCCESS;
	struct user_key *ekb_key = NULL;

	if (!key || !key_len)
		return TEE_ERROR_BAD_PARAMETERS;

	if (key_type >= EKB_USER_KEY_NUM_MAX)
		return TEE_ERROR_BAD_PARAMETERS;

	rc = jetson_pta_get_key_by_type(&ekb_key, key_type);
	if (rc)
		return rc;

	*key = ekb_key->data;
	*key_len = ekb_key->len;

	return rc;
}

static TEE_Result jetson_pta_get_ekb_key(uint32_t ptypes,
					 TEE_Param params[TEE_NUM_PARAMS])
{
	uint32_t exp_pt = TEE_PARAM_TYPES(TEE_PARAM_TYPE_VALUE_INPUT,
					  TEE_PARAM_TYPE_MEMREF_OUTPUT,
					  TEE_PARAM_TYPE_NONE,
					  TEE_PARAM_TYPE_NONE);
	TEE_Result rc = TEE_SUCCESS;
	uint32_t key_type = params[0].value.a;
	uint8_t *key = NULL;
	uint32_t key_len = 0;

	if (exp_pt != ptypes)
		return TEE_ERROR_BAD_PARAMETERS;

	rc = jetson_user_key_pta_get_user_key(&key, &key_len, key_type);
	if (rc)
		return rc;

	if (key_len <= params[1].memref.size) {
		memcpy(params[1].memref.buffer, key, key_len);
		params[1].memref.size = key_len;
	} else {
		return TEE_ERROR_SHORT_BUFFER;
	}

	return rc;
}

static TEE_Result jetson_pta_get_random(uint32_t ptypes,
					TEE_Param params[TEE_NUM_PARAMS])
{
	uint32_t exp_pt = TEE_PARAM_TYPES(TEE_PARAM_TYPE_MEMREF_OUTPUT,
					  TEE_PARAM_TYPE_NONE,
					  TEE_PARAM_TYPE_NONE,
					  TEE_PARAM_TYPE_NONE);
	TEE_Result rc = TEE_SUCCESS;

	if (exp_pt != ptypes)
		return TEE_ERROR_BAD_PARAMETERS;

	rc = tegra_se_rng_get_random(params[0].memref.buffer,
				     params[0].memref.size);
	if (rc != TEE_SUCCESS) {
		memset(params[0].memref.buffer, 0, params[0].memref.size);
		params[0].memref.size = 0;
	}

	return rc;
}

static TEE_Result jetson_pta_gen_unique_key_by_ekb(uint32_t ptypes,
					TEE_Param params[TEE_NUM_PARAMS])
{
	uint32_t exp_pt = TEE_PARAM_TYPES(TEE_PARAM_TYPE_VALUE_INPUT,
					  TEE_PARAM_TYPE_MEMREF_INPUT,
					  TEE_PARAM_TYPE_MEMREF_OUTPUT,
					  TEE_PARAM_TYPE_NONE);
	TEE_Result rc = TEE_SUCCESS;
	fuse_ecid_t *ecid;
	char ecid_str[36];
	uint8_t *key;
	uint32_t key_len;
	uint32_t key_type = params[0].value.a;

	if ((exp_pt != ptypes) ||
	    (params[1].memref.size == 0)) {
		return TEE_ERROR_BAD_PARAMETERS;
	}

	rc = jetson_user_key_pta_get_user_key(&key, &key_len, key_type);
	if (rc)
		return rc;

	ecid = tegra_fuse_get_ecid();
	snprintf(ecid_str, sizeof(ecid_str), "%08x%08x%08x%08x",
		 ecid->ecid[3], ecid->ecid[2], ecid->ecid[1],
		 ecid->ecid[0]);

	rc = nist_sp_800_108_kdf(key,
				 key_len,
				 ecid_str,
				 strlen(ecid_str),
				 params[1].memref.buffer,
				 params[1].memref.size,
				 params[2].memref.size,
				 params[2].memref.buffer);

	return rc;
}

static TEE_Result jetson_pta_gen_key(uint32_t ptypes,
				     TEE_Param params[TEE_NUM_PARAMS])
{
	uint32_t exp_pt = TEE_PARAM_TYPES(TEE_PARAM_TYPE_MEMREF_INPUT,
					  TEE_PARAM_TYPE_MEMREF_INPUT,
					  TEE_PARAM_TYPE_MEMREF_INPUT,
					  TEE_PARAM_TYPE_MEMREF_OUTPUT);
	TEE_Result rc = TEE_SUCCESS;

	if ((exp_pt != ptypes) ||
	    (params[1].memref.size == 0) ||
	    (params[2].memref.size == 0)) {
		return TEE_ERROR_BAD_PARAMETERS;
	}

	rc = nist_sp_800_108_kdf(params[0].memref.buffer,
				 params[0].memref.size,
				 params[1].memref.buffer,
				 params[1].memref.size,
				 params[2].memref.buffer,
				 params[2].memref.size,
				 params[3].memref.size,
				 params[3].memref.buffer);

	return rc;
}

static TEE_Result jetson_pta_get_flag(uint32_t ptypes,
				      TEE_Param params[TEE_NUM_PARAMS])
{
	uint32_t exp_pt = TEE_PARAM_TYPES(TEE_PARAM_TYPE_VALUE_INPUT,
					  TEE_PARAM_TYPE_VALUE_OUTPUT,
					  TEE_PARAM_TYPE_NONE,
					  TEE_PARAM_TYPE_NONE);
	TEE_Result rc = TEE_SUCCESS;

	if (exp_pt != ptypes)
		return TEE_ERROR_BAD_PARAMETERS;

	switch (params[0].value.a) {
	case LUKS_SRV_FLAG:
		break;
	default:
		return TEE_ERROR_ACCESS_DENIED;
	}

	params[1].value.a = jetson_service_flag & params[0].value.a;

	return rc;
}

static TEE_Result jetson_pta_set_flag(uint32_t ptypes,
				      TEE_Param params[TEE_NUM_PARAMS])
{
	uint32_t exp_pt = TEE_PARAM_TYPES(TEE_PARAM_TYPE_VALUE_INPUT,
					  TEE_PARAM_TYPE_NONE,
					  TEE_PARAM_TYPE_NONE,
					  TEE_PARAM_TYPE_NONE);
	TEE_Result rc = TEE_SUCCESS;

	if (exp_pt != ptypes)
		return TEE_ERROR_BAD_PARAMETERS;

	switch (params[0].value.a) {
	case LUKS_SRV_FLAG:
		break;
	default:
		return TEE_ERROR_ACCESS_DENIED;
	}
	if (params[0].value.b)
		jetson_service_flag |= params[0].value.a;
	else
		jetson_service_flag &= ~params[0].value.a;

	return rc;
}

static TEE_Result jetson_pta_is_user_key_exists(uint32_t ptypes,
				      TEE_Param params[TEE_NUM_PARAMS])
{
	uint32_t key_type;
	uint8_t *key;
	uint32_t key_len;
	uint32_t exp_pt = TEE_PARAM_TYPES(TEE_PARAM_TYPE_VALUE_INPUT,
					  TEE_PARAM_TYPE_VALUE_OUTPUT,
					  TEE_PARAM_TYPE_NONE,
					  TEE_PARAM_TYPE_NONE);
	if (exp_pt != ptypes)
		return TEE_ERROR_BAD_PARAMETERS;

	key_type = params[0].value.a;

	/* default key does not exist */
	params[1].value.a = 0;

	if (jetson_user_key_pta_get_user_key(&key, &key_len, key_type))
		goto out;

	if ((key_len != TEGRA_SE_KEY_128_SIZE)
			&& (key_len != TEGRA_SE_KEY_256_SIZE))
		goto out;

	/* key is exists */
	params[1].value.a = 1;

out:
	return TEE_SUCCESS;
}

static TEE_Result jetson_pta_decrypt_cpubl_payload(uint32_t ptypes,
				      TEE_Param params[TEE_NUM_PARAMS])
{
	TEE_Result rc = TEE_SUCCESS;
	uint32_t key_type = EKB_USER_KEY_KERNEL_ENCRYPTION;
	uint8_t *key;
	uint32_t key_len;
	static struct dec_ctx ctx = {0};
	uint32_t exp_pt = TEE_PARAM_TYPES(TEE_PARAM_TYPE_MEMREF_INOUT,
					  TEE_PARAM_TYPE_VALUE_INPUT,
					  TEE_PARAM_TYPE_NONE,
					  TEE_PARAM_TYPE_NONE);
	if ((exp_pt != ptypes) ||
	    (params[0].memref.buffer == NULL) ||
	    (params[0].memref.size == 0))
		return TEE_ERROR_BAD_PARAMETERS;

	switch (params[1].value.a) {
	case JETSON_CPUBL_PAYLOAD_DECRYPTION_INIT:
		if (ctx.dec_phase != JETSON_CPUBL_PAYLOAD_DECRYPTION_PRE_INIT) {
			rc = TEE_ERROR_BAD_STATE;
			goto error;
		}

		rc = jetson_user_key_pta_get_user_key(&key, &key_len, key_type);
		if (rc)
			goto error;

		rc = jetson_decrypt_cpubl_payload_process_params(&ctx,
				key, key_len,
				(void *)params[0].memref.buffer, &params[0].memref.size);
		if (rc)
			goto error;

		rc = jetson_decrypt_cpubl_payload_init(&ctx);
		if (rc)
			goto error;

		ctx.dec_phase = JETSON_CPUBL_PAYLOAD_DECRYPTION_INIT;

		break;

	case JETSON_CPUBL_PAYLOAD_DECRYPTION_UPDATE:
		if (ctx.dec_phase != JETSON_CPUBL_PAYLOAD_DECRYPTION_INIT
				&& ctx.dec_phase != JETSON_CPUBL_PAYLOAD_DECRYPTION_UPDATE) {
			rc = TEE_ERROR_BAD_STATE;
			goto error;
		}

		rc = jetson_decrypt_cpubl_payload_update(&ctx,
				(void *)params[0].memref.buffer, params[0].memref.size,
				(void *)params[0].memref.buffer, &params[0].memref.size);
		if (rc)
			goto error;

		ctx.dec_phase = JETSON_CPUBL_PAYLOAD_DECRYPTION_UPDATE;

		break;

	case JETSON_CPUBL_PAYLOAD_DECRYPTION_FINAL:
		if (ctx.dec_phase != JETSON_CPUBL_PAYLOAD_DECRYPTION_INIT
				&& ctx.dec_phase != JETSON_CPUBL_PAYLOAD_DECRYPTION_UPDATE) {
			rc = TEE_ERROR_BAD_STATE;
			goto error;
		}

		rc = jetson_decrypt_cpubl_payload_final(&ctx,
				(void *)params[0].memref.buffer, params[0].memref.size,
				(void *)params[0].memref.buffer, &params[0].memref.size);

		memset(&ctx, 0, sizeof(struct dec_ctx));

		break;

	default:
		rc = TEE_ERROR_BAD_PARAMETERS;
	}

	return rc;

error:
	memset(&ctx, 0, sizeof(struct dec_ctx));
	return rc;
}

static TEE_Result invoke_command(void *psess __unused,
				 uint32_t cmd, uint32_t ptypes,
				 TEE_Param params[TEE_NUM_PARAMS])
{
	if (!(jetson_service_flag & PTA_SRV_FLAG))
		return TEE_ERROR_ACCESS_DENIED;

	switch(cmd) {
	case JETSON_USER_KEY_CMD_GET_EKB_KEY:
		return jetson_pta_get_ekb_key(ptypes, params);
	case JETSON_USER_KEY_CMD_GET_RANDOM:
		return jetson_pta_get_random(ptypes, params);
	case JETSON_USER_KEY_CMD_GEN_UNIQUE_KEY_BY_EKB:
		return jetson_pta_gen_unique_key_by_ekb(ptypes, params);
	case JETSON_USER_KEY_CMD_GEN_KEY:
		return jetson_pta_gen_key(ptypes, params);
	case JETSON_USER_KEY_CMD_GET_FLAG:
		return jetson_pta_get_flag(ptypes, params);
	case JETSON_USER_KEY_CMD_SET_FLAG:
		return jetson_pta_set_flag(ptypes, params);
	case JETSON_USER_KEY_CMD_IS_KEY_EXISTS:
		return jetson_pta_is_user_key_exists(ptypes, params);
	case JETSON_USER_KEY_CMD_DECRYPT_CPUBL_PAYLOAD:
		return jetson_pta_decrypt_cpubl_payload(ptypes, params);
	default:
		return TEE_ERROR_NOT_SUPPORTED;
	}
}

static TEE_Result open_session(uint32_t param_types __unused,
			       TEE_Param params[TEE_NUM_PARAMS] __unused,
			       void **sess_ctx __unused)
{
	struct ts_session *s = NULL;

	/* Check that we're called from a user TA */
	s = ts_get_calling_session();
	if (!s || !is_user_ta_ctx(s->ctx))
		return TEE_ERROR_ACCESS_DENIED;

	return TEE_SUCCESS;
}

pseudo_ta_register(.uuid = JETSON_USER_KEY_UUID,
		   .name = JETSON_USER_KEY_PTA_NAME,
		   .flags = PTA_DEFAULT_FLAGS,
		   .open_session_entry_point = open_session,
		   .invoke_command_entry_point = invoke_command);

static const char *const jetson_ekb_dt_match_table[] = {
	"jetson-ekb-blob",
};

static TEE_Result get_oem_key_valid(bool *valid)
{
	TEE_Result rc = TEE_SUCCESS;
	uint32_t bsi = 0;

	if (!valid)
		return TEE_ERROR_BAD_PARAMETERS;

	rc = tegra_fuse_get_bsi(&bsi);
	if (rc) {
		EMSG("Read fuse BSI failed: 0x%x", rc);
		return rc;
	}

	if ((bsi & BSI_OEM_KEY_VALID_MASK) == 0)
		*valid = false;
	else
		*valid = true;

	return rc;
}

static TEE_Result get_sec_mode_enabled(bool *enabled)
{
	TEE_Result rc = TEE_SUCCESS;
	uint32_t sec_mode = 0;

	if (!enabled)
		return TEE_ERROR_BAD_PARAMETERS;

	rc = tegra_fuse_get_sec_mode(&sec_mode);
	if (rc) {
		EMSG("Read fuse security mode failed: 0x%x", rc);
		return rc;
	}

	if ((sec_mode & SECURITY_MODE_ENABLED_MASK) == 0)
		*enabled = false;
	else
		*enabled = true;

	return rc;
}

TEE_Result jetson_user_key_pta_uefi_vars_auth(uint8_t *vars, uint32_t vars_len)
{
	TEE_Result rc = TEE_SUCCESS;
	uint32_t key_type = EKB_USER_KEY_UEFI_VAR_AUTH;
	uint8_t *key;
	uint32_t key_len = 16;
	void *cmac_ctx = NULL;
	bool key_valid;
	bool sec_enabled;
	/* test UEFI variable auth key: d9f7b49e3b6264985f1326f541bb43c9 */
	uint8_t uefi_vars_auth_key_test[] = {
		0xd9, 0xf7, 0xb4, 0x9e, 0x3b, 0x62, 0x64, 0x98,
		0x5f, 0x13, 0x26, 0xf5, 0x41, 0xbb, 0x43, 0xc9
	};

	if (!vars || vars_len < TEE_AES_BLOCK_SIZE)
		return TEE_ERROR_BAD_PARAMETERS;

	/* Get UEFI variables authentication key */
	rc = jetson_user_key_pta_get_user_key(&key, &key_len, key_type);
	if (rc) {
		if ((get_oem_key_valid(&key_valid) != TEE_SUCCESS) ||
				(get_sec_mode_enabled(&sec_enabled) != TEE_SUCCESS)) {
			EMSG("Failed to get oem key valid or get security mode !");
			goto error;
		} else {
			if (key_valid || sec_enabled) {
				EMSG("UEFI variable auth key not set !");
				goto error;
			} else {
				IMSG("WARNING: Test UEFI variable auth "
						"key is being used !");
				IMSG("WARNING: UEFI variable protection "
						"is not fully enabled !");
				key = uefi_vars_auth_key_test;
			}
		}
	}

	/* Calculate AES-CMAC for uefi variables */
	rc = crypto_mac_alloc_ctx(&cmac_ctx, TEE_ALG_AES_CMAC);
	if (rc)
		goto error;

	rc = crypto_mac_init(cmac_ctx, key, key_len);
	if (rc)
		goto error;

	rc = crypto_mac_update(cmac_ctx, vars, vars_len);
	if (rc)
		goto error;

	rc = crypto_mac_final(cmac_ctx, vars, TEE_AES_BLOCK_SIZE);
	if (rc)
		goto error;

error:
	if (cmac_ctx)
		crypto_mac_free_ctx(cmac_ctx);

	return rc;
}

TEE_Result jetson_user_key_pta_query_ftpm_prop(uint32_t ftpm_ekb_id,
					       uint8_t **data,
					       uint32_t *len)
{
	TEE_Result rc = TEE_SUCCESS;
	struct user_key *ekb_key = NULL;

	if ((ftpm_ekb_id < EKB_FTPM_SN) || (ftpm_ekb_id >= EKB_FTPM_ID_MAX))
		return TEE_ERROR_ITEM_NOT_FOUND;

	rc = jetson_pta_get_key_by_type(&ekb_key, ftpm_ekb_id);
	if (rc == TEE_SUCCESS) {
		*data = ekb_key->data;
		*len = ekb_key->len;
	}

	return rc;
}

static TEE_Result jetson_ekb_dt_init(void)
{
	TEE_Result rc = TEE_SUCCESS;
	void *fdt = NULL;
	void *reg_addr = NULL;
	void *tmp_addr = NULL;
	int node = -1;
	uint32_t i = 0;
	paddr_t pbase;
	ssize_t sz;

	fdt = get_dt();
	if (!fdt) {
		EMSG("%s: DTB is not present.", __func__);
		rc = TEE_ERROR_ITEM_NOT_FOUND;
		goto fail;
	}

	for (i = 0; i < ARRAY_SIZE(jetson_ekb_dt_match_table); i++) {
		node = fdt_node_offset_by_compatible(fdt, 0,
					jetson_ekb_dt_match_table[i]);
		if (node >= 0)
			break;
	}

	if (node < 0) {
		EMSG("%s: DT not found (%x).", __func__, node);
		rc = TEE_ERROR_ITEM_NOT_FOUND;
		goto fail;
	}

	pbase = fdt_reg_base_address(fdt, node);
	if (pbase == DT_INFO_INVALID_REG) {
		rc = TEE_ERROR_GENERIC;
		goto fail;
	}

	sz = fdt_reg_size(fdt, node);
	if (sz < 0) {
		rc = TEE_ERROR_GENERIC;
		goto fail;
	}

	if (!core_mmu_add_mapping(EKB_MEMORY_AREA_TYPE, pbase, sz)) {
		EMSG("Failed to map %zu bytes at PA 0x%"PRIxPA,
		     (size_t)sz, pbase);
		rc = TEE_ERROR_GENERIC;
		goto fail;
	}

	reg_addr = (void*)phys_to_virt(pbase, EKB_MEMORY_AREA_TYPE, sz);
	if (!reg_addr) {
		EMSG("Failed to get VA for PA 0x%"PRIxPA, pbase);
		rc = TEE_ERROR_GENERIC;
		goto fail;
	}

	tmp_addr = calloc(1, sz);
	if (!tmp_addr) {
		rc = TEE_ERROR_OUT_OF_MEMORY;
		goto fail;
	}
	memcpy(tmp_addr, reg_addr, sz);
	ekb_base_addr = (vaddr_t)tmp_addr;

	/* Clear the source. */
	memset((void*)reg_addr, 0, sz);

fail:
	if (NULL != reg_addr)
		core_mmu_remove_mapping(EKB_MEMORY_AREA_TYPE, reg_addr, sz);

	return rc;
}

static TEE_Result jetson_user_key_pta_init(void)
{
	TEE_Result rc = TEE_SUCCESS;

	if (!IS_ENABLED(CFG_DT)) {
		rc = TEE_ERROR_NOT_SUPPORTED;
		goto fail;
	}

	rc = jetson_ekb_dt_init();
	if (rc) {
		EMSG("EKB dt init failed with rc = 0x%08x.", rc);
		goto fail;
	}

	rc = jetson_ekb_extraction(ekb_base_addr);
	if (rc) {
		EMSG("EKB extraction failed with rc = 0x%08x.", rc);
		goto fail;
	}

	return rc;

fail:
	jetson_service_flag &= !PTA_SRV_FLAG;
	EMSG("%s: Failed (%x).", __func__, rc);

#if defined(CFG_JETSON_MANDATORY_EKB)
	panic();
#endif

	return rc;
}

service_init_late(jetson_user_key_pta_init);
