xref: /linux/lib/crypto/tests/mldsa_kunit.c (revision f4cdf7ca9a1fdcca413157df19753f388a5a224e)
1 // SPDX-License-Identifier: GPL-2.0-or-later
2 /*
3  * KUnit tests and benchmark for ML-DSA
4  *
5  * Copyright 2025 Google LLC
6  */
7 #include <crypto/mldsa.h>
8 #include <kunit/test.h>
9 #include <linux/random.h>
10 #include <linux/unaligned.h>
11 #include "test-utils.h"
12 
13 #define Q 8380417 /* The prime q = 2^23 - 2^13 + 1 */
14 
15 /* ML-DSA parameters that the tests use */
16 static const struct {
17 	int sig_len;
18 	int pk_len;
19 	int k;
20 	int lambda;
21 	int gamma1;
22 	int beta;
23 	int omega;
24 } params[] = {
25 	[MLDSA44] = {
26 		.sig_len = MLDSA44_SIGNATURE_SIZE,
27 		.pk_len = MLDSA44_PUBLIC_KEY_SIZE,
28 		.k = 4,
29 		.lambda = 128,
30 		.gamma1 = 1 << 17,
31 		.beta = 78,
32 		.omega = 80,
33 	},
34 	[MLDSA65] = {
35 		.sig_len = MLDSA65_SIGNATURE_SIZE,
36 		.pk_len = MLDSA65_PUBLIC_KEY_SIZE,
37 		.k = 6,
38 		.lambda = 192,
39 		.gamma1 = 1 << 19,
40 		.beta = 196,
41 		.omega = 55,
42 	},
43 	[MLDSA87] = {
44 		.sig_len = MLDSA87_SIGNATURE_SIZE,
45 		.pk_len = MLDSA87_PUBLIC_KEY_SIZE,
46 		.k = 8,
47 		.lambda = 256,
48 		.gamma1 = 1 << 19,
49 		.beta = 120,
50 		.omega = 75,
51 	},
52 };
53 
54 #include "mldsa-testvecs.h"
55 
56 static void do_mldsa_and_assert_success(struct kunit *test,
57 					const struct mldsa_testvector *tv)
58 {
59 	int err = mldsa_verify(tv->alg, tv->sig, tv->sig_len, tv->msg,
60 			       tv->msg_len, tv->pk, tv->pk_len);
61 	KUNIT_ASSERT_EQ(test, err, 0);
62 }
63 
64 /*
65  * Test that changing coefficients in a valid signature's z vector results in
66  * the following behavior from mldsa_verify():
67  *
68  *  * -EBADMSG if a coefficient is changed to have an out-of-range value, i.e.
69  *    absolute value >= gamma1 - beta, corresponding to the verifier detecting
70  *    the out-of-range coefficient and rejecting the signature as malformed
71  *
72  *  * -EKEYREJECTED if a coefficient is changed to a different in-range value,
73  *    i.e. absolute value < gamma1 - beta, corresponding to the verifier
74  *    continuing to the "real" signature check and that check failing
75  */
76 static void test_mldsa_z_range(struct kunit *test,
77 			       const struct mldsa_testvector *tv)
78 {
79 	u8 *sig = memdup_buf(test, tv->sig, tv->sig_len);
80 	const int lambda = params[tv->alg].lambda;
81 	const s32 gamma1 = params[tv->alg].gamma1;
82 	const int beta = params[tv->alg].beta;
83 	/*
84 	 * We just modify the first coefficient.  The coefficient is gamma1
85 	 * minus either the first 18 or 20 bits of the u32, depending on gamma1.
86 	 *
87 	 * The layout of ML-DSA signatures is ctilde || z || h.  ctilde is
88 	 * lambda / 4 bytes, so z starts at &sig[lambda / 4].
89 	 */
90 	u8 *z_ptr = &sig[lambda / 4];
91 	const u32 z_data = get_unaligned_le32(z_ptr);
92 	const u32 mask = (gamma1 << 1) - 1;
93 	/* These are the four boundaries of the out-of-range values. */
94 	const s32 out_of_range_coeffs[] = {
95 		-gamma1 + 1,
96 		-(gamma1 - beta),
97 		gamma1,
98 		gamma1 - beta,
99 	};
100 	/*
101 	 * These are the two boundaries of the valid range, along with 0.  We
102 	 * assume that none of these matches the original coefficient.
103 	 */
104 	const s32 in_range_coeffs[] = {
105 		-(gamma1 - beta - 1),
106 		0,
107 		gamma1 - beta - 1,
108 	};
109 
110 	/* Initially the signature is valid. */
111 	do_mldsa_and_assert_success(test, tv);
112 
113 	/* Test some out-of-range coefficients. */
114 	for (int i = 0; i < ARRAY_SIZE(out_of_range_coeffs); i++) {
115 		const s32 c = out_of_range_coeffs[i];
116 
117 		put_unaligned_le32((z_data & ~mask) | (mask & (gamma1 - c)),
118 				   z_ptr);
119 		KUNIT_ASSERT_EQ(test, -EBADMSG,
120 				mldsa_verify(tv->alg, sig, tv->sig_len, tv->msg,
121 					     tv->msg_len, tv->pk, tv->pk_len));
122 	}
123 
124 	/* Test some in-range coefficients. */
125 	for (int i = 0; i < ARRAY_SIZE(in_range_coeffs); i++) {
126 		const s32 c = in_range_coeffs[i];
127 
128 		put_unaligned_le32((z_data & ~mask) | (mask & (gamma1 - c)),
129 				   z_ptr);
130 		KUNIT_ASSERT_EQ(test, -EKEYREJECTED,
131 				mldsa_verify(tv->alg, sig, tv->sig_len, tv->msg,
132 					     tv->msg_len, tv->pk, tv->pk_len));
133 	}
134 }
135 
136 /* Test that mldsa_verify() rejects malformed hint vectors with -EBADMSG. */
137 static void test_mldsa_bad_hints(struct kunit *test,
138 				 const struct mldsa_testvector *tv)
139 {
140 	const int omega = params[tv->alg].omega;
141 	const int k = params[tv->alg].k;
142 	u8 *sig = memdup_buf(test, tv->sig, tv->sig_len);
143 	/* Pointer to the encoded hint vector in the signature */
144 	u8 *hintvec = &sig[tv->sig_len - omega - k];
145 	u8 h;
146 
147 	/* Initially the signature is valid. */
148 	do_mldsa_and_assert_success(test, tv);
149 
150 	/* Cumulative hint count exceeds omega */
151 	memcpy(sig, tv->sig, tv->sig_len);
152 	hintvec[omega + k - 1] = omega + 1;
153 	KUNIT_ASSERT_EQ(test, -EBADMSG,
154 			mldsa_verify(tv->alg, sig, tv->sig_len, tv->msg,
155 				     tv->msg_len, tv->pk, tv->pk_len));
156 
157 	/* Cumulative hint count decreases */
158 	memcpy(sig, tv->sig, tv->sig_len);
159 	KUNIT_ASSERT_GE(test, hintvec[omega + k - 2], 1);
160 	hintvec[omega + k - 1] = hintvec[omega + k - 2] - 1;
161 	KUNIT_ASSERT_EQ(test, -EBADMSG,
162 			mldsa_verify(tv->alg, sig, tv->sig_len, tv->msg,
163 				     tv->msg_len, tv->pk, tv->pk_len));
164 
165 	/*
166 	 * Hint indices out of order.  To test this, swap hintvec[0] and
167 	 * hintvec[1].  This assumes that the original valid signature had at
168 	 * least two nonzero hints in the first element (asserted below).
169 	 */
170 	memcpy(sig, tv->sig, tv->sig_len);
171 	KUNIT_ASSERT_GE(test, hintvec[omega], 2);
172 	h = hintvec[0];
173 	hintvec[0] = hintvec[1];
174 	hintvec[1] = h;
175 	KUNIT_ASSERT_EQ(test, -EBADMSG,
176 			mldsa_verify(tv->alg, sig, tv->sig_len, tv->msg,
177 				     tv->msg_len, tv->pk, tv->pk_len));
178 
179 	/*
180 	 * Extra hint indices given.  For this test to work, the original valid
181 	 * signature must have fewer than omega nonzero hints (asserted below).
182 	 */
183 	memcpy(sig, tv->sig, tv->sig_len);
184 	KUNIT_ASSERT_LT(test, hintvec[omega + k - 1], omega);
185 	hintvec[omega - 1] = 0xff;
186 	KUNIT_ASSERT_EQ(test, -EBADMSG,
187 			mldsa_verify(tv->alg, sig, tv->sig_len, tv->msg,
188 				     tv->msg_len, tv->pk, tv->pk_len));
189 }
190 
191 static void test_mldsa_mutation(struct kunit *test,
192 				const struct mldsa_testvector *tv)
193 {
194 	const int sig_len = tv->sig_len;
195 	const int msg_len = tv->msg_len;
196 	const int pk_len = tv->pk_len;
197 	const int num_iter = 200;
198 	u8 *sig = memdup_buf(test, tv->sig, sig_len);
199 	u8 *msg = memdup_buf(test, tv->msg, msg_len);
200 	u8 *pk = memdup_buf(test, tv->pk, pk_len);
201 
202 	/* Initially the signature is valid. */
203 	do_mldsa_and_assert_success(test, tv);
204 
205 	/* Changing any bit in the signature should invalidate the signature */
206 	for (int i = 0; i < num_iter; i++) {
207 		size_t pos = get_random_u32_below(sig_len);
208 		u8 b = 1 << get_random_u32_below(8);
209 
210 		sig[pos] ^= b;
211 		KUNIT_ASSERT_NE(test, 0,
212 				mldsa_verify(tv->alg, sig, sig_len, msg,
213 					     msg_len, pk, pk_len));
214 		sig[pos] ^= b;
215 	}
216 
217 	/* Changing any bit in the message should invalidate the signature */
218 	for (int i = 0; i < num_iter; i++) {
219 		size_t pos = get_random_u32_below(msg_len);
220 		u8 b = 1 << get_random_u32_below(8);
221 
222 		msg[pos] ^= b;
223 		KUNIT_ASSERT_NE(test, 0,
224 				mldsa_verify(tv->alg, sig, sig_len, msg,
225 					     msg_len, pk, pk_len));
226 		msg[pos] ^= b;
227 	}
228 
229 	/* Changing any bit in the public key should invalidate the signature */
230 	for (int i = 0; i < num_iter; i++) {
231 		size_t pos = get_random_u32_below(pk_len);
232 		u8 b = 1 << get_random_u32_below(8);
233 
234 		pk[pos] ^= b;
235 		KUNIT_ASSERT_NE(test, 0,
236 				mldsa_verify(tv->alg, sig, sig_len, msg,
237 					     msg_len, pk, pk_len));
238 		pk[pos] ^= b;
239 	}
240 
241 	/* All changes should have been undone. */
242 	KUNIT_ASSERT_EQ(test, 0,
243 			mldsa_verify(tv->alg, sig, sig_len, msg, msg_len, pk,
244 				     pk_len));
245 }
246 
247 static void test_mldsa(struct kunit *test, const struct mldsa_testvector *tv)
248 {
249 	/* Valid signature */
250 	KUNIT_ASSERT_EQ(test, tv->sig_len, params[tv->alg].sig_len);
251 	KUNIT_ASSERT_EQ(test, tv->pk_len, params[tv->alg].pk_len);
252 	do_mldsa_and_assert_success(test, tv);
253 
254 	/* Signature too short */
255 	KUNIT_ASSERT_EQ(test, -EBADMSG,
256 			mldsa_verify(tv->alg, tv->sig, tv->sig_len - 1, tv->msg,
257 				     tv->msg_len, tv->pk, tv->pk_len));
258 
259 	/* Signature too long */
260 	KUNIT_ASSERT_EQ(test, -EBADMSG,
261 			mldsa_verify(tv->alg, tv->sig, tv->sig_len + 1, tv->msg,
262 				     tv->msg_len, tv->pk, tv->pk_len));
263 
264 	/* Public key too short */
265 	KUNIT_ASSERT_EQ(test, -EBADMSG,
266 			mldsa_verify(tv->alg, tv->sig, tv->sig_len, tv->msg,
267 				     tv->msg_len, tv->pk, tv->pk_len - 1));
268 
269 	/* Public key too long */
270 	KUNIT_ASSERT_EQ(test, -EBADMSG,
271 			mldsa_verify(tv->alg, tv->sig, tv->sig_len, tv->msg,
272 				     tv->msg_len, tv->pk, tv->pk_len + 1));
273 
274 	/*
275 	 * Message too short.  Error is EKEYREJECTED because it gets rejected by
276 	 * the "real" signature check rather than the well-formedness checks.
277 	 */
278 	KUNIT_ASSERT_EQ(test, -EKEYREJECTED,
279 			mldsa_verify(tv->alg, tv->sig, tv->sig_len, tv->msg,
280 				     tv->msg_len - 1, tv->pk, tv->pk_len));
281 	/*
282 	 * Can't simply try (tv->msg, tv->msg_len + 1) too, as tv->msg would be
283 	 * accessed out of bounds.  However, ML-DSA just hashes the message and
284 	 * doesn't handle different message lengths differently anyway.
285 	 */
286 
287 	/* Test the validity checks on the z vector. */
288 	test_mldsa_z_range(test, tv);
289 
290 	/* Test the validity checks on the hint vector. */
291 	test_mldsa_bad_hints(test, tv);
292 
293 	/* Test randomly mutating the inputs. */
294 	test_mldsa_mutation(test, tv);
295 }
296 
297 static void test_mldsa44(struct kunit *test)
298 {
299 	test_mldsa(test, &mldsa44_testvector);
300 }
301 
302 static void test_mldsa65(struct kunit *test)
303 {
304 	test_mldsa(test, &mldsa65_testvector);
305 }
306 
307 static void test_mldsa87(struct kunit *test)
308 {
309 	test_mldsa(test, &mldsa87_testvector);
310 }
311 
312 static s32 mod(s32 a, s32 m)
313 {
314 	a %= m;
315 	if (a < 0)
316 		a += m;
317 	return a;
318 }
319 
320 static s32 symmetric_mod(s32 a, s32 m)
321 {
322 	a = mod(a, m);
323 	if (a > m / 2)
324 		a -= m;
325 	return a;
326 }
327 
328 /* Mechanical, inefficient translation of FIPS 204 Algorithm 36, Decompose */
329 static void decompose_ref(s32 r, s32 gamma2, s32 *r0, s32 *r1)
330 {
331 	s32 rplus = mod(r, Q);
332 
333 	*r0 = symmetric_mod(rplus, 2 * gamma2);
334 	if (rplus - *r0 == Q - 1) {
335 		*r1 = 0;
336 		*r0 = *r0 - 1;
337 	} else {
338 		*r1 = (rplus - *r0) / (2 * gamma2);
339 	}
340 }
341 
342 /* Mechanical, inefficient translation of FIPS 204 Algorithm 40, UseHint */
343 static s32 use_hint_ref(u8 h, s32 r, s32 gamma2)
344 {
345 	s32 m = (Q - 1) / (2 * gamma2);
346 	s32 r0, r1;
347 
348 	decompose_ref(r, gamma2, &r0, &r1);
349 	if (h == 1 && r0 > 0)
350 		return mod(r1 + 1, m);
351 	if (h == 1 && r0 <= 0)
352 		return mod(r1 - 1, m);
353 	return r1;
354 }
355 
356 /*
357  * Test that for all possible inputs, mldsa_use_hint() gives the same output as
358  * a mechanical translation of the pseudocode from FIPS 204.
359  */
360 static void test_mldsa_use_hint(struct kunit *test)
361 {
362 	for (int i = 0; i < 2; i++) {
363 		const s32 gamma2 = (Q - 1) / (i == 0 ? 88 : 32);
364 
365 		for (u8 h = 0; h < 2; h++) {
366 			for (s32 r = 0; r < Q; r++) {
367 				KUNIT_ASSERT_EQ(test,
368 						mldsa_use_hint(h, r, gamma2),
369 						use_hint_ref(h, r, gamma2));
370 			}
371 		}
372 	}
373 }
374 
375 static void benchmark_mldsa(struct kunit *test,
376 			    const struct mldsa_testvector *tv)
377 {
378 	const int warmup_niter = 200;
379 	const int benchmark_niter = 200;
380 	u64 t0, t1;
381 
382 	if (!IS_ENABLED(CONFIG_CRYPTO_LIB_BENCHMARK))
383 		kunit_skip(test, "not enabled");
384 
385 	for (int i = 0; i < warmup_niter; i++)
386 		do_mldsa_and_assert_success(test, tv);
387 
388 	t0 = ktime_get_ns();
389 	for (int i = 0; i < benchmark_niter; i++)
390 		do_mldsa_and_assert_success(test, tv);
391 	t1 = ktime_get_ns();
392 	kunit_info(test, "%llu ops/s",
393 		   div64_u64((u64)benchmark_niter * NSEC_PER_SEC,
394 			     t1 - t0 ?: 1));
395 }
396 
397 static void benchmark_mldsa44(struct kunit *test)
398 {
399 	benchmark_mldsa(test, &mldsa44_testvector);
400 }
401 
402 static void benchmark_mldsa65(struct kunit *test)
403 {
404 	benchmark_mldsa(test, &mldsa65_testvector);
405 }
406 
407 static void benchmark_mldsa87(struct kunit *test)
408 {
409 	benchmark_mldsa(test, &mldsa87_testvector);
410 }
411 
412 static struct kunit_case mldsa_kunit_cases[] = {
413 	KUNIT_CASE(test_mldsa44),
414 	KUNIT_CASE(test_mldsa65),
415 	KUNIT_CASE(test_mldsa87),
416 	KUNIT_CASE(test_mldsa_use_hint),
417 	KUNIT_CASE(benchmark_mldsa44),
418 	KUNIT_CASE(benchmark_mldsa65),
419 	KUNIT_CASE(benchmark_mldsa87),
420 	{},
421 };
422 
423 static struct kunit_suite mldsa_kunit_suite = {
424 	.name = "mldsa",
425 	.test_cases = mldsa_kunit_cases,
426 };
427 kunit_test_suite(mldsa_kunit_suite);
428 
429 MODULE_DESCRIPTION("KUnit tests and benchmark for ML-DSA");
430 MODULE_IMPORT_NS("EXPORTED_FOR_KUNIT_TESTING");
431 MODULE_LICENSE("GPL");
432