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

#include <stdio.h>
#include <string.h>
#include <trace.h>
#include <config.h>
#include <crypto/crypto.h>
#include <drivers/tegra/tegra_fuse.h>
#include <drivers/tegra/tegra_se_kdf.h>
#include <kernel/tee_common_otp.h>
#include <sys/queue.h>
#include "jetson_ekb_extraction.h"
#include "jetson_ekb_user_key_management.h"
#include "drivers/tegra/tegra_se_keyslot.h"
#include "jetson_nist_sp_800_108_kdf.h"

#define MAX_EKB_RK			2

/* Magic of EKB header */
#define EKB_MAGIC_STR			"NVEKBP"
/* Magic of EKB content header */
#define EKB_CONTENT_MAGIC_STR		"EEKB"
/* The minimum size of EKB */
#define EKB_MIN_SIZE			1024U

/*
 * The latest EKB layout: version 2
 *
 * -----------------------------------------------------------
 * |    EKB_size   |      Magic    |    Major   |   Minor    |
 * |    (4 bytes)  |    (8 bytes)  |  (2 bytes) | (2 bytes)  |
 * -----------------------------------------------------------
 * |                         FV (16 bytes)                   |
 * -----------------------------------------------------------
 * |                         MAC (16 bytes)                  |
 * -----------------------------------------------------------
 * |  Content_size | Content_magic |        Reserved         |
 * |   (4 bytes)   |   (4 bytes)   |        (8 bytes)        |
 * -----------------------------------------------------------
 * |                         IV (16 bytes)                   |
 * -----------------------------------------------------------
 * |    Key1_tag   |    Key1_len   |         Key1            |
 * |   (4 bytes)   |   (4 bytes)   |       (Key1_len)        |
 * -----------------------------------------------------------
 * |    Key2_tag   |   Key2_len    |         Key2            |
 * |   (4 bytes)   |   (4 bytes)   |       (Key2_len)        |
 * -----------------------------------------------------------
 * |                              ...                        |
 * -----------------------------------------------------------
 * |   End_tag1(\0\0\0\0)          |    End_tag2 (\0\0\0\0)  |
 * -----------------------------------------------------------
 */

/*
 * EKB version 1: FV is hardcoded in the EKB generation tool and in
 * the EKB handling here.
 * EKB version 2: FV is randomly generated by the EKB generation tool
 * and stored right after EKB header.
 */

/*
 * Fixed vector for EKB version 1.
 */
static uint8_t fv_for_ekb[] = {
	0xba, 0xd6, 0x6e, 0xb4, 0x48, 0x49, 0x83, 0x68,
	0x4b, 0x99, 0x2f, 0xe5, 0x4a, 0x64, 0x8b, 0xb8,
};

/*
 * Random fixed vector used to derive SSK_DK (Derived Key).
 *
 * e4 20 f5 8d 1d ea b5 24 c2 70 d8 d2 3e ca 45 e8
 */
static uint8_t fv_for_ssk_dk[] = {
	0xe4, 0x20, 0xf5, 0x8d, 0x1d, 0xea, 0xb5, 0x24,
	0xc2, 0x70, 0xd8, 0xd2, 0x3e, 0xca, 0x45, 0xe8,
};

/*
 * HW unique key for tee_otp_get_hw_unique_key
 */
static struct tee_hw_unique_key *hw_unique_key = NULL;
/*
 * Fixed vector used to derive HUK_RK (Root Key).
 *
 * 89 49 b3 e0 65 eb 58 85 16 39 22 24 66 e1 48 65
 */
static uint8_t fv_for_huk_rk[] = {
	0x89, 0x49, 0xb3, 0xe0, 0x65, 0xeb, 0x58, 0x85,
	0x16, 0x39, 0x22, 0x24, 0x66, 0xe1, 0x48, 0x65,
};

static uint8_t huk_rk[TEGRA_SE_KEY_128_SIZE] = { 0 };
static bool gen_huk_rk_fail = true;

/*
 * Root keys derived from SE keyslots.
 */
static uint8_t ekb_rks[MAX_EKB_RK][TEGRA_SE_KEY_128_SIZE];
static uint8_t ekb_rks_num = 0;
static uint8_t *ekb_rk_verified = NULL;
static uint8_t ssk_rk[TEGRA_SE_KEY_128_SIZE] = { 0 };

