pk: fix library code for using the new opaque key solution

Signed-off-by: Valerio Setti <valerio.setti@nordicsemi.no>
diff --git a/library/pk.c b/library/pk.c
index 433b65f..bd6bf98 100644
--- a/library/pk.c
+++ b/library/pk.c
@@ -312,12 +312,11 @@
         return (key_usage & usage) == usage;
     }
 
-    const mbedtls_svc_key_id_t *key = (const mbedtls_svc_key_id_t *) ctx->pk_ctx;
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
     psa_algorithm_t key_alg, key_alg2;
     psa_status_t status;
 
-    status = psa_get_key_attributes(*key, &attributes);
+    status = psa_get_key_attributes(ctx->opaque_id, &attributes);
     if (status != PSA_SUCCESS) {
         return 0;
     }
@@ -698,10 +697,9 @@
     }
 
     if (mbedtls_pk_get_type(ctx) == MBEDTLS_PK_OPAQUE) {
-        const mbedtls_svc_key_id_t *key = (const mbedtls_svc_key_id_t *) ctx->pk_ctx;
         psa_status_t status;
 
-        status = psa_sign_hash(*key, PSA_ALG_RSA_PSS(psa_md_alg),
+        status = psa_sign_hash(ctx->opaque_id, PSA_ALG_RSA_PSS(psa_md_alg),
                                hash, hash_len,
                                sig, sig_size, sig_len);
         return PSA_PK_RSA_TO_MBEDTLS_ERR(status);
diff --git a/library/pk_wrap.c b/library/pk_wrap.c
index c3dea53..d9366c1 100644
--- a/library/pk_wrap.c
+++ b/library/pk_wrap.c
@@ -1505,11 +1505,10 @@
 #if defined(MBEDTLS_USE_PSA_CRYPTO)
 static size_t pk_opaque_get_bitlen(mbedtls_pk_context *pk)
 {
-    const mbedtls_svc_key_id_t *key = pk->pk_ctx;
     size_t bits;
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
 
-    if (PSA_SUCCESS != psa_get_key_attributes(*key, &attributes)) {
+    if (PSA_SUCCESS != psa_get_key_attributes(pk->opaque_id, &attributes)) {
         return 0;
     }
 
@@ -1547,7 +1546,6 @@
     ((void) p_rng);
     return MBEDTLS_ERR_PK_FEATURE_UNAVAILABLE;
 #else /* !MBEDTLS_PK_CAN_ECDSA_SIGN && !MBEDTLS_RSA_C */
-    const mbedtls_svc_key_id_t *key = pk->pk_ctx;
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
     psa_algorithm_t alg;
     psa_key_type_t type;
@@ -1557,7 +1555,7 @@
     (void) f_rng;
     (void) p_rng;
 
-    status = psa_get_key_attributes(*key, &attributes);
+    status = psa_get_key_attributes(pk->opaque_id, &attributes);
     if (status != PSA_SUCCESS) {
         return PSA_PK_TO_MBEDTLS_ERR(status);
     }
@@ -1578,7 +1576,7 @@
     return MBEDTLS_ERR_PK_FEATURE_UNAVAILABLE;
 
     /* make the signature */
-    status = psa_sign_hash(*key, alg, hash, hash_len,
+    status = psa_sign_hash(pk->opaque_id, alg, hash, hash_len,
                            sig, sig_size, sig_len);
     if (status != PSA_SUCCESS) {
 #if defined(MBEDTLS_PK_CAN_ECDSA_SIGN)
@@ -1634,14 +1632,13 @@
                                  unsigned char *output, size_t *olen, size_t osize,
                                  int (*f_rng)(void *, unsigned char *, size_t), void *p_rng)
 {
-    const mbedtls_svc_key_id_t *key = pk->pk_ctx;
     psa_status_t status;
 
     /* PSA has its own RNG */
     (void) f_rng;
     (void) p_rng;
 
-    status = psa_asymmetric_decrypt(*key, PSA_ALG_RSA_PKCS1V15_CRYPT,
+    status = psa_asymmetric_decrypt(pk->opaque_id, PSA_ALG_RSA_PKCS1V15_CRYPT,
                                     input, ilen,
                                     NULL, 0,
                                     output, osize, olen);
diff --git a/library/pkwrite.c b/library/pkwrite.c
index b83a13e..e62ad54 100644
--- a/library/pkwrite.c
+++ b/library/pkwrite.c
@@ -188,14 +188,13 @@
 #if defined(MBEDTLS_USE_PSA_CRYPTO)
     if (mbedtls_pk_get_type(key) == MBEDTLS_PK_OPAQUE) {
         size_t buffer_size;
-        mbedtls_svc_key_id_t *key_id = (mbedtls_svc_key_id_t *) key->pk_ctx;
 
         if (*p < start) {
             return MBEDTLS_ERR_PK_BAD_INPUT_DATA;
         }
 
         buffer_size = (size_t) (*p - start);
-        if (psa_export_public_key(*key_id, start, buffer_size, &len)
+        if (psa_export_public_key(key->opaque_id, start, buffer_size, &len)
             != PSA_SUCCESS) {
             return MBEDTLS_ERR_PK_BAD_INPUT_DATA;
         } else {
@@ -254,9 +253,11 @@
     if (pk_type == MBEDTLS_PK_OPAQUE) {
         psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
         psa_key_type_t key_type;
-        mbedtls_svc_key_id_t key_id;
-        key_id = *((mbedtls_svc_key_id_t *) key->pk_ctx);
-        if (PSA_SUCCESS != psa_get_key_attributes(key_id, &attributes)) {
+        psa_ecc_family_t curve;
+        size_t bits;
+
+        if (PSA_SUCCESS != psa_get_key_attributes(key->opaque_id,
+                                                  &attributes)) {
             return MBEDTLS_ERR_PLATFORM_HW_ACCEL_FAILED;
         }
         key_type = psa_get_key_type(&attributes);
diff --git a/library/ssl_tls12_server.c b/library/ssl_tls12_server.c
index 42f5fe9..3025725 100644
--- a/library/ssl_tls12_server.c
+++ b/library/ssl_tls12_server.c
@@ -2614,8 +2614,7 @@
                 return MBEDTLS_ERR_SSL_PK_TYPE_MISMATCH;
             }
 
-            ssl->handshake->ecdh_psa_privkey =
-                *((mbedtls_svc_key_id_t *) pk->pk_ctx);
+            ssl->handshake->ecdh_psa_privkey = pk->opaque_id;
 
             /* Key should not be destroyed in the TLS library */
             ssl->handshake->ecdh_psa_privkey_is_external = 1;