diff --git a/tf-psa-crypto/drivers/builtin/src/oid.c b/tf-psa-crypto/drivers/builtin/src/oid.c
index d5490d2..71c7365 100644
--- a/tf-psa-crypto/drivers/builtin/src/oid.c
+++ b/tf-psa-crypto/drivers/builtin/src/oid.c
@@ -531,7 +531,7 @@
                         mbedtls_pk_type_t,
                         pk_alg)
 
-#if defined(MBEDTLS_PK_HAVE_ECC_KEYS)
+#if defined(PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY)
 /*
  * For elliptic curves that use namedCurve inside ECParams (RFC 5480)
  */
@@ -662,7 +662,7 @@
                         oid_ecp_grp_algid,
                         mbedtls_ecp_group_id,
                         grp_id)
-#endif /* MBEDTLS_PK_HAVE_ECC_KEYS */
+#endif /* PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY */
 
 #if defined(MBEDTLS_CIPHER_C)
 /*
diff --git a/tf-psa-crypto/drivers/builtin/src/pk.c b/tf-psa-crypto/drivers/builtin/src/pk.c
index 3fe51ea..b97f438 100644
--- a/tf-psa-crypto/drivers/builtin/src/pk.c
+++ b/tf-psa-crypto/drivers/builtin/src/pk.c
@@ -20,7 +20,7 @@
 #include "mbedtls/rsa.h"
 #include "rsa_internal.h"
 #endif
-#if defined(MBEDTLS_PK_HAVE_ECC_KEYS)
+#if defined(PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY)
 #include "mbedtls/ecp.h"
 #endif
 #if defined(MBEDTLS_ECDSA_C)
@@ -118,12 +118,12 @@
         case MBEDTLS_PK_RSA:
             return &mbedtls_rsa_info;
 #endif /* MBEDTLS_RSA_C */
-#if defined(MBEDTLS_PK_HAVE_ECC_KEYS)
+#if defined(PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY)
         case MBEDTLS_PK_ECKEY:
             return &mbedtls_eckey_info;
         case MBEDTLS_PK_ECKEY_DH:
             return &mbedtls_eckeydh_info;
-#endif /* MBEDTLS_PK_HAVE_ECC_KEYS */
+#endif /* PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY */
 #if defined(MBEDTLS_PK_CAN_ECDSA_SOME)
         case MBEDTLS_PK_ECDSA:
             return &mbedtls_ecdsa_info;
@@ -174,11 +174,11 @@
     type = psa_get_key_type(&attributes);
     psa_reset_key_attributes(&attributes);
 
-#if defined(MBEDTLS_PK_HAVE_ECC_KEYS)
+#if defined(PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY)
     if (PSA_KEY_TYPE_IS_ECC_KEY_PAIR(type)) {
         info = &mbedtls_ecdsa_opaque_info;
     } else
-#endif /* MBEDTLS_PK_HAVE_ECC_KEYS */
+#endif /* PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY */
     if (type == PSA_KEY_TYPE_RSA_KEY_PAIR) {
         info = &mbedtls_rsa_opaque_info;
     } else {
@@ -457,7 +457,7 @@
         }
 #endif /* MBEDTLS_RSA_C */
 
-#if defined(MBEDTLS_PK_HAVE_ECC_KEYS)
+#if defined(PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY)
         case MBEDTLS_PK_ECKEY:
         case MBEDTLS_PK_ECKEY_DH:
         case MBEDTLS_PK_ECDSA:
@@ -512,7 +512,7 @@
             psa_set_key_algorithm(attributes, alg);
             break;
         }
-#endif /* MBEDTLS_PK_HAVE_ECC_KEYS */
+#endif /* PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY */
 
 #if defined(MBEDTLS_PK_RSA_ALT_SUPPORT)
         case MBEDTLS_PK_RSA_ALT:
@@ -671,7 +671,7 @@
         }
 #endif /* MBEDTLS_RSA_C */
 
-#if defined(MBEDTLS_PK_HAVE_ECC_KEYS)
+#if defined(PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY)
         case MBEDTLS_PK_ECKEY:
         case MBEDTLS_PK_ECKEY_DH:
         case MBEDTLS_PK_ECDSA:
