@@ -972,6 +972,10 @@ int wc_InitDhKey_ex(DhKey* key, void* heap, int devId)
972972 key -> handle = NULL ;
973973#endif
974974
975+ #ifdef WC_DH_NONBLOCK
976+ key -> nb = NULL ;
977+ #endif
978+
975979 return ret ;
976980}
977981
@@ -980,6 +984,23 @@ int wc_InitDhKey(DhKey* key)
980984 return wc_InitDhKey_ex (key , NULL , INVALID_DEVID );
981985}
982986
987+ #ifdef WC_DH_NONBLOCK
988+ int wc_DhSetNonBlock (DhKey * key , DhNb * nb )
989+ {
990+ if (key == NULL )
991+ return BAD_FUNC_ARG ;
992+
993+ if (nb != NULL ) {
994+ XMEMSET (nb , 0 , sizeof (DhNb ));
995+ }
996+
997+ /* Pass NULL to disable non-blocking mode. */
998+ key -> nb = nb ;
999+
1000+ return 0 ;
1001+ }
1002+ #endif
1003+
9831004
9841005int wc_FreeDhKey (DhKey * key )
9851006{
@@ -2030,18 +2051,80 @@ static int wc_DhAgree_Sync(DhKey* key, byte* agree, word32* agreeSz,
20302051 if (mp_iseven (& key -> p ) == MP_YES ) {
20312052 return MP_VAL ;
20322053 }
2054+
2055+ /* Non-blocking re-entry: the same wc_DhAgree call repeats until the
2056+ * SP state machine completes, so cache the per-op key validation
2057+ * results instead of re-running them each yield. The cache is
2058+ * scoped to non-blocking, non-const-time callers only. */
2059+ #ifdef WC_DH_NONBLOCK
2060+ if (key -> nb == NULL || ct || !key -> nb -> pubKeyValidated )
2061+ #endif
2062+ {
20332063#ifdef WOLFSSL_VALIDATE_FFC_IMPORT
2034- if (wc_DhCheckPrivKey (key , priv , privSz ) != 0 ) {
2035- WOLFSSL_MSG ("wc_DhAgree wc_DhCheckPrivKey failed" );
2036- return DH_CHECK_PRIV_E ;
2037- }
2064+ if (wc_DhCheckPrivKey (key , priv , privSz ) != 0 ) {
2065+ WOLFSSL_MSG ("wc_DhAgree wc_DhCheckPrivKey failed" );
2066+ return DH_CHECK_PRIV_E ;
2067+ }
2068+ #endif
2069+ /* Always validate peer public key (2 <= y <= p-2) per SP 800-56A */
2070+ if (wc_DhCheckPubKey (key , otherPub , pubSz ) != 0 ) {
2071+ WOLFSSL_MSG ("wc_DhAgree wc_DhCheckPubKey failed" );
2072+ return DH_CHECK_PUB_E ;
2073+ }
2074+ #ifdef WC_DH_NONBLOCK
2075+ if (key -> nb != NULL && !ct ) {
2076+ key -> nb -> pubKeyValidated = 1 ;
2077+ }
20382078#endif
2079+ }
20392080
2040- /* Always validate peer public key (2 <= y <= p-2) per SP 800-56A */
2041- if (wc_DhCheckPubKey (key , otherPub , pubSz ) != 0 ) {
2042- WOLFSSL_MSG ("wc_DhAgree wc_DhCheckPubKey failed" );
2043- return DH_CHECK_PUB_E ;
2081+ #if defined(WC_DH_NONBLOCK ) && defined(WOLFSSL_HAVE_SP_DH ) && \
2082+ defined(WOLFSSL_SP_NONBLOCK ) && defined(WOLFSSL_SP_SMALL ) && \
2083+ !defined(WOLFSSL_SP_FAST_MODEXP )
2084+ /* Non-blocking dispatch bypasses the mp_int dance entirely - the SP
2085+ * wrapper takes byte buffers and persists across yields. The constant-
2086+ * time fold-back (ct branch) is intentionally not applied here; nb
2087+ * callers should use the standard wc_DhAgree(). */
2088+ if (key -> nb != NULL && !ct ) {
2089+ int nb_ret = MP_OKAY ;
2090+ int dispatched = 0 ;
2091+ #ifndef WOLFSSL_SP_NO_2048
2092+ if (mp_count_bits (& key -> p ) == 2048 ) {
2093+ nb_ret = sp_DhExp_2048_nb (& key -> nb -> sp_ctx , otherPub , pubSz ,
2094+ priv , privSz , & key -> p , agree , agreeSz );
2095+ dispatched = 1 ;
2096+ }
2097+ #endif
2098+ #ifndef WOLFSSL_SP_NO_3072
2099+ if (!dispatched && mp_count_bits (& key -> p ) == 3072 ) {
2100+ nb_ret = sp_DhExp_3072_nb (& key -> nb -> sp_ctx , otherPub , pubSz ,
2101+ priv , privSz , & key -> p , agree , agreeSz );
2102+ dispatched = 1 ;
2103+ }
2104+ #endif
2105+ #ifdef WOLFSSL_SP_4096
2106+ if (!dispatched && mp_count_bits (& key -> p ) == 4096 ) {
2107+ nb_ret = sp_DhExp_4096_nb (& key -> nb -> sp_ctx , otherPub , pubSz ,
2108+ priv , privSz , & key -> p , agree , agreeSz );
2109+ dispatched = 1 ;
2110+ }
2111+ #endif
2112+ if (dispatched ) {
2113+ /* Op finished (or hit a hard error) - clear the cached
2114+ * validation so the next op on this DhNb re-runs the
2115+ * SP 800-56A peer-key check. MP_WOULDBLOCK keeps it. */
2116+ if (nb_ret != MP_WOULDBLOCK ) {
2117+ key -> nb -> pubKeyValidated = 0 ;
2118+ }
2119+ return nb_ret ;
2120+ }
2121+ /* size not nb-supported - the blocking path below completes in
2122+ * one call, so the cached validation is single-use. Clear it
2123+ * here so the next agree on this DhNb re-validates. */
2124+ key -> nb -> pubKeyValidated = 0 ;
2125+ /* fall through to blocking path */
20442126 }
2127+ #endif
20452128
20462129#if defined(WOLFSSL_SMALL_STACK ) && !defined(WOLFSSL_NO_MALLOC )
20472130 y = (mp_int * )XMALLOC (sizeof (mp_int ), key -> heap , DYNAMIC_TYPE_DH );
@@ -2304,13 +2387,21 @@ int wc_DhAgree(DhKey* key, byte* agree, word32* agreeSz, const byte* priv,
23042387 ret = KcapiDh_SharedSecret (key , otherPub , pubSz , agree , agreeSz );
23052388#else
23062389#if defined(WOLFSSL_ASYNC_CRYPT ) && defined(WC_ASYNC_ENABLE_DH )
2390+ /* Async marker takes precedence: when wolfAsync_DoSw (wolfcrypt/src/
2391+ * async.c) re-enters the compute path, wc_DhAgree_Async dispatches
2392+ * to the SP nonblock wrapper if key->nb is attached, and per-yield
2393+ * FP_WOULDBLOCK (alias of MP_WOULDBLOCK) is translated to
2394+ * WC_PENDING_E by wolfAsync_DoSw so the TLS event loop drives it. */
23072395 if (key -> asyncDev .marker == WOLFSSL_ASYNC_MARKER_DH ) {
23082396 ret = wc_DhAgree_Async (key , agree , agreeSz , priv , privSz , otherPub ,
23092397 pubSz );
23102398 }
23112399 else
23122400#endif
23132401 {
2402+ /* wc_DhAgree_Sync handles key->nb internally; no separate dispatch
2403+ * needed here. wc_DhAgree_ct (constant-time fold-back) bypasses
2404+ * this function entirely so passing ct=0 is correct. */
23142405 ret = wc_DhAgree_Sync (key , agree , agreeSz , priv , privSz , otherPub ,
23152406 pubSz , 0 );
23162407 }
0 commit comments