Merge pull request #131 from gilles-peskine-arm/psa-remove_half_filled_slots

Remove old key creation functions and update cipher and pk
diff --git a/include/mbedtls/cipher_internal.h b/include/mbedtls/cipher_internal.h
index d711339..8ea2a98 100644
--- a/include/mbedtls/cipher_internal.h
+++ b/include/mbedtls/cipher_internal.h
@@ -124,14 +124,13 @@
     MBEDTLS_CIPHER_PSA_KEY_UNSET = 0,
     MBEDTLS_CIPHER_PSA_KEY_OWNED, /* Used for PSA-based cipher contexts which */
                                   /* use raw key material internally imported */
-                                  /* into a allocated key slot, and which     */
-                                  /* hence need to destroy that key slot      */
-                                  /* when they are no longer needed.          */
+                                  /* as a volatile key, and which hence need  */
+                                  /* to destroy that key when the context is  */
+                                  /* freed.                                   */
     MBEDTLS_CIPHER_PSA_KEY_NOT_OWNED, /* Used for PSA-based cipher contexts   */
-                                      /* which use a key from a key slot      */
-                                      /* provided by the user, and which      */
-                                      /* hence should not be destroyed when   */
-                                      /* the context is no longer needed.     */
+                                      /* which use a key provided by the      */
+                                      /* user, and which hence will not be    */
+                                      /* destroyed when the context is freed. */
 } mbedtls_cipher_psa_key_ownership;
 
 typedef struct
diff --git a/include/mbedtls/pk.h b/include/mbedtls/pk.h
index 24951a6..0e24b1a 100644
--- a/include/mbedtls/pk.h
+++ b/include/mbedtls/pk.h
@@ -217,7 +217,7 @@
  *
  * \note            For contexts that have been set up with
  *                  mbedtls_pk_setup_opaque(), this does not free the underlying
- *                  key slot and you still need to call psa_destroy_key()
+ *                  PSA key and you still need to call psa_destroy_key()
  *                  independently if you want to destroy that key.
  */
 void mbedtls_pk_free( mbedtls_pk_context *ctx );
@@ -259,21 +259,21 @@
 
 #if defined(MBEDTLS_USE_PSA_CRYPTO)
 /**
- * \brief           Initialize a PK context to wrap a PSA key slot.
+ * \brief           Initialize a PK context to wrap a PSA key.
  *
  * \note            This function replaces mbedtls_pk_setup() for contexts
- *                  that wrap a (possibly opaque) PSA key slot instead of
+ *                  that wrap a (possibly opaque) PSA key instead of
  *                  storing and manipulating the key material directly.
  *
  * \param ctx       The context to initialize. It must be empty (type NONE).
- * \param key       The PSA key slot to wrap, which must hold an ECC key pair
+ * \param key       The PSA key to wrap, which must hold an ECC key pair
  *                  (see notes below).
  *
- * \note            The wrapped key slot must remain valid as long as the
+ * \note            The wrapped key must remain valid as long as the
  *                  wrapping PK context is in use, that is at least between
  *                  the point this function is called and the point
  *                  mbedtls_pk_free() is called on this context. The wrapped
- *                  key slot might then be independently used or destroyed.
+ *                  key might then be independently used or destroyed.
  *
  * \note            This function is currently only available for ECC key
  *                  pairs (that is, ECC keys containing private key material).
@@ -281,7 +281,7 @@
  *
  * \return          \c 0 on success.
  * \return          #MBEDTLS_ERR_PK_BAD_INPUT_DATA on invalid input
- *                  (context already used, invalid key slot).
+ *                  (context already used, invalid key handle).
  * \return          #MBEDTLS_ERR_PK_FEATURE_UNAVAILABLE if the key is not an
  *                  ECC key pair.
  * \return          #MBEDTLS_ERR_PK_ALLOC_FAILED on allocation failure.
@@ -788,7 +788,7 @@
 
 #if defined(MBEDTLS_USE_PSA_CRYPTO)
 /**
- * \brief           Turn an EC key into an Opaque one
+ * \brief           Turn an EC key into an opaque one.
  *
  * \warning         This is a temporary utility function for tests. It might
  *                  change or be removed at any time without notice.
@@ -796,18 +796,19 @@
  * \note            Only ECDSA keys are supported so far. Signing with the
  *                  specified hash is the only allowed use of that key.
  *
- * \param pk        Input: the EC key to transfer to a PSA key slot.
- *                  Output: a PK context wrapping that PSA key slot.
- * \param slot      Output: the chosen slot for storing the key.
- *                  It's the caller's responsibility to destroy that slot
- *                  after calling mbedtls_pk_free() on the PK context.
+ * \param pk        Input: the EC key to import to a PSA key.
+ *                  Output: a PK context wrapping that PSA key.
+ * \param handle    Output: a PSA key handle.
+ *                  It's the caller's responsibility to call
+ *                  psa_destroy_key() on that handle after calling
+ *                  mbedtls_pk_free() on the PK context.
  * \param hash_alg  The hash algorithm to allow for use with that key.
  *
  * \return          \c 0 if successful.
  * \return          An Mbed TLS error code otherwise.
  */
 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 );
 #endif /* MBEDTLS_USE_PSA_CRYPTO */
 
