Crypto: Implement PSA builtin keys

Implement builtin key driver, and add HAL apis to load HUK and IAK into
this driver. Add necessary funtions to route PSA crypto calls into this
driver. Add fixed builtin key IDs into the interface, and a mechanism to
allow platforms to add extra keys.

Change-Id: I7ffc16eb14215dd6b323baeb53b40ccb1c0ce126
Signed-off-by: Raef Coles <raef.coles@arm.com>
diff --git a/bl2/ext/mcuboot/keys.c b/bl2/ext/mcuboot/keys.c
index 46e73f3..1964ae8 100644
--- a/bl2/ext/mcuboot/keys.c
+++ b/bl2/ext/mcuboot/keys.c
@@ -27,7 +27,7 @@
 #include <stddef.h>
 #include <bootutil/sign_key.h>
 #include "mcuboot_config/mcuboot_config.h"
-#include "tfm_plat_crypto_keys.h"
+#include "tfm_plat_rotpk.h"
 
 #ifdef MCUBOOT_ENC_IMAGES
 unsigned char enc_priv_key[] = {
diff --git a/config/config_default.cmake b/config/config_default.cmake
index 81862af..ea4275e 100755
--- a/config/config_default.cmake
+++ b/config/config_default.cmake
@@ -143,6 +143,7 @@
 set(CRYPTO_IOVEC_BUFFER_SIZE            5120        CACHE STRING    "Default size of the internal scratch buffer used for PSA FF IOVec allocations")
 set(CRYPTO_NV_SEED                      ON          CACHE BOOL      "Use stored NV seed to provide entropy")
 set(CRYPTO_SINGLE_PART_FUNCS_DISABLED   OFF         CACHE BOOL      "Only enable multi-part operations in Hash, MAC, AEAD and symmetric ciphers, to optimize memory footprint in resource-constrained devices")
+set(CRYPTO_TFM_BUILTIN_KEYS_DRIVER      ON          CACHE BOOL      "Whether to allow crypto service to store builtin keys. Without this, ALL builtin keys must be stored in a platform-specific location")
 
 set(TFM_PARTITION_INITIAL_ATTESTATION   ON          CACHE BOOL      "Enable Initial Attestation partition")
 set(SYMMETRIC_INITIAL_ATTESTATION       OFF         CACHE BOOL      "Use symmetric crypto for inital attestation")
diff --git a/docs/integration_guide/services/tfm_attestation_integration_guide.rst b/docs/integration_guide/services/tfm_attestation_integration_guide.rst
index 33657f8..998bb31 100644
--- a/docs/integration_guide/services/tfm_attestation_integration_guide.rst
+++ b/docs/integration_guide/services/tfm_attestation_integration_guide.rst
@@ -216,8 +216,6 @@
 - **Crypto interface**:
     - ``t_cose_crypto.h``: Expose an API to bind the ``t_cose`` implementation
       to any cryptographic library.
-    - ``tfm_plat_crypto_keys.h``: Expose an API to get the attestation key from
-      platform layer.
 
 PSA interface
 =============
@@ -440,20 +438,10 @@
 ^^^^^^^^^^^^
 The provisioning of the initial attestation key is out of scope of the service
 and this document. It is assumed that device maker provisions the unique
-asymmetric key pair during the manufacturing process. The following API is
-defined to retrieve the attestation key pair from platform layer. Software
-integrators **must** port this interface according to their SoC design and make
-sure that key pair is available by Crypto service:
-
-- ``tfm_plat_get_initial_attest_key()``: Retrieve the initial attestation key
-  pair from platform layer.
-
-In TF-M project the attestation key is retrieved by initial attestation service.
-The key is registered and unregistered to the Crypto service by attestation
-service with ``psa_import_key()`` and ``psa_destroy_key()`` API calls for
-further usage. See in ``attestation_key.c``. In other implementation if the
-attestation key is directly retrieved by the Crypto service then this key
-handling is not necessary.
+asymmetric key pair during the manufacturing process. Software integrators
+**must** make sure that ``TFM_BUILTIN_KEY_SLOT_IAK`` is available via the Crypto
+service, which will then be used by the Attestation partition to perform the
+required signing operations via the PSA crypto interface.
 
 Symmetric key algorithm based attestation
 -----------------------------------------
diff --git a/interface/include/crypto_keys/tfm_builtin_key_ids.h b/interface/include/crypto_keys/tfm_builtin_key_ids.h
new file mode 100644
index 0000000..cb67859
--- /dev/null
+++ b/interface/include/crypto_keys/tfm_builtin_key_ids.h
@@ -0,0 +1,33 @@
+/*
+ * Copyright (c) 2022, Arm Limited. All rights reserved.
+ *
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ */
+
+#ifndef __TFM_BUILTIN_KEY_IDS_H__
+#define __TFM_BUILTIN_KEY_IDS_H__
+
+/**
+ * \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 The persistent key identifiers for TF-M builtin keys.
+ *
+ * The value of TFM_BUILTIN_KEY_ID_MIN (and therefore of the whole range) is
+ * completely arbitrary except for being inside the PSA builtin keys range.
+ *
+ */
+enum tfm_builtin_key_id_t {
+    TFM_BUILTIN_KEY_ID_MIN = 0x7fff815Bu,
+    TFM_BUILTIN_KEY_ID_HUK,
+    TFM_BUILTIN_KEY_ID_IAK,
+    TFM_BUILTIN_KEY_ID_PLAT_SPECIFIC_MIN = 0x7FFF816Bu,
+    TFM_BUILTIN_KEY_ID_MAX = 0x7FFF817Bu,
+};
+
+#endif /* __TFM_BUILTIN_KEY_IDS_H__ */
diff --git a/interface/include/tfm_crypto_defs.h b/interface/include/tfm_crypto_defs.h
index b5e65d5..5d429a8 100644
--- a/interface/include/tfm_crypto_defs.h
+++ b/interface/include/tfm_crypto_defs.h
@@ -16,6 +16,11 @@
 #include <limits.h>
 #include "tfm_api.h"
 #include "psa/crypto.h"
+#ifdef PLATFORM_DEFAULT_CRYPTO_KEYS
+#include "crypto_keys/tfm_builtin_key_ids.h"
+#else
+#include "tfm_builtin_key_ids.h"
+#endif /* PLATFORM_DEFAULT_CRYPTO_KEYS */
 
 /**
  * \brief This type is used to overcome a limitation in the number of maximum
@@ -264,19 +269,6 @@
 #define TFM_CRYPTO_INVALID_HANDLE (0x0u)
 
 /**
- * \brief The persistent key identifier that refers to the hardware unique key.
- *
- */
-#define TFM_CRYPTO_KEY_ID_HUK (0xFFFF815Bu)
-
-/**
- * \brief The algorithm identifier that refers to key derivation from the
- *        hardware unique key.
- *
- */
-#define TFM_CRYPTO_ALG_HUK_DERIVATION ((psa_algorithm_t)0xB0000F00)
-
-/**
  * \brief Define miscellaneous literal constants that are used in the service
  *
  */
diff --git a/platform/CMakeLists.txt b/platform/CMakeLists.txt
index 6418681..74da877 100755
--- a/platform/CMakeLists.txt
+++ b/platform/CMakeLists.txt
@@ -25,6 +25,10 @@
     add_library(platform_bl1_interface INTERFACE)
 endif()
 
+if (TFM_PARTITION_CRYPTO)
+    add_library(platform_crypto_keys STATIC)
+endif()
+
 set(PLATFORM_DIR ${CMAKE_CURRENT_LIST_DIR})
 
 add_subdirectory(ext/target/${TFM_PLATFORM} target)
@@ -58,7 +62,6 @@
         ext/common/tfm_hal_memory_symbols.c
         $<$<BOOL:${PLATFORM_DEFAULT_ATTEST_HAL}>:ext/common/template/attest_hal.c>
         $<$<BOOL:${PLATFORM_DEFAULT_NV_COUNTERS}>:ext/common/template/nv_counters.c>
-        $<$<AND:$<BOOL:${TFM_PARTITION_CRYPTO}>,$<BOOL:${PLATFORM_DEFAULT_CRYPTO_KEYS}>>:ext/common/template/crypto_keys.c>
         $<$<BOOL:${PLATFORM_DEFAULT_ROTPK}>:ext/common/template/tfm_rotpk.c>
         $<$<BOOL:${PLATFORM_DEFAULT_NV_SEED}>:ext/common/template/crypto_nv_seed.c>
         $<$<AND:$<NOT:$<BOOL:${SYMMETRIC_INITIAL_ATTESTATION}>>,$<BOOL:${TEST_S_ATTESTATION}>>:ext/common/template/tfm_initial_attest_pub_key.c>
@@ -88,6 +91,7 @@
         $<$<BOOL:${PLATFORM_DEFAULT_ATTEST_HAL}>:tfm_sprt>
         $<$<BOOL:${TFM_PARTITION_CRYPTO}>:crypto_service_mbedcrypto>
         $<$<BOOL:${TFM_PARTITION_INITIAL_ATTESTATION}>:tfm_attestation_defs>
+        platform_crypto_keys
 )
 
 target_compile_definitions(platform_s
@@ -103,6 +107,7 @@
         $<$<STREQUAL:${CONFIG_TFM_FP},soft>:CONFIG_TFM_FP=0>
         $<$<BOOL:${CONFIG_TFM_LAZY_STACKING}>:CONFIG_TFM_LAZY_STACKING>
         $<$<BOOL:${CONFIG_TFM_ENABLE_FPU}>:CONFIG_TFM_ENABLE_FPU>
+        $<$<BOOL:${PLATFORM_DEFAULT_CRYPTO_KEYS}>:PLATFORM_DEFAULT_CRYPTO_KEYS>
     PRIVATE
         $<$<BOOL:${SYMMETRIC_INITIAL_ATTESTATION}>:SYMMETRIC_INITIAL_ATTESTATION>
         $<$<BOOL:${PLATFORM_DEFAULT_OTP}>:PLATFORM_DEFAULT_OTP>
@@ -118,6 +123,35 @@
         ${COMPILER_CP_FLAG}
 )
 
+#========================= Platform Crypto Keys ===============================#
+
+
+if(TFM_PARTITION_CRYPTO)
+    target_sources(platform_crypto_keys
+        PRIVATE
+            $<$<BOOL:${PLATFORM_DEFAULT_CRYPTO_KEYS}>:ext/common/template/crypto_keys.c>
+    )
+
+    target_link_libraries(platform_crypto_keys
+        PRIVATE
+            crypto_service_mbedcrypto
+            tfm_psa_rot_partition_crypto
+            platform_region_defs
+            tfm_partition_defs
+    )
+
+    target_include_directories(platform_crypto_keys
+        PRIVATE
+            $<$<BOOL:${PLATFORM_DEFAULT_CRYPTO_KEYS}>:${CMAKE_SOURCE_DIR}/interface/include/crypto_keys>
+            $<$<BOOL:${PLATFORM_DEFAULT_CRYPTO_KEYS}>:${CMAKE_BINARY_DIR}/generated/interface/include>
+    )
+
+    target_compile_definitions(platform_crypto_keys
+        PRIVATE
+            $<$<BOOL:${SYMMETRIC_INITIAL_ATTESTATION}>:SYMMETRIC_INITIAL_ATTESTATION>
+    )
+endif()
+
 #========================= Platform BL2 =======================================#
 if(BL2)
     #TODO import policy
diff --git a/platform/ext/common/template/crypto_keys.c b/platform/ext/common/template/crypto_keys.c
index d4a48cb..82ebc8d 100644
--- a/platform/ext/common/template/crypto_keys.c
+++ b/platform/ext/common/template/crypto_keys.c
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2017-2021 Arm Limited. All rights reserved.
+ * Copyright (c) 2017-2022 Arm Limited. All rights reserved.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -14,164 +14,202 @@
  * limitations under the License.
  */
 
+#include "tfm_plat_crypto_keys.h"
+
+#include "tfm_builtin_key_ids.h"
+
 #include <stddef.h>
 #include <string.h>
 
-#include "psa/crypto_types.h"
-#include "tfm_plat_crypto_keys.h"
+#include "region_defs.h"
+#include "cmsis_compiler.h"
 #include "tfm_plat_otp.h"
-#include "mbedtls/hkdf.h"
+#include "psa_manifest/pid.h"
+#include "tfm_builtin_key_loader.h"
 
-#ifdef CRYPTO_HW_ACCELERATOR
-#include "crypto_hw.h"
-#endif /* CRYPTO_HW_ACCELERATOR */
+#define TFM_NS_PARTITION_ID -1
 
-enum tfm_plat_err_t tfm_plat_get_huk_derived_key(const uint8_t *label,
-                                                 size_t label_size,
-                                                 const uint8_t *context,
-                                                 size_t context_size,
-                                                 uint8_t *key,
-                                                 size_t key_size)
+#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
+
+enum tfm_plat_err_t tfm_plat_builtin_key_get_usage(psa_key_id_t key_id,
+                                                   mbedtls_key_owner_id_t user,
+                                                   psa_key_usage_t *usage)
 {
-#ifdef CRYPTO_HW_ACCELERATOR
-    return crypto_hw_accelerator_huk_derive_key(label, label_size, context,
-                                                context_size, key, key_size);
-#else
-    uint8_t huk_buf[32];
-    enum tfm_plat_err_t err;
-    int mbedtls_err;
+    *usage = 0;
 
-    if (key == NULL) {
-        return TFM_PLAT_ERR_INVALID_INPUT;
+    switch (key_id) {
+    case TFM_BUILTIN_KEY_ID_HUK:
+        /* Allow access to all partitions */
+        *usage = PSA_KEY_USAGE_DERIVE;
+        break;
+    case TFM_BUILTIN_KEY_ID_IAK:
+        switch(user) {
+#ifdef TFM_PARTITION_INITIAL_ATTESTATION
+        case TFM_SP_INITIAL_ATTESTATION:
+            *usage = PSA_KEY_USAGE_SIGN_HASH;
+#ifdef SYMMETRIC_INITIAL_ATTESTATION
+            /* Needed to calculate the instance ID */
+            *usage |= PSA_KEY_USAGE_EXPORT;
+#endif /* SYMMETRIC_INITIAL_ATTESTATION */
+            break;
+#ifdef TFM_PARTITION_TEST_SECURE_SERVICES
+        /* So that the tests can validate created tokens */
+        case TFM_SP_SECURE_TEST_PARTITION:
+        case TFM_NS_PARTITION_ID:
+            *usage = PSA_KEY_USAGE_VERIFY_HASH;
+            break;
+#endif /* TFM_PARTITION_TEST_SECURE_SERVICES */
+#endif /* TFM_PARTITION_INITIAL_ATTESTATION */
+        default:
+            return TFM_PLAT_ERR_NOT_PERMITTED;
+        }
+        break;
+    default:
+        return TFM_PLAT_ERR_UNSUPPORTED;
     }
 
-    if (label == NULL && label_size != 0) {
-        return TFM_PLAT_ERR_INVALID_INPUT;
-    }
-
-    if (context == NULL && context_size != 0) {
-        return TFM_PLAT_ERR_INVALID_INPUT;
-    }
-
-    err = tfm_plat_otp_read(PLAT_OTP_ID_HUK, sizeof(huk_buf), huk_buf);
-    if (err != TFM_PLAT_ERR_SUCCESS) {
-        goto out;
-    }
-
-    mbedtls_err = mbedtls_hkdf(mbedtls_md_info_from_type(MBEDTLS_MD_SHA256),
-                               label, label_size, huk_buf, sizeof(huk_buf),
-                               context, context_size, key, key_size);
-    if (mbedtls_err) {
-        err = TFM_PLAT_ERR_SYSTEM_ERR;
-        goto out;
-    }
-
-out:
-    memset(huk_buf, 0, sizeof(huk_buf));
-
-    return err;
-#endif /* CRYPTO_HW_ACCELERATOR */
+    return TFM_PLAT_ERR_SUCCESS;
 }
 
+enum tfm_plat_err_t tfm_plat_builtin_key_get_lifetime_and_slot(
+    mbedtls_svc_key_id_t key_id,
+    psa_key_lifetime_t *lifetime,
+    psa_drv_slot_number_t *slot_number)
+{
+    switch (MBEDTLS_SVC_KEY_ID_GET_KEY_ID(key_id)) {
+    case TFM_BUILTIN_KEY_ID_HUK:
+        *slot_number = TFM_BUILTIN_KEY_SLOT_HUK;
+        *lifetime = PSA_KEY_LIFETIME_FROM_PERSISTENCE_AND_LOCATION(PSA_KEY_LIFETIME_PERSISTENT,
+                                                                   TFM_BUILTIN_KEY_LOADER_KEY_LOCATION);
+        break;
+#ifdef TFM_PARTITION_INITIAL_ATTESTATION
+    case TFM_BUILTIN_KEY_ID_IAK:
+        *slot_number = TFM_BUILTIN_KEY_SLOT_IAK;
+        *lifetime = PSA_KEY_LIFETIME_FROM_PERSISTENCE_AND_LOCATION(PSA_KEY_LIFETIME_PERSISTENT,
+                                                                   TFM_BUILTIN_KEY_LOADER_KEY_LOCATION);
+#endif /* TFM_PARTITION_INITIAL_ATTESTATION */
+        break;
+    default:
+        return TFM_PLAT_ERR_UNSUPPORTED;
+    }
+
+    return TFM_PLAT_ERR_SUCCESS;
+}
+
+static enum tfm_plat_err_t tfm_plat_get_huk(uint8_t *buf, size_t buf_len,
+                                            size_t *key_len,
+                                            size_t *key_bits,
+                                            psa_algorithm_t *algorithm,
+                                            psa_key_type_t *type)
+{
+    enum tfm_plat_err_t err;
+
+    err = tfm_plat_otp_read(PLAT_OTP_ID_HUK, buf_len, buf);
+    if (err != TFM_PLAT_ERR_SUCCESS) {
+        return err;
+    }
+
+    err = tfm_plat_otp_get_size(PLAT_OTP_ID_HUK, key_len);
+    if (err != TFM_PLAT_ERR_SUCCESS) {
+        return err;
+    }
+
+    *key_bits = *key_len * 8;
+    *algorithm = PSA_ALG_HKDF(PSA_ALG_SHA_256);
+    *type = PSA_KEY_TYPE_DERIVE;
+
+    return TFM_PLAT_ERR_SUCCESS;
+}
+
+#ifdef TFM_PARTITION_INITIAL_ATTESTATION
+static enum tfm_plat_err_t tfm_plat_get_iak(uint8_t *buf, size_t buf_len,
+                                            size_t *key_len,
+                                            size_t *key_bits,
+                                            psa_algorithm_t *algorithm,
+                                            psa_key_type_t *type)
+{
+    enum tfm_plat_err_t err;
+#ifndef SYMMETRIC_INITIAL_ATTESTATION
+    psa_ecc_family_t curve_type;
+#endif /* SYMMETRIC_INITIAL_ATTESTATION */
+
+    err = tfm_plat_otp_read(PLAT_OTP_ID_IAK_LEN,
+                            sizeof(size_t), (uint8_t*)key_len);
+    if(err != TFM_PLAT_ERR_SUCCESS) {
+        return err;
+    }
+    *key_bits = *key_len * 8;
+
 #ifdef SYMMETRIC_INITIAL_ATTESTATION
-enum tfm_plat_err_t tfm_plat_get_symmetric_iak(uint8_t *key_buf,
-                                               size_t buf_len,
-                                               size_t *key_len,
-                                               psa_algorithm_t *key_alg)
-{
-    enum tfm_plat_err_t err;
-
-    if (key_buf == NULL || key_len == NULL) {
-        return TFM_PLAT_ERR_INVALID_INPUT;
-    }
-
-    err = tfm_plat_otp_read(PLAT_OTP_ID_IAK, buf_len, key_buf);
-    if(err != TFM_PLAT_ERR_SUCCESS) {
-        return err;
-    }
-
-    err = tfm_plat_otp_read(PLAT_OTP_ID_IAK_LEN, sizeof(size_t),
-                            (uint8_t*)key_len);
-    if(err != TFM_PLAT_ERR_SUCCESS) {
-        return err;
-    }
-
     err = tfm_plat_otp_read(PLAT_OTP_ID_IAK_TYPE,
-                            sizeof(psa_algorithm_t), (uint8_t*)key_alg);
+                            sizeof(psa_algorithm_t), (uint8_t*)algorithm);
     if(err != TFM_PLAT_ERR_SUCCESS) {
         return err;
     }
 
-    return TFM_PLAT_ERR_SUCCESS;
-
-}
-
-enum tfm_plat_err_t tfm_plat_get_symmetric_iak_id(void *kid_buf,
-                                                  size_t buf_len,
-                                                  size_t *kid_len)
-{
-    enum tfm_plat_err_t err;
-    size_t otp_size;
-
-    if (kid_buf == NULL || kid_len == NULL) {
-        return TFM_PLAT_ERR_INVALID_INPUT;
-    }
-
-    err = tfm_plat_otp_read(PLAT_OTP_ID_IAK_ID, buf_len, kid_buf);
-    if(err != TFM_PLAT_ERR_SUCCESS) {
-        return err;
-    }
-
-    err =  tfm_plat_otp_get_size(PLAT_OTP_ID_IAK_ID, &otp_size);
-    if(err != TFM_PLAT_ERR_SUCCESS) {
-        return err;
-    }
-
-    *kid_len = strlen(kid_buf) <= otp_size ? strlen(kid_buf) : otp_size;
-
-    return TFM_PLAT_ERR_SUCCESS;
-}
+    *type = PSA_KEY_TYPE_HMAC;
 #else /* SYMMETRIC_INITIAL_ATTESTATION */
-enum tfm_plat_err_t
-tfm_plat_get_initial_attest_key(uint8_t          *key_buf,
-                                uint32_t          size,
-                                struct ecc_key_t *ecc_key,
-                                psa_ecc_family_t *curve_type)
+    err = tfm_plat_otp_read(PLAT_OTP_ID_IAK_TYPE, sizeof(psa_ecc_family_t),
+                            &curve_type);
+    if(err != TFM_PLAT_ERR_SUCCESS) {
+        return err;
+    }
+
+    *algorithm = PSA_ALG_ECDSA(PSA_ALG_SHA_256);
+    *type = PSA_KEY_TYPE_ECC_KEY_PAIR(curve_type);
+#endif /* SYMMETRIC_INITIAL_ATTESTATION */
+
+    return tfm_plat_otp_read(PLAT_OTP_ID_IAK, *key_len, buf);
+}
+#endif /* TFM_PARTITION_INITIAL_ATTESTATION */
+
+enum tfm_plat_err_t tfm_plat_load_builtin_keys(void)
 {
-    uint32_t key_size;
-    enum tfm_plat_err_t err;
+    psa_status_t err;
+    mbedtls_svc_key_id_t key_id;
+    psa_key_attributes_t attr = PSA_KEY_ATTRIBUTES_INIT;
+    enum tfm_plat_err_t plat_err;
+    uint8_t buf[32];
+    size_t key_len;
+    size_t key_bits;
+    psa_algorithm_t algorithm;
+    psa_key_type_t type;
 
-    if (key_buf == NULL || ecc_key == NULL || curve_type == NULL) {
-        return TFM_PLAT_ERR_INVALID_INPUT;
+    /* HUK */
+    plat_err = tfm_plat_get_huk(buf, sizeof(buf), &key_len, &key_bits,
+                                &algorithm, &type);
+    if (plat_err != TFM_PLAT_ERR_SUCCESS) {
+        return plat_err;
+    }
+    key_id.MBEDTLS_PRIVATE(key_id) = TFM_BUILTIN_KEY_ID_HUK;
+    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);
+    err = tfm_builtin_key_loader_load_key(buf, key_len, &attr);
+    if (err != PSA_SUCCESS) {
+        return TFM_PLAT_ERR_SYSTEM_ERR;
     }
 
-    err = tfm_plat_otp_read(PLAT_OTP_ID_IAK, size, key_buf);
-    if(err != TFM_PLAT_ERR_SUCCESS) {
-        return err;
+#ifdef TFM_PARTITION_INITIAL_ATTESTATION
+    /* IAK */
+    plat_err = tfm_plat_get_iak(buf, sizeof(buf), &key_len, &key_bits,
+                                &algorithm, &type);
+    if (plat_err != TFM_PLAT_ERR_SUCCESS) {
+        return TFM_PLAT_ERR_SYSTEM_ERR;
     }
-
-    err =  tfm_plat_otp_read(PLAT_OTP_ID_IAK_LEN, sizeof(key_size),
-                             (uint8_t*)&key_size);
-    if(err != TFM_PLAT_ERR_SUCCESS) {
-        return err;
+    key_id.MBEDTLS_PRIVATE(key_id) = TFM_BUILTIN_KEY_ID_IAK;
+    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);
+    err = tfm_builtin_key_loader_load_key(buf, key_len, &attr);
+    if (err != PSA_SUCCESS) {
+        return TFM_PLAT_ERR_SYSTEM_ERR;
     }
-
-    /* Set the EC curve type which the key belongs to */
-    err = tfm_plat_otp_read(PLAT_OTP_ID_IAK_TYPE,
-                            sizeof(psa_ecc_family_t), curve_type);
-    if(err != TFM_PLAT_ERR_SUCCESS) {
-        return err;
-    }
-
-    /* Copy the private key to the buffer, it MUST be present */
-    ecc_key->priv_key = key_buf;
-    ecc_key->priv_key_size = key_size;
-
-    ecc_key->pubx_key = NULL;
-    ecc_key->pubx_key_size = 0;
-    ecc_key->puby_key = NULL;
-    ecc_key->puby_key_size = 0;
+#endif /* TFM_PARTITION_INITIAL_ATTESTATION */
 
     return TFM_PLAT_ERR_SUCCESS;
 }
-#endif /* SYMMETRIC_INITIAL_ATTESTATION */
diff --git a/platform/include/tfm_plat_crypto_keys.h b/platform/include/tfm_plat_crypto_keys.h
index 5ceb591..4754855 100644
--- a/platform/include/tfm_plat_crypto_keys.h
+++ b/platform/include/tfm_plat_crypto_keys.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2017-2020, Arm Limited. All rights reserved.
+ * Copyright (c) 2017-2022, Arm Limited. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  *
@@ -12,166 +12,52 @@
  *       SoC.
  */
 
+#define MBEDTLS_PSA_CRYPTO_KEY_ID_ENCODES_OWNER
+
+#include "tfm_mbedcrypto_include.h"
+
+#include "tfm_plat_defs.h"
+
 #include <stddef.h>
 #include <stdint.h>
-#include "tfm_plat_defs.h"
-#include "psa/crypto.h"
 
 #ifdef __cplusplus
 extern "C" {
 #endif
 
 /**
- * Structure definition to carry pointer and size information about an Elliptic
- * curve key which is stored in a buffer(key_buf) in raw format (without
- * encoding):
- *   - priv_key       Base address of the private key in key_buf. It must be
- *                    present on the device.
- *   - priv_key_size  Size of the private key in bytes.
- *   - pubx_key       Base address of x-coordinate of the public key in key_buf.
- *                    It can be empty, because it can be recomputed based on
- *                    private key.
- *   - pubx_key_size  Length of x-coordinate of the public key in key_buf.
- *                    It can be empty, because it can be recomputed based on
- *                    private key.
- *   - puby_key       Base address of y-coordinate of the public key in key_buf.
- *                    It can be empty, because either it can be recomputed based
- *                    on private key or some curve type works without it.
- *   - puby_key_size  Length of y-coordinate of the public key in key_buf.
- */
-struct ecc_key_t {
-    uint8_t  *priv_key;
-    uint32_t  priv_key_size;
-    uint8_t  *pubx_key;
-    uint32_t  pubx_key_size;
-    uint8_t  *puby_key;
-    uint32_t  puby_key_size;
-};
-
-#define ROTPK_HASH_LEN (32u) /* SHA256 */
-
-/**
- * Structure to store the hard-coded (embedded in secure firmware) hash of ROTPK
- * for firmware authentication.
+ * \brief Gets key usage for a given builtin key ID and owner.
  *
- * \note Just temporary solution, hard-coded key-hash values in firmware is not
- *       suited for use in production!
- */
-struct tfm_plat_rotpk_t {
-    const uint8_t *key_hash;
-    const uint8_t  hash_len;
-};
-
-/**
- * \brief Gets key material derived from the hardware unique key.
- *
- * \param[in]  label         Label for KDF
- * \param[in]  label_size    Size of the label
- * \param[in]  context       Context for KDF
- * \param[in]  context_size  Size of the context
- * \param[out] key           Buffer to output the derived key material
- * \param[in]  key_size      Requested size of the derived key material and
- *                           minimum size of the key buffer
+ * \param[in]  key_id        ID of key
+ * \param[in]  user          Which user to get the usage permissions for
+ * \param[out] usage         The permissions that the given user has for the key
  *
  * \return Returns error code specified in \ref tfm_plat_err_t
  */
-enum tfm_plat_err_t tfm_plat_get_huk_derived_key(const uint8_t *label,
-                                                 size_t label_size,
-                                                 const uint8_t *context,
-                                                 size_t context_size,
-                                                 uint8_t *key,
-                                                 size_t key_size);
+enum tfm_plat_err_t tfm_plat_builtin_key_get_usage(psa_key_id_t key_id,
+                                                   mbedtls_key_owner_id_t user,
+                                                   psa_key_usage_t *usage);
 
-#ifdef SYMMETRIC_INITIAL_ATTESTATION
 /**
- * \brief Get the symmetric Initial Attestation Key (IAK)
+ * \brief Gets key lifetime and slot number for a given builtin key ID.
  *
- * The device MUST contain a symmetric IAK, which is used to sign the token.
- * So far only HMAC is supported in symmetric key algorithm based Initial
- * Attestation.
- * Keys must be provided in raw format, just binary data without any encoding
- * (DER, COSE). Caller provides a buffer to copy all the raw data.
- *
- * \param[out]  key_buf     Buffer to store the initial attestation key.
- * \param[in]   buf_len     The length of buffer.
- * \param[out]  key_len     Buffer to carry the length of the initial
- *                          attestation key.
- * \param[out]  key_alg     The key algorithm. Only HMAC is supported so far.
+ * \param[in]  key_id        ID of key
+ * \param[out] lifetime      Lifetime and storage location of the key
+ * \param[out] slot_number   Index of the slot which the key is stored in
  *
  * \return Returns error code specified in \ref tfm_plat_err_t
  */
-enum tfm_plat_err_t tfm_plat_get_symmetric_iak(uint8_t *key_buf,
-                                               size_t buf_len,
-                                               size_t *key_len,
-                                               psa_algorithm_t *key_alg);
+enum tfm_plat_err_t tfm_plat_builtin_key_get_lifetime_and_slot(
+    mbedtls_svc_key_id_t key_id,
+    psa_key_lifetime_t *lifetime,
+    psa_drv_slot_number_t *slot_number);
 
-#ifdef INCLUDE_COSE_KEY_ID
 /**
- * \brief Get the key identifier of the symmetric Initial Attestation Key as the
- *        'kid' parameter in COSE Header.
- *
- * \note This `kid` parameter is included in COSE Header. Please don't confuse
- *       it with that `kid` in COSE_Key structure.
- *
- * \param[out] kid_buf  The buffer to be written with key id
- * \param[in]  buf_len  The length of kid_buf
- * \param[out] kid_len  The length of key id
- *
- * \return Returns error code specified in \ref tfm_plat_err_t.
- */
-enum tfm_plat_err_t tfm_plat_get_symmetric_iak_id(void *kid_buf,
-                                                  size_t buf_len,
-                                                  size_t *kid_len);
-#endif
-#else /* SYMMETRIC_INITIAL_ATTESTATION */
-/**
- * \brief Get the initial attestation key
- *
- * The device MUST contain an initial attestation key, which is used to sign the
- * token. Initial attestation service supports elliptic curve signing
- * algorithms. Device maker can decide whether store only the private key on the
- * device or store both (public and private) key. Public key can be recomputed
- * based on private key. Keys must be provided in raw format, just binary data
- * without any encoding (DER, COSE). Caller provides a buffer to copy all the
- * available key components to there. Key components must be copied after
- * each other to the buffer. The base address and the length of each key
- * component must be indicating in the corresponding field of ecc_key
- * (\ref struct ecc_key_t).
- * Curve_type indicates to which curve belongs the key.
- *
- *
- * Keys must be provided in
- *
- * \param[in/out]  key_buf     Buffer to store the initial attestation key.
- * \param[in]      size        Size of the buffer.
- * \param[out]     ecc_key     A structure to carry pointer and size information
- *                             about the initial attestation key, which is
- *                             stored in key_buf.
- * \param[out]     curve_type  The type of the EC curve, which the key belongs
- *                             to according to \ref psa_ecc_family_t
+ * \brief Load all builtin keys.
  *
  * \return Returns error code specified in \ref tfm_plat_err_t
  */
-enum tfm_plat_err_t
-tfm_plat_get_initial_attest_key(uint8_t          *key_buf,
-                                uint32_t          size,
-                                struct ecc_key_t *ecc_key,
-                                psa_ecc_family_t *curve_type);
-#endif /* SYMMETRIC_INITIAL_ATTESTATION */
-
-/**
- * \brief Get the hash of the corresponding Root of Trust Public Key for
- *        firmware authentication.
- *
- * \param[in]      image_id         The identifier of firmware image
- * \param[out]     rotpk_hash       Buffer to store the key-hash in
- * \param[in,out]  rotpk_hash_size  As input the size of the buffer. As output
- *                                  the actual key-hash length.
- */
-enum tfm_plat_err_t
-tfm_plat_get_rotpk_hash(uint8_t image_id,
-                        uint8_t *rotpk_hash,
-                        uint32_t *rotpk_hash_size);
+enum tfm_plat_err_t tfm_plat_load_builtin_keys(void);
 
 #ifdef __cplusplus
 }
diff --git a/platform/include/tfm_plat_defs.h b/platform/include/tfm_plat_defs.h
index 6a2a14d..bd74d13 100644
--- a/platform/include/tfm_plat_defs.h
+++ b/platform/include/tfm_plat_defs.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2017-2021, Arm Limited. All rights reserved.
+ * Copyright (c) 2017-2022, Arm Limited. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  *
@@ -21,6 +21,7 @@
     TFM_PLAT_ERR_MAX_VALUE = 0x55A3,
     TFM_PLAT_ERR_INVALID_INPUT = 0xA3C5,
     TFM_PLAT_ERR_UNSUPPORTED = 0xC35A,
+    TFM_PLAT_ERR_NOT_PERMITTED = 0xC5A3,
     /* Following entry is only to ensure the error code of int size */
     TFM_PLAT_ERR_FORCE_INT_SIZE = INT_MAX
 };
diff --git a/platform/include/tfm_plat_rotpk.h b/platform/include/tfm_plat_rotpk.h
new file mode 100644
index 0000000..3bc93f5
--- /dev/null
+++ b/platform/include/tfm_plat_rotpk.h
@@ -0,0 +1,54 @@
+/*
+ * Copyright (c) 2017-2022, Arm Limited. All rights reserved.
+ *
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ */
+
+#ifndef __TFM_PLAT_ROTPK_H__
+#define __TFM_PLAT_ROTPK_H__
+/**
+ * \note The interfaces defined in this file must be implemented for each
+ *       SoC.
+ */
+
+#include <stddef.h>
+#include <stdint.h>
+#include "tfm_plat_defs.h"
+#include "psa/crypto.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+/**
+ * Structure to store the hard-coded (embedded in secure firmware) hash of ROTPK
+ * for firmware authentication.
+ *
+ * \note Just temporary solution, hard-coded key-hash values in firmware is not
+ *       suited for use in production!
+ */
+struct tfm_plat_rotpk_t {
+    const uint8_t *key_hash;
+    const uint8_t  hash_len;
+};
+
+/**
+ * \brief Get the hash of the corresponding Root of Trust Public Key for
+ *        firmware authentication.
+ *
+ * \param[in]      image_id         The identifier of firmware image
+ * \param[out]     rotpk_hash       Buffer to store the key-hash in
+ * \param[in,out]  rotpk_hash_size  As input the size of the buffer. As output
+ *                                  the actual key-hash length.
+ */
+enum tfm_plat_err_t
+tfm_plat_get_rotpk_hash(uint8_t image_id,
+                        uint8_t *rotpk_hash,
+                        uint32_t *rotpk_hash_size);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif /* __TFM_PLAT_ROTPK_H__ */
diff --git a/platform/ns/CMakeLists.txt b/platform/ns/CMakeLists.txt
index 6fffa7f..4dc2b66 100755
--- a/platform/ns/CMakeLists.txt
+++ b/platform/ns/CMakeLists.txt
@@ -39,6 +39,7 @@
         $<$<STREQUAL:${CONFIG_TFM_FP},soft>:CONFIG_TFM_FP=0>
         $<$<BOOL:${TEST_NS_FPU}>:TEST_NS_FPU>
         $<$<BOOL:${CONFIG_TFM_ENABLE_FPU}>:CONFIG_TFM_ENABLE_FPU>
+        $<$<BOOL:${PLATFORM_DEFAULT_CRYPTO_KEYS}>:PLATFORM_DEFAULT_CRYPTO_KEYS>
     PRIVATE
         $<$<BOOL:${TEST_NS_SLIH_IRQ}>:TEST_NS_SLIH_IRQ>
 )
diff --git a/secure_fw/partitions/crypto/CMakeLists.txt b/secure_fw/partitions/crypto/CMakeLists.txt
index c792693..7721daf 100644
--- a/secure_fw/partitions/crypto/CMakeLists.txt
+++ b/secure_fw/partitions/crypto/CMakeLists.txt
@@ -27,6 +27,8 @@
         crypto_key_derivation.c
         crypto_key_management.c
         crypto_rng.c
+        tfm_mbedcrypto_builtin_keys.c
+        $<$<BOOL:CRYPTO_TFM_BUILTIN_KEYS_DRIVER>:psa_driver_api/tfm_builtin_key_loader.c>
 )
 
 # The generated sources
@@ -72,6 +74,9 @@
         $<$<BOOL:${CRYPTO_ASYM_SIGN_MODULE_DISABLED}>:TFM_CRYPTO_ASYM_SIGN_MODULE_DISABLED>
         $<$<BOOL:${CRYPTO_ASYM_ENCRYPT_MODULE_DISABLED}>:TFM_CRYPTO_ASYM_ENCRYPT_MODULE_DISABLED>
         $<$<BOOL:${CRYPTO_KEY_DERIVATION_MODULE_DISABLED}>:TFM_CRYPTO_KEY_DERIVATION_MODULE_DISABLED>
+        MBEDTLS_PSA_CRYPTO_DRIVERS
+        MBEDTLS_PSA_CRYPTO_BUILTIN_KEYS
+        $<$<BOOL:CRYPTO_BUILTIN_KEYS>:PSA_CRYPTO_DRIVER_TFM_BUILTIN_KEY>
     PRIVATE
         $<$<BOOL:${CRYPTO_ENGINE_BUF_SIZE}>:TFM_CRYPTO_ENGINE_BUF_SIZE=${CRYPTO_ENGINE_BUF_SIZE}>
         $<$<BOOL:${CRYPTO_CONC_OPER_NUM}>:TFM_CRYPTO_CONC_OPER_NUM=${CRYPTO_CONC_OPER_NUM}>
@@ -142,6 +147,10 @@
         $<$<OR:$<STREQUAL:${TFM_SYSTEM_ARCHITECTURE},armv8-m.base>,$<STREQUAL:${TFM_SYSTEM_ARCHITECTURE},armv6-m>>:MULADDC_CANNOT_USE_R7>
         $<$<BOOL:${CRYPTO_NV_SEED}>:CRYPTO_NV_SEED>
         $<$<BOOL:${PLATFORM_DEFAULT_NV_SEED}>:PLATFORM_DEFAULT_NV_SEED>
+        $<$<BOOL:${PLATFORM_DEFAULT_CRYPTO_KEYS}>:PLATFORM_DEFAULT_CRYPTO_KEYS>
+        MBEDTLS_PSA_CRYPTO_BUILTIN_KEYS
+        MBEDTLS_PSA_CRYPTO_DRIVERS
+        $<$<BOOL:CRYPTO_TFM_BUILTIN_KEYS_DRIVER>:PSA_CRYPTO_DRIVER_TFM_BUILTIN_KEY_LOADER>
 )
 cmake_policy(SET CMP0079 NEW)
 
@@ -175,6 +184,7 @@
 target_include_directories(${MBEDTLS_TARGET_PREFIX}mbedcrypto
     PUBLIC
         ${CMAKE_CURRENT_SOURCE_DIR}
+        ${CMAKE_CURRENT_SOURCE_DIR}/psa_driver_api
 )
 
 # Fix platform_s and crypto_service_mbedcrypto libraries cyclic linking
diff --git a/secure_fw/partitions/crypto/crypto_init.c b/secure_fw/partitions/crypto/crypto_init.c
index 898418e..1af64cc 100644
--- a/secure_fw/partitions/crypto/crypto_init.c
+++ b/secure_fw/partitions/crypto/crypto_init.c
@@ -12,6 +12,7 @@
 #include "tfm_crypto_defs.h"
 #include "tfm_sp_log.h"
 #include "crypto_check_config.h"
+#include "tfm_plat_crypto_keys.h"
 
 /*
  * \brief This Mbed TLS include is needed to initialise the memory allocator
@@ -345,6 +346,7 @@
 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();
@@ -353,7 +355,17 @@
     }
 
     /* Initialise the engine layer */
-    return tfm_crypto_engine_init();
+    status =  tfm_crypto_engine_init();
+    if (status != PSA_SUCCESS) {
+        return status;
+    }
+
+    plat_err = tfm_plat_load_builtin_keys();
+    if (plat_err != TFM_PLAT_ERR_SUCCESS) {
+        return PSA_ERROR_GENERIC_ERROR;
+    }
+
+    return PSA_SUCCESS;
 }
 
 #ifdef TFM_PSA_API
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
new file mode 100644
index 0000000..fd0d58e
--- /dev/null
+++ b/secure_fw/partitions/crypto/psa_driver_api/tfm_builtin_key_loader.c
@@ -0,0 +1,208 @@
+/*
+ * Copyright (c) 2022, 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 "mbedtls/hkdf.h"
+#include "psa_manifest/pid.h"
+#include "tfm_plat_crypto_keys.h"
+
+#include <string.h>
+
+#define TFM_BUILTIN_MAX_KEY_LEN 48
+#define TFM_BUILTIN_MAX_KEYS 8
+
+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;
+};
+
+static struct tfm_builtin_key_t builtin_key_slots[TFM_BUILTIN_MAX_KEYS] = {0};
+
+psa_status_t tfm_builtin_key_loader_load_key(uint8_t *buf, size_t key_len,
+                                             psa_key_attributes_t *attr)
+{
+    psa_status_t err;
+    psa_drv_slot_number_t slot_number;
+    psa_key_lifetime_t lifetime;
+    mbedtls_svc_key_id_t key_id;
+
+    /* 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_PRIVATE(owner) = 0;
+    psa_set_key_id(attr, key_id);
+
+    if (key_len > TFM_BUILTIN_MAX_KEY_LEN) {
+        return PSA_ERROR_BUFFER_TOO_SMALL;
+    }
+
+    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;
+}
+
+psa_status_t tfm_builtin_key_loader_get_key_buffer_size(
+        mbedtls_svc_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;
+    }
+
+    *len = builtin_key_slots[slot_number].key_len;
+    return PSA_SUCCESS;
+}
+
+static psa_status_t builtin_key_get_attributes(
+        struct tfm_builtin_key_t *key_slot, psa_key_attributes_t *attr)
+{
+    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);
+
+    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;
+    }
+
+    psa_set_key_id(attr, key_id);
+    psa_set_key_lifetime(attr, lifetime);
+    psa_set_key_usage_flags(attr, usage);
+
+    return PSA_SUCCESS;
+}
+
+static psa_status_t derive_subkey_into_buffer(
+        struct tfm_builtin_key_t *key_slot, mbedtls_key_owner_id_t owner,
+        uint8_t *key_buffer, size_t key_buffer_size, size_t *key_buffer_length)
+{
+    int mbedtls_err;
+
+#ifdef TFM_PARTITION_TEST_PS
+    /* Hack to allow the PS tests to work, since they directly call
+     * ps_system_prepare from the test partition which would otherwise derive a
+     * different key.
+     */
+    if (owner == TFM_SP_PS_TEST) {
+        owner = TFM_SP_PS;
+    }
+#endif /* TFM_PARTITION_TEST_PS */
+
+    /* FIXME this should be moved to using the PSA APIs once key derivation is
+     * implemented in the PSA driver wrapper. Using the external PSA apis
+     * directly creates a keyslot and we'd need to read the data from it and
+     * then destroy it, so isn't ideal. In order to avoid infinite recursion,
+     * it'll be necessary to add a special case (probably if owner == 0) to make
+     * sure the new PSA derivation request doesn't end up back here.
+     */
+    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,
+                               key_buffer_size);
+    if (mbedtls_err) {
+        return PSA_ERROR_GENERIC_ERROR;
+    }
+
+    *key_buffer_length = key_buffer_size;
+
+    return PSA_SUCCESS;
+}
+
+static psa_status_t builtin_key_copy_to_buffer(
+        struct tfm_builtin_key_t *key_slot, uint8_t *key_buffer,
+        size_t key_buffer_size, size_t *key_buffer_length)
+{
+    memcpy(key_buffer, key_slot->key, key_slot->key_len);
+    *key_buffer_length = key_slot->key_len;
+
+    return PSA_SUCCESS;
+}
+
+psa_status_t tfm_builtin_key_loader_get_key_buffer(
+        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;
+
+    if ((uint32_t)slot_number >= TFM_BUILTIN_MAX_KEYS) {
+        return PSA_ERROR_DOES_NOT_EXIST;
+    }
+
+    key_slot = &builtin_key_slots[slot_number];
+
+    if (!key_slot->is_loaded) {
+        return PSA_ERROR_DOES_NOT_EXIST;
+    }
+
+    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;
+    }
+
+    key_id = psa_get_key_id(attributes);
+
+    /* 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.
+     * 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.
+     */
+    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),
+                                        key_buffer, key_buffer_size,
+                                        key_buffer_length);
+    } else {
+        err = builtin_key_copy_to_buffer(key_slot, key_buffer, key_buffer_size,
+                                         key_buffer_length);
+    }
+
+    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
new file mode 100644
index 0000000..1220f2a
--- /dev/null
+++ b/secure_fw/partitions/crypto/psa_driver_api/tfm_builtin_key_loader.h
@@ -0,0 +1,78 @@
+/*
+ * Copyright (c) 2022, Arm Limited. All rights reserved.
+ *
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ */
+
+#ifndef TFM_BUILTIN_KEY_LOADER_H
+#define TFM_BUILTIN_KEY_LOADER_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,
+    TFM_BUILTIN_KEY_SLOT_IAK,
+    TFM_BUILTIN_KEY_SLOT_DAK_SEED,
+    TFM_BUILTIN_KEY_SLOT_MAX,
+};
+#else
+#include "platform_builtin_key_loader_ids.h"
+#endif
+
+/**
+ * \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.
+ *
+ * \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);
+
+/**
+ * \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[out] len     The length of the key.
+ *
+ * \return Returns error code specified in \ref psa_status_t
+ */
+psa_status_t tfm_builtin_key_loader_get_key_buffer_size(
+        mbedtls_svc_key_id_t key_id, size_t *len);
+
+/**
+ * \brief Returns the attributes and key material of a key from the builtin
+ *        driver
+ *
+ * \note This function is called by the psa crypto driver wrapper.
+ *
+ * \param[in] slot_number        The slot of the key
+ * \param[out] attributes        The attributes of the key.
+ * \param[out] key_buffer        The buffer to output the key material into.
+ * \param[in] key_buffer_size    The size of the key material buffer.
+ * \param[out] key_buffer_length The length of the key material returned.
+ *
+ * \return Returns error code specified in \ref psa_status_t
+ */
+psa_status_t tfm_builtin_key_loader_get_key_buffer(
+        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);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif /* TFM_BUILTIN_KEY_LOADER_H */
diff --git a/secure_fw/partitions/crypto/tfm_mbedcrypto_builtin_keys.c b/secure_fw/partitions/crypto/tfm_mbedcrypto_builtin_keys.c
new file mode 100644
index 0000000..49ce439
--- /dev/null
+++ b/secure_fw/partitions/crypto/tfm_mbedcrypto_builtin_keys.c
@@ -0,0 +1,31 @@
+/*
+ * Copyright (c) 2022, Arm Limited. All rights reserved.
+ *
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ */
+
+#include "tfm_mbedcrypto_include.h"
+#include "tfm_crypto_defs.h"
+#include "tfm_plat_crypto_keys.h"
+
+/* This function is required by MbedTLS to enable builtin key support, it maps
+ * builtin key IDs to cryptographic drivers and slots. The actual data is
+ * deferred to a platform function, as different platforms may have different
+ * key storage capabilities.
+ */
+psa_status_t mbedtls_psa_platform_get_builtin_key(
+    mbedtls_svc_key_id_t key_id,
+    psa_key_lifetime_t *lifetime,
+    psa_drv_slot_number_t *slot_number)
+{
+    enum tfm_plat_err_t plat_err;
+
+    plat_err = tfm_plat_builtin_key_get_lifetime_and_slot(key_id, lifetime,
+                                                          slot_number);
+    if (plat_err != TFM_PLAT_ERR_SUCCESS) {
+        return PSA_ERROR_DOES_NOT_EXIST;
+    }
+
+    return PSA_SUCCESS;
+}