// 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_aes.h>
#include <drivers/tegra/tegra_se_kdf.h>
#include "drivers/tegra/tegra_se_keyslot.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_ekb_extraction_t264.h"
#include "jetson_nist_sp_800_108_kdf.h"

/*
 * Layout of the EKB in version 2.1:
 *
 * -----------------------------------------------------------
 * |    EKB_size   |      Magic    |    Major   |   Minor    |
 * |    (4 bytes)  |    (8 bytes)  |  (2 bytes) | (2 bytes)  |
 * -----------------------------------------------------------
 * |                      Reserved (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.0: FV is randomly generated by the EKB generation tool
 * and stored right after EKB header.
 * EKB version 2.1: FV is not used anymore, rename it to reserved.
 *
 * NOTE: T264 currently only supports EKB version 2.1
 */

/* EKS Major and Minor Version Lower&Upper limit */
#define EKS_VERSION_MAJOR_LOWER_LIMIT		(2)
#define EKS_VERSION_MAJOR_UPPER_LIMIT		(2)
#define EKS_VERSION_MINOR_LOWER_LIMIT		(1)
#define EKS_VERSION_MINOR_UPPER_LIMIT		(1)

/* 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
 *
 * Since the minimum size of EKB is 1K, and the value of ekb_size
 * is not count in the size of ekb_size, so the minimum size of
 * ekb_size is 1k - sizeof(ekb_size) = 1020.
 * */
#define EKB_MIN_SIZE			1020U

static TEE_Result jetson_ekb_validate_header(struct ekb_header *header)
{
	if (!header)
		return TEE_ERROR_BAD_PARAMETERS;

	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 < EKS_VERSION_MAJOR_LOWER_LIMIT) ||
			(header->ver_maj > EKS_VERSION_MAJOR_UPPER_LIMIT) ||
			(header->ver_min < EKS_VERSION_MINOR_LOWER_LIMIT) ||
			(header->ver_min > EKS_VERSION_MINOR_UPPER_LIMIT))
		return TEE_ERROR_BAD_FORMAT;

	return TEE_SUCCESS;
}


static TEE_Result jetson_ekb_validate_content_header(struct ekb_content *ekb)
{
	struct ekb_content_header *content_header;

	if (!ekb)
		return TEE_ERROR_BAD_PARAMETERS;

	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;

	return TEE_SUCCESS;
}

static TEE_Result jetson_ekb_extraction_process(vaddr_t ekb_base_addr)
{
	TEE_Result rc = TEE_SUCCESS;
	struct ekb_header *header;
	struct ekb_content *ekb;
	struct ekb_content_header *content_header;
	void *decrypted_ekb_content = NULL;

	if (ekb_base_addr == 0)
		return TEE_ERROR_BAD_PARAMETERS;

	header = (struct ekb_header *)ekb_base_addr;

	/* Validate the EKB header */
	rc = jetson_ekb_validate_header(header);
	if (rc) {
		EMSG("Validate ekb header failed with rc = 0x%08x.", rc);
		goto fail;
	}

	ekb = (struct ekb_content *)(header + 1);

	/* Validate the EKB content header */
	rc = jetson_ekb_validate_content_header(ekb);
	if (rc) {
		EMSG("Validate ekb header failed with rc = 0x%08x.", rc);
		goto fail;
	}

	/* Authenticate the EKB CMAC */
	rc = jetson_ekb_authenticate_mac(ekb);
	if (rc) {
		EMSG("Authenticate ekb failed with rc = 0x%08x", rc);
		goto fail;
	}

	content_header = &ekb->content_header;

	/* Allocate buffer for decrypted EKB content */
	decrypted_ekb_content = malloc(content_header->enc_data_len);
	if (!decrypted_ekb_content) {
		rc = TEE_ERROR_OUT_OF_MEMORY;
		goto fail;
	}

	/* Decrypt the EKB content */
	rc = jetson_ekb_decrypt_content(ekb, (uint8_t *)decrypted_ekb_content);
	if (rc) {
		EMSG("Decrypt ekb content failed with rc = 0x%08x", rc);
		goto fail;
	}

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

fail:
	if (decrypted_ekb_content != NULL) {
		/* Zeroing out the sensitive data in memory */
		memset(decrypted_ekb_content, 0, content_header->enc_data_len);
		free(decrypted_ekb_content);
	}

	return rc;
}

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

	/* Derive keys*/
	rc = jetson_derive_keys();
	if (rc) {
		EMSG("Derive keys failed with rc = 0x%08x.", rc);
		return rc;
	}

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

fail:
	return rc;
}