@@ -724,7 +724,7 @@
             return ret;
 #endif /* MBEDTLS_PK_USE_PSA_EC_DATA */
         }
-#endif /* MBEDTLS_PK_HAVE_ECC_KEYS */
+#endif /* PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY */
 
 #if defined(MBEDTLS_USE_PSA_CRYPTO)
         case MBEDTLS_PK_OPAQUE:
@@ -743,7 +743,7 @@
     psa_key_type_t psa_type = psa_get_key_type(attributes);
 
 #if defined(MBEDTLS_RSA_C) ||                                           \
-    (defined(MBEDTLS_PK_HAVE_ECC_KEYS) && !defined(MBEDTLS_PK_USE_PSA_EC_DATA)) || \
+    (defined(PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY) && !defined(MBEDTLS_PK_USE_PSA_EC_DATA)) || \
     defined(MBEDTLS_USE_PSA_CRYPTO)
     unsigned char key_buffer[PSA_EXPORT_PUBLIC_KEY_MAX_SIZE];
 #endif
@@ -769,7 +769,7 @@
         }
 #endif /*MBEDTLS_RSA_C */
 
-#if defined(MBEDTLS_PK_HAVE_ECC_KEYS)
+#if defined(PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY)
         case MBEDTLS_PK_ECKEY:
         case MBEDTLS_PK_ECKEY_DH:
         case MBEDTLS_PK_ECDSA:
@@ -803,7 +803,7 @@
 #endif /* MBEDTLS_PK_USE_PSA_EC_DATA */
             break;
         }
-#endif /* MBEDTLS_PK_HAVE_ECC_KEYS */
+#endif /* PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY */
 
 #if defined(MBEDTLS_USE_PSA_CRYPTO)
         case MBEDTLS_PK_OPAQUE:
@@ -934,7 +934,7 @@
         }
     } else
 #endif /* MBEDTLS_RSA_C */
-#if defined(MBEDTLS_PK_HAVE_ECC_KEYS)
+#if defined(PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY)
     if (PSA_KEY_TYPE_IS_ECC_KEY_PAIR(key_type) ||
         PSA_KEY_TYPE_IS_ECC_PUBLIC_KEY(key_type)) {
         mbedtls_ecp_group_id grp_id;
@@ -965,7 +965,7 @@
             goto exit;
         }
     } else
