Merge pull request #7438 from valeriosetti/issue7074

Avoid parse/unparse private ECC keys in PK with USE_PSA when !ECP_C
diff --git a/include/mbedtls/pk.h b/include/mbedtls/pk.h
index ffd1b73..ec99c84 100644
--- a/include/mbedtls/pk.h
+++ b/include/mbedtls/pk.h
@@ -262,11 +262,24 @@
 typedef struct mbedtls_pk_context {
     const mbedtls_pk_info_t *MBEDTLS_PRIVATE(pk_info);    /**< Public key information         */
     void *MBEDTLS_PRIVATE(pk_ctx);                        /**< Underlying public key context  */
-    /* When MBEDTLS_PSA_CRYPTO_C is enabled then the following priv_id field is
-     * used to store the ID of the opaque key.
-     * This priv_id is guarded by MBEDTLS_PSA_CRYPTO_C and not by
-     * MBEDTLS_USE_PSA_CRYPTO because it can be used also in mbedtls_pk_sign_ext
-     * for RSA keys. */
+    /* The following field is used to store the ID of a private key in the
+     * following cases:
+     * - opaque key when MBEDTLS_PSA_CRYPTO_C is defined
+     * - normal key when MBEDTLS_PK_USE_PSA_EC_DATA is defined. In this case:
+     *    - the pk_ctx above is not not used to store the private key anymore.
+     *      Actually that field not populated at all in this case because also
+     *      the public key will be stored in raw format as explained below
+     *    - this ID is used for all private key operations (ex: sign, check
+     *      key pair, key write, etc) using PSA functions
+     *
+     * Note: this private key storing solution only affects EC keys, not the
+     *       other ones. The latters still use the pk_ctx to store their own
+     *       context.
+     *
+     * Note: this priv_id is guarded by MBEDTLS_PSA_CRYPTO_C and not by
+     *       MBEDTLS_PK_USE_PSA_EC_DATA (as the public counterpart below) because,
+     *       when working with opaque keys, it can be used also in
+     *       mbedtls_pk_sign_ext for RSA keys. */
 #if defined(MBEDTLS_PSA_CRYPTO_C)
     mbedtls_svc_key_id_t MBEDTLS_PRIVATE(priv_id);      /**< Key ID for opaque keys */
 #endif /* MBEDTLS_PSA_CRYPTO_C */
@@ -277,8 +290,7 @@
      *
      * When MBEDTLS_PK_USE_PSA_EC_DATA is enabled:
      * - the pk_ctx above is not used anymore for storing the public key
-     *   inside the ecp_keypair structure (only the private part, but also this
-     *   one is going to change in the future)
+     *   inside the ecp_keypair structure
      * - the following fields are used for all public key operations: signature
      *   verify, key pair check and key write.
      * Of course, when MBEDTLS_PK_USE_PSA_EC_DATA is not enabled, the legacy
diff --git a/library/pk.c b/library/pk.c
index 9c4aa16..8e42b8d 100644
--- a/library/pk.c
+++ b/library/pk.c
@@ -42,13 +42,6 @@
 
 #if defined(MBEDTLS_PSA_CRYPTO_C)
 #include "mbedtls/psa_util.h"
-#define PSA_PK_TO_MBEDTLS_ERR(status) psa_pk_status_to_mbedtls(status)
-#define PSA_PK_RSA_TO_MBEDTLS_ERR(status) PSA_TO_MBEDTLS_ERR_LIST(status,     \
-                                                                  psa_to_pk_rsa_errors,            \
-                                                                  psa_pk_status_to_mbedtls)
-#define PSA_PK_ECDSA_TO_MBEDTLS_ERR(status) PSA_TO_MBEDTLS_ERR_LIST(status,   \
-                                                                    psa_to_pk_ecdsa_errors,        \
-                                                                    psa_pk_status_to_mbedtls)
 #endif
 
 #include <limits.h>
@@ -85,6 +78,14 @@
         ctx->pk_info->ctx_free_func(ctx->pk_ctx);
     }
 
+#if defined(MBEDTLS_PK_USE_PSA_EC_DATA)
+    /* The ownership of the priv_id key for opaque keys is external of the PK
+     * module. It's the user responsibility to clear it after use. */
+    if ((ctx->pk_info != NULL) && (ctx->pk_info->type != MBEDTLS_PK_OPAQUE)) {
+        psa_destroy_key(ctx->priv_id);
+    }
+#endif /* MBEDTLS_PK_USE_PSA_EC_DATA */
+
     mbedtls_platform_zeroize(ctx, sizeof(mbedtls_pk_context));
 }
 
@@ -150,7 +151,7 @@
         return MBEDTLS_ERR_PK_BAD_INPUT_DATA;
     }
 
-    if ((info->ctx_alloc_func == NULL) ||
+    if ((info->ctx_alloc_func != NULL) &&
         ((ctx->pk_ctx = info->ctx_alloc_func()) == NULL)) {
         return MBEDTLS_ERR_PK_ALLOC_FAILED;
     }
@@ -911,24 +912,35 @@
 #else /* !MBEDTLS_ECP_LIGHT && !MBEDTLS_RSA_C */
 #if defined(MBEDTLS_ECP_LIGHT)
     if (mbedtls_pk_get_type(pk) == MBEDTLS_PK_ECKEY) {
-        mbedtls_ecp_keypair *ec;
-        unsigned char d[MBEDTLS_ECP_MAX_BYTES];
         size_t d_len;
         psa_ecc_family_t curve_id;
         psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
         psa_key_type_t key_type;
         size_t bits;
-        int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
         psa_status_t status;
 
         /* export the private key material in the format PSA wants */
-        ec = mbedtls_pk_ec_rw(*pk);
+#if defined(MBEDTLS_PK_USE_PSA_EC_DATA)
+        unsigned char d[MBEDTLS_PSA_MAX_EC_KEY_PAIR_LENGTH];
+        status = psa_export_key(pk->priv_id, d, sizeof(d), &d_len);
+        if (status != PSA_SUCCESS) {
+            return psa_pk_status_to_mbedtls(status);
+        }
+
+        curve_id = pk->ec_family;
+        bits = pk->ec_bits;
+#else /* MBEDTLS_PK_USE_PSA_EC_DATA */
+        unsigned char d[MBEDTLS_ECP_MAX_BYTES];
+        mbedtls_ecp_keypair *ec = mbedtls_pk_ec_rw(*pk);
+        int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
+
         d_len = PSA_BITS_TO_BYTES(ec->grp.nbits);
         if ((ret = mbedtls_ecp_write_key(ec, d, d_len)) != 0) {
             return ret;
         }
 
         curve_id = mbedtls_ecc_group_to_psa(ec->grp.id, &bits);
+#endif /* MBEDTLS_PK_USE_PSA_EC_DATA */
         key_type = PSA_KEY_TYPE_ECC_KEY_PAIR(curve_id);
 
         /* prepare the key attributes */
diff --git a/library/pk_internal.h b/library/pk_internal.h
index dbb7bc1..8d4b005 100644
--- a/library/pk_internal.h
+++ b/library/pk_internal.h
@@ -33,6 +33,17 @@
 #include "psa/crypto.h"
 #endif
 
+#if defined(MBEDTLS_PSA_CRYPTO_C)
+#include "mbedtls/psa_util.h"
+#define PSA_PK_TO_MBEDTLS_ERR(status) psa_pk_status_to_mbedtls(status)
+#define PSA_PK_RSA_TO_MBEDTLS_ERR(status) PSA_TO_MBEDTLS_ERR_LIST(status,     \
+                                                                  psa_to_pk_rsa_errors,            \
+                                                                  psa_pk_status_to_mbedtls)
+#define PSA_PK_ECDSA_TO_MBEDTLS_ERR(status) PSA_TO_MBEDTLS_ERR_LIST(status,   \
+                                                                    psa_to_pk_ecdsa_errors,        \
+                                                                    psa_pk_status_to_mbedtls)
+#endif
+
 #if defined(MBEDTLS_ECP_LIGHT)
 /**
  * Public function mbedtls_pk_ec() can be used to get direct access to the
@@ -70,24 +81,20 @@
     }
 }
 
-/* Helpers for Montgomery curves */
+static inline mbedtls_ecp_group_id mbedtls_pk_get_group_id(const mbedtls_pk_context *pk)
+{
+    mbedtls_ecp_group_id id;
+#if defined(MBEDTLS_PK_USE_PSA_EC_DATA)
+    id = mbedtls_ecc_group_of_psa(pk->ec_family, pk->ec_bits, 0);
+#else /* MBEDTLS_PK_USE_PSA_EC_DATA */
+    id = mbedtls_pk_ec_ro(*pk)->grp.id;
+#endif /* MBEDTLS_PK_USE_PSA_EC_DATA */
+    return id;
+}
+
+/* Helper for Montgomery curves */
 #if defined(MBEDTLS_ECP_DP_CURVE25519_ENABLED) || defined(MBEDTLS_ECP_DP_CURVE448_ENABLED)
 #define MBEDTLS_PK_HAVE_RFC8410_CURVES
