From 8b66d037d575e5f3ce4d35964547ff8c7e75ff8e Mon Sep 17 00:00:00 2001
From: Jakub Jelen <jjelen@redhat.com>
Date: Fri, 15 Dec 2023 12:55:54 +0100
Subject: [PATCH 17/20] CVE-2023-6918: Systematically check return values when
 calculating digests

with all crypto backends

Signed-off-by: Jakub Jelen <jjelen@redhat.com>
Reviewed-by: Andreas Schneider <asn@cryptomilk.org>
---
 include/libssh/wrapper.h |  34 ++++----
 src/kdf.c                |  96 ++++++++++++++++++-----
 src/md_crypto.c          | 161 ++++++++++++++++++++++++++++++--------
 src/md_gcrypt.c          | 107 +++++++++++++++++++++----
 src/md_mbedcrypto.c      | 165 +++++++++++++++++++++++++++++++--------
 src/session.c            |  72 ++++++++++++-----
 6 files changed, 504 insertions(+), 131 deletions(-)

diff --git a/include/libssh/wrapper.h b/include/libssh/wrapper.h
index 07e64018..b3e28eac 100644
--- a/include/libssh/wrapper.h
+++ b/include/libssh/wrapper.h
@@ -72,29 +72,33 @@ struct ssh_crypto_struct;
 
 typedef struct ssh_mac_ctx_struct *ssh_mac_ctx;
 MD5CTX md5_init(void);
-void md5_update(MD5CTX c, const void *data, size_t len);
-void md5_final(unsigned char *md,MD5CTX c);
+void md5_ctx_free(MD5CTX);
+int md5_update(MD5CTX c, const void *data, size_t len);
+int md5_final(unsigned char *md, MD5CTX c);
 
 SHACTX sha1_init(void);
-void sha1_update(SHACTX c, const void *data, size_t len);
-void sha1_final(unsigned char *md,SHACTX c);
-void sha1(const unsigned char *digest,size_t len,unsigned char *hash);
+void sha1_ctx_free(SHACTX);
+int sha1_update(SHACTX c, const void *data, size_t len);
+int sha1_final(unsigned char *md,SHACTX c);
+int sha1(const unsigned char *digest,size_t len, unsigned char *hash);
 
 SHA256CTX sha256_init(void);
-void sha256_update(SHA256CTX c, const void *data, size_t len);
-void sha256_final(unsigned char *md,SHA256CTX c);
-void sha256(const unsigned char *digest, size_t len, unsigned char *hash);
+void sha256_ctx_free(SHA256CTX);
+int sha256_update(SHA256CTX c, const void *data, size_t len);
+int sha256_final(unsigned char *md,SHA256CTX c);
+int sha256(const unsigned char *digest, size_t len, unsigned char *hash);
 
 SHA384CTX sha384_init(void);
-void sha384_update(SHA384CTX c, const void *data, size_t len);
-void sha384_final(unsigned char *md,SHA384CTX c);
-void sha384(const unsigned char *digest, size_t len, unsigned char *hash);
+void sha384_ctx_free(SHA384CTX);
+int sha384_update(SHA384CTX c, const void *data, size_t len);
+int sha384_final(unsigned char *md,SHA384CTX c);
+int sha384(const unsigned char *digest, size_t len, unsigned char *hash);
 
 SHA512CTX sha512_init(void);
-void sha512_update(SHA512CTX c, const void *data, size_t len);
-void sha512_final(unsigned char *md,SHA512CTX c);
-void sha512(const unsigned char *digest, size_t len, unsigned char *hash);
-
+void sha512_ctx_free(SHA512CTX);
+int sha512_update(SHA512CTX c, const void *data, size_t len);
+int sha512_final(unsigned char *md,SHA512CTX c);
+int sha512(const unsigned char *digest, size_t len, unsigned char *hash);
 
 HMACCTX hmac_init(const void *key,size_t len, enum ssh_hmac_e type);
 int hmac_update(HMACCTX c, const void *data, size_t len);
diff --git a/src/kdf.c b/src/kdf.c
index 987ae972..a8e534e5 100644
--- a/src/kdf.c
+++ b/src/kdf.c
@@ -77,41 +77,64 @@ static ssh_mac_ctx ssh_mac_ctx_init(enum ssh_kdf_digest type)
     }
 }
 