-#endif /* MBEDTLS_PK_HAVE_ECC_KEYS */
+#endif /* PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY */
     {
         (void) key_bits;
         return MBEDTLS_ERR_PK_BAD_INPUT_DATA;
diff --git a/tf-psa-crypto/drivers/builtin/src/pk_ecc.c b/tf-psa-crypto/drivers/builtin/src/pk_ecc.c
index 86218ff..707988d 100644
--- a/tf-psa-crypto/drivers/builtin/src/pk_ecc.c
+++ b/tf-psa-crypto/drivers/builtin/src/pk_ecc.c
@@ -12,7 +12,7 @@
 #include "mbedtls/ecp.h"
 #include "pk_internal.h"
 
-#if defined(MBEDTLS_PK_C) && defined(MBEDTLS_PK_HAVE_ECC_KEYS)
+#if defined(MBEDTLS_PK_C) && defined(PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY)
 
 int mbedtls_pk_ecc_set_group(mbedtls_pk_context *pk, mbedtls_ecp_group_id grp_id)
 {
@@ -252,4 +252,4 @@
 #endif /* MBEDTLS_PK_USE_PSA_EC_DATA */
 }
 
-#endif /* MBEDTLS_PK_C && MBEDTLS_PK_HAVE_ECC_KEYS */
+#endif /* MBEDTLS_PK_C && PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY */
diff --git a/tf-psa-crypto/drivers/builtin/src/pk_internal.h b/tf-psa-crypto/drivers/builtin/src/pk_internal.h
index b728615..ab16e90 100644
--- a/tf-psa-crypto/drivers/builtin/src/pk_internal.h
+++ b/tf-psa-crypto/drivers/builtin/src/pk_internal.h
@@ -13,7 +13,7 @@
 
 #include "mbedtls/pk.h"
 
-#if defined(MBEDTLS_PK_HAVE_ECC_KEYS)
+#if defined(PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY)
 #include "mbedtls/ecp.h"
 #endif
 
@@ -44,7 +44,7 @@
 #define PEM_BEGIN_ENCRYPTED_PRIVATE_KEY_PKCS8 "-----BEGIN ENCRYPTED PRIVATE KEY-----"
 #define PEM_END_ENCRYPTED_PRIVATE_KEY_PKCS8   "-----END ENCRYPTED PRIVATE KEY-----"
 
-#if defined(MBEDTLS_PK_HAVE_ECC_KEYS) && !defined(MBEDTLS_PK_USE_PSA_EC_DATA)
+#if defined(PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY) && !defined(MBEDTLS_PK_USE_PSA_EC_DATA)
 /**
  * Public function mbedtls_pk_ec() can be used to get direct access to the
  * wrapped ecp_keypair structure pointed to the pk_ctx. However this is not
@@ -80,9 +80,9 @@
             return NULL;
     }
 }
-#endif /* MBEDTLS_PK_HAVE_ECC_KEYS && !MBEDTLS_PK_USE_PSA_EC_DATA */
+#endif /* PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY && !MBEDTLS_PK_USE_PSA_EC_DATA */
 
-#if defined(MBEDTLS_PK_HAVE_ECC_KEYS)
+#if defined(PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY)
 static inline mbedtls_ecp_group_id mbedtls_pk_get_ec_group_id(const mbedtls_pk_context *pk)
 {
     mbedtls_ecp_group_id id;
@@ -183,7 +183,7 @@
 int mbedtls_pk_ecc_set_pubkey_from_prv(mbedtls_pk_context *pk,
                                        const unsigned char *prv, size_t prv_len,
                                        int (*f_rng)(void *, unsigned char *, size_t), void *p_rng);
-#endif /* MBEDTLS_PK_HAVE_ECC_KEYS */
+#endif /* PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY */
 
 /* Helper for (deterministic) ECDSA */
 #if defined(MBEDTLS_ECDSA_DETERMINISTIC)
diff --git a/tf-psa-crypto/drivers/builtin/src/pk_wrap.c b/tf-psa-crypto/drivers/builtin/src/pk_wrap.c
index 8710889..a9153a2 100644
--- a/tf-psa-crypto/drivers/builtin/src/pk_wrap.c
+++ b/tf-psa-crypto/drivers/builtin/src/pk_wrap.c
@@ -505,7 +505,7 @@
 };
 #endif /* MBEDTLS_RSA_C */
 
-#if defined(MBEDTLS_PK_HAVE_ECC_KEYS)
+#if defined(PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY)
 /*
  * Generic EC key
  */
@@ -1263,7 +1263,7 @@
     .debug_func = eckey_debug,        /* Compatible key structures */
 };
 #endif /* MBEDTLS_PK_CAN_ECDSA_SOME */
