psa: Call import software implementation as a driver

Signed-off-by: Ronald Cron <ronald.cron@arm.com>
diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 5463a5f..6e80cec 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -622,19 +622,6 @@
     }
     else if( PSA_KEY_TYPE_IS_ASYMMETRIC( type ) )
     {
-        status = psa_driver_wrapper_import_key( attributes,
-                                                data, data_length,
-                                                key_buffer,
-                                                key_buffer_size,
-                                                key_buffer_length,
-                                                bits );
-        if( status != PSA_ERROR_NOT_SUPPORTED )
-            return( status );
-
-        mbedtls_platform_zeroize( key_buffer, key_buffer_size );
-
-        /* Key format is not supported by any accelerator, try software fallback
-         * if present. */
 #if defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_ECC_KEY_PAIR) || \
     defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_ECC_PUBLIC_KEY)
         if( PSA_KEY_TYPE_IS_ECC( type ) )
@@ -1865,16 +1852,23 @@
     }
     else
     {
-        status = psa_allocate_buffer_to_slot( slot, data_length );
-        if( status != PSA_SUCCESS )
-            goto exit;
+       /* In the case of a transparent key or an opaque key stored in local
+        * storage (thus not in the case of generating a key in a secure element
+        * or cryptoprocessor with storage), we have to allocate a buffer to
+        * hold the generated key material. */
+        if( slot->key.data == NULL )
+        {
+            status = psa_allocate_buffer_to_slot( slot, data_length );
+            if( status != PSA_SUCCESS )
+                goto exit;
+        }
 
         size_t bits = slot->attr.bits;
-        status = psa_import_key_into_slot( attributes,
-                                           data, data_length,
-                                           slot->key.data,
-                                           slot->key.bytes,
-                                           &slot->key.bytes, &bits );
+        status = psa_driver_wrapper_import_key( attributes,
+                                                data, data_length,
+                                                slot->key.data,
+                                                slot->key.bytes,
+                                                &slot->key.bytes, &bits );
         if( status != PSA_SUCCESS )
             goto exit;
 
@@ -5191,16 +5185,18 @@
     if( status != PSA_SUCCESS )
         return( status );
 
+    slot->attr.bits = (psa_key_bits_t) bits;
     psa_key_attributes_t attributes = {
       .core = slot->attr
     };
 
-    status = psa_import_key_into_slot( &attributes,
-                                       data, bytes,
-                                       slot->key.data, slot->key.bytes,
-                                       &slot->key.bytes,
-                                       &bits );
-    slot->attr.bits = (psa_key_bits_t) bits;
+    status = psa_driver_wrapper_import_key( &attributes,
+                                            data, bytes,
+                                            slot->key.data,
+                                            slot->key.bytes,
+                                            &slot->key.bytes, &bits );
+    if( bits != slot->attr.bits )
+        status = PSA_ERROR_INVALID_ARGUMENT;
 
 exit:
     mbedtls_free( data );
diff --git a/library/psa_crypto_driver_wrappers.c b/library/psa_crypto_driver_wrappers.c
index 0cc89fe..fe54b7b 100644
--- a/library/psa_crypto_driver_wrappers.c
+++ b/library/psa_crypto_driver_wrappers.c
@@ -418,30 +418,38 @@
     size_t *key_buffer_length,
     size_t *bits )
 {
-#if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
     psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
-    /* Try accelerators in turn */
-#if defined(PSA_CRYPTO_DRIVER_TEST)
-    status = test_transparent_import_key( attributes,
-                                          data, data_length,
-                                          key_buffer, key_buffer_size,
-                                          key_buffer_length, bits );
-    /* Declared with fallback == true */
-    if( status != PSA_ERROR_NOT_SUPPORTED )
-        return( status );
-#endif /* PSA_CRYPTO_DRIVER_TEST */
+    psa_key_location_t location = PSA_KEY_LIFETIME_GET_LOCATION(
+                                      psa_get_key_lifetime( attributes ) );
 
-    return( PSA_ERROR_NOT_SUPPORTED );
-#else /* PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT */
-    (void) attributes;
-    (void) data;
-    (void) data_length;
-    (void) key_buffer;
-    (void) key_buffer_size;
-    (void) key_buffer_length;
-    (void) bits;
-    return( PSA_ERROR_NOT_SUPPORTED );
-#endif /* PSA_CRYPTO_DRIVER_PRESENT */
+    switch( location )
+    {
+        case PSA_KEY_LOCATION_LOCAL_STORAGE:
+            /* Key is stored in the slot in export representation, so
+             * cycle through all known transparent accelerators */
+#if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
+#if defined(PSA_CRYPTO_DRIVER_TEST)
+            status = test_transparent_import_key( attributes,
+                                                  data, data_length,
+                                                  key_buffer, key_buffer_size,
+                                                  key_buffer_length, bits );
+            /* Declared with fallback == true */
+            if( status != PSA_ERROR_NOT_SUPPORTED )
+                return( status );
+#endif /* PSA_CRYPTO_DRIVER_TEST */
+#endif /* PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT */
+            /* Fell through, meaning no accelerator supports this operation */
+            return( psa_import_key_into_slot( attributes,
+                                              data, data_length,
+                                              key_buffer, key_buffer_size,
+                                              key_buffer_length, bits ) );
+
+        default:
+            /* Key is declared with a lifetime not known to us */
+            (void)status;
+            return( PSA_ERROR_NOT_SUPPORTED );
+    }
+
 }
 
 psa_status_t psa_driver_wrapper_export_key(