Implement support for MBEDTLS_PSA_CRYPTO_BUILTIN_KEYS

According to the design in psa-driver-interface.md. Compiles without
issue in test_psa_crypto_drivers.

Signed-off-by: Steven Cooreman <steven.cooreman@silabs.com>
diff --git a/library/psa_crypto_driver_wrappers.c b/library/psa_crypto_driver_wrappers.c
index 536505e..70c3026 100644
--- a/library/psa_crypto_driver_wrappers.c
+++ b/library/psa_crypto_driver_wrappers.c
@@ -574,6 +574,23 @@
     }
 }
 
+psa_status_t psa_driver_wrapper_get_builtin_key(
+    psa_drv_slot_number_t slot_number,
+    psa_key_attributes_t *attributes,
+    uint8_t *key_buffer, size_t key_buffer_size, size_t *key_buffer_length )
+{
+    psa_key_location_t location = PSA_KEY_LIFETIME_GET_LOCATION( attributes->core.lifetime );
+    switch( location )
+    {
+        default:
+            (void) slot_number;
+            (void) key_buffer;
+            (void) key_buffer_size;
+            (void) key_buffer_length;
+            return( PSA_ERROR_DOES_NOT_EXIST );
+    }
+}
+
 /*
  * Cipher functions
  */
diff --git a/library/psa_crypto_driver_wrappers.h b/library/psa_crypto_driver_wrappers.h
index e499411..e82d093 100644
--- a/library/psa_crypto_driver_wrappers.h
+++ b/library/psa_crypto_driver_wrappers.h
@@ -68,6 +68,11 @@
     const psa_key_attributes_t *attributes,
     uint8_t *key_buffer, size_t key_buffer_size, size_t *key_buffer_length );
 
+psa_status_t psa_driver_wrapper_get_builtin_key(
+    psa_drv_slot_number_t slot_number,
+    psa_key_attributes_t *attributes,
+    uint8_t *key_buffer, size_t key_buffer_size, size_t *key_buffer_length );
+
 /*
  * Cipher functions
  */
diff --git a/library/psa_crypto_slot_management.c b/library/psa_crypto_slot_management.c
index cf07a36..c90ebee 100644
--- a/library/psa_crypto_slot_management.c
+++ b/library/psa_crypto_slot_management.c
@@ -274,6 +274,67 @@
 }
 #endif /* MBEDTLS_PSA_CRYPTO_STORAGE_C */
 
+#if defined(MBEDTLS_PSA_CRYPTO_BUILTIN_KEYS)
+#include "psa_crypto_driver_wrappers.h"
+
+static psa_status_t psa_load_builtin_key_into_slot( psa_key_slot_t *slot )
+{
+    /* Load keys in the 'builtin' range through their own interface */
+    if( psa_key_id_is_builtin( MBEDTLS_SVC_KEY_ID_GET_KEY_ID( slot->attr.id ) ) )
+    {
+        /* Check the platform function to see whether this key actually exists */
+        psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
+        psa_drv_slot_number_t slot_number;
+
+        psa_set_key_id(&attributes, slot->attr.id);
+        psa_status_t status = mbedtls_psa_platform_get_builtin_key(
+                                &attributes, &slot_number );
+        if( status != PSA_SUCCESS )
+            return( status );
+
+        /* If the key should exist according to the platform, load it through
+         * the driver interface. */
+        uint8_t *key_buffer = NULL;
+        size_t key_buffer_length = 0;
+
+        status = psa_driver_wrapper_get_key_buffer_size( &attributes, &key_buffer_length );
+        if( status != PSA_SUCCESS )
+            return( status );
+
+        key_buffer = mbedtls_calloc( 1, key_buffer_length );
+        if( key_buffer == NULL )
+            return( PSA_ERROR_INSUFFICIENT_MEMORY );
+
+        status = psa_driver_wrapper_get_builtin_key(
+                    slot_number, &attributes,
+                    key_buffer, key_buffer_length, &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.
+         * Use static allocations to make the compiler yell at us should one
+         * of the two structures change type. */
+        psa_core_key_attributes_t* builtin_key_core_attributes =
+            &attributes.core;
+        psa_core_key_attributes_t* slot_core_attributes =
+            &slot->attr;
+        memcpy( slot_core_attributes,
+                builtin_key_core_attributes,
+                sizeof(psa_core_key_attributes_t) );
+
+exit:
+        mbedtls_free( key_buffer );
+        return( status );
+    } else {
+        return( PSA_ERROR_DOES_NOT_EXIST );
+    }
+}
+#endif /* MBEDTLS_PSA_CRYPTO_BUILTIN_KEYS */
+
 psa_status_t psa_get_and_lock_key_slot( mbedtls_svc_key_id_t key,
                                         psa_key_slot_t **p_slot )
 {
@@ -291,17 +352,27 @@
     if( status != PSA_ERROR_DOES_NOT_EXIST )
         return( status );
 
-#if defined(MBEDTLS_PSA_CRYPTO_STORAGE_C)
     psa_key_id_t volatile_key_id;
 
     status = psa_get_empty_key_slot( &volatile_key_id, p_slot );
     if( status != PSA_SUCCESS )
         return( status );
 
-    (*p_slot)->attr.lifetime = PSA_KEY_LIFETIME_PERSISTENT;
     (*p_slot)->attr.id = key;
+    (*p_slot)->attr.lifetime = PSA_KEY_LIFETIME_PERSISTENT;
 
+    status = PSA_ERROR_DOES_NOT_EXIST;
+#if defined(MBEDTLS_PSA_CRYPTO_BUILTIN_KEYS)
+    status = psa_load_builtin_key_into_slot( *p_slot );
+    if( status == PSA_SUCCESS )
+        goto exit;
+#endif /* MBEDTLS_PSA_CRYPTO_BUILTIN_KEYS */
+
+#if defined(MBEDTLS_PSA_CRYPTO_STORAGE_C)
     status = psa_load_persistent_key_into_slot( *p_slot );
+#endif /* defined(MBEDTLS_PSA_CRYPTO_STORAGE_C) */
+
+exit:
     if( status != PSA_SUCCESS )
     {
         psa_wipe_key_slot( *p_slot );
@@ -309,9 +380,6 @@
             status = PSA_ERROR_INVALID_HANDLE;
     }
     return( status );
-#else
-    return( PSA_ERROR_INVALID_HANDLE );
-#endif /* defined(MBEDTLS_PSA_CRYPTO_STORAGE_C) */
 }
 
 psa_status_t psa_unlock_key_slot( psa_key_slot_t *slot )
diff --git a/library/version_features.c b/library/version_features.c
index 9332987..f665a23 100644
--- a/library/version_features.c
+++ b/library/version_features.c
@@ -438,6 +438,9 @@
 #if defined(MBEDTLS_PKCS1_V21)
     "MBEDTLS_PKCS1_V21",
 #endif /* MBEDTLS_PKCS1_V21 */
+#if defined(MBEDTLS_PSA_CRYPTO_BUILTIN_KEYS)
+    "MBEDTLS_PSA_CRYPTO_BUILTIN_KEYS",
+#endif /* MBEDTLS_PSA_CRYPTO_BUILTIN_KEYS */
 #if defined(MBEDTLS_PSA_CRYPTO_CLIENT)
     "MBEDTLS_PSA_CRYPTO_CLIENT",
 #endif /* MBEDTLS_PSA_CRYPTO_CLIENT */