From fffafe661a3021d1559bc38f8981723ebc1ab5fd Mon Sep 17 00:00:00 2001
From: Daniel Pouzzner <douzzer@wolfssl.com>
Date: Mon, 6 Jan 2025 14:52:42 -0600
Subject: [PATCH] wolfcrypt/src/dh.c: in _DhSetKey(), add short-circuit
 comparisons to RFC 7919 known-good moduli, preempting overhead from
 mp_prime_is_prime().

wolfcrypt/test/test.c: in dh_ffdhe_test(), when defined(HAVE_PUBLIC_FFDHE), use wc_DhSetKey_ex() rather than wc_DhSetKey() to exercise the primality check in _DhSetKey().
---
 wolfcrypt/src/dh.c    | 42 +++++++++++++++++++++++++++++++++++++++---
 wolfcrypt/test/test.c | 10 +++++++---
 2 files changed, 46 insertions(+), 6 deletions(-)

diff --git a/wolfcrypt/src/dh.c b/wolfcrypt/src/dh.c
index 5258e820e4..9fd37bc9e6 100644
--- a/wolfcrypt/src/dh.c
+++ b/wolfcrypt/src/dh.c
@@ -2544,10 +2544,46 @@ static int _DhSetKey(DhKey* key, const byte* p, word32 pSz, const byte* g,
 
     if (ret == 0 && !trusted) {
         int isPrime = 0;
-        if (rng != NULL)
-            ret = mp_prime_is_prime_ex(keyP, 8, &isPrime, rng);
+
+        /* Short-circuit the primality check for p if it is one of the named
+         * public moduli (known primes) from RFC 7919.
+         */
+        #ifdef HAVE_FFDHE_2048
+        if ((pSz == sizeof(dh_ffdhe2048_p)) && (XMEMCMP(p, dh_ffdhe2048_p, sizeof(dh_ffdhe2048_p)) == 0)) {
+            isPrime = 1;
+        }
+        else
+        #endif
+        #ifdef HAVE_FFDHE_3072
+        if ((pSz == sizeof(dh_ffdhe3072_p)) && (XMEMCMP(p, dh_ffdhe3072_p, sizeof(dh_ffdhe3072_p)) == 0)) {
+            isPrime = 1;
+        }
+        else
+        #endif
+        #ifdef HAVE_FFDHE_4096
+        if ((pSz == sizeof(dh_ffdhe4096_p)) && (XMEMCMP(p, dh_ffdhe4096_p, sizeof(dh_ffdhe4096_p)) == 0)) {
+            isPrime = 1;
+        }
         else
-            ret = mp_prime_is_prime(keyP, 8, &isPrime);
+        #endif
+        #ifdef HAVE_FFDHE_6144
+        if ((pSz == sizeof(dh_ffdhe6144_p)) && (XMEMCMP(p, dh_ffdhe6144_p, sizeof(dh_ffdhe6144_p)) == 0)) {
+            isPrime = 1;
+        }
+        else
+        #endif
+        #ifdef HAVE_FFDHE_8192
+        if ((pSz == sizeof(dh_ffdhe8192_p)) && (XMEMCMP(p, dh_ffdhe8192_p, sizeof(dh_ffdhe8192_p)) == 0)) {
+            isPrime = 1;
+        }
+        else
+        #endif
+        {
+            if (rng != NULL)
+                ret = mp_prime_is_prime_ex(keyP, 8, &isPrime, rng);
+            else
+                ret = mp_prime_is_prime(keyP, 8, &isPrime);
+        }
 
         if (ret == 0 && isPrime == 0)
             ret = DH_CHECK_PUB_E;
diff --git a/wolfcrypt/test/test.c b/wolfcrypt/test/test.c
index 9a1d501d85..093e39ca52 100644
--- a/wolfcrypt/test/test.c
+++ b/wolfcrypt/test/test.c
@@ -22714,7 +22714,11 @@ static wc_test_ret_t dh_ffdhe_test(WC_RNG *rng, int name)
         ERROR_OUT(WC_TEST_RET_ENC_EC(ret), done);
 
 #ifdef HAVE_PUBLIC_FFDHE
-    ret = wc_DhSetKey(key, params->p, params->p_len, params->g, params->g_len);
+    /* use wc_DhSetKey_ex(), not wc_DhSetKey(), so that trusted=0 is passed to
+     * _DhSetKey(), exercising the primality check on the modulus:
+     */
+    ret = wc_DhSetKey_ex(key, params->p, params->p_len, params->g,
+                         params->g_len, NULL /* q */, 0 /* qSz */);
 #else
     ret = wc_DhSetNamedKey(key, name);
 #endif
@@ -22722,8 +22726,8 @@ static wc_test_ret_t dh_ffdhe_test(WC_RNG *rng, int name)
         ERROR_OUT(WC_TEST_RET_ENC_EC(ret), done);
 
 #ifdef HAVE_PUBLIC_FFDHE
-    ret = wc_DhSetKey(key2, params->p, params->p_len, params->g,
-                                                                 params->g_len);
+    ret = wc_DhSetKey_ex(key2, params->p, params->p_len, params->g,
+                         params->g_len, NULL /* q */, 0 /* qSz */);
 #else
     ret = wc_DhSetNamedKey(key2, name);
 #endif