/*
 * Derived keys from NIST-SP 800-108.
 */
static uint8_t ekb_ek[TEGRA_SE_KEY_128_SIZE] = { 0 };
static uint8_t ekb_ak[TEGRA_SE_KEY_128_SIZE] = { 0 };

struct ekb_header {
	uint32_t ekb_size;
	char ekb_magic[8];
	uint16_t ver_maj;
	uint16_t ver_min;
};

struct ekb_content_header {
	uint32_t enc_data_len;
	char ekb_content_magic[4];
	uint8_t reserved[8];
};

struct ekb_content {
	uint8_t ekb_cmac[TEGRA_SE_AES_BLOCK_SIZE];
	struct ekb_content_header content_header;
	uint8_t random_iv[TEGRA_SE_AES_IV_SIZE];
	uint8_t ekb_ciphertext[0];
};

static TEE_Result get_ekb_content_start_addr(vaddr_t ekb_addr,
					     vaddr_t *ekb_content_addr)
{
	struct ekb_header *header;

	if (ekb_addr == 0 || !ekb_content_addr)
		return TEE_ERROR_BAD_PARAMETERS;

	header = (struct ekb_header *)ekb_addr;

	if (header->ekb_size < EKB_MIN_SIZE)
		return TEE_ERROR_BAD_FORMAT;

	if (strncmp((char *)EKB_MAGIC_STR, header->ekb_magic, strlen(EKB_MAGIC_STR)))
		return TEE_ERROR_BAD_FORMAT;

	if (header->ver_maj == 1) {
		*ekb_content_addr = ekb_addr + sizeof(struct ekb_header);
	} else if (header->ver_maj == 2) {
		memcpy(fv_for_ekb, header + 1, sizeof(fv_for_ekb));
		*ekb_content_addr = ekb_addr + sizeof(struct ekb_header)
				    + sizeof(fv_for_ekb);
	} else {
		return TEE_ERROR_BAD_FORMAT;
	}

	return TEE_SUCCESS;
}

