From 8f813ee8d82a740ebcefc8913060c02eec22495c Mon Sep 17 00:00:00 2001
From: Ronald Cron <ronald.cron@arm.com>
Date: Sat, 7 Nov 2020 18:22:09 +0100
Subject: [PATCH] psa: rsa: import: Move key buffer allocation

Move key buffer allocation from RSA specific
importation function up to psa_import_key_into_slot().

Signed-off-by: Ronald Cron <ronald.cron@arm.com>
---
 library/psa_crypto.c | 46 ++++++++++++++++++++------------------------
 1 file changed, 21 insertions(+), 25 deletions(-)

diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 5323bb3d4..960c69824 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -541,13 +541,19 @@ static psa_status_t validate_unstructured_key_bit_size( psa_key_type_t type,
  * \param[in,out] slot      The slot where to store the export representation to
  * \param[in] data          The buffer containing the import representation
  * \param[in] data_length   The amount of bytes in \p data
+ * \param[out] key_buffer   The buffer containing the export representation
+ * \param[in] key_buffer_size     The size of \p key_buffer in bytes
+ * \param[out] key_buffer_length  The length of the data written in the key
+ *                                buffer in bytes.
  */
 static psa_status_t psa_import_rsa_key( 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;
-    uint8_t* output = NULL;
     mbedtls_rsa_context *rsa = NULL;
 
     /* Parse input */
@@ -565,35 +571,17 @@ static psa_status_t psa_import_rsa_key( psa_key_slot_t *slot,
      * representation in the key slot. Export representation in case of RSA is
      * the smallest representation that's allowed as input, so a straight-up
      * allocation of the same size as the input buffer will be large enough. */
-    output = mbedtls_calloc( 1, data_length );
-    if( output == NULL )
-    {
-        status = PSA_ERROR_INSUFFICIENT_MEMORY;
-        goto exit;
-    }
-
     status = mbedtls_psa_rsa_export_key( slot->attr.type,
                                          rsa,
-                                         output,
-                                         data_length,
-                                         &data_length);
+                                         key_buffer,
+                                         key_buffer_size,
+                                         key_buffer_length );
 exit:
     /* Always free the RSA object */
     mbedtls_rsa_free( rsa );
     mbedtls_free( rsa );
 
-    /* Free the allocated buffer only on error. */
-    if( status != PSA_SUCCESS )
-    {
-        mbedtls_free( output );
-        return( status );
-    }
-
-    /* On success, store the allocated export-formatted key. */
-    slot->key.data = output;
-    slot->key.bytes = data_length;
-
-    return( PSA_SUCCESS );
+    return( status );
 }
 
 #endif /* defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_RSA_KEY_PAIR) ||
@@ -827,7 +815,15 @@ static psa_status_t psa_import_key_into_slot( psa_key_slot_t *slot,
     defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_RSA_PUBLIC_KEY)
         if( PSA_KEY_TYPE_IS_RSA( slot->attr.type ) )
         {
-            return( psa_import_rsa_key( slot, data, data_length ) );
+            status = psa_allocate_buffer_to_slot( slot, data_length );
+            if( status != PSA_SUCCESS )
+                return( status );
+
+            status = psa_import_rsa_key( slot,
+                                         data, data_length,
+                                         slot->key.data, data_length,
+                                         &slot->key.bytes );
+            return( status );
         }
 #endif /* defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_RSA_KEY_PAIR) ||
         * defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_RSA_PUBLIC_KEY) */