// SPDX-License-Identifier: BSD-2-Clause
/*
 * SPDX-FileCopyrightText: Copyright (c) 2022-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_driver_aes.h>
#include <drivers/tegra/tegra_se_keyslot.h>

TEE_Result se_aes_encrypt(uint8_t *src, size_t src_len,
			  uint8_t *dst, size_t dst_len, se_aes_op_mode mode,
			  uint32_t keyslot, uint8_t *iv)
{
	status_t ret = NO_ERROR;

	if (src == NULL || src_len == 0 || dst == NULL || dst_len == 0)
		return TEE_ERROR_BAD_PARAMETERS;

	ret = tegra_se_ccc_aes_encrypt(src, src_len,
				       dst, dst_len,
				       mode, keyslot,
				       TEGRA_SE_KEY_256_SIZE, iv);
	if (ret != NO_ERROR) {
		EMSG("%s failed with: %d", __func__, ret);
		return TEE_ERROR_GENERIC;
	}

        return TEE_SUCCESS;
}

TEE_Result se_aes_ecb_kdf(uint8_t *derived_key, size_t derived_key_len,
			  uint8_t *iv, size_t iv_len,
			  uint32_t keyslot)
{
	if ((derived_key == NULL) || (iv == NULL)) {
		EMSG("%s: Invalid arguments.", __func__);
		return TEE_ERROR_BAD_PARAMETERS;
	}

	if ((derived_key_len != TEE_AES_BLOCK_SIZE) ||
	    (iv_len != TEE_AES_BLOCK_SIZE)) {
		EMSG("%s: Invalid size arguments.", __func__);
		return TEE_ERROR_BAD_PARAMETERS;
	}

	return se_aes_encrypt(iv, iv_len,
			      derived_key, derived_key_len,
			      SE_AES_OP_MODE_ECB, keyslot,
			      NULL);
}

TEE_Result se_nist_sp_800_108_cmac_kdf(se_aes_keyslot_t keyslot,
				       uint32_t key_len,
				       char const *context,
				       uint32_t context_len,
				       char const *label,
				       uint32_t label_len,
				       uint32_t dk_len,
				       uint8_t *out_dk)
{
	uint8_t counter[] = { 1 }, zero_byte[] = { 0 };
	uint32_t L[] = { TEE_U32_TO_BIG_ENDIAN(dk_len * 8) };
	uint8_t *message = NULL, *mptr;
	uint32_t msg_len, i, n;
	status_t ret = NO_ERROR;
	TEE_Result rc = TEE_SUCCESS;

	if ((key_len != TEGRA_SE_KEY_128_SIZE) &&
	    (key_len != TEGRA_SE_KEY_256_SIZE))
		return TEE_ERROR_BAD_PARAMETERS;

	if ((dk_len % TEE_AES_BLOCK_SIZE) != 0)
		return TEE_ERROR_BAD_PARAMETERS;

	if (!context || !label || !out_dk)
		return TEE_ERROR_BAD_PARAMETERS;

	/*
	 *  Regarding to NIST-SP 800-108
	 *  message = counter || label || 0 || context || L
	 *
	 *  A || B = The concatenation of binary strings A and B.
	 */
	msg_len = context_len + label_len + 2 + sizeof(L);
	message = calloc(1, msg_len);
	if (message == NULL)
		return TEE_ERROR_OUT_OF_MEMORY;

	/* Concatenate the messages */
	mptr = message;
	memcpy(mptr, counter, sizeof(counter));
	mptr += sizeof(counter);
	memcpy(mptr, label, label_len);
	mptr += label_len;
	memcpy(mptr, zero_byte, sizeof(zero_byte));
	mptr += sizeof(zero_byte);
	memcpy(mptr, context, context_len);
	mptr += context_len;
	memcpy(mptr, L, sizeof(L));

	/* SE AES-CMAC */
	/* n: iterations of the PRF count */
	n = dk_len / TEE_AES_BLOCK_SIZE;
	for (i = 0; i < n; i++) {
		/* Update the counter */
		message[0] = i + 1;

		ret = tegra_se_ccc_aes_encrypt(message, msg_len,
					       (out_dk + (i * TEE_AES_BLOCK_SIZE)),
					       TEE_AES_BLOCK_SIZE,
					       SE_AES_OP_MODE_CMAC, keyslot,
					       key_len, NULL);
		if (ret != NO_ERROR) {
			EMSG("CMAC-AES failed with: %d", ret);
			rc = TEE_ERROR_GENERIC;
			goto kdf_error;
		}
	}

kdf_error:
	free(message);
	return rc;
}

TEE_Result se_write_aes_keyslot(uint8_t *input_key, uint32_t keylen,
				uint32_t key_quad_sel,
				se_aes_keyslot_t keyslot)
{
	status_t ret = NO_ERROR;

	if ((keylen != TEGRA_SE_KEY_128_SIZE) &&
	    (keylen != TEGRA_SE_KEY_256_SIZE)) {
		EMSG("%s: Invalid key size.", __func__);
		return TEE_ERROR_BAD_PARAMETERS;
	}

	if ((key_quad_sel != TEGRA_SE_AES_QUAD_KEYS_128) &&
	    (key_quad_sel != TEGRA_SE_AES_QUAD_KEYS_256) &&
	    (key_quad_sel != TEGRA_SE_AES_QUAD_ORG_IV)) {
		EMSG("%s: Invalid key QUAD selection.", __func__);
		return TEE_ERROR_BAD_PARAMETERS;
	}

	ret = tegra_se_ccc_set_aes_keyslot(keyslot, input_key, keylen,
					   SE_AES_OP_MODE_CBC);
	if (ret != NO_ERROR) {
		EMSG("%s failed with: %d", __func__, ret);
		return TEE_ERROR_GENERIC;
	}

	return TEE_SUCCESS;
}

TEE_Result se_clear_aes_keyslots(void)
{
	status_t ret = NO_ERROR;
	uint32_t i;

	for (i = SE_AES_KEYSLOT_11; i < SE_AES_KEYSLOT_COUNT; i++) {
		if (i == SE_AES_KEYSLOT_11 + 2) {
			/* Skip keyslot #13 since we don't use it */
			continue;
		}

		ret = tegra_se_ccc_clear_aes_keyslot(i);
		if (ret != NO_ERROR) {
			EMSG("%s failed with: %d", __func__, ret);
			return TEE_ERROR_GENERIC;
		}
	}

	return TEE_SUCCESS;
}
