Convert cipher and pk to PSA attribute-based key creation

This fixes the build under MBEDTLS_USE_PSA_CRYPTO.
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;
diff --git a/library/pk.c b/library/pk.c
index bcf7e0a..e93ccfd 100644
--- a/library/pk.c
+++ b/library/pk.c
@@ -158,14 +158,17 @@
 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_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
     psa_key_handle_t *pk_ctx;
     psa_key_type_t type;
 
     if( ctx == NULL || ctx->pk_info != NULL )
         return( MBEDTLS_ERR_PK_BAD_INPUT_DATA );
 
-    if( PSA_SUCCESS != psa_get_key_information( key, &type, NULL ) )
+    if( PSA_SUCCESS != psa_get_key_attributes( key, &attributes ) )
         return( MBEDTLS_ERR_PK_BAD_INPUT_DATA );
+    type = psa_get_key_type( &attributes );
+    psa_reset_key_attributes( &attributes );
 
     /* Current implementation of can_do() relies on this. */
     if( ! PSA_KEY_TYPE_IS_ECC_KEY_PAIR( type ) )
@@ -589,19 +592,18 @@
  * Currently only works for EC private keys.
  */
 int mbedtls_pk_wrap_as_opaque( mbedtls_pk_context *pk,
-                               psa_key_handle_t *slot,
+                               psa_key_handle_t *handle,
                                psa_algorithm_t hash_alg )
 {
 #if !defined(MBEDTLS_ECP_C)
     return( MBEDTLS_ERR_PK_TYPE_MISMATCH );
 #else
-    psa_key_handle_t key;
     const mbedtls_ecp_keypair *ec;
     unsigned char d[MBEDTLS_ECP_MAX_BYTES];
     size_t d_len;
     psa_ecc_curve_t curve_id;
+    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
     psa_key_type_t key_type;
-    psa_key_policy_t policy;
     int ret;
 
     /* export the private key material in the format PSA wants */
@@ -617,29 +619,20 @@
     key_type = PSA_KEY_TYPE_ECC_KEY_PAIR(
                                  mbedtls_psa_parse_tls_ecc_group ( curve_id ) );
 
-    /* allocate a key slot */
-    if( PSA_SUCCESS != psa_allocate_key( &key ) )
-        return( MBEDTLS_ERR_PK_HW_ACCEL_FAILED );
+    /* prepare the key attributes */
+    psa_set_key_type( &attributes, key_type );
+    psa_set_key_usage_flags( &attributes, PSA_KEY_USAGE_SIGN );
+    psa_set_key_algorithm( &attributes, PSA_ALG_ECDSA(hash_alg) );
 
-    /* set policy */
-    policy = psa_key_policy_init();
-    psa_key_policy_set_usage( &policy, PSA_KEY_USAGE_SIGN,
-                                       PSA_ALG_ECDSA(hash_alg) );
-    if( PSA_SUCCESS != psa_set_key_policy( key, &policy ) )
+    /* import private key into PSA */
+    if( PSA_SUCCESS != psa_import_key( &attributes, d, d_len, handle ) )
         return( MBEDTLS_ERR_PK_HW_ACCEL_FAILED );
 
-    /* import private key in slot */
-    if( PSA_SUCCESS != psa_import_key_to_handle( key, key_type, d, d_len ) )
-        return( MBEDTLS_ERR_PK_HW_ACCEL_FAILED );
-
-    /* remember slot number to be destroyed later by caller */
-    *slot = key;
-
     /* make PK context wrap the key slot */
     mbedtls_pk_free( pk );
     mbedtls_pk_init( pk );
 
-    return( mbedtls_pk_setup_opaque( pk, key ) );
+    return( mbedtls_pk_setup_opaque( pk, *handle ) );
 #endif /* MBEDTLS_ECP_C */
 }
 #endif /* MBEDTLS_USE_PSA_CRYPTO */
diff --git a/library/pk_wrap.c b/library/pk_wrap.c
index 0c74825..5a699c0 100644
--- a/library/pk_wrap.c
+++ b/library/pk_wrap.c
@@ -546,9 +546,9 @@
                        const unsigned char *sig, size_t sig_len )
 {
     int ret;
-    psa_key_handle_t key_slot;
-    psa_key_policy_t policy;
-    psa_key_type_t psa_type;
+    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
+    psa_key_handle_t key_handle = 0;
+    psa_status_t status;
     mbedtls_pk_context key;
     int key_len;
     /* see ECP_PUB_DER_MAX_BYTES in pkwrite.c */
@@ -576,23 +576,17 @@
     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( &key_slot ) ) != PSA_SUCCESS )
