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

#include <stdio.h>
#include <string.h>
#include "utee_defines.h"
#include "crypto/crypto.h"
#include "drivers/tegra/tegra_se_keyslot.h"
#include "jetson_nist_sp_800_108_kdf.h"

/*
 * A software-based NIST-SP 800-108 KDF.
 * derives keys from a key in a key buffer.
 *
 * key		[in] input key for derivation.
 * key_len	[in] length in bytes of the input key.
 * context	[in] a pointer to a NIST-SP 800-108 context string.
 * context_len	[in] length in bytes of the contexct.
 * label	[in] a pointer to a NIST-SP 800-108 label string.
 * label_len	[in] length in bytes of the label.
 * dk_len	[in] length of the derived key in bytes;
 *		     may be 16 (128 bits) or any multiple of 16.
 * out_dk 	[out] a pointer to the derived key. The function stores
 *		      its result in this location.
 */
TEE_Result nist_sp_800_108_hmac_kdf(const uint8_t *key,
				    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)
{
	TEE_Result res = TEE_SUCCESS;
	uint8_t zero_byte[] = { 0U };
	uint8_t *p, tmp_buf[TEE_SHA256_HASH_SIZE] = { 0U };
	uint32_t tmp_len;
	uint32_t counter[] = { TEE_U32_TO_BIG_ENDIAN( 1 ) };
	uint32_t L[] = { TEE_U32_TO_BIG_ENDIAN(dk_len * 8) };
	uint8_t *message = NULL;
	uint8_t *mptr;
	uint32_t msg_len;
	void *hmac_ctx = NULL;

	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 (!key || !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 = sizeof(counter) + context_len + label_len + sizeof(L) +
		  sizeof(zero_byte);
	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));

	/* HMAC-SHA256 */
	res = crypto_mac_alloc_ctx(&hmac_ctx, TEE_ALG_HMAC_SHA256);
	if (res)
		goto kdf_error;

	tmp_len = dk_len;
	p = out_dk;

	while (tmp_len > 0) {
		crypto_mac_init(hmac_ctx, key, key_len);
		crypto_mac_update(hmac_ctx, message, msg_len);

		if (tmp_len >= TEE_SHA256_HASH_SIZE) {
			crypto_mac_final(hmac_ctx, p, TEE_SHA256_HASH_SIZE);
			tmp_len -= TEE_SHA256_HASH_SIZE;
			p += TEE_SHA256_HASH_SIZE;
		} else {
			crypto_mac_final(hmac_ctx, tmp_buf, TEE_SHA256_HASH_SIZE);
			memcpy(p, tmp_buf, tmp_len);
			break;
		}

		/* Update the counter for next use */
		message[0] += 1;
	}

	crypto_mac_free_ctx(hmac_ctx);

kdf_error:
	free(message);
	return res;
}

#if defined(PLATFORM_FLAVOR_t234)
/*
 * A software-based NIST-SP 800-108 KDF.
 *  derives keys from a key in a key buffer.
 *
 * key		[in] input key for derivation.
 * key_len	[in] length in bytes of the input key.
 * context	[in] a pointer to a NIST-SP 800-108 context string.
 * context_len	[in] length of the context string.
 * label	[in] a pointer to a NIST-SP 800-108 label string.
 * label_len	[in] length of the label string.
 * dk_len	[in] length of the derived key in bytes;
 *		     may be 16 (128 bits) or any multiple of 16.
 * out_dk 	[out] a pointer to the derived key. The function stores
 *		      its result in this location.
 */
TEE_Result nist_sp_800_108_cmac_kdf(uint8_t *key,
				    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;
	void *cmac_ctx = NULL;
	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 (!key || !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));

	/* AES-CMAC */
	rc = crypto_mac_alloc_ctx(&cmac_ctx, TEE_ALG_AES_CMAC);
	if (rc)
		goto kdf_error;

	/* 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;

		crypto_mac_init(cmac_ctx, key, key_len);
		crypto_mac_update(cmac_ctx, message, msg_len);
		crypto_mac_final(cmac_ctx, (out_dk + (i * TEE_AES_BLOCK_SIZE)),
				 TEE_AES_BLOCK_SIZE);
	}

	crypto_mac_free_ctx(cmac_ctx);

kdf_error:
	free(message);
	return rc;
}
#endif