static TEE_Result ekb_extraction_process(vaddr_t ekb_content_addr)
{
	TEE_Result rc = TEE_SUCCESS;
	struct ekb_content *ekb;
	struct ekb_content_header *content_header;
	void *decrypted_ekb_content;
	void *crypt_ctx = NULL;
	uint8_t ekb_cmac_verify[TEGRA_SE_AES_BLOCK_SIZE] = { 0 };
	int i = 0;

	if (ekb_content_addr == 0)
		return TEE_ERROR_BAD_PARAMETERS;

	ekb = (struct ekb_content *)ekb_content_addr;
	content_header = &ekb->content_header;
	if (strncmp((char *)EKB_CONTENT_MAGIC_STR,
		    content_header->ekb_content_magic,
		    strlen(EKB_CONTENT_MAGIC_STR)))
		return TEE_ERROR_BAD_PARAMETERS;


	for (i = 0; i < ekb_rks_num; i++) {
		crypt_ctx = NULL;

		/* Derive EKB_EK */
		rc = nist_sp_800_108_cmac_kdf(ekb_rks[i], TEGRA_SE_KEY_128_SIZE,
					      "ekb", strlen("ekb"),
					      "encryption", strlen("encryption"),
					      TEGRA_SE_KEY_128_SIZE, ekb_ek);
		if (rc) {
			DMSG("Derive EKB_EK from #%d EKB_RK failed: 0x%x.", i, rc);
			continue;
		}

		/* Derive EKB_AK */
		rc = nist_sp_800_108_cmac_kdf(ekb_rks[i], TEGRA_SE_KEY_128_SIZE,
					      "ekb", strlen("ekb"),
					      "authentication", strlen("authentication"),
					      TEGRA_SE_KEY_128_SIZE, ekb_ak);
		if (rc) {
			DMSG("Derive EKB_AK from #%d EKB_RK failed: 0x%x.", i, rc);
			continue;
		}

		/* Calculate the EKB CMAC */
		rc = crypto_mac_alloc_ctx(&crypt_ctx, TEE_ALG_AES_CMAC);
		if (rc)
			continue;

		crypto_mac_init(crypt_ctx, ekb_ak, sizeof(ekb_ak));
		crypto_mac_update(crypt_ctx, (uint8_t *)content_header,
				  sizeof(struct ekb_content_header) + sizeof(ekb->random_iv)
				  + content_header->enc_data_len);
		crypto_mac_final(crypt_ctx, ekb_cmac_verify, TEGRA_SE_AES_BLOCK_SIZE);
		crypto_mac_free_ctx(crypt_ctx);

		/* Verify the CMAC */
		if (memcmp(ekb->ekb_cmac, ekb_cmac_verify, TEGRA_SE_AES_BLOCK_SIZE)) {
			DMSG("EKB signature mismatch when using #%d EKB_RK.", i);
			continue;
		}

		/* Decrypt the EKB ciphertext */
		rc = crypto_cipher_alloc_ctx(&crypt_ctx, TEE_ALG_AES_CBC_NOPAD);
		if (rc)
			continue;

		decrypted_ekb_content = malloc(content_header->enc_data_len);
		if (!decrypted_ekb_content) {
			crypto_cipher_free_ctx(crypt_ctx);
			continue;
		}

		crypto_cipher_init(crypt_ctx, TEE_MODE_DECRYPT,
				   ekb_ek, TEGRA_SE_KEY_128_SIZE,
				   NULL, 0,
				   ekb->random_iv, TEGRA_SE_AES_BLOCK_SIZE);
		crypto_cipher_update(crypt_ctx, TEE_MODE_DECRYPT, false,
				     ekb->ekb_ciphertext,
				     content_header->enc_data_len,
				     (uint8_t*)decrypted_ekb_content);
		crypto_cipher_final(crypt_ctx);
		crypto_cipher_free_ctx(crypt_ctx);


		rc = jetson_ekb_extraction_parse_key(decrypted_ekb_content);
		if (rc)
			EMSG("Parse ekb content failed with rc = 0x%08x", rc);

		memset(ekb_cmac_verify, 0, sizeof(ekb_cmac_verify));

		/* Zeroing out the sensitive data in memory */
		memset(ekb_ek, 0, sizeof(ekb_ek));
		memset(ekb_ak, 0, sizeof(ekb_ak));
		memset(decrypted_ekb_content, 0, content_header->enc_data_len);
		free(decrypted_ekb_content);

		if (rc == TEE_SUCCESS) {
			DMSG("Verified #%d EKB_RK works. Skip other EKB_RK keys.", i);
			ekb_rk_verified = ekb_rks[i];
			break;
		}
	}

	if (i == ekb_rks_num) {
		EMSG("Tried all EKB_RKs but still can't extract the EKB image.");
		return TEE_ERROR_SECURITY;
	}

	return TEE_SUCCESS;
}

