tls/x509: minor enhancement for using the new private key format

Signed-off-by: Valerio Setti <valerio.setti@nordicsemi.no>
diff --git a/library/ssl_tls.c b/library/ssl_tls.c
index fe666e8..2babb04 100644
--- a/library/ssl_tls.c
+++ b/library/ssl_tls.c
@@ -7388,13 +7388,18 @@
             /* and in the unlikely case the above assumption no longer holds
              * we are making sure that pk_ec() here does not return a NULL
              */
+            mbedtls_ecp_group_id grp_id;
+#if defined(MBEDTLS_PK_USE_PSA_EC_DATA)
+            grp_id = mbedtls_ecc_group_of_psa(pk->ec_family, pk->ec_bits, 0);
+#else /* MBEDTLS_PK_USE_PSA_EC_DATA */
             const mbedtls_ecp_keypair *ec = mbedtls_pk_ec_ro(*pk);
             if (ec == NULL) {
                 MBEDTLS_SSL_DEBUG_MSG(1, ("mbedtls_pk_ec_ro() returned NULL"));
                 return MBEDTLS_ERR_SSL_INTERNAL_ERROR;
             }
-
-            if (mbedtls_ssl_check_curve(ssl, ec->grp.id) != 0) {
+            grp_id = ec->grp.id;
+#endif /* MBEDTLS_PK_USE_PSA_EC_DATA */
+            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 aa3e306..a442b37 100644
--- a/library/ssl_tls12_server.c
+++ b/library/ssl_tls12_server.c
@@ -2600,7 +2600,10 @@
     psa_ecc_family_t ecc_family;
     size_t key_len;
     mbedtls_pk_context *pk;
+    mbedtls_ecp_group_id grp_id;
+#if !defined(MBEDTLS_PK_USE_PSA_EC_DATA)
     mbedtls_ecp_keypair *key;
+#endif /* !MBEDTLS_PK_USE_PSA_EC_DATA */
 
     pk = mbedtls_ssl_own_key(ssl);
 
@@ -2636,12 +2639,16 @@
         case MBEDTLS_PK_ECKEY:
         case MBEDTLS_PK_ECKEY_DH:
         case MBEDTLS_PK_ECDSA:
+#if defined(MBEDTLS_PK_USE_PSA_EC_DATA)
+            grp_id = mbedtls_ecc_group_of_psa(pk->ec_family, pk->ec_bits, 0);
+#else /* MBEDTLS_PK_USE_PSA_EC_DATA */
             key = mbedtls_pk_ec_rw(*pk);
             if (key == NULL) {
                 return MBEDTLS_ERR_ECP_BAD_INPUT_DATA;
             }
-
-            tls_id = mbedtls_ssl_get_tls_id_from_ecp_group_id(key->grp.id);
+            grp_id = key->grp.id;
+#endif /* MBEDTLS_PK_USE_PSA_EC_DATA */
+            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;
@@ -2661,11 +2668,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;