-static void ssh_mac_update(ssh_mac_ctx ctx, const void *data, size_t len)
+static void ssh_mac_ctx_free(ssh_mac_ctx ctx)
 {
+    if (ctx == NULL) {
+        return;
+    }
+
     switch (ctx->digest_type) {
     case SSH_KDF_SHA1:
-        sha1_update(ctx->ctx.sha1_ctx, data, len);
+        sha1_ctx_free(ctx->ctx.sha1_ctx);
         break;
     case SSH_KDF_SHA256:
-        sha256_update(ctx->ctx.sha256_ctx, data, len);
+        sha256_ctx_free(ctx->ctx.sha256_ctx);
         break;
     case SSH_KDF_SHA384:
-        sha384_update(ctx->ctx.sha384_ctx, data, len);
+        sha384_ctx_free(ctx->ctx.sha384_ctx);
         break;
     case SSH_KDF_SHA512:
-        sha512_update(ctx->ctx.sha512_ctx, data, len);
+        sha512_ctx_free(ctx->ctx.sha512_ctx);
         break;
     }
+    SAFE_FREE(ctx);
+}
+
+static int ssh_mac_update(ssh_mac_ctx ctx, const void *data, size_t len)
+{
+    switch (ctx->digest_type) {
+    case SSH_KDF_SHA1:
+        return sha1_update(ctx->ctx.sha1_ctx, data, len);
+    case SSH_KDF_SHA256:
+        return sha256_update(ctx->ctx.sha256_ctx, data, len);
+    case SSH_KDF_SHA384:
+        return sha384_update(ctx->ctx.sha384_ctx, data, len);
+    case SSH_KDF_SHA512:
+        return sha512_update(ctx->ctx.sha512_ctx, data, len);
+    }
+    return SSH_ERROR;
 }
 
