diff --git a/drivers/auth/mbedtls/mbedtls_common.c b/drivers/auth/mbedtls/mbedtls_common.c
index 7095fde..dbf45ba 100644
--- a/drivers/auth/mbedtls/mbedtls_common.c
+++ b/drivers/auth/mbedtls/mbedtls_common.c
@@ -4,26 +4,15 @@
  * SPDX-License-Identifier: BSD-3-Clause
  */
 
+#include <assert.h>
 #include <debug.h>
-#include <stdlib.h>
-#include <stdio.h>
-
 /* mbed TLS headers */
 #include <mbedtls/memory_buffer_alloc.h>
 #include <mbedtls/platform.h>
-#include <mbedtls_config.h>
 #include <mbedtls_common.h>
-
-/*
- * mbed TLS heap
- */
-#if (TF_MBEDTLS_KEY_ALG_ID == TF_MBEDTLS_ECDSA) \
-	|| (TF_MBEDTLS_KEY_ALG_ID == TF_MBEDTLS_RSA_AND_ECDSA)
-#define MBEDTLS_HEAP_SIZE		(13*1024)
-#elif (TF_MBEDTLS_KEY_ALG_ID == TF_MBEDTLS_RSA)
-#define MBEDTLS_HEAP_SIZE		(7*1024)
-#endif
-static unsigned char heap[MBEDTLS_HEAP_SIZE];
+#include <mbedtls_config.h>
+#include <platform.h>
+#include <stddef.h>
 
 static void cleanup(void)
 {
@@ -37,13 +26,25 @@
 void mbedtls_init(void)
 {
 	static int ready;
+	void *heap_addr;
+	size_t heap_size = 0;
+	int err;
 
 	if (!ready) {
 		if (atexit(cleanup))
 			panic();
+
+		err = plat_get_mbedtls_heap(&heap_addr, &heap_size);
+
+		/* Ensure heap setup is proper */
+		if (err < 0) {
+			ERROR("Mbed TLS failed to get a heap\n");
+			panic();
+		}
+		assert(heap_size >= TF_MBEDTLS_HEAP_SIZE);
 
 		/* Initialize the mbed TLS heap */
-		mbedtls_memory_buffer_alloc_init(heap, MBEDTLS_HEAP_SIZE);
+		mbedtls_memory_buffer_alloc_init(heap_addr, heap_size);
 
 #ifdef MBEDTLS_PLATFORM_SNPRINTF_ALT
 		mbedtls_platform_set_snprintf(snprintf);