psa: se: Create key context in SE key slots

In key slots containing the description of a key of a
dynamically registered Secure Element (SE), store the
key slot number in a key context as defined in the
PSA driver interface for opaque drivers.

That way transparent key data and slot numbers are
, in a key slot, both stored in a dynamically allocated
buffer. The `data` union in structures of type
psa_key_slot_t to distinguish between the storage of
transparent key data and slot numbers is consequently
not necessary anymore and thus removed.

This alignement of some part of the code dedicated to
dynamically registered SE with the PSA driver interface
specification is done to ease the support of both
dynamically registered and statically defined secure
elements.

Signed-off-by: Ronald Cron <ronald.cron@arm.com>
diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 39144a3..fccb800 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -749,8 +749,8 @@
     }
 
     /* On success, store the allocated export-formatted key. */
-    slot->data.key.data = output;
-    slot->data.key.bytes = data_length;
+    slot->key.data = output;
+    slot->key.bytes = data_length;
 
     return( PSA_SUCCESS );
 }
@@ -983,8 +983,8 @@
     }
 
     /* On success, store the allocated export-formatted key. */
-    slot->data.key.data = output;
-    slot->data.key.bytes = data_length;
+    slot->key.data = output;
+    slot->key.bytes = data_length;
 
     return( PSA_SUCCESS );
 }
@@ -1017,14 +1017,14 @@
 static psa_status_t psa_allocate_buffer_to_slot( psa_key_slot_t *slot,
                                                  size_t buffer_length )
 {
-    if( slot->data.key.data != NULL )
+    if( slot->key.data != NULL )
         return( PSA_ERROR_ALREADY_EXISTS );
 
-    slot->data.key.data = mbedtls_calloc( 1, buffer_length );
-    if( slot->data.key.data == NULL )
+    slot->key.data = mbedtls_calloc( 1, buffer_length );
+    if( slot->key.data == NULL )
         return( PSA_ERROR_INSUFFICIENT_MEMORY );
 
-    slot->data.key.bytes = buffer_length;
+    slot->key.bytes = buffer_length;
     return( PSA_SUCCESS );
 }
 
@@ -1037,7 +1037,7 @@
     if( status != PSA_SUCCESS )
         return( status );
 
-    memcpy( slot->data.key.data, data, data_length );
+    memcpy( slot->key.data, data, data_length );
     return( PSA_SUCCESS );
 }
 
@@ -1342,23 +1342,14 @@
 /** Wipe key data from a slot. Preserve metadata such as the policy. */
 static psa_status_t psa_remove_key_data_from_memory( psa_key_slot_t *slot )
 {
-#if defined(MBEDTLS_PSA_CRYPTO_SE_C)
-    if( psa_get_se_driver( slot->attr.lifetime, NULL, NULL ) &&
-        psa_key_slot_is_external( slot ) )
-    {
-        /* No key material to clean. */
-    }
-    else
-#endif /* MBEDTLS_PSA_CRYPTO_SE_C */
-    {
-        /* Data pointer will always be either a valid pointer or NULL in an
-         * initialized slot, so we can just free it. */
-        if( slot->data.key.data != NULL )
-            mbedtls_platform_zeroize( slot->data.key.data, slot->data.key.bytes);
-        mbedtls_free( slot->data.key.data );
-        slot->data.key.data = NULL;
-        slot->data.key.bytes = 0;
-    }
+    /* Data pointer will always be either a valid pointer or NULL in an
+     * initialized slot, so we can just free it. */
+    if( slot->key.data != NULL )
+        mbedtls_platform_zeroize( slot->key.data, slot->key.bytes);
+
+    mbedtls_free( slot->key.data );
+    slot->key.data = NULL;
+    slot->key.bytes = 0;
 
     return( PSA_SUCCESS );
 }
@@ -1443,7 +1434,7 @@
          * three actions. */
         psa_crypto_prepare_transaction( PSA_CRYPTO_TRANSACTION_DESTROY_KEY );
         psa_crypto_transaction.key.lifetime = slot->attr.lifetime;