diff --git a/include/psa/crypto_extra.h b/include/psa/crypto_extra.h
index b3ec54f..56e0536 100644
--- a/include/psa/crypto_extra.h
+++ b/include/psa/crypto_extra.h
@@ -247,248 +247,6 @@
 /* FIXME Deprecated. Remove this as soon as all the tests are updated. */
 #define PSA_ALG_SELECT_RAW                      ((psa_algorithm_t)0x31000001)
 
-/** \defgroup policy Key policies
- * @{
- *
- * The functions in this section are legacy interfaces where the properties
- * of a key object are set after allocating a handle, in constrast with the
- * preferred interface where key objects are created atomically from
- * a structure that represents the properties.
- */
-
-/** \def PSA_KEY_POLICY_INIT
- *
- * This macro returns a suitable initializer for a key policy object of type
- * #psa_key_policy_t.
- */
-#ifdef __DOXYGEN_ONLY__
-/* This is an example definition for documentation purposes.
- * Implementations should define a suitable value in `crypto_struct.h`.
- */
-#define PSA_KEY_POLICY_INIT {0}
-#endif
-
-/** Return an initial value for a key policy that forbids all usage of the key.
- */
-static psa_key_policy_t psa_key_policy_init(void);
-
-/** \brief Set the standard fields of a policy structure.
- *
- * Note that this function does not make any consistency check of the
- * parameters. The values are only checked when applying the policy to
- * a key with psa_set_key_policy().
- *
- * \param[in,out] policy The key policy to modify. It must have been
- *                       initialized as per the documentation for
- *                       #psa_key_policy_t.
- * \param usage          The permitted uses for the key.
- * \param alg            The algorithm that the key may be used for.
- */
-void psa_key_policy_set_usage(psa_key_policy_t *policy,
-                              psa_key_usage_t usage,
-                              psa_algorithm_t alg);
-
-/** \brief Retrieve the usage field of a policy structure.
- *
- * \param[in] policy    The policy object to query.
- *
- * \return The permitted uses for a key with this policy.
- */
-psa_key_usage_t psa_key_policy_get_usage(const psa_key_policy_t *policy);
-
-/** \brief Retrieve the algorithm field of a policy structure.
- *
- * \param[in] policy    The policy object to query.
- *
- * \return The permitted algorithm for a key with this policy.
- */
-psa_algorithm_t psa_key_policy_get_algorithm(const psa_key_policy_t *policy);
-
-/** \brief Set the usage policy for a key.
- *
- * This function must be called on a key handle before importing,
- * generating or creating a key. Changing the policy of an
- * existing key is not permitted.
- *
- * Implementations may set restrictions on supported key policies
- * depending on the key type.
- *
- * \param handle        Handle to the key whose policy is to be changed.
- * \param[in] policy    The policy object to query.
- *
- * \retval #PSA_SUCCESS
- *         Success.
- *         If the key is persistent, it is implementation-defined whether
- *         the policy has been saved to persistent storage. Implementations
- *         may defer saving the policy until the key material is created.
- * \retval #PSA_ERROR_INVALID_HANDLE
- * \retval #PSA_ERROR_ALREADY_EXISTS
- * \retval #PSA_ERROR_NOT_SUPPORTED
- * \retval #PSA_ERROR_INVALID_ARGUMENT
- * \retval #PSA_ERROR_COMMUNICATION_FAILURE
- * \retval #PSA_ERROR_HARDWARE_FAILURE
- * \retval #PSA_ERROR_CORRUPTION_DETECTED
- * \retval #PSA_ERROR_BAD_STATE
- *         The library has not been previously initialized by psa_crypto_init().
- *         It is implementation-dependent whether a failure to initialize
- *         results in this error code.
- */
-psa_status_t psa_set_key_policy(psa_key_handle_t handle,
-                                const psa_key_policy_t *policy);
-
-/** \brief Get the usage policy for a key.
- *
- * \param handle        Handle to the key whose policy is being queried.
- * \param[out] policy   On success, the key's policy.
- *
- * \retval #PSA_SUCCESS
- * \retval #PSA_ERROR_INVALID_HANDLE
- * \retval #PSA_ERROR_COMMUNICATION_FAILURE
- * \retval #PSA_ERROR_HARDWARE_FAILURE
- * \retval #PSA_ERROR_CORRUPTION_DETECTED
- * \retval #PSA_ERROR_BAD_STATE
- *         The library has not been previously initialized by psa_crypto_init().
- *         It is implementation-dependent whether a failure to initialize
- *         results in this error code.
- */
-psa_status_t psa_get_key_policy(psa_key_handle_t handle,
-                                psa_key_policy_t *policy);
-
-/**@}*/
-
-/** \defgroup to_handle Key creation to allocated handle
- * @{
- *
- * The functions in this section are legacy interfaces where the properties
- * of a key object are set after allocating a handle, in constrast with the
- * preferred interface where key objects are created atomically from
- * a structure that represents the properties.
- */
-
-/** Create a new persistent key.
- *
- * Create a new persistent key and return a handle to it. The handle
- * remains valid until the application calls psa_close_key() or terminates.
- * The application can open the key again with psa_open_key() until it
- * removes the key by calling psa_destroy_key().
- *
- * \param lifetime      The lifetime of the key. This designates a storage
- *                      area where the key material is stored. This must not
- *                      be #PSA_KEY_LIFETIME_VOLATILE.
- * \param id            The persistent identifier of the key.
- * \param[out] handle   On success, a handle to the newly created key.
- *                      When key material is later created in this key,
- *                      it will be saved to the specified persistent location.
- *
- * \retval #PSA_SUCCESS
- *         Success. The application can now use the value of `*handle`
- *         for key operations.
- * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
- * \retval #PSA_ERROR_INSUFFICIENT_STORAGE
- * \retval #PSA_ERROR_ALREADY_EXISTS
- *         There is already a key with the identifier \p id in the storage
- *         area designated by \p lifetime.
- * \retval #PSA_ERROR_INVALID_ARGUMENT
- *         \p lifetime is invalid, for example #PSA_KEY_LIFETIME_VOLATILE.
- * \retval #PSA_ERROR_INVALID_ARGUMENT
- *         \p id is invalid for the specified lifetime.
- * \retval #PSA_ERROR_NOT_SUPPORTED
- *         \p lifetime is not supported.
- * \retval #PSA_ERROR_NOT_PERMITTED
- *         \p lifetime is valid, but the application does not have the
- *         permission to create a key there.
- */
-psa_status_t psa_create_key(psa_key_lifetime_t lifetime,
-                            psa_key_id_t id,
-                            psa_key_handle_t *handle);
-
-/** Allocate space for a transient key, i.e. a key which is only stored
- * in volatile memory.
- *
- * The allocated key and its handle remain valid until the
- * application calls psa_close_key() or psa_destroy_key() or until the
- * application terminates.
- *
- * \param[out] handle   On success, a handle to a volatile key.
- *
- * \retval #PSA_SUCCESS
- *         Success. The application can now use the value of `*handle`
- *         to refer to the key.
- * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
- *         There was not enough memory, or the maximum number of transient keys
- *         has been reached.
- */
-psa_status_t psa_allocate_key(psa_key_handle_t *handle);
-
-/**
- * \brief Get basic metadata about a key.
- *
- * \param handle        Handle to the key to query.
- * \param[out] type     On success, the key type (a \c PSA_KEY_TYPE_XXX value).
- *                      This may be a null pointer, in which case the key type
- *                      is not written.
- * \param[out] bits     On success, the key size in bits.
- *                      This may be a null pointer, in which case the key size
- *                      is not written.
- *
- * \retval #PSA_SUCCESS
- * \retval #PSA_ERROR_INVALID_HANDLE
- * \retval #PSA_ERROR_DOES_NOT_EXIST
- *         The handle does not contain a key.
- * \retval #PSA_ERROR_COMMUNICATION_FAILURE
- * \retval #PSA_ERROR_HARDWARE_FAILURE
- * \retval #PSA_ERROR_CORRUPTION_DETECTED
- * \retval #PSA_ERROR_BAD_STATE
- *         The library has not been previously initialized by psa_crypto_init().
- *         It is implementation-dependent whether a failure to initialize
- *         results in this error code.
- */
-psa_status_t psa_get_key_information(psa_key_handle_t handle,
-                                     psa_key_type_t *type,
-                                     size_t *bits);
-
-/** \brief Retrieve the lifetime of an open key.
- *
- * \param handle        Handle to query.
- * \param[out] lifetime On success, the lifetime value.
- *
- * \retval #PSA_SUCCESS
- *         Success.
- * \retval #PSA_ERROR_INVALID_HANDLE
- * \retval #PSA_ERROR_COMMUNICATION_FAILURE
- * \retval #PSA_ERROR_HARDWARE_FAILURE
- * \retval #PSA_ERROR_CORRUPTION_DETECTED
- * \retval #PSA_ERROR_BAD_STATE
- *         The library has not been previously initialized by psa_crypto_init().
- *         It is implementation-dependent whether a failure to initialize
- *         results in this error code.
- */
-psa_status_t psa_get_key_lifetime_from_handle(psa_key_handle_t handle,
-                                  psa_key_lifetime_t *lifetime);
-
-psa_status_t psa_import_key_to_handle(psa_key_handle_t handle,
-                            psa_key_type_t type,
-                            const uint8_t *data,
-                            size_t data_length);
-
-psa_status_t psa_copy_key_to_handle(psa_key_handle_t source_handle,
-                          psa_key_handle_t target_handle,
-                          const psa_key_policy_t *constraint);
-
-psa_status_t psa_generate_derived_key_to_handle(psa_key_handle_t handle,
-                                      psa_key_type_t type,
-                                      size_t bits,
-                                      psa_key_derivation_operation_t *operation);
-
-psa_status_t psa_generate_key_to_handle(psa_key_handle_t handle,
-                              psa_key_type_t type,
-                              size_t bits,
-                              const void *extra,
-                              size_t extra_size);
-
-/**@}*/
-
-
 /** \addtogroup crypto_types
  * @{
  */
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/library/psa_crypto.c b/library/psa_crypto.c
index 768410c..b3be261 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -739,27 +739,6 @@
     return( status );
 }
 
