Adapt to the new key allocation mechanism

diff --git a/library/cipher.c b/library/cipher.c
index 1cc0beb..e479b9c 100644
--- a/library/cipher.c
+++ b/library/cipher.c
@@ -302,38 +302,39 @@
         if( cipher_psa->slot_state != MBEDTLS_CIPHER_PSA_KEY_UNSET )
             return( MBEDTLS_ERR_CIPHER_BAD_INPUT_DATA );
 
-        /* Find a fresh key slot to use. */
-        status = mbedtls_psa_get_free_key_slot( &cipher_psa->slot );
+        key_type = mbedtls_psa_translate_cipher_type(
+            ctx->cipher_info->type );
+        if( key_type == 0 )
+            return( MBEDTLS_ERR_CIPHER_FEATURE_UNAVAILABLE );
+
+        /* Allocate a key slot to use. */
+        status = psa_allocate_key( key_type, key_bitlen, &cipher_psa->slot );
         if( status != PSA_SUCCESS )
             return( MBEDTLS_ERR_CIPHER_HW_ACCEL_FAILED );
-         /* Indicate that we own the key slot and need to
-          * destroy it in mbedtls_cipher_free(). */
-        cipher_psa->slot_state = MBEDTLS_CIPHER_PSA_KEY_OWNED;
-
-        /* From that point on, the responsibility for destroying the
-         * key slot is on mbedtls_cipher_free(). This includes the case
-         * where the policy setup or key import below fail, as
-         * mbedtls_cipher_free() needs to be called in any case. */
 
         /* Setup policy for the new key slot. */
         psa_key_policy_init( &key_policy );
 
         /* Mbed TLS' cipher layer doesn't enforce the mode of operation
-         * (encrypt vs. decrypt): it is possible to setup a key for encryption
-         * and use it for AEAD decryption. Until tests relying on this
-         * are changed, allow any usage in PSA. */
+        * (encrypt vs. decrypt): it is possible to setup a key for encryption
+        * and use it for AEAD decryption. Until tests relying on this
+        * are changed, allow any usage in PSA. */
         /* key_usage = mbedtls_psa_translate_cipher_operation( operation ); */
         key_usage = PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT;
         psa_key_policy_set_usage( &key_policy, key_usage, cipher_psa->alg );
         status = psa_set_key_policy( cipher_psa->slot, &key_policy );
         if( status != PSA_SUCCESS )
             return( MBEDTLS_ERR_CIPHER_HW_ACCEL_FAILED );
+        /* Indicate that we own the key slot and need to
+         * destroy it in mbedtls_cipher_free(). */
+        cipher_psa->slot_state = MBEDTLS_CIPHER_PSA_KEY_OWNED;
+
+        /* From that point on, the responsibility for destroying the
+        * key slot is on mbedtls_cipher_free(). This includes the case
+        * where the policy setup or key import below fail, as
+        * mbedtls_cipher_free() needs to be called in any case. */
 
         /* Populate new key slot. */
-        key_type = mbedtls_psa_translate_cipher_type(
-            ctx->cipher_info->type );
-        if( key_type == 0 )
-            return( MBEDTLS_ERR_CIPHER_FEATURE_UNAVAILABLE );
         status = psa_import_key( cipher_psa->slot,
                                  key_type, key, key_bytelen );
         if( status != PSA_SUCCESS )
diff --git a/library/pk.c b/library/pk.c
index 989ed09..8f649f9 100644
--- a/library/pk.c
+++ b/library/pk.c
@@ -147,10 +147,10 @@
 /*
  * Initialise a PSA-wrapping context
  */