-        psa_crypto_transaction.key.slot = slot->data.se.slot_number;
+        psa_crypto_transaction.key.slot = psa_key_slot_get_slot_number( slot );
         psa_crypto_transaction.key.id = slot->attr.id;
         status = psa_crypto_save_transaction( );
         if( status != PSA_SUCCESS )
@@ -1460,7 +1451,8 @@
             goto exit;
         }
 
-        status = psa_destroy_se_key( driver, slot->data.se.slot_number );
+        status = psa_destroy_se_key( driver,
+                                     psa_key_slot_get_slot_number( slot ) );
         if( overall_status == PSA_SUCCESS )
             overall_status = status;
     }
@@ -1616,7 +1608,8 @@
 
 #if defined(MBEDTLS_PSA_CRYPTO_SE_C)
     if( psa_key_slot_is_external( slot ) )
-        psa_set_key_slot_number( attributes, slot->data.se.slot_number );
+        psa_set_key_slot_number( attributes,
+                                 psa_key_slot_get_slot_number( slot ) );
 #endif /* MBEDTLS_PSA_CRYPTO_SE_C */
 
     switch( slot->attr.type )
@@ -1637,8 +1630,8 @@
                 mbedtls_rsa_context *rsa = NULL;
 
                 status = psa_load_rsa_representation( slot->attr.type,
-                                                      slot->data.key.data,
-                                                      slot->data.key.bytes,
+                                                      slot->key.data,
+                                                      slot->key.bytes,
                                                       &rsa );
                 if( status != PSA_SUCCESS )
                     break;
@@ -1684,12 +1677,12 @@
                                                     size_t data_size,
                                                     size_t *data_length )
 {
-    if( slot->data.key.bytes > data_size )
+    if( slot->key.bytes > data_size )
         return( PSA_ERROR_BUFFER_TOO_SMALL );
-    memcpy( data, slot->data.key.data, slot->data.key.bytes );
-    memset( data + slot->data.key.bytes, 0,
-            data_size - slot->data.key.bytes );
-    *data_length = slot->data.key.bytes;
+    memcpy( data, slot->key.data, slot->key.bytes );
+    memset( data + slot->key.bytes, 0,
+            data_size - slot->key.bytes );
+    *data_length = slot->key.bytes;
     return( PSA_SUCCESS );
 }
 
@@ -1727,7 +1720,7 @@
         if( method == NULL )
             return( PSA_ERROR_NOT_SUPPORTED );
         return( method( drv_context,
-                        slot->data.se.slot_number,
+                        psa_key_slot_get_slot_number( slot ),
                         data, data_size, data_length ) );
     }
 #endif /* MBEDTLS_PSA_CRYPTO_SE_C */
@@ -1768,8 +1761,8 @@
             mbedtls_rsa_context *rsa = NULL;
             status = psa_load_rsa_representation(
                                     slot->attr.type,
-                                    slot->data.key.data,
-                                    slot->data.key.bytes,
+                                    slot->key.data,
+                                    slot->key.bytes,
                                     &rsa );
             if( status != PSA_SUCCESS )
                 return( status );
@@ -1797,8 +1790,8 @@
             mbedtls_ecp_keypair *ecp = NULL;
             status = psa_load_ecp_representation(
                                     slot->attr.type,
-                                    slot->data.key.data,
-                                    slot->data.key.bytes,
+                                    slot->key.data,
+                                    slot->key.bytes,
                                     &ecp );
             if( status != PSA_SUCCESS )
                 return( status );
@@ -2071,8 +2064,9 @@
      * we can roll back to a state where the key doesn't exist. */
     if( *p_drv != NULL )
     {
+        psa_key_slot_number_t slot_number;
         status = psa_find_se_slot_for_key( attributes, method, *p_drv,
-                                           &slot->data.se.slot_number );
+                                           &slot_number );
         if( status != PSA_SUCCESS )
             return( status );
 
