Add return value to dh_build_k().

git-svn-id: svn+ssh://svn.berlios.de/svnroot/repos/libssh/trunk@506 7dcaeef0-15fb-0310-b436-a5af3365683c

4 files changed, 46 insertions, 20 deletions

diff --git a/include/libssh/priv.h b/include/libssh/priv.h
index 4fe54a9e..704f0b1c 100644
--- a/include/libssh/priv.h
+++ b/include/libssh/priv.h
@@ -528,7 +528,7 @@ STRING *dh_get_f(SSH_SESSION *session);
 int dh_import_f(SSH_SESSION *session,STRING *f_string);
 int dh_import_e(SSH_SESSION *session, STRING *e_string);
 void dh_import_pubkey(SSH_SESSION *session,STRING *pubkey_string);
-void dh_build_k(SSH_SESSION *session);
+int dh_build_k(SSH_SESSION *session);
 int make_sessionid(SSH_SESSION *session);
 /* add data for the final cookie */
 int hashbufin_add_cookie(SSH_SESSION *session, unsigned char *cookie);
diff --git a/libssh/client.c b/libssh/client.c
index 166eb3a7..f8d95d9b 100644
--- a/libssh/client.c
+++ b/libssh/client.c
@@ -266,7 +266,11 @@ static int dh_handshake(SSH_SESSION *session) {
         goto error;
       }
       session->dh_server_signature = signature;
-      dh_build_k(session);
+      if (dh_build_k(session) < 0) {
+        ssh_set_error(session, SSH_FATAL, "Cannot build k number");
+        rc = SSH_ERROR;
+        goto error;
+      }
 
       /* Send the MSG_NEWKEYS */
       if (buffer_add_u8(session->out_buffer, SSH2_MSG_NEWKEYS) < 0) {
diff --git a/libssh/dh.c b/libssh/dh.c
index 4f37615f..fd581cb2 100644
--- a/libssh/dh.c
+++ b/libssh/dh.c
@@ -407,34 +407,54 @@ int dh_import_e(SSH_SESSION *session, STRING *e_string) {
   return 0;
 }
 
-void dh_build_k(SSH_SESSION *session){
+int dh_build_k(SSH_SESSION *session) {
 #ifdef HAVE_LIBCRYPTO
-    bignum_CTX ctx=bignum_ctx_new();
+  bignum_CTX ctx = bignum_ctx_new();
+  if (ctx == NULL) {
+    return -1;
+  }
+#endif
+
+  session->next_crypto->k = bignum_new();
+  if (session->next_crypto->k == NULL) {
+#ifdef HAVE_LIBCRYPTO
+    bignum_ctx_free(ctx);
 #endif
-    session->next_crypto->k=bignum_new();
+    return -1;
+  }
+
     /* the server and clients don't use the same numbers */
 #ifdef HAVE_LIBGCRYPT
-    if(session->client){
-        bignum_mod_exp(session->next_crypto->k,session->next_crypto->f,session->next_crypto->x,p);
-    } else {
-        bignum_mod_exp(session->next_crypto->k,session->next_crypto->e,session->next_crypto->y,p);
-    }
+  if(session->client) {
+    bignum_mod_exp(session->next_crypto->k, session->next_crypto->f,
+        session->next_crypto->x, p);
+  } else {
+    bignum_mod_exp(session->next_crypto->k, session->next_crypto->e,
+        session->next_crypto->y, p);
+  }
 #elif defined HAVE_LIBCRYPTO
-    if(session->client){
-        bignum_mod_exp(session->next_crypto->k,session->next_crypto->f,session->next_crypto->x,p,ctx);
-    } else {
-        bignum_mod_exp(session->next_crypto->k,session->next_crypto->e,session->next_crypto->y,p,ctx);
-    }
+  if (session->client) {
+    bignum_mod_exp(session->next_crypto->k, session->next_crypto->f,
+        session->next_crypto->x, p, ctx);
+  } else {
+    bignum_mod_exp(session->next_crypto->k, session->next_crypto->e,
+        session->next_crypto->y, p, ctx);
+  }
 #endif
+
 #ifdef DEBUG_CRYPTO
-    ssh_print_hexa("session server cookie",session->server_kex.cookie,16);
-    ssh_print_hexa("session client cookie",session->client_kex.cookie,16);
-    ssh_print_bignum("shared secret key",session->next_crypto->k);
+  ssh_print_hexa("Session server cookie", session->server_kex.cookie, 16);
+  ssh_print_hexa("Session client cookie", session->client_kex.cookie, 16);
+  ssh_print_bignum("Shared secret key", session->next_crypto->k);
 #endif
+
 #ifdef HAVE_LIBCRYPTO
-    bignum_ctx_free(ctx);
+  bignum_ctx_free(ctx);
 #endif
+
+  return 0;
 }
+
 /*
 static void sha_add(STRING *str,SHACTX ctx){
     sha1_update(ctx,str,string_len(str)+4);
diff --git a/libssh/server.c b/libssh/server.c
index 7b3602b9..0a2f13fe 100644
--- a/libssh/server.c
+++ b/libssh/server.c
@@ -300,7 +300,9 @@ static int dh_handshake_server(SSH_SESSION *session){
     pubkey=publickey_to_string(pub);
     publickey_free(pub);
     dh_import_pubkey(session,pubkey);
-    dh_build_k(session);
+    if (dh_build_k(session) < 0) {
+      return -1;
+    }
     if (make_sessionid(session) != SSH_OK) {
       return -1;
     }