-int mbedtls_pk_setup_opaque( mbedtls_pk_context *ctx, const psa_key_slot_t key )
+int mbedtls_pk_setup_opaque( mbedtls_pk_context *ctx, const psa_key_handle_t key )
 {
     const mbedtls_pk_info_t * const info = &mbedtls_pk_opaque_info;
-    psa_key_slot_t *pk_ctx;
+    psa_key_handle_t *pk_ctx;
     psa_key_type_t type;
 
     if( ctx == NULL || ctx->pk_info != NULL )
@@ -168,7 +168,7 @@
 
     ctx->pk_info = info;
 
-    pk_ctx = (psa_key_slot_t *) ctx->pk_ctx;
+    pk_ctx = (psa_key_handle_t *) ctx->pk_ctx;
     *pk_ctx = key;
 
     return( 0 );
@@ -547,13 +547,13 @@
  * Currently only works for EC private keys.
  */
 int mbedtls_pk_wrap_as_opaque( mbedtls_pk_context *pk,
-                               psa_key_slot_t *slot,
+                               psa_key_handle_t *slot,
                                psa_algorithm_t hash_alg )
 {
 #if !defined(MBEDTLS_ECP_C)
     return( MBEDTLS_ERR_PK_TYPE_MISMATCH );
 #else
-    psa_key_slot_t key;
+    psa_key_handle_t key;
     const mbedtls_ecp_keypair *ec;
     unsigned char d[MBEDTLS_ECP_MAX_BYTES];
     size_t d_len;
@@ -572,9 +572,10 @@
         return( ret );
 
     curve_id = mbedtls_ecp_curve_info_from_grp_id( ec->grp.id )->tls_id;
+    key_type = PSA_KEY_TYPE_ECC_KEYPAIR(curve_id);
 
-    /* find a free key slot */
-    if( PSA_SUCCESS != mbedtls_psa_get_free_key_slot( &key ) )
+    /* allocate a key slot */
+    if( PSA_SUCCESS != psa_allocate_key( key_type, d_len * 8, &key ) )
         return( MBEDTLS_ERR_PK_HW_ACCEL_FAILED );
 
     /* set policy */
@@ -585,7 +586,6 @@
         return( MBEDTLS_ERR_PK_HW_ACCEL_FAILED );
 
     /* import private key in slot */
-    key_type = PSA_KEY_TYPE_ECC_KEYPAIR(curve_id);
     if( PSA_SUCCESS != psa_import_key( key, key_type, d, d_len ) )
         return( MBEDTLS_ERR_PK_HW_ACCEL_FAILED );
 
diff --git a/library/pk_wrap.c b/library/pk_wrap.c
index 6aacba8..9023457 100644
--- a/library/pk_wrap.c
+++ b/library/pk_wrap.c
@@ -546,7 +546,7 @@
                        const unsigned char *sig, size_t sig_len )
 {
     int ret;
-    psa_key_slot_t key_slot;
+    psa_key_handle_t key_slot;
     psa_key_policy_t policy;
     psa_key_type_t psa_type;
     mbedtls_pk_context key;
@@ -571,15 +571,15 @@
     if( key_len <= 0 )
         return( MBEDTLS_ERR_PK_BAD_INPUT_DATA );
 
-    if( ( ret = mbedtls_psa_get_free_key_slot( &key_slot ) ) != PSA_SUCCESS )
-        return( mbedtls_psa_err_translate_pk( ret ) );
-
     psa_md = mbedtls_psa_translate_md( md_alg );
     if( psa_md == 0 )
         return( MBEDTLS_ERR_PK_BAD_INPUT_DATA );
     psa_sig_md = PSA_ALG_ECDSA( psa_md );
     psa_type = PSA_KEY_TYPE_ECC_PUBLIC_KEY( curve );
 
+    if( ( ret = psa_allocate_key( psa_type, key_len * 8, &key_slot ) ) != PSA_SUCCESS )
+          return( mbedtls_psa_err_translate_pk( ret ) );
+
     psa_key_policy_init( &policy );
     psa_key_policy_set_usage( &policy, PSA_KEY_USAGE_VERIFY, psa_sig_md );
     if( ( ret = psa_set_key_policy( key_slot, &policy ) ) != PSA_SUCCESS )
@@ -879,7 +879,7 @@
 
 static void *pk_opaque_alloc_wrap( void )
 {
-    void *ctx = mbedtls_calloc( 1, sizeof( psa_key_slot_t ) );
+    void *ctx = mbedtls_calloc( 1, sizeof( psa_key_handle_t ) );
 
     /* no _init() function to call, an calloc() already zeroized */
 
@@ -888,13 +888,13 @@
 
 static void pk_opaque_free_wrap( void *ctx )
 {
-    mbedtls_platform_zeroize( ctx, sizeof( psa_key_slot_t ) );
+    mbedtls_platform_zeroize( ctx, sizeof( psa_key_handle_t ) );
     mbedtls_free( ctx );
 }
 
 static size_t pk_opaque_get_bitlen( const void *ctx )
 {
-    const psa_key_slot_t *key = (const psa_key_slot_t *) ctx;
+    const psa_key_handle_t *key = (const psa_key_handle_t *) ctx;
     size_t bits;
 
     if( PSA_SUCCESS != psa_get_key_information( *key, NULL, &bits ) )
@@ -999,7 +999,7 @@
                    unsigned char *sig, size_t *sig_len,
                    int (*f_rng)(void *, unsigned char *, size_t), void *p_rng )
 {
-    const psa_key_slot_t *key = (const psa_key_slot_t *) ctx;
+    const psa_key_handle_t *key = (const psa_key_handle_t *) ctx;
     psa_algorithm_t alg = PSA_ALG_ECDSA( mbedtls_psa_translate_md( md_alg ) );
     size_t bits, buf_len;
     psa_status_t status;
diff --git a/library/pkwrite.c b/library/pkwrite.c
index d34714b..11a2a61 100644
--- a/library/pkwrite.c
+++ b/library/pkwrite.c
@@ -168,7 +168,7 @@
     if( mbedtls_pk_get_type( key ) == MBEDTLS_PK_OPAQUE )
     {
         size_t buffer_size;
-        psa_key_slot_t* key_slot = (psa_key_slot_t*) key->pk_ctx;
+        psa_key_handle_t* key_slot = (psa_key_handle_t*) key->pk_ctx;
 
         if ( *p < start )
             return( MBEDTLS_ERR_PK_BAD_INPUT_DATA );
diff --git a/library/ssl_tls.c b/library/ssl_tls.c
index d14434a..8fe9314 100644
--- a/library/ssl_tls.c
+++ b/library/ssl_tls.c
@@ -798,7 +798,7 @@
             psa_status_t status;
             psa_algorithm_t alg;
             psa_crypto_generator_t generator = PSA_CRYPTO_GENERATOR_INIT;
-            psa_key_slot_t psk;
+            psa_key_handle_t psk;
 
             MBEDTLS_SSL_DEBUG_MSG( 2, ( "perform PSA-based PSK-to-MS expansion" ) );
 
@@ -7617,7 +7617,7 @@
 
 #if defined(MBEDTLS_USE_PSA_CRYPTO)
 int mbedtls_ssl_conf_psk_opaque( mbedtls_ssl_config *conf,
-                                 psa_key_slot_t psk_slot,
+                                 psa_key_handle_t psk_slot,
                                  const unsigned char *psk_identity,
                                  size_t psk_identity_len )
 {
@@ -7640,7 +7640,7 @@
 }
 
 int mbedtls_ssl_set_hs_psk_opaque( mbedtls_ssl_context *ssl,
-                                   psa_key_slot_t psk_slot )
+                                   psa_key_handle_t psk_slot )
 {
     if( psk_slot == 0 || ssl->handshake == NULL )
         return( MBEDTLS_ERR_SSL_BAD_INPUT_DATA );