From b359dc581ac7ea2dbcd988af646642a3af37d0d7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tobias=20Frauenschl=C3=A4ger?= <tobias@wolfssl.com>
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/wc_mlkem.c                      | 124 ++++++---
 wolfcrypt/src/wc_mlkem_poly.c                 | 255 +++++++++++++-----
 wolfssl/wolfcrypt/wc_mlkem.h                  |   8 +-
 4 files changed, 314 insertions(+), 130 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/wc_mlkem.c b/wolfcrypt/src/wc_mlkem.c
index 1192caffe63..1c369f71def 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 };
@@ -905,7 +913,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
@@ -913,6 +922,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;
 
@@ -946,7 +957,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));
@@ -969,9 +987,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;
@@ -1132,20 +1149,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
 
@@ -1364,14 +1378,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;
@@ -1389,11 +1407,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) {
@@ -1449,7 +1470,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) {
@@ -1515,8 +1540,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);
@@ -1527,19 +1553,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;
@@ -1573,7 +1608,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);
             });
@@ -1581,6 +1616,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
@@ -1915,7 +1953,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)
@@ -1924,7 +1963,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
     }
@@ -2227,6 +2267,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))
@@ -2234,6 +2275,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. */
@@ -2321,6 +2365,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) {
@@ -2330,6 +2375,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
@@ -2360,12 +2406,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)
@@ -2401,6 +2455,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
@@ -2414,6 +2469,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 3ac34bf6071..91e8f58ce53 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;
 }
@@ -4401,25 +4491,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.
@@ -4428,8 +4528,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;
@@ -4486,8 +4586,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)) {
@@ -4495,10 +4594,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);
+            }
         }
     }
 
@@ -6419,10 +6528,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;
@@ -6431,14 +6541,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
@@ -6497,10 +6607,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
@@ -6508,7 +6619,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
@@ -6517,13 +6628,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 34781c5a769..2b5228cad39 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);

From 75c0b32f1b5557e160989955c17f2226ff930632 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tobias=20Frauenschl=C3=A4ger?= <tobias@wolfssl.com>
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 1c369f71def..70213f0d1d1 100644
--- a/wolfcrypt/src/wc_mlkem.c
+++ b/wolfcrypt/src/wc_mlkem.c
@@ -1638,6 +1638,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 6e8f81d6798c0ddde349874e8674b19790dc887f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tobias=20Frauenschl=C3=A4ger?= <tobias@wolfssl.com>
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.
---
 .../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 +-
 4 files changed, 68 insertions(+), 35 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 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 70213f0d1d1..1e448201a66 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) && \
@@ -917,9 +912,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. */
@@ -1111,7 +1105,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
@@ -1414,6 +1408,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. */
@@ -1569,7 +1571,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 91e8f58ce53..fc74a1834ca 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);
@@ -4518,7 +4541,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.
@@ -4529,7 +4556,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;
@@ -4582,11 +4610,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)) {
@@ -4594,7 +4623,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);
                 }
             }
@@ -4603,10 +4632,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 2b5228cad39..0d9f6d404cb 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 26d65096785df770ee427fb041d9459eb34c24d5 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tobias=20Frauenschl=C3=A4ger?= <tobias@wolfssl.com>
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 32a3c50f9fc..ed0435d9ce9 100644
--- a/.github/configs/pq-all.json
+++ b/.github/configs/pq-all.json
@@ -233,6 +233,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 fc74a1834ca..0b8250b502e 100644
--- a/wolfcrypt/src/wc_mlkem_poly.c
+++ b/wolfcrypt/src/wc_mlkem_poly.c
@@ -4628,14 +4628,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 ff02ebaa7a2..f1706a8e247 100644
--- a/wolfcrypt/test/test.c
+++ b/wolfcrypt/test/test.c
@@ -70238,6 +70238,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;
@@ -70560,6 +70631,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 5cbe4206ea700405114f6cd9dd89f99e857da390 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tobias=20Frauenschl=C3=A4ger?= <tobias@wolfssl.com>
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 ed0435d9ce9..e2952ffe55d 100644
--- a/.github/configs/pq-all.json
+++ b/.github/configs/pq-all.json
@@ -233,6 +233,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",
@@ -240,6 +244,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",
@@ -258,6 +266,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 d256bd1dae3..0ecc8a82a2b 100644
--- a/configure.ac
+++ b/configure.ac
@@ -8788,6 +8788,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 1e448201a66..761ca19b75c 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 f1706a8e247..60c335e3a78 100644
--- a/wolfcrypt/test/test.c
+++ b/wolfcrypt/test/test.c
@@ -70246,53 +70246,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);
 