-static void ssh_mac_final(unsigned char *md, ssh_mac_ctx ctx)
+static int ssh_mac_final(unsigned char *md, ssh_mac_ctx ctx)
 {
+    int rc = SSH_ERROR;
+
     switch (ctx->digest_type) {
     case SSH_KDF_SHA1:
-        sha1_final(md, ctx->ctx.sha1_ctx);
+        rc = sha1_final(md, ctx->ctx.sha1_ctx);
         break;
     case SSH_KDF_SHA256:
-        sha256_final(md, ctx->ctx.sha256_ctx);
+        rc = sha256_final(md, ctx->ctx.sha256_ctx);
         break;
     case SSH_KDF_SHA384:
-        sha384_final(md, ctx->ctx.sha384_ctx);
+        rc = sha384_final(md, ctx->ctx.sha384_ctx);
         break;
     case SSH_KDF_SHA512:
-        sha512_final(md, ctx->ctx.sha512_ctx);
+        rc = sha512_final(md, ctx->ctx.sha512_ctx);
         break;
     }
     SAFE_FREE(ctx);
+    return rc;
 }
 
 int sshkdf_derive_key(struct ssh_crypto_struct *crypto,
@@ -126,6 +149,7 @@ int sshkdf_derive_key(struct ssh_crypto_struct *crypto,
     unsigned char digest[DIGEST_MAX_LEN];
     size_t output_len = crypto->digest_len;
     ssh_mac_ctx ctx;
+    int rc;
 
     if (DIGEST_MAX_LEN < crypto->digest_len) {
         return -1;
@@ -136,11 +160,30 @@ int sshkdf_derive_key(struct ssh_crypto_struct *crypto,
         return -1;
     }
 
-    ssh_mac_update(ctx, key, key_len);
-    ssh_mac_update(ctx, crypto->secret_hash, crypto->digest_len);
-    ssh_mac_update(ctx, &key_type, 1);
-    ssh_mac_update(ctx, crypto->session_id, crypto->session_id_len);
-    ssh_mac_final(digest, ctx);
+    rc = ssh_mac_update(ctx, key, key_len);
+    if (rc != SSH_OK) {
+        ssh_mac_ctx_free(ctx);
+        return -1;
+    }
+    rc = ssh_mac_update(ctx, crypto->secret_hash, crypto->digest_len);
+    if (rc != SSH_OK) {
+        ssh_mac_ctx_free(ctx);
+        return -1;
+    }
+    rc = ssh_mac_update(ctx, &key_type, 1);
+    if (rc != SSH_OK) {
+        ssh_mac_ctx_free(ctx);
+        return -1;
+    }
+    rc = ssh_mac_update(ctx, crypto->session_id, crypto->session_id_len);
+    if (rc != SSH_OK) {
+        ssh_mac_ctx_free(ctx);
+        return -1;
+    }
+    rc = ssh_mac_final(digest, ctx);
+    if (rc != SSH_OK) {
+        return -1;
+    }
 
     if (requested_len < output_len) {
         output_len = requested_len;
@@ -152,10 +195,25 @@ int sshkdf_derive_key(struct ssh_crypto_struct *crypto,
         if (ctx == NULL) {
             return -1;
         }
-        ssh_mac_update(ctx, key, key_len);
-        ssh_mac_update(ctx, crypto->secret_hash, crypto->digest_len);
-        ssh_mac_update(ctx, output, output_len);
-        ssh_mac_final(digest, ctx);
+        rc = ssh_mac_update(ctx, key, key_len);
+        if (rc != SSH_OK) {
+            ssh_mac_ctx_free(ctx);
+            return -1;
+        }
+        rc = ssh_mac_update(ctx, crypto->secret_hash, crypto->digest_len);
+        if (rc != SSH_OK) {
+            ssh_mac_ctx_free(ctx);
+            return -1;
+        }
+        rc = ssh_mac_update(ctx, output, output_len);
+        if (rc != SSH_OK) {
+            ssh_mac_ctx_free(ctx);
+            return -1;
+        }
+        rc = ssh_mac_final(digest, ctx);
+        if (rc != SSH_OK) {
+            return -1;
+        }
         if (requested_len < output_len + crypto->digest_len) {
             memcpy(output + output_len, digest, requested_len - output_len);
         } else {
diff --git a/src/md_crypto.c b/src/md_crypto.c
index f5104f04..f7cda8dd 100644
--- a/src/md_crypto.c
+++ b/src/md_crypto.c
@@ -25,6 +25,7 @@
 #include "libssh/crypto.h"
 #include "libssh/wrapper.h"
 
+#include <openssl/crypto.h>
 #include <openssl/evp.h>
 #include <openssl/md5.h>
 #include <openssl/sha.h>
@@ -46,28 +47,49 @@ sha1_init(void)
 }
 
 void
+sha1_ctx_free(SHACTX c)
+{
+    EVP_MD_CTX_free(c);
+}
+
+int
 sha1_update(SHACTX c, const void *data, size_t len)
 {
-    EVP_DigestUpdate(c, data, len);
+    int rc = EVP_DigestUpdate(c, data, len);
+    if (rc != 1) {
+        return SSH_ERROR;
+    }
+    return SSH_OK;
 }
 
-void
+int
 sha1_final(unsigned char *md, SHACTX c)
 {
     unsigned int mdlen = 0;
+    int rc = EVP_DigestFinal(c, md, &mdlen);
 
-    EVP_DigestFinal(c, md, &mdlen);
     EVP_MD_CTX_free(c);
+    if (rc != 1) {
+        return SSH_ERROR;
+    }
+    return SSH_OK;
 }
 
-void
+int
 sha1(const unsigned char *digest, size_t len, unsigned char *hash)
 {
     SHACTX c = sha1_init();
-    if (c != NULL) {
-        sha1_update(c, digest, len);
-        sha1_final(hash, c);
+    int rc;
+
+    if (c == NULL) {
+        return SSH_ERROR;
     }
+    rc = sha1_update(c, digest, len);
+    if (rc != SSH_OK) {
+        EVP_MD_CTX_free(c);
+        return SSH_ERROR;
+    }
+    return sha1_final(hash, c);
 }
 
 SHA256CTX
@@ -87,28 +109,49 @@ sha256_init(void)
 }
 
 void
+sha256_ctx_free(SHA256CTX c)
+{
+    EVP_MD_CTX_free(c);
+}
+
+int
 sha256_update(SHA256CTX c, const void *data, size_t len)
 {
-    EVP_DigestUpdate(c, data, len);
+    int rc = EVP_DigestUpdate(c, data, len);
+    if (rc != 1) {
+        return SSH_ERROR;
+    }
+    return SSH_OK;
 }
 
-void
+int
 sha256_final(unsigned char *md, SHA256CTX c)
 {
     unsigned int mdlen = 0;
+    int rc = EVP_DigestFinal(c, md, &mdlen);
 
-    EVP_DigestFinal(c, md, &mdlen);
     EVP_MD_CTX_free(c);
+    if (rc != 1) {
+        return SSH_ERROR;
+    }
+    return SSH_OK;
 }
 
-void
+int
 sha256(const unsigned char *digest, size_t len, unsigned char *hash)
 {
     SHA256CTX c = sha256_init();
-    if (c != NULL) {
-        sha256_update(c, digest, len);
-        sha256_final(hash, c);
+    int rc;
+
+    if (c == NULL) {
+        return SSH_ERROR;
+    }
+    rc = sha256_update(c, digest, len);
+    if (rc != SSH_OK) {
+        EVP_MD_CTX_free(c);
+        return SSH_ERROR;
     }
+    return sha256_final(hash, c);
 }
 
 SHA384CTX
@@ -128,28 +171,49 @@ sha384_init(void)
 }
 
 void
+sha384_ctx_free(SHA384CTX c)
+{
+    EVP_MD_CTX_free(c);
+}
+
+int
 sha384_update(SHA384CTX c, const void *data, size_t len)
 {
-    EVP_DigestUpdate(c, data, len);
+    int rc = EVP_DigestUpdate(c, data, len);
+    if (rc != 1) {
+        return SSH_ERROR;
+    }
+    return SSH_OK;
 }
 
-void
+int
 sha384_final(unsigned char *md, SHA384CTX c)
 {
     unsigned int mdlen = 0;
+    int rc = EVP_DigestFinal(c, md, &mdlen);
 
-    EVP_DigestFinal(c, md, &mdlen);
     EVP_MD_CTX_free(c);
+    if (rc != 1) {
+        return SSH_ERROR;
+    }
+    return SSH_OK;
 }
 
-void
+int
 sha384(const unsigned char *digest, size_t len, unsigned char *hash)
 {
     SHA384CTX c = sha384_init();
-    if (c != NULL) {
-        sha384_update(c, digest, len);
-        sha384_final(hash, c);
+    int rc;
+
+    if (c == NULL) {
+        return SSH_ERROR;
     }
+    rc = sha384_update(c, digest, len);
+    if (rc != SSH_OK) {
+        EVP_MD_CTX_free(c);
+        return SSH_ERROR;
+    }
+    return sha384_final(hash, c);
 }
 
 SHA512CTX
@@ -169,28 +233,49 @@ sha512_init(void)
 }
 
 void
+sha512_ctx_free(SHA512CTX c)
+{
+    EVP_MD_CTX_free(c);
+}
+
+int
 sha512_update(SHA512CTX c, const void *data, size_t len)
 {
-    EVP_DigestUpdate(c, data, len);
+    int rc = EVP_DigestUpdate(c, data, len);
+    if (rc != 1) {
+        return SSH_ERROR;
+    }
+    return SSH_OK;
 }
 
-void
+int
 sha512_final(unsigned char *md, SHA512CTX c)
 {
     unsigned int mdlen = 0;
+    int rc = EVP_DigestFinal(c, md, &mdlen);
 
-    EVP_DigestFinal(c, md, &mdlen);
     EVP_MD_CTX_free(c);
+    if (rc != 1) {
+        return SSH_ERROR;
+    }
+    return SSH_OK;
 }
 
-void
+int
 sha512(const unsigned char *digest, size_t len, unsigned char *hash)
 {
     SHA512CTX c = sha512_init();
-    if (c != NULL) {
-        sha512_update(c, digest, len);
-        sha512_final(hash, c);
+    int rc;
+
+    if (c == NULL) {
+        return SSH_ERROR;
+    }
+    rc = sha512_update(c, digest, len);
+    if (rc != SSH_OK) {
+        EVP_MD_CTX_free(c);
+        return SSH_ERROR;
     }
+    return sha512_final(hash, c);
 }
 
 MD5CTX
@@ -210,16 +295,30 @@ md5_init(void)
 }
 
 void
+md5_ctx_free(MD5CTX c)
+{
+    EVP_MD_CTX_free(c);
+}
+
+int
 md5_update(MD5CTX c, const void *data, size_t len)
 {
-    EVP_DigestUpdate(c, data, len);
+    int rc = EVP_DigestUpdate(c, data, len);
+    if (rc != 1) {
+        return SSH_ERROR;
+    }
+    return SSH_OK;
 }
 
-void
+int
 md5_final(unsigned char *md, MD5CTX c)
 {
     unsigned int mdlen = 0;
+    int rc = EVP_DigestFinal(c, md, &mdlen);
 
-    EVP_DigestFinal(c, md, &mdlen);
     EVP_MD_CTX_free(c);
+    if (rc != 1) {
+        return SSH_ERROR;
+    }
+    return SSH_OK;
 }
diff --git a/src/md_gcrypt.c b/src/md_gcrypt.c
index 1f0a71f3..93c7b0d9 100644
--- a/src/md_gcrypt.c
+++ b/src/md_gcrypt.c
@@ -36,24 +36,40 @@ sha1_init(void)
     return ctx;
 }
 
-void
+int
 sha1_update(SHACTX c, const void *data, size_t len)
 {
     gcry_md_write(c, data, len);
+    return SSH_OK;
 }
 
 void
+sha1_ctx_free(SHACTX c)
+{
+    gcry_md_close(c);
+}
+
+int
 sha1_final(unsigned char *md, SHACTX c)
 {
+    unsigned char *tmp = NULL;
+
     gcry_md_final(c);
-    memcpy(md, gcry_md_read(c, 0), SHA_DIGEST_LEN);
+    tmp = gcry_md_read(c, 0);
+    if (tmp == NULL) {
+        gcry_md_close(c);
+        return SSH_ERROR;
+    }
+    memcpy(md, tmp, SHA_DIGEST_LEN);
     gcry_md_close(c);
+    return SSH_OK;
 }
 
-void
+int
 sha1(const unsigned char *digest, size_t len, unsigned char *hash)
 {
     gcry_md_hash_buffer(GCRY_MD_SHA1, hash, digest, len);
+    return SSH_OK;
 }
 
 SHA256CTX
@@ -66,23 +82,39 @@ sha256_init(void)
 }
 
 void
+sha256_ctx_free(SHA256CTX c)
+{
+    gcry_md_close(c);
+}
+
+int
 sha256_update(SHACTX c, const void *data, size_t len)
 {
     gcry_md_write(c, data, len);
+    return SSH_OK;
 }
 
-void
+int
 sha256_final(unsigned char *md, SHACTX c)
 {
+    unsigned char *tmp = NULL;
+
     gcry_md_final(c);
-    memcpy(md, gcry_md_read(c, 0), SHA256_DIGEST_LEN);
+    tmp = gcry_md_read(c, 0);
+    if (tmp == NULL) {
+        gcry_md_close(c);
+        return SSH_ERROR;
+    }
+    memcpy(md, tmp, SHA256_DIGEST_LEN);
     gcry_md_close(c);
+    return SSH_OK;
 }
 
-void
+int
 sha256(const unsigned char *digest, size_t len, unsigned char *hash)
 {
     gcry_md_hash_buffer(GCRY_MD_SHA256, hash, digest, len);
+    return SSH_OK;
 }
 
 SHA384CTX
@@ -95,23 +127,39 @@ sha384_init(void)
 }
 
 void
+sha384_ctx_free(SHA384CTX c)
+{
+    gcry_md_close(c);
+}
+
+int
 sha384_update(SHACTX c, const void *data, size_t len)
 {
     gcry_md_write(c, data, len);
+    return SSH_OK;
 }
 
-void
+int
 sha384_final(unsigned char *md, SHACTX c)
 {
+    unsigned char *tmp = NULL;
+
     gcry_md_final(c);
-    memcpy(md, gcry_md_read(c, 0), SHA384_DIGEST_LEN);
+    tmp = gcry_md_read(c, 0);
+    if (tmp == NULL) {
+        gcry_md_close(c);
+        return SSH_ERROR;
+    }
+    memcpy(md, tmp, SHA384_DIGEST_LEN);
     gcry_md_close(c);
+    return SSH_OK;
 }
 
-void
+int
 sha384(const unsigned char *digest, size_t len, unsigned char *hash)
 {
     gcry_md_hash_buffer(GCRY_MD_SHA384, hash, digest, len);
+    return SSH_OK;
 }
 
 SHA512CTX
@@ -124,23 +172,39 @@ sha512_init(void)
 }
 
 void
+sha512_ctx_free(SHA512CTX c)
+{
+    gcry_md_close(c);
+}
+
+int
 sha512_update(SHACTX c, const void *data, size_t len)
 {
     gcry_md_write(c, data, len);
+    return SSH_OK;
 }
 
-void
+int
 sha512_final(unsigned char *md, SHACTX c)
 {
+    unsigned char *tmp = NULL;
+
     gcry_md_final(c);
-    memcpy(md, gcry_md_read(c, 0), SHA512_DIGEST_LEN);
+    tmp = gcry_md_read(c, 0);
+    if (tmp == NULL) {
+        gcry_md_close(c);
+        return SSH_ERROR;
+    }
+    memcpy(md, tmp, SHA512_DIGEST_LEN);
     gcry_md_close(c);
+    return SSH_OK;
 }
 
-void
+int
 sha512(const unsigned char *digest, size_t len, unsigned char *hash)
 {
     gcry_md_hash_buffer(GCRY_MD_SHA512, hash, digest, len);
+    return SSH_OK;
 }
 
 MD5CTX
@@ -153,15 +217,30 @@ md5_init(void)
 }
 
 void
