lib/crypto/tests/mldsa_kunit.c - third_party/linux - Git at Google

 // SPDX-License-Identifier: GPL-2.0-or-later
 /*
  * KUnit tests and benchmark for ML-DSA
  *
  * Copyright 2025 Google LLC
  */
 #include <crypto/mldsa.h>
 #include <kunit/test.h>
 #include <linux/random.h>
 #include <linux/unaligned.h>

 #define Q 8380417 /* The prime q = 2^23 - 2^13 + 1 */

 /* ML-DSA parameters that the tests use */
 static const struct {
 	int sig_len;
 	int pk_len;
 	int k;
 	int lambda;
 	int gamma1;
 	int beta;
 	int omega;
 } params[] = {
 	[MLDSA44] = {
 		.sig_len = MLDSA44_SIGNATURE_SIZE,
 		.pk_len = MLDSA44_PUBLIC_KEY_SIZE,
 		.k = 4,
 		.lambda = 128,
 		.gamma1 = 1 << 17,
 		.beta = 78,
 		.omega = 80,
 	},
 	[MLDSA65] = {
 		.sig_len = MLDSA65_SIGNATURE_SIZE,
 		.pk_len = MLDSA65_PUBLIC_KEY_SIZE,
 		.k = 6,
 		.lambda = 192,
 		.gamma1 = 1 << 19,
 		.beta = 196,
 		.omega = 55,
 	},
 	[MLDSA87] = {
 		.sig_len = MLDSA87_SIGNATURE_SIZE,
 		.pk_len = MLDSA87_PUBLIC_KEY_SIZE,
 		.k = 8,
 		.lambda = 256,
 		.gamma1 = 1 << 19,
 		.beta = 120,
 		.omega = 75,
 	},
 };

 #include "mldsa-testvecs.h"

 static void do_mldsa_and_assert_success(struct kunit *test,
 					const struct mldsa_testvector *tv)
 {
 	int err = mldsa_verify(tv->alg, tv->sig, tv->sig_len, tv->msg,
 			       tv->msg_len, tv->pk, tv->pk_len);
 	KUNIT_ASSERT_EQ(test, err, 0);
 }

 static u8 *kunit_kmemdup_or_fail(struct kunit *test, const u8 *src, size_t len)
 {
 	u8 *dst = kunit_kmalloc(test, len, GFP_KERNEL);

 	KUNIT_ASSERT_NOT_NULL(test, dst);
 	return memcpy(dst, src, len);
 }

 /*
  * Test that changing coefficients in a valid signature's z vector results in
  * the following behavior from mldsa_verify():
  *
  *  * -EBADMSG if a coefficient is changed to have an out-of-range value, i.e.
  *    absolute value >= gamma1 - beta, corresponding to the verifier detecting
  *    the out-of-range coefficient and rejecting the signature as malformed
  *
  *  * -EKEYREJECTED if a coefficient is changed to a different in-range value,
  *    i.e. absolute value < gamma1 - beta, corresponding to the verifier
  *    continuing to the "real" signature check and that check failing
  */
 static void test_mldsa_z_range(struct kunit *test,
 			       const struct mldsa_testvector *tv)
 {
 	u8 *sig = kunit_kmemdup_or_fail(test, tv->sig, tv->sig_len);
 	const int lambda = params[tv->alg].lambda;
 	const s32 gamma1 = params[tv->alg].gamma1;
 	const int beta = params[tv->alg].beta;
 	/*
 	 * We just modify the first coefficient.  The coefficient is gamma1
 	 * minus either the first 18 or 20 bits of the u32, depending on gamma1.
 	 *
 	 * The layout of ML-DSA signatures is ctilde || z || h.  ctilde is
 	 * lambda / 4 bytes, so z starts at &sig[lambda / 4].
 	 */
 	u8 *z_ptr = &sig[lambda / 4];
 	const u32 z_data = get_unaligned_le32(z_ptr);
 	const u32 mask = (gamma1 << 1) - 1;
 	/* These are the four boundaries of the out-of-range values. */
 	const s32 out_of_range_coeffs[] = {
 		-gamma1 + 1,
 		-(gamma1 - beta),
 		gamma1,
 		gamma1 - beta,
 	};
 	/*
 	 * These are the two boundaries of the valid range, along with 0.  We
 	 * assume that none of these matches the original coefficient.
 	 */
 	const s32 in_range_coeffs[] = {
 		-(gamma1 - beta - 1),
 		0,
 		gamma1 - beta - 1,
 	};

 	/* Initially the signature is valid. */
 	do_mldsa_and_assert_success(test, tv);

 	/* Test some out-of-range coefficients. */
 	for (int i = 0; i < ARRAY_SIZE(out_of_range_coeffs); i++) {
 		const s32 c = out_of_range_coeffs[i];

 		put_unaligned_le32((z_data & ~mask) | (mask & (gamma1 - c)),
 				   z_ptr);
 		KUNIT_ASSERT_EQ(test, -EBADMSG,
 				mldsa_verify(tv->alg, sig, tv->sig_len, tv->msg,
 					     tv->msg_len, tv->pk, tv->pk_len));
 	}

 	/* Test some in-range coefficients. */
 	for (int i = 0; i < ARRAY_SIZE(in_range_coeffs); i++) {
 		const s32 c = in_range_coeffs[i];

 		put_unaligned_le32((z_data & ~mask) | (mask & (gamma1 - c)),
 				   z_ptr);
 		KUNIT_ASSERT_EQ(test, -EKEYREJECTED,
 				mldsa_verify(tv->alg, sig, tv->sig_len, tv->msg,
 					     tv->msg_len, tv->pk, tv->pk_len));
 	}
 }

 /* Test that mldsa_verify() rejects malformed hint vectors with -EBADMSG. */
 static void test_mldsa_bad_hints(struct kunit *test,
 				 const struct mldsa_testvector *tv)
 {
 	const int omega = params[tv->alg].omega;
 	const int k = params[tv->alg].k;
 	u8 *sig = kunit_kmemdup_or_fail(test, tv->sig, tv->sig_len);
 	/* Pointer to the encoded hint vector in the signature */
 	u8 *hintvec = &sig[tv->sig_len - omega - k];
 	u8 h;

 	/* Initially the signature is valid. */
 	do_mldsa_and_assert_success(test, tv);

 	/* Cumulative hint count exceeds omega */
 	memcpy(sig, tv->sig, tv->sig_len);
 	hintvec[omega + k - 1] = omega + 1;
 	KUNIT_ASSERT_EQ(test, -EBADMSG,
 			mldsa_verify(tv->alg, sig, tv->sig_len, tv->msg,
 				     tv->msg_len, tv->pk, tv->pk_len));

 	/* Cumulative hint count decreases */
 	memcpy(sig, tv->sig, tv->sig_len);
 	KUNIT_ASSERT_GE(test, hintvec[omega + k - 2], 1);
 	hintvec[omega + k - 1] = hintvec[omega + k - 2] - 1;
 	KUNIT_ASSERT_EQ(test, -EBADMSG,
 			mldsa_verify(tv->alg, sig, tv->sig_len, tv->msg,
 				     tv->msg_len, tv->pk, tv->pk_len));

 	/*
 	 * Hint indices out of order.  To test this, swap hintvec[0] and
 	 * hintvec[1].  This assumes that the original valid signature had at
 	 * least two nonzero hints in the first element (asserted below).
 	 */
 	memcpy(sig, tv->sig, tv->sig_len);
 	KUNIT_ASSERT_GE(test, hintvec[omega], 2);
 	h = hintvec[0];
 	hintvec[0] = hintvec[1];
 	hintvec[1] = h;
 	KUNIT_ASSERT_EQ(test, -EBADMSG,
 			mldsa_verify(tv->alg, sig, tv->sig_len, tv->msg,
 				     tv->msg_len, tv->pk, tv->pk_len));

 	/*
 	 * Extra hint indices given.  For this test to work, the original valid
 	 * signature must have fewer than omega nonzero hints (asserted below).
 	 */
 	memcpy(sig, tv->sig, tv->sig_len);
 	KUNIT_ASSERT_LT(test, hintvec[omega + k - 1], omega);
 	hintvec[omega - 1] = 0xff;
 	KUNIT_ASSERT_EQ(test, -EBADMSG,
 			mldsa_verify(tv->alg, sig, tv->sig_len, tv->msg,
 				     tv->msg_len, tv->pk, tv->pk_len));
 }

 static void test_mldsa_mutation(struct kunit *test,
 				const struct mldsa_testvector *tv)
 {
 	const int sig_len = tv->sig_len;
 	const int msg_len = tv->msg_len;
 	const int pk_len = tv->pk_len;
 	const int num_iter = 200;
 	u8 *sig = kunit_kmemdup_or_fail(test, tv->sig, sig_len);
 	u8 *msg = kunit_kmemdup_or_fail(test, tv->msg, msg_len);
 	u8 *pk = kunit_kmemdup_or_fail(test, tv->pk, pk_len);

 	/* Initially the signature is valid. */
 	do_mldsa_and_assert_success(test, tv);

 	/* Changing any bit in the signature should invalidate the signature */
 	for (int i = 0; i < num_iter; i++) {
 		size_t pos = get_random_u32_below(sig_len);
 		u8 b = 1 << get_random_u32_below(8);

 		sig[pos] ^= b;
 		KUNIT_ASSERT_NE(test, 0,
 				mldsa_verify(tv->alg, sig, sig_len, msg,
 					     msg_len, pk, pk_len));
 		sig[pos] ^= b;
 	}

 	/* Changing any bit in the message should invalidate the signature */
 	for (int i = 0; i < num_iter; i++) {
 		size_t pos = get_random_u32_below(msg_len);
 		u8 b = 1 << get_random_u32_below(8);

 		msg[pos] ^= b;
 		KUNIT_ASSERT_NE(test, 0,
 				mldsa_verify(tv->alg, sig, sig_len, msg,
 					     msg_len, pk, pk_len));
 		msg[pos] ^= b;
 	}

 	/* Changing any bit in the public key should invalidate the signature */
 	for (int i = 0; i < num_iter; i++) {
 		size_t pos = get_random_u32_below(pk_len);
 		u8 b = 1 << get_random_u32_below(8);

 		pk[pos] ^= b;
 		KUNIT_ASSERT_NE(test, 0,
 				mldsa_verify(tv->alg, sig, sig_len, msg,
 					     msg_len, pk, pk_len));
 		pk[pos] ^= b;
 	}

 	/* All changes should have been undone. */
 	KUNIT_ASSERT_EQ(test, 0,
 			mldsa_verify(tv->alg, sig, sig_len, msg, msg_len, pk,
 				     pk_len));
 }

 static void test_mldsa(struct kunit *test, const struct mldsa_testvector *tv)
 {
 	/* Valid signature */
 	KUNIT_ASSERT_EQ(test, tv->sig_len, params[tv->alg].sig_len);
 	KUNIT_ASSERT_EQ(test, tv->pk_len, params[tv->alg].pk_len);
 	do_mldsa_and_assert_success(test, tv);

 	/* Signature too short */
 	KUNIT_ASSERT_EQ(test, -EBADMSG,
 			mldsa_verify(tv->alg, tv->sig, tv->sig_len - 1, tv->msg,
 				     tv->msg_len, tv->pk, tv->pk_len));

 	/* Signature too long */
 	KUNIT_ASSERT_EQ(test, -EBADMSG,
 			mldsa_verify(tv->alg, tv->sig, tv->sig_len + 1, tv->msg,
 				     tv->msg_len, tv->pk, tv->pk_len));

 	/* Public key too short */
 	KUNIT_ASSERT_EQ(test, -EBADMSG,
 			mldsa_verify(tv->alg, tv->sig, tv->sig_len, tv->msg,
 				     tv->msg_len, tv->pk, tv->pk_len - 1));

 	/* Public key too long */
 	KUNIT_ASSERT_EQ(test, -EBADMSG,
 			mldsa_verify(tv->alg, tv->sig, tv->sig_len, tv->msg,
 				     tv->msg_len, tv->pk, tv->pk_len + 1));

 	/*
 	 * Message too short.  Error is EKEYREJECTED because it gets rejected by
 	 * the "real" signature check rather than the well-formedness checks.
 	 */
 	KUNIT_ASSERT_EQ(test, -EKEYREJECTED,
 			mldsa_verify(tv->alg, tv->sig, tv->sig_len, tv->msg,
 				     tv->msg_len - 1, tv->pk, tv->pk_len));
 	/*
 	 * Can't simply try (tv->msg, tv->msg_len + 1) too, as tv->msg would be
 	 * accessed out of bounds.  However, ML-DSA just hashes the message and
 	 * doesn't handle different message lengths differently anyway.
 	 */

 	/* Test the validity checks on the z vector. */
 	test_mldsa_z_range(test, tv);

 	/* Test the validity checks on the hint vector. */
 	test_mldsa_bad_hints(test, tv);

 	/* Test randomly mutating the inputs. */
 	test_mldsa_mutation(test, tv);
 }

 static void test_mldsa44(struct kunit *test)
 {
 	test_mldsa(test, &mldsa44_testvector);
 }

 static void test_mldsa65(struct kunit *test)
 {
 	test_mldsa(test, &mldsa65_testvector);
 }

 static void test_mldsa87(struct kunit *test)
 {
 	test_mldsa(test, &mldsa87_testvector);
 }

 static s32 mod(s32 a, s32 m)
 {
 	a %= m;
 	if (a < 0)
 		a += m;
 	return a;
 }

 static s32 symmetric_mod(s32 a, s32 m)
 {
 	a = mod(a, m);
 	if (a > m / 2)
 		a -= m;
 	return a;
 }

 /* Mechanical, inefficient translation of FIPS 204 Algorithm 36, Decompose */
 static void decompose_ref(s32 r, s32 gamma2, s32 *r0, s32 *r1)
 {
 	s32 rplus = mod(r, Q);

 	*r0 = symmetric_mod(rplus, 2 * gamma2);
 	if (rplus - *r0 == Q - 1) {
 		*r1 = 0;
 		*r0 = *r0 - 1;
 	} else {
 		*r1 = (rplus - *r0) / (2 * gamma2);
 	}
 }

 /* Mechanical, inefficient translation of FIPS 204 Algorithm 40, UseHint */
 static s32 use_hint_ref(u8 h, s32 r, s32 gamma2)
 {
 	s32 m = (Q - 1) / (2 * gamma2);
 	s32 r0, r1;

 	decompose_ref(r, gamma2, &r0, &r1);
 	if (h == 1 && r0 > 0)
 		return mod(r1 + 1, m);
 	if (h == 1 && r0 <= 0)
 		return mod(r1 - 1, m);
 	return r1;
 }

 /*
  * Test that for all possible inputs, mldsa_use_hint() gives the same output as
  * a mechanical translation of the pseudocode from FIPS 204.
  */
 static void test_mldsa_use_hint(struct kunit *test)
 {
 	for (int i = 0; i < 2; i++) {
 		const s32 gamma2 = (Q - 1) / (i == 0 ? 88 : 32);

 		for (u8 h = 0; h < 2; h++) {
 			for (s32 r = 0; r < Q; r++) {
 				KUNIT_ASSERT_EQ(test,
 						mldsa_use_hint(h, r, gamma2),
 						use_hint_ref(h, r, gamma2));
 			}
 		}
 	}
 }

 static void benchmark_mldsa(struct kunit *test,
 			    const struct mldsa_testvector *tv)
 {
 	const int warmup_niter = 200;
 	const int benchmark_niter = 200;
 	u64 t0, t1;

 	if (!IS_ENABLED(CONFIG_CRYPTO_LIB_BENCHMARK))
 		kunit_skip(test, "not enabled");

 	for (int i = 0; i < warmup_niter; i++)
 		do_mldsa_and_assert_success(test, tv);

 	t0 = ktime_get_ns();
 	for (int i = 0; i < benchmark_niter; i++)
 		do_mldsa_and_assert_success(test, tv);
 	t1 = ktime_get_ns();
 	kunit_info(test, "%llu ops/s",
 		   div64_u64((u64)benchmark_niter * NSEC_PER_SEC,
 			     t1 - t0 ?: 1));
 }

 static void benchmark_mldsa44(struct kunit *test)
 {
 	benchmark_mldsa(test, &mldsa44_testvector);
 }

 static void benchmark_mldsa65(struct kunit *test)
 {
 	benchmark_mldsa(test, &mldsa65_testvector);
 }

 static void benchmark_mldsa87(struct kunit *test)
 {
 	benchmark_mldsa(test, &mldsa87_testvector);
 }

 static struct kunit_case mldsa_kunit_cases[] = {
 	KUNIT_CASE(test_mldsa44),
 	KUNIT_CASE(test_mldsa65),
 	KUNIT_CASE(test_mldsa87),
 	KUNIT_CASE(test_mldsa_use_hint),
 	KUNIT_CASE(benchmark_mldsa44),
 	KUNIT_CASE(benchmark_mldsa65),
 	KUNIT_CASE(benchmark_mldsa87),
 	{},
 };

 static struct kunit_suite mldsa_kunit_suite = {
 	.name = "mldsa",
 	.test_cases = mldsa_kunit_cases,
 };
 kunit_test_suite(mldsa_kunit_suite);

 MODULE_DESCRIPTION("KUnit tests and benchmark for ML-DSA");
 MODULE_IMPORT_NS("EXPORTED_FOR_KUNIT_TESTING");
 MODULE_LICENSE("GPL");
	// SPDX-License-Identifier: GPL-2.0-or-later
	/*
	* KUnit tests and benchmark for ML-DSA
	*
	* Copyright 2025 Google LLC
	*/
	#include <crypto/mldsa.h>
	#include <kunit/test.h>
	#include <linux/random.h>
	#include <linux/unaligned.h>

	#define Q 8380417 /* The prime q = 2^23 - 2^13 + 1 */

	/* ML-DSA parameters that the tests use */
	static const struct {
	int sig_len;
	int pk_len;
	int k;
	int lambda;
	int gamma1;
	int beta;
	int omega;
	} params[] = {
	[MLDSA44] = {
	.sig_len = MLDSA44_SIGNATURE_SIZE,
	.pk_len = MLDSA44_PUBLIC_KEY_SIZE,
	.k = 4,
	.lambda = 128,
	.gamma1 = 1 << 17,
	.beta = 78,
	.omega = 80,
	},
	[MLDSA65] = {
	.sig_len = MLDSA65_SIGNATURE_SIZE,
	.pk_len = MLDSA65_PUBLIC_KEY_SIZE,
	.k = 6,
	.lambda = 192,
	.gamma1 = 1 << 19,
	.beta = 196,
	.omega = 55,
	},
	[MLDSA87] = {
	.sig_len = MLDSA87_SIGNATURE_SIZE,
	.pk_len = MLDSA87_PUBLIC_KEY_SIZE,
	.k = 8,
	.lambda = 256,
	.gamma1 = 1 << 19,
	.beta = 120,
	.omega = 75,
	},
	};

	#include "mldsa-testvecs.h"

	static void do_mldsa_and_assert_success(struct kunit *test,
	const struct mldsa_testvector *tv)
	{
	int err = mldsa_verify(tv->alg, tv->sig, tv->sig_len, tv->msg,
	tv->msg_len, tv->pk, tv->pk_len);
	KUNIT_ASSERT_EQ(test, err, 0);
	}

	static u8 kunit_kmemdup_or_fail(struct kunit test, const u8 *src, size_t len)
	{
	u8 *dst = kunit_kmalloc(test, len, GFP_KERNEL);

	KUNIT_ASSERT_NOT_NULL(test, dst);
	return memcpy(dst, src, len);
	}

	/*
	* Test that changing coefficients in a valid signature's z vector results in
	* the following behavior from mldsa_verify():
	*
	* * -EBADMSG if a coefficient is changed to have an out-of-range value, i.e.
	* absolute value >= gamma1 - beta, corresponding to the verifier detecting
	* the out-of-range coefficient and rejecting the signature as malformed
	*
	* * -EKEYREJECTED if a coefficient is changed to a different in-range value,
	* i.e. absolute value < gamma1 - beta, corresponding to the verifier
	* continuing to the "real" signature check and that check failing
	*/
	static void test_mldsa_z_range(struct kunit *test,
	const struct mldsa_testvector *tv)
	{
	u8 *sig = kunit_kmemdup_or_fail(test, tv->sig, tv->sig_len);
	const int lambda = params[tv->alg].lambda;
	const s32 gamma1 = params[tv->alg].gamma1;
	const int beta = params[tv->alg].beta;
	/*
	* We just modify the first coefficient. The coefficient is gamma1
	* minus either the first 18 or 20 bits of the u32, depending on gamma1.
	*
	* The layout of ML-DSA signatures is ctilde \|\| z \|\| h. ctilde is
	* lambda / 4 bytes, so z starts at &sig[lambda / 4].
	*/
	u8 *z_ptr = &sig[lambda / 4];
	const u32 z_data = get_unaligned_le32(z_ptr);
	const u32 mask = (gamma1 << 1) - 1;
	/* These are the four boundaries of the out-of-range values. */
	const s32 out_of_range_coeffs[] = {
	-gamma1 + 1,
	-(gamma1 - beta),
	gamma1,
	gamma1 - beta,
	};
	/*
	* These are the two boundaries of the valid range, along with 0. We
	* assume that none of these matches the original coefficient.
	*/
	const s32 in_range_coeffs[] = {
	-(gamma1 - beta - 1),
	0,
	gamma1 - beta - 1,
	};

	/* Initially the signature is valid. */
	do_mldsa_and_assert_success(test, tv);

	/* Test some out-of-range coefficients. */
	for (int i = 0; i < ARRAY_SIZE(out_of_range_coeffs); i++) {
	const s32 c = out_of_range_coeffs[i];

	put_unaligned_le32((z_data & ~mask) \| (mask & (gamma1 - c)),
	z_ptr);
	KUNIT_ASSERT_EQ(test, -EBADMSG,
	mldsa_verify(tv->alg, sig, tv->sig_len, tv->msg,
	tv->msg_len, tv->pk, tv->pk_len));
	}

	/* Test some in-range coefficients. */
	for (int i = 0; i < ARRAY_SIZE(in_range_coeffs); i++) {
	const s32 c = in_range_coeffs[i];

	put_unaligned_le32((z_data & ~mask) \| (mask & (gamma1 - c)),
	z_ptr);
	KUNIT_ASSERT_EQ(test, -EKEYREJECTED,
	mldsa_verify(tv->alg, sig, tv->sig_len, tv->msg,
	tv->msg_len, tv->pk, tv->pk_len));
	}
	}

	/* Test that mldsa_verify() rejects malformed hint vectors with -EBADMSG. */
	static void test_mldsa_bad_hints(struct kunit *test,
	const struct mldsa_testvector *tv)
	{
	const int omega = params[tv->alg].omega;
	const int k = params[tv->alg].k;
	u8 *sig = kunit_kmemdup_or_fail(test, tv->sig, tv->sig_len);
	/* Pointer to the encoded hint vector in the signature */
	u8 *hintvec = &sig[tv->sig_len - omega - k];
	u8 h;

	/* Initially the signature is valid. */
	do_mldsa_and_assert_success(test, tv);

	/* Cumulative hint count exceeds omega */
	memcpy(sig, tv->sig, tv->sig_len);
	hintvec[omega + k - 1] = omega + 1;
	KUNIT_ASSERT_EQ(test, -EBADMSG,
	mldsa_verify(tv->alg, sig, tv->sig_len, tv->msg,
	tv->msg_len, tv->pk, tv->pk_len));

	/* Cumulative hint count decreases */
	memcpy(sig, tv->sig, tv->sig_len);
	KUNIT_ASSERT_GE(test, hintvec[omega + k - 2], 1);
	hintvec[omega + k - 1] = hintvec[omega + k - 2] - 1;
	KUNIT_ASSERT_EQ(test, -EBADMSG,
	mldsa_verify(tv->alg, sig, tv->sig_len, tv->msg,
	tv->msg_len, tv->pk, tv->pk_len));

	/*
	* Hint indices out of order. To test this, swap hintvec[0] and
	* hintvec[1]. This assumes that the original valid signature had at
	* least two nonzero hints in the first element (asserted below).
	*/
	memcpy(sig, tv->sig, tv->sig_len);
	KUNIT_ASSERT_GE(test, hintvec[omega], 2);
	h = hintvec[0];
	hintvec[0] = hintvec[1];
	hintvec[1] = h;
	KUNIT_ASSERT_EQ(test, -EBADMSG,
	mldsa_verify(tv->alg, sig, tv->sig_len, tv->msg,
	tv->msg_len, tv->pk, tv->pk_len));

	/*
	* Extra hint indices given. For this test to work, the original valid
	* signature must have fewer than omega nonzero hints (asserted below).
	*/
	memcpy(sig, tv->sig, tv->sig_len);
	KUNIT_ASSERT_LT(test, hintvec[omega + k - 1], omega);
	hintvec[omega - 1] = 0xff;
	KUNIT_ASSERT_EQ(test, -EBADMSG,
	mldsa_verify(tv->alg, sig, tv->sig_len, tv->msg,
	tv->msg_len, tv->pk, tv->pk_len));
	}

	static void test_mldsa_mutation(struct kunit *test,
	const struct mldsa_testvector *tv)
	{
	const int sig_len = tv->sig_len;
	const int msg_len = tv->msg_len;
	const int pk_len = tv->pk_len;
	const int num_iter = 200;
	u8 *sig = kunit_kmemdup_or_fail(test, tv->sig, sig_len);
	u8 *msg = kunit_kmemdup_or_fail(test, tv->msg, msg_len);
	u8 *pk = kunit_kmemdup_or_fail(test, tv->pk, pk_len);

	/* Initially the signature is valid. */
	do_mldsa_and_assert_success(test, tv);

	/* Changing any bit in the signature should invalidate the signature */
	for (int i = 0; i < num_iter; i++) {
	size_t pos = get_random_u32_below(sig_len);
	u8 b = 1 << get_random_u32_below(8);

	sig[pos] ^= b;
	KUNIT_ASSERT_NE(test, 0,
	mldsa_verify(tv->alg, sig, sig_len, msg,
	msg_len, pk, pk_len));
	sig[pos] ^= b;
	}

	/* Changing any bit in the message should invalidate the signature */
	for (int i = 0; i < num_iter; i++) {
	size_t pos = get_random_u32_below(msg_len);
	u8 b = 1 << get_random_u32_below(8);

	msg[pos] ^= b;
	KUNIT_ASSERT_NE(test, 0,
	mldsa_verify(tv->alg, sig, sig_len, msg,
	msg_len, pk, pk_len));
	msg[pos] ^= b;
	}

	/* Changing any bit in the public key should invalidate the signature */
	for (int i = 0; i < num_iter; i++) {
	size_t pos = get_random_u32_below(pk_len);
	u8 b = 1 << get_random_u32_below(8);

	pk[pos] ^= b;
	KUNIT_ASSERT_NE(test, 0,
	mldsa_verify(tv->alg, sig, sig_len, msg,
	msg_len, pk, pk_len));
	pk[pos] ^= b;
	}

	/* All changes should have been undone. */
	KUNIT_ASSERT_EQ(test, 0,
	mldsa_verify(tv->alg, sig, sig_len, msg, msg_len, pk,
	pk_len));
	}

	static void test_mldsa(struct kunit test, const struct mldsa_testvector tv)
	{
	/* Valid signature */
	KUNIT_ASSERT_EQ(test, tv->sig_len, params[tv->alg].sig_len);
	KUNIT_ASSERT_EQ(test, tv->pk_len, params[tv->alg].pk_len);
	do_mldsa_and_assert_success(test, tv);

	/* Signature too short */
	KUNIT_ASSERT_EQ(test, -EBADMSG,
	mldsa_verify(tv->alg, tv->sig, tv->sig_len - 1, tv->msg,
	tv->msg_len, tv->pk, tv->pk_len));

	/* Signature too long */
	KUNIT_ASSERT_EQ(test, -EBADMSG,
	mldsa_verify(tv->alg, tv->sig, tv->sig_len + 1, tv->msg,
	tv->msg_len, tv->pk, tv->pk_len));

	/* Public key too short */
	KUNIT_ASSERT_EQ(test, -EBADMSG,
	mldsa_verify(tv->alg, tv->sig, tv->sig_len, tv->msg,
	tv->msg_len, tv->pk, tv->pk_len - 1));

	/* Public key too long */
	KUNIT_ASSERT_EQ(test, -EBADMSG,
	mldsa_verify(tv->alg, tv->sig, tv->sig_len, tv->msg,
	tv->msg_len, tv->pk, tv->pk_len + 1));

	/*
	* Message too short. Error is EKEYREJECTED because it gets rejected by
	* the "real" signature check rather than the well-formedness checks.
	*/
	KUNIT_ASSERT_EQ(test, -EKEYREJECTED,
	mldsa_verify(tv->alg, tv->sig, tv->sig_len, tv->msg,
	tv->msg_len - 1, tv->pk, tv->pk_len));
	/*
	* Can't simply try (tv->msg, tv->msg_len + 1) too, as tv->msg would be
	* accessed out of bounds. However, ML-DSA just hashes the message and
	* doesn't handle different message lengths differently anyway.
	*/

	/* Test the validity checks on the z vector. */
	test_mldsa_z_range(test, tv);

	/* Test the validity checks on the hint vector. */
	test_mldsa_bad_hints(test, tv);

	/* Test randomly mutating the inputs. */
	test_mldsa_mutation(test, tv);
	}

	static void test_mldsa44(struct kunit *test)
	{
	test_mldsa(test, &mldsa44_testvector);
	}

	static void test_mldsa65(struct kunit *test)
	{
	test_mldsa(test, &mldsa65_testvector);
	}

	static void test_mldsa87(struct kunit *test)
	{
	test_mldsa(test, &mldsa87_testvector);
	}

	static s32 mod(s32 a, s32 m)
	{
	a %= m;
	if (a < 0)
	a += m;
	return a;
	}

	static s32 symmetric_mod(s32 a, s32 m)
	{
	a = mod(a, m);
	if (a > m / 2)
	a -= m;
	return a;
	}

	/* Mechanical, inefficient translation of FIPS 204 Algorithm 36, Decompose */
	static void decompose_ref(s32 r, s32 gamma2, s32 r0, s32 r1)
	{
	s32 rplus = mod(r, Q);

	r0 = symmetric_mod(rplus, 2 gamma2);
	if (rplus - *r0 == Q - 1) {
	*r1 = 0;
	r0 = r0 - 1;
	} else {
	r1 = (rplus - r0) / (2 * gamma2);
	}
	}

	/* Mechanical, inefficient translation of FIPS 204 Algorithm 40, UseHint */
	static s32 use_hint_ref(u8 h, s32 r, s32 gamma2)
	{
	s32 m = (Q - 1) / (2 * gamma2);
	s32 r0, r1;

	decompose_ref(r, gamma2, &r0, &r1);
	if (h == 1 && r0 > 0)
	return mod(r1 + 1, m);
	if (h == 1 && r0 <= 0)
	return mod(r1 - 1, m);
	return r1;
	}

	/*
	* Test that for all possible inputs, mldsa_use_hint() gives the same output as
	* a mechanical translation of the pseudocode from FIPS 204.
	*/
	static void test_mldsa_use_hint(struct kunit *test)
	{
	for (int i = 0; i < 2; i++) {
	const s32 gamma2 = (Q - 1) / (i == 0 ? 88 : 32);

	for (u8 h = 0; h < 2; h++) {
	for (s32 r = 0; r < Q; r++) {
	KUNIT_ASSERT_EQ(test,
	mldsa_use_hint(h, r, gamma2),
	use_hint_ref(h, r, gamma2));
	}
	}
	}
	}

	static void benchmark_mldsa(struct kunit *test,
	const struct mldsa_testvector *tv)
	{
	const int warmup_niter = 200;
	const int benchmark_niter = 200;
	u64 t0, t1;

	if (!IS_ENABLED(CONFIG_CRYPTO_LIB_BENCHMARK))
	kunit_skip(test, "not enabled");

	for (int i = 0; i < warmup_niter; i++)
	do_mldsa_and_assert_success(test, tv);

	t0 = ktime_get_ns();
	for (int i = 0; i < benchmark_niter; i++)
	do_mldsa_and_assert_success(test, tv);
	t1 = ktime_get_ns();
	kunit_info(test, "%llu ops/s",
	div64_u64((u64)benchmark_niter * NSEC_PER_SEC,
	t1 - t0 ?: 1));
	}

	static void benchmark_mldsa44(struct kunit *test)
	{
	benchmark_mldsa(test, &mldsa44_testvector);
	}

	static void benchmark_mldsa65(struct kunit *test)
	{
	benchmark_mldsa(test, &mldsa65_testvector);
	}

	static void benchmark_mldsa87(struct kunit *test)
	{
	benchmark_mldsa(test, &mldsa87_testvector);
	}

	static struct kunit_case mldsa_kunit_cases[] = {
	KUNIT_CASE(test_mldsa44),
	KUNIT_CASE(test_mldsa65),
	KUNIT_CASE(test_mldsa87),
	KUNIT_CASE(test_mldsa_use_hint),
	KUNIT_CASE(benchmark_mldsa44),
	KUNIT_CASE(benchmark_mldsa65),
	KUNIT_CASE(benchmark_mldsa87),
	{},
	};

	static struct kunit_suite mldsa_kunit_suite = {
	.name = "mldsa",
	.test_cases = mldsa_kunit_cases,
	};
	kunit_test_suite(mldsa_kunit_suite);

	MODULE_DESCRIPTION("KUnit tests and benchmark for ML-DSA");
	MODULE_IMPORT_NS("EXPORTED_FOR_KUNIT_TESTING");
	MODULE_LICENSE("GPL");