From 48ee35b8317be64bcf808ba9d4503e1eacf06cbc Mon Sep 17 00:00:00 2001
From: Emma Stensland <emma.stensland@wolfssl.com>
Date: Wed, 9 Sep 2026 10:43:40 -0600
Subject: [PATCH] Fix sword32 overflow before inverse NTT in ML-DSA small-mem

---
 wolfcrypt/src/wc_mldsa.c | 25 +++++++++++++++++++++++++
 1 file changed, 25 insertions(+)

diff --git a/wolfcrypt/src/wc_mldsa.c b/wolfcrypt/src/wc_mldsa.c
index b673dcb477..830c33ad90 100644
--- a/wolfcrypt/src/wc_mldsa.c
+++ b/wolfcrypt/src/wc_mldsa.c
@@ -6728,6 +6728,9 @@ static sword32 mldsa_mont_red(sword64 a)
 
 #if !defined(WOLFSSL_MLDSA_SMALL) || \
     (!defined(WOLFSSL_MLDSA_NO_SIGN) || \
+     (!defined(WOLFSSL_MLDSA_SMALL_MEM_POLY64) && \
+      !defined(WOLFSSL_MLDSA_NO_VERIFY) && \
+      defined(WOLFSSL_MLDSA_VERIFY_SMALL_MEM)) || \
      (defined(WOLFSSL_MLDSA_SMALL) && \
       (!defined(WOLFSSL_MLDSA_NO_MAKE_KEY) || \
        (!defined(WOLFSSL_MLDSA_NO_VERIFY) && \
@@ -8331,6 +8334,10 @@ static int mldsa_invntt(sword32* r)
 #endif
 
 /* Inverse Number-Theoretic Transform.
+ *
+ * Cannot overflow when |r[i]| < Q: the sum lane is left unreduced across
+ * all 8 levels, so it reaches 256 times the input bound, and
+ * 256 * Q < 2^31.
  *
  * @param [in, out] r  Polynomial to transform.
  */
@@ -8831,6 +8838,11 @@ static int mldsa_vec_mul(sword32* r, sword32* a, sword32* b, byte l)
 #endif
 
 #if !defined(WOLFSSL_MLDSA_NO_SIGN) || \
+    (!defined(WOLFSSL_MLDSA_SMALL_MEM_POLY64) && \
+     ((!defined(WOLFSSL_MLDSA_NO_MAKE_KEY) && \
+       defined(WOLFSSL_MLDSA_MAKE_KEY_SMALL_MEM)) || \
+      (!defined(WOLFSSL_MLDSA_NO_VERIFY) && \
+       defined(WOLFSSL_MLDSA_VERIFY_SMALL_MEM)))) || \
     (defined(WOLFSSL_MLDSA_SMALL) && \
      (!defined(WOLFSSL_MLDSA_NO_MAKE_KEY) || \
       (!defined(WOLFSSL_MLDSA_NO_VERIFY) && \
@@ -9673,6 +9685,12 @@ static int mldsa_make_key_from_seed(wc_MlDsaKey* key, const byte* seed)
             for (e = 0; e < MLDSA_N; e++) {
                 tt[e] = mldsa_mont_red(t64[e]);
             }
+        #else
+            /* Sum of l Montgomery products reaches l*Q; bring it back
+             * inside |x| < Q, where mldsa_invntt_full() cannot overflow. */
+            if (ret == 0) {
+                ret = mldsa_poly_red(tt);
+            }
         #endif
             if (ret == 0) {
                 ret = mldsa_invntt_full(tt);
@@ -12441,6 +12459,13 @@ static int mldsa_verify_with_mu(wc_MlDsaKey* key, const byte* mu,
             for (e = 0; e < MLDSA_N; e++) {
                 w[e] = mldsa_mont_red(t64[e]);
             }
+        #else
+            /* Sum of l + 1 Montgomery products reaches (l + 1)*Q; bring it
+             * back inside |x| < Q, where mldsa_invntt_full() cannot
+             * overflow. */
+            if (ret == 0) {
+                ret = mldsa_poly_red(w);
+            }
         #endif
 
             /* Step 10: w = NTT-1(A o NTT(z) - NTT(c) o NTT(t1)) */