+md5_ctx_free(MD5CTX c)
+{
+    gcry_md_close(c);
+}
+
+int
 md5_update(MD5CTX c, const void *data, size_t len)
 {
     gcry_md_write(c, data, len);
+    return SSH_OK;
 }
 
-void
+int
 md5_final(unsigned char *md, MD5CTX c)
 {
+    unsigned char *tmp = NULL;
+
     gcry_md_final(c);
-    memcpy(md, gcry_md_read(c, 0), MD5_DIGEST_LEN);
+    tmp = gcry_md_read(c, 0);
+    if (tmp == NULL) {
+        gcry_md_close(c);
+        return SSH_ERROR;
+    }
+    memcpy(md, tmp, MD5_DIGEST_LEN);
     gcry_md_close(c);
+    return SSH_OK;
 }
diff --git a/src/md_mbedcrypto.c b/src/md_mbedcrypto.c
index 227e20ab..b3529b4b 100644
--- a/src/md_mbedcrypto.c
+++ b/src/md_mbedcrypto.c
@@ -64,27 +64,48 @@ sha1_init(void)
 }
 
 void
+sha1_ctx_free(SHACTX c)
+{
+    mbedtls_md_free(c);
+    SAFE_FREE(c);
+}
+
+int
 sha1_update(SHACTX c, const void *data, size_t len)
 {
-    mbedtls_md_update(c, data, len);
+    int rc = mbedtls_md_update(c, data, len);
+    if (rc != 0) {
+        return SSH_ERROR;
+    }
+    return SSH_OK;
 }
 