@@ -2080,7 +2074,7 @@
         {
             psa_crypto_prepare_transaction( PSA_CRYPTO_TRANSACTION_CREATE_KEY );
             psa_crypto_transaction.key.lifetime = slot->attr.lifetime;
-            psa_crypto_transaction.key.slot = slot->data.se.slot_number;
+            psa_crypto_transaction.key.slot = slot_number;
             psa_crypto_transaction.key.id = slot->attr.id;
             status = psa_crypto_save_transaction( );
             if( status != PSA_SUCCESS )
@@ -2089,6 +2083,9 @@
                 return( status );
             }
         }
+
+        status = psa_copy_key_material_into_slot(
+            slot, (uint8_t *)( &slot_number ), sizeof( slot_number ) );
     }
 
     if( *p_drv == NULL && method == PSA_KEY_CREATION_REGISTER )
@@ -2140,13 +2137,15 @@
         if( driver != NULL )
         {
             psa_se_key_data_storage_t data;
+            psa_key_slot_number_t slot_number =
+                psa_key_slot_get_slot_number( slot ) ;
+
 #if defined(static_assert)
-            static_assert( sizeof( slot->data.se.slot_number ) ==
+            static_assert( sizeof( slot_number ) ==
                            sizeof( data.slot_number ),
                            "Slot number size does not match psa_se_key_data_storage_t" );
 #endif
-            memcpy( &data.slot_number, &slot->data.se.slot_number,
-                    sizeof( slot->data.se.slot_number ) );
+            memcpy( &data.slot_number, &slot_number, sizeof( slot_number ) );
             status = psa_save_persistent_key( &slot->attr,
                                               (uint8_t*) &data,
                                               sizeof( data ) );
@@ -2157,8 +2156,8 @@
             /* Key material is saved in export representation in the slot, so
              * just pass the slot buffer for storage. */
             status = psa_save_persistent_key( &slot->attr,
-                                              slot->data.key.data,
-                                              slot->data.key.bytes );
+                                              slot->key.data,
+                                              slot->key.bytes );
         }
     }
 #endif /* defined(MBEDTLS_PSA_CRYPTO_STORAGE_C) */
@@ -2264,8 +2263,8 @@
 
             psa_status_t status = psa_load_rsa_representation(
                                     slot->attr.type,
-                                    slot->data.key.data,
-                                    slot->data.key.bytes,
+                                    slot->key.data,
+                                    slot->key.bytes,
                                     &rsa );
             if( status != PSA_SUCCESS )
                 return( status );
@@ -2345,8 +2344,8 @@
         }
         status = drv->key_management->p_import(
             psa_get_se_driver_context( driver ),
-            slot->data.se.slot_number, attributes, data, data_length,
-            &bits );
+            psa_key_slot_get_slot_number( slot ),
+            attributes, data, data_length, &bits );
         if( status != PSA_SUCCESS )
             goto exit;
         if( bits > PSA_MAX_KEY_BITS )
@@ -2423,8 +2422,8 @@
                                            psa_key_slot_t *target )
 {
     psa_status_t status = psa_copy_key_material_into_slot( target,
-                                                           source->data.key.data,
-                                                           source->data.key.bytes );
+                                                           source->key.data,
+                                                           source->key.bytes );
     if( status != PSA_SUCCESS )
         return( status );
 
@@ -3236,7 +3235,7 @@
         return( ret );
 
     ret = mbedtls_cipher_cmac_starts( &operation->ctx.cmac,
-                                      slot->data.key.data,
+                                      slot->key.data,
                                       key_bits );
     return( ret );
 }
@@ -3382,8 +3381,8 @@
         }
 
         status = psa_hmac_setup_internal( &operation->ctx.hmac,
-                                          slot->data.key.data,
-                                          slot->data.key.bytes,
+                                          slot->key.data,
+                                          slot->key.bytes,
                                           hash_alg );
     }
     else
