Implement atomic-creation psa_import_key

Implement the new, attribute-based psa_import_key and some basic
functions to access psa_key_attributes_t. Replace
psa_import_key_to_handle by psa_import_key in a few test functions.

This commit does not handle persistence attributes yet.
diff --git a/include/psa/crypto.h b/include/psa/crypto.h
index 564dd87..74a36b0 100644
--- a/include/psa/crypto.h
+++ b/include/psa/crypto.h
@@ -109,6 +109,39 @@
  */
 typedef struct psa_key_attributes_s psa_key_attributes_t;
 
+static void psa_make_key_persistent(psa_key_attributes_t *attributes,
+                                    psa_key_id_t id,
+                                    psa_key_lifetime_t lifetime);
+
+static psa_key_id_t psa_get_key_id(const psa_key_attributes_t *attributes);
+
+static psa_key_lifetime_t psa_get_key_lifetime(
+    const psa_key_attributes_t *attributes);
+
+static void psa_set_key_usage_flags(psa_key_attributes_t *attributes,
+                                    psa_key_usage_t usage_flags);
+
+static psa_key_usage_t psa_get_key_usage_flags(
+    const psa_key_attributes_t *attributes);
+
+static void psa_set_key_algorithm(psa_key_attributes_t *attributes,
+                                  psa_algorithm_t alg);
+
+static psa_algorithm_t psa_get_key_algorithm(
+    const psa_key_attributes_t *attributes);
+
+static void psa_set_key_type(psa_key_attributes_t *attributes,
+                             psa_key_type_t type);
+
+static psa_key_type_t psa_get_key_type(const psa_key_attributes_t *attributes);
+
+static size_t psa_get_key_bits(const psa_key_attributes_t *attributes);
+
+psa_status_t psa_get_key_attributes(psa_key_handle_t handle,
+                                    psa_key_attributes_t *attributes);
+
+psa_status_t psa_reset_key_attributes(psa_key_attributes_t *attributes);
+
 /**@}*/
 
 /** \defgroup policy Key policies
@@ -380,7 +413,6 @@
  */
 psa_status_t psa_import_key(const psa_key_attributes_t *attributes,
                             psa_key_handle_t *handle,
-                            psa_key_type_t type,
                             const uint8_t *data,
                             size_t data_length);
 
@@ -2970,7 +3002,6 @@
  */
 psa_status_t psa_generator_import_key(const psa_key_attributes_t *attributes,
                                       psa_key_handle_t *handle,
-                                      psa_key_type_t type,
                                       size_t bits,
                                       psa_crypto_generator_t *generator);
 
@@ -3363,7 +3394,6 @@
  */
 psa_status_t psa_generate_key(const psa_key_attributes_t *attributes,
                               psa_key_handle_t *handle,
-                              psa_key_type_t type,
                               size_t bits,
                               const void *extra,
                               size_t extra_size);
diff --git a/include/psa/crypto_struct.h b/include/psa/crypto_struct.h
index 6eed259..16674d4 100644
--- a/include/psa/crypto_struct.h
+++ b/include/psa/crypto_struct.h
@@ -260,4 +260,56 @@
     return( v );
 }
 
+struct psa_key_attributes_s
+{
+    psa_key_id_t id;
+    psa_key_lifetime_t lifetime;
+    psa_key_policy_t policy;
+    psa_key_type_t type;
+    size_t bits;
+};
+
+#define PSA_KEY_ATTRIBUTES_INIT {0, 0, {0, 0}, 0, 0}
+static inline struct psa_key_attributes_s psa_key_attributes_init( void )
+{
+    const struct psa_key_attributes_s v = PSA_KEY_ATTRIBUTES_INIT;
+    return( v );
+}
+
+static inline void psa_set_key_usage_flags(psa_key_attributes_t *attributes,
+                                           psa_key_usage_t usage_flags)
+{
+    attributes->policy.usage = usage_flags;
+}
+
+static inline psa_key_usage_t psa_get_key_usage_flags(
+    const psa_key_attributes_t *attributes)
+{
+    return( attributes->policy.usage );
+}
+
+static inline void psa_set_key_algorithm(psa_key_attributes_t *attributes,
+                                         psa_algorithm_t alg)
+{
+    attributes->policy.alg = alg;
+}
+
+static inline psa_algorithm_t psa_get_key_algorithm(
+    const psa_key_attributes_t *attributes)
+{
+    return( attributes->policy.alg );
+}
+
+static inline void psa_set_key_type(psa_key_attributes_t *attributes,
+                                    psa_key_type_t type)
+{
+    attributes->type = type;
+}
+
+static inline psa_key_type_t psa_get_key_type(
+    const psa_key_attributes_t *attributes)
+{
+    return( attributes->type );
+}
+
 #endif /* PSA_CRYPTO_STRUCT_H */
diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 2fab91c..9b43d13 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -1212,6 +1212,140 @@
 }
 #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 )
