diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index c35b2a644..adc730b57 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -1758,10 +1758,12 @@ static psa_status_t psa_internal_export_key( const psa_key_slot_t *slot,
 
         /* Need to export the public part of a private key,
          * so conversion is needed. Try the accelerators first. */
-        psa_status_t status = psa_driver_wrapper_export_public_key( slot,
-                                                                    data,
-                                                                    data_size,
-                                                                    data_length );
+        psa_key_attributes_t attributes = {
+          .core = slot->attr
+        };
+        psa_status_t status = psa_driver_wrapper_export_public_key(
+            &attributes, slot->key.data, slot->key.bytes,
+            data, data_size, data_length );
 
         if( status != PSA_ERROR_NOT_SUPPORTED ||
             psa_key_lifetime_is_external( slot->attr.lifetime ) )
diff --git a/library/psa_crypto_driver_wrappers.c b/library/psa_crypto_driver_wrappers.c
index 0562756e2..a58967a1d 100644
--- a/library/psa_crypto_driver_wrappers.c
+++ b/library/psa_crypto_driver_wrappers.c
@@ -444,17 +444,16 @@ psa_status_t psa_driver_wrapper_import_key(
 #endif /* PSA_CRYPTO_DRIVER_PRESENT */
 }
 
-psa_status_t psa_driver_wrapper_export_public_key( const psa_key_slot_t *slot,
-                                                   uint8_t *data,
-                                                   size_t data_size,
-                                                   size_t *data_length )
+psa_status_t psa_driver_wrapper_export_public_key(
+    const psa_key_attributes_t *attributes,
+    const uint8_t *key_buffer, size_t key_buffer_size,
+    uint8_t *data, size_t data_size, size_t *data_length )
+
 {
 #if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
     psa_status_t status = PSA_ERROR_INVALID_ARGUMENT;
-    psa_key_location_t location = PSA_KEY_LIFETIME_GET_LOCATION(slot->attr.lifetime);
-    psa_key_attributes_t attributes = {
-      .core = slot->attr
-    };
+    psa_key_location_t location = PSA_KEY_LIFETIME_GET_LOCATION(
+                                      psa_get_key_lifetime( attributes ) );
 
     switch( location )
     {
@@ -462,9 +461,9 @@ psa_status_t psa_driver_wrapper_export_public_key( const psa_key_slot_t *slot,
             /* Key is stored in the slot in export representation, so
              * cycle through all known transparent accelerators */
 #if defined(PSA_CRYPTO_DRIVER_TEST)
-            status = test_transparent_export_public_key( &attributes,
-                                                         slot->key.data,
-                                                         slot->key.bytes,
+            status = test_transparent_export_public_key( attributes,
+                                                         key_buffer,
+                                                         key_buffer_size,
                                                          data,
                                                          data_size,
                                                          data_length );
@@ -477,9 +476,9 @@ psa_status_t psa_driver_wrapper_export_public_key( const psa_key_slot_t *slot,
         /* Add cases for opaque driver here */
 #if defined(PSA_CRYPTO_DRIVER_TEST)
         case PSA_CRYPTO_TEST_DRIVER_LIFETIME:
-            return( test_opaque_export_public_key( &attributes,
-                                                   slot->key.data,
-                                                   slot->key.bytes,
+            return( test_opaque_export_public_key( attributes,
+                                                   key_buffer,
+                                                   key_buffer_size,
                                                    data,
                                                    data_size,
                                                    data_length ) );
@@ -489,7 +488,9 @@ psa_status_t psa_driver_wrapper_export_public_key( const psa_key_slot_t *slot,
             return( status );
     }
 #else /* PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT */
-    (void) slot;
+    (void) attributes;
+    (void) key_buffer;
+    (void) key_buffer_size;
     (void) data;
     (void) data_size;
     (void) data_length;
diff --git a/library/psa_crypto_driver_wrappers.h b/library/psa_crypto_driver_wrappers.h
index 4c6cce95b..fb32b8d14 100644
--- a/library/psa_crypto_driver_wrappers.h
+++ b/library/psa_crypto_driver_wrappers.h
@@ -47,19 +47,20 @@ psa_status_t psa_driver_wrapper_verify_hash( psa_key_slot_t *slot,
  * Key handling functions
  */
 
-psa_status_t psa_driver_wrapper_generate_key( const psa_key_attributes_t *attributes,
-                                              psa_key_slot_t *slot );
-
 psa_status_t psa_driver_wrapper_import_key(
     const psa_key_attributes_t *attributes,
     const uint8_t *data, size_t data_length,
     uint8_t *key_buffer, size_t key_buffer_size,
     size_t *key_buffer_length, size_t *bits );
 
-psa_status_t psa_driver_wrapper_export_public_key( const psa_key_slot_t *slot,
-                                                   uint8_t *data,
-                                                   size_t data_size,
-                                                   size_t *data_length );
+psa_status_t psa_driver_wrapper_export_public_key(
+    const psa_key_attributes_t *attributes,
+    const uint8_t *key_buffer, size_t key_buffer_size,
+    uint8_t *data, size_t data_size, size_t *data_length );
+
+psa_status_t psa_driver_wrapper_generate_key(
+    const psa_key_attributes_t *attributes,
+    psa_key_slot_t *slot );
 
 /*
  * Cipher functions