pk: manage allocate and free space when working with PSA private key

Allocation does not need to perform any action since the priv_id field
is already present on the pk_context.
Free should destroy the key. Of course this is true only if the key
is not opaque (because in that case it's the user responsibility
to do so).

Signed-off-by: Valerio Setti <valerio.setti@nordicsemi.no>
diff --git a/library/pk.c b/library/pk.c
index 5ed485b..77012e1 100644
--- a/library/pk.c
+++ b/library/pk.c
@@ -78,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));
 }
 
@@ -143,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;
     }
diff --git a/library/pk_wrap.c b/library/pk_wrap.c
index 45ded6e..7f5e751 100644
--- a/library/pk_wrap.c
+++ b/library/pk_wrap.c
@@ -1214,6 +1214,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));
@@ -1230,6 +1231,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)
 {
@@ -1267,8 +1269,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,
@@ -1299,8 +1306,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,
@@ -1389,8 +1401,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,