// SPDX-License-Identifier: BSD-2-Clause
/*
 * SPDX-FileCopyrightText: Copyright (c) 2023-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_decrypt_cpubl_payload.h"
#include "jetson_nist_sp_800_108_kdf.h"

static TEE_Result get_one_and_zeros_padding(uint8_t *input,
					    uint32_t input_len,
					    size_t *data_len)
{
	uint32_t i;
	unsigned char done = 0, prev_done, bad;

	if (NULL == input || NULL == data_len)
		return TEE_ERROR_BAD_PARAMETERS;

	if ((input_len % TEGRA_SE_AES_BLOCK_SIZE) != 0)
		return TEE_ERROR_BAD_PARAMETERS;

	bad = 0x80;
	*data_len = 0;
	for (i = input_len; i > 0; i--) {
		prev_done = done;
		done |= (input[i - 1] != 0);
		*data_len |= (i - 1) * (done != prev_done);
		bad ^= input[i - 1] * (done != prev_done);
	}

	if (bad != 0)
		*data_len = input_len;

	return TEE_SUCCESS;
}

TEE_Result jetson_decrypt_cpubl_payload_process_params(struct dec_ctx *dec_ctx,
						       uint8_t *key,
						       uint32_t key_len,
						       void *img,
						       size_t *img_size)
{
	TEE_Result res;
	uint8_t *bch;

	if (!dec_ctx)
		return TEE_ERROR_BAD_PARAMETERS;

	if (!key || (key_len == 0))
		return TEE_ERROR_BAD_PARAMETERS;

	if (!img || !img_size)
		return TEE_ERROR_BAD_PARAMETERS;

	if (*img_size < BCH_HEADER_SIZE)
		return TEE_ERROR_BAD_PARAMETERS;

	bch = img;

	if (memcmp(bch + BCH_HEADER_MAGIC_OFFSET, BCH_HEADER_MAGIC,
		   BCH_HEADER_MAGIC_SIZE))
		return TEE_ERROR_BAD_PARAMETERS;

	memcpy(&(dec_ctx->payload_size),
	       bch + BCH_HEADER_BINARY_LEN_OFFSET,
	       BCH_HEADER_BINARY_LEN_SIZE);
	memcpy(dec_ctx->aad_in_hdr, bch + BCH_HEADER_AAD_OFFSET,
	       BCH_HEADER_AAD_SIZE);
	memcpy(dec_ctx->iv_in_hdr, bch + BCH_HEADER_IV_OFFSET,
	       BCH_HEADER_IV_SIZE);
	memcpy(dec_ctx->tag_in_hdr, bch + BCH_HEADER_TAG_OFFSET,
	       BCH_HEADER_TAG_SIZE);
	dec_ctx->key_len = TEGRA_SE_KEY_256_SIZE;

	res = nist_sp_800_108_hmac_kdf(key,
				       key_len,
				       (char *)bch + BCH_HEADER_VER_OFFSET,
				       BCH_HEADER_VER_SIZE,
				       (char *)bch + BCH_HEADER_DERIVE_STR_OFFSET,
				       BCH_HEADER_DERIVE_STR_SIZE,
				       dec_ctx->key_len,
				       dec_ctx->key);
	if (res != TEE_SUCCESS)
		EMSG("Derive DK failed with res(%x)", res);

	return res;
}

TEE_Result jetson_decrypt_cpubl_payload_init(struct dec_ctx *dec_ctx)
{
	TEE_Result res = TEE_ERROR_GENERIC;

	if (!dec_ctx)
		return TEE_ERROR_BAD_PARAMETERS;

	res = crypto_authenc_alloc_ctx(&(dec_ctx->ctx), TEE_ALG_AES_GCM);
	if (res != TEE_SUCCESS) {
		EMSG("crypto_authenc_alloc_ctx failed with res(%x)", res);
		goto out;
	}

	res = crypto_authenc_init(dec_ctx->ctx,
				  TEE_MODE_DECRYPT,
				  dec_ctx->key,
				  dec_ctx->key_len,
				  dec_ctx->iv_in_hdr,
				  BCH_HEADER_IV_SIZE,
				  BCH_HEADER_TAG_SIZE,
				  BCH_HEADER_AAD_SIZE,
				  dec_ctx->payload_size);
	if (res != TEE_SUCCESS) {
		EMSG("crypto_authenc_init failed with res(%x)", res);
		goto out;
	}

	res = crypto_authenc_update_aad(dec_ctx->ctx,
					TEE_MODE_DECRYPT,
					dec_ctx->aad_in_hdr,
					BCH_HEADER_AAD_SIZE);
	if (res != TEE_SUCCESS)
		EMSG("crypto_authenc_update_aad failed with res(%x)", res);

out:
	if ((res != TEE_SUCCESS) && (dec_ctx->ctx != NULL))
		crypto_authenc_free_ctx(dec_ctx->ctx);

	return res;
}

TEE_Result jetson_decrypt_cpubl_payload_update(struct dec_ctx *dec_ctx,
					       const uint8_t *src_data,
					       uint32_t src_len,
					       uint8_t *dst_data,
					       size_t *dst_len)
{
	TEE_Result res = TEE_ERROR_GENERIC;

	if (!dec_ctx)
		return TEE_ERROR_BAD_PARAMETERS;

	if (!src_data || src_len <= 0)
		return TEE_ERROR_BAD_PARAMETERS;

	if (!dst_data || !dst_len)
		return TEE_ERROR_BAD_PARAMETERS;

	res = crypto_authenc_update_payload(dec_ctx->ctx, TEE_MODE_DECRYPT,
					    src_data, src_len,
					    dst_data, dst_len);
	if (res != TEE_SUCCESS)
		EMSG("crypto_authenc_update_payload failed with res(%x)", res);

	if ((res != TEE_SUCCESS) && (dec_ctx->ctx != NULL))
		crypto_authenc_free_ctx(dec_ctx->ctx);

	return res;
}


TEE_Result jetson_decrypt_cpubl_payload_final(struct dec_ctx *dec_ctx,
					      const uint8_t *src_data,
					      uint32_t src_len,
					      uint8_t *dst_data,
					      size_t *dst_len)
{
	TEE_Result res = TEE_ERROR_GENERIC;
	size_t final_len;

	if (!dec_ctx)
		return TEE_ERROR_BAD_PARAMETERS;

	if (!src_data || src_len <= 0)
		return TEE_ERROR_BAD_PARAMETERS;

	if (!dst_data || !dst_len)
		return TEE_ERROR_BAD_PARAMETERS;

	final_len = *dst_len;

	res = crypto_authenc_dec_final(dec_ctx->ctx,
				       src_data, src_len,
				       dst_data, &final_len,
				       dec_ctx->tag_in_hdr,
				       BCH_HEADER_TAG_SIZE);
	if (res != TEE_SUCCESS) {
		EMSG("crypto_authenc_dec_final failed with res(%x)", res);
		goto out;
	}

	res = get_one_and_zeros_padding(dst_data, final_len, dst_len);
	if (res != TEE_SUCCESS)
		EMSG("unpadding failed with res(%x)", res);

out:
	if (dec_ctx->ctx)
		crypto_authenc_free_ctx(dec_ctx->ctx);

	return res;
}