@@ -3969,8 +3968,8 @@
         mbedtls_rsa_context *rsa = NULL;
 
         status = psa_load_rsa_representation( slot->attr.type,
-                                              slot->data.key.data,
-                                              slot->data.key.bytes,
+                                              slot->key.data,
+                                              slot->key.bytes,
                                               &rsa );
         if( status != PSA_SUCCESS )
             goto exit;
@@ -4001,8 +4000,8 @@
         {
             mbedtls_ecp_keypair *ecp = NULL;
             status = psa_load_ecp_representation( slot->attr.type,
-                                                  slot->data.key.data,
-                                                  slot->data.key.bytes,
+                                                  slot->key.data,
+                                                  slot->key.bytes,
                                                   &ecp );
             if( status != PSA_SUCCESS )
                 goto exit;
@@ -4079,8 +4078,8 @@
         mbedtls_rsa_context *rsa = NULL;
 
         status = psa_load_rsa_representation( slot->attr.type,
-                                              slot->data.key.data,
-                                              slot->data.key.bytes,
+                                              slot->key.data,
+                                              slot->key.bytes,
                                               &rsa );
         if( status != PSA_SUCCESS )
             goto exit;
@@ -4104,8 +4103,8 @@
         {
             mbedtls_ecp_keypair *ecp = NULL;
             status = psa_load_ecp_representation( slot->attr.type,
-                                                  slot->data.key.data,
-                                                  slot->data.key.bytes,
+                                                  slot->key.data,
+                                                  slot->key.bytes,
                                                   &ecp );
             if( status != PSA_SUCCESS )
                 goto exit;
@@ -4188,8 +4187,8 @@
     {
         mbedtls_rsa_context *rsa = NULL;
         status = psa_load_rsa_representation( slot->attr.type,
-                                              slot->data.key.data,
-                                              slot->data.key.bytes,
+                                              slot->key.data,
+                                              slot->key.bytes,
                                               &rsa );
         if( status != PSA_SUCCESS )
             goto rsa_exit;
@@ -4294,8 +4293,8 @@
     {
         mbedtls_rsa_context *rsa = NULL;
         status = psa_load_rsa_representation( slot->attr.type,
-                                              slot->data.key.data,
-                                              slot->data.key.bytes,
+                                              slot->key.data,
+                                              slot->key.bytes,
                                               &rsa );
         if( status != PSA_SUCCESS )
             goto exit;
@@ -4455,8 +4454,8 @@
     {
         /* Two-key Triple-DES is 3-key Triple-DES with K1=K3 */
         uint8_t keys[24];
-        memcpy( keys, slot->data.key.data, 16 );
-        memcpy( keys + 16, slot->data.key.data, 8 );
+        memcpy( keys, slot->key.data, 16 );
+        memcpy( keys + 16, slot->key.data, 8 );
         ret = mbedtls_cipher_setkey( &operation->ctx.cipher,
                                      keys,
                                      192, cipher_operation );
@@ -4465,7 +4464,7 @@
 #endif
     {
         ret = mbedtls_cipher_setkey( &operation->ctx.cipher,
-                                     slot->data.key.data,
+                                     slot->key.data,
                                      (int) key_bits, cipher_operation );
     }
     if( ret != 0 )
@@ -4964,7 +4963,7 @@
             mbedtls_ccm_init( &operation->ctx.ccm );
             status = mbedtls_to_psa_error(
                 mbedtls_ccm_setkey( &operation->ctx.ccm, cipher_id,
-                                    operation->slot->data.key.data,
+                                    operation->slot->key.data,
                                     (unsigned int) key_bits ) );
             if( status != 0 )
                 goto cleanup;
@@ -4986,7 +4985,7 @@
             mbedtls_gcm_init( &operation->ctx.gcm );
             status = mbedtls_to_psa_error(
                 mbedtls_gcm_setkey( &operation->ctx.gcm, cipher_id,
-                                    operation->slot->data.key.data,
+                                    operation->slot->key.data,
                                     (unsigned int) key_bits ) );
             if( status != 0 )
                 goto cleanup;
@@ -5006,7 +5005,7 @@
             mbedtls_chachapoly_init( &operation->ctx.chachapoly );
             status = mbedtls_to_psa_error(
                 mbedtls_chachapoly_setkey( &operation->ctx.chachapoly,
-                                           operation->slot->data.key.data ) );
+                                           operation->slot->key.data ) );
             if( status != 0 )
                 goto cleanup;
             break;
@@ -6129,8 +6128,8 @@
 
     status = psa_key_derivation_input_internal( operation,
                                                 step, slot->attr.type,
-                                                slot->data.key.data,
-                                                slot->data.key.bytes );
+                                                slot->key.data,
+                                                slot->key.bytes );
 
     unlock_status = psa_unlock_key_slot( slot );
 
@@ -6215,8 +6214,8 @@
             mbedtls_ecp_keypair *ecp = NULL;
             psa_status_t status = psa_load_ecp_representation(
                                     private_key->attr.type,
-                                    private_key->data.key.data,
-                                    private_key->data.key.bytes,
+                                    private_key->key.data,
+                                    private_key->key.bytes,
                                     &ecp );
             if( status != PSA_SUCCESS )
                 return( status );
@@ -6558,16 +6557,16 @@
         if( status != PSA_SUCCESS )
             return( status );
 
-        status = psa_generate_random( slot->data.key.data,
-                                      slot->data.key.bytes );
+        status = psa_generate_random( slot->key.data,
+                                      slot->key.bytes );
         if( status != PSA_SUCCESS )
             return( status );
 
         slot->attr.bits = (psa_key_bits_t) bits;
 #if defined(MBEDTLS_DES_C)
         if( type == PSA_KEY_TYPE_DES )
-            psa_des_set_key_parity( slot->data.key.data,
-                                    slot->data.key.bytes );
+            psa_des_set_key_parity( slot->key.data,
+                                    slot->key.bytes );
 #endif /* MBEDTLS_DES_C */
     }
     else
@@ -6611,9 +6610,9 @@
 
         status = psa_export_rsa_key( type,
                                      &rsa,
-                                     slot->data.key.data,
+                                     slot->key.data,
                                      bytes,
-                                     &slot->data.key.bytes );
+                                     &slot->key.bytes );
         mbedtls_rsa_free( &rsa );
         if( status != PSA_SUCCESS )
             psa_remove_key_data_from_memory( slot );
@@ -6657,11 +6656,11 @@
         }
 
         status = mbedtls_to_psa_error(
-            mbedtls_ecp_write_key( &ecp, slot->data.key.data, bytes ) );
+            mbedtls_ecp_write_key( &ecp, slot->key.data, bytes ) );
 
         mbedtls_ecp_keypair_free( &ecp );
         if( status != PSA_SUCCESS ) {
-            memset( slot->data.key.data, 0, bytes );
+            memset( slot->key.data, 0, bytes );
             psa_remove_key_data_from_memory( slot );
         }
         return( status );
diff --git a/library/psa_crypto_core.h b/library/psa_crypto_core.h
index f61ef95..3f266f3 100644
--- a/library/psa_crypto_core.h
+++ b/library/psa_crypto_core.h
@@ -62,23 +62,13 @@
      */
     size_t lock_count;
 
-    union
+    /* Dynamically allocated key data buffer.
+     * Format as specified in psa_export_key(). */
+    struct key_data
     {
-        /* Dynamically allocated key data buffer.
-         * Format as specified in psa_export_key(). */
-        struct key_data
-        {
-            uint8_t *data;
-            size_t bytes;
-        } key;
-#if defined(MBEDTLS_PSA_CRYPTO_SE_C)
-        /* Any key type in a secure element */
-        struct se
-        {
-            psa_key_slot_number_t slot_number;
-        } se;
-#endif /* MBEDTLS_PSA_CRYPTO_SE_C */
-    } data;
+        uint8_t *data;
+        size_t bytes;
+    } key;
 } psa_key_slot_t;
 
 /* A mask of key attribute flags used only internally.
@@ -163,6 +153,20 @@
     slot->attr.flags &= ~mask;
 }
 
+#if defined(MBEDTLS_PSA_CRYPTO_SE_C)
+/** Get the SE slot number of a key from the key slot storing its description.
+ *
+ * \param[in]  slot  The key slot to query. This must be a key slot storing
+ *                   the description of a key of a dynamically registered
+ *                   secure element, otherwise the behaviour is undefined.
+ */
+static inline psa_key_slot_number_t psa_key_slot_get_slot_number(
+    const psa_key_slot_t *slot )
+{
+    return( *( (psa_key_slot_number_t *)( slot->key.data ) ) );
+}
+#endif
+
 /** Completely wipe a slot in memory, including its policy.
  *
  * Persistent storage is not affected.
diff --git a/library/psa_crypto_driver_wrappers.c b/library/psa_crypto_driver_wrappers.c
index 1243bd3..2d43383 100644
--- a/library/psa_crypto_driver_wrappers.c
+++ b/library/psa_crypto_driver_wrappers.c
@@ -80,7 +80,7 @@
             return( PSA_ERROR_NOT_SUPPORTED );
         }
         return( drv->asymmetric->p_sign( drv_context,
-                                         slot->data.se.slot_number,
+                                         psa_key_slot_get_slot_number( slot ),
                                          alg,
                                          hash, hash_length,
                                          signature, signature_size,
@@ -103,8 +103,8 @@
              * cycle through all known transparent accelerators */
 #if defined(PSA_CRYPTO_DRIVER_TEST)
             status = test_transparent_signature_sign_hash( &attributes,
-                                                           slot->data.key.data,
-                                                           slot->data.key.bytes,
+                                                           slot->key.data,
+                                                           slot->key.bytes,
                                                            alg,
                                                            hash,
                                                            hash_length,
@@ -121,8 +121,8 @@
 #if defined(PSA_CRYPTO_DRIVER_TEST)
         case PSA_CRYPTO_TEST_DRIVER_LIFETIME:
             return( test_opaque_signature_sign_hash( &attributes,
-                                                     slot->data.key.data,
-                                                     slot->data.key.bytes,
+                                                     slot->key.data,
+                                                     slot->key.bytes,
                                                      alg,
                                                      hash,
                                                      hash_length,
@@ -172,7 +172,7 @@
             return( PSA_ERROR_NOT_SUPPORTED );
         }
         return( drv->asymmetric->p_verify( drv_context,
-                                           slot->data.se.slot_number,
+                                           psa_key_slot_get_slot_number( slot ),
                                            alg,
                                            hash, hash_length,
                                            signature, signature_length ) );
@@ -194,8 +194,8 @@
              * cycle through all known transparent accelerators */
 #if defined(PSA_CRYPTO_DRIVER_TEST)
             status = test_transparent_signature_verify_hash( &attributes,
-                                                             slot->data.key.data,
-                                                             slot->data.key.bytes,
+                                                             slot->key.data,
+                                                             slot->key.bytes,
                                                              alg,
                                                              hash,
                                                              hash_length,
@@ -211,8 +211,8 @@
 #if defined(PSA_CRYPTO_DRIVER_TEST)
         case PSA_CRYPTO_TEST_DRIVER_LIFETIME:
             return( test_opaque_signature_verify_hash( &attributes,
-                                                       slot->data.key.data,
-                                                       slot->data.key.bytes,
+                                                       slot->key.data,
+                                                       slot->key.bytes,
                                                        alg,
                                                        hash,
                                                        hash_length,
@@ -330,9 +330,8 @@
             return( PSA_ERROR_NOT_SUPPORTED );
         }
         return( drv->key_management->p_generate(
-            drv_context,
-            slot->data.se.slot_number, attributes,
-            NULL, 0, &pubkey_length ) );
+            drv_context, psa_key_slot_get_slot_number( slot ),
+            attributes, NULL, 0, &pubkey_length ) );
     }
 #endif /* MBEDTLS_PSA_CRYPTO_SE_C */
 
@@ -346,10 +345,10 @@
     if( status != PSA_SUCCESS )
         return( status );
 
-    slot->data.key.data = mbedtls_calloc(1, export_size);
-    if( slot->data.key.data == NULL )
+    slot->key.data = mbedtls_calloc(1, export_size);
+    if( slot->key.data == NULL )
         return( PSA_ERROR_INSUFFICIENT_MEMORY );
-    slot->data.key.bytes = export_size;
+    slot->key.bytes = export_size;
 
     switch( location )
     {
@@ -365,9 +364,9 @@
             }
 #if defined(PSA_CRYPTO_DRIVER_TEST)
             status = test_transparent_generate_key( attributes,
-                                                    slot->data.key.data,
-                                                    slot->data.key.bytes,
-                                                    &slot->data.key.bytes );
+                                                    slot->key.data,
+                                                    slot->key.bytes,
+                                                    &slot->key.bytes );
             /* Declared with fallback == true */
             if( status != PSA_ERROR_NOT_SUPPORTED )
                 break;
@@ -379,9 +378,9 @@
 #if defined(PSA_CRYPTO_DRIVER_TEST)
         case PSA_CRYPTO_TEST_DRIVER_LIFETIME:
             status = test_opaque_generate_key( attributes,
-                                               slot->data.key.data,
-                                               slot->data.key.bytes,
-                                               &slot->data.key.bytes );
+                                               slot->key.data,
+                                               slot->key.bytes,
+                                               &slot->key.bytes );
             break;
 #endif /* PSA_CRYPTO_DRIVER_TEST */
         default:
@@ -393,9 +392,9 @@
     if( status != PSA_SUCCESS )
     {
         /* free allocated buffer */
-        mbedtls_free( slot->data.key.data );
-        slot->data.key.data = NULL;
-        slot->data.key.bytes = 0;
+        mbedtls_free( slot->key.data );
+        slot->key.data = NULL;
+        slot->key.bytes = 0;
     }
 
     return( status );
@@ -457,8 +456,8 @@
              * cycle through all known transparent accelerators */
 #if defined(PSA_CRYPTO_DRIVER_TEST)
             status = test_transparent_export_public_key( &attributes,
-                                                         slot->data.key.data,
-                                                         slot->data.key.bytes,
+                                                         slot->key.data,
+                                                         slot->key.bytes,
                                                          data,
                                                          data_size,
                                                          data_length );
@@ -472,8 +471,8 @@
 #if defined(PSA_CRYPTO_DRIVER_TEST)
         case PSA_CRYPTO_TEST_DRIVER_LIFETIME:
             return( test_opaque_export_public_key( &attributes,
-                                                   slot->data.key.data,
-                                                   slot->data.key.bytes,
+                                                   slot->key.data,
+                                                   slot->key.bytes,
                                                    data,
                                                    data_size,
                                                    data_length ) );
@@ -517,8 +516,8 @@
              * cycle through all known transparent accelerators */
 #if defined(PSA_CRYPTO_DRIVER_TEST)
             status = test_transparent_cipher_encrypt( &attributes,
-                                                      slot->data.key.data,
-                                                      slot->data.key.bytes,
+                                                      slot->key.data,
+                                                      slot->key.bytes,
                                                       alg,
                                                       input,
                                                       input_length,
@@ -535,8 +534,8 @@
 #if defined(PSA_CRYPTO_DRIVER_TEST)
         case PSA_CRYPTO_TEST_DRIVER_LIFETIME:
             return( test_opaque_cipher_encrypt( &attributes,
-                                                slot->data.key.data,
-                                                slot->data.key.bytes,
+                                                slot->key.data,
+                                                slot->key.bytes,
                                                 alg,
                                                 input,
                                                 input_length,
@@ -584,8 +583,8 @@
              * cycle through all known transparent accelerators */
 #if defined(PSA_CRYPTO_DRIVER_TEST)
             status = test_transparent_cipher_decrypt( &attributes,
-                                                      slot->data.key.data,
-                                                      slot->data.key.bytes,
+                                                      slot->key.data,
+                                                      slot->key.bytes,
                                                       alg,
                                                       input,
                                                       input_length,
@@ -602,8 +601,8 @@
 #if defined(PSA_CRYPTO_DRIVER_TEST)
         case PSA_CRYPTO_TEST_DRIVER_LIFETIME:
             return( test_opaque_cipher_decrypt( &attributes,
-                                                slot->data.key.data,
-                                                slot->data.key.bytes,
+                                                slot->key.data,
+                                                slot->key.bytes,
                                                 alg,
                                                 input,
                                                 input_length,
@@ -652,8 +651,8 @@
 
             status = test_transparent_cipher_encrypt_setup( operation->ctx,
                                                             &attributes,
-                                                            slot->data.key.data,
-                                                            slot->data.key.bytes,
+                                                            slot->key.data,
+                                                            slot->key.bytes,
                                                             alg );
             /* Declared with fallback == true */
             if( status == PSA_SUCCESS )
@@ -680,8 +679,8 @@
 
             status = test_opaque_cipher_encrypt_setup( operation->ctx,
                                                        &attributes,
-                                                       slot->data.key.data,
-                                                       slot->data.key.bytes,
+                                                       slot->key.data,
+                                                       slot->key.bytes,
                                                        alg );
             if( status == PSA_SUCCESS )
                 operation->id = PSA_CRYPTO_OPAQUE_TEST_DRIVER_ID;
@@ -733,8 +732,8 @@
 
             status = test_transparent_cipher_decrypt_setup( operation->ctx,
                                                             &attributes,
-                                                            slot->data.key.data,
-                                                            slot->data.key.bytes,
+                                                            slot->key.data,
+                                                            slot->key.bytes,
                                                             alg );
             /* Declared with fallback == true */
             if( status == PSA_SUCCESS )
@@ -761,8 +760,8 @@
 
             status = test_opaque_cipher_decrypt_setup( operation->ctx,
                                                        &attributes,
-                                                       slot->data.key.data,
-                                                       slot->data.key.bytes,
+                                                       slot->key.data,
+                                                       slot->key.bytes,
                                                        alg );
             if( status == PSA_SUCCESS )
                 operation->id = PSA_CRYPTO_OPAQUE_TEST_DRIVER_ID;
diff --git a/library/psa_crypto_slot_management.c b/library/psa_crypto_slot_management.c
index 39d6dbb..1188ce6 100644
--- a/library/psa_crypto_slot_management.c
+++ b/library/psa_crypto_slot_management.c
@@ -255,16 +255,15 @@
     if( psa_get_se_driver( slot->attr.lifetime, &drv, &drv_context ) )
     {
         psa_se_key_data_storage_t *data;
+
         if( key_data_length != sizeof( *data ) )
         {
             status = PSA_ERROR_STORAGE_FAILURE;
             goto exit;
         }
         data = (psa_se_key_data_storage_t *) key_data;
-        memcpy( &slot->data.se.slot_number, &data->slot_number,
-                sizeof( slot->data.se.slot_number ) );
-
-        status = PSA_SUCCESS;
+        status = psa_copy_key_material_into_slot(
+                     slot, data->slot_number, sizeof( data->slot_number ) );
         goto exit;
     }
 #endif /* MBEDTLS_PSA_CRYPTO_SE_C */