-
-static inline int mbedtls_pk_is_rfc8410_curve(mbedtls_ecp_group_id id)
-{
-#if defined(MBEDTLS_ECP_DP_CURVE25519_ENABLED)
-    if (id == MBEDTLS_ECP_DP_CURVE25519) {
-        return 1;
-    }
-#endif
-#if defined(MBEDTLS_ECP_DP_CURVE448_ENABLED)
-    if (id == MBEDTLS_ECP_DP_CURVE448) {
-        return 1;
-    }
-#endif
-    return 0;
-}
 #endif /* MBEDTLS_ECP_DP_CURVE25519_ENABLED || MBEDTLS_ECP_DP_CURVE448_ENABLED */
 #endif /* MBEDTLS_ECP_LIGHT */
 
diff --git a/library/pk_wrap.c b/library/pk_wrap.c
index 3a3d399..92937c8 100644
--- a/library/pk_wrap.c
+++ b/library/pk_wrap.c
@@ -43,13 +43,6 @@
 
 #if defined(MBEDTLS_PSA_CRYPTO_C)
 #include "mbedtls/psa_util.h"
-#define PSA_PK_TO_MBEDTLS_ERR(status) psa_pk_status_to_mbedtls(status)
-#define PSA_PK_RSA_TO_MBEDTLS_ERR(status) PSA_TO_MBEDTLS_ERR_LIST(status,     \
-                                                                  psa_to_pk_rsa_errors,            \
-                                                                  psa_pk_status_to_mbedtls)
-#define PSA_PK_ECDSA_TO_MBEDTLS_ERR(status) PSA_TO_MBEDTLS_ERR_LIST(status,   \
-                                                                    psa_to_pk_ecdsa_errors,        \
-                                                                    psa_pk_status_to_mbedtls)
 #endif
 
 #if defined(MBEDTLS_USE_PSA_CRYPTO)
@@ -932,12 +925,9 @@
                            unsigned char *sig, size_t sig_size, size_t *sig_len,
                            int (*f_rng)(void *, unsigned char *, size_t), void *p_rng)
 {
-    mbedtls_ecp_keypair *ctx = pk->pk_ctx;
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
-    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
     mbedtls_svc_key_id_t key_id = MBEDTLS_SVC_KEY_ID_INIT;
     psa_status_t status;
-    unsigned char buf[MBEDTLS_PSA_MAX_EC_KEY_PAIR_LENGTH];
 #if defined(MBEDTLS_ECDSA_DETERMINISTIC)
     psa_algorithm_t psa_sig_md =
         PSA_ALG_DETERMINISTIC_ECDSA(mbedtls_hash_info_psa_from_md(md_alg));
@@ -945,10 +935,17 @@
     psa_algorithm_t psa_sig_md =
         PSA_ALG_ECDSA(mbedtls_hash_info_psa_from_md(md_alg));
 #endif
+#if defined(MBEDTLS_PK_USE_PSA_EC_DATA)
+    psa_ecc_family_t curve = pk->ec_family;
+#else /* MBEDTLS_PK_USE_PSA_EC_DATA */
+    mbedtls_ecp_keypair *ctx = pk->pk_ctx;
+    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
+    unsigned char buf[MBEDTLS_PSA_MAX_EC_KEY_PAIR_LENGTH];
     size_t curve_bits;
     psa_ecc_family_t curve =
         mbedtls_ecc_group_to_psa(ctx->grp.id, &curve_bits);
     size_t key_len = PSA_BITS_TO_BYTES(curve_bits);
+#endif /* MBEDTLS_PK_USE_PSA_EC_DATA */
 
     /* PSA has its own RNG */
     ((void) f_rng);
@@ -958,6 +955,12 @@
         return MBEDTLS_ERR_PK_BAD_INPUT_DATA;
     }
 
+#if defined(MBEDTLS_PK_USE_PSA_EC_DATA)
+    if (MBEDTLS_SVC_KEY_ID_GET_KEY_ID(pk->priv_id) == PSA_KEY_ID_NULL) {
+        return MBEDTLS_ERR_PK_BAD_INPUT_DATA;
+    }
+    key_id = pk->priv_id;
+#else
     if (key_len > sizeof(buf)) {
         return MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
     }
@@ -977,6 +980,7 @@
         ret = PSA_PK_TO_MBEDTLS_ERR(status);
         goto cleanup;
     }
+#endif /* MBEDTLS_PK_USE_PSA_EC_DATA */
 
     status = psa_sign_hash(key_id, psa_sig_md, hash, hash_len,
                            sig, sig_size, sig_len);
@@ -988,8 +992,11 @@
     ret = pk_ecdsa_sig_asn1_from_psa(sig, sig_len, sig_size);
 
 cleanup:
+
+#if !defined(MBEDTLS_PK_USE_PSA_EC_DATA)
     mbedtls_platform_zeroize(buf, sizeof(buf));
     status = psa_destroy_key(key_id);
+#endif /* MBEDTLS_PK_USE_PSA_EC_DATA */
     if (ret == 0 && status != PSA_SUCCESS) {
         ret = PSA_PK_TO_MBEDTLS_ERR(status);
     }
