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

#include <string.h>
#include <stdlib.h>
#include <types_ext.h>
#include <tee_api_types.h>
#include <trace.h>
#include <initcall.h>
#include <utee_defines.h>
#include <tegra_se_ccc.h>
#include <tegra_se_ccc_aes.h>
#include <tegra_se_ccc_errno.h>
#include <tegra_se_ccc_kds.h>
#include <tegra_se_driver_aes.h>
#include <tegra_se_driver_rng.h>
#include <tegra_driver_srv_intf.h>
#include <drivers/tegra/tegra_se_keyslot.h>
#include <drivers/tegra/tegra_se_kdf.h>

/*
 * As the root key of TZ, if TZ RK is derived from a zero-value fuse key,
 * security will be reduced. We assume that the fuse key is zero, and use
 * a fixed KDF function (NIST-SP-800-108 KDF with HAMC-SHA256) with a
 * fixed label ("IS_ZERO_VALUE") and context (0x00) to get a fixed key as
 * below.
 *
 * "c0061316a18707584385ff1bed0fef98b1266caeaddcf6628ce4b012a2de5f40"
 *
 * This fixed key is used to check whether TZ RK is derived from zero-value
 * OEM_KDK1 fuse key.
 */
static const uint8_t kdk1_zero_derived_key[] = {
	0xc0, 0x06, 0x13, 0x16, 0xa1, 0x87, 0x07, 0x58,
	0x43, 0x85, 0xff, 0x1b, 0xed, 0x0f, 0xef, 0x98,
	0xb1, 0x26, 0x6c, 0xae, 0xad, 0xdc, 0xf6, 0x62,
	0x8c, 0xe4, 0xb0, 0x12, 0xa2, 0xde, 0x5f, 0x40
};

static bool is_oem_kdk1_zero_value(void)
{
	TEE_Result rc = TEE_SUCCESS;
	uint8_t context[1] = { 0 };
	uint8_t derived_test_key[TEGRA_SE_KEY_256_SIZE] = { 0 };

	rc = se_nist_sp_800_108_hmac_kdf_to_keyslot(SE_AES_KEYSLOT_TZ_RK,
						    (char *)"IS_ZERO_VALUE",
						    strlen("IS_ZERO_VALUE"),
						    (char *)context, 1,
						    KEYSIZE_256_TZ_HMAC_MANIFEST,
						    SE_AES_KEYSLOT_EKB_RK);
	if (rc != TEE_SUCCESS) {
		EMSG("%s: Failed to derive temporary test key (%x)", __func__, rc);
		goto out;
	}

	rc = se_nist_sp_800_108_hmac_kdf_to_memory(SE_AES_KEYSLOT_EKB_RK,
						   TEGRA_SE_KEY_256_SIZE,
						   (char *)"IS_ZERO_VALUE",
						   strlen("IS_ZERO_VALUE"),
						   (char *)context, 1,
						   TEGRA_SE_KEY_256_SIZE,
						   derived_test_key);
	if (rc != TEE_SUCCESS) {
		EMSG("%s: Failed to derive final test key (%x)", __func__, rc);
		goto out;
	}

out:
	if ((rc == TEE_SUCCESS) &&
	    (!memcmp(kdk1_zero_derived_key, derived_test_key, TEGRA_SE_KEY_256_SIZE)))
			return true;

	return false;
}

static TEE_Result tegra_se_init(void)
{
	status_t ret;
	tegra_drv_srv_intf_t *drv_srv_intf;

	TEE_Result err = TEE_SUCCESS;
	uint32_t bsi = 0;
	uint32_t sec_mode = 0;

	/* TZ RK: b711d9a0d22fe5a8f38a63e2498e1617b2bc0b918fc9433f000a626fc091fb63 */
	const uint8_t test_tz_rk[] = {
		0xb7, 0x11, 0xd9, 0xa0, 0xd2, 0x2f, 0xe5, 0xa8,
		0xf3, 0x8a, 0x63, 0xe2, 0x49, 0x8e, 0x16, 0x17,
		0xb2, 0xbc, 0x0b, 0x91, 0x8f, 0xc9, 0x43, 0x3f,
		0x00, 0x0a, 0x62, 0x6f, 0xc0, 0x91, 0xfb, 0x63
	};

#ifndef TEGRA_SE_CRYPTO_INITIALIZED
#define TEGRA_SE_CRYPTO_INITIALIZED
	/* This function will call tegra_se_map_device to map devices */
	ret = tegra_init_crypto_devices();
	if (ret != NO_ERROR) {
		return TEE_ERROR_GENERIC;
	}
#endif

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

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

	if ((bsi & BSI_OEM_KEY_VALID_MASK) || (sec_mode & SECURITY_MODE_ENABLED_MASK)) {
		ret = tegra_se_ccc_move_key_from_kds_to_keyslot(TZ_KDS_SRC_KEYID,
								TZ_KDS_TABLE_ID,
								TZ_KDS_REGION_ID,
								SE_AES_KEYSLOT_TZ_RK);
		if (ret != NO_ERROR) {
			EMSG("Move tz root key from kds to keyslot failed: %d", ret);
			return TEE_ERROR_GENERIC;
		}

		if (is_oem_kdk1_zero_value())
			IMSG("The OEM KDK1 is zero-value key. This is insecure for shipping products!");
	} else {
		ret = tegra_se_ccc_set_test_tz_rk(test_tz_rk);
		if (ret != NO_ERROR) {
			EMSG("Setting test TZ root key failed: %d", ret);
			return TEE_ERROR_GENERIC;
		}
		IMSG("Test TZ root key is being used. This is insecure for shipping products!");
	}

	drv_srv_intf = tegra_drv_srv_intf_get();
	if (drv_srv_intf != NULL) {
		drv_srv_intf->aes_encrypt = se_aes_encrypt;
		drv_srv_intf->aes_cmac= se_aes_cmac;
		drv_srv_intf->nist_sp_800_108_hmac_kdf_to_mem = se_nist_sp_800_108_hmac_kdf_to_memory;
		drv_srv_intf->nist_sp_800_108_hmac_kdf_to_ks = se_nist_sp_800_108_hmac_kdf_to_keyslot;
		drv_srv_intf->clear_aes_keyslots = se_clear_aes_keyslots;
#if defined(CFG_TEGRA_SE_RNG1)
		drv_srv_intf->rng_get_random = se_rng1_get_random;
#endif
	} else {
		return TEE_ERROR_NOT_SUPPORTED;
	}

	return TEE_SUCCESS;
}
service_init_late(tegra_se_init);