static TEE_Result hwkey_derivation_process(void)
{
	TEE_Result rc = TEE_SUCCESS;
	size_t i = 0;
	se_aes_keyslot_t fuse_key_for_ekb[] = { SE_AES_KEYSLOT_OEM_K1, SE_AES_KEYSLOT_OEM_K2 };
	se_aes_keyslot_t fuse_key_for_ssk = SE_AES_KEYSLOT_OEM_K2;
	se_aes_keyslot_t fuse_key_for_huk_rk = SE_AES_KEYSLOT_OEM_K1;
	static_assert(ARRAY_SIZE(fuse_key_for_ekb) <= MAX_EKB_RK);

	memset(ekb_rks, 0, sizeof(ekb_rks));
	for (i = 0; i < ARRAY_SIZE(fuse_key_for_ekb); i++) {
		/*
		 * Derive root keys by performing AES-ECB encryption with the fixed
		 * vector (fv) and keys in SE keyslots.
		 */
		rc = tegra_se_aes_ecb_kdf(ekb_rks[ekb_rks_num], sizeof(ekb_rks[ekb_rks_num]),
					  fv_for_ekb, sizeof(fv_for_ekb),
					  fuse_key_for_ekb[i]);
		if (rc == TEE_SUCCESS) {
			DMSG("Derived EKB_RK from #%lu fuse key succeeded", i);
			ekb_rks_num++;
		} else {
			IMSG("Derive EKB_RK from #%lu fuse key failed with rc = 0x%08x, ignored.", i, rc);
			rc = TEE_SUCCESS;
		}
	}

	if (ekb_rks_num == 0) {
		EMSG("%s: Failed to derive EKB root key from all fuse keys(%x)", __func__, rc);
		goto key_derivation_fail;
	}

	rc = tegra_se_aes_ecb_kdf(ssk_rk, sizeof(ssk_rk),
				  fv_for_ssk_dk, sizeof(fv_for_ssk_dk),
				  fuse_key_for_ssk);
	if (rc != TEE_SUCCESS) {
		EMSG("%s: Failed to derive SSK root key (%x)", __func__, rc);
		/*
		 * Failing to derive SSK is not a critical issue, to not halt OP-TEE boots,
		 * reset the rc to SUCCESS so that OP-TEE can keep booting.
		 */
		rc = TEE_SUCCESS;
	}

	rc = tegra_se_aes_ecb_kdf(huk_rk, sizeof(huk_rk),
				  fv_for_huk_rk, sizeof(fv_for_huk_rk),
				  fuse_key_for_huk_rk);
	if (rc != TEE_SUCCESS) {
		EMSG("%s: Failed to derive HUK root key (%x)", __func__, rc);
		/*
		 * Failing to derive HUK root key is not a critical issue, to not halt
		 * OP-TEE boots, reset the rc to SUCCESS so that OP-TEE can keep booting.
		 */
		rc = TEE_SUCCESS;
	} else {
		gen_huk_rk_fail = false;
	}

key_derivation_fail:
	return rc;
}

TEE_Result jetson_ekb_extraction(vaddr_t ekb_base_addr)
{
	TEE_Result rc = TEE_SUCCESS;
	vaddr_t ekb_content_addr;

	/* Get the start address of ekb content. */
	rc = get_ekb_content_start_addr(ekb_base_addr, &ekb_content_addr);
	if (rc) {
		EMSG("Get ekb content start address failed with rc = 0x%08x.", rc);
		goto fail;
	}

	/* Derive keys from SE keyslots. */
	rc = hwkey_derivation_process();
	if (rc) {

		EMSG("Derive key from hwkey failed with rc = 0x%08x.", rc);
		goto fail;
	}

	/* EKB extraction */
	rc = ekb_extraction_process(ekb_content_addr);
	if (rc)
		EMSG("EKB extract failed with rc = 0x%08x.", rc);

fail:
	return rc;
}

/* Override the weak tee_otp_get_hw_unique_key */
TEE_Result tee_otp_get_hw_unique_key(struct tee_hw_unique_key *hwkey)
{
	TEE_Result ret = TEE_SUCCESS;
	fuse_ecid_t *ecid;
	char ecid_str[33];

	if (hwkey == NULL)
		return TEE_ERROR_BAD_PARAMETERS;
	if (gen_huk_rk_fail)
		return TEE_ERROR_NO_DATA;

	if (hw_unique_key == NULL) {
		hw_unique_key = calloc(1, sizeof(struct tee_hw_unique_key));
		if (hw_unique_key == NULL) {
			ret = TEE_ERROR_OUT_OF_MEMORY;
			goto out;
		}
		ecid = tegra_fuse_get_ecid();
		if (ecid == NULL) {
			ret = TEE_ERROR_NO_DATA;
			goto out;
		}

		memset(ecid_str, 0, sizeof(ecid_str));
		snprintf(ecid_str, sizeof(ecid_str), "%08x%08x%08x%08x",
			 ecid->ecid[3], ecid->ecid[2], ecid->ecid[1],
			 ecid->ecid[0]);
		ret = nist_sp_800_108_cmac_kdf(huk_rk,
					       sizeof(huk_rk),
					       ecid_str,
					       strlen(ecid_str),
					       "tee-hw-unique-key",
					       strlen("tee-hw-unique-key"),
					       HW_UNIQUE_KEY_LENGTH,
					       hw_unique_key->data);
	}

out:
	if (ret == TEE_SUCCESS) {
		memcpy(hwkey->data, hw_unique_key->data, HW_UNIQUE_KEY_LENGTH);
	} else {
		if (hw_unique_key != NULL) {
			free(hw_unique_key);
			hw_unique_key = NULL;
		}
	}

	return ret;
}