-/* Retrieve an empty key slot (slot with no key data, but possibly
- * with some metadata such as a policy or domain parameters). */
-static psa_status_t psa_get_empty_key_slot( psa_key_handle_t handle,
-                                            psa_key_slot_t **p_slot )
-{
-    psa_status_t status;
-    psa_key_slot_t *slot = NULL;
-
-    *p_slot = NULL;
-
-    status = psa_get_key_slot( handle, &slot );
-    if( status != PSA_SUCCESS )
-        return( status );
-
-    if( slot->type != PSA_KEY_TYPE_NONE )
-        return( PSA_ERROR_ALREADY_EXISTS );
-
-    *p_slot = slot;
-    return( status );
-}
-
 /** Calculate the intersection of two algorithm usage policies.
  *
  * Return 0 (which allows no operation) on incompatibility.
@@ -938,45 +917,6 @@
     return( status );
 }
 
-psa_status_t psa_import_key_to_handle( psa_key_handle_t handle,
-                             psa_key_type_t type,
-                             const uint8_t *data,
-                             size_t data_length )
-{
-    psa_key_slot_t *slot;
-    psa_status_t status;
-
-    status = psa_get_empty_key_slot( handle, &slot );
-    if( status != PSA_SUCCESS )
-        return( status );
-
-    slot->type = type;
-
-    status = psa_import_key_into_slot( slot, data, data_length );
-    if( status != PSA_SUCCESS )
-    {
-        slot->type = PSA_KEY_TYPE_NONE;
-        return( status );
-    }
-
-#if defined(MBEDTLS_PSA_CRYPTO_STORAGE_C)
-    if( slot->lifetime == PSA_KEY_LIFETIME_PERSISTENT )
-    {
-        /* Store in file location */
-        status = psa_save_persistent_key( slot->persistent_storage_id,
-                                          slot->type, &slot->policy, data,
-                                          data_length );
-        if( status != PSA_SUCCESS )
-        {
-            (void) psa_remove_key_data_from_memory( slot );
-            slot->type = PSA_KEY_TYPE_NONE;
-        }
-    }
-#endif /* defined(MBEDTLS_PSA_CRYPTO_STORAGE_C) */
-
-    return( status );
-}
-
 psa_status_t psa_destroy_key( psa_key_handle_t handle )
 {
     psa_key_slot_t *slot;
@@ -1143,30 +1083,6 @@
     return( status );
 }
 
