From 7609b1ff6cea2d25347a395f4579f4ab03f48a3b Mon Sep 17 00:00:00 2001
From: Steven Cooreman <steven.cooreman@silabs.com>
Date: Tue, 6 Apr 2021 16:45:06 +0200
Subject: [PATCH] leverage psa_allocate_buffer_to_slot from slot management

It makes the implementation of psa_load_builtin_key_into_slot a lot
cleaner.

Signed-off-by: Steven Cooreman <steven.cooreman@silabs.com>
---
 library/psa_crypto.c                 | 16 ++--------------
 library/psa_crypto_core.h            | 15 +++++++++++++++
 library/psa_crypto_slot_management.c | 26 ++++++++++++--------------
 3 files changed, 29 insertions(+), 28 deletions(-)

diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 32568b322..068990a7a 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -604,20 +604,8 @@ MBEDTLS_STATIC_TESTABLE psa_status_t psa_mac_key_can_do(
     return( PSA_ERROR_INVALID_ARGUMENT );
 }
 
-/** Try to allocate a buffer to an empty key slot.
- *
- * \param[in,out] slot          Key slot to attach buffer to.
- * \param[in] buffer_length     Requested size of the buffer.
- *
- * \retval #PSA_SUCCESS
- *         The buffer has been successfully allocated.
- * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
- *         Not enough memory was available for allocation.
- * \retval #PSA_ERROR_ALREADY_EXISTS
- *         Trying to allocate a buffer to a non-empty key slot.
- */
-static psa_status_t psa_allocate_buffer_to_slot( psa_key_slot_t *slot,
-                                                 size_t buffer_length )
+psa_status_t psa_allocate_buffer_to_slot( psa_key_slot_t *slot,
+                                          size_t buffer_length )
 {
     if( slot->key.data != NULL )
         return( PSA_ERROR_ALREADY_EXISTS );
diff --git a/library/psa_crypto_core.h b/library/psa_crypto_core.h
index f949c7188..eeb0105e3 100644
--- a/library/psa_crypto_core.h
+++ b/library/psa_crypto_core.h
@@ -180,6 +180,21 @@ static inline psa_key_slot_number_t psa_key_slot_get_slot_number(
  */
 psa_status_t psa_wipe_key_slot( psa_key_slot_t *slot );
 
+/** Try to allocate a buffer to an empty key slot.
+ *
+ * \param[in,out] slot          Key slot to attach buffer to.
+ * \param[in] buffer_length     Requested size of the buffer.
+ *
+ * \retval #PSA_SUCCESS
+ *         The buffer has been successfully allocated.
+ * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
+ *         Not enough memory was available for allocation.
+ * \retval #PSA_ERROR_ALREADY_EXISTS
+ *         Trying to allocate a buffer to a non-empty key slot.
+ */
+psa_status_t psa_allocate_buffer_to_slot( psa_key_slot_t *slot,
+                                          size_t buffer_length );
+
 /** Copy key data (in export format) into an empty key slot.
  *
  * This function assumes that the slot does not contain
diff --git a/library/psa_crypto_slot_management.c b/library/psa_crypto_slot_management.c
index de20fa137..bdb45eed8 100644
--- a/library/psa_crypto_slot_management.c
+++ b/library/psa_crypto_slot_management.c
@@ -283,7 +283,6 @@ static psa_status_t psa_load_builtin_key_into_slot( psa_key_slot_t *slot )
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
     psa_key_lifetime_t lifetime = PSA_KEY_LIFETIME_VOLATILE;
     psa_drv_slot_number_t slot_number = 0;
-    uint8_t *key_buffer = NULL;
     size_t key_buffer_size = 0;
     size_t key_buffer_length = 0;
 
@@ -303,33 +302,32 @@ static psa_status_t psa_load_builtin_key_into_slot( psa_key_slot_t *slot )
     /* Set mapped lifetime on the attributes */
     psa_set_key_lifetime( &attributes, lifetime );
 
-    /* If the key should exist according to the platform, load it through the
-     * driver interface. */
+    /* If the key should exist according to the platform, then ask the driver
+     * what its expected size is. */
     status = psa_driver_wrapper_get_key_buffer_size( &attributes,
                                                      &key_buffer_size );
     if( status != PSA_SUCCESS )
         return( status );
 
-    key_buffer = mbedtls_calloc( 1, key_buffer_size );
-    if( key_buffer == NULL )
-        return( PSA_ERROR_INSUFFICIENT_MEMORY );
+    /* Allocate a buffer of the required size and load the builtin key directly
+     * into the slot buffer. */
+    status = psa_allocate_buffer_to_slot( slot, key_buffer_size );
+    if( status != PSA_SUCCESS )
+        return( status );
 
     status = psa_driver_wrapper_get_builtin_key(
                 slot_number, &attributes,
-                key_buffer, key_buffer_size, &key_buffer_length );
+                slot->key.data, slot->key.bytes, &key_buffer_length );
     if( status != PSA_SUCCESS )
         goto exit;
 
-    status = psa_copy_key_material_into_slot(
-                slot, key_buffer, key_buffer_length );
-    if( status != PSA_SUCCESS )
-        goto exit;
-
-    /* Copy core attributes into the slot on success */
+    /* Copy actual key length and core attributes into the slot on success */
+    slot->key.bytes = key_buffer_length;
     slot->attr = attributes.core;
 
 exit:
-    mbedtls_free( key_buffer );
+    if( status != PSA_SUCCESS )
+        psa_wipe_key_slot( slot );
     return( status );
 }
 #endif /* MBEDTLS_PSA_CRYPTO_BUILTIN_KEYS */