@@ -6813,6 +6813,7 @@ struct wolfSSH_sigKeyBlockFull {
68136813int SendKexDhReply (WOLFSSH * ssh )
68146814{
68156815 int ret = WS_SUCCESS ;
6816+ void * heap = NULL ;
68166817 byte * f_ptr = NULL , * sig_ptr = NULL ;
68176818#ifndef WOLFSSH_NO_ECDH
68186819 byte * r_ptr = NULL , * s_ptr = NULL ;
@@ -6838,12 +6839,7 @@ int SendKexDhReply(WOLFSSH* ssh)
68386839 word32 generatorSz = 0 ;
68396840#endif
68406841 struct wolfSSH_sigKeyBlockFull * sigKeyBlock_ptr = NULL ;
6841- #ifdef WOLFSSH_SMALL_STACK
6842- f_ptr = (byte * )WMALLOC (KEX_F_SIZE , ssh -> ctx -> heap , DYNTYPE_BUFFER );
6843- sig_ptr = (byte * )WMALLOC (KEX_SIG_SIZE , ssh -> ctx -> heap , DYNTYPE_BUFFER );
6844- if (f_ptr == NULL || sig_ptr == NULL )
6845- ret = WS_MEMORY_E ;
6846- #else
6842+ #ifndef WOLFSSH_SMALL_STACK
68476843 byte f_s [KEX_F_SIZE ];
68486844 byte sig_s [KEX_SIG_SIZE ];
68496845
@@ -6853,14 +6849,24 @@ int SendKexDhReply(WOLFSSH* ssh)
68536849 WLOG (WS_LOG_DEBUG , "Entering SendKexDhReply()" );
68546850
68556851 if (ret == WS_SUCCESS ) {
6856- if (ssh == NULL || ssh -> handshake == NULL ) {
6852+ if (ssh == NULL || ssh -> ctx == NULL || ssh -> handshake == NULL ) {
68576853 ret = WS_BAD_ARGUMENT ;
68586854 }
68596855 }
68606856
6857+ if (ret == WS_SUCCESS ) {
6858+ heap = ssh -> ctx -> heap ;
6859+ }
6860+
6861+ #ifdef WOLFSSH_SMALL_STACK
6862+ f_ptr = (byte * )WMALLOC (KEX_F_SIZE , heap , DYNTYPE_BUFFER );
6863+ sig_ptr = (byte * )WMALLOC (KEX_SIG_SIZE , heap , DYNTYPE_BUFFER );
6864+ if (f_ptr == NULL || sig_ptr == NULL )
6865+ ret = WS_MEMORY_E ;
6866+ #endif
6867+
68616868 sigKeyBlock_ptr = (struct wolfSSH_sigKeyBlockFull * )WMALLOC (
6862- sizeof (struct wolfSSH_sigKeyBlockFull ),
6863- ssh -> ctx -> heap , DYNTYPE_PRIVKEY );
6869+ sizeof (struct wolfSSH_sigKeyBlockFull ), heap , DYNTYPE_PRIVKEY );
68646870 if (sigKeyBlock_ptr == NULL )
68656871 ret = WS_MEMORY_E ;
68666872
@@ -6934,7 +6940,7 @@ int SendKexDhReply(WOLFSSH* ssh)
69346940 /* Decode the user-configured RSA private key. */
69356941 sigKeyBlock_ptr -> sk .rsa .eSz = sizeof (sigKeyBlock_ptr -> sk .rsa .e );
69366942 sigKeyBlock_ptr -> sk .rsa .nSz = sizeof (sigKeyBlock_ptr -> sk .rsa .n );
6937- ret = wc_InitRsaKey (& sigKeyBlock_ptr -> sk .rsa .key , ssh -> ctx -> heap );
6943+ ret = wc_InitRsaKey (& sigKeyBlock_ptr -> sk .rsa .key , heap );
69386944 if (ret == 0 )
69396945 ret = wc_RsaPrivateKeyDecode (ssh -> ctx -> privateKey , & scratch ,
69406946 & sigKeyBlock_ptr -> sk .rsa .key ,
@@ -7040,8 +7046,8 @@ int SendKexDhReply(WOLFSSH* ssh)
70407046
70417047 /* Decode the user-configured ECDSA private key. */
70427048 sigKeyBlock_ptr -> sk .ecc .qSz = sizeof (sigKeyBlock_ptr -> sk .ecc .q );
7043- ret = wc_ecc_init_ex (& sigKeyBlock_ptr -> sk .ecc .key , ssh -> ctx -> heap ,
7044- INVALID_DEVID );
7049+ ret = wc_ecc_init_ex (& sigKeyBlock_ptr -> sk .ecc .key , heap ,
7050+ INVALID_DEVID );
70457051 scratch = 0 ;
70467052 if (ret == 0 )
70477053 ret = wc_EccPrivateKeyDecode (ssh -> ctx -> privateKey , & scratch ,
@@ -7209,7 +7215,7 @@ int SendKexDhReply(WOLFSSH* ssh)
72097215 DhKey privKey ;
72107216 word32 ySz = MAX_KEX_KEY_SZ ;
72117217#ifdef WOLFSSH_SMALL_STACK
7212- y_ptr = (byte * )WMALLOC (ySz , ssh -> ctx -> heap , DYNTYPE_PRIVKEY );
7218+ y_ptr = (byte * )WMALLOC (ySz , heap , DYNTYPE_PRIVKEY );
72137219 if (y_ptr == NULL )
72147220 ret = WS_MEMORY_E ;
72157221#else
@@ -7243,11 +7249,9 @@ int SendKexDhReply(WOLFSSH* ssh)
72437249 ret = WS_INVALID_PRIME_CURVE ;
72447250
72457251 if (ret == 0 )
7246- ret = wc_ecc_init_ex (& pubKey , ssh -> ctx -> heap ,
7247- INVALID_DEVID );
7252+ ret = wc_ecc_init_ex (& pubKey , heap , INVALID_DEVID );
72487253 if (ret == 0 )
7249- ret = wc_ecc_init_ex (& privKey , ssh -> ctx -> heap ,
7250- INVALID_DEVID );
7254+ ret = wc_ecc_init_ex (& privKey , heap , INVALID_DEVID );
72517255#ifdef HAVE_WC_ECC_SET_RNG
72527256 if (ret == 0 )
72537257 ret = wc_ecc_set_rng (& privKey , ssh -> rng );
@@ -7391,8 +7395,8 @@ int SendKexDhReply(WOLFSSH* ssh)
73917395 byte rPad ;
73927396 byte sPad ;
73937397#ifdef WOLFSSH_SMALL_STACK
7394- r_ptr = (byte * )WMALLOC (rSz , ssh -> ctx -> heap , DYNTYPE_BUFFER );
7395- s_ptr = (byte * )WMALLOC (sSz , ssh -> ctx -> heap , DYNTYPE_BUFFER );
7398+ r_ptr = (byte * )WMALLOC (rSz , heap , DYNTYPE_BUFFER );
7399+ s_ptr = (byte * )WMALLOC (sSz , heap , DYNTYPE_BUFFER );
73967400 if (r_ptr == NULL || r_ptr == NULL )
73977401 ret = WS_MEMORY_E ;
73987402#else
@@ -7532,20 +7536,20 @@ int SendKexDhReply(WOLFSSH* ssh)
75327536
75337537 WLOG (WS_LOG_DEBUG , "Leaving SendKexDhReply(), ret = %d" , ret );
75347538 if (sigKeyBlock_ptr )
7535- WFREE (sigKeyBlock_ptr , ssh -> ctx -> heap , DYNTYPE_PRIVKEY );
7539+ WFREE (sigKeyBlock_ptr , heap , DYNTYPE_PRIVKEY );
75367540#ifdef WOLFSSH_SMALL_STACK
75377541 if (f_ptr )
7538- WFREE (f_ptr , ssh -> ctx -> heap , DYNTYPE_BUFFER );
7542+ WFREE (f_ptr , heap , DYNTYPE_BUFFER );
75397543 if (sig_ptr )
7540- WFREE (sig_ptr , ssh -> ctx -> heap , DYNTYPE_BUFFER );
7544+ WFREE (sig_ptr , heap , DYNTYPE_BUFFER );
75417545#ifndef WOLFSSH_NO_DH
75427546 if (y_ptr )
7543- WFREE (r_ptr , ssh -> ctx -> heap , DYNTYPE_PRIVKEY );
7547+ WFREE (y_ptr , heap , DYNTYPE_PRIVKEY );
75447548#endif
75457549 if (r_ptr )
7546- WFREE (r_ptr , ssh -> ctx -> heap , DYNTYPE_BUFFER );
7550+ WFREE (r_ptr , heap , DYNTYPE_BUFFER );
75477551 if (s_ptr )
7548- WFREE (s_ptr , ssh -> ctx -> heap , DYNTYPE_BUFFER );
7552+ WFREE (s_ptr , heap , DYNTYPE_BUFFER );
75497553#endif
75507554 return ret ;
75517555}
0 commit comments