-psa_status_t psa_get_key_information( psa_key_handle_t handle,
-                                      psa_key_type_t *type,
-                                      size_t *bits )
-{
-    psa_key_slot_t *slot;
-    psa_status_t status;
-
-    if( type != NULL )
-        *type = 0;
-    if( bits != NULL )
-        *bits = 0;
-    status = psa_get_key_slot( handle, &slot );
-    if( status != PSA_SUCCESS )
-        return( status );
-
-    if( slot->type == PSA_KEY_TYPE_NONE )
-        return( PSA_ERROR_DOES_NOT_EXIST );
-    if( type != NULL )
-        *type = slot->type;
-    if( bits != NULL )
-        *bits = psa_get_key_slot_bits( slot );
-    return( PSA_SUCCESS );
-}
-
 #if defined(MBEDTLS_RSA_C) || defined(MBEDTLS_ECP_C)
 static int pk_write_pubkey_simple( mbedtls_pk_context *key,
                                    unsigned char *buf, size_t size )
@@ -1341,39 +1257,6 @@
                                      data_length, 1 ) );
 }
 
-#if defined(MBEDTLS_PSA_CRYPTO_STORAGE_C)
-static psa_status_t psa_save_generated_persistent_key( psa_key_slot_t *slot,
-                                                       size_t bits )
-{
-    psa_status_t status;
-    uint8_t *data;
-    size_t key_length;
-    size_t data_size = PSA_KEY_EXPORT_MAX_SIZE( slot->type, bits );
-    data = mbedtls_calloc( 1, data_size );
-    if( data == NULL )
-        return( PSA_ERROR_INSUFFICIENT_MEMORY );
-    /* Get key data in export format */
-    status = psa_internal_export_key( slot, data, data_size, &key_length, 0 );
-    if( status != PSA_SUCCESS )
-    {
-        slot->type = PSA_KEY_TYPE_NONE;
-        goto exit;
-    }
-    /* Store in file location */
-    status = psa_save_persistent_key( slot->persistent_storage_id,
-                                      slot->type, &slot->policy,
-                                      data, key_length );
-    if( status != PSA_SUCCESS )
-    {
-        slot->type = PSA_KEY_TYPE_NONE;
-    }
-exit:
-    mbedtls_platform_zeroize( data, key_length );
-    mbedtls_free( data );
-    return( status );
-}
-#endif /* defined(MBEDTLS_PSA_CRYPTO_STORAGE_C) */
-
 static psa_status_t psa_set_key_policy_internal(
     psa_key_slot_t *slot,
     const psa_key_policy_t *policy )