-void
+int
 sha1_final(unsigned char *md, SHACTX c)
 {
-    mbedtls_md_finish(c, md);
-    mbedtls_md_free(c);
-    SAFE_FREE(c);
+    int rc = mbedtls_md_finish(c, md);
+    sha1_ctx_free(c);
+    if (rc != 0) {
+        return SSH_ERROR;
+    }
+    return SSH_OK;
 }
 
-void
+int
 sha1(const unsigned char *digest, size_t len, unsigned char *hash)
 {
     const mbedtls_md_info_t *md_info =
         mbedtls_md_info_from_type(MBEDTLS_MD_SHA1);
-    if (md_info != NULL) {
-        mbedtls_md(md_info, digest, len, hash);
+    int rc;
+
+    if (md_info == NULL) {
+        return SSH_ERROR;
+    }
+    rc = mbedtls_md(md_info, digest, len, hash);
+    if (rc != 0) {
+        return SSH_ERROR;
     }
+    return SSH_OK;
 }
 
 SHA256CTX
@@ -122,27 +143,48 @@ sha256_init(void)
 }
 
 void
+sha256_ctx_free(SHA256CTX c)
+{
+    mbedtls_md_free(c);
+    SAFE_FREE(c);
+}
+
+int
 sha256_update(SHA256CTX c, const void *data, size_t len)
 {
-    mbedtls_md_update(c, data, len);
+    int rc = mbedtls_md_update(c, data, len);
+    if (rc != 0) {
+        return SSH_ERROR;
+    }
+    return SSH_OK;
 }
 
