3 files changed, 58 insertions, 24 deletions

diff --git a/include/libssh/libssh.h b/include/libssh/libssh.h
index 20c057d..c89e31b 100644
--- a/include/libssh/libssh.h
+++ b/include/libssh/libssh.h
@@ -237,7 +237,7 @@ void ssh_print_hexa(const char *descr, const unsigned char *what, size_t len);
 int ssh_get_random(void *where,int len,int strong);
 
 /* this one can be called by the client to see the hash of the public key before accepting it */
-int ssh_get_pubkey_hash(SSH_SESSION *session,unsigned char hash[MD5_DIGEST_LEN]);
+int ssh_get_pubkey_hash(SSH_SESSION *session, unsigned char **hash);
 STRING *ssh_get_pubkey(SSH_SESSION *session);
 
 /* in connect.c */
diff --git a/libssh/dh.c b/libssh/dh.c
index c89c7c8..ad571dd 100644
--- a/libssh/dh.c
+++ b/libssh/dh.c
@@ -496,6 +496,7 @@ void generate_session_keys(SSH_SESSION *session){
     if(session->next_crypto->out_cipher->keysize > SHA_DIGEST_LEN*8){
         ctx=sha1_init();
         if (ctx == NULL) {
+          leave_function();
           return;
         }
         sha1_update(ctx,k_string,string_len(k_string)+4);
@@ -533,27 +534,52 @@ void generate_session_keys(SSH_SESSION *session){
 
 /** \addtogroup ssh_session
  * @{ */
-/** \brief get the md5 hash of the server public key
- * \param session ssh session
- * \param hash destination for the md5 hash
- * \return size of the hash in bytes
- * \warning it is very important that you verify at some moment that the hash matches
- * a known server. If you don't do it, cryptography won't help you at making things secure
- * \see ssh_is_server_known()
+/**
+ * @brief Allocates a buffer with the MD5 hash of the server public key.
+ *
+ * @param  session      The SSH session to use.
+ *
+ * @param  hash         The buffer to allocate.
+ *
+ * @return The bytes allocated or < 0 on error.
+ *
+ * @warning It is very important that you verify at some moment that the hash
+ *          matches a known server. If you don't do it, cryptography wont help
+ *          you at making things secure
+ *
+ * @see ssh_is_server_known()
+ * @see ssh_get_hexa()
+ * @see ssh_print_hexa()
  */
-int ssh_get_pubkey_hash(SSH_SESSION *session,unsigned char hash[MD5_DIGEST_LEN]){
-    STRING *pubkey=session->current_crypto->server_pubkey;
-    MD5CTX ctx;
-    int len=string_len(pubkey);
+int ssh_get_pubkey_hash(SSH_SESSION *session, unsigned char **hash) {
+  STRING *pubkey;
+  MD5CTX ctx;
+  unsigned char *h;
 
-    ctx=md5_init();
-    if (ctx == NULL) {
-      return 0;
-    }
+  if (session == NULL || hash == NULL) {
+    return -1;
+  }
+
+  *hash = NULL;
+
+  h = malloc(sizeof(unsigned char *) * MD5_DIGEST_LEN);
+  if (h == NULL) {
+    return -1;
+  }
+
+  ctx = md5_init();
+  if (ctx == NULL) {
+    return -1;
+  }
+
+  pubkey = session->current_crypto->server_pubkey;
+
+  md5_update(ctx, pubkey->string, string_len(pubkey));
+  md5_final(h, ctx);
+
+  *hash = h;
 
-    md5_update(ctx,pubkey->string,len);
-    md5_final(hash,ctx);
-    return MD5_DIGEST_LEN;
+  return MD5_DIGEST_LEN;
 }
 
 STRING *ssh_get_pubkey(SSH_SESSION *session){
diff --git a/sample.c b/sample.c
index 43a33db..002c326 100644
--- a/sample.c
+++ b/sample.c
@@ -406,7 +406,8 @@ int main(int argc, char **argv){
     char *hexa;
     int state;
     char buf[10];
-    unsigned char hash[MD5_DIGEST_LEN];
+    unsigned char *hash = NULL;
+    int hlen;
 
     options=ssh_options_new();
     if(ssh_options_getopt(options,&argc, argv)){
@@ -427,13 +428,20 @@ int main(int argc, char **argv){
         return 1;
     }
     state=ssh_is_server_known(session);
+
+    hlen = ssh_get_pubkey_hash(session, &hash);
+    if (hlen < 0) {
+      ssh_disconnect(session);
+      ssh_finalize();
+      return 1;
+    }
     switch(state){
         case SSH_SERVER_KNOWN_OK:
             break; /* ok */
         case SSH_SERVER_KNOWN_CHANGED:
             fprintf(stderr,"Host key for server changed : server's one is now :\n");
-            ssh_get_pubkey_hash(session,hash);
-            ssh_print_hexa("Public key hash",hash,MD5_DIGEST_LEN);
+            ssh_print_hexa("Public key hash",hash, hlen);
+            free(hash);
             fprintf(stderr,"For security reason, connection will be stopped\n");
             ssh_disconnect(session);
 	        ssh_finalize();
@@ -447,8 +455,8 @@ int main(int argc, char **argv){
 		    ssh_finalize();
             exit(-1);
         case SSH_SERVER_NOT_KNOWN:
-            ssh_get_pubkey_hash(session, hash);
-            hexa = ssh_get_hexa(hash, MD5_DIGEST_LEN);
+            hexa = ssh_get_hexa(hash, hlen);
+            free(hash);
             fprintf(stderr,"The server is unknown. Do you trust the host key ?\n");
             fprintf(stderr, "Public key hash: %s\n", hexa);
             free(hexa);