/*
 * Copyright (C) 2022 - 2026, Stephan Mueller <smueller@chronox.de>
 *
 * License: see LICENSE file in root directory
 *
 * THIS SOFTWARE IS PROVIDED ``AS IS'' AND ANY EXPRESS OR IMPLIED
 * WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
 * OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE, ALL OF
 * WHICH ARE HEREBY DISCLAIMED.  IN NO EVENT SHALL THE AUTHOR BE
 * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
 * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT
 * OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR
 * BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
 * LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE
 * USE OF THIS SOFTWARE, EVEN IF NOT ADVISED OF THE POSSIBILITY OF SUCH
 * DAMAGE.
 */

#include "aes_aesni.h"
#include "aes_armce.h"
#include "aes_c.h"
#include "aes_riscv64.h"
#include "aes_internal.h"
#include "lc_aes.h"
#include "lc_init.h"
#include "compare.h"
#include "ret_checkers.h"
#include "test_helper_common.h"
#include "timecop.h"
#include "visibility.h"

#define LC_EXEC_ONE_TEST(aes_impl)                                             \
	if (aes_impl)                                                          \
	ret += test_kw(aes_impl, #aes_impl)

static const uint8_t key128[] = { 0x75, 0x75, 0xda, 0x3a, 0x93, 0x60,
				  0x7c, 0xc2, 0xbf, 0xd8, 0xce, 0xc7,
				  0xaa, 0xdf, 0xd9, 0xa6 };
static const uint8_t pt128[] = {
	0x42, 0x13, 0x6d, 0x3c, 0x38, 0x4a, 0x3e, 0xea,
	0xc9, 0x5a, 0x06, 0x6f, 0xd2, 0x8f, 0xed, 0x3f
};
static const uint8_t ct128[] = {
	0xf6, 0x85, 0x94, 0x81, 0x6f, 0x64, 0xca, 0xa3,
	0xf5, 0x6f, 0xab, 0xea, 0x25, 0x48, 0xf5, 0xfb
};
static const uint8_t iv128[] = {
	0x03, 0x1f, 0x6b, 0xd7, 0xe6, 0x1e, 0x64, 0x3d
};

static const uint8_t key256[] = { 0x80, 0xaa, 0x99, 0x73, 0x27, 0xa4, 0x80,
				  0x6b, 0x6a, 0x7a, 0x41, 0xa5, 0x2b, 0x86,
				  0xc3, 0x71, 0x03, 0x86, 0xf9, 0x32, 0x78,
				  0x6e, 0xf7, 0x96, 0x76, 0xfa, 0xfb, 0x90,
				  0xb8, 0x26, 0x3c, 0x5f };
static const uint8_t pt256[] = {
	0x0a, 0x25, 0x6b, 0xa7, 0x5c, 0xfa, 0x03, 0xaa,
	0xa0, 0x2b, 0xa9, 0x42, 0x03, 0xf1, 0x5b, 0xaa
};
static const uint8_t ct256[] = {
	0xd3, 0x3d, 0x3d, 0x97, 0x7b, 0xf0, 0xa9, 0x15,
	0x59, 0xf9, 0x9c, 0x8a, 0xcd, 0x29, 0x3d, 0x43
};
static const uint8_t iv256[] = {
	0x42, 0x3c, 0x96, 0x0d, 0x8a, 0x2a, 0xc4, 0xc1
};

static int test_encrypt_kw_one(struct lc_sym_ctx *ctx, const uint8_t *key,
			       size_t keylen, const uint8_t *pt, size_t ptlen,
			       const uint8_t *ct, const uint8_t *iv)
{
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wvla"
	uint8_t out[ptlen + 8], out2[ptlen], ivout[8], extiv[8];
#pragma GCC diagnostic pop
	int ret, rc;

	/* Unpoison key to let implementation poison it */
	unpoison(key, keylen);

	/* Encrypt */
	CKINT(lc_sym_init(ctx));
	CKINT(lc_sym_setkey(ctx, key, keylen));
	lc_aes_kw_encrypt(ctx, pt, out, ptlen);
	rc = lc_compare(out + 8, ct, ptlen, "AES-KW encrypt ciphertext");
	rc += lc_compare(out, iv, 8, "AES-KW encrypt tag");

	/* Encrypt with external IV */
	memset(extiv, 0, sizeof(extiv));
	CKINT(lc_sym_init_iv(ctx, extiv, sizeof(extiv)));
	CKINT(lc_sym_encrypt_iv(ctx, pt, out + 8, ptlen, extiv, sizeof(extiv)));
	rc += lc_compare(out + 8, ct, ptlen,
			 "AES-KW encrypt external IV ciphertext");
	rc += lc_compare(extiv, iv, sizeof(extiv),
			 "AES-KW external IV encrypt tag");

	/* Decrypt */
	ret = lc_aes_kw_decrypt(ctx, out, out2, sizeof(out));
	if (ret) {
		rc++;
		printf("AES-KW Decryption error\n");
	}
	rc += lc_compare(out2, pt, ptlen, "AES-KW decrypt plaintext");

	/* Decrypt with external IV */
	CKINT(lc_sym_init_iv(ctx, extiv, sizeof(extiv)));
	CKINT(lc_sym_decrypt_iv(ctx, out + 8, out2, ptlen, extiv,
				sizeof(extiv)));
	rc += lc_compare(out2, pt, ptlen,
			 "AES-KW external IV decrypt plaintext");
	memset(ivout, 0xa6, sizeof(ivout));
	rc += lc_compare(extiv, ivout, sizeof(extiv),
			 "AES-KW decrypt external IV");

	/* Decrypt with error */
	out[0] = (uint8_t)((out[0] + 1) & 0xff);
	ret = lc_aes_kw_decrypt(ctx, out, out2, sizeof(out));
	if (ret != -EBADMSG) {
		rc++;
		printf("AES-KW Decryption error not caught\n");
	}
	ret = 0;

out:
	return ret ? !!ret : rc;
}

static int test_kw(const struct lc_sym *aes, const char *name)
{
	int ret;
	LC_SYM_CTX_ON_STACK(aes_kw, aes);

	printf("AES KW ctx %s (%s implementation) len %u\n", name,
	       aes == lc_aes_kw_c ? "C" : "accelerated",
	       (unsigned int)LC_SYM_CTX_SIZE);

	ret = test_encrypt_kw_one(aes_kw, key128, sizeof(key128), pt128,
				  sizeof(pt128), ct128, iv128);
	lc_sym_zero(aes_kw);

	ret += test_encrypt_kw_one(aes_kw, key256, sizeof(key256), pt256,
				   sizeof(pt256), ct256, iv256);
	lc_sym_zero(aes_kw);

	return ret;
}

LC_TEST_FUNC(int, main, int argc, char *argv[])
{
	int ret = 0;

	(void)argv;

	if (argc >= 2)
		lc_init(LC_INIT_AES_SBOX);

	LC_EXEC_ONE_TEST(lc_aes_kw);
	LC_EXEC_ONE_TEST(lc_aes_kw_aesni);
	LC_EXEC_ONE_TEST(lc_aes_kw_armce);
	LC_EXEC_ONE_TEST(lc_aes_kw_c);
	LC_EXEC_ONE_TEST(lc_aes_kw_riscv64);

	ret = test_validate_status(ret, lc_sym_alg_status(lc_aes_kw), 1);
	ret = test_validate_status(ret, lc_sym_alg_status(lc_aes_kw_aesni),
				   lc_aes_kw == lc_aes_kw_aesni);
	ret = test_validate_status(ret, lc_sym_alg_status(lc_aes_kw_armce),
				   lc_aes_kw == lc_aes_kw_armce);
	ret = test_validate_status(ret, lc_sym_alg_status(lc_aes_kw_c),
				   lc_aes_kw == lc_aes_kw_c);
	ret = test_validate_status(ret, lc_sym_alg_status(lc_aes_kw_riscv64),
				   lc_aes_kw == lc_aes_kw_riscv64);
	ret += test_print_status();

	return ret;
}