-#endif /* MBEDTLS_PK_HAVE_ECC_KEYS */
+#endif /* PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY */
 
 #if defined(MBEDTLS_PK_RSA_ALT_SUPPORT)
 /*
@@ -1414,7 +1414,7 @@
     return bits;
 }
 
-#if defined(MBEDTLS_PK_HAVE_ECC_KEYS)
+#if defined(PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY)
 static int ecdsa_opaque_can_do(mbedtls_pk_type_t type)
 {
     return type == MBEDTLS_PK_ECKEY ||
@@ -1449,7 +1449,7 @@
     .ctx_free_func = NULL,
     .debug_func = NULL,
 };
-#endif /* MBEDTLS_PK_HAVE_ECC_KEYS */
+#endif /* PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY */
 
 static int rsa_opaque_can_do(mbedtls_pk_type_t type)
 {
diff --git a/tf-psa-crypto/drivers/builtin/src/pk_wrap.h b/tf-psa-crypto/drivers/builtin/src/pk_wrap.h
index be096da..5f5f6b0 100644
--- a/tf-psa-crypto/drivers/builtin/src/pk_wrap.h
+++ b/tf-psa-crypto/drivers/builtin/src/pk_wrap.h
@@ -108,7 +108,7 @@
 extern const mbedtls_pk_info_t mbedtls_rsa_info;
 #endif
 
-#if defined(MBEDTLS_PK_HAVE_ECC_KEYS)
+#if defined(PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY)
 extern const mbedtls_pk_info_t mbedtls_eckey_info;
 extern const mbedtls_pk_info_t mbedtls_eckeydh_info;
 #endif
diff --git a/tf-psa-crypto/drivers/builtin/src/pkparse.c b/tf-psa-crypto/drivers/builtin/src/pkparse.c
index 4f6ee13..3419ad9 100644
--- a/tf-psa-crypto/drivers/builtin/src/pkparse.c
+++ b/tf-psa-crypto/drivers/builtin/src/pkparse.c
@@ -42,7 +42,7 @@
 #include "mbedtls/pkcs12.h"
 #endif
 
-#if defined(MBEDTLS_PK_HAVE_ECC_KEYS)
+#if defined(PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY)
 
 /***********************************************************************
  *
@@ -460,7 +460,7 @@
 }
 #endif /* MBEDTLS_PK_HAVE_RFC8410_CURVES */
 
-#endif /* MBEDTLS_PK_HAVE_ECC_KEYS */
+#endif /* PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY */
 
 /* Get a PK algorithm identifier
  *
@@ -483,7 +483,7 @@
     }
 
     ret = mbedtls_oid_get_pk_alg(&alg_oid, pk_alg);
-#if defined(MBEDTLS_PK_HAVE_ECC_KEYS)
+#if defined(PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY)
     if (ret == MBEDTLS_ERR_OID_NOT_FOUND) {
         ret = mbedtls_oid_get_ec_grp_algid(&alg_oid, ec_grp_id);
         if (ret == 0) {
@@ -567,7 +567,7 @@
         }
     } else
 #endif /* MBEDTLS_RSA_C */
-#if defined(MBEDTLS_PK_HAVE_ECC_KEYS)
+#if defined(PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY)
     if (pk_alg == MBEDTLS_PK_ECKEY_DH || pk_alg == MBEDTLS_PK_ECKEY) {
 #if defined(MBEDTLS_PK_HAVE_RFC8410_CURVES)
         if (MBEDTLS_PK_IS_RFC8410_GROUP_ID(ec_grp_id)) {
@@ -582,7 +582,7 @@
             *p += end - *p;
         }
     } else
-#endif /* MBEDTLS_PK_HAVE_ECC_KEYS */
+#endif /* PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY */
     ret = MBEDTLS_ERR_PK_UNKNOWN_PK_ALG;
 
     if (ret == 0 && *p != end) {
@@ -597,7 +597,7 @@
     return ret;
 }
 
-#if defined(MBEDTLS_PK_HAVE_ECC_KEYS)
+#if defined(PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY)
 /*
  * Parse a SEC1 encoded private EC key
  */
@@ -718,7 +718,7 @@
 
     return 0;
 }
-#endif /* MBEDTLS_PK_HAVE_ECC_KEYS */
+#endif /* PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY */
 
 /***********************************************************************
  *
@@ -753,7 +753,7 @@
     mbedtls_ecp_group_id ec_grp_id = MBEDTLS_ECP_DP_NONE;
     const mbedtls_pk_info_t *pk_info;
 
-#if !defined(MBEDTLS_PK_HAVE_ECC_KEYS)
+#if !defined(PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY)
     (void) f_rng;
     (void) p_rng;
 #endif
@@ -818,7 +818,7 @@
         }
     } else
 #endif /* MBEDTLS_RSA_C */
-#if defined(MBEDTLS_PK_HAVE_ECC_KEYS)
+#if defined(PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY)
     if (pk_alg == MBEDTLS_PK_ECKEY || pk_alg == MBEDTLS_PK_ECKEY_DH) {
 #if defined(MBEDTLS_PK_HAVE_RFC8410_CURVES)
         if (MBEDTLS_PK_IS_RFC8410_GROUP_ID(ec_grp_id)) {
@@ -840,7 +840,7 @@
             }
         }
     } else
-#endif /* MBEDTLS_PK_HAVE_ECC_KEYS */
+#endif /* PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY */
     return MBEDTLS_ERR_PK_UNKNOWN_PK_ALG;
 
     end = p + len;
@@ -1018,7 +1018,7 @@
     }
 #endif /* MBEDTLS_RSA_C */
 
