pk: minor code fixes/enhancements

Signed-off-by: Valerio Setti <valerio.setti@nordicsemi.no>
diff --git a/library/pk.c b/library/pk.c
index cccadb1..8e42b8d 100644
--- a/library/pk.c
+++ b/library/pk.c
@@ -912,7 +912,6 @@
 #else /* !MBEDTLS_ECP_LIGHT && !MBEDTLS_RSA_C */
 #if defined(MBEDTLS_ECP_LIGHT)
     if (mbedtls_pk_get_type(pk) == MBEDTLS_PK_ECKEY) {
-        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;
@@ -922,6 +921,7 @@
 
         /* export the private key material in the format PSA wants */
 #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);
@@ -930,6 +930,7 @@
         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;
 
diff --git a/library/pkwrite.c b/library/pkwrite.c
index c8a96d9..8df6394 100644
--- a/library/pkwrite.c
+++ b/library/pkwrite.c
@@ -395,22 +395,21 @@
     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, 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 defined(MBEDTLS_PK_USE_PSA_EC_DATA)
-    mbedtls_ecp_group_id grp_id = mbedtls_ecc_group_of_psa(pk->ec_family,
-                                                           pk->ec_bits, 0);
     if ((ret = mbedtls_oid_get_oid_by_ec_grp_algid(grp_id, &oid, &oid_len)) != 0) {
         return ret;
     }
 #else /* MBEDTLS_PK_USE_PSA_EC_DATA */
-    mbedtls_ecp_keypair *ec = mbedtls_pk_ec_rw(*pk);
-    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;
     }
 #endif /* MBEDTLS_PK_USE_PSA_EC_DATA */