-void
+int
 sha256_final(unsigned char *md, SHA256CTX c)
 {
-    mbedtls_md_finish(c, md);
+    int rc = mbedtls_md_finish(c, md);
     mbedtls_md_free(c);
     SAFE_FREE(c);
+    if (rc != 0) {
+        return SSH_ERROR;
+    }
+    return SSH_OK;
 }
 
-void
+int
 sha256(const unsigned char *digest, size_t len, unsigned char *hash)
 {
+    int rc;
     const mbedtls_md_info_t *md_info =
         mbedtls_md_info_from_type(MBEDTLS_MD_SHA256);
-    if (md_info != NULL) {
-        mbedtls_md(md_info, digest, len, hash);
+    if (md_info == NULL) {
+        return SSH_ERROR;
+    }
+    rc = mbedtls_md(md_info, digest, len, hash);
+    if (rc != 0) {
+        return SSH_ERROR;
     }
+    return SSH_OK;
 }
 
 SHA384CTX
@@ -180,27 +222,48 @@ sha384_init(void)
 }
 
 void
+sha384_ctx_free(SHA384CTX c)
+{
+    mbedtls_md_free(c);
+    SAFE_FREE(c);
+}
+
+int
 sha384_update(SHA384CTX c, const void *data, size_t len)
 {
-    mbedtls_md_update(c, data, len);
+    int rc = mbedtls_md_update(c, data, len);
+    if (rc != 0) {
+        return SSH_ERROR;
+    }
+    return SSH_OK;
 }
 