-#if defined(MBEDTLS_PK_HAVE_ECC_KEYS)
+#if defined(PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY)
     /* Avoid calling mbedtls_pem_read_buffer() on non-null-terminated string */
     if (key[keylen - 1] != '\0') {
         ret = MBEDTLS_ERR_PEM_NO_HEADER_FOOTER_PRESENT;
@@ -1047,7 +1047,7 @@
     } else if (ret != MBEDTLS_ERR_PEM_NO_HEADER_FOOTER_PRESENT) {
         return ret;
     }
-#endif /* MBEDTLS_PK_HAVE_ECC_KEYS */
+#endif /* PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY */
 
     /* Avoid calling mbedtls_pem_read_buffer() on non-null-terminated string */
     if (key[keylen - 1] != '\0') {
@@ -1151,7 +1151,7 @@
     mbedtls_pk_init(pk);
 #endif /* MBEDTLS_RSA_C */
 
-#if defined(MBEDTLS_PK_HAVE_ECC_KEYS)
+#if defined(PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY)
     pk_info = mbedtls_pk_info_from_type(MBEDTLS_PK_ECKEY);
     if (mbedtls_pk_setup(pk, pk_info) == 0 &&
         pk_parse_key_sec1_der(pk,
@@ -1159,13 +1159,13 @@
         return 0;
     }
     mbedtls_pk_free(pk);
-#endif /* MBEDTLS_PK_HAVE_ECC_KEYS */
+#endif /* PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY */
 
-    /* If MBEDTLS_RSA_C is defined but MBEDTLS_PK_HAVE_ECC_KEYS isn't,
+    /* If MBEDTLS_RSA_C is defined but PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY isn't,
      * it is ok to leave the PK context initialized but not
      * freed: It is the caller's responsibility to call pk_init()
      * before calling this function, and to call pk_free()
-     * when it fails. If MBEDTLS_PK_HAVE_ECC_KEYS is defined but MBEDTLS_RSA_C
+     * when it fails. If PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY is defined but MBEDTLS_RSA_C
      * isn't, this leads to mbedtls_pk_free() being called
      * twice, once here and once by the caller, but this is
      * also ok and in line with the mbedtls_pk_free() calls
diff --git a/tf-psa-crypto/drivers/builtin/src/pkwrite.c b/tf-psa-crypto/drivers/builtin/src/pkwrite.c
index 5e009c5..8c01b44 100644
--- a/tf-psa-crypto/drivers/builtin/src/pkwrite.c
+++ b/tf-psa-crypto/drivers/builtin/src/pkwrite.c
@@ -23,10 +23,10 @@
 #include "mbedtls/ecp.h"
 #include "mbedtls/platform_util.h"
 #endif
-#if defined(MBEDTLS_PK_HAVE_ECC_KEYS)
+#if defined(PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY)
 #include "pk_internal.h"
 #endif
-#if defined(MBEDTLS_RSA_C) || defined(MBEDTLS_PK_HAVE_ECC_KEYS)
+#if defined(MBEDTLS_RSA_C) || defined(PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY)
 #include "pkwrite.h"
 #endif
 #if defined(MBEDTLS_PEM_WRITE_C)
@@ -85,7 +85,7 @@
 /******************************************************************************
  * Internal functions for EC keys.
  ******************************************************************************/
-#if defined(MBEDTLS_PK_HAVE_ECC_KEYS)
+#if defined(PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY)
 #if defined(MBEDTLS_PK_USE_PSA_EC_DATA)
 static int pk_write_ec_pubkey(unsigned char **p, unsigned char *start,
                               const mbedtls_pk_context *pk)
@@ -343,7 +343,7 @@
 
     return (int) len;
 }
-#endif /* MBEDTLS_PK_HAVE_ECC_KEYS */
+#endif /* PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY */
 
 /******************************************************************************
  * Internal functions for Opaque keys.
@@ -419,7 +419,7 @@
         MBEDTLS_ASN1_CHK_ADD(len, mbedtls_rsa_write_pubkey(mbedtls_pk_rsa(*key), start, p));
     } else
 #endif
-#if defined(MBEDTLS_PK_HAVE_ECC_KEYS)
+#if defined(PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY)
     if (mbedtls_pk_get_type(key) == MBEDTLS_PK_ECKEY) {
         MBEDTLS_ASN1_CHK_ADD(len, pk_write_ec_pubkey(p, start, key));
     } else
@@ -468,7 +468,7 @@
 
     pk_type = pk_get_type_ext(key);
 
-#if defined(MBEDTLS_PK_HAVE_ECC_KEYS)
+#if defined(PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY)
     if (pk_get_type_ext(key) == MBEDTLS_PK_ECKEY) {
         mbedtls_ecp_group_id ec_grp_id = mbedtls_pk_get_ec_group_id(key);
         if (MBEDTLS_PK_IS_RFC8410_GROUP_ID(ec_grp_id)) {
@@ -481,7 +481,7 @@
             MBEDTLS_ASN1_CHK_ADD(par_len, pk_write_ec_param(&c, buf, ec_grp_id));
         }
     }
-#endif /* MBEDTLS_PK_HAVE_ECC_KEYS */
+#endif /* PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY */
 
     /* At this point oid_len is not null only for EC Montgomery keys. */
     if (oid_len == 0) {
@@ -516,7 +516,7 @@
         return pk_write_rsa_der(&c, buf, key);
     } else
 #endif /* MBEDTLS_RSA_C */
-#if defined(MBEDTLS_PK_HAVE_ECC_KEYS)
+#if defined(PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY)
     if (pk_get_type_ext(key) == MBEDTLS_PK_ECKEY) {
 #if defined(MBEDTLS_PK_HAVE_RFC8410_CURVES)
         if (mbedtls_pk_is_rfc8410(key)) {
@@ -525,7 +525,7 @@
 #endif /* MBEDTLS_PK_HAVE_RFC8410_CURVES */
         return pk_write_ec_der(&c, buf, key);
     } else
-#endif /* MBEDTLS_PK_HAVE_ECC_KEYS */
+#endif /* PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY */
     return MBEDTLS_ERR_PK_FEATURE_UNAVAILABLE;
 }
 
@@ -589,7 +589,7 @@
         end = PEM_END_PRIVATE_KEY_RSA "\n";
     } else
 #endif
