diff --git a/drivers/auth/mbedtls/mbedtls_common.c b/drivers/auth/mbedtls/mbedtls_common.c
index 053bf1a..1d2df5a 100644
--- a/drivers/auth/mbedtls/mbedtls_common.c
+++ b/drivers/auth/mbedtls/mbedtls_common.c
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2015, ARM Limited and Contributors. All rights reserved.
+ * Copyright (c) 2015-2017, ARM Limited and Contributors. All rights reserved.
  *
  * Redistribution and use in source and binary forms, with or without
  * modification, are permitted provided that the following conditions are met:
@@ -36,9 +36,9 @@
 /*
  * mbed TLS heap
  */
-#if (MBEDTLS_KEY_ALG_ID == MBEDTLS_ECDSA)
+#if (TBBR_KEY_ALG_ID == TBBR_ECDSA)
 #define MBEDTLS_HEAP_SIZE		(14*1024)
-#elif (MBEDTLS_KEY_ALG_ID == MBEDTLS_RSA)
+#elif (TBBR_KEY_ALG_ID == TBBR_RSA)
 #define MBEDTLS_HEAP_SIZE		(8*1024)
 #endif
 static unsigned char heap[MBEDTLS_HEAP_SIZE];