fallensigh/ciciec2026
Archived
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

mlkem_test: add mlkem-native

Browse Source
This commit is contained in:
FallenSigh
2026-07-05 21:26:19 +08:00
parent b8ddb2ef08
commit cd039e2e2e
126 changed files with 37929 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

367
sdk/software/examples/mlkem_test/mlkem_native/src/sampling.c Normal file
View File

@@ -0,0 +1,367 @@
/*
* 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