@@ -1129,25 +1136,32 @@
  */
 static int eckey_check_pair_psa(mbedtls_pk_context *pub, mbedtls_pk_context *prv)
 {
-    psa_status_t status, destruction_status;
-    psa_key_attributes_t key_attr = PSA_KEY_ATTRIBUTES_INIT;
+    psa_status_t status;
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
-    /* We are using MBEDTLS_PSA_MAX_EC_PUBKEY_LENGTH for the size of this
-     * buffer because it will be used to hold the private key at first and
-     * then its public part (but not at the same time). */
     uint8_t prv_key_buf[MBEDTLS_PSA_MAX_EC_PUBKEY_LENGTH];
     size_t prv_key_len;
-    mbedtls_svc_key_id_t key_id = MBEDTLS_SVC_KEY_ID_INIT;
 #if defined(MBEDTLS_PK_USE_PSA_EC_DATA)
-    const psa_ecc_family_t curve = prv->ec_family;
-    const size_t curve_bits = prv->ec_bits;
+    mbedtls_svc_key_id_t key_id = prv->priv_id;
+
+    status = psa_export_public_key(key_id, prv_key_buf, sizeof(prv_key_buf),
+                                   &prv_key_len);
+    ret = PSA_PK_TO_MBEDTLS_ERR(status);
+    if (ret != 0) {
+        return ret;
+    }
+
+    if (memcmp(prv_key_buf, pub->pub_raw, pub->pub_raw_len) != 0) {
+        return MBEDTLS_ERR_PK_BAD_INPUT_DATA;
+    }
 #else /* !MBEDTLS_PK_USE_PSA_EC_DATA */
+    psa_status_t destruction_status;
+    mbedtls_svc_key_id_t key_id = MBEDTLS_SVC_KEY_ID_INIT;
+    psa_key_attributes_t key_attr = PSA_KEY_ATTRIBUTES_INIT;
     uint8_t pub_key_buf[MBEDTLS_PSA_MAX_EC_PUBKEY_LENGTH];
     size_t pub_key_len;
     size_t curve_bits;
     const psa_ecc_family_t curve =
         mbedtls_ecc_group_to_psa(mbedtls_pk_ec_ro(*prv)->grp.id, &curve_bits);
-#endif /* !MBEDTLS_PK_USE_PSA_EC_DATA */
     const size_t curve_bytes = PSA_BITS_TO_BYTES(curve_bits);
 
     if (curve == 0) {
@@ -1181,11 +1195,6 @@
         return PSA_PK_TO_MBEDTLS_ERR(destruction_status);
     }
 
-#if defined(MBEDTLS_PK_USE_PSA_EC_DATA)
-    if (memcmp(prv_key_buf, pub->pub_raw, pub->pub_raw_len) != 0) {
-        return MBEDTLS_ERR_PK_BAD_INPUT_DATA;
-    }
-#else
     ret = mbedtls_ecp_point_write_binary(&mbedtls_pk_ec_rw(*pub)->grp,
                                          &mbedtls_pk_ec_rw(*pub)->Q,
                                          MBEDTLS_ECP_PF_UNCOMPRESSED,
@@ -1221,6 +1230,7 @@
 #endif
 }
 
+#if !defined(MBEDTLS_PK_USE_PSA_EC_DATA)
 static void *eckey_alloc_wrap(void)
 {
     void *ctx = mbedtls_calloc(1, sizeof(mbedtls_ecp_keypair));
@@ -1237,6 +1247,7 @@
     mbedtls_ecp_keypair_free((mbedtls_ecp_keypair *) ctx);
     mbedtls_free(ctx);
 }
+#endif /* MBEDTLS_PK_USE_PSA_EC_DATA */
 
 static void eckey_debug(mbedtls_pk_context *pk, mbedtls_pk_debug_item *items)
 {
@@ -1274,8 +1285,13 @@
     NULL,
     NULL,
     eckey_check_pair,
+#if defined(MBEDTLS_PK_USE_PSA_EC_DATA)
+    NULL,
+    NULL,
+#else /* MBEDTLS_PK_USE_PSA_EC_DATA */
     eckey_alloc_wrap,
     eckey_free_wrap,
+#endif /* MBEDTLS_PK_USE_PSA_EC_DATA */
 #if defined(MBEDTLS_ECDSA_C) && defined(MBEDTLS_ECP_RESTARTABLE)
     eckey_rs_alloc,
     eckey_rs_free,
@@ -1306,8 +1322,13 @@
     NULL,
     NULL,
     eckey_check_pair,
-    eckey_alloc_wrap,       /* Same underlying key structure */
-    eckey_free_wrap,        /* Same underlying key structure */
+#if defined(MBEDTLS_PK_USE_PSA_EC_DATA)
+    NULL,
+    NULL,
+#else /* MBEDTLS_PK_USE_PSA_EC_DATA */
+    eckey_alloc_wrap,   /* Same underlying key structure */
+    eckey_free_wrap,    /* Same underlying key structure */
+#endif /* MBEDTLS_PK_USE_PSA_EC_DATA */
 #if defined(MBEDTLS_ECDSA_C) && defined(MBEDTLS_ECP_RESTARTABLE)
     NULL,
     NULL,
@@ -1396,8 +1417,13 @@
     NULL,
     NULL,
     eckey_check_pair,   /* Compatible key structures */
+#if defined(MBEDTLS_PK_USE_PSA_EC_DATA)
+    NULL,
+    NULL,
+#else /* MBEDTLS_PK_USE_PSA_EC_DATA */
     eckey_alloc_wrap,   /* Compatible key structures */
     eckey_free_wrap,   /* Compatible key structures */
+#endif /* MBEDTLS_PK_USE_PSA_EC_DATA */
 #if defined(MBEDTLS_ECDSA_C) && defined(MBEDTLS_ECP_RESTARTABLE)
     ecdsa_rs_alloc,
     ecdsa_rs_free,
diff --git a/library/pkparse.c b/library/pkparse.c
index 9bc8801..07fce5c 100644
--- a/library/pkparse.c
+++ b/library/pkparse.c
@@ -63,6 +63,12 @@
 
 #include "mbedtls/platform.h"
 
+/* Helper for Montgomery curves */
+#if defined(MBEDTLS_ECP_LIGHT) && defined(MBEDTLS_PK_HAVE_RFC8410_CURVES)
+#define MBEDTLS_PK_IS_RFC8410_GROUP_ID(id)  \
+    ((id == MBEDTLS_ECP_DP_CURVE25519) || (id == MBEDTLS_ECP_DP_CURVE448))
+#endif /* MBEDTLS_ECP_LIGHT && MBEDTLS_PK_HAVE_RFC8410_CURVES */
+
 #if defined(MBEDTLS_FS_IO)
 /*
  * Load all data from a file into a given buffer.
@@ -508,6 +514,9 @@
 #endif
     }
 
+#if defined(MBEDTLS_PK_USE_PSA_EC_DATA)
+    ret = pk_update_psa_ecparams(pk, grp_id);
+#else
     /* grp may already be initialized; if so, make sure IDs match */
     if (mbedtls_pk_ec_ro(*pk)->grp.id != MBEDTLS_ECP_DP_NONE &&
         mbedtls_pk_ec_ro(*pk)->grp.id != grp_id) {
@@ -518,8 +527,6 @@
                                       grp_id)) != 0) {
         return ret;
     }
-#if defined(MBEDTLS_PK_USE_PSA_EC_DATA)
-    ret = pk_update_psa_ecparams(pk, grp_id);
 #endif /* MBEDTLS_PK_USE_PSA_EC_DATA */
 
     return ret;
