From cee46b02f7bd715e9d7b7c2764a7f747a8191576 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tobias=20Frauenschl=C3=A4ger?= Date: Sun, 6 Sep 2026 14:16:49 +0200 Subject: [PATCH 01/11] Reduce ML-KEM small memory heap footprint Key generation and encapsulation held whole vectors that were only ever used a polynomial at a time. mlkem_gen_matrix_i() generated a full vector of matrix A before any of it was multiplied in. It is replaced by mlkem_gen_matrix_i_acc(), which multiplies each polynomial into the result as soon as it has been sampled and reuses the one polynomial for the next. Key generation now needs a single polynomial of temporary memory for every parameter set. Each row of the public key is now accumulated in place, so a hash failure partway through a row would leave A_hat o s_hat without its noise in the key, which gives away s. mlkem_keygen_seeds() clears the public vector when it fails. The hash-fault white-box drives the new generator under the same condition the library compiles it with. Encapsulation additionally held all of u until the cipher text was encoded, even though u[i] is complete before u[i+1] starts. Each polynomial is now compressed into the cipher text as it is calculated, so v reuses u's polynomial, the message polynomial reuses y, and e_2 reuses the matrix scratch. Temporary memory drops from 3k to k + 2 polynomials. Decapsulation kept a second copy of the whole cipher text to compare against. mlkem_encapsulate_seeds() takes a buffer to compare with instead and checks each encoded block, constant time, as it is produced. mlkem_vec_compress_11() gains the polynomial count parameter that mlkem_vec_compress_10() already had. mlkem_pointwise_acc_mont() is now guarded as small memory key generation no longer calls it and it would be unused when encapsulation and decapsulation are both disabled. Cipher texts and shared secrets are unchanged. Peak heap for ML-KEM-1024 falls from 2048 to 512 bytes for key generation, 6144 to 3072 for encapsulation and 7712 to 3424 for decapsulation, with no change in operation count. --- .../test_mlkem_poly_hash_fault_whitebox.c | 57 ++-- wolfcrypt/src/port/arm/thumb2-mlkem-asm.S | 4 +- wolfcrypt/src/port/arm/thumb2-mlkem-asm_c.c | 20 +- wolfcrypt/src/wc_mlkem.c | 124 ++++++--- wolfcrypt/src/wc_mlkem_poly.c | 260 +++++++++++++----- wolfssl/wolfcrypt/wc_mlkem.h | 11 +- 6 files changed, 333 insertions(+), 143 deletions(-) diff --git a/tests/unit-mcdc/test_mlkem_poly_hash_fault_whitebox.c b/tests/unit-mcdc/test_mlkem_poly_hash_fault_whitebox.c index ebf9449f0a3..a2bad25c297 100644 --- a/tests/unit-mcdc/test_mlkem_poly_hash_fault_whitebox.c +++ b/tests/unit-mcdc/test_mlkem_poly_hash_fault_whitebox.c @@ -29,7 +29,8 @@ * mlkem_hash512() if ((ret == 0) && (data2 != NULL) && (data2Len > 0)) * mlkem_gen_matrix_c() for (i = 0; (ret == 0) && (i < k); i++, ...) * for (j = 0; (ret == 0) && (j < k); j++) - * mlkem_gen_matrix_i() for (j = 0; (ret == 0) && (j < k); j++) + * mlkem_gen_matrix_i_acc() + * for (j = 0; (ret == 0) && (j < k); j++) * mlkem_get_noise_c() for (i = 0; (ret == 0) && (i < k); i++) * if ((ret == 0) && (vec2 != NULL)) * for (i = 0; (ret == 0) && (i < k); i++) @@ -348,20 +349,27 @@ static void wb_gen_matrix_c_continuation(void) #endif /* ------------------------------------------------------------------------- * - * mlkem_gen_matrix_i(): the small-memory single-row generator. Only compiled - * when one of the small-mem arms is on. + * mlkem_gen_matrix_i_acc(): the small-memory single-row generator, which + * multiplies each polynomial into the result as it is generated. Only + * compiled when one of the small-mem arms is on. * ------------------------------------------------------------------------- */ #if defined(WOLFSSL_HAVE_MLKEM) && \ !(defined(WOLFSSL_ARMASM) && defined(__aarch64__)) && \ - (defined(WOLFSSL_MLKEM_MAKEKEY_SMALL_MEM) || \ - defined(WOLFSSL_MLKEM_ENCAPSULATE_SMALL_MEM)) + ((defined(WOLFSSL_MLKEM_MAKEKEY_SMALL_MEM) && \ + !defined(WOLFSSL_MLKEM_NO_MAKE_KEY)) || \ + (defined(WOLFSSL_MLKEM_ENCAPSULATE_SMALL_MEM) && \ + (!defined(WOLFSSL_MLKEM_NO_ENCAPSULATE) || \ + !defined(WOLFSSL_MLKEM_NO_DECAPSULATE)))) static void wb_gen_matrix_i_rows(void) { MLKEM_PRF_T prf; static sword16 a[WB_K * MLKEM_N]; + static sword16 v[WB_K * MLKEM_N]; + static sword16 r[MLKEM_N]; byte seed[WC_ML_KEM_SYM_SZ + 2]; long n; + unsigned int j; if (mlkem_prf_new(&prf, NULL, INVALID_DEVID) != 0) { WB_NOTE("mlkem_prf_new failed"); @@ -369,32 +377,37 @@ static void wb_gen_matrix_i_rows(void) return; } XMEMSET(seed, 0x5c, sizeof(seed)); + for (j = 0; j < WB_K * MLKEM_N; j++) { + v[j] = (sword16)(j & 0x7); + } mcdc_fh_disarm(); - if (mlkem_gen_matrix_i(&prf, a, WB_K, seed, 0, 0) != 0) { - WB_NOTE("baseline mlkem_gen_matrix_i failed"); + if (mlkem_gen_matrix_i_acc(&prf, r, a, v, WB_K, seed, 0, 0) != 0) { + WB_NOTE("baseline mlkem_gen_matrix_i_acc failed"); wb_fail = 1; } for (n = 1; n <= (long)WB_SWEEP; n++) { mcdc_fh_arm(n); - (void)mlkem_gen_matrix_i(&prf, a, WB_K, seed, 0, 0); + (void)mlkem_gen_matrix_i_acc(&prf, r, a, v, WB_K, seed, 0, 0); mcdc_fh_disarm(); mcdc_fh_arm(n); - (void)mlkem_gen_matrix_i(&prf, a, WB_K, seed, 1, 1); + (void)mlkem_gen_matrix_i_acc(&prf, r, a, v, WB_K, seed, 1, 1); mcdc_fh_disarm(); } mlkem_prf_free(&prf); - WB_NOTE("mlkem_gen_matrix_i success-chain rows exercised"); + WB_NOTE("mlkem_gen_matrix_i_acc success-chain rows exercised"); } -/* mlkem_gen_matrix_i(): a failing continuation squeeze must be returned, not - * ignored with the stale buffer sampled again. */ +/* mlkem_gen_matrix_i_acc(): a failing continuation squeeze must be returned, + * not ignored with the stale buffer sampled again. */ static void wb_gen_matrix_i_continuation(void) { MLKEM_PRF_T prf; static sword16 a[MLKEM_N]; + static sword16 v[MLKEM_N]; + static sword16 r[MLKEM_N]; byte seed[WC_ML_KEM_SYM_SZ]; word32 s; int i; @@ -411,8 +424,8 @@ static void wb_gen_matrix_i_continuation(void) seed[0] = (byte)s; seed[1] = (byte)(s >> 8); mcdc_fh_disarm(); - if (mlkem_gen_matrix_i(&prf, a, 1, seed, 0, 0) != 0) { - WB_NOTE("mlkem_gen_matrix_i seed search failed"); + if (mlkem_gen_matrix_i_acc(&prf, r, a, v, 1, seed, 0, 0) != 0) { + WB_NOTE("mlkem_gen_matrix_i_acc seed search failed"); wb_fail = 1; break; } @@ -429,32 +442,36 @@ static void wb_gen_matrix_i_continuation(void) a[i] = -1; } mcdc_fh_arm(WB_XOF_CALLS + 1); - ret = mlkem_gen_matrix_i(&prf, a, 1, seed, 0, 0); + ret = mlkem_gen_matrix_i_acc(&prf, r, a, v, 1, seed, 0, 0); mcdc_fh_disarm(); if (ret != MCDC_FH_ERR) { - WB_NOTE("mlkem_gen_matrix_i ignored a continuation squeeze error"); + WB_NOTE("mlkem_gen_matrix_i_acc ignored a continuation squeeze " + "error"); wb_fail = 1; } if (!WB_POLY_INCOMPLETE(a)) { - WB_NOTE("mlkem_gen_matrix_i sampled after a failed squeeze"); + WB_NOTE("mlkem_gen_matrix_i_acc sampled after a failed squeeze"); wb_fail = 1; } } mlkem_prf_free(&prf); - WB_NOTE("mlkem_gen_matrix_i continuation-squeeze error row exercised"); + WB_NOTE("mlkem_gen_matrix_i_acc continuation-squeeze error row " + "exercised"); } #else static void wb_gen_matrix_i_rows(void) { - WB_NOTE("mlkem_gen_matrix_i arm not compiled in this variant; skipped"); + WB_NOTE("mlkem_gen_matrix_i_acc arm not compiled in this variant; " + "skipped"); } static void wb_gen_matrix_i_continuation(void) { - WB_NOTE("mlkem_gen_matrix_i arm not compiled in this variant; skipped"); + WB_NOTE("mlkem_gen_matrix_i_acc arm not compiled in this variant; " + "skipped"); } #endif diff --git a/wolfcrypt/src/port/arm/thumb2-mlkem-asm.S b/wolfcrypt/src/port/arm/thumb2-mlkem-asm.S index a32e66c9bd0..bef67429ac6 100644 --- a/wolfcrypt/src/port/arm/thumb2-mlkem-asm.S +++ b/wolfcrypt/src/port/arm/thumb2-mlkem-asm.S @@ -3706,7 +3706,7 @@ L_mlkem_thumb2_vec_compress_10_loop: .type mlkem_thumb2_vec_compress_11, %function mlkem_thumb2_vec_compress_11: PUSH {r4, r5, r6, r7, r8, r9, r10, r11, lr} - ADD r2, r1, #0x800 + ADD r2, r1, r2, LSL #9 #ifndef WOLFSSL_ARM_ARCH_7M MOV r8, #0xdbb6 MOVT r8, #0x9d7 @@ -4362,7 +4362,7 @@ L_mlkem_thumb2_vec_compress_11_loop: #endif #endif /* !WOLFSSL_ARM_ARCH_7M */ POP {r4, r5, r6, r7, r8, r9, r10, r11, pc} - /* Cycle Count = 807 */ + /* Cycle Count = 808 */ .size mlkem_thumb2_vec_compress_11,.-mlkem_thumb2_vec_compress_11 .text .align 4 diff --git a/wolfcrypt/src/port/arm/thumb2-mlkem-asm_c.c b/wolfcrypt/src/port/arm/thumb2-mlkem-asm_c.c index b3342b9b9fd..b8bdb026246 100644 --- a/wolfcrypt/src/port/arm/thumb2-mlkem-asm_c.c +++ b/wolfcrypt/src/port/arm/thumb2-mlkem-asm_c.c @@ -3989,16 +3989,19 @@ WC_OMIT_FRAME_POINTER void mlkem_thumb2_vec_compress_10(byte* r, sword16* v, } #ifndef WOLFSSL_NO_VAR_ASSIGN_REG -WC_OMIT_FRAME_POINTER void mlkem_thumb2_vec_compress_11(byte* r_p, sword16* v_p) +WC_OMIT_FRAME_POINTER void mlkem_thumb2_vec_compress_11(byte* r_p, sword16* v_p, + unsigned int k_p) #else -WC_OMIT_FRAME_POINTER void mlkem_thumb2_vec_compress_11(byte* r, sword16* v) +WC_OMIT_FRAME_POINTER void mlkem_thumb2_vec_compress_11(byte* r, sword16* v, + unsigned int k) #endif /* !WOLFSSL_NO_VAR_ASSIGN_REG */ { #ifndef WOLFSSL_NO_VAR_ASSIGN_REG register byte* r __asm__ ("r0") = (byte*)r_p; register sword16* v __asm__ ("r1") = (sword16*)v_p; + register unsigned int k __asm__ ("r2") = (unsigned int)k_p; #else - void* L_asm_args[2] = {(void*)(size_t)r, (void*)(size_t)v + void* L_asm_args[3] = {(void*)(size_t)r, (void*)(size_t)v, (void*)(size_t)k }; void** L_asm_args_p = L_asm_args; #endif /* !WOLFSSL_NO_VAR_ASSIGN_REG */ @@ -4006,9 +4009,9 @@ WC_OMIT_FRAME_POINTER void mlkem_thumb2_vec_compress_11(byte* r, sword16* v) __asm__ __volatile__ ( #ifdef WOLFSSL_NO_VAR_ASSIGN_REG "PUSH {%[L_asm_args]}\n\t" - "LDM %[L_asm_args], {r0, r1}\n\t" + "LDM %[L_asm_args], {r0, r1, r2}\n\t" #endif /* WOLFSSL_NO_VAR_ASSIGN_REG */ - "ADD r2, r1, #0x800\n\t" + "ADD r2, r1, r2, LSL #9\n\t" #ifndef WOLFSSL_ARM_ARCH_7M "MOV r8, #0xdbb6\n\t" "MOVT r8, #0x9d7\n\t" @@ -4679,12 +4682,12 @@ WC_OMIT_FRAME_POINTER void mlkem_thumb2_vec_compress_11(byte* r, sword16* v) #endif /* !WOLFSSL_ARM_ARCH_7M */ #ifdef WOLFSSL_NO_VAR_ASSIGN_REG "POP {%[L_asm_args]}\n\t" - "STM %[L_asm_args], {r0, r1}\n\t" + "STM %[L_asm_args], {r0, r1, r2}\n\t" #endif /* WOLFSSL_NO_VAR_ASSIGN_REG */ #ifndef WOLFSSL_NO_VAR_ASSIGN_REG - : [r] "+r" (r), [v] "+r" (v) + : [r] "+r" (r), [v] "+r" (v), [k] "+r" (k) : - : "memory", "cc", "r2", "r3", "r4", "r5", "r6", "r7", "r8", "r9", "r10", + : "memory", "cc", "r3", "r4", "r5", "r6", "r7", "r8", "r9", "r10", "r11" #else : [L_asm_args] "+r" (L_asm_args_p) @@ -4696,6 +4699,7 @@ WC_OMIT_FRAME_POINTER void mlkem_thumb2_vec_compress_11(byte* r, sword16* v) #ifdef WOLFSSL_NO_VAR_ASSIGN_REG r = (byte*)(size_t)L_asm_args[0]; v = (sword16*)(size_t)L_asm_args[1]; + k = (unsigned int)(size_t)L_asm_args[2]; #endif /* WOLFSSL_NO_VAR_ASSIGN_REG */ } diff --git a/wolfcrypt/src/wc_mlkem.c b/wolfcrypt/src/wc_mlkem.c index bb18473259f..40a5f6fe122 100644 --- a/wolfcrypt/src/wc_mlkem.c +++ b/wolfcrypt/src/wc_mlkem.c @@ -32,12 +32,14 @@ * * WOLFSSL_MLKEM_MAKEKEY_SMALL_MEM Default: OFF * Uses less dynamic memory to perform key generation. + * Matrix A is generated a polynomial at a time. * Has a small performance trade-off. * Only usable with C implementation. * * WOLFSSL_MLKEM_ENCAPSULATE_SMALL_MEM Default: OFF * Uses less dynamic memory to perform encapsulation. * Affects decapsulation too as encapsulation called. + * Matrix A and u are handled a polynomial at a time. * Has a small performance trade-off. * Only usable with C implementation. * @@ -146,6 +148,12 @@ #error "Cannot use dynamic key buffers without malloc" #endif +#ifdef WOLFSSL_WC_ML_KEM_1024 + #define MLKEM_MAX_COMP_POLY_SZ MLKEM_POLY_COMPRESSED_SZ(11) +#else + #define MLKEM_MAX_COMP_POLY_SZ MLKEM_POLY_COMPRESSED_SZ(10) +#endif + #if FIPS_VERSION3_GE(7,0,0) const unsigned int wolfCrypt_FIPS_mlkem_ro_sanity[2] = { 0x1a2b3c4d, 0x00000019 }; @@ -906,7 +914,8 @@ int wc_MlKemKey_MakeKeyWithRandom(MlKemKey* key, const unsigned char* rand, sword16 e[WC_ML_KEM_MAX_K * MLKEM_N]; #endif #else - sword16 e[WC_ML_KEM_MAX_K * MLKEM_N]; + /* Small memory generation only holds one polynomial at a time. */ + sword16 e[MLKEM_N]; #endif #endif #ifndef WOLFSSL_MLKEM_MAKEKEY_SMALL_MEM @@ -914,6 +923,8 @@ int wc_MlKemKey_MakeKeyWithRandom(MlKemKey* key, const unsigned char* rand, #endif sword16* s = NULL; sword16* t = NULL; + /* Number of bytes of e holding secret noise, to be zeroized. */ + size_t eSz = 0; int ret = 0; int k = 0; @@ -947,7 +958,14 @@ int wc_MlKemKey_MakeKeyWithRandom(MlKemKey* key, const unsigned char* rand, if (k == 0) { ret = NOT_COMPILED_IN; } - else if (hadPriv) { + else { +#ifndef WOLFSSL_MLKEM_MAKEKEY_SMALL_MEM + eSz = (size_t)(k * MLKEM_N) * sizeof(sword16); +#else + eSz = (size_t)MLKEM_N * sizeof(sword16); +#endif + } + if ((ret == 0) && hadPriv) { /* Do not write the new key over the old one. */ #ifndef WOLFSSL_MLKEM_DYNAMIC_KEYS ForceZero(key->priv, (size_t)k * MLKEM_N * sizeof(sword16)); @@ -970,9 +988,8 @@ int wc_MlKemKey_MakeKeyWithRandom(MlKemKey* key, const unsigned char* rand, key->heap, DYNAMIC_TYPE_TMP_BUFFER); #endif #else - /* e (v) */ - e = (sword16*)XMALLOC((size_t)(k * MLKEM_N) * sizeof(sword16), - key->heap, DYNAMIC_TYPE_TMP_BUFFER); + /* e (p) - matrix A is generated a polynomial at a time into it. */ + e = (sword16*)XMALLOC(eSz, key->heap, DYNAMIC_TYPE_TMP_BUFFER); #endif if (e == NULL) { ret = MEMORY_E; @@ -1133,20 +1150,17 @@ int wc_MlKemKey_MakeKeyWithRandom(MlKemKey* key, const unsigned char* rand, #ifndef WOLFSSL_NO_MALLOC /* Free dynamic memory allocated in function. */ if (e != NULL) { - /* e holds the secret noise vector; zeroize before release. The - * (public) matrix A may follow it in the same allocation but does - * not need clearing. */ - ForceZero(e, (size_t)(k * MLKEM_N) * sizeof(sword16)); + ForceZero(e, eSz); XFREE(e, key->heap, DYNAMIC_TYPE_TMP_BUFFER); } #else - /* e is a stack buffer holding the secret noise vector; zeroize it. */ + /* e is a stack buffer holding the secret noise; zeroize it. */ #ifdef WOLFSSL_CHECK_MEM_ZERO - wc_MemZero_Add("mlkem keygen e", e, (size_t)(k * MLKEM_N) * sizeof(sword16)); + wc_MemZero_Add("mlkem keygen e", e, eSz); #endif - ForceZero(e, (size_t)(k * MLKEM_N) * sizeof(sword16)); + ForceZero(e, eSz); #ifdef WOLFSSL_CHECK_MEM_ZERO - wc_MemZero_Check(e, (size_t)(k * MLKEM_N) * sizeof(sword16)); + wc_MemZero_Check(e, eSz); #endif #endif @@ -1365,14 +1379,18 @@ int wc_MlKemKey_SharedSecretSize(MlKemKey* key, word32* len) * 23: c_2 <- ByteEncode_d_v(Compress_d_v(v)) * 24: return c <- (c_1||c_2) * - * @param [in] key ML-KEM key object. - * @param [in] m Random bytes. - * @param [in] r Seed to feed to PRF when generating y, e1 and e2. - * @param [out] c Calculated cipher text. + * @param [in] key ML-KEM key object. + * @param [in] m Random bytes. + * @param [in] r Seed to feed to PRF when generating y, e1 and e2. + * @param [out] c Calculated cipher text. NULL when comparing. + * @param [in] cmp Cipher text to compare against. May be NULL. + * @param [in, out] fail Set to -1 when cipher text does not match cmp. + * Only used when cmp is not NULL. * @return 0 on success. * @return NOT_COMPILED_IN when key type is not supported. */ -static int mlkemkey_encapsulate(MlKemKey* key, const byte* m, byte* r, byte* c) +static int mlkemkey_encapsulate(MlKemKey* key, const byte* m, byte* r, byte* c, + const byte* cmp, int* fail) { int ret = 0; sword16* a = NULL; @@ -1390,11 +1408,14 @@ static int mlkemkey_encapsulate(MlKemKey* key, const byte* m, byte* r, byte* c) #ifndef WOLFSSL_MLKEM_ENCAPSULATE_SMALL_MEM sword16 y[((WC_ML_KEM_MAX_K + 3) * WC_ML_KEM_MAX_K + 3) * MLKEM_N]; #else - sword16 y[3 * WC_ML_KEM_MAX_K * MLKEM_N]; + sword16 y[(WC_ML_KEM_MAX_K + 2) * MLKEM_N]; + byte block[MLKEM_MAX_COMP_POLY_SZ]; #endif #endif sword16* u = 0; +#ifndef WOLFSSL_MLKEM_ENCAPSULATE_SMALL_MEM sword16* v = 0; +#endif /* Establish parameters based on key type. */ switch (key->type) { @@ -1450,7 +1471,11 @@ static int mlkemkey_encapsulate(MlKemKey* key, const byte* m, byte* r, byte* c) #ifndef WOLFSSL_MLKEM_ENCAPSULATE_SMALL_MEM yAllocSz = ((k + 3) * k + 3) * MLKEM_N * sizeof(sword16); #else - yAllocSz = 3 * k * MLKEM_N * sizeof(sword16); + yAllocSz = (k + 2) * MLKEM_N * sizeof(sword16); + if (cmp != NULL) { + /* One block of cipher text to compare against. */ + yAllocSz += MLKEM_MAX_COMP_POLY_SZ; + } #endif y = (sword16*)XMALLOC(yAllocSz, key->heap, DYNAMIC_TYPE_TMP_BUFFER); if (y == NULL) { @@ -1516,8 +1541,9 @@ static int mlkemkey_encapsulate(MlKemKey* key, const byte* m, byte* r, byte* c) #else /* WOLFSSL_MLKEM_ENCAPSULATE_SMALL_MEM */ if (ret == 0) { /* Assign allocated dynamic memory to pointers. - * y (v) | a (v) | u (v) */ - a = y + MLKEM_N * k; + * y (v) | u (p) | a (p) [| block] */ + u = y + MLKEM_N * k; + a = u + MLKEM_N; /* Initialize the PRF for use in the noise generation. */ ret = mlkem_prf_reset(&key->prf); @@ -1528,19 +1554,28 @@ static int mlkemkey_encapsulate(MlKemKey* key, const byte* m, byte* r, byte* c) ret = mlkem_get_noise(&key->prf, (int)k, y, NULL, NULL, r); } if (ret == 0) { - /* Assign remaining allocated dynamic memory to pointers. - * y (v) | at (v) | u (v) */ - u = a + MLKEM_N * k; - v = a; + byte* cb = c; - /* Perform encapsulation maths. + if (cmp != NULL) { + /* Cipher text is compared a block at a time as it is calculated. */ +#ifndef WOLFSSL_NO_MALLOC + cb = (byte*)(a + MLKEM_N); +#else + cb = block; +#endif + } + + /* Perform encapsulation maths and encode the cipher text. * Steps 13-17: generate e_1 and e_2 - * Steps 18-19, 21: calculate u and v */ - ret = mlkem_encapsulate_seeds(key->pub, &key->prf, u, a, y, (int)k, m, - key->pubSeed, r); + * Steps 18-19, 21: calculate u and v + * Steps 22-24: c <- (c_1||c_2) */ + ret = mlkem_encapsulate_seeds(key->pub, &key->prf, cb, cmp, fail, u, a, + y, (int)k, m, key->pubSeed, r); } + (void)compVecSz; #endif /* WOLFSSL_MLKEM_ENCAPSULATE_SMALL_MEM */ +#ifndef WOLFSSL_MLKEM_ENCAPSULATE_SMALL_MEM if (ret == 0) { byte* c1 = c; byte* c2 = c + compVecSz; @@ -1574,7 +1609,7 @@ static int mlkemkey_encapsulate(MlKemKey* key, const byte* m, byte* r, byte* c) /* Step 22: c_1 <- ByteEncode_d_u(Compress_d_u(u)) * Step 23: c_2 <- ByteEncode_d_v(Compress_d_v(v)) */ MLKEM_ARM64_SVR({ - ret = mlkem_vec_compress_11(c1, u); + ret = mlkem_vec_compress_11(c1, u, k); if (ret == 0) ret = mlkem_compress_5(c2, v); }); @@ -1582,6 +1617,9 @@ static int mlkemkey_encapsulate(MlKemKey* key, const byte* m, byte* r, byte* c) } #endif } + (void)cmp; + (void)fail; +#endif /* !WOLFSSL_MLKEM_ENCAPSULATE_SMALL_MEM */ #ifndef WOLFSSL_NO_MALLOC /* Dispose of dynamic memory allocated in function. The buffer holds secret @@ -1916,7 +1954,8 @@ int wc_MlKemKey_EncapsulateWithRandom(MlKemKey* key, unsigned char* ct, #endif #ifdef WOLFSSL_MLKEM_KYBER { - ret = mlkemkey_encapsulate(key, msg, kr + WC_ML_KEM_SYM_SZ, ct); + ret = mlkemkey_encapsulate(key, msg, kr + WC_ML_KEM_SYM_SZ, ct, + NULL, NULL); } #endif #if defined(WOLFSSL_MLKEM_KYBER) && !defined(WOLFSSL_NO_ML_KEM) @@ -1925,7 +1964,8 @@ int wc_MlKemKey_EncapsulateWithRandom(MlKemKey* key, unsigned char* ct, #ifndef WOLFSSL_NO_ML_KEM { /* Step 2: c <- K-PKE.Encrypt(ek,m,r) */ - ret = mlkemkey_encapsulate(key, rand, kr + WC_ML_KEM_SYM_SZ, ct); + ret = mlkemkey_encapsulate(key, rand, kr + WC_ML_KEM_SYM_SZ, ct, + NULL, NULL); } #endif } @@ -2228,6 +2268,7 @@ int wc_MlKemKey_Decapsulate(MlKemKey* key, unsigned char* ss, byte msg[WC_ML_KEM_SYM_SZ]; byte kr[2 * WC_ML_KEM_SYM_SZ + 1]; unsigned int i = 0; +#ifndef WOLFSSL_MLKEM_ENCAPSULATE_SMALL_MEM int fail = -1; /* mismatch until mlkem_cmp() says otherwise */ #if defined(WOLFSSL_SMALL_STACK) || \ (!defined(USE_INTEL_SPEEDUP) && !defined(WOLFSSL_NO_MALLOC)) @@ -2235,6 +2276,9 @@ int wc_MlKemKey_Decapsulate(MlKemKey* key, unsigned char* ss, #else byte cmp[WC_ML_KEM_MAX_CIPHER_TEXT_SIZE]; #endif +#else + int fail = 0; +#endif #endif /* Validate parameters. */ @@ -2322,6 +2366,7 @@ int wc_MlKemKey_Decapsulate(MlKemKey* key, unsigned char* ss, } #else +#ifndef WOLFSSL_MLKEM_ENCAPSULATE_SMALL_MEM #if defined(WOLFSSL_SMALL_STACK) || \ (!defined(USE_INTEL_SPEEDUP) && !defined(WOLFSSL_NO_MALLOC)) if (ret == 0) { @@ -2331,6 +2376,7 @@ int wc_MlKemKey_Decapsulate(MlKemKey* key, unsigned char* ss, ret = MEMORY_E; } } +#endif #endif /* msg and kr hold secret decapsulation material; baseline-zero and register @@ -2361,12 +2407,20 @@ int wc_MlKemKey_Decapsulate(MlKemKey* key, unsigned char* ss, } if (ret == 0) { /* Encapsulate the message. */ - ret = mlkemkey_encapsulate(key, msg, kr + WC_ML_KEM_SYM_SZ, cmp); +#ifndef WOLFSSL_MLKEM_ENCAPSULATE_SMALL_MEM + ret = mlkemkey_encapsulate(key, msg, kr + WC_ML_KEM_SYM_SZ, cmp, NULL, + NULL); +#else + ret = mlkemkey_encapsulate(key, msg, kr + WC_ML_KEM_SYM_SZ, NULL, ct, + &fail); +#endif } +#ifndef WOLFSSL_MLKEM_ENCAPSULATE_SMALL_MEM if (ret == 0) { /* Compare generated cipher text with that passed in. */ MLKEM_ARM64_SVR(ret = mlkem_cmp(ct, cmp, (int)ctSz, &fail)); } +#endif if (ret == 0) { #if defined(WOLFSSL_MLKEM_KYBER) && !defined(WOLFSSL_NO_ML_KEM) if (key->type & MLKEM_KYBER) @@ -2402,6 +2456,7 @@ int wc_MlKemKey_Decapsulate(MlKemKey* key, unsigned char* ss, #endif } +#ifndef WOLFSSL_MLKEM_ENCAPSULATE_SMALL_MEM #if defined(WOLFSSL_SMALL_STACK) || \ (!defined(USE_INTEL_SPEEDUP) && !defined(WOLFSSL_NO_MALLOC)) /* Dispose of dynamic memory allocated in function. cmp holds the @@ -2415,6 +2470,7 @@ int wc_MlKemKey_Decapsulate(MlKemKey* key, unsigned char* ss, #else /* cmp is a stack buffer holding the re-encrypted ciphertext; zeroize it. */ ForceZero(cmp, sizeof(cmp)); +#endif #endif ForceZero(msg, sizeof(msg)); diff --git a/wolfcrypt/src/wc_mlkem_poly.c b/wolfcrypt/src/wc_mlkem_poly.c index 2ab6159c187..dd4a48bfee0 100644 --- a/wolfcrypt/src/wc_mlkem_poly.c +++ b/wolfcrypt/src/wc_mlkem_poly.c @@ -126,8 +126,8 @@ !defined(WOLFSSL_MLKEM_NO_MAKE_KEY)) || \ (defined(WOLFSSL_MLKEM_ENCAPSULATE_SMALL_MEM) && \ !defined(WOLFSSL_MLKEM_NO_ENCAPSULATE)) -static int mlkem_gen_matrix_i(MLKEM_PRF_T* prf, sword16* a, int k, byte* seed, - int i, int transposed); +static int mlkem_gen_matrix_i_acc(MLKEM_PRF_T* prf, sword16* r, sword16* a, + const sword16* v, int k, byte* seed, int i, int transposed); static int mlkem_get_noise_i(MLKEM_PRF_T* prf, int k, sword16* vec2, byte* seed, int i, int make); static int mlkem_get_noise_eta2_c(MLKEM_PRF_T* prf, sword16* p, @@ -1250,7 +1250,8 @@ static void mlkem_basemul_mont_add(sword16* WC_RESTRICT r, * !WOLFSSL_MLKEM_NO_DECAPSULATE */ #endif -#if !defined(WOLFSSL_MLKEM_NO_MAKE_KEY) || \ +#if (!defined(WOLFSSL_MLKEM_NO_MAKE_KEY) && \ + !defined(WOLFSSL_MLKEM_MAKEKEY_SMALL_MEM)) || \ !defined(WOLFSSL_MLKEM_NO_ENCAPSULATE) || \ !defined(WOLFSSL_MLKEM_NO_DECAPSULATE) /* Pointwise multiply elements of a and b, into r, and multiply by 2^-16. @@ -1277,8 +1278,8 @@ static void mlkem_pointwise_acc_mont(sword16* r, const sword16* a, mlkem_basemul_mont_add(r, a + (k - 1) * MLKEM_N, b + (k - 1) * MLKEM_N); #endif } -#endif /* !WOLFSSL_MLKEM_NO_MAKE_KEY || !WOLFSSL_MLKEM_NO_ENCAPSULATE || - * !WOLFSSL_MLKEM_NO_DECAPSULATE */ +#endif /* (!WOLFSSL_MLKEM_NO_MAKE_KEY && !WOLFSSL_MLKEM_MAKEKEY_SMALL_MEM) || + * !WOLFSSL_MLKEM_NO_ENCAPSULATE || !WOLFSSL_MLKEM_NO_DECAPSULATE */ /******************************************************************************/ @@ -2026,7 +2027,7 @@ int mlkem_keygen(sword16* s, sword16* t, sword16* e, const sword16* a, int k) * @param [in, out] s Private key vector of polynomials. * @param [out] t Public key vector of polynomials. * @param [in, out] prf XOF object. - * @param [in] tv Temporary vector of polynomials. + * @param [in] tp Temporary polynomial. * @param [in] k Number of polynomials in vector. * @param [in] rho Random seed to generate matrix A from. * @param [in, out] sigma Random seed to generate noise from. @@ -2036,12 +2037,12 @@ int mlkem_keygen(sword16* s, sword16* t, sword16* e, const sword16* a, int k) * @return Other negative value when a hash error occurred. */ int mlkem_keygen_seeds(sword16* s, sword16* t, MLKEM_PRF_T* prf, - sword16* tv, int k, byte* rho, byte* sigma) + sword16* tp, int k, byte* rho, byte* sigma) { int i; int ret = 0; - sword16* ai = tv; - sword16* e = tv; + sword16* ai = tp; + sword16* e = tp; /* Transform private key. All of result used in public key calculation * Step 16: s_hat = NTT(s) */ @@ -2054,16 +2055,15 @@ int mlkem_keygen_seeds(sword16* s, sword16* t, MLKEM_PRF_T* prf, for (i = 0; i < k; ++i) { int j; - /* Generate a vector of matrix A. - * Steps 4-6: generate A[i] */ - ret = mlkem_gen_matrix_i(prf, ai, k, rho, i, 0); + /* Generate a vector of matrix A, a polynomial at a time, multiplying + * each into the public polynomial. + * Steps 4-6: generate A[i] + * Step 18: ... A_hat o s_hat ... */ + ret = mlkem_gen_matrix_i_acc(prf, t + i * MLKEM_N, ai, s, k, rho, i, 0); if (ret != 0) { break; } - /* Multiply a by private into public polynomial. - * Step 18: ... A_hat o s_hat ... */ - mlkem_pointwise_acc_mont(t + i * MLKEM_N, ai, s, (unsigned int)k); /* Convert public polynomial to Montgomery form. * Step 18: ... MontRed(A_hat o s_hat) ... */ for (j = 0; j < MLKEM_N; ++j) { @@ -2094,6 +2094,10 @@ int mlkem_keygen_seeds(sword16* s, sword16* t, MLKEM_PRF_T* prf, #endif } + if (ret != 0) { + ForceZero(t, (size_t)k * MLKEM_N * sizeof(sword16)); + } + return ret; } @@ -2227,11 +2231,15 @@ int mlkem_encapsulate(const sword16* pub, sword16* u, sword16* v, #else -/* Encapsulate message. +/* Encapsulate message and encode the cipher text a polynomial of u at a time. * * @param [in] pub Public key vector of polynomials. * @param [in, out] prf XOF object. - * @param [out] u Vector of polynomials. + * @param [out] c Cipher text, or one block of it when comparing. + * @param [in] cmp Cipher text to compare against. May be NULL. + * @param [in, out] fail Set to -1 when cipher text does not match cmp. + * Only used when cmp is not NULL. + * @param [in, out] u Polynomial. * @param [in, out] tp Polynomial. * @param [in, out] y Vector of polynomials. * @param [in] k Number of polynomials in vector. @@ -2241,18 +2249,37 @@ int mlkem_encapsulate(const sword16* pub, sword16* u, sword16* v, * @return 0 on success. * @return MEMORY_E when dynamic memory allocation fails. Only possible when * WOLFSSL_SMALL_STACK is defined. + * @return BAD_FUNC_ARG when cmp is not NULL and decapsulation is not compiled + * in. * @return Other negative value when a hash error occurred. */ -int mlkem_encapsulate_seeds(const sword16* pub, MLKEM_PRF_T* prf, sword16* u, - sword16* tp, sword16* y, int k, const byte* msg, byte* seed, byte* coins) +int mlkem_encapsulate_seeds(const sword16* pub, MLKEM_PRF_T* prf, byte* c, + const byte* cmp, int* fail, sword16* u, sword16* tp, sword16* y, int k, + const byte* msg, byte* seed, byte* coins) { int ret = 0; int i; sword16* a = tp; sword16* e1 = tp; - sword16* v = tp; - sword16* e2 = tp + MLKEM_N; + sword16* e2 = tp; + sword16* v = u; sword16* m = y; + byte* cb = c; + /* Number of bytes a polynomial of u is compressed into. */ + unsigned int blockSz = MLKEM_POLY_COMPRESSED_SZ(MLKEM_COMP_10BITS); + +#if defined(WOLFSSL_KYBER1024) || defined(WOLFSSL_WC_ML_KEM_1024) + if (k == WC_ML_KEM_1024_K) { + blockSz = MLKEM_POLY_COMPRESSED_SZ(MLKEM_COMP_11BITS); + } +#endif + +#ifdef WOLFSSL_MLKEM_NO_DECAPSULATE + (void)fail; + if (cmp != NULL) { + return BAD_FUNC_ARG; + } +#endif /* Transform y. All of result used in calculation of u and v. */ for (i = 0; i < k; ++i) { @@ -2263,16 +2290,14 @@ int mlkem_encapsulate_seeds(const sword16* pub, MLKEM_PRF_T* prf, sword16* u, for (i = 0; i < k; ++i) { int j; - /* Generate a vector of matrix A. */ - ret = mlkem_gen_matrix_i(prf, a, k, seed, i, 1); + /* Generate a vector of matrix A, a polynomial at a time, multiplying + * each into the u polynomial. */ + ret = mlkem_gen_matrix_i_acc(prf, u, a, y, k, seed, i, 1); if (ret != 0) { break; } - - /* Multiply at by y into u polynomial. */ - mlkem_pointwise_acc_mont(u + i * MLKEM_N, a, y, (unsigned int)k); /* Inverse transform u polynomial. */ - mlkem_invntt(u + i * MLKEM_N); + mlkem_invntt(u); /* Generate noise using PRF. */ ret = mlkem_get_noise_i(prf, k, e1, coins, i, 0); @@ -2282,43 +2307,78 @@ int mlkem_encapsulate_seeds(const sword16* pub, MLKEM_PRF_T* prf, sword16* u, /* Add errors to u and reduce. */ #if defined(WOLFSSL_MLKEM_SMALL) || defined(WOLFSSL_MLKEM_NO_LARGE_CODE) for (j = 0; j < MLKEM_N; ++j) { - sword16 t = (sword16)(u[i * MLKEM_N + j] + e1[j]); - u[i * MLKEM_N + j] = MLKEM_BARRETT_RED(t); + sword16 t = (sword16)(u[j] + e1[j]); + u[j] = MLKEM_BARRETT_RED(t); } #else for (j = 0; j < MLKEM_N; j += 8) { - sword16 t0 = (sword16)(u[i * MLKEM_N + j + 0] + e1[j + 0]); - sword16 t1 = (sword16)(u[i * MLKEM_N + j + 1] + e1[j + 1]); - sword16 t2 = (sword16)(u[i * MLKEM_N + j + 2] + e1[j + 2]); - sword16 t3 = (sword16)(u[i * MLKEM_N + j + 3] + e1[j + 3]); - sword16 t4 = (sword16)(u[i * MLKEM_N + j + 4] + e1[j + 4]); - sword16 t5 = (sword16)(u[i * MLKEM_N + j + 5] + e1[j + 5]); - sword16 t6 = (sword16)(u[i * MLKEM_N + j + 6] + e1[j + 6]); - sword16 t7 = (sword16)(u[i * MLKEM_N + j + 7] + e1[j + 7]); - u[i * MLKEM_N + j + 0] = MLKEM_BARRETT_RED(t0); - u[i * MLKEM_N + j + 1] = MLKEM_BARRETT_RED(t1); - u[i * MLKEM_N + j + 2] = MLKEM_BARRETT_RED(t2); - u[i * MLKEM_N + j + 3] = MLKEM_BARRETT_RED(t3); - u[i * MLKEM_N + j + 4] = MLKEM_BARRETT_RED(t4); - u[i * MLKEM_N + j + 5] = MLKEM_BARRETT_RED(t5); - u[i * MLKEM_N + j + 6] = MLKEM_BARRETT_RED(t6); - u[i * MLKEM_N + j + 7] = MLKEM_BARRETT_RED(t7); + sword16 t0 = (sword16)(u[j + 0] + e1[j + 0]); + sword16 t1 = (sword16)(u[j + 1] + e1[j + 1]); + sword16 t2 = (sword16)(u[j + 2] + e1[j + 2]); + sword16 t3 = (sword16)(u[j + 3] + e1[j + 3]); + sword16 t4 = (sword16)(u[j + 4] + e1[j + 4]); + sword16 t5 = (sword16)(u[j + 5] + e1[j + 5]); + sword16 t6 = (sword16)(u[j + 6] + e1[j + 6]); + sword16 t7 = (sword16)(u[j + 7] + e1[j + 7]); + u[j + 0] = MLKEM_BARRETT_RED(t0); + u[j + 1] = MLKEM_BARRETT_RED(t1); + u[j + 2] = MLKEM_BARRETT_RED(t2); + u[j + 3] = MLKEM_BARRETT_RED(t3); + u[j + 4] = MLKEM_BARRETT_RED(t4); + u[j + 5] = MLKEM_BARRETT_RED(t5); + u[j + 6] = MLKEM_BARRETT_RED(t6); + u[j + 7] = MLKEM_BARRETT_RED(t7); + } +#endif + + /* Encode the polynomial of u into a block of cipher text. + * Step 22: c_1 <- ByteEncode_d_u(Compress_d_u(u)) */ +#if defined(WOLFSSL_KYBER512) || defined(WOLFSSL_WC_ML_KEM_512) + if (k == WC_ML_KEM_512_K) { + ret = mlkem_vec_compress_10(cb, u, 1); } #endif - } +#if defined(WOLFSSL_KYBER768) || defined(WOLFSSL_WC_ML_KEM_768) + if (k == WC_ML_KEM_768_K) { + ret = mlkem_vec_compress_10(cb, u, 1); + } +#endif +#if defined(WOLFSSL_KYBER1024) || defined(WOLFSSL_WC_ML_KEM_1024) + if (k == WC_ML_KEM_1024_K) { + ret = mlkem_vec_compress_11(cb, u, 1); + } +#endif + if (ret != 0) { + break; + } + if (cmp == NULL) { + cb += blockSz; + } +#ifndef WOLFSSL_MLKEM_NO_DECAPSULATE + else { + int blockFail; - /* Multiply public key by y into v polynomial. */ - mlkem_pointwise_acc_mont(v, pub, y, (unsigned int)k); - /* Inverse transform v. */ - mlkem_invntt(v); + ret = mlkem_cmp(cb, cmp, (int)blockSz, &blockFail); + *fail |= blockFail; + if (ret != 0) { + break; + } + cmp += blockSz; + } +#endif + } if (ret == 0) { + /* Multiply public key by y into v polynomial. */ + mlkem_pointwise_acc_mont(v, pub, y, (unsigned int)k); + /* Inverse transform v. */ + mlkem_invntt(v); + ret = mlkem_from_msg(m, msg); } - - /* Generate noise using PRF. */ - coins[WC_ML_KEM_SYM_SZ] = WC_OCTET(2 * k); if (ret == 0) { + /* Generate noise using PRF. */ + coins[WC_ML_KEM_SYM_SZ] = WC_OCTET(2 * k); ret = mlkem_get_noise_eta2_c(prf, e2, coins); } if (ret == 0) { @@ -2348,7 +2408,37 @@ int mlkem_encapsulate_seeds(const sword16* pub, MLKEM_PRF_T* prf, sword16* u, v[i + 7] = MLKEM_BARRETT_RED(t7); } #endif + + /* Encode v into the last block of cipher text. + * Step 23: c_2 <- ByteEncode_d_v(Compress_d_v(v)) */ + #if defined(WOLFSSL_KYBER512) || defined(WOLFSSL_WC_ML_KEM_512) + if (k == WC_ML_KEM_512_K) { + ret = mlkem_compress_4(cb, v); + blockSz = MLKEM_POLY_COMPRESSED_SZ(MLKEM_COMP_4BITS); + } + #endif + #if defined(WOLFSSL_KYBER768) || defined(WOLFSSL_WC_ML_KEM_768) + if (k == WC_ML_KEM_768_K) { + ret = mlkem_compress_4(cb, v); + blockSz = MLKEM_POLY_COMPRESSED_SZ(MLKEM_COMP_4BITS); + } + #endif + #if defined(WOLFSSL_KYBER1024) || defined(WOLFSSL_WC_ML_KEM_1024) + if (k == WC_ML_KEM_1024_K) { + ret = mlkem_compress_5(cb, v); + blockSz = MLKEM_POLY_COMPRESSED_SZ(MLKEM_COMP_5BITS); + } + #endif + } +#ifndef WOLFSSL_MLKEM_NO_DECAPSULATE + if ((ret == 0) && (cmp != NULL)) { + int blockFail; + + ret = mlkem_cmp(cb, cmp, (int)blockSz, &blockFail); + *fail |= blockFail; } +#endif + /* Step 24: return c <- (c_1||c_2) */ return ret; } @@ -4438,25 +4528,35 @@ int mlkem_gen_matrix(MLKEM_PRF_T* prf, sword16* a, int k, byte* seed, (defined(WOLFSSL_MLKEM_ENCAPSULATE_SMALL_MEM) && \ !defined(WOLFSSL_MLKEM_NO_ENCAPSULATE)) -/* Deterministically generate a matrix (or transpose) of uniform integers mod q. +/* Deterministically generate a vector of a matrix (or transpose) of uniform + * integers mod q and multiply it into a vector of polynomials. * * Seed used with XOF to generate random bytes. * + * Each polynomial of the vector of the matrix is multiplied into the result as + * soon as it has been generated. + * * FIPS 203, Algorithm 13: K-PKE.KeyGen(d) * ... * 4: for (j <- 0; j < k; j++) * 5: A_hat[i,j] <- SampleNTT(rho||j||i) * 6: end for * ... + * 18: t_hat <- A_hat o s_hat + e_hat + * ... * FIPS 203, Algorithm 14: K-PKE.Encrypt(ek_PKE,m,r) * ... * 5: for (j <- 0; j < k; j++) * 6: A_hat[i,j] <- SampleNTT(rho||j||i) (Transposed is rho||i||j) * 7: end for * ... + * 18: u <- InvNTT(A_hat_trans o y_hat) + e_1 + * ... * * @param [in, out] prf XOF object. - * @param [out] a Matrix of uniform integers. + * @param [out] r Polynomial holding the multiplied result. + * @param [out] a Polynomial of uniform integers. Scratch. + * @param [in] v Vector of polynomials to multiply with. * @param [in] k Number of dimensions. k x k polynomials. * @param [in] seed Bytes to seed XOF generation. * @param [in] i Index of vector to generate. @@ -4465,8 +4565,8 @@ int mlkem_gen_matrix(MLKEM_PRF_T* prf, sword16* a, int k, byte* seed, * @return MEMORY_E when dynamic memory allocation fails. Only possible when * WOLFSSL_SMALL_STACK is defined. */ -static int mlkem_gen_matrix_i(MLKEM_PRF_T* prf, sword16* a, int k, byte* seed, - int i, int transposed) +static int mlkem_gen_matrix_i_acc(MLKEM_PRF_T* prf, sword16* r, sword16* a, + const sword16* v, int k, byte* seed, int i, int transposed) { #if defined(WOLFSSL_SMALL_STACK) && !defined(WOLFSSL_NO_MALLOC) byte* rand; @@ -4523,8 +4623,7 @@ static int mlkem_gen_matrix_i(MLKEM_PRF_T* prf, sword16* a, int k, byte* seed, /* Sample random bytes to create a polynomial. * Alg 7, Step 3 - implicitly counter is 0. * Alg 7, Step 4-16. */ - ctr = mlkem_rej_uniform_c(a + j * MLKEM_N, MLKEM_N, rand, - GEN_MATRIX_SIZE); + ctr = mlkem_rej_uniform_c(a, MLKEM_N, rand, GEN_MATRIX_SIZE); /* Create more blocks if too many rejected. * Alg 7, Step 4. */ while ((ret == 0) && (ctr < MLKEM_N)) { @@ -4532,10 +4631,20 @@ static int mlkem_gen_matrix_i(MLKEM_PRF_T* prf, sword16* a, int k, byte* seed, ret = mlkem_xof_squeezeblocks(prf, rand, 1); if (ret == 0) { /* Alg 7, Step 4-16. */ - ctr += mlkem_rej_uniform_c(a + j * MLKEM_N + ctr, - MLKEM_N - ctr, rand, XOF_BLOCK_SIZE); + ctr += mlkem_rej_uniform_c(a + ctr, MLKEM_N - ctr, rand, + XOF_BLOCK_SIZE); } } + + /* Multiply the polynomial of the matrix into the result. + * Alg 13, Step 18: ... A_hat o s_hat ... + * Alg 14, Step 18: ... A_hat_trans o y_hat ... */ + if (j == 0) { + mlkem_basemul_mont(r, a, v); + } + else { + mlkem_basemul_mont_add(r, a, v + j * MLKEM_N); + } } } @@ -6487,8 +6596,9 @@ int mlkem_vec_compress_10(byte* r, sword16* v, unsigned int k) * * Coefficients must be in range 0..2q-1 and are not modified. * - * @param [out] r Array of bytes. - * @param [in, out] v Vector of polynomials. + * @param [out] r Array of bytes. + * @param [in, out] v Vector of polynomials. + * @param [in] kp Number of polynomials in vector. */ #define mlkem_vec_compress_11_c mlkem_thumb2_vec_compress_11 #else @@ -6496,10 +6606,11 @@ int mlkem_vec_compress_10(byte* r, sword16* v, unsigned int k) * * FIPS 203, Section 4.2.1, Compression and decompression * - * @param [out] r Array of bytes. - * @param [in, out] v Vector of polynomials. + * @param [out] r Array of bytes. + * @param [in, out] v Vector of polynomials. + * @param [in] kp Number of polynomials in vector. */ -static void mlkem_vec_compress_11_c(byte* r, sword16* v) +static void mlkem_vec_compress_11_c(byte* r, sword16* v, unsigned int kp) { unsigned int i; unsigned int j; @@ -6508,14 +6619,14 @@ static void mlkem_vec_compress_11_c(byte* r, sword16* v) sword16 t[8]; #endif - for (i = 0; i < 4; i++) { + for (i = 0; i < kp; i++) { /* Reduce each coefficient to mod q. */ mlkem_csubq_c(v + i * MLKEM_N); /* All values are now positive. */ } /* Each polynomial. */ - for (i = 0; i < 4; i++) { + for (i = 0; i < kp; i++) { /* Each 8 polynomial coefficients. */ for (j = 0; j < MLKEM_N; j += 8) { #ifdef WOLFSSL_MLKEM_SMALL @@ -6576,10 +6687,11 @@ static void mlkem_vec_compress_11_c(byte* r, sword16* v) * * FIPS 203, Section 4.2.1, Compression and decompression * - * @param [out] r Array of bytes. - * @param [in, out] v Vector of polynomials. + * @param [out] r Array of bytes. + * @param [in, out] v Vector of polynomials. + * @param [in] kp Number of polynomials in vector. */ -int mlkem_vec_compress_11(byte* r, sword16* v) +int mlkem_vec_compress_11(byte* r, sword16* v, unsigned int kp) { #ifdef USE_INTEL_SPEEDUP #ifdef WOLFSSL_MLKEM_HAVE_INTEL_AVX512 @@ -6587,7 +6699,7 @@ int mlkem_vec_compress_11(byte* r, sword16* v) int svr_ret = SAVE_VECTOR_REGISTERS2(); if (svr_ret != 0) return svr_ret; - mlkem_compress_11_avx512(r, v, 4); + mlkem_compress_11_avx512(r, v, (int)kp); RESTORE_VECTOR_REGISTERS(); } else @@ -6596,13 +6708,13 @@ int mlkem_vec_compress_11(byte* r, sword16* v) int svr_ret = SAVE_VECTOR_REGISTERS2(); if (svr_ret != 0) return svr_ret; - mlkem_compress_11_avx2(r, v, 4); + mlkem_compress_11_avx2(r, v, (int)kp); RESTORE_VECTOR_REGISTERS(); } else #endif { - mlkem_vec_compress_11_c(r, v); + mlkem_vec_compress_11_c(r, v, kp); } return 0; } diff --git a/wolfssl/wolfcrypt/wc_mlkem.h b/wolfssl/wolfcrypt/wc_mlkem.h index 78774d86683..81481439e3c 100644 --- a/wolfssl/wolfcrypt/wc_mlkem.h +++ b/wolfssl/wolfcrypt/wc_mlkem.h @@ -565,9 +565,9 @@ int mlkem_encapsulate(const sword16* pub, sword16* bp, sword16* v, const sword16* m, int kp); #else WOLFSSL_LOCAL -int mlkem_encapsulate_seeds(const sword16* pub, MLKEM_PRF_T* prf, sword16* bp, - sword16* tp, sword16* sp, int kp, const byte* msg, byte* seed, - byte* coins); +int mlkem_encapsulate_seeds(const sword16* pub, MLKEM_PRF_T* prf, byte* c, + const byte* cmp, int* fail, sword16* bp, sword16* tp, sword16* sp, int kp, + const byte* msg, byte* seed, byte* coins); #endif WOLFSSL_LOCAL int mlkem_decapsulate(const sword16* priv, sword16* mp, sword16* bp, @@ -612,7 +612,7 @@ int mlkem_cmp(const byte* a, const byte* b, int sz, int* fail); WOLFSSL_LOCAL int mlkem_vec_compress_10(byte* r, sword16* v, unsigned int kp); WOLFSSL_LOCAL -int mlkem_vec_compress_11(byte* r, sword16* v); +int mlkem_vec_compress_11(byte* r, sword16* v, unsigned int kp); WOLFSSL_LOCAL int mlkem_vec_decompress_10(sword16* v, const unsigned char* b, unsigned int kp); @@ -888,7 +888,8 @@ WOLFSSL_LOCAL void mlkem_thumb2_basemul_mont_add(sword16* r, const sword16* a, WOLFSSL_LOCAL void mlkem_thumb2_csubq(sword16* p); WOLFSSL_LOCAL void mlkem_thumb2_vec_compress_10(byte* r, sword16* v, unsigned int k); -WOLFSSL_LOCAL void mlkem_thumb2_vec_compress_11(byte* r, sword16* v); +WOLFSSL_LOCAL void mlkem_thumb2_vec_compress_11(byte* r, sword16* v, + unsigned int k); WOLFSSL_LOCAL unsigned int mlkem_thumb2_rej_uniform(sword16* p, unsigned int len, const byte* r, unsigned int rLen); #elif defined(WOLFSSL_ARMASM) From 3e78d9fb0121cbaab84d9bbf8d622f4e0f0a08b5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tobias=20Frauenschl=C3=A4ger?= Date: Sun, 6 Sep 2026 14:16:50 +0200 Subject: [PATCH 02/11] Zeroize the ML-KEM cipher text compare buffer without malloc Small memory decapsulation compares the re-encapsulated cipher text a block at a time in a scratch buffer. With malloc that buffer sits inside the y allocation and is covered by the ForceZero of the whole allocation. Without malloc it is a separate stack array and was left set, unlike the full cipher text copy it replaced. Zeroize it on the way out, matching the handling of the buffer it replaced (FIPS 203 section 3.3). --- wolfcrypt/src/wc_mlkem.c | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/wolfcrypt/src/wc_mlkem.c b/wolfcrypt/src/wc_mlkem.c index 40a5f6fe122..06afadd90e8 100644 --- a/wolfcrypt/src/wc_mlkem.c +++ b/wolfcrypt/src/wc_mlkem.c @@ -1639,6 +1639,11 @@ static int mlkemkey_encapsulate(MlKemKey* key, const byte* m, byte* r, byte* c, #ifdef WOLFSSL_CHECK_MEM_ZERO wc_MemZero_Check(y, sizeof(y)); #endif +#ifdef WOLFSSL_MLKEM_ENCAPSULATE_SMALL_MEM + /* block holds the cipher text re-encapsulated from the secret decrypted + * message. With malloc it sits in the y allocation and is covered above. */ + ForceZero(block, sizeof(block)); +#endif #endif return ret; From c0fcb8450dc69103ee8ad7a4ca19bf4b22203805 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tobias=20Frauenschl=C3=A4ger?= Date: Sun, 6 Sep 2026 14:16:50 +0200 Subject: [PATCH 03/11] Allow the ML-KEM matrix cache with the small memory implementations WOLFSSL_MLKEM_CACHE_A refused to build alongside the small memory key generation and encapsulation. Both stream matrix A a polynomial at a time, so there was nowhere for the cache to be filled from or read. Key generation now writes each polynomial of A into the cache as it generates it rather than reusing one temporary, and encapsulation reads the transpose of the cache when key generation left one, since A_trans[i,j] is A[j,i]. Neither needs more temporary memory: the generator takes the stride between polynomials, which is zero when the matrix is not being kept. Encapsulation and decapsulation are unchanged when there is no cached matrix, which is the case for a public key that was decoded rather than generated. Cipher texts and shared secrets are unchanged. For ML-KEM-768 the matrix generation was 27 of the 44 Keccak permutations in decapsulation, so decapsulation drops 38 percent and encapsulation with a locally generated key drops 50 percent. Add a multi-arch row that runs both small memory options with the cache on the Thumb-2 assembly, the only one whose compressor the small memory encapsulation calls a polynomial at a time. --- .github/configs/multi-arch.json | 8 +++ .../test_mlkem_poly_hash_fault_whitebox.c | 16 ++--- wolfcrypt/src/wc_mlkem.c | 22 ++++--- wolfcrypt/src/wc_mlkem_poly.c | 61 ++++++++++++++----- wolfssl/wolfcrypt/wc_mlkem.h | 4 +- 5 files changed, 76 insertions(+), 35 deletions(-) diff --git a/.github/configs/multi-arch.json b/.github/configs/multi-arch.json index 023bb4492d1..db92f7a04f0 100644 --- a/.github/configs/multi-arch.json +++ b/.github/configs/multi-arch.json @@ -203,6 +203,14 @@ "CFLAGS=-O2"], "check": false, "run": [["env", "QEMU_LD_PREFIX=/usr/arm-linux-gnueabi", "./wolfcrypt/test/testwolfcrypt"]]}, +{"name": "thumb2-mlkem-small-mem", "minutes": 0.5, + "cc": "ccache arm-linux-gnueabi-gcc", + "comment": "Small memory ML-KEM compresses one polynomial at a time through the Thumb-2 assembly.", + "configure": ["--host=armv7m-linux-gnueabi", "--enable-cryptonly", + "--enable-mlkem", "--enable-armasm", + "CFLAGS=-O2 -DWOLFSSL_MLKEM_MAKEKEY_SMALL_MEM -DWOLFSSL_MLKEM_ENCAPSULATE_SMALL_MEM -DWOLFSSL_MLKEM_CACHE_A"], + "check": false, + "run": [["env", "QEMU_LD_PREFIX=/usr/arm-linux-gnueabi", "./wolfcrypt/test/testwolfcrypt"]]}, {"name": "armv6-gcm-table", "minutes": 3, "cc": "ccache arm-linux-gnueabi-gcc", "comment": "armv6 host_cpu selects the 32-bit ARM assembly with WOLFSSL_ARMASM_NO_HW_CRYPTO, the only configuration that calls the armv8-32 GCM_gmult_len.", diff --git a/tests/unit-mcdc/test_mlkem_poly_hash_fault_whitebox.c b/tests/unit-mcdc/test_mlkem_poly_hash_fault_whitebox.c index a2bad25c297..9e9882d55ee 100644 --- a/tests/unit-mcdc/test_mlkem_poly_hash_fault_whitebox.c +++ b/tests/unit-mcdc/test_mlkem_poly_hash_fault_whitebox.c @@ -350,8 +350,10 @@ static void wb_gen_matrix_c_continuation(void) /* ------------------------------------------------------------------------- * * mlkem_gen_matrix_i_acc(): the small-memory single-row generator, which - * multiplies each polynomial into the result as it is generated. Only - * compiled when one of the small-mem arms is on. + * multiplies each polynomial in as it is generated. Both stride forms are + * driven: aStride 0 regenerates into one scratch polynomial (the encapsulate + * caller) and aStride MLKEM_N keeps the whole vector (key generation with the + * matrix cached). * ------------------------------------------------------------------------- */ #if defined(WOLFSSL_HAVE_MLKEM) && \ !(defined(WOLFSSL_ARMASM) && defined(__aarch64__)) && \ @@ -382,17 +384,17 @@ static void wb_gen_matrix_i_rows(void) } mcdc_fh_disarm(); - if (mlkem_gen_matrix_i_acc(&prf, r, a, v, WB_K, seed, 0, 0) != 0) { + if (mlkem_gen_matrix_i_acc(&prf, r, a, 0, v, WB_K, seed, 0, 0) != 0) { WB_NOTE("baseline mlkem_gen_matrix_i_acc failed"); wb_fail = 1; } for (n = 1; n <= (long)WB_SWEEP; n++) { mcdc_fh_arm(n); - (void)mlkem_gen_matrix_i_acc(&prf, r, a, v, WB_K, seed, 0, 0); + (void)mlkem_gen_matrix_i_acc(&prf, r, a, 0, v, WB_K, seed, 0, 0); mcdc_fh_disarm(); mcdc_fh_arm(n); - (void)mlkem_gen_matrix_i_acc(&prf, r, a, v, WB_K, seed, 1, 1); + (void)mlkem_gen_matrix_i_acc(&prf, r, a, MLKEM_N, v, WB_K, seed, 1, 1); mcdc_fh_disarm(); } @@ -424,7 +426,7 @@ static void wb_gen_matrix_i_continuation(void) seed[0] = (byte)s; seed[1] = (byte)(s >> 8); mcdc_fh_disarm(); - if (mlkem_gen_matrix_i_acc(&prf, r, a, v, 1, seed, 0, 0) != 0) { + if (mlkem_gen_matrix_i_acc(&prf, r, a, 0, v, 1, seed, 0, 0) != 0) { WB_NOTE("mlkem_gen_matrix_i_acc seed search failed"); wb_fail = 1; break; @@ -442,7 +444,7 @@ static void wb_gen_matrix_i_continuation(void) a[i] = -1; } mcdc_fh_arm(WB_XOF_CALLS + 1); - ret = mlkem_gen_matrix_i_acc(&prf, r, a, v, 1, seed, 0, 0); + ret = mlkem_gen_matrix_i_acc(&prf, r, a, 0, v, 1, seed, 0, 0); mcdc_fh_disarm(); if (ret != MCDC_FH_ERR) { WB_NOTE("mlkem_gen_matrix_i_acc ignored a continuation squeeze " diff --git a/wolfcrypt/src/wc_mlkem.c b/wolfcrypt/src/wc_mlkem.c index 06afadd90e8..d26c2858b37 100644 --- a/wolfcrypt/src/wc_mlkem.c +++ b/wolfcrypt/src/wc_mlkem.c @@ -63,6 +63,7 @@ * performing decapsulation. * MlKemKey is 8KB larger but decapsulation is significantly faster. * Turn on when performing make key and decapsulation with same object. + * Also usable with the small memory options. * * WOLFSSL_MLKEM_DYNAMIC_KEYS Default: OFF * Dynamically allocates private and public key buffers instead of using @@ -132,12 +133,6 @@ #error "Can't use small memory with assembly optimized code" #endif #endif -#if defined(WOLFSSL_MLKEM_CACHE_A) - #if defined(WOLFSSL_MLKEM_MAKEKEY_SMALL_MEM) || \ - defined(WOLFSSL_MLKEM_ENCAPSULATE_SMALL_MEM) - #error "Can't cache A with small memory code" - #endif -#endif #if defined(WOLFSSL_MLKEM_NO_MAKE_KEY) && \ defined(WOLFSSL_MLKEM_NO_ENCAPSULATE) && \ @@ -918,9 +913,8 @@ int wc_MlKemKey_MakeKeyWithRandom(MlKemKey* key, const unsigned char* rand, sword16 e[MLKEM_N]; #endif #endif -#ifndef WOLFSSL_MLKEM_MAKEKEY_SMALL_MEM + /* Matrix A. NULL when it is streamed and not kept. */ sword16* a = NULL; -#endif sword16* s = NULL; sword16* t = NULL; /* Number of bytes of e holding secret noise, to be zeroized. */ @@ -1112,7 +1106,7 @@ int wc_MlKemKey_MakeKeyWithRandom(MlKemKey* key, const unsigned char* rand, * Alg 13: Steps 12-15: generate e * Alg 13: Steps 16-18: calculate t_hat from A_hat, s and e */ - ret = mlkem_keygen_seeds(s, t, &key->prf, e, k, rho, sigma); + ret = mlkem_keygen_seeds(s, t, &key->prf, e, a, k, rho, sigma); } if (ret == 0) { #endif @@ -1415,6 +1409,14 @@ static int mlkemkey_encapsulate(MlKemKey* key, const byte* m, byte* r, byte* c, sword16* u = 0; #ifndef WOLFSSL_MLKEM_ENCAPSULATE_SMALL_MEM sword16* v = 0; +#else + /* Matrix A from key generation, when it was kept. */ +#ifdef WOLFSSL_MLKEM_CACHE_A + const sword16* cachedA = ((key->flags & MLKEM_FLAG_A_SET) != 0) ? + key->a : NULL; +#else + const sword16* cachedA = NULL; +#endif #endif /* Establish parameters based on key type. */ @@ -1570,7 +1572,7 @@ static int mlkemkey_encapsulate(MlKemKey* key, const byte* m, byte* r, byte* c, * Steps 18-19, 21: calculate u and v * Steps 22-24: c <- (c_1||c_2) */ ret = mlkem_encapsulate_seeds(key->pub, &key->prf, cb, cmp, fail, u, a, - y, (int)k, m, key->pubSeed, r); + y, (int)k, m, key->pubSeed, r, cachedA); } (void)compVecSz; #endif /* WOLFSSL_MLKEM_ENCAPSULATE_SMALL_MEM */ diff --git a/wolfcrypt/src/wc_mlkem_poly.c b/wolfcrypt/src/wc_mlkem_poly.c index dd4a48bfee0..4ea79f06049 100644 --- a/wolfcrypt/src/wc_mlkem_poly.c +++ b/wolfcrypt/src/wc_mlkem_poly.c @@ -127,7 +127,8 @@ (defined(WOLFSSL_MLKEM_ENCAPSULATE_SMALL_MEM) && \ !defined(WOLFSSL_MLKEM_NO_ENCAPSULATE)) static int mlkem_gen_matrix_i_acc(MLKEM_PRF_T* prf, sword16* r, sword16* a, - const sword16* v, int k, byte* seed, int i, int transposed); + unsigned int aStride, const sword16* v, int k, byte* seed, int i, + int transposed); static int mlkem_get_noise_i(MLKEM_PRF_T* prf, int k, sword16* vec2, byte* seed, int i, int make); static int mlkem_get_noise_eta2_c(MLKEM_PRF_T* prf, sword16* p, @@ -2028,6 +2029,7 @@ int mlkem_keygen(sword16* s, sword16* t, sword16* e, const sword16* a, int k) * @param [out] t Public key vector of polynomials. * @param [in, out] prf XOF object. * @param [in] tp Temporary polynomial. + * @param [out] cacheA Matrix to keep the generated A in. May be NULL. * @param [in] k Number of polynomials in vector. * @param [in] rho Random seed to generate matrix A from. * @param [in, out] sigma Random seed to generate noise from. @@ -2037,12 +2039,13 @@ int mlkem_keygen(sword16* s, sword16* t, sword16* e, const sword16* a, int k) * @return Other negative value when a hash error occurred. */ int mlkem_keygen_seeds(sword16* s, sword16* t, MLKEM_PRF_T* prf, - sword16* tp, int k, byte* rho, byte* sigma) + sword16* tp, sword16* cacheA, int k, byte* rho, byte* sigma) { int i; int ret = 0; - sword16* ai = tp; sword16* e = tp; + sword16* a = (cacheA != NULL) ? cacheA : tp; + unsigned int aStride = (cacheA != NULL) ? MLKEM_N : 0; /* Transform private key. All of result used in public key calculation * Step 16: s_hat = NTT(s) */ @@ -2059,7 +2062,9 @@ int mlkem_keygen_seeds(sword16* s, sword16* t, MLKEM_PRF_T* prf, * each into the public polynomial. * Steps 4-6: generate A[i] * Step 18: ... A_hat o s_hat ... */ - ret = mlkem_gen_matrix_i_acc(prf, t + i * MLKEM_N, ai, s, k, rho, i, 0); + ret = mlkem_gen_matrix_i_acc(prf, t + i * MLKEM_N, + a + (unsigned int)i * (unsigned int)k * aStride, aStride, s, k, + rho, i, 0); if (ret != 0) { break; } @@ -2246,6 +2251,8 @@ int mlkem_encapsulate(const sword16* pub, sword16* u, sword16* v, * @param [in] msg Message to encapsulate. * @param [in] seed Random seed to generate matrix A from. * @param [in, out] coins Random seed to generate noise from. + * @param [in] cacheA Cached matrix A to transpose in place of + * generating it. May be NULL. * @return 0 on success. * @return MEMORY_E when dynamic memory allocation fails. Only possible when * WOLFSSL_SMALL_STACK is defined. @@ -2255,7 +2262,7 @@ int mlkem_encapsulate(const sword16* pub, sword16* u, sword16* v, */ int mlkem_encapsulate_seeds(const sword16* pub, MLKEM_PRF_T* prf, byte* c, const byte* cmp, int* fail, sword16* u, sword16* tp, sword16* y, int k, - const byte* msg, byte* seed, byte* coins) + const byte* msg, byte* seed, byte* coins, const sword16* cacheA) { int ret = 0; int i; @@ -2290,11 +2297,27 @@ int mlkem_encapsulate_seeds(const sword16* pub, MLKEM_PRF_T* prf, byte* c, for (i = 0; i < k; ++i) { int j; - /* Generate a vector of matrix A, a polynomial at a time, multiplying - * each into the u polynomial. */ - ret = mlkem_gen_matrix_i_acc(prf, u, a, y, k, seed, i, 1); - if (ret != 0) { - break; + if (cacheA != NULL) { + /* A_trans[i,j] is A[j,i] of the cached matrix. */ + for (j = 0; j < k; j++) { + const sword16* aij = cacheA + + ((unsigned int)j * (unsigned int)k + (unsigned int)i) * + MLKEM_N; + if (j == 0) { + mlkem_basemul_mont(u, aij, y); + } + else { + mlkem_basemul_mont_add(u, aij, y + j * MLKEM_N); + } + } + } + else { + /* Generate a vector of matrix A, a polynomial at a time, + * multiplying each into the u polynomial. */ + ret = mlkem_gen_matrix_i_acc(prf, u, a, 0, y, k, seed, i, 1); + if (ret != 0) { + break; + } } /* Inverse transform u polynomial. */ mlkem_invntt(u); @@ -4555,7 +4578,11 @@ int mlkem_gen_matrix(MLKEM_PRF_T* prf, sword16* a, int k, byte* seed, * * @param [in, out] prf XOF object. * @param [out] r Polynomial holding the multiplied result. - * @param [out] a Polynomial of uniform integers. Scratch. + * @param [out] a Polynomial of uniform integers. Scratch when + * aStride is 0, otherwise the vector to keep + * the generated polynomials in. + * @param [in] aStride Number of elements between polynomials of a. + * 0 to generate each into the same polynomial. * @param [in] v Vector of polynomials to multiply with. * @param [in] k Number of dimensions. k x k polynomials. * @param [in] seed Bytes to seed XOF generation. @@ -4566,7 +4593,8 @@ int mlkem_gen_matrix(MLKEM_PRF_T* prf, sword16* a, int k, byte* seed, * WOLFSSL_SMALL_STACK is defined. */ static int mlkem_gen_matrix_i_acc(MLKEM_PRF_T* prf, sword16* r, sword16* a, - const sword16* v, int k, byte* seed, int i, int transposed) + unsigned int aStride, const sword16* v, int k, byte* seed, int i, + int transposed) { #if defined(WOLFSSL_SMALL_STACK) && !defined(WOLFSSL_NO_MALLOC) byte* rand; @@ -4619,11 +4647,12 @@ static int mlkem_gen_matrix_i_acc(MLKEM_PRF_T* prf, sword16* r, sword16* a, } if (ret == 0) { unsigned int ctr; + sword16* aj = a + (unsigned int)j * aStride; /* Sample random bytes to create a polynomial. * Alg 7, Step 3 - implicitly counter is 0. * Alg 7, Step 4-16. */ - ctr = mlkem_rej_uniform_c(a, MLKEM_N, rand, GEN_MATRIX_SIZE); + ctr = mlkem_rej_uniform_c(aj, MLKEM_N, rand, GEN_MATRIX_SIZE); /* Create more blocks if too many rejected. * Alg 7, Step 4. */ while ((ret == 0) && (ctr < MLKEM_N)) { @@ -4631,7 +4660,7 @@ static int mlkem_gen_matrix_i_acc(MLKEM_PRF_T* prf, sword16* r, sword16* a, ret = mlkem_xof_squeezeblocks(prf, rand, 1); if (ret == 0) { /* Alg 7, Step 4-16. */ - ctr += mlkem_rej_uniform_c(a + ctr, MLKEM_N - ctr, rand, + ctr += mlkem_rej_uniform_c(aj + ctr, MLKEM_N - ctr, rand, XOF_BLOCK_SIZE); } } @@ -4640,10 +4669,10 @@ static int mlkem_gen_matrix_i_acc(MLKEM_PRF_T* prf, sword16* r, sword16* a, * Alg 13, Step 18: ... A_hat o s_hat ... * Alg 14, Step 18: ... A_hat_trans o y_hat ... */ if (j == 0) { - mlkem_basemul_mont(r, a, v); + mlkem_basemul_mont(r, aj, v); } else { - mlkem_basemul_mont_add(r, a, v + j * MLKEM_N); + mlkem_basemul_mont_add(r, aj, v + j * MLKEM_N); } } } diff --git a/wolfssl/wolfcrypt/wc_mlkem.h b/wolfssl/wolfcrypt/wc_mlkem.h index 81481439e3c..f2e7ce8fe05 100644 --- a/wolfssl/wolfcrypt/wc_mlkem.h +++ b/wolfssl/wolfcrypt/wc_mlkem.h @@ -556,7 +556,7 @@ int mlkem_keygen(sword16* priv, sword16* pub, sword16* e, const sword16* a, #else WOLFSSL_LOCAL int mlkem_keygen_seeds(sword16* priv, sword16* pub, MLKEM_PRF_T* prf, - sword16* e, int kp, byte* seed, byte* noiseSeed); + sword16* e, sword16* cacheA, int kp, byte* seed, byte* noiseSeed); #endif #ifndef WOLFSSL_MLKEM_ENCAPSULATE_SMALL_MEM WOLFSSL_LOCAL @@ -567,7 +567,7 @@ int mlkem_encapsulate(const sword16* pub, sword16* bp, sword16* v, WOLFSSL_LOCAL int mlkem_encapsulate_seeds(const sword16* pub, MLKEM_PRF_T* prf, byte* c, const byte* cmp, int* fail, sword16* bp, sword16* tp, sword16* sp, int kp, - const byte* msg, byte* seed, byte* coins); + const byte* msg, byte* seed, byte* coins, const sword16* cacheA); #endif WOLFSSL_LOCAL int mlkem_decapsulate(const sword16* priv, sword16* mp, sword16* bp, From e2c6a9a7c16fdaacc594b1904a86fdbd251053ef Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tobias=20Frauenschl=C3=A4ger?= Date: Sun, 6 Sep 2026 14:16:51 +0200 Subject: [PATCH 04/11] Skip the ML-KEM multiply after a failed squeeze and test the matrix cache mlkem_gen_matrix_i_acc() still multiplied the polynomial into the result after a rejection sampling squeeze failed, accumulating a partly sampled polynomial. Skip the multiply once the squeeze fails. Add mlkem_cache_a_decode_test(), which decodes a peer's public key into an object that already cached matrix A for another public seed and checks the peer recovers the shared secret, so the small memory options are covered too. Add functional and -Wconversion CI rows for WOLFSSL_MLKEM_CACHE_A with the small memory options, which nothing built. --- .github/configs/pq-all.json | 4 + .github/workflows/wolfCrypt-Wconversion.yml | 8 ++ wolfcrypt/src/wc_mlkem_poly.c | 18 ++-- wolfcrypt/test/test.c | 91 +++++++++++++++++++++ 4 files changed, 113 insertions(+), 8 deletions(-) diff --git a/.github/configs/pq-all.json b/.github/configs/pq-all.json index e760486a6c1..30901361320 100644 --- a/.github/configs/pq-all.json +++ b/.github/configs/pq-all.json @@ -238,6 +238,10 @@ "--enable-dilithium=yes,small,verify-only", "CPPFLAGS=-pedantic -Wdeclaration-after-statement -Wnull-dereference -DWOLFSSL_MLDSA_VERIFY_SMALLEST_MEM -DWOLFSSL_MLDSA_NO_ASN1 -DWOLFSSL_MLDSA_NO_LARGE_CODE"], "run": [["sh", "-c", "out=$(./wolfcrypt/benchmark/benchmark -ml-dsa 2>&1) || exit 1; printf '%s\\n' \"$out\"; if printf '%s\\n' \"$out\" | grep -qE 'VerifyCtx failed|Failed to load (public|private) key'; then echo 'ML-DSA verify or key load failed' >&2; exit 1; fi; for lvl in 44 65 87; do printf '%s\\n' \"$out\" | grep -qE \"^ML-DSA +$lvl +verify\" || { echo \"no ML-DSA-$lvl verify row reported\" >&2; exit 1; }; done"]]}, +{"name": "mlkem-cache-a-small-mem", "minutes": 2.1, + "comment": "Cached matrix A combined with the small memory ML-KEM key generation and encapsulation, which used to be rejected outright", + "configure": ["--enable-cryptonly", "--enable-mlkem", + "CPPFLAGS=-DWOLFSSL_MLKEM_CACHE_A -DWOLFSSL_MLKEM_MAKEKEY_SMALL_MEM -DWOLFSSL_MLKEM_ENCAPSULATE_SMALL_MEM"]}, {"name": "mldsa-dynamic-keys", "minutes": 1.8, "configure": ["--enable-intelasm", "--enable-sp-asm", "--enable-dilithium=yes", "CPPFLAGS=-DWOLFSSL_MLDSA_DYNAMIC_KEYS"]}, diff --git a/.github/workflows/wolfCrypt-Wconversion.yml b/.github/workflows/wolfCrypt-Wconversion.yml index 789886ee0f7..280aed04122 100644 --- a/.github/workflows/wolfCrypt-Wconversion.yml +++ b/.github/workflows/wolfCrypt-Wconversion.yml @@ -114,6 +114,14 @@ jobs: "--enable-xmss", "CPPFLAGS=-DWOLFSSL_MLDSA_SIGN_SMALL_MEM_PRECALC_A=1 -Wconversion -Warith-conversion -Wenum-conversion -Wfloat-conversion -Wsign-conversion -Wcast-qual -Wdeclaration-after-statement"], "check": false}, + {"name": "mlkem-cache-a-small-mem", "minutes": 1, + "configure": ["--enable-cryptonly", "--enable-all-crypto", + "--disable-examples", "--disable-benchmark", + "--disable-crypttests", "--enable-mlkem", + "--enable-slhdsa=yes,sha2", "--enable-mldsa", "--enable-lms", + "--enable-xmss", + "CPPFLAGS=-DWOLFSSL_MLKEM_CACHE_A -DWOLFSSL_MLKEM_MAKEKEY_SMALL_MEM -DWOLFSSL_MLKEM_ENCAPSULATE_SMALL_MEM -Wconversion -Warith-conversion -Wenum-conversion -Wfloat-conversion -Wsign-conversion -Wcast-qual"], + "check": false}, {"name": "precalc-a-no-int128", "minutes": 1, "configure": ["--enable-cryptonly", "--enable-all-crypto", "--disable-examples", "--disable-benchmark", diff --git a/wolfcrypt/src/wc_mlkem_poly.c b/wolfcrypt/src/wc_mlkem_poly.c index 4ea79f06049..83ec512ef3f 100644 --- a/wolfcrypt/src/wc_mlkem_poly.c +++ b/wolfcrypt/src/wc_mlkem_poly.c @@ -4665,14 +4665,16 @@ static int mlkem_gen_matrix_i_acc(MLKEM_PRF_T* prf, sword16* r, sword16* a, } } - /* Multiply the polynomial of the matrix into the result. - * Alg 13, Step 18: ... A_hat o s_hat ... - * Alg 14, Step 18: ... A_hat_trans o y_hat ... */ - if (j == 0) { - mlkem_basemul_mont(r, aj, v); - } - else { - mlkem_basemul_mont_add(r, aj, v + j * MLKEM_N); + if (ret == 0) { + /* Multiply the polynomial of the matrix into the result. + * Alg 13, Step 18: ... A_hat o s_hat ... + * Alg 14, Step 18: ... A_hat_trans o y_hat ... */ + if (j == 0) { + mlkem_basemul_mont(r, aj, v); + } + else { + mlkem_basemul_mont_add(r, aj, v + j * MLKEM_N); + } } } } diff --git a/wolfcrypt/test/test.c b/wolfcrypt/test/test.c index 75f2d90b683..c53ac4d522c 100644 --- a/wolfcrypt/test/test.c +++ b/wolfcrypt/test/test.c @@ -71060,6 +71060,77 @@ static wc_test_ret_t mlkem_decode_pub_test(void) } #endif +#if defined(WC_MLKEM_HAVE_NATIVE) && defined(WOLFSSL_MLKEM_CACHE_A) && \ + !defined(WOLFSSL_MLKEM_NO_MAKE_KEY) && \ + !defined(WOLFSSL_MLKEM_NO_ENCAPSULATE) && \ + !defined(WOLFSSL_MLKEM_NO_DECAPSULATE) && !defined(WC_NO_RNG) +/* Decoding a key must drop a matrix A cached for a different public seed. */ +static wc_test_ret_t mlkem_cache_a_decode_test(int type, WC_RNG* rng) +{ + wc_test_ret_t ret; + MlKemKey peer; + MlKemKey key; + byte pub[WC_ML_KEM_MAX_PUBLIC_KEY_SIZE]; + byte ct[WC_ML_KEM_MAX_CIPHER_TEXT_SIZE]; + byte ss[WC_ML_KEM_SS_SZ]; + byte ssDec[WC_ML_KEM_SS_SZ]; + word32 pubLen; + word32 ctLen; + int peerInit = 0; + int keyInit = 0; + + ret = wc_MlKemKey_Init(&peer, type, HEAP_HINT, INVALID_DEVID); + if (ret != 0) + ERROR_OUT(WC_TEST_RET_ENC_EC(ret), out); + peerInit = 1; + ret = wc_MlKemKey_Init(&key, type, HEAP_HINT, INVALID_DEVID); + if (ret != 0) + ERROR_OUT(WC_TEST_RET_ENC_EC(ret), out); + keyInit = 1; + + /* The peer owns the key pair the shared secret is checked against. */ + ret = wc_MlKemKey_MakeKey(&peer, rng); + if (ret != 0) + ERROR_OUT(WC_TEST_RET_ENC_EC(ret), out); + ret = wc_MlKemKey_PublicKeySize(&peer, &pubLen); + if (ret != 0) + ERROR_OUT(WC_TEST_RET_ENC_EC(ret), out); + ret = wc_MlKemKey_EncodePublicKey(&peer, pub, pubLen); + if (ret != 0) + ERROR_OUT(WC_TEST_RET_ENC_EC(ret), out); + + /* Populate this object's matrix A cache for an unrelated seed, then decode + * the peer's public key over the top of it. */ + ret = wc_MlKemKey_MakeKey(&key, rng); + if (ret != 0) + ERROR_OUT(WC_TEST_RET_ENC_EC(ret), out); + ret = wc_MlKemKey_DecodePublicKey(&key, pub, pubLen); + if (ret != 0) + ERROR_OUT(WC_TEST_RET_ENC_EC(ret), out); + + ret = wc_MlKemKey_CipherTextSize(&key, &ctLen); + if (ret != 0) + ERROR_OUT(WC_TEST_RET_ENC_EC(ret), out); + ret = wc_MlKemKey_Encapsulate(&key, ct, ss, rng); + if (ret != 0) + ERROR_OUT(WC_TEST_RET_ENC_EC(ret), out); + ret = wc_MlKemKey_Decapsulate(&peer, ssDec, ct, ctLen); + if (ret != 0) + ERROR_OUT(WC_TEST_RET_ENC_EC(ret), out); + + if (XMEMCMP(ss, ssDec, WC_ML_KEM_SS_SZ) != 0) + ERROR_OUT(WC_TEST_RET_ENC_NC, out); + + ret = 0; +out: + if (keyInit) + wc_MlKemKey_Free(&key); + if (peerInit) + wc_MlKemKey_Free(&peer); + return ret; +} +#endif + WOLFSSL_TEST_SUBROUTINE wc_test_ret_t mlkem_test(void) { wc_test_ret_t ret; @@ -71388,6 +71459,26 @@ WOLFSSL_TEST_SUBROUTINE wc_test_ret_t mlkem_test(void) goto out; #endif +#if defined(WOLFSSL_MLKEM_CACHE_A) && !defined(WOLFSSL_MLKEM_NO_MAKE_KEY) && \ + !defined(WOLFSSL_MLKEM_NO_ENCAPSULATE) && \ + !defined(WOLFSSL_MLKEM_NO_DECAPSULATE) && !defined(WC_NO_RNG) +#if !defined(WOLFSSL_NO_KYBER512) && !defined(WOLFSSL_NO_ML_KEM_512) + ret = mlkem_cache_a_decode_test(WC_ML_KEM_512, &rng); + if (ret != 0) + goto out; +#endif +#if !defined(WOLFSSL_NO_KYBER768) && !defined(WOLFSSL_NO_ML_KEM_768) + ret = mlkem_cache_a_decode_test(WC_ML_KEM_768, &rng); + if (ret != 0) + goto out; +#endif +#if !defined(WOLFSSL_NO_KYBER1024) && !defined(WOLFSSL_NO_ML_KEM_1024) + ret = mlkem_cache_a_decode_test(WC_ML_KEM_1024, &rng); + if (ret != 0) + goto out; +#endif +#endif + out: if (key_inited) From 5f890dffa0afeb214eef40e92a196285716f03eb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tobias=20Frauenschl=C3=A4ger?= Date: Sun, 6 Sep 2026 14:16:52 +0200 Subject: [PATCH 05/11] Default the ML-KEM parameter sets and operations for a mode-only option --enable-mlkem=original, =all and =ml-kem selected no size and no operation, so configure disabled every one of them and the build stopped at "No Kyber key size chosen". Default each group unless the user named a member of it. Give the new cache invalidation test the algorithm family guard and the small stack treatment the rest of the file uses, and align MLKEM_MAX_COMP_POLY_SZ with the 1024 parameter guard used everywhere else. Add CI rows for the configurations none of this was covered by: the original Kyber mode, a Kyber only cache-A build, and a small stack cache-A build. --- .github/configs/pq-all.json | 11 ++++ configure.ac | 18 ++++++ wolfcrypt/src/wc_mlkem.c | 2 +- wolfcrypt/test/test.c | 111 ++++++++++++++++++++++++++++++------ 4 files changed, 124 insertions(+), 18 deletions(-) diff --git a/.github/configs/pq-all.json b/.github/configs/pq-all.json index 30901361320..a66cc4d7aa3 100644 --- a/.github/configs/pq-all.json +++ b/.github/configs/pq-all.json @@ -238,6 +238,10 @@ "--enable-dilithium=yes,small,verify-only", "CPPFLAGS=-pedantic -Wdeclaration-after-statement -Wnull-dereference -DWOLFSSL_MLDSA_VERIFY_SMALLEST_MEM -DWOLFSSL_MLDSA_NO_ASN1 -DWOLFSSL_MLDSA_NO_LARGE_CODE"], "run": [["sh", "-c", "out=$(./wolfcrypt/benchmark/benchmark -ml-dsa 2>&1) || exit 1; printf '%s\\n' \"$out\"; if printf '%s\\n' \"$out\" | grep -qE 'VerifyCtx failed|Failed to load (public|private) key'; then echo 'ML-DSA verify or key load failed' >&2; exit 1; fi; for lvl in 44 65 87; do printf '%s\\n' \"$out\" | grep -qE \"^ML-DSA +$lvl +verify\" || { echo \"no ML-DSA-$lvl verify row reported\" >&2; exit 1; }; done"]]}, +{"name": "mlkem-cache-a-kyber-only", "minutes": 2.1, + "comment": "Cached matrix A with the original Kyber parameter set only, where the ML-KEM type ids are not compiled", + "configure": ["--enable-cryptonly", "--enable-mlkem", + "CPPFLAGS=-DWOLFSSL_MLKEM_KYBER -DWOLFSSL_NO_ML_KEM -DWOLFSSL_MLKEM_CACHE_A"]}, {"name": "mlkem-cache-a-small-mem", "minutes": 2.1, "comment": "Cached matrix A combined with the small memory ML-KEM key generation and encapsulation, which used to be rejected outright", "configure": ["--enable-cryptonly", "--enable-mlkem", @@ -245,6 +249,10 @@ {"name": "mldsa-dynamic-keys", "minutes": 1.8, "configure": ["--enable-intelasm", "--enable-sp-asm", "--enable-dilithium=yes", "CPPFLAGS=-DWOLFSSL_MLDSA_DYNAMIC_KEYS"]}, +{"name": "mlkem-cache-a-small-stack", "minutes": 1.8, + "comment": "Cached matrix A on a small stack build, where the key objects must come off the heap", + "configure": ["--enable-cryptonly", "--enable-mlkem", "--enable-smallstack", + "CPPFLAGS=-DWOLFSSL_MLKEM_CACHE_A"]}, {"name": "mldsa-verify-only-dynamic-keys", "minutes": 1.7, "configure": ["--disable-intelasm", "--enable-dilithium=44,65,87,verify-only", @@ -263,6 +271,9 @@ "--enable-experimental", "--enable-frodokem=no-976,no-1344,no-aes", "CPPFLAGS=-Wdeclaration-after-statement"]}, +{"name": "mlkem-original-mode", "minutes": 1.5, + "comment": "Selecting the original Kyber mode by name, which selects no parameter set or operation of its own", + "configure": ["--enable-cryptonly", "--enable-mlkem=original"]}, {"name": "frodokem-aarch64-armasm", "minutes": 0.8, "check": false, "cc": "ccache aarch64-linux-gnu-gcc", "comment": "Cross-build only (no ARM runner): catches ARM asm link/compile breaks. sve+sme+aes A-gen all live in the .S.", diff --git a/configure.ac b/configure.ac index 8504097da7f..79604ee3137 100644 --- a/configure.ac +++ b/configure.ac @@ -8819,6 +8819,24 @@ do esac done +# Default the parameter sets and operations when only a mode is named. +if test "$ENABLED_MLKEM" != "no" && test "$ENABLED_MLKEM512" = "" && \ + test "$ENABLED_MLKEM768" = "" && test "$ENABLED_MLKEM1024" = "" +then + ENABLED_MLKEM512=yes + ENABLED_MLKEM768=yes + ENABLED_MLKEM1024=yes +fi +if test "$ENABLED_MLKEM" != "no" && \ + test "$ENABLED_MLKEM_MAKE_KEY" = "no" && \ + test "$ENABLED_MLKEM_ENCAPSULATE" = "no" && \ + test "$ENABLED_MLKEM_DECAPSULATE" = "no" +then + ENABLED_MLKEM_MAKE_KEY=yes + ENABLED_MLKEM_ENCAPSULATE=yes + ENABLED_MLKEM_DECAPSULATE=yes +fi + # Selftest uses its own random.c which doesn't support SHA-512 DRBG # or runtime DRBG disable/enable APIs AS_IF([test "x$ENABLED_SELFTEST" = "xyes"], diff --git a/wolfcrypt/src/wc_mlkem.c b/wolfcrypt/src/wc_mlkem.c index d26c2858b37..bb36076286b 100644 --- a/wolfcrypt/src/wc_mlkem.c +++ b/wolfcrypt/src/wc_mlkem.c @@ -143,7 +143,7 @@ #error "Cannot use dynamic key buffers without malloc" #endif -#ifdef WOLFSSL_WC_ML_KEM_1024 +#if defined(WOLFSSL_KYBER1024) || defined(WOLFSSL_WC_ML_KEM_1024) #define MLKEM_MAX_COMP_POLY_SZ MLKEM_POLY_COMPRESSED_SZ(11) #else #define MLKEM_MAX_COMP_POLY_SZ MLKEM_POLY_COMPRESSED_SZ(10) diff --git a/wolfcrypt/test/test.c b/wolfcrypt/test/test.c index c53ac4d522c..9d261623436 100644 --- a/wolfcrypt/test/test.c +++ b/wolfcrypt/test/test.c @@ -71068,53 +71068,104 @@ static wc_test_ret_t mlkem_decode_pub_test(void) static wc_test_ret_t mlkem_cache_a_decode_test(int type, WC_RNG* rng) { wc_test_ret_t ret; - MlKemKey peer; - MlKemKey key; +#ifdef WOLFSSL_SMALL_STACK + MlKemKey* peer = NULL; + MlKemKey* key = NULL; + byte* pub = NULL; + byte* priv = NULL; + byte* ct = NULL; +#else + MlKemKey peer[1]; + MlKemKey key[1]; byte pub[WC_ML_KEM_MAX_PUBLIC_KEY_SIZE]; + byte priv[WC_ML_KEM_MAX_PRIVATE_KEY_SIZE]; byte ct[WC_ML_KEM_MAX_CIPHER_TEXT_SIZE]; +#endif byte ss[WC_ML_KEM_SS_SZ]; byte ssDec[WC_ML_KEM_SS_SZ]; word32 pubLen; + word32 privLen; word32 ctLen; int peerInit = 0; int keyInit = 0; - ret = wc_MlKemKey_Init(&peer, type, HEAP_HINT, INVALID_DEVID); +#ifdef WOLFSSL_SMALL_STACK + peer = (MlKemKey*)XMALLOC(sizeof(MlKemKey), HEAP_HINT, + DYNAMIC_TYPE_TMP_BUFFER); + key = (MlKemKey*)XMALLOC(sizeof(MlKemKey), HEAP_HINT, + DYNAMIC_TYPE_TMP_BUFFER); + pub = (byte*)XMALLOC(WC_ML_KEM_MAX_PUBLIC_KEY_SIZE, HEAP_HINT, + DYNAMIC_TYPE_TMP_BUFFER); + priv = (byte*)XMALLOC(WC_ML_KEM_MAX_PRIVATE_KEY_SIZE, HEAP_HINT, + DYNAMIC_TYPE_TMP_BUFFER); + ct = (byte*)XMALLOC(WC_ML_KEM_MAX_CIPHER_TEXT_SIZE, HEAP_HINT, + DYNAMIC_TYPE_TMP_BUFFER); + if (peer == NULL || key == NULL || pub == NULL || priv == NULL || + ct == NULL) + ERROR_OUT(WC_TEST_RET_ENC_NC, out); +#endif + + ret = wc_MlKemKey_Init(peer, type, HEAP_HINT, INVALID_DEVID); if (ret != 0) ERROR_OUT(WC_TEST_RET_ENC_EC(ret), out); peerInit = 1; - ret = wc_MlKemKey_Init(&key, type, HEAP_HINT, INVALID_DEVID); + ret = wc_MlKemKey_Init(key, type, HEAP_HINT, INVALID_DEVID); if (ret != 0) ERROR_OUT(WC_TEST_RET_ENC_EC(ret), out); keyInit = 1; /* The peer owns the key pair the shared secret is checked against. */ - ret = wc_MlKemKey_MakeKey(&peer, rng); + ret = wc_MlKemKey_MakeKey(peer, rng); if (ret != 0) ERROR_OUT(WC_TEST_RET_ENC_EC(ret), out); - ret = wc_MlKemKey_PublicKeySize(&peer, &pubLen); + ret = wc_MlKemKey_PublicKeySize(peer, &pubLen); if (ret != 0) ERROR_OUT(WC_TEST_RET_ENC_EC(ret), out); - ret = wc_MlKemKey_EncodePublicKey(&peer, pub, pubLen); + ret = wc_MlKemKey_EncodePublicKey(peer, pub, pubLen); if (ret != 0) ERROR_OUT(WC_TEST_RET_ENC_EC(ret), out); /* Populate this object's matrix A cache for an unrelated seed, then decode * the peer's public key over the top of it. */ - ret = wc_MlKemKey_MakeKey(&key, rng); + ret = wc_MlKemKey_MakeKey(key, rng); if (ret != 0) ERROR_OUT(WC_TEST_RET_ENC_EC(ret), out); - ret = wc_MlKemKey_DecodePublicKey(&key, pub, pubLen); + ret = wc_MlKemKey_DecodePublicKey(key, pub, pubLen); if (ret != 0) ERROR_OUT(WC_TEST_RET_ENC_EC(ret), out); - ret = wc_MlKemKey_CipherTextSize(&key, &ctLen); + ret = wc_MlKemKey_CipherTextSize(key, &ctLen); if (ret != 0) ERROR_OUT(WC_TEST_RET_ENC_EC(ret), out); - ret = wc_MlKemKey_Encapsulate(&key, ct, ss, rng); + ret = wc_MlKemKey_Encapsulate(key, ct, ss, rng); if (ret != 0) ERROR_OUT(WC_TEST_RET_ENC_EC(ret), out); - ret = wc_MlKemKey_Decapsulate(&peer, ssDec, ct, ctLen); + ret = wc_MlKemKey_Decapsulate(peer, ssDec, ct, ctLen); + if (ret != 0) + ERROR_OUT(WC_TEST_RET_ENC_EC(ret), out); + + if (XMEMCMP(ss, ssDec, WC_ML_KEM_SS_SZ) != 0) + ERROR_OUT(WC_TEST_RET_ENC_NC, out); + + /* Decoding a private key must drop the cached matrix too. */ + ret = wc_MlKemKey_PrivateKeySize(peer, &privLen); + if (ret != 0) + ERROR_OUT(WC_TEST_RET_ENC_EC(ret), out); + ret = wc_MlKemKey_EncodePrivateKey(peer, priv, privLen); + if (ret != 0) + ERROR_OUT(WC_TEST_RET_ENC_EC(ret), out); + + ret = wc_MlKemKey_MakeKey(key, rng); + if (ret != 0) + ERROR_OUT(WC_TEST_RET_ENC_EC(ret), out); + ret = wc_MlKemKey_DecodePrivateKey(key, priv, privLen); + if (ret != 0) + ERROR_OUT(WC_TEST_RET_ENC_EC(ret), out); + + ret = wc_MlKemKey_Encapsulate(peer, ct, ss, rng); + if (ret != 0) + ERROR_OUT(WC_TEST_RET_ENC_EC(ret), out); + ret = wc_MlKemKey_Decapsulate(key, ssDec, ct, ctLen); if (ret != 0) ERROR_OUT(WC_TEST_RET_ENC_EC(ret), out); @@ -71124,9 +71175,16 @@ static wc_test_ret_t mlkem_cache_a_decode_test(int type, WC_RNG* rng) ret = 0; out: if (keyInit) - wc_MlKemKey_Free(&key); + wc_MlKemKey_Free(key); if (peerInit) - wc_MlKemKey_Free(&peer); + wc_MlKemKey_Free(peer); +#ifdef WOLFSSL_SMALL_STACK + XFREE(ct, HEAP_HINT, DYNAMIC_TYPE_TMP_BUFFER); + XFREE(priv, HEAP_HINT, DYNAMIC_TYPE_TMP_BUFFER); + XFREE(pub, HEAP_HINT, DYNAMIC_TYPE_TMP_BUFFER); + XFREE(key, HEAP_HINT, DYNAMIC_TYPE_TMP_BUFFER); + XFREE(peer, HEAP_HINT, DYNAMIC_TYPE_TMP_BUFFER); +#endif return ret; } #endif @@ -71462,22 +71520,41 @@ WOLFSSL_TEST_SUBROUTINE wc_test_ret_t mlkem_test(void) #if defined(WOLFSSL_MLKEM_CACHE_A) && !defined(WOLFSSL_MLKEM_NO_MAKE_KEY) && \ !defined(WOLFSSL_MLKEM_NO_ENCAPSULATE) && \ !defined(WOLFSSL_MLKEM_NO_DECAPSULATE) && !defined(WC_NO_RNG) -#if !defined(WOLFSSL_NO_KYBER512) && !defined(WOLFSSL_NO_ML_KEM_512) +#ifndef WOLFSSL_NO_ML_KEM +#ifndef WOLFSSL_NO_ML_KEM_512 ret = mlkem_cache_a_decode_test(WC_ML_KEM_512, &rng); if (ret != 0) goto out; #endif -#if !defined(WOLFSSL_NO_KYBER768) && !defined(WOLFSSL_NO_ML_KEM_768) +#ifndef WOLFSSL_NO_ML_KEM_768 ret = mlkem_cache_a_decode_test(WC_ML_KEM_768, &rng); if (ret != 0) goto out; #endif -#if !defined(WOLFSSL_NO_KYBER1024) && !defined(WOLFSSL_NO_ML_KEM_1024) +#ifndef WOLFSSL_NO_ML_KEM_1024 ret = mlkem_cache_a_decode_test(WC_ML_KEM_1024, &rng); if (ret != 0) goto out; #endif #endif +#ifdef WOLFSSL_MLKEM_KYBER +#ifndef WOLFSSL_NO_KYBER512 + ret = mlkem_cache_a_decode_test(KYBER512, &rng); + if (ret != 0) + goto out; +#endif +#ifndef WOLFSSL_NO_KYBER768 + ret = mlkem_cache_a_decode_test(KYBER768, &rng); + if (ret != 0) + goto out; +#endif +#ifndef WOLFSSL_NO_KYBER1024 + ret = mlkem_cache_a_decode_test(KYBER1024, &rng); + if (ret != 0) + goto out; +#endif +#endif +#endif out: From 48d8ca313d0e529fd4310b1a3381cdef4f4d809d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tobias=20Frauenschl=C3=A4ger?= Date: Tue, 6 Oct 2026 16:45:34 +0200 Subject: [PATCH 06/11] Own the ML-KEM compare accumulator and prove its block is zeroed Register the encrypt block that small memory decapsulation compares in with the memory zero checker, so a build with WOLFSSL_CHECK_MEM_ZERO proves it is cleared. Own the comparison accumulator inside mlkem_encapsulate_seeds() rather than relying on its one caller to have zeroed it, and reject a null accumulator when a cipher text to compare against is supplied. The accumulator is only written once every block has been compared, and small memory decapsulation now starts it at mismatch like the full path does, so a comparison that never ran selects the implicit rejection secret. The accumulator and the per-block result are zeroized before returning, as the caller does with its own copy. The new -Wconversion row for cached matrix A passed --enable-all-crypto without --disable-intelasm. On x86_64 that turns on the Intel assembly by default and trips the guard that forbids it alongside the small memory options, as both sibling rows in that file already account for. Add a row for small memory encapsulation without decapsulation, which nothing built. Quote the parameter set and operation tests added to configure.ac so an unset variable cannot turn the comparison into a syntax error. --- .github/configs/pq-all.json | 16 ++++++++++++++ .github/workflows/wolfCrypt-Wconversion.yml | 5 +++-- configure.ac | 12 +++++------ wolfcrypt/src/wc_mlkem.c | 12 +++++++---- wolfcrypt/src/wc_mlkem_poly.c | 23 ++++++++++++++------- 5 files changed, 49 insertions(+), 19 deletions(-) diff --git a/.github/configs/pq-all.json b/.github/configs/pq-all.json index a66cc4d7aa3..7b4ba0241b4 100644 --- a/.github/configs/pq-all.json +++ b/.github/configs/pq-all.json @@ -238,6 +238,14 @@ "--enable-dilithium=yes,small,verify-only", "CPPFLAGS=-pedantic -Wdeclaration-after-statement -Wnull-dereference -DWOLFSSL_MLDSA_VERIFY_SMALLEST_MEM -DWOLFSSL_MLDSA_NO_ASN1 -DWOLFSSL_MLDSA_NO_LARGE_CODE"], "run": [["sh", "-c", "out=$(./wolfcrypt/benchmark/benchmark -ml-dsa 2>&1) || exit 1; printf '%s\\n' \"$out\"; if printf '%s\\n' \"$out\" | grep -qE 'VerifyCtx failed|Failed to load (public|private) key'; then echo 'ML-DSA verify or key load failed' >&2; exit 1; fi; for lvl in 44 65 87; do printf '%s\\n' \"$out\" | grep -qE \"^ML-DSA +$lvl +verify\" || { echo \"no ML-DSA-$lvl verify row reported\" >&2; exit 1; }; done"]]}, +{"name": "mlkem-small-mem-no-malloc", "minutes": 2.1, + "comment": "Heapless ML-KEM with both small memory options, the only build that uses the resized stack arrays including the cipher text compare block", + "configure": ["--enable-cryptonly", "--enable-mlkem", + "CPPFLAGS=-DWOLFSSL_NO_MALLOC -DWOLFSSL_MLKEM_MAKEKEY_SMALL_MEM -DWOLFSSL_MLKEM_ENCAPSULATE_SMALL_MEM -DWOLFSSL_MLKEM_CACHE_A"]}, +{"name": "mlkem-cache-a-makekey-small-mem", "minutes": 2.1, + "comment": "Cached matrix A written a polynomial at a time by the small memory key generation and read back by the full encapsulation, the other mixed layout", + "configure": ["--enable-cryptonly", "--enable-mlkem", + "CPPFLAGS=-DWOLFSSL_MLKEM_CACHE_A -DWOLFSSL_MLKEM_MAKEKEY_SMALL_MEM"]}, {"name": "mlkem-cache-a-kyber-only", "minutes": 2.1, "comment": "Cached matrix A with the original Kyber parameter set only, where the ML-KEM type ids are not compiled", "configure": ["--enable-cryptonly", "--enable-mlkem", @@ -249,6 +257,10 @@ {"name": "mldsa-dynamic-keys", "minutes": 1.8, "configure": ["--enable-intelasm", "--enable-sp-asm", "--enable-dilithium=yes", "CPPFLAGS=-DWOLFSSL_MLDSA_DYNAMIC_KEYS"]}, +{"name": "mlkem-cache-a-encapsulate-small-mem", "minutes": 1.8, + "comment": "Cached matrix A written by the full generator and read back by the small memory encapsulation, one of the two mixed layouts", + "configure": ["--enable-cryptonly", "--enable-mlkem", + "CPPFLAGS=-DWOLFSSL_MLKEM_CACHE_A -DWOLFSSL_MLKEM_ENCAPSULATE_SMALL_MEM"]}, {"name": "mlkem-cache-a-small-stack", "minutes": 1.8, "comment": "Cached matrix A on a small stack build, where the key objects must come off the heap", "configure": ["--enable-cryptonly", "--enable-mlkem", "--enable-smallstack", @@ -274,6 +286,10 @@ {"name": "mlkem-original-mode", "minutes": 1.5, "comment": "Selecting the original Kyber mode by name, which selects no parameter set or operation of its own", "configure": ["--enable-cryptonly", "--enable-mlkem=original"]}, +{"name": "mlkem-encapsulate-only-small-mem", "minutes": 1.2, + "comment": "Small memory encapsulation without decapsulation, where the cipher text block compare is not compiled", + "configure": ["--enable-cryptonly", "--enable-mlkem=make,enc", + "CPPFLAGS=-DWOLFSSL_MLKEM_ENCAPSULATE_SMALL_MEM"]}, {"name": "frodokem-aarch64-armasm", "minutes": 0.8, "check": false, "cc": "ccache aarch64-linux-gnu-gcc", "comment": "Cross-build only (no ARM runner): catches ARM asm link/compile breaks. sve+sme+aes A-gen all live in the .S.", diff --git a/.github/workflows/wolfCrypt-Wconversion.yml b/.github/workflows/wolfCrypt-Wconversion.yml index 280aed04122..204eda90a48 100644 --- a/.github/workflows/wolfCrypt-Wconversion.yml +++ b/.github/workflows/wolfCrypt-Wconversion.yml @@ -115,12 +115,13 @@ jobs: "CPPFLAGS=-DWOLFSSL_MLDSA_SIGN_SMALL_MEM_PRECALC_A=1 -Wconversion -Warith-conversion -Wenum-conversion -Wfloat-conversion -Wsign-conversion -Wcast-qual -Wdeclaration-after-statement"], "check": false}, {"name": "mlkem-cache-a-small-mem", "minutes": 1, - "configure": ["--enable-cryptonly", "--enable-all-crypto", + "configure": ["--disable-intelasm", "--enable-cryptonly", + "--enable-all-crypto", "--disable-examples", "--disable-benchmark", "--disable-crypttests", "--enable-mlkem", "--enable-slhdsa=yes,sha2", "--enable-mldsa", "--enable-lms", "--enable-xmss", - "CPPFLAGS=-DWOLFSSL_MLKEM_CACHE_A -DWOLFSSL_MLKEM_MAKEKEY_SMALL_MEM -DWOLFSSL_MLKEM_ENCAPSULATE_SMALL_MEM -Wconversion -Warith-conversion -Wenum-conversion -Wfloat-conversion -Wsign-conversion -Wcast-qual"], + "CPPFLAGS=-DWOLFSSL_MLKEM_CACHE_A -DWOLFSSL_MLKEM_MAKEKEY_SMALL_MEM -DWOLFSSL_MLKEM_ENCAPSULATE_SMALL_MEM -Wconversion -Warith-conversion -Wenum-conversion -Wfloat-conversion -Wsign-conversion -Wcast-qual -Wdeclaration-after-statement"], "check": false}, {"name": "precalc-a-no-int128", "minutes": 1, "configure": ["--enable-cryptonly", "--enable-all-crypto", diff --git a/configure.ac b/configure.ac index 79604ee3137..9d5287ac2c3 100644 --- a/configure.ac +++ b/configure.ac @@ -8820,17 +8820,17 @@ do done # Default the parameter sets and operations when only a mode is named. -if test "$ENABLED_MLKEM" != "no" && test "$ENABLED_MLKEM512" = "" && \ - test "$ENABLED_MLKEM768" = "" && test "$ENABLED_MLKEM1024" = "" +if test "x$ENABLED_MLKEM" != "xno" && test "x$ENABLED_MLKEM512" = "x" && \ + test "x$ENABLED_MLKEM768" = "x" && test "x$ENABLED_MLKEM1024" = "x" then ENABLED_MLKEM512=yes ENABLED_MLKEM768=yes ENABLED_MLKEM1024=yes fi -if test "$ENABLED_MLKEM" != "no" && \ - test "$ENABLED_MLKEM_MAKE_KEY" = "no" && \ - test "$ENABLED_MLKEM_ENCAPSULATE" = "no" && \ - test "$ENABLED_MLKEM_DECAPSULATE" = "no" +if test "x$ENABLED_MLKEM" != "xno" && \ + test "x$ENABLED_MLKEM_MAKE_KEY" = "xno" && \ + test "x$ENABLED_MLKEM_ENCAPSULATE" = "xno" && \ + test "x$ENABLED_MLKEM_DECAPSULATE" = "xno" then ENABLED_MLKEM_MAKE_KEY=yes ENABLED_MLKEM_ENCAPSULATE=yes diff --git a/wolfcrypt/src/wc_mlkem.c b/wolfcrypt/src/wc_mlkem.c index bb36076286b..eb5fee081c6 100644 --- a/wolfcrypt/src/wc_mlkem.c +++ b/wolfcrypt/src/wc_mlkem.c @@ -1378,7 +1378,7 @@ int wc_MlKemKey_SharedSecretSize(MlKemKey* key, word32* len) * @param [in] r Seed to feed to PRF when generating y, e1 and e2. * @param [out] c Calculated cipher text. NULL when comparing. * @param [in] cmp Cipher text to compare against. May be NULL. - * @param [in, out] fail Set to -1 when cipher text does not match cmp. + * @param [out] fail Set to -1 when cipher text does not match cmp. * Only used when cmp is not NULL. * @return 0 on success. * @return NOT_COMPILED_IN when key type is not supported. @@ -1644,7 +1644,13 @@ static int mlkemkey_encapsulate(MlKemKey* key, const byte* m, byte* r, byte* c, #ifdef WOLFSSL_MLKEM_ENCAPSULATE_SMALL_MEM /* block holds the cipher text re-encapsulated from the secret decrypted * message. With malloc it sits in the y allocation and is covered above. */ +#ifdef WOLFSSL_CHECK_MEM_ZERO + wc_MemZero_Add("mlkem encrypt block", block, sizeof(block)); +#endif ForceZero(block, sizeof(block)); +#ifdef WOLFSSL_CHECK_MEM_ZERO + wc_MemZero_Check(block, sizeof(block)); +#endif #endif #endif @@ -2275,16 +2281,14 @@ int wc_MlKemKey_Decapsulate(MlKemKey* key, unsigned char* ss, byte msg[WC_ML_KEM_SYM_SZ]; byte kr[2 * WC_ML_KEM_SYM_SZ + 1]; unsigned int i = 0; -#ifndef WOLFSSL_MLKEM_ENCAPSULATE_SMALL_MEM int fail = -1; /* mismatch until mlkem_cmp() says otherwise */ +#ifndef WOLFSSL_MLKEM_ENCAPSULATE_SMALL_MEM #if defined(WOLFSSL_SMALL_STACK) || \ (!defined(USE_INTEL_SPEEDUP) && !defined(WOLFSSL_NO_MALLOC)) byte* cmp = NULL; #else byte cmp[WC_ML_KEM_MAX_CIPHER_TEXT_SIZE]; #endif -#else - int fail = 0; #endif #endif diff --git a/wolfcrypt/src/wc_mlkem_poly.c b/wolfcrypt/src/wc_mlkem_poly.c index 83ec512ef3f..ba6353d4958 100644 --- a/wolfcrypt/src/wc_mlkem_poly.c +++ b/wolfcrypt/src/wc_mlkem_poly.c @@ -2274,6 +2274,10 @@ int mlkem_encapsulate_seeds(const sword16* pub, MLKEM_PRF_T* prf, byte* c, byte* cb = c; /* Number of bytes a polynomial of u is compressed into. */ unsigned int blockSz = MLKEM_POLY_COMPRESSED_SZ(MLKEM_COMP_10BITS); +#ifndef WOLFSSL_MLKEM_NO_DECAPSULATE + int cmpFail = 0; + int blockFail = 0; +#endif #if defined(WOLFSSL_KYBER1024) || defined(WOLFSSL_WC_ML_KEM_1024) if (k == WC_ML_KEM_1024_K) { @@ -2281,8 +2285,12 @@ int mlkem_encapsulate_seeds(const sword16* pub, MLKEM_PRF_T* prf, byte* c, } #endif + /* Comparing against a cipher text requires somewhere to accumulate the + * result. */ + if ((cmp != NULL) && (fail == NULL)) { + return BAD_FUNC_ARG; + } #ifdef WOLFSSL_MLKEM_NO_DECAPSULATE - (void)fail; if (cmp != NULL) { return BAD_FUNC_ARG; } @@ -2379,10 +2387,8 @@ int mlkem_encapsulate_seeds(const sword16* pub, MLKEM_PRF_T* prf, byte* c, } #ifndef WOLFSSL_MLKEM_NO_DECAPSULATE else { - int blockFail; - ret = mlkem_cmp(cb, cmp, (int)blockSz, &blockFail); - *fail |= blockFail; + cmpFail |= blockFail; if (ret != 0) { break; } @@ -2455,11 +2461,14 @@ int mlkem_encapsulate_seeds(const sword16* pub, MLKEM_PRF_T* prf, byte* c, } #ifndef WOLFSSL_MLKEM_NO_DECAPSULATE if ((ret == 0) && (cmp != NULL)) { - int blockFail; - ret = mlkem_cmp(cb, cmp, (int)blockSz, &blockFail); - *fail |= blockFail; + if (ret == 0) { + /* Only a comparison of every block may clear the caller's fail. */ + *fail = cmpFail | blockFail; + } } + ForceZero(&cmpFail, sizeof(cmpFail)); + ForceZero(&blockFail, sizeof(blockFail)); #endif /* Step 24: return c <- (c_1||c_2) */ From 4b6dfed20c1ac3bda1377c22fe9bccf5e6bdbaaf Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tobias=20Frauenschl=C3=A4ger?= Date: Sun, 6 Sep 2026 14:16:53 +0200 Subject: [PATCH 07/11] Require a cipher text buffer in the small memory ML-KEM encapsulation The cipher text is written or compared a block at a time, so c is dereferenced in both modes. Reject a null one rather than faulting, and carry the failure through the remaining steps instead of returning from the middle of the function. Record the invariant the encapsulation relies on when it compresses straight into the caller's cipher text, which carries no slack after its last block, and the aliasing each shared temporary polynomial requires: mlkem_keygen_seeds()'s tp must not alias s or t, and mlkem_encapsulate_seeds() shares u and tp between values that are each only live for part of the calculation. --- wolfcrypt/src/wc_mlkem_poly.c | 26 +++++++++++++------------- 1 file changed, 13 insertions(+), 13 deletions(-) diff --git a/wolfcrypt/src/wc_mlkem_poly.c b/wolfcrypt/src/wc_mlkem_poly.c index ba6353d4958..e68c08981b5 100644 --- a/wolfcrypt/src/wc_mlkem_poly.c +++ b/wolfcrypt/src/wc_mlkem_poly.c @@ -2028,7 +2028,7 @@ int mlkem_keygen(sword16* s, sword16* t, sword16* e, const sword16* a, int k) * @param [in, out] s Private key vector of polynomials. * @param [out] t Public key vector of polynomials. * @param [in, out] prf XOF object. - * @param [in] tp Temporary polynomial. + * @param [out] tp Temporary polynomial. Must not alias s or t. * @param [out] cacheA Matrix to keep the generated A in. May be NULL. * @param [in] k Number of polynomials in vector. * @param [in] rho Random seed to generate matrix A from. @@ -2242,10 +2242,10 @@ int mlkem_encapsulate(const sword16* pub, sword16* u, sword16* v, * @param [in, out] prf XOF object. * @param [out] c Cipher text, or one block of it when comparing. * @param [in] cmp Cipher text to compare against. May be NULL. - * @param [in, out] fail Set to -1 when cipher text does not match cmp. + * @param [out] fail Set to -1 when cipher text does not match cmp. * Only used when cmp is not NULL. - * @param [in, out] u Polynomial. - * @param [in, out] tp Polynomial. + * @param [in, out] u Polynomial. Holds v once u is encoded. + * @param [out] tp Temporary polynomial. Must not alias u or y. * @param [in, out] y Vector of polynomials. * @param [in] k Number of polynomials in vector. * @param [in] msg Message to encapsulate. @@ -2285,24 +2285,24 @@ int mlkem_encapsulate_seeds(const sword16* pub, MLKEM_PRF_T* prf, byte* c, } #endif - /* Comparing against a cipher text requires somewhere to accumulate the - * result. */ - if ((cmp != NULL) && (fail == NULL)) { - return BAD_FUNC_ARG; + if ((c == NULL) || ((cmp != NULL) && (fail == NULL))) { + ret = BAD_FUNC_ARG; } #ifdef WOLFSSL_MLKEM_NO_DECAPSULATE if (cmp != NULL) { - return BAD_FUNC_ARG; + ret = BAD_FUNC_ARG; } #endif - /* Transform y. All of result used in calculation of u and v. */ - for (i = 0; i < k; ++i) { - mlkem_ntt(y + i * MLKEM_N); + if (ret == 0) { + /* Transform y. All of result used in calculation of u and v. */ + for (i = 0; i < k; ++i) { + mlkem_ntt(y + i * MLKEM_N); + } } /* For each polynomial in the vectors. */ - for (i = 0; i < k; ++i) { + for (i = 0; (ret == 0) && (i < k); ++i) { int j; if (cacheA != NULL) { From b81f6de0f58c312269ded9e6db5427d0c6fc32d9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tobias=20Frauenschl=C3=A4ger?= Date: Sun, 6 Sep 2026 14:16:53 +0200 Subject: [PATCH 08/11] Reject an unsupported cipher text comparison and cover the block compare mlkemkey_encapsulate() takes cmp and fail, but only honours them when WOLFSSL_MLKEM_ENCAPSULATE_SMALL_MEM is defined; the default build discarded both and stored the cipher text as before. No caller gets this wrong today, but the failure mode is silent: a comparison that never runs leaves fail untouched, which reads as a match and would accept a tampered cipher text where the Fujisaki-Okamoto transform requires implicit rejection. Return BAD_FUNC_ARG instead. Tamper in a later block of u and in the trailing c_2 in test_wc_mlkem_decap_fo_reject. Decapsulation now compares the re-encapsulated cipher text with k + 1 separate calls rather than one over the whole buffer, and the only offsets the test used, bytes 0 and 32, both land in the first block of u: disabling the trailing c_2 comparison left the test green. Both new cases check for exactly K_bar = J(z || c), as the byte 32 case does, so a mask other than 0 or -1 mixing K and K_bar is caught too. Add a pq-all row that builds the small memory options into a library that links unit.test. The existing small memory rows are --enable-cryptonly, which does not build the API tests, so nothing ran that comparison through them. mlkem_test now tampers the same three places for every compiled parameter set, checking for K_bar with ML-KEM, so the cryptonly rows run the rejection path too, including the heapless one that compares in a stack block. Document the BAD_FUNC_ARG return that the argument checks now produce, and the 0 or -1 that fail is set to once every block has been compared. --- .github/configs/pq-all.json | 4 + tests/api/test_mlkem.c | 26 +++++- wolfcrypt/src/wc_mlkem.c | 13 ++- wolfcrypt/src/wc_mlkem_poly.c | 8 +- wolfcrypt/test/test.c | 153 ++++++++++++++++++++++++++++++++++ 5 files changed, 195 insertions(+), 9 deletions(-) diff --git a/.github/configs/pq-all.json b/.github/configs/pq-all.json index 7b4ba0241b4..f6cc26f0693 100644 --- a/.github/configs/pq-all.json +++ b/.github/configs/pq-all.json @@ -163,6 +163,10 @@ "configure": ["--enable-intelasm", "--enable-sp-asm", "--enable-dilithium=yes", "--enable-opensslextra", "CPPFLAGS=-DWOLFSSL_MLDSA_NO_ASN1"]}, +{"name": "mlkem-small-mem-api-tests", "minutes": 2.7, + "comment": "Small memory ML-KEM in a build that links unit.test: the cryptonly rows only run testwolfcrypt, so this is the only entry that exercises the per-block cipher text comparison through test_wc_mlkem_decap_fo_reject", + "configure": ["--enable-mlkem", + "CPPFLAGS=-DWOLFSSL_MLKEM_MAKEKEY_SMALL_MEM -DWOLFSSL_MLKEM_ENCAPSULATE_SMALL_MEM"]}, {"name": "mldsa-small-dynamic-keys", "minutes": 2.4, "configure": ["--disable-intelasm", "--enable-dilithium=yes,small", "CPPFLAGS=-DWOLFSSL_MLDSA_DYNAMIC_KEYS"]}, diff --git a/tests/api/test_mlkem.c b/tests/api/test_mlkem.c index 37adf7a67fe..a89e077dcb0 100644 --- a/tests/api/test_mlkem.c +++ b/tests/api/test_mlkem.c @@ -3981,6 +3981,7 @@ int test_wc_mlkem_decap_fo_reject(void) byte ssDec[WC_ML_KEM_SS_SZ]; byte ssTampered[WC_ML_KEM_SS_SZ]; word32 ctLen = 0; + int i; #ifdef WOLFSSL_SHAKE256 byte priv[WC_ML_KEM_MAX_PRIVATE_KEY_SIZE]; byte ssExpected[WC_ML_KEM_SS_SZ]; @@ -4044,7 +4045,6 @@ int test_wc_mlkem_decap_fo_reject(void) ExpectIntEQ(wc_Shake256_Final(&shake, ssExpected, WC_ML_KEM_SS_SZ), 0); wc_Shake256_Free(&shake); ExpectIntEQ(XMEMCMP(ssTampered, ssExpected, WC_ML_KEM_SS_SZ), 0); - ForceZero(priv, sizeof(priv)); #endif /* Tamper at byte 0: decapsulation must still return 0. We do NOT assert @@ -4059,6 +4059,30 @@ int test_wc_mlkem_decap_fo_reject(void) ExpectIntEQ(wc_MlKemKey_Decapsulate(key, ssTampered, ctTampered, ctLen), 0); PRIVATE_KEY_LOCK(); + /* Tamper in a later block of u and in c_2. */ + for (i = 0; (i < 2) && EXPECT_SUCCESS(); i++) { + XMEMCPY(ctTampered, ct, ctLen); + ctTampered[(i == 0) ? (ctLen - 200) : (ctLen - 1)] ^= 0x01; + XMEMSET(ssTampered, 0, sizeof(ssTampered)); + PRIVATE_KEY_UNLOCK(); + ExpectIntEQ(wc_MlKemKey_Decapsulate(key, ssTampered, ctTampered, ctLen), + 0); + PRIVATE_KEY_LOCK(); + ExpectIntNE(XMEMCMP(ssTampered, ss, WC_ML_KEM_SS_SZ), 0); + #ifdef WOLFSSL_SHAKE256 + ExpectIntEQ(wc_InitShake256(&shake, NULL, INVALID_DEVID), 0); + ExpectIntEQ(wc_Shake256_Update(&shake, + priv + privLen - WC_ML_KEM_SYM_SZ, WC_ML_KEM_SYM_SZ), 0); + ExpectIntEQ(wc_Shake256_Update(&shake, ctTampered, ctLen), 0); + ExpectIntEQ(wc_Shake256_Final(&shake, ssExpected, WC_ML_KEM_SS_SZ), 0); + wc_Shake256_Free(&shake); + ExpectIntEQ(XMEMCMP(ssTampered, ssExpected, WC_ML_KEM_SS_SZ), 0); + #endif + } +#ifdef WOLFSSL_SHAKE256 + ForceZero(priv, sizeof(priv)); +#endif + DoExpectIntEQ(wc_FreeRng(&rng), 0); wc_MlKemKey_Free(key); XFREE(key, NULL, DYNAMIC_TYPE_TMP_BUFFER); diff --git a/wolfcrypt/src/wc_mlkem.c b/wolfcrypt/src/wc_mlkem.c index eb5fee081c6..dde602f808a 100644 --- a/wolfcrypt/src/wc_mlkem.c +++ b/wolfcrypt/src/wc_mlkem.c @@ -1376,11 +1376,12 @@ int wc_MlKemKey_SharedSecretSize(MlKemKey* key, word32* len) * @param [in] key ML-KEM key object. * @param [in] m Random bytes. * @param [in] r Seed to feed to PRF when generating y, e1 and e2. - * @param [out] c Calculated cipher text. NULL when comparing. + * @param [out] c Calculated cipher text. Unused when comparing. * @param [in] cmp Cipher text to compare against. May be NULL. - * @param [out] fail Set to -1 when cipher text does not match cmp. - * Only used when cmp is not NULL. + * @param [out] fail Set to 0 on a match and -1 on a mismatch, once every + * block is compared. Only used when cmp is not NULL. * @return 0 on success. + * @return BAD_FUNC_ARG when cmp is not NULL and comparing is not compiled in. * @return NOT_COMPILED_IN when key type is not supported. */ static int mlkemkey_encapsulate(MlKemKey* key, const byte* m, byte* r, byte* c, @@ -1466,6 +1467,11 @@ static int mlkemkey_encapsulate(MlKemKey* key, const byte* m, byte* r, byte* c, ret = NOT_COMPILED_IN; break; } +#ifndef WOLFSSL_MLKEM_ENCAPSULATE_SMALL_MEM + if ((ret == 0) && (cmp != NULL)) { + ret = BAD_FUNC_ARG; + } +#endif #ifndef WOLFSSL_NO_MALLOC if (ret == 0) { @@ -1619,7 +1625,6 @@ static int mlkemkey_encapsulate(MlKemKey* key, const byte* m, byte* r, byte* c, } #endif } - (void)cmp; (void)fail; #endif /* !WOLFSSL_MLKEM_ENCAPSULATE_SMALL_MEM */ diff --git a/wolfcrypt/src/wc_mlkem_poly.c b/wolfcrypt/src/wc_mlkem_poly.c index e68c08981b5..7e2e1e28a02 100644 --- a/wolfcrypt/src/wc_mlkem_poly.c +++ b/wolfcrypt/src/wc_mlkem_poly.c @@ -2242,8 +2242,8 @@ int mlkem_encapsulate(const sword16* pub, sword16* u, sword16* v, * @param [in, out] prf XOF object. * @param [out] c Cipher text, or one block of it when comparing. * @param [in] cmp Cipher text to compare against. May be NULL. - * @param [out] fail Set to -1 when cipher text does not match cmp. - * Only used when cmp is not NULL. + * @param [out] fail Set to 0 on a match and -1 on a mismatch, once + * every block is compared. Only used when cmp is set. * @param [in, out] u Polynomial. Holds v once u is encoded. * @param [out] tp Temporary polynomial. Must not alias u or y. * @param [in, out] y Vector of polynomials. @@ -2254,10 +2254,10 @@ int mlkem_encapsulate(const sword16* pub, sword16* u, sword16* v, * @param [in] cacheA Cached matrix A to transpose in place of * generating it. May be NULL. * @return 0 on success. + * @return BAD_FUNC_ARG when c is NULL, or cmp is not NULL and either fail is + * NULL or decapsulation is not compiled in. * @return MEMORY_E when dynamic memory allocation fails. Only possible when * WOLFSSL_SMALL_STACK is defined. - * @return BAD_FUNC_ARG when cmp is not NULL and decapsulation is not compiled - * in. * @return Other negative value when a hash error occurred. */ int mlkem_encapsulate_seeds(const sword16* pub, MLKEM_PRF_T* prf, byte* c, diff --git a/wolfcrypt/test/test.c b/wolfcrypt/test/test.c index 9d261623436..854c443163d 100644 --- a/wolfcrypt/test/test.c +++ b/wolfcrypt/test/test.c @@ -71189,6 +71189,120 @@ static wc_test_ret_t mlkem_cache_a_decode_test(int type, WC_RNG* rng) } #endif +#if defined(WC_MLKEM_HAVE_NATIVE) && !defined(WOLFSSL_MLKEM_NO_MAKE_KEY) && \ + !defined(WOLFSSL_MLKEM_NO_ENCAPSULATE) && \ + !defined(WOLFSSL_MLKEM_NO_DECAPSULATE) && !defined(WC_NO_RNG) +/* A tampered cipher text must decapsulate to the implicit rejection secret. */ +static wc_test_ret_t mlkem_decap_reject_test(int type, WC_RNG* rng) +{ + wc_test_ret_t ret; +#ifdef WOLFSSL_SMALL_STACK + MlKemKey* key = NULL; + byte* priv = NULL; + byte* ct = NULL; + byte* ctBad = NULL; +#else + MlKemKey key[1]; + byte priv[WC_ML_KEM_MAX_PRIVATE_KEY_SIZE]; + byte ct[WC_ML_KEM_MAX_CIPHER_TEXT_SIZE]; + byte ctBad[WC_ML_KEM_MAX_CIPHER_TEXT_SIZE]; +#endif + byte ss[WC_ML_KEM_SS_SZ]; + byte ssDec[WC_ML_KEM_SS_SZ]; +#ifdef WOLFSSL_SHAKE256 + wc_Shake shake; + byte kBar[WC_ML_KEM_SS_SZ]; +#endif + word32 privLen = 0; + word32 ctLen = 0; + word32 pos[3]; + int keyInit = 0; + int i; + +#ifdef WOLFSSL_SMALL_STACK + key = (MlKemKey*)XMALLOC(sizeof(MlKemKey), HEAP_HINT, + DYNAMIC_TYPE_TMP_BUFFER); + priv = (byte*)XMALLOC(WC_ML_KEM_MAX_PRIVATE_KEY_SIZE, HEAP_HINT, + DYNAMIC_TYPE_TMP_BUFFER); + ct = (byte*)XMALLOC(WC_ML_KEM_MAX_CIPHER_TEXT_SIZE, HEAP_HINT, + DYNAMIC_TYPE_TMP_BUFFER); + ctBad = (byte*)XMALLOC(WC_ML_KEM_MAX_CIPHER_TEXT_SIZE, HEAP_HINT, + DYNAMIC_TYPE_TMP_BUFFER); + if (key == NULL || priv == NULL || ct == NULL || ctBad == NULL) + ERROR_OUT(WC_TEST_RET_ENC_NC, out); +#endif + + ret = wc_MlKemKey_Init(key, type, HEAP_HINT, INVALID_DEVID); + if (ret != 0) + ERROR_OUT(WC_TEST_RET_ENC_EC(ret), out); + keyInit = 1; + ret = wc_MlKemKey_MakeKey(key, rng); + if (ret != 0) + ERROR_OUT(WC_TEST_RET_ENC_EC(ret), out); + ret = wc_MlKemKey_PrivateKeySize(key, &privLen); + if (ret != 0) + ERROR_OUT(WC_TEST_RET_ENC_EC(ret), out); + PRIVATE_KEY_UNLOCK(); + ret = wc_MlKemKey_EncodePrivateKey(key, priv, privLen); + PRIVATE_KEY_LOCK(); + if (ret != 0) + ERROR_OUT(WC_TEST_RET_ENC_EC(ret), out); + ret = wc_MlKemKey_CipherTextSize(key, &ctLen); + if (ret != 0) + ERROR_OUT(WC_TEST_RET_ENC_EC(ret), out); + ret = wc_MlKemKey_Encapsulate(key, ct, ss, rng); + if (ret != 0) + ERROR_OUT(WC_TEST_RET_ENC_EC(ret), out); + + /* The first and a later block of u, and c_2. */ + pos[0] = 32; + pos[1] = ctLen - 200; + pos[2] = ctLen - 1; + for (i = 0; i < 3; i++) { + XMEMCPY(ctBad, ct, ctLen); + ctBad[pos[i]] ^= 0x01; + PRIVATE_KEY_UNLOCK(); + ret = wc_MlKemKey_Decapsulate(key, ssDec, ctBad, ctLen); + PRIVATE_KEY_LOCK(); + if (ret != 0) + ERROR_OUT(WC_TEST_RET_ENC_EC(ret), out); + if (XMEMCMP(ssDec, ss, WC_ML_KEM_SS_SZ) == 0) + ERROR_OUT(WC_TEST_RET_ENC_I(i), out); + #ifdef WOLFSSL_SHAKE256 + if ((type != WC_ML_KEM_512) && (type != WC_ML_KEM_768) && + (type != WC_ML_KEM_1024)) + continue; + /* K_bar = SHAKE256(z || c) with z the last 32 bytes of the key. */ + ret = wc_InitShake256(&shake, HEAP_HINT, INVALID_DEVID); + if (ret != 0) + ERROR_OUT(WC_TEST_RET_ENC_EC(ret), out); + ret = wc_Shake256_Update(&shake, priv + privLen - WC_ML_KEM_SYM_SZ, + WC_ML_KEM_SYM_SZ); + if (ret == 0) + ret = wc_Shake256_Update(&shake, ctBad, ctLen); + if (ret == 0) + ret = wc_Shake256_Final(&shake, kBar, WC_ML_KEM_SS_SZ); + wc_Shake256_Free(&shake); + if (ret != 0) + ERROR_OUT(WC_TEST_RET_ENC_EC(ret), out); + if (XMEMCMP(ssDec, kBar, WC_ML_KEM_SS_SZ) != 0) + ERROR_OUT(WC_TEST_RET_ENC_I(i), out); + #endif + } + +out: + if (keyInit) + wc_MlKemKey_Free(key); +#ifdef WOLFSSL_SMALL_STACK + XFREE(ctBad, HEAP_HINT, DYNAMIC_TYPE_TMP_BUFFER); + XFREE(ct, HEAP_HINT, DYNAMIC_TYPE_TMP_BUFFER); + XFREE(priv, HEAP_HINT, DYNAMIC_TYPE_TMP_BUFFER); + XFREE(key, HEAP_HINT, DYNAMIC_TYPE_TMP_BUFFER); +#endif + return ret; +} +#endif + WOLFSSL_TEST_SUBROUTINE wc_test_ret_t mlkem_test(void) { wc_test_ret_t ret; @@ -71556,6 +71670,45 @@ WOLFSSL_TEST_SUBROUTINE wc_test_ret_t mlkem_test(void) #endif #endif +#if !defined(WOLFSSL_MLKEM_NO_MAKE_KEY) && \ + !defined(WOLFSSL_MLKEM_NO_ENCAPSULATE) && \ + !defined(WOLFSSL_MLKEM_NO_DECAPSULATE) && !defined(WC_NO_RNG) +#ifndef WOLFSSL_NO_ML_KEM +#ifndef WOLFSSL_NO_ML_KEM_512 + ret = mlkem_decap_reject_test(WC_ML_KEM_512, &rng); + if (ret != 0) + goto out; +#endif +#ifndef WOLFSSL_NO_ML_KEM_768 + ret = mlkem_decap_reject_test(WC_ML_KEM_768, &rng); + if (ret != 0) + goto out; +#endif +#ifndef WOLFSSL_NO_ML_KEM_1024 + ret = mlkem_decap_reject_test(WC_ML_KEM_1024, &rng); + if (ret != 0) + goto out; +#endif +#endif +#ifdef WOLFSSL_MLKEM_KYBER +#ifndef WOLFSSL_NO_KYBER512 + ret = mlkem_decap_reject_test(KYBER512, &rng); + if (ret != 0) + goto out; +#endif +#ifndef WOLFSSL_NO_KYBER768 + ret = mlkem_decap_reject_test(KYBER768, &rng); + if (ret != 0) + goto out; +#endif +#ifndef WOLFSSL_NO_KYBER1024 + ret = mlkem_decap_reject_test(KYBER1024, &rng); + if (ret != 0) + goto out; +#endif +#endif +#endif + out: if (key_inited) From 9fbb767da78c4312ba5745068e1eb492982b18dd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tobias=20Frauenschl=C3=A4ger?= Date: Sun, 6 Sep 2026 18:24:06 +0200 Subject: [PATCH 09/11] Keep the matrix generated during encapsulation when caching is enabled WOLFSSL_MLKEM_CACHE_A only ever populated the cache in key generation: MLKEM_FLAG_A_SET is set by wc_MlKemKey_MakeKeyWithRandom and cleared by both decode paths, so a key object built by DecodePublicKey or DecodePrivateKey never had a cached matrix and regenerated all k*k polynomials of A, each one a SHAKE-128 absorb and squeeze, on every encapsulation and decapsulation. That is the common shape for a peer's key. The small memory generator already takes a stride, so keeping the matrix costs nothing extra: at row i of the transpose the polynomials are A[j,i], which land at the same offsets key generation would have written. Pass the key's matrix as the destination when nothing is cached yet, and flag it only once the whole matrix has been generated without error. With a decoded ML-KEM-768 key, repeated encapsulation goes from 101,000 to 210,000 operations per second and decapsulation from 86,000 to 154,000, with no change in peak heap. The full memory encapsulation generated the transposed matrix into its own scratch. When nothing is cached and the key has room for the matrix, generate A into the key instead, flag it, and take the existing transpose copy. Repeated decapsulation with a decoded ML-KEM-768 key goes from 84,000 to 147,000 operations per second, the same as with a key that made its own key pair. Reserve and zeroize the cipher text comparison block only where it can be used. It is written when a cipher text is supplied to compare against, which only decapsulation does, but the heapless layout declared it unconditionally and ForceZeroed all 360 bytes of it on every encapsulation - in the build whose whole purpose is a smaller footprint. With WOLFSSL_MLKEM_DYNAMIC_KEYS the matrix is only kept when the key already has one allocated by key generation. A key that never had one streams the matrix as before, so encapsulating to a peer's key does not grow the heap. The decode test now runs two rounds per decoded key, so the second reads the matrix the first kept, and the white-box drives every stride the generator's callers pass. A pq-all row builds the dynamic key combination. --- .github/configs/pq-all.json | 4 ++ .../test_mlkem_poly_hash_fault_whitebox.c | 17 +++--- wolfcrypt/src/wc_mlkem.c | 58 ++++++++++++++----- wolfcrypt/src/wc_mlkem_poly.c | 16 ++++- wolfcrypt/test/test.c | 38 +++++++----- wolfssl/wolfcrypt/wc_mlkem.h | 7 ++- 6 files changed, 99 insertions(+), 41 deletions(-) diff --git a/.github/configs/pq-all.json b/.github/configs/pq-all.json index f6cc26f0693..b2250e0bdd0 100644 --- a/.github/configs/pq-all.json +++ b/.github/configs/pq-all.json @@ -294,6 +294,10 @@ "comment": "Small memory encapsulation without decapsulation, where the cipher text block compare is not compiled", "configure": ["--enable-cryptonly", "--enable-mlkem=make,enc", "CPPFLAGS=-DWOLFSSL_MLKEM_ENCAPSULATE_SMALL_MEM"]}, +{"name": "mlkem-cache-a-dynamic-keys-small-mem", "minutes": 1.2, + "comment": "Cached matrix A with dynamic keys and the small memory options: a decoded key refills a matrix the key already allocated and otherwise streams it", + "configure": ["--enable-cryptonly", "--enable-mlkem", + "CPPFLAGS=-DWOLFSSL_MLKEM_CACHE_A -DWOLFSSL_MLKEM_DYNAMIC_KEYS -DWOLFSSL_MLKEM_MAKEKEY_SMALL_MEM -DWOLFSSL_MLKEM_ENCAPSULATE_SMALL_MEM"]}, {"name": "frodokem-aarch64-armasm", "minutes": 0.8, "check": false, "cc": "ccache aarch64-linux-gnu-gcc", "comment": "Cross-build only (no ARM runner): catches ARM asm link/compile breaks. sve+sme+aes A-gen all live in the .S.", diff --git a/tests/unit-mcdc/test_mlkem_poly_hash_fault_whitebox.c b/tests/unit-mcdc/test_mlkem_poly_hash_fault_whitebox.c index 9e9882d55ee..293102df694 100644 --- a/tests/unit-mcdc/test_mlkem_poly_hash_fault_whitebox.c +++ b/tests/unit-mcdc/test_mlkem_poly_hash_fault_whitebox.c @@ -349,11 +349,7 @@ static void wb_gen_matrix_c_continuation(void) #endif /* ------------------------------------------------------------------------- * - * mlkem_gen_matrix_i_acc(): the small-memory single-row generator, which - * multiplies each polynomial in as it is generated. Both stride forms are - * driven: aStride 0 regenerates into one scratch polynomial (the encapsulate - * caller) and aStride MLKEM_N keeps the whole vector (key generation with the - * matrix cached). + * mlkem_gen_matrix_i_acc(): the small-memory single-row generator. * ------------------------------------------------------------------------- */ #if defined(WOLFSSL_HAVE_MLKEM) && \ !(defined(WOLFSSL_ARMASM) && defined(__aarch64__)) && \ @@ -366,7 +362,7 @@ static void wb_gen_matrix_c_continuation(void) static void wb_gen_matrix_i_rows(void) { MLKEM_PRF_T prf; - static sword16 a[WB_K * MLKEM_N]; + static sword16 a[WB_K * WB_K * MLKEM_N]; static sword16 v[WB_K * MLKEM_N]; static sword16 r[MLKEM_N]; byte seed[WC_ML_KEM_SYM_SZ + 2]; @@ -394,7 +390,14 @@ static void wb_gen_matrix_i_rows(void) (void)mlkem_gen_matrix_i_acc(&prf, r, a, 0, v, WB_K, seed, 0, 0); mcdc_fh_disarm(); mcdc_fh_arm(n); - (void)mlkem_gen_matrix_i_acc(&prf, r, a, MLKEM_N, v, WB_K, seed, 1, 1); + (void)mlkem_gen_matrix_i_acc(&prf, r, a, 0, v, WB_K, seed, 1, 1); + mcdc_fh_disarm(); + mcdc_fh_arm(n); + (void)mlkem_gen_matrix_i_acc(&prf, r, a, MLKEM_N, v, WB_K, seed, 1, 0); + mcdc_fh_disarm(); + mcdc_fh_arm(n); + (void)mlkem_gen_matrix_i_acc(&prf, r, a, WB_K * MLKEM_N, v, WB_K, + seed, 1, 1); mcdc_fh_disarm(); } diff --git a/wolfcrypt/src/wc_mlkem.c b/wolfcrypt/src/wc_mlkem.c index dde602f808a..862e70eb32a 100644 --- a/wolfcrypt/src/wc_mlkem.c +++ b/wolfcrypt/src/wc_mlkem.c @@ -64,6 +64,7 @@ * MlKemKey is 8KB larger but decapsulation is significantly faster. * Turn on when performing make key and decapsulation with same object. * Also usable with the small memory options. + * Encapsulation also keeps A for a decoded key that has room for it. * * WOLFSSL_MLKEM_DYNAMIC_KEYS Default: OFF * Dynamically allocates private and public key buffers instead of using @@ -1404,8 +1405,10 @@ static int mlkemkey_encapsulate(MlKemKey* key, const byte* m, byte* r, byte* c, sword16 y[((WC_ML_KEM_MAX_K + 3) * WC_ML_KEM_MAX_K + 3) * MLKEM_N]; #else sword16 y[(WC_ML_KEM_MAX_K + 2) * MLKEM_N]; +#ifndef WOLFSSL_MLKEM_NO_DECAPSULATE byte block[MLKEM_MAX_COMP_POLY_SZ]; #endif +#endif #endif sword16* u = 0; #ifndef WOLFSSL_MLKEM_ENCAPSULATE_SMALL_MEM @@ -1415,8 +1418,11 @@ static int mlkemkey_encapsulate(MlKemKey* key, const byte* m, byte* r, byte* c, #ifdef WOLFSSL_MLKEM_CACHE_A const sword16* cachedA = ((key->flags & MLKEM_FLAG_A_SET) != 0) ? key->a : NULL; + /* A decoded key has no cached matrix, so keep the one generated here. */ + sword16* fillA = NULL; #else const sword16* cachedA = NULL; + sword16* fillA = NULL; #endif #endif @@ -1516,6 +1522,16 @@ static int mlkemkey_encapsulate(MlKemKey* key, const byte* m, byte* r, byte* c, ret = mlkem_get_noise(&key->prf, (int)k, y, e1, e2, r); } #ifdef WOLFSSL_MLKEM_CACHE_A + if ((ret == 0) && ((key->flags & MLKEM_FLAG_A_SET) == 0)) { + sword16* fillA = key->a; + + if (fillA != NULL) { + ret = mlkem_gen_matrix(&key->prf, fillA, (int)k, key->pubSeed, 0); + } + if ((ret == 0) && (fillA != NULL)) { + key->flags |= MLKEM_FLAG_A_SET; + } + } if ((ret == 0) && ((key->flags & MLKEM_FLAG_A_SET) != 0)) { unsigned int i; /* Transpose matrix. @@ -1568,8 +1584,10 @@ static int mlkemkey_encapsulate(MlKemKey* key, const byte* m, byte* r, byte* c, /* Cipher text is compared a block at a time as it is calculated. */ #ifndef WOLFSSL_NO_MALLOC cb = (byte*)(a + MLKEM_N); -#else +#elif !defined(WOLFSSL_MLKEM_NO_DECAPSULATE) cb = block; +#else + ret = BAD_FUNC_ARG; #endif } @@ -1577,8 +1595,20 @@ static int mlkemkey_encapsulate(MlKemKey* key, const byte* m, byte* r, byte* c, * Steps 13-17: generate e_1 and e_2 * Steps 18-19, 21: calculate u and v * Steps 22-24: c <- (c_1||c_2) */ - ret = mlkem_encapsulate_seeds(key->pub, &key->prf, cb, cmp, fail, u, a, - y, (int)k, m, key->pubSeed, r, cachedA); +#ifdef WOLFSSL_MLKEM_CACHE_A + if ((ret == 0) && (cachedA == NULL)) { + fillA = key->a; + } +#endif + if (ret == 0) { + ret = mlkem_encapsulate_seeds(key->pub, &key->prf, cb, cmp, fail, + u, a, y, (int)k, m, key->pubSeed, r, cachedA, fillA); + } +#ifdef WOLFSSL_MLKEM_CACHE_A + if ((ret == 0) && (fillA != NULL)) { + key->flags |= MLKEM_FLAG_A_SET; + } +#endif } (void)compVecSz; #endif /* WOLFSSL_MLKEM_ENCAPSULATE_SMALL_MEM */ @@ -1646,16 +1676,18 @@ static int mlkemkey_encapsulate(MlKemKey* key, const byte* m, byte* r, byte* c, #ifdef WOLFSSL_CHECK_MEM_ZERO wc_MemZero_Check(y, sizeof(y)); #endif -#ifdef WOLFSSL_MLKEM_ENCAPSULATE_SMALL_MEM - /* block holds the cipher text re-encapsulated from the secret decrypted - * message. With malloc it sits in the y allocation and is covered above. */ -#ifdef WOLFSSL_CHECK_MEM_ZERO - wc_MemZero_Add("mlkem encrypt block", block, sizeof(block)); -#endif - ForceZero(block, sizeof(block)); -#ifdef WOLFSSL_CHECK_MEM_ZERO - wc_MemZero_Check(block, sizeof(block)); -#endif +#if defined(WOLFSSL_MLKEM_ENCAPSULATE_SMALL_MEM) && \ + !defined(WOLFSSL_MLKEM_NO_DECAPSULATE) + /* block holds cipher text re-encapsulated from the secret message. */ + if (cmp != NULL) { + #ifdef WOLFSSL_CHECK_MEM_ZERO + wc_MemZero_Add("mlkem encrypt block", block, sizeof(block)); + #endif + ForceZero(block, sizeof(block)); + #ifdef WOLFSSL_CHECK_MEM_ZERO + wc_MemZero_Check(block, sizeof(block)); + #endif + } #endif #endif diff --git a/wolfcrypt/src/wc_mlkem_poly.c b/wolfcrypt/src/wc_mlkem_poly.c index 7e2e1e28a02..7d768c91566 100644 --- a/wolfcrypt/src/wc_mlkem_poly.c +++ b/wolfcrypt/src/wc_mlkem_poly.c @@ -2253,6 +2253,9 @@ int mlkem_encapsulate(const sword16* pub, sword16* u, sword16* v, * @param [in, out] coins Random seed to generate noise from. * @param [in] cacheA Cached matrix A to transpose in place of * generating it. May be NULL. + * @param [out] fillA Matrix to keep the generated polynomials of A in, + * laid out as key generation leaves them. May be + * NULL. Ignored when cacheA is not NULL. * @return 0 on success. * @return BAD_FUNC_ARG when c is NULL, or cmp is not NULL and either fail is * NULL or decapsulation is not compiled in. @@ -2262,7 +2265,8 @@ int mlkem_encapsulate(const sword16* pub, sword16* u, sword16* v, */ int mlkem_encapsulate_seeds(const sword16* pub, MLKEM_PRF_T* prf, byte* c, const byte* cmp, int* fail, sword16* u, sword16* tp, sword16* y, int k, - const byte* msg, byte* seed, byte* coins, const sword16* cacheA) + const byte* msg, byte* seed, byte* coins, const sword16* cacheA, + sword16* fillA) { int ret = 0; int i; @@ -2320,9 +2324,17 @@ int mlkem_encapsulate_seeds(const sword16* pub, MLKEM_PRF_T* prf, byte* c, } } else { + sword16* arow = a; + unsigned int aStride = 0; + + if (fillA != NULL) { + arow = fillA + (unsigned int)i * MLKEM_N; + aStride = (unsigned int)k * MLKEM_N; + } /* Generate a vector of matrix A, a polynomial at a time, * multiplying each into the u polynomial. */ - ret = mlkem_gen_matrix_i_acc(prf, u, a, 0, y, k, seed, i, 1); + ret = mlkem_gen_matrix_i_acc(prf, u, arow, aStride, y, k, seed, i, + 1); if (ret != 0) { break; } diff --git a/wolfcrypt/test/test.c b/wolfcrypt/test/test.c index 854c443163d..9fabf2da026 100644 --- a/wolfcrypt/test/test.c +++ b/wolfcrypt/test/test.c @@ -71088,6 +71088,7 @@ static wc_test_ret_t mlkem_cache_a_decode_test(int type, WC_RNG* rng) word32 ctLen; int peerInit = 0; int keyInit = 0; + int round; #ifdef WOLFSSL_SMALL_STACK peer = (MlKemKey*)XMALLOC(sizeof(MlKemKey), HEAP_HINT, @@ -71137,15 +71138,18 @@ static wc_test_ret_t mlkem_cache_a_decode_test(int type, WC_RNG* rng) ret = wc_MlKemKey_CipherTextSize(key, &ctLen); if (ret != 0) ERROR_OUT(WC_TEST_RET_ENC_EC(ret), out); - ret = wc_MlKemKey_Encapsulate(key, ct, ss, rng); - if (ret != 0) - ERROR_OUT(WC_TEST_RET_ENC_EC(ret), out); - ret = wc_MlKemKey_Decapsulate(peer, ssDec, ct, ctLen); - if (ret != 0) - ERROR_OUT(WC_TEST_RET_ENC_EC(ret), out); + /* Round 2 reads the matrix round 1 kept. */ + for (round = 0; round < 2; round++) { + ret = wc_MlKemKey_Encapsulate(key, ct, ss, rng); + if (ret != 0) + ERROR_OUT(WC_TEST_RET_ENC_EC(ret), out); + ret = wc_MlKemKey_Decapsulate(peer, ssDec, ct, ctLen); + if (ret != 0) + ERROR_OUT(WC_TEST_RET_ENC_EC(ret), out); - if (XMEMCMP(ss, ssDec, WC_ML_KEM_SS_SZ) != 0) - ERROR_OUT(WC_TEST_RET_ENC_NC, out); + if (XMEMCMP(ss, ssDec, WC_ML_KEM_SS_SZ) != 0) + ERROR_OUT(WC_TEST_RET_ENC_NC, out); + } /* Decoding a private key must drop the cached matrix too. */ ret = wc_MlKemKey_PrivateKeySize(peer, &privLen); @@ -71162,15 +71166,17 @@ static wc_test_ret_t mlkem_cache_a_decode_test(int type, WC_RNG* rng) if (ret != 0) ERROR_OUT(WC_TEST_RET_ENC_EC(ret), out); - ret = wc_MlKemKey_Encapsulate(peer, ct, ss, rng); - if (ret != 0) - ERROR_OUT(WC_TEST_RET_ENC_EC(ret), out); - ret = wc_MlKemKey_Decapsulate(key, ssDec, ct, ctLen); - if (ret != 0) - ERROR_OUT(WC_TEST_RET_ENC_EC(ret), out); + for (round = 0; round < 2; round++) { + ret = wc_MlKemKey_Encapsulate(peer, ct, ss, rng); + if (ret != 0) + ERROR_OUT(WC_TEST_RET_ENC_EC(ret), out); + ret = wc_MlKemKey_Decapsulate(key, ssDec, ct, ctLen); + if (ret != 0) + ERROR_OUT(WC_TEST_RET_ENC_EC(ret), out); - if (XMEMCMP(ss, ssDec, WC_ML_KEM_SS_SZ) != 0) - ERROR_OUT(WC_TEST_RET_ENC_NC, out); + if (XMEMCMP(ss, ssDec, WC_ML_KEM_SS_SZ) != 0) + ERROR_OUT(WC_TEST_RET_ENC_NC, out); + } ret = 0; out: diff --git a/wolfssl/wolfcrypt/wc_mlkem.h b/wolfssl/wolfcrypt/wc_mlkem.h index f2e7ce8fe05..68b57cde5c5 100644 --- a/wolfssl/wolfcrypt/wc_mlkem.h +++ b/wolfssl/wolfcrypt/wc_mlkem.h @@ -432,10 +432,10 @@ struct MlKemKey { byte z[WC_ML_KEM_SYM_SZ]; #ifdef WOLFSSL_MLKEM_CACHE_A #ifndef WOLFSSL_MLKEM_DYNAMIC_KEYS - /* A matrix from key generation. */ + /* Cached matrix A. */ sword16 a[WC_ML_KEM_MAX_K * WC_ML_KEM_MAX_K * MLKEM_N]; #else - /* A matrix from key generation (dynamically allocated). */ + /* Cached matrix A (dynamically allocated). */ sword16* a; #endif #endif @@ -567,7 +567,8 @@ int mlkem_encapsulate(const sword16* pub, sword16* bp, sword16* v, WOLFSSL_LOCAL int mlkem_encapsulate_seeds(const sword16* pub, MLKEM_PRF_T* prf, byte* c, const byte* cmp, int* fail, sword16* bp, sword16* tp, sword16* sp, int kp, - const byte* msg, byte* seed, byte* coins, const sword16* cacheA); + const byte* msg, byte* seed, byte* coins, const sword16* cacheA, + sword16* fillA); #endif WOLFSSL_LOCAL int mlkem_decapsulate(const sword16* priv, sword16* mp, sword16* bp, From ca244309b9dae8ebebc8ca424606c92f884eee8a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tobias=20Frauenschl=C3=A4ger?= Date: Tue, 6 Oct 2026 17:02:02 +0200 Subject: [PATCH 10/11] Build the ML-KEM small memory sampling helpers for decapsulation only Decapsulation re-encrypts through mlkem_encapsulate_seeds() when WOLFSSL_MLKEM_ENCAPSULATE_SMALL_MEM is set, but mlkem_gen_matrix_i_acc(), mlkem_get_noise_i() and mlkem_get_noise_eta2_c() were only compiled when encapsulation was enabled. A build with decapsulation and without encapsulation, such as --enable-mlkem=dec with the small memory option, failed to compile. Compile them when either operation is enabled, and add a CI row for that build. --- .github/configs/pq-all.json | 4 ++++ wolfcrypt/src/wc_mlkem_poly.c | 9 ++++++--- 2 files changed, 10 insertions(+), 3 deletions(-) diff --git a/.github/configs/pq-all.json b/.github/configs/pq-all.json index b2250e0bdd0..fcf45260533 100644 --- a/.github/configs/pq-all.json +++ b/.github/configs/pq-all.json @@ -298,6 +298,10 @@ "comment": "Cached matrix A with dynamic keys and the small memory options: a decoded key refills a matrix the key already allocated and otherwise streams it", "configure": ["--enable-cryptonly", "--enable-mlkem", "CPPFLAGS=-DWOLFSSL_MLKEM_CACHE_A -DWOLFSSL_MLKEM_DYNAMIC_KEYS -DWOLFSSL_MLKEM_MAKEKEY_SMALL_MEM -DWOLFSSL_MLKEM_ENCAPSULATE_SMALL_MEM"]}, +{"name": "mlkem-decapsulate-only-small-mem", "minutes": 1.1, + "comment": "Small memory decapsulation without encapsulation, which still re-encrypts through the small memory sampling helpers", + "configure": ["--enable-cryptonly", "--enable-mlkem=dec", + "CPPFLAGS=-DWOLFSSL_MLKEM_ENCAPSULATE_SMALL_MEM"]}, {"name": "frodokem-aarch64-armasm", "minutes": 0.8, "check": false, "cc": "ccache aarch64-linux-gnu-gcc", "comment": "Cross-build only (no ARM runner): catches ARM asm link/compile breaks. sve+sme+aes A-gen all live in the .S.", diff --git a/wolfcrypt/src/wc_mlkem_poly.c b/wolfcrypt/src/wc_mlkem_poly.c index 7d768c91566..b8a880b4313 100644 --- a/wolfcrypt/src/wc_mlkem_poly.c +++ b/wolfcrypt/src/wc_mlkem_poly.c @@ -125,7 +125,8 @@ #if (defined(WOLFSSL_MLKEM_MAKEKEY_SMALL_MEM) && \ !defined(WOLFSSL_MLKEM_NO_MAKE_KEY)) || \ (defined(WOLFSSL_MLKEM_ENCAPSULATE_SMALL_MEM) && \ - !defined(WOLFSSL_MLKEM_NO_ENCAPSULATE)) + (!defined(WOLFSSL_MLKEM_NO_ENCAPSULATE) || \ + !defined(WOLFSSL_MLKEM_NO_DECAPSULATE))) static int mlkem_gen_matrix_i_acc(MLKEM_PRF_T* prf, sword16* r, sword16* a, unsigned int aStride, const sword16* v, int k, byte* seed, int i, int transposed); @@ -4570,7 +4571,8 @@ int mlkem_gen_matrix(MLKEM_PRF_T* prf, sword16* a, int k, byte* seed, #if (defined(WOLFSSL_MLKEM_MAKEKEY_SMALL_MEM) && \ !defined(WOLFSSL_MLKEM_NO_MAKE_KEY)) || \ (defined(WOLFSSL_MLKEM_ENCAPSULATE_SMALL_MEM) && \ - !defined(WOLFSSL_MLKEM_NO_ENCAPSULATE)) + (!defined(WOLFSSL_MLKEM_NO_ENCAPSULATE) || \ + !defined(WOLFSSL_MLKEM_NO_DECAPSULATE))) /* Deterministically generate a vector of a matrix (or transpose) of uniform * integers mod q and multiply it into a vector of polynomials. @@ -6186,7 +6188,8 @@ int mlkem_get_noise(MLKEM_PRF_T* prf, int k, sword16* vec1, sword16* vec2, #if (defined(WOLFSSL_MLKEM_MAKEKEY_SMALL_MEM) && \ !defined(WOLFSSL_MLKEM_NO_MAKE_KEY)) || \ (defined(WOLFSSL_MLKEM_ENCAPSULATE_SMALL_MEM) && \ - !defined(WOLFSSL_MLKEM_NO_ENCAPSULATE)) + (!defined(WOLFSSL_MLKEM_NO_ENCAPSULATE) || \ + !defined(WOLFSSL_MLKEM_NO_DECAPSULATE))) /* Get the noise/error by calculating random bytes and sampling to a binomial * distribution. * From 6c647a1d1556a0cddbfe84d59cffc20839e3c341 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tobias=20Frauenschl=C3=A4ger?= Date: Mon, 2 Nov 2026 13:53:59 +0100 Subject: [PATCH 11/11] Build the ML-KEM matrix allocator only when key generation is built With WOLFSSL_MLKEM_CACHE_A and WOLFSSL_MLKEM_DYNAMIC_KEYS, mlkemkey_alloc_a() is only called from key generation, so a build that also defines WOLFSSL_MLKEM_NO_MAKE_KEY fails with -Werror on an unused static function. Guard it on key generation as well. --- wolfcrypt/src/wc_mlkem.c | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/wolfcrypt/src/wc_mlkem.c b/wolfcrypt/src/wc_mlkem.c index 862e70eb32a..b16d4ed972c 100644 --- a/wolfcrypt/src/wc_mlkem.c +++ b/wolfcrypt/src/wc_mlkem.c @@ -323,7 +323,7 @@ static int mlkemkey_alloc_pub(MlKemKey* key, unsigned int k) return 0; } -#ifdef WOLFSSL_MLKEM_CACHE_A +#if defined(WOLFSSL_MLKEM_CACHE_A) && !defined(WOLFSSL_MLKEM_NO_MAKE_KEY) /* Allocate (or reallocate) the A matrix buffer, right-sized for k. * * @param [in, out] key ML-KEM key object. @@ -347,7 +347,7 @@ static int mlkemkey_alloc_a(MlKemKey* key, unsigned int k) return ret; } -#endif /* WOLFSSL_MLKEM_CACHE_A */ +#endif /* WOLFSSL_MLKEM_CACHE_A && !WOLFSSL_MLKEM_NO_MAKE_KEY */ #endif /* WOLFSSL_MLKEM_DYNAMIC_KEYS */ /******************************************************************************/