diff --git a/library/pk.c b/library/pk.c
index ae1966b..433b65f 100644
--- a/library/pk.c
+++ b/library/pk.c
@@ -60,6 +60,9 @@
 {
     ctx->pk_info = NULL;
     ctx->pk_ctx = NULL;
+#if defined(MBEDTLS_PSA_CRYPTO_C)
+    ctx->opaque_id = MBEDTLS_SVC_KEY_ID_INIT;
+#endif /* MBEDTLS_PSA_CRYPTO_C */
 }
 
 /*
@@ -71,7 +74,7 @@
         return;
     }
 
-    if (ctx->pk_info != NULL) {
+    if ((ctx->pk_info != NULL) && (ctx->pk_info->ctx_free_func != NULL)) {
         ctx->pk_info->ctx_free_func(ctx->pk_ctx);
     }
 
@@ -140,7 +143,8 @@
         return MBEDTLS_ERR_PK_BAD_INPUT_DATA;
     }
 
-    if ((ctx->pk_ctx = 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;
     }
 
@@ -158,7 +162,6 @@
 {
     const mbedtls_pk_info_t *info = NULL;
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
-    mbedtls_svc_key_id_t *pk_ctx;
     psa_key_type_t type;
 
     if (ctx == NULL || ctx->pk_info != NULL) {
@@ -179,14 +182,8 @@
         return MBEDTLS_ERR_PK_FEATURE_UNAVAILABLE;
     }
 
-    if ((ctx->pk_ctx = info->ctx_alloc_func()) == NULL) {
-        return MBEDTLS_ERR_PK_ALLOC_FAILED;
-    }
-
     ctx->pk_info = info;
-
-    pk_ctx = (mbedtls_svc_key_id_t *) ctx->pk_ctx;
-    *pk_ctx = key;
+    ctx->opaque_id = key;
 
     return 0;
 }
diff --git a/library/pk_wrap.c b/library/pk_wrap.c
index 6c9f97b..c3dea53 100644
--- a/library/pk_wrap.c
+++ b/library/pk_wrap.c
@@ -1503,22 +1503,6 @@
 #endif /* MBEDTLS_PK_RSA_ALT_SUPPORT */
 
 #if defined(MBEDTLS_USE_PSA_CRYPTO)
-
-static void *pk_opaque_alloc_wrap(void)
-{
-    void *ctx = mbedtls_calloc(1, sizeof(mbedtls_svc_key_id_t));
-
-    /* no _init() function to call, as calloc() already zeroized */
-
-    return ctx;
-}
-
-static void pk_opaque_free_wrap(void *ctx)
-{
-    mbedtls_platform_zeroize(ctx, sizeof(mbedtls_svc_key_id_t));
-    mbedtls_free(ctx);
-}
-
 static size_t pk_opaque_get_bitlen(mbedtls_pk_context *pk)
 {
     const mbedtls_svc_key_id_t *key = pk->pk_ctx;
@@ -1635,8 +1619,8 @@
     NULL, /* decrypt - not relevant */
     NULL, /* encrypt - not relevant */
     NULL, /* check_pair - could be done later or left NULL */
-    pk_opaque_alloc_wrap,
-    pk_opaque_free_wrap,
+    NULL, /* alloc - no need to allocate new data dynamically */
+    NULL, /* free - as for the alloc, there is no data to free */
 #if defined(MBEDTLS_ECDSA_C) && defined(MBEDTLS_ECP_RESTARTABLE)
     NULL, /* restart alloc - not relevant */
     NULL, /* restart free - not relevant */
@@ -1687,8 +1671,8 @@
 #endif /* PSA_WANT_KEY_TYPE_RSA_PUBLIC_KEY */
     NULL, /* encrypt - will be done later */
     NULL, /* check_pair - could be done later or left NULL */
-    pk_opaque_alloc_wrap,
-    pk_opaque_free_wrap,
+    NULL, /* alloc - no need to allocate new data dynamically */
+    NULL, /* free - as for the alloc, there is no data to free */
 #if defined(MBEDTLS_ECDSA_C) && defined(MBEDTLS_ECP_RESTARTABLE)
     NULL, /* restart alloc - not relevant */
     NULL, /* restart free - not relevant */