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_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);