-void
+int
 sha384_final(unsigned char *md, SHA384CTX c)
 {
-    mbedtls_md_finish(c, md);
-    mbedtls_md_free(c);
-    SAFE_FREE(c);
+    int rc = mbedtls_md_finish(c, md);
+    sha384_ctx_free(c);
+    if (rc != 0) {
+        return SSH_ERROR;
+    }
+    return SSH_OK;
 }
 
-void
+int
 sha384(const unsigned char *digest, size_t len, unsigned char *hash)
 {
     const mbedtls_md_info_t *md_info =
         mbedtls_md_info_from_type(MBEDTLS_MD_SHA384);
-    if (md_info != NULL) {
-        mbedtls_md(md_info, digest, len, hash);
+    int rc;
+
+    if (md_info == NULL) {
+        return SSH_ERROR;
+    }
+    rc = mbedtls_md(md_info, digest, len, hash);
+    if (rc != 0) {
+        return SSH_ERROR;
     }
+    return SSH_OK;
 }
 
 SHA512CTX
@@ -237,27 +300,48 @@ sha512_init(void)
 }
 
 void
+sha512_ctx_free(SHA512CTX c)
+{
+    mbedtls_md_free(c);
+    SAFE_FREE(c);
+}
+
+int
 sha512_update(SHA512CTX c, const void *data, size_t len)
 {
-    mbedtls_md_update(c, data, len);
+    int rc = mbedtls_md_update(c, data, len);
+    if (rc != 0) {
+        return SSH_ERROR;
+    }
+    return SSH_OK;
 }
 
-void
+int
 sha512_final(unsigned char *md, SHA512CTX c)
 {
-    mbedtls_md_finish(c, md);
-    mbedtls_md_free(c);
-    SAFE_FREE(c);
+    int rc = mbedtls_md_finish(c, md);
+    sha512_ctx_free(c);
+    if (rc != 0) {
+        return SSH_ERROR;
+    }
+    return SSH_OK;
 }
 
-void
+int
 sha512(const unsigned char *digest, size_t len, unsigned char *hash)
 {
     const mbedtls_md_info_t *md_info =
         mbedtls_md_info_from_type(MBEDTLS_MD_SHA512);
-    if (md_info != NULL) {
-        mbedtls_md(md_info, digest, len, hash);
+    int rc;
+
+    if (md_info == NULL) {
+        return SSH_ERROR;
     }
+    rc = mbedtls_md(md_info, digest, len, hash);
+    if (rc != 0) {
+        return SSH_ERROR;
+    }
+    return SSH_OK;
 }
 
 MD5CTX
@@ -294,15 +378,30 @@ md5_init(void)
 }
 
 void
+md5_ctx_free(MD5CTX c)
+{
+    mbedtls_md_free(c);
+    SAFE_FREE(c);
+}
+
+int
 md5_update(MD5CTX c, const void *data, size_t len)
 {
-    mbedtls_md_update(c, data, len);
+    int rc = mbedtls_md_update(c, data, len);
+    if (rc != 0) {
+        return SSH_ERROR;
+    }
+    return SSH_OK;
 }
 
-void
+int
 md5_final(unsigned char *md, MD5CTX c)
 {
-    mbedtls_md_finish(c, md);
+    int rc = mbedtls_md_finish(c, md);
     mbedtls_md_free(c);
     SAFE_FREE(c);
+    if (rc != 0) {
+        return SSH_ERROR;
+    }
+    return SSH_OK;
 }
