diff --git a/drivers/auth/crypto_mod.c b/drivers/auth/crypto_mod.c
index c63ff08..127eb0d 100644
--- a/drivers/auth/crypto_mod.c
+++ b/drivers/auth/crypto_mod.c
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2015-2020, ARM Limited and Contributors. All rights reserved.
+ * Copyright (c) 2015-2021, Arm Limited and Contributors. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -114,8 +114,9 @@
  *   data_ptr, data_len: data to be hashed
  *   output: resulting hash
  */
-int crypto_mod_calc_hash(unsigned int alg, void *data_ptr,
-			 unsigned int data_len, unsigned char *output)
+int crypto_mod_calc_hash(enum crypto_md_algo alg, void *data_ptr,
+			 unsigned int data_len,
+			 unsigned char output[CRYPTO_MD_MAX_SIZE])
 {
 	assert(data_ptr != NULL);
 	assert(data_len != 0);
diff --git a/drivers/auth/mbedtls/mbedtls_common.mk b/drivers/auth/mbedtls/mbedtls_common.mk
index 53ebe30..54c819c 100644
--- a/drivers/auth/mbedtls/mbedtls_common.mk
+++ b/drivers/auth/mbedtls/mbedtls_common.mk
@@ -1,5 +1,5 @@
 #
-# Copyright (c) 2015-2020, Arm Limited. All rights reserved.
+# Copyright (c) 2015-2021, Arm Limited. All rights reserved.
 #
 # SPDX-License-Identifier: BSD-3-Clause
 #
@@ -96,6 +96,18 @@
     TF_MBEDTLS_USE_AES_GCM	:=	0
 endif
 
+ifeq ($(MEASURED_BOOT),1)
+    ifeq (${TPM_HASH_ALG}, sha256)
+        TF_MBEDTLS_TPM_HASH_ALG_ID	:=	TF_MBEDTLS_SHA256
+    else ifeq (${TPM_HASH_ALG}, sha384)
+        TF_MBEDTLS_TPM_HASH_ALG_ID	:=	TF_MBEDTLS_SHA384
+    else ifeq (${TPM_HASH_ALG}, sha512)
+        TF_MBEDTLS_TPM_HASH_ALG_ID	:=	TF_MBEDTLS_SHA512
+    else
+        $(error "TPM_HASH_ALG not defined.")
+    endif
+endif
+
 # Needs to be set to drive mbed TLS configuration correctly
 $(eval $(call add_defines,\
     $(sort \
@@ -105,6 +117,10 @@
         TF_MBEDTLS_USE_AES_GCM \
 )))
 
+ifeq ($(MEASURED_BOOT),1)
+  $(eval $(call add_define,TF_MBEDTLS_TPM_HASH_ALG_ID))
+endif
+
 $(eval $(call MAKE_LIB,mbedtls))
 
 endif
diff --git a/drivers/auth/mbedtls/mbedtls_crypto.c b/drivers/auth/mbedtls/mbedtls_crypto.c
index 6d6efb5..114e6ad 100644
--- a/drivers/auth/mbedtls/mbedtls_crypto.c
+++ b/drivers/auth/mbedtls/mbedtls_crypto.c
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2015-2020, ARM Limited and Contributors. All rights reserved.
+ * Copyright (c) 2015-2021, Arm Limited and Contributors. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -24,6 +24,16 @@
 
 #define LIB_NAME		"mbed TLS"
 
+#if MEASURED_BOOT
+/*
+ * CRYPTO_MD_MAX_SIZE value is as per current stronger algorithm available
+ * so make sure that mbed TLS MD maximum size must be lesser than this.
+ */
+CASSERT(CRYPTO_MD_MAX_SIZE >= MBEDTLS_MD_MAX_SIZE,
+	assert_mbedtls_md_size_overflow);
+
+#endif /* MEASURED_BOOT */
+
 /*
  * AlgorithmIdentifier  ::=  SEQUENCE  {
  *     algorithm               OBJECT IDENTIFIER,
@@ -211,21 +221,45 @@
 
 #if MEASURED_BOOT
 /*
+ * Map a generic crypto message digest algorithm to the corresponding macro used
+ * by Mbed TLS.
+ */
+static inline mbedtls_md_type_t md_type(enum crypto_md_algo algo)
+{
+	switch (algo) {
+	case CRYPTO_MD_SHA512:
+		return MBEDTLS_MD_SHA512;
+	case CRYPTO_MD_SHA384:
+		return MBEDTLS_MD_SHA384;
+	case CRYPTO_MD_SHA256:
+		return MBEDTLS_MD_SHA256;
+	default:
+		/* Invalid hash algorithm. */
+		return MBEDTLS_MD_NONE;
+	}
+}
+
+/*
  * Calculate a hash
  *
  * output points to the computed hash
  */
-int calc_hash(unsigned int alg, void *data_ptr,
-	      unsigned int data_len, unsigned char *output)
+static int calc_hash(enum crypto_md_algo md_algo, void *data_ptr,
+		     unsigned int data_len,
+		     unsigned char output[CRYPTO_MD_MAX_SIZE])
 {
 	const mbedtls_md_info_t *md_info;
 
-	md_info = mbedtls_md_info_from_type((mbedtls_md_type_t)alg);
+	md_info = mbedtls_md_info_from_type(md_type(md_algo));
 	if (md_info == NULL) {
 		return CRYPTO_ERR_HASH;
 	}
 
-	/* Calculate the hash of the data */
+	/*
+	 * Calculate the hash of the data, it is safe to pass the
+	 * 'output' hash buffer pointer considering its size is always
+	 * bigger than or equal to MBEDTLS_MD_MAX_SIZE.
+	 */
 	return mbedtls_md(md_info, data_ptr, data_len, output);
 }
 #endif /* MEASURED_BOOT */