Re: [PATCH v9 2/9] crypto: Add ML-DSA/Dilithium verify support

[Eric Biggers <ebiggers@kernel.org>]

On Mon, Nov 17, 2025 at 02:55:51PM +0000, David Howells wrote:
> +/*
> + * @brief poly_uniform - Sample polynomial with uniformly random coefficients
> + *			 in [0,Q-1] by performing rejection sampling on the
> + *			 output stream of SHAKE128(seed|nonce).
> + *
> + * @param [out] a pointer to output polynomial
> + * @param [in] seed byte array with seed of length DILITHIUM_SEEDBYTES
> + * @param [in] nonce 2-byte nonce
> + */
> +void poly_uniform(poly *a, const uint8_t seed[DILITHIUM_SEEDBYTES],
> +		  __le16 nonce, void *ws_buf)
> +{
> +	struct shake_ctx hash_ctx;
> +	unsigned int i, ctr, off;
> +	unsigned int buflen = POLY_UNIFORM_NBLOCKS * SHAKE128_BLOCK_SIZE;
> +	uint8_t *buf = ws_buf;
> +
> +	shake128_init(&hash_ctx);
> +	shake_update(&hash_ctx, seed, DILITHIUM_SEEDBYTES);
> +	shake_update(&hash_ctx, (uint8_t *)&nonce, sizeof(nonce));
> +	shake_squeeze(&hash_ctx, buf, buflen);
> +
> +	ctr = rej_uniform(a->coeffs, DILITHIUM_N, buf, buflen);
> +
> +	while (ctr < DILITHIUM_N) {
> +		off = buflen % 3;
> +		for (i = 0; i < off; ++i)
> +			buf[i] = buf[buflen - off + i];
> +
> +		shake_squeeze(&hash_ctx, buf + off, SHAKE128_BLOCK_SIZE);
> +		buflen = DILITHIUM_SEEDBYTES + off;
> +		ctr += rej_uniform(a->coeffs + ctr, DILITHIUM_N - ctr, buf,
> +				   buflen);
> +	}
> +
> +	shake_zeroize_ctx(&hash_ctx);
> +}

By the way, the above has a bug.  In the second and later squeezes, it
squeezes SHAKE128_BLOCK_SIZE (168) bytes, but then it uses only the
first DILITHIUM_SEEDBYTES (32) bytes.

Now, that 32 is on top of the 840-byte first squeeze, so there are 872
correct bytes which is enough for 290 samples.  So an incorrect matrix
would be generated only if more than 290 samples happen to be required
to get the 256 coefficients.  q / 2^23 = ~99.9% of coefficients are
accepted, so that number of rejections would be pretty unlikely.

Still, it's a bug.  Anyway, we're not going to use this code (we'll use
my code that does this correctly and in a simpler way), but I thought
I'd point it out so that Stephan can fix it.  This seems to be a
"leancrypto" specific bug.

Note: this feedback should not be taken as implying that I've reviewed
the entire 4800 lines of code.  I just happened to notice this.

- Eric