-#if defined(MBEDTLS_PK_HAVE_ECC_KEYS)
+#if defined(PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY)
     if (pk_get_type_ext(key) == MBEDTLS_PK_ECKEY) {
         if (mbedtls_pk_is_rfc8410(key)) {
             begin = PEM_BEGIN_PRIVATE_KEY_PKCS8 "\n";
@@ -599,7 +599,7 @@
             end = PEM_END_PRIVATE_KEY_EC "\n";
         }
     } else
-#endif /* MBEDTLS_PK_HAVE_ECC_KEYS */
+#endif /* PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY */
     {
         ret = MBEDTLS_ERR_PK_FEATURE_UNAVAILABLE;
         goto cleanup;
diff --git a/tf-psa-crypto/drivers/builtin/src/pkwrite.h b/tf-psa-crypto/drivers/builtin/src/pkwrite.h
index 01dc3d2..f7fb18a 100644
--- a/tf-psa-crypto/drivers/builtin/src/pkwrite.h
+++ b/tf-psa-crypto/drivers/builtin/src/pkwrite.h
@@ -65,7 +65,7 @@
 
 #endif /* MBEDTLS_RSA_C */
 
-#if defined(MBEDTLS_PK_HAVE_ECC_KEYS)
+#if defined(PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY)
 
 /* Find the maximum number of bytes necessary to store an EC point. When USE_PSA
  * is defined this means looking for the maximum between PSA and built-in
@@ -103,12 +103,12 @@
  */
 #define MBEDTLS_PK_ECP_PRV_DER_MAX_BYTES    (29 + 3 * MBEDTLS_PK_MAX_ECC_BYTES)
 
-#else /* MBEDTLS_PK_HAVE_ECC_KEYS */
+#else /* PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY */
 
 #define MBEDTLS_PK_ECP_PUB_DER_MAX_BYTES   0
 #define MBEDTLS_PK_ECP_PRV_DER_MAX_BYTES   0
 
-#endif /* MBEDTLS_PK_HAVE_ECC_KEYS */
+#endif /* PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY */
 
 /* Define the maximum available public key DER length based on the supported
  * key types (EC and/or RSA). */