psa: import: Move key buffer allocation

Move key buffer allocation from psa_import_key_into_slot()
function up to the two functions calling it.

Signed-off-by: Ronald Cron <ronald.cron@arm.com>
diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index d55599e..411010c 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -591,12 +591,18 @@
  *
  * Persistent storage is not affected.
  *
- * \param[in,out] slot  The key slot to import data into.
- *                      Its `type` field must have previously been set to
- *                      the desired key type.
- *                      It must not contain any key material yet.
- * \param[in] data      Buffer containing the key material to parse and import.
- * \param data_length   Size of \p data in bytes.
+ * \param[in,out] slot     The key slot to import data into.
+ *                         Its `type` field must have previously been set to
+ *                         the desired key type.
+ *                         It must not contain any key material yet.
+ * \param[in] data         Buffer containing the key material to parse and
+ *                         import.
+ * \param data_length      Size of \p data in bytes.
+ * \param[out] key_buffer  The buffer containing the export representation.
+ * \param[in]  key_buffer_size    The size of \p key_buffer in bytes. The size
+ *                                is greater or equal to \p data_length.
+ * \param[out] key_buffer_length  The length of the data written in \p
+ *                                key_buffer in bytes.
  *
  * \retval #PSA_SUCCESS
  * \retval #PSA_ERROR_INVALID_ARGUMENT
@@ -605,7 +611,10 @@
  */
 static psa_status_t psa_import_key_into_slot( psa_key_slot_t *slot,
                                               const uint8_t *data,
-                                              size_t data_length )
+                                              size_t data_length,
+                                              uint8_t *key_buffer,
+                                              size_t key_buffer_size,
+                                              size_t *key_buffer_length )
 {
     psa_status_t status = PSA_SUCCESS;
     size_t bit_size;
@@ -631,10 +640,10 @@
         if( status != PSA_SUCCESS )
             return( status );
 
-        /* Allocate memory for the key */
-        status = psa_copy_key_material_into_slot( slot, data, data_length );
-        if( status != PSA_SUCCESS )
-            return( status );
+        /* Copy the key material. */
+        memcpy( key_buffer, data, data_length );
+        *key_buffer_length = data_length;
+        (void)key_buffer_size;
 
         /* Write the actual key size to the slot.
          * psa_start_key_creation() wrote the size declared by the
@@ -650,16 +659,12 @@
           .core = slot->attr
         };
 
-        status = psa_allocate_buffer_to_slot( slot, data_length );
-        if( status != PSA_SUCCESS )
-            return( status );
-
         bit_size = slot->attr.bits;
         status = psa_driver_wrapper_import_key( &attributes,
                                                 data, data_length,
-                                                slot->key.data,
-                                                slot->key.bytes,
-                                                &slot->key.bytes,
+                                                key_buffer,
+                                                key_buffer_size,
+                                                key_buffer_length,
                                                 &bit_size );
         if( status == PSA_SUCCESS )
         {
@@ -676,10 +681,7 @@
                 return( status );
         }
 
-        mbedtls_platform_zeroize( slot->key.data, data_length );
-        mbedtls_free( slot->key.data );
-        slot->key.data = NULL;
-        slot->key.bytes = 0;
+        mbedtls_platform_zeroize( key_buffer, key_buffer_size );
 
         /* Key format is not supported by any accelerator, try software fallback
          * if present. */
@@ -687,14 +689,10 @@
     defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_ECC_PUBLIC_KEY)
         if( PSA_KEY_TYPE_IS_ECC( slot->attr.type ) )
         {
-            status = psa_allocate_buffer_to_slot( slot, data_length );
-            if( status != PSA_SUCCESS )
-                return( status );
-
             status = mbedtls_psa_ecp_import_key( &attributes,
                                                  data, data_length,
-                                                 slot->key.data, data_length,
-                                                 &slot->key.bytes,
+                                                 key_buffer, key_buffer_size,
+                                                 key_buffer_length,
                                                  &bit_size );
             slot->attr.bits = (psa_key_bits_t) bit_size;
             return( status );
@@ -705,14 +703,10 @@
     defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_RSA_PUBLIC_KEY)
         if( PSA_KEY_TYPE_IS_RSA( slot->attr.type ) )
         {
-            status = psa_allocate_buffer_to_slot( slot, data_length );
-            if( status != PSA_SUCCESS )
-                return( status );
-
             status = mbedtls_psa_rsa_import_key( &attributes,
                                                  data, data_length,
-                                                 slot->key.data, data_length,
-                                                 &slot->key.bytes,
+                                                 key_buffer, key_buffer_size,
+                                                 key_buffer_length,
                                                  &bit_size );
             slot->attr.bits = (psa_key_bits_t) bit_size;
             return( status );
@@ -1931,7 +1925,14 @@
     }
     else
     {
-        status = psa_import_key_into_slot( slot, data, data_length );
+        status = psa_allocate_buffer_to_slot( slot, data_length );
+        if( status != PSA_SUCCESS )
+            goto exit;
+
+        status = psa_import_key_into_slot( slot, data, data_length,
+                                           slot->key.data,
+                                           slot->key.bytes,
+                                           &slot->key.bytes );
         if( status != PSA_SUCCESS )
             goto exit;
     }
@@ -5234,7 +5235,14 @@
     if( slot->attr.type == PSA_KEY_TYPE_DES )
         psa_des_set_key_parity( data, bytes );
 #endif /* MBEDTLS_DES_C */
-    status = psa_import_key_into_slot( slot, data, bytes );
+
+    status = psa_allocate_buffer_to_slot( slot, bytes );
+    if( status != PSA_SUCCESS )
+        return( status );
+
+    status = psa_import_key_into_slot( slot, data, bytes,
+                                       slot->key.data, slot->key.bytes,
+                                       &slot->key.bytes );
 
 exit:
     mbedtls_free( data );