// SPDX-License-Identifier: GPL-2.0-only
/*
 * Unit tests for NVMe authentication functions
 *
 * Copyright 2026 Google LLC
 */

#include <crypto/sha2.h>
#include <kunit/test.h>
#include <linux/nvme.h>
#include <linux/nvme-auth.h>
#include <linux/slab.h>

struct nvme_auth_test_values {
	u8 hmac_id;
	size_t hash_len;
	u8 expected_psk[NVME_AUTH_MAX_DIGEST_SIZE];
	char *expected_psk_digest;
	u8 expected_tls_psk[NVME_AUTH_MAX_DIGEST_SIZE];
};

static void kfree_action(void *ptr)
{
	kfree(ptr);
}

static void kunit_add_kfree_action(struct kunit *test, void *ptr)
{
	KUNIT_ASSERT_EQ(test, 0,
			kunit_add_action_or_reset(test, kfree_action, ptr));
}

/*
 * Test the derivation of a TLS PSK from the initial skey.  The vals parameter
 * gives the expected value of tls_psk as well as the intermediate values psk
 * and psk_digest.  The inputs are implicitly the fixed values set below.
 */
static void
test_nvme_auth_derive_tls_psk(struct kunit *test,
			      const struct nvme_auth_test_values *vals)
{
	const u8 hmac_id = vals->hmac_id;
	const size_t hash_len = vals->hash_len;
	const size_t skey_len = hash_len;
	u8 skey[NVME_AUTH_MAX_DIGEST_SIZE];
	u8 c1[NVME_AUTH_MAX_DIGEST_SIZE];
	u8 c2[NVME_AUTH_MAX_DIGEST_SIZE];
	const char *subsysnqn = "subsysnqn";
	const char *hostnqn = "hostnqn";
	u8 *psk = NULL, *tls_psk = NULL;
	char *psk_digest = NULL;
	size_t psk_len;
	int ret;

	for (int i = 0; i < NVME_AUTH_MAX_DIGEST_SIZE; i++) {
		skey[i] = 'A' + i;
		c1[i] = i;
		c2[i] = 0xff - i;
	}

	ret = nvme_auth_generate_psk(hmac_id, skey, skey_len, c1, c2, hash_len,
				     &psk, &psk_len);
	kunit_add_kfree_action(test, psk);
	KUNIT_ASSERT_EQ(test, 0, ret);
	KUNIT_ASSERT_EQ(test, hash_len, psk_len);
	KUNIT_ASSERT_MEMEQ(test, vals->expected_psk, psk, psk_len);

	ret = nvme_auth_generate_digest(hmac_id, psk, psk_len, subsysnqn,
					hostnqn, &psk_digest);
	kunit_add_kfree_action(test, psk_digest);
	if (vals->expected_psk_digest == NULL) {
		/*
		 * Algorithm has an ID assigned but is not supported by
		 * nvme_auth_generate_digest().
		 */
		KUNIT_ASSERT_EQ(test, -EINVAL, ret);
		return;
	}
	KUNIT_ASSERT_EQ(test, 0, ret);
	KUNIT_ASSERT_STREQ(test, vals->expected_psk_digest, psk_digest);

	ret = nvme_auth_derive_tls_psk(hmac_id, psk, psk_len, psk_digest,
				       &tls_psk);
	kunit_add_kfree_action(test, tls_psk);
	KUNIT_ASSERT_EQ(test, 0, ret);
	KUNIT_ASSERT_MEMEQ(test, vals->expected_tls_psk, tls_psk, psk_len);
}

static void test_nvme_auth_derive_tls_psk_hmac_sha256(struct kunit *test)
{
	static const struct nvme_auth_test_values vals = {
		.hmac_id = NVME_AUTH_HASH_SHA256,
		.hash_len = SHA256_DIGEST_SIZE,
		.expected_psk = {
			0x17, 0x33, 0xc5, 0x9f, 0xa7, 0xf4, 0x8f, 0xcf,
			0x37, 0xf5, 0xf2, 0x6f, 0xc4, 0xff, 0x02, 0x68,
			0xad, 0x4f, 0x78, 0xe0, 0x30, 0xf4, 0xf3, 0xb0,
			0xbf, 0xd1, 0xd4, 0x7e, 0x7b, 0xb1, 0x44, 0x7a,
		},
		.expected_psk_digest = "OldoKuTfKddMuyCznAZojkWD7P4D9/AtzDzLimtOxqI=",
		.expected_tls_psk = {
			0x3c, 0x17, 0xda, 0x62, 0x84, 0x74, 0xa0, 0x4d,
			0x22, 0x47, 0xc4, 0xca, 0xb4, 0x79, 0x68, 0xc9,
			0x15, 0x38, 0x81, 0x93, 0xf7, 0xc0, 0x71, 0xbd,
			0x94, 0x89, 0xcc, 0x36, 0x66, 0xcd, 0x7c, 0xc8,
		},
	};

	test_nvme_auth_derive_tls_psk(test, &vals);
}