-          return( mbedtls_psa_err_translate_pk( ret ) );
+    psa_set_key_type( &attributes, PSA_KEY_TYPE_ECC_PUBLIC_KEY( curve ) );
+    psa_set_key_usage_flags( &attributes, PSA_KEY_USAGE_VERIFY );
+    psa_set_key_algorithm( &attributes, psa_sig_md );
 
-    policy = psa_key_policy_init();
-    psa_key_policy_set_usage( &policy, PSA_KEY_USAGE_VERIFY, psa_sig_md );
-    if( ( ret = psa_set_key_policy( key_slot, &policy ) ) != PSA_SUCCESS )
+    status = psa_import_key( &attributes,
+                             buf + sizeof( buf ) - key_len, key_len,
+                             &key_handle );
+    if( status != PSA_SUCCESS )
     {
-        ret = mbedtls_psa_err_translate_pk( ret );
-        goto cleanup;
-    }
-
-    if( psa_import_key_to_handle( key_slot, psa_type, buf + sizeof( buf ) - key_len, key_len )
-         != PSA_SUCCESS )
-    {
-        ret = MBEDTLS_ERR_PK_BAD_INPUT_DATA;
+        ret = mbedtls_psa_err_translate_pk( status );
         goto cleanup;
     }
 
@@ -611,7 +605,7 @@
         goto cleanup;
     }
 
