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/include/mbedtls/config.h b/include/mbedtls/config.h
index d370dbf..62d89c9 100644
--- a/include/mbedtls/config.h
+++ b/include/mbedtls/config.h
@@ -1338,6 +1338,22 @@
  */
 #define MBEDTLS_PKCS1_V21
 
+/** \def MBEDTLS_PSA_CRYPTO_BUILTIN_KEYS
+ *
+ * Enable support for platform built-in keys. If you enable this feature,
+ * you must implement the function mbedtls_psa_platform_get_builtin_key().
+ * See the documentation of that function for more information.
+ *
+ * Built-in keys are typically derived from a hardware unique key or
+ * stored in a secure element.
+ *
+ * Requires: MBEDTLS_PSA_CRYPTO_C.
+ *
+ * \warning This interface is experimental and may change or be removed
+ * without notice.
+ */
+//#define MBEDTLS_PSA_CRYPTO_BUILTIN_KEYS
+
 /** \def MBEDTLS_PSA_CRYPTO_CLIENT
  *
  * Enable support for PSA crypto client.
diff --git a/include/psa/crypto_extra.h b/include/psa/crypto_extra.h
index e01d827..f9a9aee 100644
--- a/include/psa/crypto_extra.h
+++ b/include/psa/crypto_extra.h
@@ -713,6 +713,99 @@
 
 /**@}*/
 
+/** \defgroup psa_builtin_keys Built-in keys
+ * @{
+ */
+
+/** The minimum value for a key identifier that is built into the
+ * implementation.
+ *
+ * The range of key identifiers from #MBEDTLS_PSA_KEY_ID_BUILTIN_MIN
+ * to #MBEDTLS_PSA_KEY_ID_BUILTIN_MAX within the range from
+ * #PSA_KEY_ID_VENDOR_MIN and #PSA_KEY_ID_VENDOR_MAX and must not intersect
+ * with any other set of implementation-chosen key identifiers.
+ *
+ * This value is part of the library's ABI since changing it would invalidate
+ * the values of built-in key identifiers in applications.
+ */
+#define MBEDTLS_PSA_KEY_ID_BUILTIN_MIN          ((psa_key_id_t)0x7fff0000)
+
+/** The maximum value for a key identifier that is built into the
+ * implementation.
+ *
+ * See #MBEDTLS_PSA_KEY_ID_BUILTIN_MIN for more information.
+ */
+#define MBEDTLS_PSA_KEY_ID_BUILTIN_MAX          ((psa_key_id_t)0x7fffefff)
+
+/** A slot number identifying a key in a driver.
+ *
+ * Values of this type are used to identify built-in keys.
+ */
+typedef uint64_t psa_drv_slot_number_t;
+
+#if defined(MBEDTLS_PSA_CRYPTO_BUILTIN_KEYS)
+/** Test whether a key identifier belongs to the builtin key range.
+ *
+ * \param key_id  Key identifier to test.
+ *
+ * \retval 1
+ *         The key identifier is a builtin key identifier.
+ * \retval 0
+ *         The key identifier is not a builtin key identifier.
+ */
+static inline int psa_key_id_is_builtin( psa_key_id_t key_id )
+{
+    return( ( key_id >= MBEDTLS_PSA_KEY_ID_BUILTIN_MIN ) &&
+            ( key_id <= MBEDTLS_PSA_KEY_ID_BUILTIN_MAX ) );
+}
+
+/** Platform function to obtain the data of a built-in key.
+ *
+ * An application-specific implementation of this function must be provided if
+ * #MBEDTLS_PSA_CRYPTO_BUILTIN_KEYS is enabled. This would typically provided
+ * as part of a platform's system image.
+ *
+ * Call psa_get_key_id(\p attributes) to obtain the key identifier \c key_id.
+ * #MBEDTLS_SVC_KEY_ID_GET_KEY_ID(\p key_id) is in the range from
+ * #MBEDTLS_PSA_KEY_ID_BUILTIN_MIN to #MBEDTLS_PSA_KEY_ID_BUILTIN_MAX.
+ *
+ * In a multi-application configuration
+ * (\c MBEDTLS_PSA_CRYPTO_KEY_ID_ENCODES_OWNER is defined),
+ * this function should check that #MBEDTLS_SVC_KEY_ID_GET_OWNER_ID(\p key_id)
+ * is allowed to use the given key.
+ *
+ * \param[in,out] attributes    On entry, this is #PSA_KEY_ATTRIBUTES_INIT or
+ *                              an equivalent value, except that the key
+ *                              identifier field is set.
+ *                              On successful return, this function must set
+ *                              the attributes of the key: lifetime, type,
+ *                              bit-size, usage policy.
+ * \param[out] slot_number      On successful return, this function must
+ *                              this to the slot number known to the driver for
+ *                              the lifetime location reported through
+ *                              \p attributes which corresponds to the
+ *                              requested built-in key.
+ *
+ * \retval #PSA_SUCCESS
+ *         The requested key identifier designates a built-in key.
+ *         In a multi-application configuration, the requested owner
+ *         is allowed to access it.
+ * \retval #PSA_ERROR_DOES_NOT_EXIST
+ *         The requested key identifier is not a built-in key which is known
+ *         to this function. If a key exists in the key storage with this
+ *         identifier, the data from the storage will be used.
+ * \retval (any other error)
+ *         Any other error is propagated to the function that requested the key.
+ *         Common errors include:
+ *         - #PSA_ERROR_NOT_PERMITTED: the key exists but the requested owner
+ *           is not allowed to access it.
+ */
+psa_status_t mbedtls_psa_platform_get_builtin_key(
+    psa_key_attributes_t *attributes, psa_drv_slot_number_t *slot_number );
+#endif /* MBEDTLS_PSA_CRYPTO_BUILTIN_KEYS */
+
+/** @} */
+
 #ifdef __cplusplus
 }
 #endif
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 */
diff --git a/programs/test/query_config.c b/programs/test/query_config.c
index b9105f8..9760f62 100644
--- a/programs/test/query_config.c
+++ b/programs/test/query_config.c
@@ -1226,6 +1226,14 @@
     }
 #endif /* MBEDTLS_PKCS1_V21 */
 
