diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 5323bb3..960c698 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -541,13 +541,19 @@
  * \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 @@
      * 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 @@
     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) */