-    if( psa_asymmetric_verify( key_slot, psa_sig_md,
+    if( psa_asymmetric_verify( key_handle, psa_sig_md,
                                hash, hash_len,
                                buf, 2 * signature_part_size )
          != PSA_SUCCESS )
@@ -628,7 +622,7 @@
     ret = 0;
 
 cleanup:
-    psa_destroy_key( key_slot );
+    psa_destroy_key( key_handle );
     return( ret );
 }
 #else /* MBEDTLS_USE_PSA_CRYPTO */
@@ -898,10 +892,13 @@
 {
     const psa_key_handle_t *key = (const psa_key_handle_t *) ctx;
     size_t bits;
+    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
 
-    if( PSA_SUCCESS != psa_get_key_information( *key, NULL, &bits ) )
+    if( PSA_SUCCESS != psa_get_key_attributes( *key, &attributes ) )
         return( 0 );
 
+    bits = psa_get_key_bits( &attributes );
+    psa_reset_key_attributes( &attributes );
     return( bits );
 }
 
@@ -1002,8 +999,9 @@
                    int (*f_rng)(void *, unsigned char *, size_t), void *p_rng )
 {
     const psa_key_handle_t *key = (const psa_key_handle_t *) ctx;
+    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
     psa_algorithm_t alg = PSA_ALG_ECDSA( mbedtls_psa_translate_md( md_alg ) );
-    size_t bits, buf_len;
+    size_t buf_len;
     psa_status_t status;
 
     /* PSA has its own RNG */
@@ -1014,11 +1012,11 @@
      * that information. Assume that the buffer is large enough for a
      * maximal-length signature with that key (otherwise the application is
      * buggy anyway). */
-    status = psa_get_key_information( *key, NULL, &bits );
+    status = psa_get_key_attributes( *key, &attributes );
     if( status != PSA_SUCCESS )
         return( mbedtls_psa_err_translate_pk( status ) );
-
-    buf_len = MBEDTLS_ECDSA_MAX_SIG_LEN( bits );
+    buf_len = MBEDTLS_ECDSA_MAX_SIG_LEN( psa_get_key_bits( &attributes ) );
+    psa_reset_key_attributes( &attributes );
 
     /* make the signature */
     status = psa_asymmetric_sign( *key, alg, hash, hash_len,
diff --git a/library/pkwrite.c b/library/pkwrite.c
index b87f81b..4388160 100644
--- a/library/pkwrite.c
+++ b/library/pkwrite.c
@@ -246,17 +246,16 @@
 #if defined(MBEDTLS_USE_PSA_CRYPTO)
     if( pk_type == MBEDTLS_PK_OPAQUE )
     {
-        psa_status_t status;
+        psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
         psa_key_type_t key_type;
         psa_key_handle_t handle;
         psa_ecc_curve_t curve;
 
         handle = *((psa_key_handle_t*) key->pk_ctx );
-
-        status = psa_get_key_information( handle, &key_type,
-                                          NULL /* bitsize not needed */ );
-        if( status != PSA_SUCCESS )
+        if( PSA_SUCCESS != psa_get_key_attributes( handle, &attributes ) )
             return( MBEDTLS_ERR_PK_HW_ACCEL_FAILED );
+        key_type = psa_get_key_type( &attributes );
+        psa_reset_key_attributes( &attributes );
 
         curve = PSA_KEY_TYPE_GET_CURVE( key_type );
         if( curve == 0 )
diff --git a/tests/suites/test_suite_pk.function b/tests/suites/test_suite_pk.function
index de90b47..fd923c2 100644
--- a/tests/suites/test_suite_pk.function
+++ b/tests/suites/test_suite_pk.function
@@ -69,37 +69,26 @@
 
 #include "mbedtls/psa_util.h"
 
-#define PK_PSA_INVALID_SLOT 0 /* guaranteed invalid */
-
 /*
- * Generate a key in a free key slot and return this key slot,
- * or PK_PSA_INVALID_SLOT if no slot was available.
+ * Generate a key using PSA and return a handle to that key,
+ * or 0 if the key generation failed.
  * The key uses NIST P-256 and is usable for signing with SHA-256.
  */
 psa_key_handle_t pk_psa_genkey( void )
 {
     psa_key_handle_t key;
-
+    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
     const int curve = PSA_ECC_CURVE_SECP256R1;
     const psa_key_type_t type = PSA_KEY_TYPE_ECC_KEY_PAIR(curve);
     const size_t bits = 256;
-    psa_key_policy_t policy;
 
-    /* Allocate a key slot */
-    if( PSA_SUCCESS != psa_allocate_key( &key ) )
-        return( PK_PSA_INVALID_SLOT );
+    psa_set_key_usage_flags( &attributes, PSA_KEY_USAGE_SIGN );
+    psa_set_key_algorithm( &attributes, PSA_ALG_ECDSA(PSA_ALG_SHA_256) );
+    psa_set_key_type( &attributes, type );
+    psa_set_key_bits( &attributes, bits );
+    PSA_ASSERT( psa_generate_key( &attributes, &key ) );
 
-    /* set up policy on key slot */
-    policy = psa_key_policy_init();
-    psa_key_policy_set_usage( &policy, PSA_KEY_USAGE_SIGN,
-                                      PSA_ALG_ECDSA(PSA_ALG_SHA_256) );
-    if( PSA_SUCCESS != psa_set_key_policy( key, &policy ) )
-        return( PK_PSA_INVALID_SLOT );
-
-    /* generate key */
-    if( PSA_SUCCESS != psa_generate_key_to_handle( key, type, bits, NULL, 0 ) )
-        return( PK_PSA_INVALID_SLOT );
-
+exit:
     return( key );
 }
 #endif /* MBEDTLS_USE_PSA_CRYPTO */
@@ -115,6 +104,7 @@
 {
     mbedtls_pk_context pk, pk2;
     psa_key_handle_t key;
+    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
 
     const char * const name = "Opaque";
     const size_t bitlen = 256; /* harcoded in genkey() */
@@ -136,7 +126,8 @@
     mbedtls_pk_init( &pk );
 
     key = pk_psa_genkey();
-    TEST_ASSERT( key != 0 );
+    if( key == 0 )
+        goto exit;
 
     TEST_ASSERT( mbedtls_pk_setup_opaque( &pk, key ) == 0 );
 
@@ -173,7 +164,7 @@
 
     /* test that freeing the context does not destroy the key */
     mbedtls_pk_free( &pk );
-    TEST_ASSERT( PSA_SUCCESS == psa_get_key_information( key, NULL, NULL ) );
+    TEST_ASSERT( PSA_SUCCESS == psa_get_key_attributes( key, &attributes ) );
     TEST_ASSERT( PSA_SUCCESS == psa_destroy_key( key ) );
 
 exit:
@@ -1233,7 +1224,6 @@
     pkey_legacy_start = pkey_legacy + sizeof( pkey_legacy ) - klen_legacy;
 
     /* Turn PK context into an opaque one. */
-    TEST_ASSERT( psa_allocate_key( &handle ) == PSA_SUCCESS );
     TEST_ASSERT( mbedtls_pk_wrap_as_opaque( &pk, &handle,
                                             PSA_ALG_SHA_256 ) == 0 );