diff --git a/library/pk.c b/library/pk.c
index b2f6812..331ed6c 100644
--- a/library/pk.c
+++ b/library/pk.c
@@ -146,6 +146,7 @@
 int mbedtls_pk_setup_psa( mbedtls_pk_context *ctx, const psa_key_slot_t key )
 {
     const mbedtls_pk_info_t * const info = &mbedtls_pk_opaque_psa_info;
+    psa_key_slot_t *pk_ctx;
 
     if( ctx == NULL || ctx->pk_info != NULL )
         return( MBEDTLS_ERR_PK_BAD_INPUT_DATA );
@@ -153,11 +154,11 @@
     if( ( ctx->pk_ctx = info->ctx_alloc_func() ) == NULL )
         return( MBEDTLS_ERR_PK_ALLOC_FAILED );
 
-    /* coming soon: remember key */
-    (void) key;
-
     ctx->pk_info = info;
 
+    pk_ctx = (psa_key_slot_t *) ctx->pk_ctx;
+    *pk_ctx = key;
+
     return( 0 );
 }
 #endif /* MBEDTLS_USE_PSA_CRYPTO */
diff --git a/library/pk_wrap.c b/library/pk_wrap.c
index 4885c49..0e12d05 100644
--- a/library/pk_wrap.c
+++ b/library/pk_wrap.c
@@ -718,6 +718,21 @@
 
 #if defined(MBEDTLS_USE_PSA_CRYPTO)
 
+static void *pk_psa_alloc_wrap( void )
+{
+    void *ctx = mbedtls_calloc( 1, sizeof( psa_key_slot_t ) );
+
+    /* no _init() function to call, an calloc() already zeroized */
+
+    return( ctx );
+}
+
+static void pk_psa_free_wrap( void *ctx )
+{
+    mbedtls_platform_zeroize( ctx, sizeof( psa_key_slot_t ) );
+    mbedtls_free( ctx );
+}
+
 const mbedtls_pk_info_t mbedtls_pk_opaque_psa_info = {
     MBEDTLS_PK_OPAQUE_PSA,
     "Opaque (PSA)",
@@ -732,8 +747,8 @@
     NULL, /* decrypt - will be done later */
     NULL, /* encrypt - will be done later */
     NULL, /* check_pair - could be done later or left NULL */
-    NULL, /* coming soon: alloc */
-    NULL, /* coming soon: free */
+    pk_psa_alloc_wrap,
+    pk_psa_free_wrap,
 #if defined(MBEDTLS_ECDSA_C) && defined(MBEDTLS_ECP_RESTARTABLE)
     NULL, /* restart alloc - not relevant */
     NULL, /* restart free - not relevant */