/* * Copyright (c) The mlkem-native project authors * SPDX-License-Identifier: Apache-2.0 OR ISC OR MIT */ /* References * ========== * * - [FIPS203] * FIPS 203 Module-Lattice-Based Key-Encapsulation Mechanism Standard * National Institute of Standards and Technology * https://csrc.nist.gov/pubs/fips/203/final * * - [REF] * CRYSTALS-Kyber C reference implementation * Bos, Ducas, Kiltz, Lepoint, Lyubashevsky, Schanck, Schwabe, Seiler, Stehlé * https://github.com/pq-crystals/kyber/tree/main/ref */ #include "common.h" #if !defined(MLK_CONFIG_MULTILEVEL_NO_SHARED) #include "debug.h" #include "sampling.h" #include "symmetric.h" #include "verify.h" /* Reference: `rej_uniform()` in the reference implementation @[REF]. * - Our signature differs from the reference implementation * in that it adds the offset and always expects the base of the * target buffer. This avoids shifting the buffer base in the * caller, which appears tricky to reason about. */ MLK_STATIC_TESTABLE unsigned mlk_rej_uniform_c(int16_t *r, unsigned target, unsigned offset, const uint8_t *buf, unsigned buflen) __contract__( requires(offset <= target && target <= 4096 && buflen <= 4096 && buflen % 3 == 0) requires(memory_no_alias(r, sizeof(int16_t) * target)) requires(memory_no_alias(buf, buflen)) requires(array_bound(r, 0, offset, 0, MLKEM_Q)) assigns(memory_slice(r, sizeof(int16_t) * target)) ensures(offset <= return_value && return_value <= target) ensures(array_bound(r, 0, return_value, 0, MLKEM_Q))) { unsigned ctr, pos; int16_t val0, val1; mlk_assert_bound(r, offset, 0, MLKEM_Q); ctr = offset; pos = 0; /* pos + 3 cannot overflow due to the assumption buflen <= 4096 */ while (ctr < target && pos + 3 <= buflen) __loop__( invariant(offset <= ctr && ctr <= target && pos <= buflen) invariant(array_bound(r, 0, ctr, 0, MLKEM_Q)) decreases(buflen - pos)) { /* Safety: * - The explicit cast to uint16_t ensures that << 8 does * not signed-overflow even on a 16-bit system. * - The conversion to int16_t is safe due to the explicit 0xFFF * truncation. */ val0 = (int16_t)(((buf[pos + 0] >> 0) | ((uint16_t)buf[pos + 1] << 8)) & 0xFFF); val1 = (int16_t)(((buf[pos + 1] >> 4) | (buf[pos + 2] << 4)) & 0xFFF); pos += 3; if (val0 < MLKEM_Q) { r[ctr++] = val0; } if (ctr < target && val1 < MLKEM_Q) { r[ctr++] = val1; } } mlk_assert_bound(r, ctr, 0, MLKEM_Q); return ctr; } /** * Run rejection sampling on uniform random bytes to generate uniform random * integers mod MLKEM_Q. * * @reference{`rej_uniform()` in the reference implementation @[REF]. Our * signature differs from the reference in that it adds the offset and always * expects the base of the target buffer; this avoids shifting the buffer * base in the caller, which is tricky to reason about. Has an optional * fallback to a native implementation.} * * @param[out] r Output buffer. * @param target Requested number of 16-bit integers (uniform mod MLKEM_Q). * Must be <= 4096. * @param offset Number of 16-bit integers that have already been * sampled. Must be <= @p target. * @param[in] buf Input buffer (assumed to be uniform random bytes). * @param buflen Length of input buffer in bytes. Must be <= 4096 and a * multiple of 3. * * @note Strictly speaking, only a few values of @p buflen near UINT_MAX need * excluding. The limit of 4096 is somewhat arbitrary but sufficient * for all uses of this function. Similarly, the actual limit for * @p target is UINT_MAX/2. * * @return New offset of sampled 16-bit integers, at most @p target and at * least the initial @p offset. If the new offset is strictly less * than @p target, the entire input buffer is guaranteed to have been * consumed; otherwise no information is provided on how many bytes * of the input buffer have been consumed. */ static unsigned mlk_rej_uniform(int16_t *r, unsigned target, unsigned offset, const uint8_t *buf, unsigned buflen) __contract__( requires(offset <= target && target <= 4096 && buflen <= 4096 && buflen % 3 == 0) requires(memory_no_alias(r, sizeof(int16_t) * target)) requires(memory_no_alias(buf, buflen)) requires(array_bound(r, 0, offset, 0, MLKEM_Q)) assigns(memory_slice(r, sizeof(int16_t) * target)) ensures(offset <= return_value && return_value <= target) ensures(array_bound(r, 0, return_value, 0, MLKEM_Q)) ) { #if defined(MLK_USE_NATIVE_REJ_UNIFORM) if (offset == 0) { int ret; ret = mlk_rej_uniform_native(r, target, buf, buflen); if (ret != MLK_NATIVE_FUNC_FALLBACK) { unsigned res = (unsigned)ret; mlk_assert_bound(r, res, 0, MLKEM_Q); return res; } } #endif /* MLK_USE_NATIVE_REJ_UNIFORM */ return mlk_rej_uniform_c(r, target, offset, buf, buflen); } #ifndef MLKEM_GEN_MATRIX_NBLOCKS #define MLKEM_GEN_MATRIX_NBLOCKS \ ((12 * MLKEM_N / 8 * ((uint32_t)1 << 12) / MLKEM_Q + MLK_XOF_RATE) / \ MLK_XOF_RATE) #endif #if !defined(MLK_CONFIG_SERIAL_FIPS202_ONLY) /* Reference: Does not exist in the reference implementation @[REF]. * - x4-batched version of `rej_uniform()` from the * reference implementation, leveraging x4-batched Keccak-f1600. */ MLK_INTERNAL_API void mlk_poly_rej_uniform_x4(mlk_poly *vec0, mlk_poly *vec1, mlk_poly *vec2, mlk_poly *vec3, uint8_t seed[4][MLK_ALIGN_UP(MLKEM_SYMBYTES + 2)]) { /* Temporary buffers for XOF output before rejection sampling */ MLK_ALIGN uint8_t buf[4][MLK_ALIGN_UP(MLKEM_GEN_MATRIX_NBLOCKS * MLK_XOF_RATE)]; /* Tracks the number of coefficients we have already sampled */ unsigned ctr[4]; mlk_xof_x4_ctx statex; unsigned buflen; mlk_xof_x4_init(&statex); mlk_xof_x4_absorb(&statex, seed, MLKEM_SYMBYTES + 2); /* * Initially, squeeze heuristic number of MLKEM_GEN_MATRIX_NBLOCKS. * This should generate the matrix entries with high probability. */ mlk_xof_x4_squeezeblocks(buf, MLKEM_GEN_MATRIX_NBLOCKS, &statex); buflen = MLKEM_GEN_MATRIX_NBLOCKS * MLK_XOF_RATE; ctr[0] = mlk_rej_uniform(vec0->coeffs, MLKEM_N, 0, buf[0], buflen); ctr[1] = mlk_rej_uniform(vec1->coeffs, MLKEM_N, 0, buf[1], buflen); ctr[2] = mlk_rej_uniform(vec2->coeffs, MLKEM_N, 0, buf[2], buflen); ctr[3] = mlk_rej_uniform(vec3->coeffs, MLKEM_N, 0, buf[3], buflen); /* * So long as not all matrix entries have been generated, squeeze * one more block a time until we're done. */ buflen = MLK_XOF_RATE; while (ctr[0] < MLKEM_N || ctr[1] < MLKEM_N || ctr[2] < MLKEM_N || ctr[3] < MLKEM_N) __loop__( assigns(ctr, statex, memory_slice(vec0, sizeof(mlk_poly)), memory_slice(vec1, sizeof(mlk_poly)), memory_slice(vec2, sizeof(mlk_poly)), memory_slice(vec3, sizeof(mlk_poly)), object_whole(buf)) invariant(ctr[0] <= MLKEM_N && ctr[1] <= MLKEM_N) invariant(ctr[2] <= MLKEM_N && ctr[3] <= MLKEM_N) invariant(array_bound(vec0->coeffs, 0, ctr[0], 0, MLKEM_Q)) invariant(array_bound(vec1->coeffs, 0, ctr[1], 0, MLKEM_Q)) invariant(array_bound(vec2->coeffs, 0, ctr[2], 0, MLKEM_Q)) invariant(array_bound(vec3->coeffs, 0, ctr[3], 0, MLKEM_Q))) { mlk_xof_x4_squeezeblocks(buf, 1, &statex); ctr[0] = mlk_rej_uniform(vec0->coeffs, MLKEM_N, ctr[0], buf[0], buflen); ctr[1] = mlk_rej_uniform(vec1->coeffs, MLKEM_N, ctr[1], buf[1], buflen); ctr[2] = mlk_rej_uniform(vec2->coeffs, MLKEM_N, ctr[2], buf[2], buflen); ctr[3] = mlk_rej_uniform(vec3->coeffs, MLKEM_N, ctr[3], buf[3], buflen); } mlk_xof_x4_release(&statex); /* Specification: Partially implements * @[FIPS203, Section 3.3, Destruction of intermediate values] */ mlk_zeroize(buf, sizeof(buf)); } #endif /* !MLK_CONFIG_SERIAL_FIPS202_ONLY */ MLK_INTERNAL_API void mlk_poly_rej_uniform(mlk_poly *entry, uint8_t seed[MLKEM_SYMBYTES + 2]) { mlk_xof_ctx state; MLK_ALIGN uint8_t buf[MLKEM_GEN_MATRIX_NBLOCKS * MLK_XOF_RATE]; unsigned ctr, buflen; mlk_xof_init(&state); mlk_xof_absorb(&state, seed, MLKEM_SYMBYTES + 2); /* Initially, squeeze + sample heuristic number of MLKEM_GEN_MATRIX_NBLOCKS. */ /* This should generate the matrix entry with high probability. */ mlk_xof_squeezeblocks(buf, MLKEM_GEN_MATRIX_NBLOCKS, &state); buflen = MLKEM_GEN_MATRIX_NBLOCKS * MLK_XOF_RATE; ctr = mlk_rej_uniform(entry->coeffs, MLKEM_N, 0, buf, buflen); /* Squeeze + sample one more block a time until we're done */ buflen = MLK_XOF_RATE; while (ctr < MLKEM_N) __loop__( assigns(ctr, state, memory_slice(entry, sizeof(mlk_poly)), object_whole(buf)) invariant(ctr <= MLKEM_N) invariant(array_bound(entry->coeffs, 0, ctr, 0, MLKEM_Q))) { mlk_xof_squeezeblocks(buf, 1, &state); ctr = mlk_rej_uniform(entry->coeffs, MLKEM_N, ctr, buf, buflen); } mlk_xof_release(&state); /* Specification: Partially implements * @[FIPS203, Section 3.3, Destruction of intermediate values] */ mlk_zeroize(buf, sizeof(buf)); } /** * Load 4 bytes into a 32-bit integer in little-endian order. * * @reference{`load32_littleendian()` in the reference implementation @[REF].} * * @param[in] x Input byte array. * * @return 32-bit unsigned integer loaded from @p x. */ static uint32_t mlk_load32_littleendian(const uint8_t x[4]) { uint32_t r; r = (uint32_t)x[0]; r |= (uint32_t)x[1] << 8; r |= (uint32_t)x[2] << 16; r |= (uint32_t)x[3] << 24; return r; } /* Reference: `cbd2()` in the reference implementation @[REF]. */ MLK_INTERNAL_API void mlk_poly_cbd2(mlk_poly *r, const uint8_t buf[2 * MLKEM_N / 4]) { unsigned i; for (i = 0; i < MLKEM_N / 8; i++) __loop__( invariant(i <= MLKEM_N / 8) invariant(array_abs_bound(r->coeffs, 0, 8 * i, 3)) decreases(MLKEM_N / 8 - i)) { unsigned j; uint32_t t = mlk_load32_littleendian(buf + 4 * i); uint32_t d = t & 0x55555555; d += (t >> 1) & 0x55555555; for (j = 0; j < 8; j++) __loop__( invariant(i <= MLKEM_N / 8 && j <= 8) invariant(array_abs_bound(r->coeffs, 0, 8 * i + j, 3)) decreases(8 - j)) { /* Safety: The & 0x3 masks each value to 2 bits (range [0, 3]), so the * truncation and subsequent subtraction in int16_t is lossless. */ const int16_t a = (int16_t)((d >> (4 * j + 0)) & 0x3); const int16_t b = (int16_t)((d >> (4 * j + 2)) & 0x3); r->coeffs[8 * i + j] = (int16_t)(a - b); } } } #if defined(MLK_CONFIG_MULTILEVEL_WITH_SHARED) || MLKEM_ETA1 == 3 /** * Load 3 bytes into a 32-bit integer in little-endian order. * * This function is only needed for ML-KEM-512. * * @reference{`load24_littleendian()` in the reference implementation @[REF].} * * @param[in] x Input byte array. * * @return 32-bit unsigned integer loaded from @p x (most significant byte * is zero). */ static uint32_t mlk_load24_littleendian(const uint8_t x[3]) { uint32_t r; r = (uint32_t)x[0]; r |= (uint32_t)x[1] << 8; r |= (uint32_t)x[2] << 16; return r; } /* Reference: `cbd3()` in the reference implementation @[REF]. */ MLK_INTERNAL_API void mlk_poly_cbd3(mlk_poly *r, const uint8_t buf[3 * MLKEM_N / 4]) { unsigned i; for (i = 0; i < MLKEM_N / 4; i++) __loop__( invariant(i <= MLKEM_N / 4) invariant(array_abs_bound(r->coeffs, 0, 4 * i, 4)) decreases(MLKEM_N / 4 - i)) { unsigned j; const uint32_t t = mlk_load24_littleendian(buf + 3 * i); uint32_t d = t & 0x00249249; d += (t >> 1) & 0x00249249; d += (t >> 2) & 0x00249249; for (j = 0; j < 4; j++) __loop__( invariant(i <= MLKEM_N / 4 && j <= 4) invariant(array_abs_bound(r->coeffs, 0, 4 * i + j, 4)) decreases(4 - j)) { /* Safety: The & 0x7 masks each value to 3 bits (range [0, 7]), so the * truncation and subsequent subtraction in int16_t is lossless. */ const int16_t a = (int16_t)((d >> (6 * j + 0)) & 0x7); const int16_t b = (int16_t)((d >> (6 * j + 3)) & 0x7); r->coeffs[4 * i + j] = (int16_t)(a - b); } } } #endif /* MLK_CONFIG_MULTILEVEL_WITH_SHARED || MLKEM_ETA1 == 3 */ #else /* !MLK_CONFIG_MULTILEVEL_NO_SHARED */ MLK_EMPTY_CU(sampling) #endif /* MLK_CONFIG_MULTILEVEL_NO_SHARED */ /* To facilitate single-compilation-unit (SCU) builds, undefine all macros. * Don't modify by hand -- this is auto-generated by scripts/autogen. */ #undef MLKEM_GEN_MATRIX_NBLOCKS