+{
+    if( ( policy->usage & ~( PSA_KEY_USAGE_EXPORT |
+                             PSA_KEY_USAGE_ENCRYPT |
+                             PSA_KEY_USAGE_DECRYPT |
+                             PSA_KEY_USAGE_SIGN |
+                             PSA_KEY_USAGE_VERIFY |
+                             PSA_KEY_USAGE_DERIVE ) ) != 0 )
+        return( PSA_ERROR_INVALID_ARGUMENT );
+
+    slot->policy = *policy;
+    return( PSA_SUCCESS );
+}
+
+/** Prepare a key slot to receive key material.
+ *
+ * This function allocates a key slot and sets its metadata.
+ *
+ * If this function fails, call psa_fail_key_creation().
+ *
+ * \param attributes    Key attributes for the new key.
+ * \param handle        On success, the allocated handle.
+ * \param p_slot        On success, a pointer to the prepared slot.
+ */
+static psa_status_t psa_start_key_creation(
+    const psa_key_attributes_t *attributes,
+    psa_key_handle_t *handle,
+    psa_key_slot_t **p_slot )
+{
+    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 );
+    if( status != PSA_SUCCESS )
+        return( status );
+    slot = *p_slot;
+
+    status = psa_set_key_policy_internal( slot, &attributes->policy );
+    if( status != PSA_SUCCESS )
+        return( status );
+    slot->lifetime = attributes->lifetime;
+    if( attributes->lifetime != PSA_KEY_LIFETIME_VOLATILE )
+        slot->persistent_storage_id = attributes->id;
+    slot->type = attributes->type;
+
+    return( status );
+}
+
+/** Finalize the creation of a key once its key material has been set.
+ *
+ * This entails writing the key to persistent storage.
+ *
+ * If this function fails, call psa_fail_key_creation().
+ *
+ * \param slot          Pointer to the slot with key material.
+ */
+static psa_status_t psa_finish_key_creation( psa_key_slot_t *slot )
+{
+    psa_status_t status = PSA_SUCCESS;
+
+#if defined(MBEDTLS_PSA_CRYPTO_STORAGE_C)
+    if( slot->lifetime == PSA_KEY_LIFETIME_PERSISTENT )
+    {
+        uint8_t *buffer = NULL;
+        size_t buffer_size = 0;
+        size_t length;
+
+        buffer_size = PSA_KEY_EXPORT_MAX_SIZE( slot->type,
+                                               psa_get_key_bits( slot ) );
+        buffer = mbedtls_calloc( 1, buffer_size );
+        if( buffer == NULL && buffer_size != 0 )
+            return( PSA_ERROR_INSUFFICIENT_MEMORY );
+        status = psa_internal_export_key( slot,
+                                          buffer, buffer_size, &length,
+                                          0 );
+
+        if( status == PSA_SUCCESS )
+        {
+            status = psa_save_persistent_key( slot->persistent_storage_id,
+                                              slot->type, &slot->policy,
+                                              buffer, length );
+        }
+
+        if( buffer_size != 0 )
+            mbedtls_platform_zeroize( buffer, buffer_size );
+        mbedtls_free( buffer );
+    }
+#endif /* defined(MBEDTLS_PSA_CRYPTO_STORAGE_C) */
+
+    return( status );
+}
+
+/** Abort the creation of a key.
+ *
+ * You may call this function after calling psa_start_key_creation(),
+ * or after psa_finish_key_creation() fails. In other circumstances, this
+ * function may not clean up persistent storage.
+ *
+ * \param slot          Pointer to the slot with key material.
+ */
+static void psa_fail_key_creation( psa_key_slot_t *slot )
+{
+    if( slot == NULL )
+        return;
+    psa_wipe_key_slot( slot );
+}
+
+psa_status_t psa_import_key( const psa_key_attributes_t *attributes,
+                             psa_key_handle_t *handle,
+                             const uint8_t *data,
+                             size_t data_length )
+{
+    psa_status_t status;
+    psa_key_slot_t *slot = NULL;
+    status = psa_start_key_creation( attributes, handle, &slot );
+    if( status == PSA_SUCCESS )
+    {
+        status = psa_import_key_into_slot( slot, data, data_length );
+    }
+    if( status == PSA_SUCCESS )
+        status = psa_finish_key_creation( slot );
+    if( status != PSA_SUCCESS )
+    {
+        psa_fail_key_creation( slot );
+        *handle = 0;
+    }
+    return( status );
+}
+
 static psa_status_t psa_copy_key_material( const psa_key_slot_t *source,
                                            psa_key_handle_t target )
 {
@@ -3240,17 +3374,7 @@
     if( status != PSA_SUCCESS )
         return( status );
 
-    if( ( policy->usage & ~( PSA_KEY_USAGE_EXPORT |
-                             PSA_KEY_USAGE_ENCRYPT |
-                             PSA_KEY_USAGE_DECRYPT |
-                             PSA_KEY_USAGE_SIGN |
-                             PSA_KEY_USAGE_VERIFY |
-                             PSA_KEY_USAGE_DERIVE ) ) != 0 )
-        return( PSA_ERROR_INVALID_ARGUMENT );
-
-    slot->policy = *policy;
-
-    return( PSA_SUCCESS );
+    return( psa_set_key_policy_internal( slot, policy ) );
 }
 
 psa_status_t psa_get_key_policy( psa_key_handle_t handle,
diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index 7972597..b9f0d5f 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -1111,14 +1111,15 @@
 /* BEGIN_CASE */
 void import( data_t *data, int type, int expected_status_arg )
 {
+    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
     psa_key_handle_t handle = 0;
     psa_status_t expected_status = expected_status_arg;
     psa_status_t status;
 
     PSA_ASSERT( psa_crypto_init( ) );
 
-    PSA_ASSERT( psa_allocate_key( &handle ) );
-    status = psa_import_key_to_handle( handle, type, data->x, data->len );
+    psa_set_key_type( &attributes, type );
+    status = psa_import_key( &attributes, &handle, data->x, data->len );
     TEST_EQUAL( status, expected_status );
     if( status == PSA_SUCCESS )
         PSA_ASSERT( psa_destroy_key( handle ) );
@@ -1226,7 +1227,7 @@
     size_t reexported_length;
     psa_key_type_t got_type;
     size_t got_bits;
-    psa_key_policy_t policy = PSA_KEY_POLICY_INIT;
+    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
 
     export_size = (ptrdiff_t) data->len + export_size_delta;
     ASSERT_ALLOC( exported, export_size );
@@ -1234,16 +1235,12 @@
         ASSERT_ALLOC( reexported, export_size );
     PSA_ASSERT( psa_crypto_init( ) );
 
-    PSA_ASSERT( psa_allocate_key( &handle ) );
-    psa_key_policy_set_usage( &policy, usage_arg, alg );
-    PSA_ASSERT( psa_set_key_policy( handle, &policy ) );
-
-    TEST_EQUAL( psa_get_key_information( handle, NULL, NULL ),
-                PSA_ERROR_DOES_NOT_EXIST );
+    psa_set_key_usage_flags( &attributes, usage_arg );
+    psa_set_key_algorithm( &attributes, alg );
+    psa_set_key_type( &attributes, type );
 
     /* Import the key */
-    PSA_ASSERT( psa_import_key_to_handle( handle, type,
-                                data->x, data->len ) );
+    PSA_ASSERT( psa_import_key( &attributes, &handle, data->x, data->len ) );
 
     /* Test the key information */
     PSA_ASSERT( psa_get_key_information( handle,
@@ -1280,12 +1277,8 @@
     else
     {
         psa_key_handle_t handle2;
-        PSA_ASSERT( psa_allocate_key( &handle2 ) );
-        PSA_ASSERT( psa_set_key_policy( handle2, &policy ) );
-
-        PSA_ASSERT( psa_import_key_to_handle( handle2, type,
-                                    exported,
-                                    exported_length ) );
+        PSA_ASSERT( psa_import_key( &attributes, &handle2,
+                                    exported, exported_length ) );
         PSA_ASSERT( psa_export_key( handle2,
                                     reexported,
                                     export_size,
@@ -1525,17 +1518,16 @@
     unsigned char *exported = NULL;
     size_t export_size = expected_public_key->len + export_size_delta;
     size_t exported_length = INVALID_EXPORT_LENGTH;
-    psa_key_policy_t policy = PSA_KEY_POLICY_INIT;
+    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
 
     PSA_ASSERT( psa_crypto_init( ) );
 
-    PSA_ASSERT( psa_allocate_key( &handle ) );
-    psa_key_policy_set_usage( &policy, PSA_KEY_USAGE_EXPORT, alg );
-    PSA_ASSERT( psa_set_key_policy( handle, &policy ) );
+    psa_set_key_usage_flags( &attributes, PSA_KEY_USAGE_EXPORT );
+    psa_set_key_algorithm( &attributes, alg );
+    psa_set_key_type( &attributes, type );
 
     /* Import the key */
-    PSA_ASSERT( psa_import_key_to_handle( handle, type,
-                                data->x, data->len ) );
+    PSA_ASSERT( psa_import_key( &attributes, &handle, data->x, data->len ) );
 
     /* Export the public key */
     ASSERT_ALLOC( exported, export_size );
@@ -1572,20 +1564,18 @@
     size_t bits = bits_arg;
     psa_algorithm_t alg = alg_arg;
     psa_key_usage_t usage = usage_to_exercise( type, alg );
-    psa_key_policy_t policy = PSA_KEY_POLICY_INIT;
+    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
     psa_key_type_t got_type;
     size_t got_bits;
-    psa_status_t status;
 
     PSA_ASSERT( psa_crypto_init( ) );
 
-    PSA_ASSERT( psa_allocate_key( &handle ) );
-    psa_key_policy_set_usage( &policy, usage, alg );
-    PSA_ASSERT( psa_set_key_policy( handle, &policy ) );
+    psa_set_key_usage_flags( &attributes, usage );
+    psa_set_key_algorithm( &attributes, alg );
+    psa_set_key_type( &attributes, type );
 
     /* Import the key */
-    status = psa_import_key_to_handle( handle, type, data->x, data->len );
-    PSA_ASSERT( status );
+    PSA_ASSERT( psa_import_key( &attributes, &handle, data->x, data->len ) );
 
     /* Test the key information */
     PSA_ASSERT( psa_get_key_information( handle,