diff --git a/library/cipher.c b/library/cipher.c
index 1f74fcc..69079aa 100644
--- a/library/cipher.c
+++ b/library/cipher.c
@@ -297,8 +297,7 @@
 
         psa_status_t status;
         psa_key_type_t key_type;
-        psa_key_usage_t key_usage;
-        psa_key_policy_t key_policy;
+        psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
 
         /* PSA Crypto API only accepts byte-aligned keys. */
         if( key_bitlen % 8 != 0 )
@@ -312,40 +311,33 @@
             ctx->cipher_info->type );
         if( key_type == 0 )
             return( MBEDTLS_ERR_CIPHER_FEATURE_UNAVAILABLE );
-
-        /* Allocate a key slot to use. */
-        status = psa_allocate_key( &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. */
-        key_policy = psa_key_policy_init();
+        psa_set_key_type( &attributes, key_type );
 
         /* 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. */
-        /* 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 );
+        psa_set_key_usage_flags( &attributes,
+                                 /* mbedtls_psa_translate_cipher_operation( operation ); */
+                                 PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT );
+        psa_set_key_algorithm( &attributes, cipher_psa->alg );
 
-        /* Populate new key slot. */
-        status = psa_import_key_to_handle( cipher_psa->slot,
-                                 key_type, key, key_bytelen );
-        if( status != PSA_SUCCESS )
-            return( MBEDTLS_ERR_CIPHER_HW_ACCEL_FAILED );
+        status = psa_import_key( &attributes, key, key_bytelen,
+                                 &cipher_psa->slot );
+        switch( status )
+        {
+            case PSA_SUCCESS:
+                break;
+            case PSA_ERROR_INSUFFICIENT_MEMORY:
+                return( MBEDTLS_ERR_CIPHER_ALLOC_FAILED );
+            case PSA_ERROR_NOT_SUPPORTED:
+                return( MBEDTLS_ERR_CIPHER_FEATURE_UNAVAILABLE );
+            default:
+                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;
 
         ctx->key_bitlen = key_bitlen;
         ctx->operation = operation;