@@ -70302,9 +70353,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
@@ -70634,22 +70692,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 54cb5bce9be036baec750de5fd8ff03a176710b4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tobias=20Frauenschl=C3=A4ger?= <tobias@wolfssl.com>
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 e2952ffe55d..8d5716058c9 100644
--- a/.github/configs/pq-all.json
+++ b/.github/configs/pq-all.json
@@ -233,6 +233,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",
@@ -244,6 +252,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",
@@ -269,6 +281,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 0ecc8a82a2b..19cb88101a5 100644
--- a/configure.ac
+++ b/configure.ac
@@ -8789,17 +8789,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 761ca19b75c..8692bc0bb73 100644
--- a/wolfcrypt/src/wc_mlkem.c
+++ b/wolfcrypt/src/wc_mlkem.c
@@ -1377,7 +1377,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.
@@ -1643,7 +1643,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
 
@@ -2274,16 +2280,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 0b8250b502e..aaed3a092af 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 7112c21965771f8f50dd2b6a61b151dbcaa03c6e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tobias=20Frauenschl=C3=A4ger?= <tobias@wolfssl.com>
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 aaed3a092af..8decaa22eb4 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 89728e8de81c71eb6dd1548eaa1f9962728227d5 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tobias=20Frauenschl=C3=A4ger?= <tobias@wolfssl.com>
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         | 149 ++++++++++++++++++++++++++++++++++
 5 files changed, 191 insertions(+), 9 deletions(-)

diff --git a/.github/configs/pq-all.json b/.github/configs/pq-all.json
index 8d5716058c9..37ca46d564d 100644
--- a/.github/configs/pq-all.json
+++ b/.github/configs/pq-all.json
@@ -158,6 +158,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 ba335cb6393..77cb5537521 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 8692bc0bb73..5b9f3bb90dc 100644
--- a/wolfcrypt/src/wc_mlkem.c
+++ b/wolfcrypt/src/wc_mlkem.c
@@ -1375,11 +1375,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,
@@ -1465,6 +1466,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) {
@@ -1618,7 +1624,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 8decaa22eb4..7cafea4f343 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 60c335e3a78..ac90cb245ac 100644
--- a/wolfcrypt/test/test.c
+++ b/wolfcrypt/test/test.c
@@ -70367,6 +70367,116 @@ 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);
+    ret = wc_MlKemKey_EncodePrivateKey(key, priv, privLen);
+    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;
+        ret = wc_MlKemKey_Decapsulate(key, ssDec, ctBad, ctLen);
+        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;
@@ -70728,6 +70838,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 19e609657e479b431523db094186f05aff6bb2ad Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tobias=20Frauenschl=C3=A4ger?= <tobias@wolfssl.com>
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 37ca46d564d..d0837c16333 100644
--- a/.github/configs/pq-all.json
+++ b/.github/configs/pq-all.json
@@ -289,6 +289,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 5b9f3bb90dc..854338708a6 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
@@ -1403,8 +1404,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
@@ -1414,8 +1417,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
 
@@ -1515,6 +1521,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.
@@ -1567,8 +1583,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
         }
 
@@ -1576,8 +1594,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 */
@@ -1645,16 +1675,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 7cafea4f343..617cc38dde0 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 ac90cb245ac..6204eab80c3 100644
--- a/wolfcrypt/test/test.c
+++ b/wolfcrypt/test/test.c
@@ -70266,6 +70266,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,
@@ -70315,15 +70316,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);
@@ -70340,15 +70344,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 0d9f6d404cb..3597d42ff6a 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 a212d38786a9762aa4b28fef341eab4d4f61f115 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tobias=20Frauenschl=C3=A4ger?= <tobias@wolfssl.com>
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 d0837c16333..010b3cdabe6 100644
--- a/.github/configs/pq-all.json
+++ b/.github/configs/pq-all.json
@@ -293,6 +293,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 617cc38dde0..4e7096d3366 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);
@@ -4533,7 +4534,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.
@@ -6138,7 +6140,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 eb10c91d119d75200a20830f35040bee8935625c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tobias=20Frauenschl=C3=A4ger?= <tobias@wolfssl.com>
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 854338708a6..b50f1d18943 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 */
 
 /******************************************************************************/