@@ -1422,10 +1305,7 @@
     psa_status_t status;
     psa_key_slot_t *slot;
 
-    status = psa_allocate_key( handle );
-    if( status != PSA_SUCCESS )
-        return( status );
-    status = psa_get_key_slot( *handle, p_slot );
+    status = psa_internal_allocate_key_slot( handle, p_slot );
     if( status != PSA_SUCCESS )
         return( status );
     slot = *p_slot;
@@ -1623,41 +1503,6 @@
     return( status );
 }
 
-psa_status_t psa_copy_key_to_handle(psa_key_handle_t source_handle,
-                          psa_key_handle_t target_handle,
-                          const psa_key_policy_t *constraint)
-{
-    psa_key_slot_t *source_slot = NULL;
-    psa_key_slot_t *target_slot = NULL;
-    psa_key_policy_t new_policy;
-    psa_status_t status;
-    status = psa_get_key_from_slot( source_handle, &source_slot,
-                                    PSA_KEY_USAGE_COPY, 0 );
-    if( status != PSA_SUCCESS )
-        return( status );
-    status = psa_get_empty_key_slot( target_handle, &target_slot );
-    if( status != PSA_SUCCESS )
-        return( status );
-
-    new_policy = target_slot->policy;
-    status = psa_restrict_key_policy( &new_policy, &source_slot->policy );
-    if( status != PSA_SUCCESS )
-        return( status );
-    if( constraint != NULL )
-    {
-        status = psa_restrict_key_policy( &new_policy, constraint );
-        if( status != PSA_SUCCESS )
-            return( status );
-    }
-
-    status = psa_copy_key_material( source_slot, target_slot );
-    if( status != PSA_SUCCESS )
-        return( status );
-
-    target_slot->policy = new_policy;
-    return( PSA_SUCCESS );
-}
-
 psa_status_t psa_copy_key( psa_key_handle_t source_handle,
                            const psa_key_attributes_t *specified_attributes,
                            psa_key_handle_t *target_handle )
@@ -3644,86 +3489,6 @@
 
 
 
-/****************************************************************/
-/* Key Policy */
-/****************************************************************/
-
-#if !defined(MBEDTLS_PSA_CRYPTO_SPM)
-void psa_key_policy_set_usage( psa_key_policy_t *policy,
-                               psa_key_usage_t usage,
-                               psa_algorithm_t alg )
-{
-    policy->usage = usage;
-    policy->alg = alg;
-}
-
-psa_key_usage_t psa_key_policy_get_usage( const psa_key_policy_t *policy )
-{
-    return( policy->usage );
-}
-
-psa_algorithm_t psa_key_policy_get_algorithm( const psa_key_policy_t *policy )
-{
-    return( policy->alg );
-}
-#endif /* !defined(MBEDTLS_PSA_CRYPTO_SPM) */
-
-psa_status_t psa_set_key_policy( psa_key_handle_t handle,
-                                 const psa_key_policy_t *policy )
-{
-    psa_key_slot_t *slot;
-    psa_status_t status;
-
-    if( policy == NULL )
-        return( PSA_ERROR_INVALID_ARGUMENT );
-
-    status = psa_get_empty_key_slot( handle, &slot );
-    if( status != PSA_SUCCESS )
-        return( status );
-
-    return( psa_set_key_policy_internal( slot, policy ) );
-}
-
-psa_status_t psa_get_key_policy( psa_key_handle_t handle,
-                                 psa_key_policy_t *policy )
-{
-    psa_key_slot_t *slot;
-    psa_status_t status;
-
-    if( policy == NULL )
-        return( PSA_ERROR_INVALID_ARGUMENT );
-
-    status = psa_get_key_slot( handle, &slot );
-    if( status != PSA_SUCCESS )
-        return( status );
-
-    *policy = slot->policy;
-
-    return( PSA_SUCCESS );
-}
-
-
-
-/****************************************************************/
-/* Key Lifetime */
-/****************************************************************/
-
-psa_status_t psa_get_key_lifetime_from_handle( psa_key_handle_t handle,
-                                   psa_key_lifetime_t *lifetime )
-{
-    psa_key_slot_t *slot;
-    psa_status_t status;
-
-    status = psa_get_key_slot( handle, &slot );
-    if( status != PSA_SUCCESS )
-        return( status );
-
-    *lifetime = slot->lifetime;
-
-    return( PSA_SUCCESS );
-}
-
-
 
 /****************************************************************/
 /* AEAD */
