Reduce ML-KEM small memory heap use and allow caching matrix A with it

.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"]},
@@ -233,9 +237,33 @@
   "--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",
+  "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",
+  "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"]},
+{"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",
+  "CPPFLAGS=-DWOLFSSL_MLKEM_CACHE_A"]},
 {"name": "mldsa-verify-only-dynamic-keys", "minutes": 1.7,
  "configure": ["--disable-intelasm",
   "--enable-dilithium=44,65,87,verify-only",
@@ -254,6 +282,21 @@
   "--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": "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": "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": "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.",

.github/workflows/wolfCrypt-Wconversion.yml

@@ -114,6 +114,15 @@ 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": ["--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 -Wdeclaration-after-statement"],
+           "check": false},
           {"name": "precalc-a-no-int128", "minutes": 1,
            "configure": ["--enable-cryptonly", "--enable-all-crypto",
             "--disable-examples", "--disable-benchmark",

configure.ac

@@ -8788,6 +8788,24 @@ do
   esac
 done
 
+# Default the parameter sets and operations when only a mode is named.
+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 "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
+    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"],

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);

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,53 +349,70 @@ 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.
  * ------------------------------------------------------------------------- */
 #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 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];
     long           n;
+    unsigned int   j;
 
     if (mlkem_prf_new(&prf, NULL, INVALID_DEVID) != 0) {
         WB_NOTE("mlkem_prf_new failed");
         wb_fail = 1;
         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, 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(&prf, a, 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, 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(&prf, a, WB_K, seed, 1, 1);
+        (void)mlkem_gen_matrix_i_acc(&prf, r, a, WB_K * MLKEM_N, 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 +429,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, 0, v, 1, seed, 0, 0) != 0) {
+            WB_NOTE("mlkem_gen_matrix_i_acc seed search failed");
             wb_fail = 1;
             break;
         }
@@ -429,32 +447,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, 0, 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