@@ -533,20 +540,26 @@
                                 int (*f_rng)(void *, unsigned char *, size_t), void *p_rng)
 {
     int ret;
-    mbedtls_ecp_keypair *eck = (mbedtls_ecp_keypair *) pk->pk_ctx;
 #if defined(MBEDTLS_USE_PSA_CRYPTO)
-    psa_status_t status, destruction_status;
+    psa_status_t status;
+    (void) f_rng;
+    (void) p_rng;
+#if defined(MBEDTLS_PK_USE_PSA_EC_DATA)
+    (void) d;
+    (void) d_len;
+
+    status = psa_export_public_key(pk->priv_id, pk->pub_raw, sizeof(pk->pub_raw),
+                                   &pk->pub_raw_len);
+    ret = psa_pk_status_to_mbedtls(status);
+#else /* MBEDTLS_PK_USE_PSA_EC_DATA */
+    mbedtls_ecp_keypair *eck = (mbedtls_ecp_keypair *) pk->pk_ctx;
+    unsigned char key_buf[MBEDTLS_PSA_MAX_EC_PUBKEY_LENGTH];
+    size_t key_len;
+    mbedtls_svc_key_id_t key_id = MBEDTLS_SVC_KEY_ID_INIT;
     psa_key_attributes_t key_attr = PSA_KEY_ATTRIBUTES_INIT;
     size_t curve_bits;
     psa_ecc_family_t curve = mbedtls_ecc_group_to_psa(eck->grp.id, &curve_bits);
-#if !defined(MBEDTLS_PK_USE_PSA_EC_DATA)
-    unsigned char key_buf[MBEDTLS_PSA_MAX_EC_PUBKEY_LENGTH];
-    size_t key_len;
-#endif /* !MBEDTLS_PK_USE_PSA_EC_DATA */
-    mbedtls_svc_key_id_t key_id = MBEDTLS_SVC_KEY_ID_INIT;
-
-    (void) f_rng;
-    (void) p_rng;
+    psa_status_t destruction_status;
 
     psa_set_key_type(&key_attr, PSA_KEY_TYPE_ECC_KEY_PAIR(curve));
     psa_set_key_usage_flags(&key_attr, PSA_KEY_USAGE_EXPORT);
@@ -557,12 +570,7 @@
         return ret;
     }
 
-#if defined(MBEDTLS_PK_USE_PSA_EC_DATA)
-    status = psa_export_public_key(key_id, pk->pub_raw, sizeof(pk->pub_raw),
-                                   &pk->pub_raw_len);
-#else /* MBEDTLS_PK_USE_PSA_EC_DATA */
     status = psa_export_public_key(key_id, key_buf, sizeof(key_buf), &key_len);
-#endif /* MBEDTLS_PK_USE_PSA_EC_DATA */
     ret = psa_pk_status_to_mbedtls(status);
     destruction_status = psa_destroy_key(key_id);
     if (ret != 0) {
@@ -570,10 +578,10 @@
     } else if (destruction_status != PSA_SUCCESS) {
         return psa_pk_status_to_mbedtls(destruction_status);
     }
-#if !defined(MBEDTLS_PK_USE_PSA_EC_DATA)
     ret = mbedtls_ecp_point_read_binary(&eck->grp, &eck->Q, key_buf, key_len);
-#endif /* !MBEDTLS_PK_USE_PSA_EC_DATA */
+#endif /* MBEDTLS_PK_USE_PSA_EC_DATA */
 #else /* MBEDTLS_USE_PSA_CRYPTO */
+    mbedtls_ecp_keypair *eck = (mbedtls_ecp_keypair *) pk->pk_ctx;
     (void) d;
     (void) d_len;
 
@@ -591,21 +599,21 @@
                                    mbedtls_ecp_group_id grp_id,
                                    mbedtls_pk_context *pk)
 {
-    mbedtls_ecp_keypair *ecp = mbedtls_pk_ec_rw(*pk);
     int ret;
 
     if (params->tag != 0 || params->len != 0) {
         return MBEDTLS_ERR_PK_KEY_INVALID_FORMAT;
     }
 
+#if defined(MBEDTLS_PK_USE_PSA_EC_DATA)
+    ret = pk_update_psa_ecparams(pk, grp_id);
+#else
+    mbedtls_ecp_keypair *ecp = mbedtls_pk_ec_rw(*pk);
     ret = mbedtls_ecp_group_load(&(ecp->grp), grp_id);
     if (ret != 0) {
         return ret;
     }
-
-#if defined(MBEDTLS_PK_USE_PSA_EC_DATA)
-    ret = pk_update_psa_ecparams(pk, grp_id);
-#endif /* MBEDTLS_PK_USE_PSA_EC_DATA */
+#endif
     return ret;
 }
 
@@ -618,7 +626,6 @@
                                     unsigned char *key, size_t keylen, const unsigned char *end,
                                     int (*f_rng)(void *, unsigned char *, size_t), void *p_rng)
 {
-    mbedtls_ecp_keypair *eck = mbedtls_pk_ec_rw(*pk);
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
     size_t len;
 
@@ -630,23 +637,47 @@
         return MBEDTLS_ERR_PK_KEY_INVALID_FORMAT;
     }
 
+#if defined(MBEDTLS_PK_USE_PSA_EC_DATA)
+    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
+    psa_status_t status;
+
+    psa_set_key_type(&attributes, PSA_KEY_TYPE_ECC_KEY_PAIR(pk->ec_family));
+    psa_set_key_usage_flags(&attributes, PSA_KEY_USAGE_EXPORT |
+                            PSA_KEY_USAGE_DERIVE);
+    psa_set_key_algorithm(&attributes, PSA_ALG_ECDH);
+
+    status = psa_import_key(&attributes, key, len, &pk->priv_id);
+    if (status != PSA_SUCCESS) {
+        ret = psa_pk_status_to_mbedtls(status);
+        return ret;
+    }
+#else /* MBEDTLS_PK_USE_PSA_EC_DATA */
+    mbedtls_ecp_keypair *eck = mbedtls_pk_ec_rw(*pk);
+
     if ((ret = mbedtls_mpi_read_binary_le(&eck->d, key, len)) != 0) {
         mbedtls_ecp_keypair_free(eck);
         return MBEDTLS_ERROR_ADD(MBEDTLS_ERR_PK_KEY_INVALID_FORMAT, ret);
     }
+#endif /* MBEDTLS_PK_USE_PSA_EC_DATA */
 
     /* pk_parse_key_pkcs8_unencrypted_der() only supports version 1 PKCS8 keys,
      * which never contain a public key. As such, derive the public key
      * unconditionally. */
     if ((ret = pk_derive_public_key(pk, key, len, f_rng, p_rng)) != 0) {
+#if !defined(MBEDTLS_PK_USE_PSA_EC_DATA)
         mbedtls_ecp_keypair_free(eck);
+#endif /* !MBEDTLS_PK_USE_PSA_EC_DATA */
         return ret;
     }
 
+    /* When MBEDTLS_PK_USE_PSA_EC_DATA the key is checked while importing it
+     * into PSA. */
+#if !defined(MBEDTLS_PK_USE_PSA_EC_DATA)
     if ((ret = mbedtls_ecp_check_privkey(&eck->grp, &eck->d)) != 0) {
         mbedtls_ecp_keypair_free(eck);
         return ret;
     }
+#endif /* !MBEDTLS_PK_USE_PSA_EC_DATA */
 
     return 0;
 }
