crypto: move key setup in newkeys handler

3 files changed, 16 insertions, 11 deletions

diff --git a/src/packet_cb.c b/src/packet_cb.c
index 4fe6402a..f2162304 100644
--- a/src/packet_cb.c
+++ b/src/packet_cb.c
@@ -237,6 +237,14 @@ SSH_PACKET_CALLBACK(ssh_packet_newkeys){
     }
     memcpy(session->next_crypto->session_id, session->current_crypto->session_id,
             session->current_crypto->digest_len);
+    if (session->current_crypto->in_cipher->set_decrypt_key(session->current_crypto->in_cipher, session->current_crypto->decryptkey,
+        session->current_crypto->decryptIV) < 0) {
+      goto error;
+    }
+    if (session->current_crypto->out_cipher->set_encrypt_key(session->current_crypto->out_cipher, session->current_crypto->encryptkey,
+        session->current_crypto->encryptIV) < 0) {
+      goto error;
+    }
   }
   session->dh_handshake_state = DH_STATE_FINISHED;
   session->ssh_connection_callback(session);
diff --git a/src/packet_crypt.c b/src/packet_crypt.c
index 914727e0..83a9bb3d 100644
--- a/src/packet_crypt.c
+++ b/src/packet_crypt.c
@@ -72,11 +72,6 @@ int packet_decrypt(ssh_session session, void *data,uint32_t len) {
     return -1;
   }
 
-  if (crypto->set_decrypt_key(crypto, session->current_crypto->decryptkey,
-        session->current_crypto->decryptIV) < 0) {
-    SAFE_FREE(out);
-    return -1;
-  }
   crypto->decrypt(crypto,data,out,len);
 
   memcpy(data,out,len);
@@ -111,12 +106,6 @@ unsigned char *packet_encrypt(ssh_session session, void *data, uint32_t len) {
   seq = ntohl(session->send_seq);
   crypto = session->current_crypto->out_cipher;
 
-  if (crypto->set_encrypt_key(crypto, session->current_crypto->encryptkey,
-      session->current_crypto->encryptIV) < 0) {
-    SAFE_FREE(out);
-    return NULL;
-  }
-
   if (session->version == 2) {
     ctx = hmac_init(session->current_crypto->encryptMAC, hmac_digest_len(type), type);
     if (ctx == NULL) {
diff --git a/src/server.c b/src/server.c
index a078e7af..951e1602 100644
--- a/src/server.c
+++ b/src/server.c
@@ -457,6 +457,14 @@ static void ssh_server_connection_callback(ssh_session session){
 			}
 			memcpy(session->next_crypto->session_id, session->current_crypto->session_id,
 			    session->current_crypto->digest_len);
+		    if (session->current_crypto->in_cipher->set_decrypt_key(session->current_crypto->in_cipher, session->current_crypto->decryptkey,
+		        session->current_crypto->decryptIV) < 0) {
+		      goto error;
+		    }
+		    if (session->current_crypto->out_cipher->set_encrypt_key(session->current_crypto->out_cipher, session->current_crypto->encryptkey,
+		        session->current_crypto->encryptIV) < 0) {
+		      goto error;
+		    }
 
 			    set_status(session,1.0f);
 			    session->connected = 1;