@@ -4539,37 +4304,6 @@
     return( status );
 }
 
-psa_status_t psa_generate_derived_key_to_handle( psa_key_handle_t handle,
-                                       psa_key_type_t type,
-                                       size_t bits,
-                                       psa_key_derivation_operation_t *operation )
-{
-    uint8_t *data = NULL;
-    size_t bytes = PSA_BITS_TO_BYTES( bits );
-    psa_status_t status;
-
-    if( ! key_type_is_raw_bytes( type ) )
-        return( PSA_ERROR_INVALID_ARGUMENT );
-    if( bits % 8 != 0 )
-        return( PSA_ERROR_INVALID_ARGUMENT );
-    data = mbedtls_calloc( 1, bytes );
-    if( data == NULL )
-        return( PSA_ERROR_INSUFFICIENT_MEMORY );
-
-    status = psa_key_derivation_output_bytes( operation, data, bytes );
-    if( status != PSA_SUCCESS )
-        goto exit;
-#if defined(MBEDTLS_DES_C)
-    if( type == PSA_KEY_TYPE_DES )
-        psa_des_set_key_parity( data, bytes );
-#endif /* MBEDTLS_DES_C */
-    status = psa_import_key_to_handle( handle, type, data, bytes );
-
-exit:
-    mbedtls_free( data );
-    return( status );
-}
-
 
 
 /****************************************************************/
@@ -5426,40 +5160,6 @@
     return( PSA_SUCCESS );
 }
 
-psa_status_t psa_generate_key_to_handle( psa_key_handle_t handle,
-                               psa_key_type_t type,
-                               size_t bits,
-                               const void *extra,
-                               size_t extra_size )
-{
-    psa_key_slot_t *slot;
-    psa_status_t status;
-
-#if defined(MBEDTLS_RSA_C) && defined(MBEDTLS_GENPRIME)
-    /* The old public exponent encoding is no longer supported. */
-    if( extra_size != 0 )
-        return( PSA_ERROR_NOT_SUPPORTED );
-#endif
-
-    status = psa_get_empty_key_slot( handle, &slot );
-    if( status != PSA_SUCCESS )
-        return( status );
-
-    slot->type = type;
-    status = psa_generate_key_internal( slot, bits, extra, extra_size );
-    if( status != PSA_SUCCESS )
-        slot->type = 0;
-
-#if defined(MBEDTLS_PSA_CRYPTO_STORAGE_C)
-    if( slot->lifetime == PSA_KEY_LIFETIME_PERSISTENT )
-    {
-        return( psa_save_generated_persistent_key( slot, bits ) );
-    }
-#endif /* defined(MBEDTLS_PSA_CRYPTO_STORAGE_C) */
-
-    return( status );
-}
-
 psa_status_t psa_generate_key( const psa_key_attributes_t *attributes,
                                psa_key_handle_t *handle )
 {
diff --git a/library/psa_crypto_slot_management.c b/library/psa_crypto_slot_management.c
index 3876f4b..0ffc2aa 100644
--- a/library/psa_crypto_slot_management.c
+++ b/library/psa_crypto_slot_management.c
@@ -99,56 +99,25 @@
     global_data.key_slots_initialized = 0;
 }
 
-/** Find a free key slot and mark it as in use.
- *
- * \param[out] handle   On success, a slot number that is not in use. This
- *                      value can be used as a handle to the slot.
- *
- * \retval #PSA_SUCCESS
- * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
- */
-static psa_status_t psa_internal_allocate_key_slot( psa_key_handle_t *handle )
+psa_status_t psa_internal_allocate_key_slot( psa_key_handle_t *handle,
+                                             psa_key_slot_t **p_slot )
 {
+    if( ! global_data.key_slots_initialized )
+        return( PSA_ERROR_BAD_STATE );
+
     for( *handle = PSA_KEY_SLOT_COUNT; *handle != 0; --( *handle ) )
     {
-        psa_key_slot_t *slot = &global_data.key_slots[*handle - 1];
-        if( ! slot->allocated )
+        *p_slot = &global_data.key_slots[*handle - 1];
+        if( ! ( *p_slot )->allocated )
         {
-            slot->allocated = 1;
+            ( *p_slot )->allocated = 1;
             return( PSA_SUCCESS );
         }
     }
+    *p_slot = NULL;
     return( PSA_ERROR_INSUFFICIENT_MEMORY );
 }
 