@@ -920,7 +951,7 @@
 #if defined(MBEDTLS_ECP_LIGHT)
     if (pk_alg == MBEDTLS_PK_ECKEY_DH || pk_alg == MBEDTLS_PK_ECKEY) {
 #if defined(MBEDTLS_PK_HAVE_RFC8410_CURVES)
-        if (mbedtls_pk_is_rfc8410_curve(ec_grp_id)) {
+        if (MBEDTLS_PK_IS_RFC8410_GROUP_ID(ec_grp_id)) {
             ret = pk_use_ecparams_rfc8410(&alg_params, ec_grp_id, pk);
         } else
 #endif
@@ -1152,6 +1183,10 @@
     unsigned char *end = p + keylen;
     unsigned char *end2;
     mbedtls_ecp_keypair *eck = mbedtls_pk_ec_rw(*pk);
+#if defined(MBEDTLS_PK_USE_PSA_EC_DATA)
+    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
+    psa_status_t status;
+#endif /* MBEDTLS_PK_USE_PSA_EC_DATA */
 
     /*
      * RFC 5915, or SEC1 Appendix C.4
@@ -1184,10 +1219,13 @@
 
     d = p;
     d_len = len;
+
+#if !defined(MBEDTLS_PK_USE_PSA_EC_DATA)
     if ((ret = mbedtls_mpi_read_binary(&eck->d, p, len)) != 0) {
         mbedtls_ecp_keypair_free(eck);
         return MBEDTLS_ERROR_ADD(MBEDTLS_ERR_PK_KEY_INVALID_FORMAT, ret);
     }
+#endif
 
     p += len;
 
@@ -1246,6 +1284,27 @@
         }
     }
 
+#if defined(MBEDTLS_PK_USE_PSA_EC_DATA)
+    psa_set_key_type(&attributes, PSA_KEY_TYPE_ECC_KEY_PAIR(pk->ec_family));
+    /* Setting largest masks for usage and key algorithms */
+    psa_set_key_usage_flags(&attributes, PSA_KEY_USAGE_SIGN_HASH |
+                            PSA_KEY_USAGE_SIGN_MESSAGE |
+                            PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_DERIVE);
+#if defined(MBEDTLS_ECDSA_DETERMINISTIC)
+    psa_set_key_algorithm(&attributes,
+                          PSA_ALG_DETERMINISTIC_ECDSA(PSA_ALG_ANY_HASH));
+#else
+    psa_set_key_algorithm(&attributes, PSA_ALG_ECDSA(PSA_ALG_ANY_HASH));
+#endif
+    psa_set_key_enrollment_algorithm(&attributes, PSA_ALG_ECDH);
+
+    status = psa_import_key(&attributes, d, d_len, &pk->priv_id);
+    if (status != PSA_SUCCESS) {
+        ret = psa_pk_status_to_mbedtls(status);
+        return ret;
+    }
+#endif /* MBEDTLS_PK_USE_PSA_EC_DATA */
+
     if (!pubkey_done) {
         if ((ret = pk_derive_public_key(pk, d, d_len, f_rng, p_rng)) != 0) {
             mbedtls_ecp_keypair_free(eck);
@@ -1253,10 +1312,12 @@
         }
     }
 
+#if !defined(MBEDTLS_PK_USE_PSA_EC_DATA)
     if ((ret = mbedtls_ecp_check_privkey(&eck->grp, &eck->d)) != 0) {
         mbedtls_ecp_keypair_free(eck);
         return ret;
     }
+#endif /* !MBEDTLS_PK_USE_PSA_EC_DATA */
 
     return 0;
 }
@@ -1357,7 +1418,7 @@
 #if defined(MBEDTLS_ECP_LIGHT)
     if (pk_alg == MBEDTLS_PK_ECKEY || pk_alg == MBEDTLS_PK_ECKEY_DH) {
 #if defined(MBEDTLS_PK_HAVE_RFC8410_CURVES)
-        if (mbedtls_pk_is_rfc8410_curve(ec_grp_id)) {
+        if (MBEDTLS_PK_IS_RFC8410_GROUP_ID(ec_grp_id)) {
             if ((ret =
                      pk_use_ecparams_rfc8410(&params, ec_grp_id, pk)) != 0 ||
                 (ret =
diff --git a/library/pkwrite.c b/library/pkwrite.c
index 3577fa1..d684815 100644
--- a/library/pkwrite.c
+++ b/library/pkwrite.c
@@ -57,6 +57,26 @@
 #endif
 #include "mbedtls/platform.h"
 
+/* Helper for Montgomery curves */
+#if defined(MBEDTLS_ECP_LIGHT) && defined(MBEDTLS_PK_HAVE_RFC8410_CURVES)
+static inline int mbedtls_pk_is_rfc8410(const mbedtls_pk_context *pk)
+{
+    mbedtls_ecp_group_id id = mbedtls_pk_get_group_id(pk);
+
+#if defined(MBEDTLS_ECP_DP_CURVE25519_ENABLED)
+    if (id == MBEDTLS_ECP_DP_CURVE25519) {
+        return 1;
+    }
+#endif
+#if defined(MBEDTLS_ECP_DP_CURVE448_ENABLED)
+    if (id == MBEDTLS_ECP_DP_CURVE448) {
+        return 1;
+    }
+#endif
+    return 0;
+}
+#endif /* MBEDTLS_ECP_LIGHT && MBEDTLS_PK_HAVE_RFC8410_CURVES */
+
 #if defined(MBEDTLS_RSA_C)
 /*
  *  RSAPublicKey ::= SEQUENCE {
@@ -165,20 +185,33 @@
  * privateKey  OCTET STRING -- always of length ceil(log2(n)/8)
  */
 static int pk_write_ec_private(unsigned char **p, unsigned char *start,
-                               mbedtls_ecp_keypair *ec)
+                               const mbedtls_pk_context *pk)
 {
+    size_t byte_length;
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
-    size_t byte_length = (ec->grp.pbits + 7) / 8;
+
+#if defined(MBEDTLS_PK_USE_PSA_EC_DATA)
+    unsigned char tmp[MBEDTLS_PSA_MAX_EC_KEY_PAIR_LENGTH];
+    psa_status_t status;
+
+    status = psa_export_key(pk->priv_id, tmp, sizeof(tmp), &byte_length);
+    if (status != PSA_SUCCESS) {
+        ret = PSA_PK_ECDSA_TO_MBEDTLS_ERR(status);
+        goto exit;
+    }
+#else /* MBEDTLS_PK_USE_PSA_EC_DATA */
     unsigned char tmp[MBEDTLS_ECP_MAX_BYTES];
+    mbedtls_ecp_keypair *ec = mbedtls_pk_ec_rw(*pk);
+    byte_length = (ec->grp.pbits + 7) / 8;
 
     ret = mbedtls_ecp_write_key(ec, tmp, byte_length);
     if (ret != 0) {
         goto exit;
     }
+#endif /* MBEDTLS_PK_USE_PSA_EC_DATA */
     ret = mbedtls_asn1_write_octet_string(p, start, tmp, byte_length);
-
 exit:
-    mbedtls_platform_zeroize(tmp, byte_length);
+    mbedtls_platform_zeroize(tmp, sizeof(tmp));
     return ret;
 }
 #endif /* MBEDTLS_ECP_LIGHT */
@@ -260,7 +293,11 @@
     pk_type = mbedtls_pk_get_type(key);
 #if defined(MBEDTLS_ECP_LIGHT)
     if (pk_type == MBEDTLS_PK_ECKEY) {
+#if defined(MBEDTLS_ECP_C)
         ec_grp_id = mbedtls_pk_ec_ro(*key)->grp.id;
+#else /* MBEDTLS_ECP_C */
+        ec_grp_id = mbedtls_ecc_group_of_psa(key->ec_family, key->ec_bits, 0);
+#endif /* MBEDTLS_ECP_C */
     }
 #endif /* MBEDTLS_ECP_LIGHT */
 #if defined(MBEDTLS_USE_PSA_CRYPTO)
@@ -352,20 +389,22 @@
  * CurvePrivateKey ::= OCTET STRING
  */
 static int pk_write_ec_rfc8410_der(unsigned char **p, unsigned char *buf,
-                                   mbedtls_ecp_keypair *ec)
+                                   const mbedtls_pk_context *pk)
 {
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
     size_t len = 0;
     size_t oid_len = 0;
     const char *oid;
+    mbedtls_ecp_group_id grp_id;
 
     /* privateKey */
-    MBEDTLS_ASN1_CHK_ADD(len, pk_write_ec_private(p, buf, ec));
+    MBEDTLS_ASN1_CHK_ADD(len, pk_write_ec_private(p, buf, pk));
     MBEDTLS_ASN1_CHK_ADD(len, mbedtls_asn1_write_len(p, buf, len));
     MBEDTLS_ASN1_CHK_ADD(len, mbedtls_asn1_write_tag(p, buf, MBEDTLS_ASN1_OCTET_STRING));
 
+    grp_id = mbedtls_pk_get_group_id(pk);
     /* privateKeyAlgorithm */
-    if ((ret = mbedtls_oid_get_oid_by_ec_grp_algid(ec->grp.id, &oid, &oid_len)) != 0) {
+    if ((ret = mbedtls_oid_get_oid_by_ec_grp_algid(grp_id, &oid, &oid_len)) != 0) {
         return ret;
     }
     MBEDTLS_ASN1_CHK_ADD(len,
@@ -388,6 +427,9 @@
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
     unsigned char *c;
     size_t len = 0;
+#if defined(MBEDTLS_ECP_LIGHT)
+    mbedtls_ecp_group_id grp_id;
+#endif
 
     if (size == 0) {
         return MBEDTLS_ERR_ASN1_BUF_TOO_SMALL;
@@ -483,12 +525,11 @@
 #endif /* MBEDTLS_RSA_C */
 #if defined(MBEDTLS_ECP_LIGHT)
     if (mbedtls_pk_get_type(key) == MBEDTLS_PK_ECKEY) {
-        mbedtls_ecp_keypair *ec = mbedtls_pk_ec_rw(*key);
         size_t pub_len = 0, par_len = 0;
 
 #if defined(MBEDTLS_PK_HAVE_RFC8410_CURVES)
-        if (mbedtls_pk_is_rfc8410_curve(ec->grp.id)) {
-            return pk_write_ec_rfc8410_der(&c, buf, ec);
+        if (mbedtls_pk_is_rfc8410(key)) {
+            return pk_write_ec_rfc8410_der(&c, buf, key);
         }
 #endif
 
@@ -522,7 +563,8 @@
         len += pub_len;
 
         /* parameters */
-        MBEDTLS_ASN1_CHK_ADD(par_len, pk_write_ec_param(&c, buf, ec->grp.id));
+        grp_id = mbedtls_pk_get_group_id(key);
+        MBEDTLS_ASN1_CHK_ADD(par_len, pk_write_ec_param(&c, buf, grp_id));
 
         MBEDTLS_ASN1_CHK_ADD(par_len, mbedtls_asn1_write_len(&c, buf, par_len));
         MBEDTLS_ASN1_CHK_ADD(par_len, mbedtls_asn1_write_tag(&c, buf,
@@ -531,7 +573,7 @@
         len += par_len;
 
         /* privateKey */
-        MBEDTLS_ASN1_CHK_ADD(len, pk_write_ec_private(&c, buf, ec));
+        MBEDTLS_ASN1_CHK_ADD(len, pk_write_ec_private(&c, buf, key));
 
         /* version */
         MBEDTLS_ASN1_CHK_ADD(len, mbedtls_asn1_write_int(&c, buf, 1));
@@ -605,7 +647,7 @@
 #if defined(MBEDTLS_ECP_LIGHT)
     if (mbedtls_pk_get_type(key) == MBEDTLS_PK_ECKEY) {
 #if defined(MBEDTLS_PK_HAVE_RFC8410_CURVES)
-        if (mbedtls_pk_is_rfc8410_curve(mbedtls_pk_ec_ro(*key)->grp.id)) {
+        if (mbedtls_pk_is_rfc8410(key)) {
             begin = PEM_BEGIN_PRIVATE_KEY_PKCS8;
             end = PEM_END_PRIVATE_KEY_PKCS8;
         } else
diff --git a/library/ssl_tls.c b/library/ssl_tls.c
index 7afb352..c928ccd 100644
--- a/library/ssl_tls.c
+++ b/library/ssl_tls.c
@@ -7367,13 +7367,12 @@
             /* and in the unlikely case the above assumption no longer holds
              * we are making sure that pk_ec() here does not return a NULL
              */
-            const mbedtls_ecp_keypair *ec = mbedtls_pk_ec_ro(*pk);
-            if (ec == NULL) {
-                MBEDTLS_SSL_DEBUG_MSG(1, ("mbedtls_pk_ec_ro() returned NULL"));
+            mbedtls_ecp_group_id grp_id = mbedtls_pk_get_group_id(pk);
+            if (grp_id == MBEDTLS_ECP_DP_NONE) {
+                MBEDTLS_SSL_DEBUG_MSG(1, ("invalid group ID"));
                 return MBEDTLS_ERR_SSL_INTERNAL_ERROR;
             }
-
-            if (mbedtls_ssl_check_curve(ssl, ec->grp.id) != 0) {
+            if (mbedtls_ssl_check_curve(ssl, grp_id) != 0) {
                 ssl->session_negotiate->verify_result |=
                     MBEDTLS_X509_BADCERT_BAD_KEY;
 
diff --git a/library/ssl_tls12_client.c b/library/ssl_tls12_client.c
index 070583b..691fa62 100644
--- a/library/ssl_tls12_client.c
+++ b/library/ssl_tls12_client.c
@@ -1986,7 +1986,6 @@
 static int ssl_get_ecdh_params_from_cert(mbedtls_ssl_context *ssl)
 {
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
-    const mbedtls_ecp_keypair *peer_key;
     mbedtls_pk_context *peer_pk;
 
 #if !defined(MBEDTLS_SSL_KEEP_PEER_CERTIFICATE)
@@ -2007,21 +2006,24 @@
         return MBEDTLS_ERR_SSL_PK_TYPE_MISMATCH;
     }
 
-    peer_key = mbedtls_pk_ec_ro(*peer_pk);
+#if defined(MBEDTLS_ECP_C)
+    const mbedtls_ecp_keypair *peer_key = mbedtls_pk_ec_ro(*peer_pk);
+#endif /* MBEDTLS_ECP_C */
 
 #if defined(MBEDTLS_USE_PSA_CRYPTO)
     uint16_t tls_id = 0;
     psa_ecc_family_t ecc_family;
+    mbedtls_ecp_group_id grp_id = mbedtls_pk_get_group_id(peer_pk);
 
-    if (mbedtls_ssl_check_curve(ssl, peer_key->grp.id) != 0) {
+    if (mbedtls_ssl_check_curve(ssl, grp_id) != 0) {
         MBEDTLS_SSL_DEBUG_MSG(1, ("bad server certificate (ECDH curve)"));
         return MBEDTLS_ERR_SSL_BAD_CERTIFICATE;
     }
 
-    tls_id = mbedtls_ssl_get_tls_id_from_ecp_group_id(peer_key->grp.id);
+    tls_id = mbedtls_ssl_get_tls_id_from_ecp_group_id(grp_id);
     if (tls_id == 0) {
         MBEDTLS_SSL_DEBUG_MSG(1, ("ECC group %u not suported",
-                                  peer_key->grp.id));
+                                  grp_id));
         return MBEDTLS_ERR_SSL_ILLEGAL_PARAMETER;
     }
 
@@ -2037,7 +2039,7 @@
     memcpy(ssl->handshake->ecdh_psa_peerkey, peer_pk->pub_raw, peer_pk->pub_raw_len);
     ssl->handshake->ecdh_psa_peerkey_len = peer_pk->pub_raw_len;
     ret = 0;
-#else
+#else /* MBEDTLS_PK_USE_PSA_EC_DATA */
     size_t olen = 0;
     ret = mbedtls_ecp_point_write_binary(&peer_key->grp, &peer_key->Q,
                                          MBEDTLS_ECP_PF_UNCOMPRESSED, &olen,
@@ -2049,8 +2051,8 @@
         return ret;
     }
     ssl->handshake->ecdh_psa_peerkey_len = olen;
-#endif /* MBEDTLS_ECP_C */
-#else
+#endif /* MBEDTLS_PK_USE_PSA_EC_DATA */
+#else /* MBEDTLS_USE_PSA_CRYPTO */
     if ((ret = mbedtls_ecdh_get_params(&ssl->handshake->ecdh_ctx, peer_key,
                                        MBEDTLS_ECDH_THEIRS)) != 0) {
         MBEDTLS_SSL_DEBUG_RET(1, ("mbedtls_ecdh_get_params"), ret);
@@ -2061,7 +2063,7 @@
         MBEDTLS_SSL_DEBUG_MSG(1, ("bad server certificate (ECDH curve)"));
         return MBEDTLS_ERR_SSL_BAD_CERTIFICATE;
     }
-#endif
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 #if !defined(MBEDTLS_SSL_KEEP_PEER_CERTIFICATE)
     /* We don't need the peer's public key anymore. Free it,
      * so that more RAM is available for upcoming expensive
diff --git a/library/ssl_tls12_server.c b/library/ssl_tls12_server.c
index a377d80..3b8710e 100644
--- a/library/ssl_tls12_server.c
+++ b/library/ssl_tls12_server.c
@@ -2597,7 +2597,7 @@
     psa_ecc_family_t ecc_family;
     size_t key_len;
     mbedtls_pk_context *pk;
-    mbedtls_ecp_keypair *key;
+    mbedtls_ecp_group_id grp_id;
 
     pk = mbedtls_ssl_own_key(ssl);
 
@@ -2605,6 +2605,10 @@
         return MBEDTLS_ERR_ECP_BAD_INPUT_DATA;
     }
 
+#if !defined(MBEDTLS_PK_USE_PSA_EC_DATA)
+    mbedtls_ecp_keypair *key = mbedtls_pk_ec_rw(*pk);
+#endif /* !MBEDTLS_PK_USE_PSA_EC_DATA */
+
     switch (mbedtls_pk_get_type(pk)) {
         case MBEDTLS_PK_OPAQUE:
             if (!mbedtls_pk_can_do(pk, MBEDTLS_PK_ECKEY)) {
@@ -2633,12 +2637,11 @@
         case MBEDTLS_PK_ECKEY:
         case MBEDTLS_PK_ECKEY_DH:
         case MBEDTLS_PK_ECDSA:
-            key = mbedtls_pk_ec_rw(*pk);
-            if (key == NULL) {
+            grp_id = mbedtls_pk_get_group_id(pk);
+            if (grp_id == MBEDTLS_ECP_DP_NONE) {
                 return MBEDTLS_ERR_ECP_BAD_INPUT_DATA;
             }
-
-            tls_id = mbedtls_ssl_get_tls_id_from_ecp_group_id(key->grp.id);
+            tls_id = mbedtls_ssl_get_tls_id_from_ecp_group_id(grp_id);
             if (tls_id == 0) {
                 /* This elliptic curve is not supported */
                 return MBEDTLS_ERR_SSL_HANDSHAKE_FAILURE;
@@ -2658,11 +2661,19 @@
                              PSA_KEY_TYPE_ECC_KEY_PAIR(ssl->handshake->ecdh_psa_type));
             psa_set_key_bits(&key_attributes, ssl->handshake->ecdh_bits);
 
+#if defined(MBEDTLS_PK_USE_PSA_EC_DATA)
+            status = psa_export_key(pk->priv_id, buf, sizeof(buf), &key_len);
+            if (status != PSA_SUCCESS) {
+                ret = PSA_TO_MBEDTLS_ERR(status);
+                goto cleanup;
+            }
+#else /* MBEDTLS_PK_USE_PSA_EC_DATA */
             key_len = PSA_BITS_TO_BYTES(key->grp.pbits);
             ret = mbedtls_ecp_write_key(key, buf, key_len);
             if (ret != 0) {
                 goto cleanup;
             }
+#endif /* MBEDTLS_PK_USE_PSA_EC_DATA */
 
             status = psa_import_key(&key_attributes, buf, key_len,
                                     &ssl->handshake->ecdh_psa_privkey);
diff --git a/library/x509_crt.c b/library/x509_crt.c
index cba30aa..b658f7c 100644
--- a/library/x509_crt.c
+++ b/library/x509_crt.c
@@ -238,7 +238,7 @@
     if (pk_alg == MBEDTLS_PK_ECDSA ||
         pk_alg == MBEDTLS_PK_ECKEY ||
         pk_alg == MBEDTLS_PK_ECKEY_DH) {
-        const mbedtls_ecp_group_id gid = mbedtls_pk_ec_ro(*pk)->grp.id;
+        const mbedtls_ecp_group_id gid = mbedtls_pk_get_group_id(pk);
 
         if (gid == MBEDTLS_ECP_DP_NONE) {
             return -1;
diff --git a/tests/scripts/analyze_outcomes.py b/tests/scripts/analyze_outcomes.py
index 293459b..0238555 100755
--- a/tests/scripts/analyze_outcomes.py
+++ b/tests/scripts/analyze_outcomes.py
@@ -247,7 +247,21 @@
                     'ECP test vectors secp256r1 rfc 5114',
                     'ECP test vectors secp384r1 rfc 5114',
                     'ECP test vectors secp521r1 rfc 5114',
-                ]
+                ],
+                'test_suite_pkparse': [
+                    # This is a known difference for Montgomery curves: in
+                    # reference component private keys are parsed using
+                    # mbedtls_mpi_read_binary_le(), while in driver version they
+                    # they are imported in PSA and there the parsing is done
+                    # through mbedtls_ecp_read_key(). Unfortunately the latter
+                    # fixes the errors which are intentionally set on the parsed
+                    # key and therefore the following test case is not failing
+                    # as expected.
+                    # This cause the following test to be guarded by ECP_C and
+                    # not being executed on the driver version.
+                    ('Key ASN1 (OneAsymmetricKey X25519, doesn\'t match masking '
+                     'requirements, from RFC8410 Appendix A but made into version 0)'),
+                ],
             }
         }
     },
diff --git a/tests/suites/test_suite_ecp.function b/tests/suites/test_suite_ecp.function
index 53da2fc..34d0288 100644
--- a/tests/suites/test_suite_ecp.function
+++ b/tests/suites/test_suite_ecp.function
@@ -1086,7 +1086,7 @@
 }
 /* END_CASE */
 
-/* BEGIN_CASE depends_on:MBEDTLS_TEST_HOOKS:MBEDTLS_ECP_MONTGOMERY_ENABLED:MBBEDTLS_ECP_C */
+/* BEGIN_CASE depends_on:MBEDTLS_TEST_HOOKS:MBEDTLS_ECP_MONTGOMERY_ENABLED:MBEDTLS_ECP_LIGHT */
 void genkey_mx_known_answer(int bits, data_t *seed, data_t *expected)
 {
     mbedtls_test_rnd_buf_info rnd_info;
diff --git a/tests/suites/test_suite_pk.function b/tests/suites/test_suite_pk.function
index 7227f92..a5b50de 100644
--- a/tests/suites/test_suite_pk.function
+++ b/tests/suites/test_suite_pk.function
@@ -29,13 +29,9 @@
 static int pk_genkey_ec(mbedtls_pk_context *pk, mbedtls_ecp_group_id grp_id)
 {
     psa_status_t status;
-    mbedtls_ecp_keypair *eck = mbedtls_pk_ec_rw(*pk);
     psa_key_attributes_t key_attr = PSA_KEY_ATTRIBUTES_INIT;
-    mbedtls_svc_key_id_t key_id = MBEDTLS_SVC_KEY_ID_INIT;
     size_t curve_bits;
     psa_ecc_family_t curve = mbedtls_ecc_group_to_psa(grp_id, &curve_bits);
-    unsigned char key_buf[MBEDTLS_PSA_MAX_EC_PUBKEY_LENGTH];
-    size_t key_len;
     int ret;
 
     if (curve == 0) {
@@ -44,25 +40,21 @@
 
     psa_set_key_type(&key_attr, PSA_KEY_TYPE_ECC_KEY_PAIR(curve));
     psa_set_key_bits(&key_attr, curve_bits);
-    psa_set_key_usage_flags(&key_attr, PSA_KEY_USAGE_EXPORT);
+    psa_set_key_usage_flags(&key_attr, PSA_KEY_USAGE_EXPORT |
+                            PSA_KEY_USAGE_SIGN_HASH |
+                            PSA_KEY_USAGE_SIGN_MESSAGE);
+#if defined(MBEDTLS_ECDSA_DETERMINISTIC)
+    psa_set_key_algorithm(&key_attr, PSA_ALG_DETERMINISTIC_ECDSA(PSA_ALG_ANY_HASH));
+#else
+    psa_set_key_algorithm(&key_attr, PSA_ALG_ECDSA(PSA_ALG_ANY_HASH));
+#endif
 
-    status = psa_generate_key(&key_attr, &key_id);
+    status = psa_generate_key(&key_attr, &pk->priv_id);
     if (status != PSA_SUCCESS) {
         return MBEDTLS_ERR_PK_FEATURE_UNAVAILABLE;
     }
 
-    status = psa_export_key(key_id, key_buf, sizeof(key_buf), &key_len);
-    if (status != PSA_SUCCESS) {
-        ret = MBEDTLS_ERR_PK_FEATURE_UNAVAILABLE;
-        goto exit;
-    }
-
-    ret = mbedtls_mpi_read_binary(&eck->d, key_buf, key_len);
-    if (ret != 0) {
-        goto exit;
-    }
-
-    status = psa_export_public_key(key_id, pk->pub_raw, sizeof(pk->pub_raw),
+    status = psa_export_public_key(pk->priv_id, pk->pub_raw, sizeof(pk->pub_raw),
                                    &pk->pub_raw_len);
     if (status != PSA_SUCCESS) {
         ret = MBEDTLS_ERR_PK_FEATURE_UNAVAILABLE;
@@ -72,15 +64,10 @@
     pk->ec_family = curve;
     pk->ec_bits = curve_bits;
 
-    status = psa_destroy_key(key_id);
-    if (status != PSA_SUCCESS) {
-        return psa_pk_status_to_mbedtls(status);
-    }
-
     return 0;
 
 exit:
-    status = psa_destroy_key(key_id);
+    status = psa_destroy_key(pk->priv_id);
     return (ret != 0) ? ret : psa_pk_status_to_mbedtls(status);
 }
 #endif /* MBEDTLS_PK_USE_PSA_EC_DATA */
@@ -114,10 +101,16 @@
         mbedtls_pk_get_type(pk) == MBEDTLS_PK_ECDSA) {
         int ret;
 
+#if defined(MBEDTLS_ECP_C)
         ret = mbedtls_ecp_group_load(&mbedtls_pk_ec_rw(*pk)->grp, parameter);
         if (ret != 0) {
             return ret;
         }
+        return mbedtls_ecp_gen_keypair(&mbedtls_pk_ec_rw(*pk)->grp,
+                                       &mbedtls_pk_ec_rw(*pk)->d,
+                                       &mbedtls_pk_ec_rw(*pk)->Q,
+                                       mbedtls_test_rnd_std_rand, NULL);
+#endif /* MBEDTLS_ECP_C */
 
 #if defined(MBEDTLS_PK_USE_PSA_EC_DATA)
         mbedtls_ecp_group grp;
@@ -136,12 +129,6 @@
 
         return 0;
 #endif /* MBEDTLS_PK_USE_PSA_EC_DATA */
-#if defined(MBEDTLS_ECP_C)
-        return mbedtls_ecp_gen_keypair(&mbedtls_pk_ec_rw(*pk)->grp,
-                                       &mbedtls_pk_ec_rw(*pk)->d,
-                                       &mbedtls_pk_ec_rw(*pk)->Q,
-                                       mbedtls_test_rnd_std_rand, NULL);
-#endif /* MBEDTLS_ECP_C */
 
     }
 #endif /* MBEDTLS_ECP_LIGHT */
diff --git a/tests/suites/test_suite_pkparse.data b/tests/suites/test_suite_pkparse.data
index 02a6ae7..978439a 100644
--- a/tests/suites/test_suite_pkparse.data
+++ b/tests/suites/test_suite_pkparse.data
@@ -1197,7 +1197,7 @@
 pk_parse_key:"30070201010400a000":MBEDTLS_ERR_PK_KEY_INVALID_FORMAT
 
 Key ASN1 (OneAsymmetricKey X25519, doesn't match masking requirements, from RFC8410 Appendix A but made into version 0)
-depends_on:MBEDTLS_ECP_LIGHT
+depends_on:MBEDTLS_ECP_C
 pk_parse_key:"302e020100300506032b656e04220420f8ffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff3f":MBEDTLS_ERR_PK_KEY_INVALID_FORMAT
 
 Key ASN1 (OneAsymmetricKey X25519, with invalid optional AlgorithIdentifier parameters)
diff --git a/tests/suites/test_suite_pkparse.function b/tests/suites/test_suite_pkparse.function
index a49b6d3..6fa78c1 100644
--- a/tests/suites/test_suite_pkparse.function
+++ b/tests/suites/test_suite_pkparse.function
@@ -117,10 +117,13 @@
     TEST_ASSERT(res == result);
 
     if (res == 0) {
-        const mbedtls_ecp_keypair *eckey;
         TEST_ASSERT(mbedtls_pk_can_do(&ctx, MBEDTLS_PK_ECKEY));
-        eckey = mbedtls_pk_ec_ro(ctx);
+#if defined(MBEDTLS_ECP_C)
+        const mbedtls_ecp_keypair *eckey = mbedtls_pk_ec_ro(ctx);
         TEST_ASSERT(mbedtls_ecp_check_privkey(&eckey->grp, &eckey->d) == 0);
+#else
+        /* PSA keys are already checked on import so nothing to do here. */
+#endif
     }
 
 exit: