diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index e048e9f..84b6911 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -1469,20 +1469,30 @@
     (void) driver;
 
 #if defined(MBEDTLS_PSA_CRYPTO_STORAGE_C)
-    if( slot->lifetime == PSA_KEY_LIFETIME_PERSISTENT )
+    if( slot->lifetime != PSA_KEY_LIFETIME_VOLATILE )
     {
         uint8_t *buffer = NULL;
         size_t buffer_size = 0;
-        size_t length;
+        size_t length = 0;
 
-        buffer_size = PSA_KEY_EXPORT_MAX_SIZE( slot->type,
-                                               psa_get_key_slot_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 defined(MBEDTLS_PSA_CRYPTO_SE_C)
+        if( driver != NULL )
+        {
+            buffer = (uint8_t*) &slot->data.se.slot_number;
+            length = sizeof( slot->data.se.slot_number );
+        }
+        else
+#endif /* MBEDTLS_PSA_CRYPTO_SE_C */
+        {
+            buffer_size = PSA_KEY_EXPORT_MAX_SIZE( slot->type,
+                                                   psa_get_key_slot_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 )
         {
@@ -1491,9 +1501,14 @@
             status = psa_save_persistent_key( &attributes, buffer, length );
         }
 
-        if( buffer_size != 0 )
-            mbedtls_platform_zeroize( buffer, buffer_size );
-        mbedtls_free( buffer );
+#if defined(MBEDTLS_PSA_CRYPTO_SE_C)
+        if( driver == NULL )
+#endif /* MBEDTLS_PSA_CRYPTO_SE_C */
+        {
+            if( buffer_size != 0 )
+                mbedtls_platform_zeroize( buffer, buffer_size );
+            mbedtls_free( buffer );
+        }
     }
 #endif /* defined(MBEDTLS_PSA_CRYPTO_STORAGE_C) */
 
diff --git a/library/psa_crypto_slot_management.c b/library/psa_crypto_slot_management.c
index 5326fbd..6b87ea0 100644
--- a/library/psa_crypto_slot_management.c
+++ b/library/psa_crypto_slot_management.c
@@ -131,10 +131,28 @@
                                       &key_data, &key_data_length );
     if( status != PSA_SUCCESS )
         goto exit;
+    p_slot->lifetime = psa_get_key_lifetime( &attributes );
     p_slot->type = psa_get_key_type( &attributes );
     p_slot->policy = attributes.policy;
-    status = psa_import_key_into_slot( p_slot,
-                                       key_data, key_data_length );
+
+#if defined(MBEDTLS_PSA_CRYPTO_SE_C)
+    if( psa_key_lifetime_is_external( p_slot->lifetime ) )
+    {
+        if( key_data_length != sizeof( p_slot->data.se.slot_number ) )
+        {
+            status = PSA_ERROR_STORAGE_FAILURE;
+            goto exit;
+        }
+        memcpy( &p_slot->data.se.slot_number, key_data,
+                sizeof( p_slot->data.se.slot_number ) );
+    }
+    else
+#endif /* MBEDTLS_PSA_CRYPTO_SE_C */
+    {
+        status = psa_import_key_into_slot( p_slot,
+                                           key_data, key_data_length );
+    }
+
 exit:
     psa_free_persistent_key_data( key_data, key_data_length );
     return( status );
diff --git a/tests/suites/test_suite_psa_crypto_se_driver_hal.data b/tests/suites/test_suite_psa_crypto_se_driver_hal.data
index 28c7d75..cb21ab5 100644
--- a/tests/suites/test_suite_psa_crypto_se_driver_hal.data
+++ b/tests/suites/test_suite_psa_crypto_se_driver_hal.data
@@ -28,7 +28,13 @@
 register_max:
 
 Key creation smoke test (p_allocate allows all slots)
-key_creation_import_export:0
+key_creation_import_export:0:0
 
 Key creation smoke test (p_allocate allows 1 slot)
-key_creation_import_export:ARRAY_LENGTH( ram_slots ) - 1
+key_creation_import_export:ARRAY_LENGTH( ram_slots ) - 1:0
+
+Key creation smoke test, check after restart (slot 0)
+key_creation_import_export:0:1
+
+Key creation smoke test, check after restart (slot 3)
+key_creation_import_export:3:1
diff --git a/tests/suites/test_suite_psa_crypto_se_driver_hal.function b/tests/suites/test_suite_psa_crypto_se_driver_hal.function
index 2e2a648..5a2ebe7 100644
--- a/tests/suites/test_suite_psa_crypto_se_driver_hal.function
+++ b/tests/suites/test_suite_psa_crypto_se_driver_hal.function
@@ -3,6 +3,7 @@
 #include "psa/crypto_se_driver.h"
 
 #include "psa_crypto_se.h"
+#include "psa_crypto_storage.h"
 
 /** The minimum valid lifetime value for a secure element driver. */
 #define MIN_DRIVER_LIFETIME 2
@@ -115,6 +116,18 @@
     return( PSA_ERROR_INSUFFICIENT_STORAGE );
 }
 
+#define MAX_KEY_ID_FOR_TEST 10
+void psa_purge_storage( void )
+{
+    psa_key_id_t i;
+    /* The tests may have potentially created key ids from 1 to
+     * MAX_KEY_ID_FOR_TEST. In addition, run the destroy function on key id
+     * 0, which file-based storage uses as a temporary file. */
+    for( i = 0; i <= MAX_KEY_ID_FOR_TEST; i++ )
+        psa_destroy_persistent_key( i );
+    psa_crypto_stop_transaction( );
+}
+
 /* END_HEADER */
 
 /* BEGIN_DEPENDENCIES
@@ -188,7 +201,7 @@
 /* END_CASE */
 
 /* BEGIN_CASE */
-void key_creation_import_export( int min_slot )
+void key_creation_import_export( int min_slot, int restart )
 {
     psa_drv_se_t driver;
     psa_drv_se_key_management_t key_management;
@@ -223,6 +236,15 @@
                                 key_material, sizeof( key_material ),
                                 &handle ) );
 
+    /* Maybe restart, to check that the information is saved correctly. */
+    if( restart )
+    {
+        mbedtls_psa_crypto_free( );
+        PSA_ASSERT( psa_register_se_driver( lifetime, &driver ) );
+        PSA_ASSERT( psa_crypto_init( ) );
+        PSA_ASSERT( psa_open_key( id, &handle ) );
+    }
+
     /* Test that the key was created in the expected slot. */
     TEST_ASSERT( ram_slots[min_slot].type == PSA_KEY_TYPE_RAW_DATA );
 
@@ -240,5 +262,6 @@
 exit:
     PSA_DONE( );
     ram_slots_reset( );
+    psa_purge_storage( );
 }
 /* END_CASE */