-/** Wipe a key slot and mark it as available.
- *
- * This does not affect persistent storage.
- *
- * \param handle        The handle to the key slot to release.
- *
- * \retval #PSA_SUCCESS
- * \retval #PSA_ERROR_INVALID_ARGUMENT
- * \retval #PSA_ERROR_CORRUPTION_DETECTED
- */
-static psa_status_t psa_internal_release_key_slot( psa_key_handle_t handle )
-{
-    psa_key_slot_t *slot;
-    psa_status_t status;
-
-    status = psa_get_key_slot( handle, &slot );
-    if( status != PSA_SUCCESS )
-        return( status );
-
-    return( psa_wipe_key_slot( slot ) );
-}
-
-psa_status_t psa_allocate_key( psa_key_handle_t *handle )
-{
-    *handle = 0;
-    return( psa_internal_allocate_key_slot( handle ) );
-}
-
 #if defined(MBEDTLS_PSA_CRYPTO_STORAGE_C)
 static psa_status_t psa_load_persistent_key_into_slot( psa_key_slot_t *p_slot )
 {
@@ -194,41 +163,6 @@
     else
         return( 0 );
 }
-
-/** Declare a slot as persistent and load it from storage.
- *
- * This function may only be called immediately after a successful call
- * to psa_internal_allocate_key_slot().
- *
- * \param handle        A handle to a key slot freshly allocated with
- *                      psa_internal_allocate_key_slot().
- *
- * \retval #PSA_SUCCESS
- *         The slot content was loaded successfully.
- * \retval #PSA_ERROR_DOES_NOT_EXIST
- *         There is no content for this slot in persistent storage.
- * \retval #PSA_ERROR_INVALID_HANDLE
- * \retval #PSA_ERROR_INVALID_ARGUMENT
- *         \p id is not acceptable.
- * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
- * \retval #PSA_ERROR_STORAGE_FAILURE
- */
-static psa_status_t psa_internal_make_key_persistent( psa_key_handle_t handle,
-                                                      psa_key_file_id_t id )
-{
-    psa_key_slot_t *slot;
-    psa_status_t status;
-
-    status = psa_get_key_slot( handle, &slot );
-    if( status != PSA_SUCCESS )
-        return( status );
-
-    slot->lifetime = PSA_KEY_LIFETIME_PERSISTENT;
-    slot->persistent_storage_id = id;
-    status = psa_load_persistent_key_into_slot( slot );
-
-    return( status );
-}
 #endif /* defined(MBEDTLS_PSA_CRYPTO_STORAGE_C) */
 
 psa_status_t psa_validate_persistent_key_parameters(
@@ -251,63 +185,51 @@
 #endif /* !MBEDTLS_PSA_CRYPTO_STORAGE_C */
 }
 
-static psa_status_t persistent_key_setup( psa_key_lifetime_t lifetime,
-                                          psa_key_file_id_t id,
-                                          psa_key_handle_t *handle,
-                                          int creating )
+psa_status_t psa_open_key( psa_key_file_id_t id, psa_key_handle_t *handle )
 {
+#if defined(MBEDTLS_PSA_CRYPTO_STORAGE_C)
     psa_status_t status;
-    psa_status_t wanted_load_status =
-        ( creating ? PSA_ERROR_DOES_NOT_EXIST : PSA_SUCCESS );
+    psa_key_slot_t *slot;
 
     *handle = 0;
 
-    status = psa_validate_persistent_key_parameters( lifetime, id, creating );
+    status = psa_validate_persistent_key_parameters(
+        PSA_KEY_LIFETIME_PERSISTENT, id, 0 );
     if( status != PSA_SUCCESS )
         return( status );
 
-#if defined(MBEDTLS_PSA_CRYPTO_STORAGE_C)
-    status = psa_internal_allocate_key_slot( handle );
+    status = psa_internal_allocate_key_slot( handle, &slot );
     if( status != PSA_SUCCESS )
         return( status );
 
-    status = psa_internal_make_key_persistent( *handle, id );
-    if( status != wanted_load_status )
+    slot->lifetime = PSA_KEY_LIFETIME_PERSISTENT;
+    slot->persistent_storage_id = id;
+
+    status = psa_load_persistent_key_into_slot( slot );
+    if( status != PSA_SUCCESS )
     {
-        psa_internal_release_key_slot( *handle );
+        psa_wipe_key_slot( slot );
         *handle = 0;
     }
     return( status );
+
 #else /* defined(MBEDTLS_PSA_CRYPTO_STORAGE_C) */
-    (void) wanted_load_status;
+    (void) id;
+    *handle = 0;
     return( PSA_ERROR_NOT_SUPPORTED );
 #endif /* !defined(MBEDTLS_PSA_CRYPTO_STORAGE_C) */
 }
 