static void test_nvme_auth_derive_tls_psk_hmac_sha384(struct kunit *test)
{
	static const struct nvme_auth_test_values vals = {
		.hmac_id = NVME_AUTH_HASH_SHA384,
		.hash_len = SHA384_DIGEST_SIZE,
		.expected_psk = {
			0xf1, 0x4b, 0x2d, 0xd3, 0x23, 0x4c, 0x45, 0x96,
			0x94, 0xd3, 0xbc, 0x63, 0xf8, 0x96, 0x8b, 0xd6,
			0xb3, 0x7c, 0x2c, 0x6d, 0xe8, 0x49, 0xe2, 0x2e,
			0x11, 0x87, 0x49, 0x00, 0x1c, 0xe4, 0xbb, 0xe8,
			0x64, 0x0b, 0x9e, 0x3a, 0x74, 0x8c, 0xb1, 0x1c,
			0xe4, 0xb1, 0xd7, 0x1d, 0x35, 0x9c, 0xce, 0x39,
		},
		.expected_psk_digest = "cffMWk8TSS7HOQebjgYEIkrPrjWPV4JE5cdPB8WhEvY4JBW5YynKyv66XscN4A9n",
		.expected_tls_psk = {
			0x27, 0x74, 0x75, 0x32, 0x33, 0x53, 0x7b, 0x3f,
			0xa5, 0x0e, 0xb7, 0xd1, 0x6a, 0x8e, 0x43, 0x45,
			0x7d, 0x85, 0xf4, 0x90, 0x6c, 0x00, 0x5b, 0x22,
			0x36, 0x61, 0x6c, 0x5d, 0x80, 0x93, 0x9d, 0x08,
			0x98, 0xff, 0xf1, 0x5b, 0xb8, 0xb7, 0x71, 0x19,
			0xd2, 0xbe, 0x0a, 0xac, 0x42, 0x3e, 0x75, 0x90,
		},
	};

	test_nvme_auth_derive_tls_psk(test, &vals);
}

static void test_nvme_auth_derive_tls_psk_hmac_sha512(struct kunit *test)
{
	static const struct nvme_auth_test_values vals = {
		.hmac_id = NVME_AUTH_HASH_SHA512,
		.hash_len = SHA512_DIGEST_SIZE,
		.expected_psk = {
			0x9c, 0x9f, 0x08, 0x9a, 0x61, 0x8b, 0x47, 0xd2,
			0xd7, 0x5f, 0x4b, 0x6c, 0x28, 0x07, 0x04, 0x24,
			0x48, 0x7b, 0x44, 0x5d, 0xd9, 0x6e, 0x70, 0xc4,
			0xc0, 0x9b, 0x55, 0xe8, 0xb6, 0x00, 0x01, 0x52,
			0xa3, 0x36, 0x3c, 0x34, 0x54, 0x04, 0x3f, 0x38,
			0xf0, 0xb8, 0x50, 0x36, 0xde, 0xd4, 0x06, 0x55,
			0x35, 0x0a, 0xa8, 0x7b, 0x8b, 0x6a, 0x28, 0x2b,
			0x5c, 0x1a, 0xca, 0xe1, 0x62, 0x33, 0xdd, 0x5b,
		},
		/* nvme_auth_generate_digest() doesn't support SHA-512 yet. */
		.expected_psk_digest = NULL,
	};

	test_nvme_auth_derive_tls_psk(test, &vals);
}

static struct kunit_case nvme_auth_test_cases[] = {
	KUNIT_CASE(test_nvme_auth_derive_tls_psk_hmac_sha256),
	KUNIT_CASE(test_nvme_auth_derive_tls_psk_hmac_sha384),
	KUNIT_CASE(test_nvme_auth_derive_tls_psk_hmac_sha512),
	{},
};

static struct kunit_suite nvme_auth_test_suite = {
	.name = "nvme-auth",
	.test_cases = nvme_auth_test_cases,
};
kunit_test_suite(nvme_auth_test_suite);

MODULE_DESCRIPTION("Unit tests for NVMe authentication functions");
MODULE_LICENSE("GPL");