Crypto: Refactor the tfm_builtin_key_loader and HAL interaction

This patch refactors the entry points of the tfm_builtin_key_loader
driver to simplify its interaction with crypto keys HAL layer and
the rest of the TF-M Crypto service and PSA Crypto core layer.

* Decouple as much as possible each module from mbed TLS specific
  types, and makes sure all library interactions are abstracted in
  the crypto_library module of the crypto service
* Simplify the HAL requirements to provide platform builtin keys
  as described in tfm_plat_crypto_keys.h
* Update the documentation to reflect the design change
* Fix minor issues and typos, include paths, etc
* Regenerate mbed TLS patches on top of 3.3.0 tag to be applied

Signed-off-by: Antonio de Angelis <antonio.deangelis@arm.com>
Change-Id: Id26ff0b88da87075490d17a4c8af8f209bb34a08
diff --git a/secure_fw/partitions/crypto/crypto_init.c b/secure_fw/partitions/crypto/crypto_init.c
index 84bdcf9..741a5f6 100644
--- a/secure_fw/partitions/crypto/crypto_init.c
+++ b/secure_fw/partitions/crypto/crypto_init.c
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2018-2022, Arm Limited. All rights reserved.
+ * Copyright (c) 2018-2023, Arm Limited. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  *
@@ -301,7 +301,6 @@
 psa_status_t tfm_crypto_init(void)
 {
     psa_status_t status;
-    enum tfm_plat_err_t plat_err;
 
     /* Initialise other modules of the service */
     status = tfm_crypto_module_init();
@@ -315,11 +314,6 @@
         return status;
     }
 
-    plat_err = tfm_plat_load_builtin_keys();
-    if (plat_err != TFM_PLAT_ERR_SUCCESS) {
-        return PSA_ERROR_GENERIC_ERROR;
-    }
-
     return PSA_SUCCESS;
 }
 
diff --git a/secure_fw/partitions/crypto/crypto_library.c b/secure_fw/partitions/crypto/crypto_library.c
index f69126d..014f6e5 100644
--- a/secure_fw/partitions/crypto/crypto_library.c
+++ b/secure_fw/partitions/crypto/crypto_library.c
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022, Arm Limited. All rights reserved.
+ * Copyright (c) 2022-2023, Arm Limited. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  *
@@ -167,18 +167,17 @@
     psa_key_lifetime_t *lifetime,
     psa_drv_slot_number_t *slot_number)
 {
-    enum tfm_plat_err_t plat_err;
-    struct tfm_crypto_key_id_s tfm_key_id = {
-        .key_id = MBEDTLS_SVC_KEY_ID_GET_KEY_ID(key_id),
-        .owner = MBEDTLS_SVC_KEY_ID_GET_OWNER_ID(key_id)};
+    const tfm_plat_builtin_key_descriptor_t *desc_table = NULL;
+    size_t number_of_keys = tfm_plat_builtin_key_get_desc_table_ptr(&desc_table);
 
-    plat_err = tfm_plat_builtin_key_get_lifetime_and_slot(tfm_key_id,
-                                                          lifetime,
-                                                          slot_number);
-    if (plat_err != TFM_PLAT_ERR_SUCCESS) {
-        return PSA_ERROR_DOES_NOT_EXIST;
+    for (size_t idx = 0; idx < number_of_keys; idx++) {
+        if (desc_table[idx].key_id == MBEDTLS_SVC_KEY_ID_GET_KEY_ID(key_id)) {
+            *lifetime = desc_table[idx].lifetime;
+            *slot_number = desc_table[idx].slot_number;
+            return PSA_SUCCESS;
+        }
     }
 
-    return PSA_SUCCESS;
+    return PSA_ERROR_DOES_NOT_EXIST;
 }
 /*!@}*/
diff --git a/secure_fw/partitions/crypto/crypto_library.h b/secure_fw/partitions/crypto/crypto_library.h
index 1ee9414..87c7626 100644
--- a/secure_fw/partitions/crypto/crypto_library.h
+++ b/secure_fw/partitions/crypto/crypto_library.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022, Arm Limited. All rights reserved.
+ * Copyright (c) 2022-2023, Arm Limited. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  *
@@ -32,6 +32,12 @@
 #define CRYPTO_LIBRARY_GET_KEY_ID(encoded_key_library) MBEDTLS_SVC_KEY_ID_GET_KEY_ID(encoded_key_library)
 
 /**
+ * \brief This macro extracts the owner from the library encoded key passed as parameter
+ *
+ */
+#define CRYPTO_LIBRARY_GET_OWNER(encoded_key_library) MBEDTLS_SVC_KEY_ID_GET_OWNER_ID(encoded_key_library)
+
+/**
  * \brief The following typedef must be defined to the type associated to the key_id in the underlying library
  *
  */
diff --git a/secure_fw/partitions/crypto/psa_driver_api/tfm_builtin_key_loader.c b/secure_fw/partitions/crypto/psa_driver_api/tfm_builtin_key_loader.c
index 1156961..2027033 100644
--- a/secure_fw/partitions/crypto/psa_driver_api/tfm_builtin_key_loader.c
+++ b/secure_fw/partitions/crypto/psa_driver_api/tfm_builtin_key_loader.c
@@ -1,70 +1,120 @@
 /*
- * Copyright (c) 2022, Arm Limited. All rights reserved.
+ * Copyright (c) 2022-2023, Arm Limited. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  *
  */
-
-#include "tfm_builtin_key_loader.h"
-
-#include "psa/error.h"
-#include "tfm_mbedcrypto_include.h"
-#include "tfm_crypto_defs.h"
+#include <string.h>
 #include "mbedtls/hkdf.h"
+#include "tfm_builtin_key_loader.h"
 #include "psa_manifest/pid.h"
 #include "tfm_plat_crypto_keys.h"
-
-#include <string.h>
+#include "crypto_library.h"
 
 #ifndef TFM_BUILTIN_MAX_KEY_LEN
-#define TFM_BUILTIN_MAX_KEY_LEN 48
+#define TFM_BUILTIN_MAX_KEY_LEN (48)
 #endif /* TFM_BUILTIN_MAX_KEY_LEN */
 
 #ifndef TFM_BUILTIN_MAX_KEYS
-#define TFM_BUILTIN_MAX_KEYS 8
+#define TFM_BUILTIN_MAX_KEYS (TFM_BUILTIN_KEY_SLOT_MAX)
 #endif /* TFM_BUILTIN_MAX_KEYS */
 
-#ifndef MBEDTLS_PSA_CRYPTO_KEY_ID_ENCODES_OWNER
-#error "MBEDTLS_PSA_CRYPTO_KEY_ID_ENCODES_OWNER must be selected in Mbed TLS config file"
-#endif
+#define NUMBER_OF_ELEMENTS_OF(x) sizeof(x)/sizeof(*x)
 
+/*!
+ * \brief A structure which describes a builtin key slot
+ */
 struct tfm_builtin_key_t {
-    uint8_t key[TFM_BUILTIN_MAX_KEY_LEN];
-    size_t key_len;
-    psa_key_attributes_t attr;
-    uint32_t is_loaded;
-    psa_drv_slot_number_t slot_number;
+    uint8_t __attribute__((aligned(4))) key[TFM_BUILTIN_MAX_KEY_LEN]; /*!< Raw key material, 4-byte aligned */
+    size_t key_len;                       /*!< Size of the key material held in the key buffer */
+    psa_key_attributes_t attr;            /*!< Key attributes associated to the key */
+    uint32_t is_loaded;                   /*!< Boolean indicating whether the slot is being used */
 };
 
-static struct tfm_builtin_key_t builtin_key_slots[TFM_BUILTIN_MAX_KEYS] = {0};
+/*!
+ * \brief The below array is used by the driver to store key material in secure memory, in order
+ *        for the keys and relevant metadata to be accessible by the PSA Crypto core layer
+ */
+static struct tfm_builtin_key_t g_builtin_key_slots[TFM_BUILTIN_MAX_KEYS] = {0};
 
-static psa_status_t builtin_key_get_attributes(
-        struct tfm_builtin_key_t *key_slot, psa_key_attributes_t *attr)
+/*!
+ * \brief This functions returns the slot associated to a key id interrogating the
+ *        platform HAL table
+ */
+static psa_status_t builtin_key_get_slot(psa_key_id_t key_id, psa_drv_slot_number_t *slot_ptr)
 {
-    enum tfm_plat_err_t plat_err;
-    /* Preserve the key id and lifetime */
-    mbedtls_svc_key_id_t key_id = psa_get_key_id(attr);
-    psa_key_lifetime_t lifetime = psa_get_key_lifetime(attr);
-    psa_key_usage_t usage = psa_get_key_usage_flags(attr);
+    const tfm_plat_builtin_key_descriptor_t *desc_table = NULL;
+    size_t number_of_keys = tfm_plat_builtin_key_get_desc_table_ptr(&desc_table);
+    psa_drv_slot_number_t slot_number = TFM_BUILTIN_KEY_SLOT_MAX;
 
-    memcpy(attr, &(key_slot->attr), sizeof(psa_key_attributes_t));
-
-    plat_err = tfm_plat_builtin_key_get_usage(
-            MBEDTLS_SVC_KEY_ID_GET_KEY_ID(key_id),
-            MBEDTLS_SVC_KEY_ID_GET_OWNER_ID(key_id), &usage);
-    if (plat_err != TFM_PLAT_ERR_SUCCESS) {
-        return PSA_ERROR_NOT_PERMITTED;
+    for (size_t idx = 0; idx < number_of_keys; idx++) {
+        if (desc_table[idx].key_id == key_id) {
+            slot_number = desc_table[idx].slot_number;
+            break;
+        }
     }
 
-    psa_set_key_id(attr, key_id);
-    psa_set_key_lifetime(attr, lifetime);
+    if (slot_number == TFM_BUILTIN_KEY_SLOT_MAX) {
+        *slot_ptr = TFM_BUILTIN_KEY_SLOT_MAX;
+        return PSA_ERROR_DOES_NOT_EXIST;
+    }
+
+    *slot_ptr = slot_number;
+    return PSA_SUCCESS;
+}
+
+/*!
+ * \brief This functions returns the attributes of the key interrogating the
+ *        platform HAL
+ */
+static psa_status_t builtin_key_get_attributes(
+        struct tfm_builtin_key_t *key_slot, int32_t user, psa_key_id_t key_id, psa_key_attributes_t *attr)
+{
+    psa_key_usage_t usage = 0x0;
+
+    /* Retrieve the usage policy based on the key_id and the user of the key */
+    const tfm_plat_builtin_key_policy_t *policy_table = NULL;
+    size_t number_of_keys = tfm_plat_builtin_key_get_policy_table_ptr(&policy_table);
+
+    for (size_t idx = 0; idx < number_of_keys; idx++) {
+        if (policy_table[idx].key_id == key_id) {
+            if (policy_table[idx].per_user_policy == 0) {
+                usage = policy_table[idx].usage;
+            } else {
+                /* The policy depedends also on the user of the key */
+                size_t num_users = policy_table[idx].per_user_policy;
+                const tfm_plat_builtin_key_per_user_policy_t *p_policy = policy_table[idx].policy_ptr;
+
+                for (size_t j = 0; j < num_users; j++) {
+                    if (p_policy[j].user == user) {
+                        usage = p_policy[j].usage;
+                        break;
+                    }
+                }
+            }
+            break;
+        }
+    }
+
+    /* A normal copy is enough to copy all the fields as there no pointers */
+    memcpy(attr, &(key_slot->attr), sizeof(psa_key_attributes_t));
+    /* The stored attributes have an owner == 0, but we need to preserve the
+     * user received from the caller, which is stored in user_id
+     */
+    psa_set_key_id(attr, tfm_crypto_library_key_id_init(user, key_id));
+    /* Set the flags according to the policy for the user for this key_id */
     psa_set_key_usage_flags(attr, usage);
 
     return PSA_SUCCESS;
 }
 
+/*!
+ * \brief This function derives a key into a provided buffer, to make sure that
+ *        keys that are returned for usage to the PSA Crypto core, are differentiated
+ *        based on the partition user. The derived keys are described as platform keys
+ */
 static psa_status_t derive_subkey_into_buffer(
-        struct tfm_builtin_key_t *key_slot, mbedtls_key_owner_id_t owner,
+        struct tfm_builtin_key_t *key_slot, int32_t user,
         uint8_t *key_buffer, size_t key_buffer_size, size_t *key_buffer_length)
 {
     int mbedtls_err;
@@ -74,8 +124,8 @@
      * ps_system_prepare from the test partition which would otherwise derive a
      * different key.
      */
-    if (owner == TFM_SP_PS_TEST) {
-        owner = TFM_SP_PS;
+    if (user == TFM_SP_PS_TEST) {
+        user = TFM_SP_PS;
     }
 #endif /* TFM_PARTITION_TEST_PS */
 
@@ -88,7 +138,7 @@
      */
     mbedtls_err = mbedtls_hkdf(mbedtls_md_info_from_type(MBEDTLS_MD_SHA256),
                                NULL, 0, key_slot->key, key_slot->key_len,
-                               (uint8_t *)&owner, sizeof(owner), key_buffer,
+                               (uint8_t *)&user, sizeof(user), key_buffer,
                                key_buffer_size);
     if (mbedtls_err) {
         return PSA_ERROR_GENERIC_ERROR;
@@ -114,103 +164,140 @@
  *
  */
 /*!@{*/
-psa_status_t tfm_builtin_key_loader_load_key(uint8_t *buf, size_t key_len,
-                                             psa_key_attributes_t *attr)
+psa_status_t tfm_builtin_key_loader_init(void)
 {
-    psa_status_t err;
-    psa_drv_slot_number_t slot_number;
-    psa_key_lifetime_t lifetime;
-    mbedtls_svc_key_id_t key_id;
+    psa_status_t err = PSA_ERROR_CORRUPTION_DETECTED;
+    const tfm_plat_builtin_key_descriptor_t *desc_table = NULL;
+    size_t number_of_keys = tfm_plat_builtin_key_get_desc_table_ptr(&desc_table);
 
-    /* Set the owner to 0, as we handle permissions on a granular basis. Having
-     * builtin keys being defined with different owners seems to cause a memory
-     * leak in the MbedTLS core.
-     */
-    key_id = psa_get_key_id(attr);
-    key_id = mbedtls_svc_key_id_make(0, MBEDTLS_SVC_KEY_ID_GET_KEY_ID(key_id));
-    psa_set_key_id(attr, key_id);
+    /* These properties and key material are filled by the loaders */
+    uint8_t buf[TFM_BUILTIN_MAX_KEY_LEN];
+    size_t key_len;
+    size_t key_bits;
+    psa_algorithm_t algorithm;
+    psa_key_type_t type;
 
-    if (key_len > TFM_BUILTIN_MAX_KEY_LEN) {
-        return PSA_ERROR_BUFFER_TOO_SMALL;
+    for (size_t key = 0; key < number_of_keys; key++) {
+        if (desc_table[key].lifetime != TFM_BUILTIN_KEY_LOADER_LIFETIME) {
+            /* If the key is not bound to this driver, just don't load it */
+            continue;
+        }
+        psa_drv_slot_number_t slot_number = desc_table[key].slot_number;
+        /* The owner of a builtin key is set to 0 */
+        tfm_crypto_library_key_id_t key_id = tfm_crypto_library_key_id_init(0, desc_table[key].key_id);
+        psa_key_attributes_t attr = PSA_KEY_ATTRIBUTES_INIT;
+        enum tfm_plat_err_t plat_err = desc_table[key].loader_key_func(
+                                    &buf[0], TFM_BUILTIN_MAX_KEY_LEN, &key_len, &key_bits, &algorithm, &type);
+        if (plat_err != TFM_PLAT_ERR_SUCCESS) {
+            err = PSA_ERROR_HARDWARE_FAILURE;
+            goto wrap_up;
+        }
+
+        /* Build the attributes with the metadata retrieved from the platform and desc table */
+        psa_set_key_id(&attr, key_id);
+        psa_set_key_bits(&attr, key_bits);
+        psa_set_key_algorithm(&attr, algorithm);
+        psa_set_key_type(&attr, type);
+        psa_set_key_lifetime(&attr, desc_table[key].lifetime);
+
+        /* Populate the internal table of the tfm_builtin_key_loader driver with key and metadata */
+        memcpy(&(g_builtin_key_slots[slot_number].attr), &attr, sizeof(psa_key_attributes_t));
+        memcpy(&(g_builtin_key_slots[slot_number].key), buf, key_len);
+        g_builtin_key_slots[slot_number].key_len = key_len;
+        g_builtin_key_slots[slot_number].is_loaded = 1;
     }
+    /* At this point the discovered keys have been loaded successfully into the driver */
+    err = PSA_SUCCESS;
 
-    err = mbedtls_psa_platform_get_builtin_key(psa_get_key_id(attr), &lifetime,
-                                               &slot_number);
-    if (err != PSA_SUCCESS) {
-        return err;
-    }
-
-    memcpy(&(builtin_key_slots[slot_number].attr), attr,
-           sizeof(psa_key_attributes_t));
-    memcpy(&(builtin_key_slots[slot_number].key), buf, key_len);
-    builtin_key_slots[slot_number].key_len = key_len;
-    builtin_key_slots[slot_number].is_loaded = 1;
-
-    return PSA_SUCCESS;
+wrap_up:
+    return err;
 }
 
 psa_status_t tfm_builtin_key_loader_get_key_buffer_size(
-        mbedtls_svc_key_id_t key_id, size_t *len)
+        tfm_crypto_library_key_id_t key_id, size_t *len)
 {
     psa_status_t err;
     psa_drv_slot_number_t slot_number;
-    psa_key_lifetime_t lifetime;
 
-    err = mbedtls_psa_platform_get_builtin_key(key_id, &lifetime, &slot_number);
-    if (err != PSA_SUCCESS) {
-        return err;
+    if (len == NULL) {
+        err = PSA_ERROR_INVALID_ARGUMENT;
+        goto wrap_up;
     }
 
-    *len = builtin_key_slots[slot_number].key_len;
-    return PSA_SUCCESS;
+    *len = 0;
+
+    err = builtin_key_get_slot(CRYPTO_LIBRARY_GET_KEY_ID(key_id), &slot_number);
+    if (err != PSA_SUCCESS) {
+        goto wrap_up;
+    }
+
+    *len = g_builtin_key_slots[slot_number].key_len;
+    err = PSA_SUCCESS;
+
+wrap_up:
+    return err;
 }
 
-psa_status_t tfm_builtin_key_loader_get_key_buffer(
+psa_status_t tfm_builtin_key_loader_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_status_t err;
     struct tfm_builtin_key_t *key_slot;
-    mbedtls_svc_key_id_t key_id;
+    tfm_crypto_library_key_id_t key_id;
 
-    if ((uint32_t)slot_number >= TFM_BUILTIN_MAX_KEYS) {
-        return PSA_ERROR_DOES_NOT_EXIST;
+    if ((uint32_t)slot_number >= NUMBER_OF_ELEMENTS_OF(g_builtin_key_slots)) {
+        err = PSA_ERROR_INVALID_ARGUMENT;
+        goto wrap_up;
     }
 
-    key_slot = &builtin_key_slots[slot_number];
+    key_slot = &g_builtin_key_slots[slot_number];
 
+    /* The request is for a valid slot that has not been loaded*/
     if (!key_slot->is_loaded) {
-        return PSA_ERROR_DOES_NOT_EXIST;
+        err = PSA_ERROR_DOES_NOT_EXIST;
+        goto wrap_up;
     }
 
     if (attributes == NULL) {
-        return PSA_ERROR_INVALID_ARGUMENT;
-    }
-
-    err = builtin_key_get_attributes(key_slot, attributes);
-    if (err != PSA_SUCCESS) {
-        return err;
-    }
-
-    if (key_buffer_size < key_slot->key_len) {
-        return PSA_ERROR_BUFFER_TOO_SMALL;
-    }
-
-    if (key_buffer == NULL || key_buffer_length == NULL) {
-        return PSA_ERROR_INVALID_ARGUMENT;
+        err = PSA_ERROR_INVALID_ARGUMENT;
+        goto wrap_up;
     }
 
     key_id = psa_get_key_id(attributes);
+    err = builtin_key_get_attributes(key_slot,
+                                     CRYPTO_LIBRARY_GET_OWNER(key_id),
+                                     CRYPTO_LIBRARY_GET_KEY_ID(key_id),
+                                     attributes);
+    if (err != PSA_SUCCESS) {
+        goto wrap_up;
+    }
+
+    /* Note that the core layer might use this function just to probe, for a loaded
+     * key, the attributes of the key. In this case, key_buffer and key_buffer_length
+     * could be NULL hence the validation of those pointers must happen after this
+     * check here. In fact, PSA_ERROR_BUFFER_TOO_SMALL is considered in the crypto
+     * core as a safe failure value that lets the caller continue.
+     */
+    if (key_buffer_size < key_slot->key_len) {
+        err = PSA_ERROR_BUFFER_TOO_SMALL;
+        goto wrap_up;
+    }
+
+    if (key_buffer == NULL || key_buffer_length == NULL) {
+        err = PSA_ERROR_INVALID_ARGUMENT;
+        goto wrap_up;
+    }
 
     /* If a key can be used for derivation, we derive a further subkey for each
-     * owner to prevent multiple owners deriving the same keys as each other.
+     * users to prevent multiple users deriving the same keys as each other.
      * For keys for encryption and signing, it's assumed that if multiple
      * partitions have access to the key, there is a good reason and therefore
-     * they both need access to the raw builtin key.
+     * they all need access to the raw builtin key.
      */
+    int32_t user = CRYPTO_LIBRARY_GET_OWNER(key_id);
     if (psa_get_key_usage_flags(attributes) & PSA_KEY_USAGE_DERIVE) {
-        err = derive_subkey_into_buffer(key_slot,
-                                        MBEDTLS_SVC_KEY_ID_GET_OWNER_ID(key_id),
+        err = derive_subkey_into_buffer(key_slot, user,
                                         key_buffer, key_buffer_size,
                                         key_buffer_length);
     } else {
@@ -218,6 +305,7 @@
                                          key_buffer_length);
     }
 
+wrap_up:
     return err;
 }
 /*!@}*/
diff --git a/secure_fw/partitions/crypto/psa_driver_api/tfm_builtin_key_loader.h b/secure_fw/partitions/crypto/psa_driver_api/tfm_builtin_key_loader.h
index 488f5600a..bb328a5 100644
--- a/secure_fw/partitions/crypto/psa_driver_api/tfm_builtin_key_loader.h
+++ b/secure_fw/partitions/crypto/psa_driver_api/tfm_builtin_key_loader.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022, Arm Limited. All rights reserved.
+ * Copyright (c) 2022-2023, Arm Limited. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  *
@@ -8,12 +8,12 @@
 #ifndef TFM_BUILTIN_KEY_LOADER_H
 #define TFM_BUILTIN_KEY_LOADER_H
 
+#include <psa/crypto.h>
+
 #ifdef __cplusplus
 extern "C" {
 #endif
 
-#include "tfm_mbedcrypto_include.h"
-
 #ifdef PLATFORM_DEFAULT_CRYPTO_KEYS
 enum psa_drv_slot_number_t {
     TFM_BUILTIN_KEY_SLOT_HUK = 0,
@@ -27,27 +27,49 @@
 #include "platform_builtin_key_loader_ids.h"
 #endif
 
+#ifdef __DOXYGEN_ONLY__
 /**
- * \brief Load a key into the builtin key driver
- *
- * \note This should be called for all slots that are required on initial boot,
- *       from the tfm_plat_load_builtin_keys function.
- *
- * \param[in]  buf     Buffer containing key material.
- * \param[in]  key_len Size of the key in bytes.
- * \param[in]  attr    The attributes of the key.
+ *  \brief Enables the tfm_builtin_key_loader driver in the PSA Crypto
+ *         core subsystem
+ */
+#define PSA_CRYPTO_DRIVER_TFM_BUILTIN_KEY_LOADER
+#endif /* __DOXYGEN_ONLY__ */
+
+/**
+ * \brief The PSA driver location for TF-M builtin keys. Arbitrary within the
+ *        ranges documented at
+ *        https://armmbed.github.io/mbed-crypto/html/api/keys/lifetimes.html#c.psa_key_location_t
+ */
+#define TFM_BUILTIN_KEY_LOADER_KEY_LOCATION ((psa_key_location_t)0x800001)
+
+/**
+ * \brief This macro defines the lifetime associated to TF-M builtin keys as
+ *        persistent and as an ad-hoc location associated to the TFM_BUILTIN_KEY_LOADER
+ *        driver. To be handled by the tfm_builtin_ker_loader driver, the lifetime of
+ *        the platform keys must be set equal to this particular lifetime value
+ */
+#define TFM_BUILTIN_KEY_LOADER_LIFETIME PSA_KEY_LIFETIME_FROM_PERSISTENCE_AND_LOCATION( \
+                                            PSA_KEY_LIFETIME_PERSISTENT, TFM_BUILTIN_KEY_LOADER_KEY_LOCATION)
+
+/**
+ * \brief This is the initialisation function for the tfm_builtin_key_laoder driver,
+ *        to be called from the PSA core initialisation subsystem. It discovers the
+ *        keys available in the underlying hardware platform and loads them in
+ *        memory visible to the PSA Crypto subsystem to be used to the normal APIs
  *
  * \return Returns error code specified in \ref psa_status_t
  */
-psa_status_t tfm_builtin_key_loader_load_key(uint8_t *buf, size_t key_len,
-                                             psa_key_attributes_t *attr);
+psa_status_t tfm_builtin_key_loader_init(void);
 
 /**
  * \brief Returns the length of a key from the builtin driver.
  *
  * \note This function is called by the psa crypto driver wrapper.
  *
- * \param[in]  key_id  The ID of the key to return the length of
+ * \param[in]  key_id  The ID of the key to return the length of. The type of this
+ *                     must match the expected type of the underlying library that
+ *                     provides the key management for the PSA Crypto core, and
+ *                     must support encoding the owner in addition to the key_id.
  * \param[out] len     The length of the key.
  *
  * \return Returns error code specified in \ref psa_status_t
@@ -57,7 +79,7 @@
 
 /**
  * \brief Returns the attributes and key material of a key from the builtin
- *        driver
+ *        driver to be used by the PSA Crypto core
  *
  * \note This function is called by the psa crypto driver wrapper.
  *
@@ -69,7 +91,7 @@
  *
  * \return Returns error code specified in \ref psa_status_t
  */
-psa_status_t tfm_builtin_key_loader_get_key_buffer(
+psa_status_t tfm_builtin_key_loader_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);