-psa_status_t psa_open_key( psa_key_file_id_t id, psa_key_handle_t *handle )
-{
-    return( persistent_key_setup( PSA_KEY_LIFETIME_PERSISTENT,
-                                  id, handle, 0 ) );
-}
-
-psa_status_t psa_create_key( psa_key_lifetime_t lifetime,
-                             psa_key_file_id_t id,
-                             psa_key_handle_t *handle )
-{
-    psa_status_t status;
-
-    status = persistent_key_setup( lifetime, id, handle, 1 );
-    switch( status )
-    {
-        case PSA_SUCCESS: return( PSA_ERROR_ALREADY_EXISTS );
-        case PSA_ERROR_DOES_NOT_EXIST: return( PSA_SUCCESS );
-        default: return( status );
-    }
-}
-
 psa_status_t psa_close_key( psa_key_handle_t handle )
 {
-    return( psa_internal_release_key_slot( handle ) );
+    psa_status_t status;
+    psa_key_slot_t *slot;
+
+    status = psa_get_key_slot( handle, &slot );
+    if( status != PSA_SUCCESS )
+        return( status );
+
+    return( psa_wipe_key_slot( slot ) );
 }
 
 #endif /* MBEDTLS_PSA_CRYPTO_C */
diff --git a/library/psa_crypto_slot_management.h b/library/psa_crypto_slot_management.h
index 2e459d1..5c1bde1 100644
--- a/library/psa_crypto_slot_management.h
+++ b/library/psa_crypto_slot_management.h
@@ -55,6 +55,22 @@
  * This does not affect persistent storage. */
 void psa_wipe_all_key_slots( void );
 
+/** Find a free key slot and mark it as in use.
+ *
+ * \param[out] handle   On success, a slot number that can be used as a
+ *                      handle to the slot. The selected slot was not
+ *                      in use before. This function marks it as in use
+ *                      and otherwise leaves it in a freshly-initialized
+ *                      state.
+ * \param[out] p_slot   On success, a pointer to the slot.
+ *
+ * \retval #PSA_SUCCESS
+ * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
+ * \retval #PSA_ERROR_BAD_STATE
+ */
+psa_status_t psa_internal_allocate_key_slot( psa_key_handle_t *handle,
+                                             psa_key_slot_t **p_slot );
+
 /** Test whether the given parameters are acceptable for a persistent key.
  *
  * This function does not access the storage in any way. It only tests
diff --git a/programs/psa/crypto_examples.c b/programs/psa/crypto_examples.c
index 1a81f45..f156b7b 100644
--- a/programs/psa/crypto_examples.c
+++ b/programs/psa/crypto_examples.c
@@ -206,9 +206,6 @@
     status = psa_generate_random( input, sizeof( input ) );
     ASSERT_STATUS( status, PSA_SUCCESS );
 
-    status = psa_allocate_key( &key_handle );
-    ASSERT_STATUS( status, PSA_SUCCESS );
-
     psa_set_key_usage_flags( &attributes,
                              PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT );
     psa_set_key_algorithm( &attributes, alg );
diff --git a/programs/psa/key_ladder_demo.c b/programs/psa/key_ladder_demo.c
index 36d7b5d..af7be1e 100644
--- a/programs/psa/key_ladder_demo.c
+++ b/programs/psa/key_ladder_demo.c
@@ -260,7 +260,7 @@
     mbedtls_platform_zeroize( key_data, sizeof( key_data ) );
     if( status != PSA_SUCCESS )
     {
-        /* If psa_allocate_key hasn't been called yet or has failed,
+        /* If the key creation hasn't happened yet or has failed,
          * *master_key_handle is 0. psa_destroy_key(0) is guaranteed to do
          * nothing and return PSA_ERROR_INVALID_HANDLE. */
         (void) psa_destroy_key( *master_key_handle );
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 );
 
diff --git a/tests/suites/test_suite_psa_crypto_slot_management.function b/tests/suites/test_suite_psa_crypto_slot_management.function
index 3bc9f78..db46328 100644
--- a/tests/suites/test_suite_psa_crypto_slot_management.function
+++ b/tests/suites/test_suite_psa_crypto_slot_management.function
@@ -584,8 +584,6 @@
 {
     psa_key_handle_t handle1 = 0;
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
-    psa_key_type_t read_type;
-    size_t read_bits;
     uint8_t material[1] = "a";
 
     PSA_ASSERT( psa_crypto_init( ) );
@@ -608,9 +606,10 @@
     TEST_EQUAL( psa_destroy_key( handle1 + 1 ), PSA_ERROR_INVALID_HANDLE );
 
     /* After all this, check that the original handle is intact. */
-    PSA_ASSERT( psa_get_key_information( handle1, &read_type, &read_bits ) );
-    TEST_EQUAL( read_type, PSA_KEY_TYPE_RAW_DATA );
-    TEST_EQUAL( read_bits, PSA_BYTES_TO_BITS( sizeof( material ) ) );
+    PSA_ASSERT( psa_get_key_attributes( handle1, &attributes ) );
+    TEST_EQUAL( psa_get_key_type( &attributes ), PSA_KEY_TYPE_RAW_DATA );
+    TEST_EQUAL( psa_get_key_bits( &attributes ),
+                PSA_BYTES_TO_BITS( sizeof( material ) ) );
     PSA_ASSERT( psa_close_key( handle1 ) );
 
 exit: