diff --git a/tf-psa-crypto/core/psa_crypto.c b/tf-psa-crypto/core/psa_crypto.c
index 5a11b2b57d..df4fc274e8 100644
--- a/tf-psa-crypto/core/psa_crypto.c
+++ b/tf-psa-crypto/core/psa_crypto.c
@@ -8184,7 +8184,7 @@ psa_status_t psa_generate_key_iop_complete(
 {
 #if defined(MBEDTLS_ECP_RESTARTABLE)
     psa_status_t status;
-    uint8_t key_data[MBEDTLS_ECP_MAX_BYTES] = { 0 };
+    uint8_t key_data[MBEDTLS_ECP_MAX_MPI] = { 0 };
     size_t key_len = 0;
 
     if (operation->id == 0 || operation->error_occurred) {
@@ -8210,7 +8210,7 @@ exit:
         psa_generate_key_iop_abort_internal(operation);
     }
 
-    mbedtls_platform_zeroize(key_data, MBEDTLS_ECP_MAX_BYTES);
+    mbedtls_platform_zeroize(key_data, sizeof(key_data));
     return status;
 #else
     (void) operation;
diff --git a/tf-psa-crypto/drivers/builtin/include/mbedtls/ecp.h b/tf-psa-crypto/drivers/builtin/include/mbedtls/ecp.h
index b3406142fe..533482c0db 100644
--- a/tf-psa-crypto/drivers/builtin/include/mbedtls/ecp.h
+++ b/tf-psa-crypto/drivers/builtin/include/mbedtls/ecp.h
@@ -345,6 +345,8 @@ mbedtls_ecp_group;
 
 #define MBEDTLS_ECP_MAX_BYTES    ((MBEDTLS_ECP_MAX_BITS + 7) / 8)
 #define MBEDTLS_ECP_MAX_PT_LEN   (2 * MBEDTLS_ECP_MAX_BYTES + 1)
+#define MBEDTLS_ECP_MAX_MPI      ((MBEDTLS_ECP_MAX_BYTES + sizeof(mbedtls_mpi_uint)) & \
+                                  ~(sizeof(mbedtls_mpi_uint)-1))
 
 #if defined(MBEDTLS_ECP_RESTARTABLE)