diff --git a/src/session.c b/src/session.c
index fe998dee..8c509699 100644
--- a/src/session.c
+++ b/src/session.c
@@ -1024,7 +1024,18 @@ int ssh_get_pubkey_hash(ssh_session session, unsigned char **hash)
     *hash = NULL;
     if (session->current_crypto == NULL ||
         session->current_crypto->server_pubkey == NULL) {
-        ssh_set_error(session,SSH_FATAL,"No current cryptographic context");
+        ssh_set_error(session, SSH_FATAL, "No current cryptographic context");
+        return SSH_ERROR;
+    }
+
+    rc = ssh_get_server_publickey(session, &pubkey);
+    if (rc != SSH_OK) {
+        return SSH_ERROR;
+    }
+
+    rc = ssh_pki_export_pubkey_blob(pubkey, &pubkey_blob);
+    ssh_key_free(pubkey);
+    if (rc != SSH_OK) {
         return SSH_ERROR;
     }
 
@@ -1039,25 +1050,21 @@ int ssh_get_pubkey_hash(ssh_session session, unsigned char **hash)
         return SSH_ERROR;
     }
 
-    rc = ssh_get_server_publickey(session, &pubkey);
+    rc = md5_update(ctx,
+                    ssh_string_data(pubkey_blob),
+                    ssh_string_len(pubkey_blob));
     if (rc != SSH_OK) {
-        md5_final(h, ctx);
+        md5_ctx_free(ctx);
         SAFE_FREE(h);
-        return SSH_ERROR;
+        return rc;
     }
-
-    rc = ssh_pki_export_pubkey_blob(pubkey, &pubkey_blob);
-    ssh_key_free(pubkey);
+    SSH_STRING_FREE(pubkey_blob);
+    rc = md5_final(h, ctx);
     if (rc != SSH_OK) {
-        md5_final(h, ctx);
         SAFE_FREE(h);
-        return SSH_ERROR;
+        return rc;
     }
 
-    md5_update(ctx, ssh_string_data(pubkey_blob), ssh_string_len(pubkey_blob));
-    SSH_STRING_FREE(pubkey_blob);
-    md5_final(h, ctx);
-
     *hash = h;
 
     return MD5_DIGEST_LEN;
@@ -1177,8 +1184,17 @@ int ssh_get_publickey_hash(const ssh_key key,
                 goto out;
             }
 
-            sha1_update(ctx, ssh_string_data(blob), ssh_string_len(blob));
-            sha1_final(h, ctx);
+            rc = sha1_update(ctx, ssh_string_data(blob), ssh_string_len(blob));
+            if (rc != SSH_OK) {
+                free(h);
+                sha1_ctx_free(ctx);
+                goto out;
+            }
+            rc = sha1_final(h, ctx);
+            if (rc != SSH_OK) {
+                free(h);
+                goto out;
+            }
 
             *hlen = SHA_DIGEST_LEN;
         }
@@ -1200,8 +1216,17 @@ int ssh_get_publickey_hash(const ssh_key key,
                 goto out;
             }
 
-            sha256_update(ctx, ssh_string_data(blob), ssh_string_len(blob));
-            sha256_final(h, ctx);
+            rc = sha256_update(ctx, ssh_string_data(blob), ssh_string_len(blob));
+            if (rc != SSH_OK) {
+                free(h);
+                sha256_ctx_free(ctx);
+                goto out;
+            }
+            rc = sha256_final(h, ctx);
+            if (rc != SSH_OK) {
+                free(h);
+                goto out;
+            }
 
             *hlen = SHA256_DIGEST_LEN;
         }
@@ -1231,8 +1256,17 @@ int ssh_get_publickey_hash(const ssh_key key,
                 goto out;
             }
 
-            md5_update(ctx, ssh_string_data(blob), ssh_string_len(blob));
-            md5_final(h, ctx);
+            rc = md5_update(ctx, ssh_string_data(blob), ssh_string_len(blob));
+            if (rc != SSH_OK) {
+                free(h);
+                md5_ctx_free(ctx);
+                goto out;
+            }
+            rc = md5_final(h, ctx);
+            if (rc != SSH_OK) {
+                free(h);
+                goto out;
+            }
 
             *hlen = MD5_DIGEST_LEN;
         }
-- 
2.33.0