+#if defined(MBEDTLS_PSA_CRYPTO_BUILTIN_KEYS)
+    if( strcmp( "MBEDTLS_PSA_CRYPTO_BUILTIN_KEYS", config ) == 0 )
+    {
+        MACRO_EXPANSION_TO_STR( MBEDTLS_PSA_CRYPTO_BUILTIN_KEYS );
+        return( 0 );
+    }
+#endif /* MBEDTLS_PSA_CRYPTO_BUILTIN_KEYS */
+
 #if defined(MBEDTLS_PSA_CRYPTO_CLIENT)
     if( strcmp( "MBEDTLS_PSA_CRYPTO_CLIENT", config ) == 0 )
     {
diff --git a/tests/scripts/all.sh b/tests/scripts/all.sh
index f768e1e..a85c7ce 100755
--- a/tests/scripts/all.sh
+++ b/tests/scripts/all.sh
@@ -2267,6 +2267,7 @@
     msg "build: MBEDTLS_PSA_CRYPTO_DRIVERS w/ driver hooks"
     scripts/config.py full
     scripts/config.py set MBEDTLS_PSA_CRYPTO_DRIVERS
+    scripts/config.py set MBEDTLS_PSA_CRYPTO_BUILTIN_KEYS
     # Need to define the correct symbol and include the test driver header path in order to build with the test driver
     loc_cflags="$ASAN_CFLAGS -DPSA_CRYPTO_DRIVER_TEST"
     loc_cflags="${loc_cflags} -DMBEDTLS_PSA_ACCEL_KEY_TYPE_AES"
diff --git a/tests/src/helpers.c b/tests/src/helpers.c
index e323275..c282edc 100644
--- a/tests/src/helpers.c
+++ b/tests/src/helpers.c
@@ -282,3 +282,38 @@
     }
 }
 #endif /* MBEDTLS_CHECK_PARAMS */
+
+#if defined(MBEDTLS_PSA_CRYPTO_BUILTIN_KEYS)
+#include <psa/crypto.h>
+typedef struct
+{
+    psa_key_id_t builtin_key_id;
+    psa_key_location_t location;
+    psa_drv_slot_number_t slot_number;
+} mbedtls_psa_builtin_key_description_t;
+static const mbedtls_psa_builtin_key_description_t builtin_keys[] = {
+    // TODO: declare some keys
+    {0, 0, 0},
+};
+psa_status_t mbedtls_psa_platform_get_builtin_key(
+    psa_key_attributes_t *attributes, psa_drv_slot_number_t *slot_number )
+{
+    mbedtls_svc_key_id_t svc_key_id = psa_get_key_id( attributes );
+    psa_key_id_t app_key_id = MBEDTLS_SVC_KEY_ID_GET_KEY_ID( svc_key_id );
+
+    for( size_t i = 0; i < ( sizeof( builtin_keys ) / sizeof( builtin_keys[0] ) ); i++ )
+    {
+        if( builtin_keys[i].builtin_key_id == app_key_id )
+        {
+            psa_set_key_lifetime( attributes,
+                          PSA_KEY_LIFETIME_FROM_PERSISTENCE_AND_LOCATION(
+                              PSA_KEY_PERSISTENCE_READ_ONLY,
+                              builtin_keys[i].location ) );
+            *slot_number = builtin_keys[i].slot_number;
+            return( PSA_SUCCESS );
+        }
+    }
+
+    return( PSA_ERROR_DOES_NOT_EXIST );
+}
+#endif /* MBEDTLS_PSA_CRYPTO_BUILTIN_KEYS */