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/docs/design_docs/tfm_builtin_keys.rst b/docs/design_docs/tfm_builtin_keys.rst
index c27d2ca..d68f51f 100644
--- a/docs/design_docs/tfm_builtin_keys.rst
+++ b/docs/design_docs/tfm_builtin_keys.rst
@@ -10,85 +10,141 @@
 Introduction
 ************
 
-TF-M has several keys that are bound to the device itself instead of a secure
-partition. These keys must be accessed through a HAL function, either loading
-them from OTP or another platform-specific location. These keys are henceforth
-referred to as "builtin keys", and include (but are not limited to):
+When TF-M is integrated on a platform, the platform itself can provide several
+keys for usage which are bound to the device instead of being owned by a
+specific secure partition. When those keys are readable by the secure
+processing environment, the platform must provide a function to load these keys
+in the HAL layer, either loading them from OTP or another platform specific
+location or subsystem. These keys are henceforth referred to as "builtin keys",
+and might include (but are not limited to) the following:
 
-1) The Hardware Unique Key (HUK)
-2) The Initial Attestation Key (IAK)
+  1. The Hardware Unique Key (HUK)
+  2. The Initial Attestation Key (IAK)
 
-Currently, the IAK is loaded by the attestation partition as a transient key,
-which requires some key-loading logic to be implemented by that partition. The
-HUK is not loaded in the crypto service, and is instead used by an
-implementation of a TF-M specific KDF algorithm which then loads the key and
-invokes Mbed TLS directly.
+The ``tfm_builtin_key_loader`` component implements a mechanism to discover
+those keys and make them available to the PSA Crypto APIs for usage. Note that
+if a key is stored in a subsystem which can't be read by the secure
+processing environment, a full opaque driver must be used to be able to use
+those keys through the PSA Crypto APIs. This document focuses only on the case
+where the keys can be read by the SPE and processed by the driver described in
+this document.
+
+In TF-M's legacy solution, the IAK is loaded by the attestation service as a
+volatile key, which requires some key-loading logic to be implemented by that
+partition. The HUK is not loaded in the crypto service, and is instead used by
+an implementation of a TF-M specific KDF algorithm which then loads the key and
+invokes Mbed TLS directly. Both solutions are far from ideal as they require an
+effort to load (and duplicate) the keys per each user, or require a dedicated
+code in TF-M that directly interacts with the HAL layer and invokes functions
+from the crypto library bypassing the PSA Crypto interface. The aim of the
+``tfm_builtin_key_loader`` driver is to provide a uniform interface to use the
+builtin keys available in a platform.
+
+Implementing a driver to deal with builtin keys allows to expand the legacy
+solution for dealing with this type of keys in several ways. For example, it
+avoid the need for partitions to implement dedicated mechanisms for probing the
+HAL layer for builtin keys and load them as volatile. It removes the need to
+have implementation specific call flows for deriving keys from the HUK (as that
+requirement now is fulfilled by the driver itself transparently). It allows
+uniform access to the keys also from the NS world (in any case, subject to the
+policy dictated by the platform). Correctly abstracts away details of the key
+handling mechanism from TF-M partitions into the PSA Crypto core key management
+subsystem.
 
 ****************
 PSA builtin keys
 ****************
 
 The PSA Cryptographic API provides a mechanism for accessing keys that are
-stored in platform-specific locations (often hardware accelerators or OTP). One
-of the properties of builtin keys is that they are accessed via a predefined
-handle, which can be leveraged to allow TF-M to define a set of handles for the
-builtin keys that it provides.
+stored in platform-specific locations (often hardware accelerators or OTP). A
+builtin key is assigned a specific key_id, i.e. a handle, which is hardcoded at
+build time and must be selected in the range [MBEDTLS_PSA_KEY_ID_BUILTIN_MIN,
+MBEDLTS_PSA_KEY_ID_BUILTIN_MAX]. A user of the PSA Crypto API can reference
+those keys directly by using these handles. It is up to the platform to specify
+policy restrictions for specific users of the keys, e.g. an NS entity, or a
+secure partition. The PSA Crypto core will then check those policies to grant
+or deny access to the builtin key for that user.
 
-Defining these constant handles allows these keys to be used by secure partition
-and non-secure callers (subject to access policy), via the standard PSA crypto
-interfaces.
+******************************
+PSA cryptoprocessor driver API
+******************************
 
-Ideally, it would be possible to just have PSA builtin keys that are stored in
-crypto service RAM, in the same way that volatile keys are. Mbed TLS does not
-support this and only supports builtin keys as part of the code flow that
-interfaces with hardware accelerators.
+The PSA specification allows the PSA Crypto APIs to defer their operation to an
+accelerator driver, through a mechanism described in the PSA cryptoprocessor
+driver interface specification [1]_. This specification defines the concept of
+PSA Crypto core, i.e. "An implementation of the PSA Cryptography API is
+composed of a core and zero or more drivers". The PSA specification also has
+the concept of storage locations for keys, through the type
+``psa_key_location_t``, which is used to access keys that don't have local
+storage [2]_. This is leveraged mainly by opaque drivers mainly that use keys
+for which the key material is not readable by the PSA crypto core layer.
 
-*********************
-PSA crypto driver API
-*********************
+TF-M defines a software driver called ``tfm_builtin_key_loader`` that provides
+no acceleration but just defines a dedicated key location defined through the
+``TFM_BUILTIN_KEY_LOADER_KEY_LOCATION`` define. By resorting to the entry points
+provided by this driver, the PSA Crypto core slot management subsystem can
+access keys stored in the underlying platform, validate key usage policies and
+allow PSA Crypto APIs to uniformly access builtin keys using the same call flows
+used with traditional local-storage keys.
 
-The PSA crypto driver API allows most PSA Crypto APIs to defer their operation
-to an accelerator driver in preference of the software implementation. It also
-adds the concept of storage locations for keys, which is used to access keys
-stored on hardware accelerators.
+This is implemented by hooking the entry points defined by the driver into the
+``library/psa_crypto_driver_wrappers.c`` file provided by the PSA Crypto core.
+This is currently done manually but eventually could be just autogenerated by
+parsing a description of the driver entry points in the JSON format. These entry
+points are:
 
-The TF-M builtin keys code leverages the PSA crypto driver API by creating a new
-driver that provides no acceleration, only a key storage location. This storage
-location is not backed by hardware, but is instead inside the RAM of the crypto
-partition.
+  1. ``tfm_builtin_key_loader_init``
+  2. ``tfm_builtin_key_loader_get_key_buffer_size``
+  3. ``tfm_builtin_key_loader_get_builtin_key``
 
-This is done by hooking two functions into the
-``library/psa_crypto_driver_wrappers.c`` file. These functions are:
+The call flow for the entry points from the driver wrapper layer is as follows:
 
-1) ``tfm_key_loader_get_builtin_key``
-2) ``tfm_key_loader_get_builtin_key_len``
-
-The flow for these functions being used is:
-
-1) A request is made to a PSA Crypto API that references a key by a key
-   handle.
-2) The PSA Crypto core layer checks that the handle is inside the builtin keys
-    region, and then if the key has not yet been loaded into a transient Mbed
-    TLS keyslot calls ``tfm_plat_builtin_key_get_lifetime_and_slot`` (which is a
-    wrapper around ``mbedtls_psa_platform_get_builtin_key``), which is defined
-    in ``crypto_keys.h``. This function maps each builtin key to a driver, which
-    in most cases is the default ``tfm_builtin_key_loader`` via
-    ``TFM_BUILTIN_KEY_LOADER_KEY_LOCATION``. The function also returns a slot
-    number, which is a driver-specific index to specify the key.
-3) This location and slot index then calls
-   ``psa_driver_wrapper_get_builtin_key``, which for the key location
-   ``TFM_BUILTIN_KEY_LOADER_KEY_LOCATION`` (the new location value that is bound
-   to the TF-M builtin keys driver) calls the previously hooked function
-   ``tfm_key_loader_get_builtin_key``.
-4) This function, along with its counterpart
-   ``tfm_key_loader_get_builtin_key_len``, allow Mbed TLS to copy the
-   key material into an internal keyslot, which is then used whenever further
-   calls to using that same builtin key ID are made.
-
-In order to load the keys into the tfm_key_loader memory (in the crypto
-partition), ``crypto_keys.h`` defines a function ``tfm_plat_load_builtin_keys``
-which is responsible for loading all builtin keys into and driver that requires
-loading.
+  1. During the driver initialisation phase started by the PSA Crypto Core init,
+     the ``tfm_builtin_key_loader_init`` function is called to probe the
+     platform HAL and retrieve the builtin keys. Those keys are described by
+     the types defined in the HAL layer interface header
+     ``tfm_plat_crypto_keys.h``. In particular global tables describing key
+     properties and user policies must be implemented by each platform and are
+     retrieved by the two accessor functions
+     ``tfm_plat_builtin_key_get_desc_table_ptr()`` and
+     ``tfm_plat_builtin_key_get_policy_table_ptr()``. The keys are loaded from
+     the platform in secure RAM in the TF-M Crypto partition with associated
+     metadata. It's worth to note that the keys are loaded through callback
+     functions which the platform lists in the key descriptor table.
+  2. Once the TF-M Crypto service is initialised, at runtime it might receive a
+     request through an API call to use one of the builtin key IDs. Those IDs
+     are described in the ``tfm_builtin_key_ids.h`` header. Platforms can
+     override the default values providing their own header.
+  3. When the PSA Crypto core in the ``psa_get_and_lock_key_slot()`` function
+     checks that the ``key_id`` being requested is in the builtin range region,
+     it probes the platform through the function
+     ``mbedtls_psa_platform_get_builtin_key`` which must be implemented by the
+     TF-M Crypto service library abstraction layer. This function just checks
+     the key descriptor table from the HAL to understand if such ``key_id`` is
+     available in the platform and what are the corresponding ``slot`` and
+     ``lifetime`` values associated to it. The lifetime contains the location of
+     the key, which determines which driver is responsible for dealing with it
+     (for the ``tfm_builtin_key_loader`` driver, that is
+     ``TFM_BUILTIN_KEY_LOADER_KEY_LOCATION``), while the ``slot`` value maps the
+     key to the slot in the driver internal storage.
+  4. At this point, the PSA Crypto core knows that the key exists and is bound
+     to the location associated to the driver. By calling into the driver
+     wrappers layer is then able to retrieve the key attributes stored in the
+     platform for that key ID, and the size required to allocate in its
+     internal slot management system in order to load the key material in the
+     core. This is done by calling ``tfm_builtin_key_loader_get_builtin_key``
+     just with a valid key attributes pointer (and nothing else), to retrieve
+     the attributes. Once the attributes are available, the required size is
+     retrieved through the driver wrapper by calling
+     ``tfm_builtin_key_loader_get_key_buffer_size``.
+  5. At this stage, the slot management subsystem calls again into the driver
+     wrapper layer through ``tfm_builtin_key_loader_get_builtin_key`` with a
+     valid buffer to hold the key material returned by the
+     ``tfm_builtin_key_loader`` driver. When loading the key, the user
+     requiring that key_id is validated by the driver code against the policies
+     defined by the platform. If the policies match, the builtin key material
+     and metadata is loaded and is used like a transparent key available to the
+     PSA Crypto core slot management subsystem.
 
 *****************
 Technical details
@@ -98,22 +154,25 @@
 Builtin key IDs and overriding
 ------------------------------
 
-TF-M builtin key IDs are defined in ``interface/include/tfm_crypto_defs.h`` by
-the enum ``tfm_key_id_builtin_t``. They are allocated inside the range that PSA
-considers to be builtin keys. A platform can specify extra builtin key IDs by
-setting the ``PLATFORM_DEFAULT_CRYPTO_KEYS`` variable to ``OFF``, creating the
-header ``platform_builtin_key_ids.h``, and specifying new keys and IDs.
+TF-M builtin key IDs are defined in
+``interface/include/crypto_keys/tfm_builtin_key_ids.h`` through the enum
+``tfm_key_id_builtin_t``. They are allocated inside the range that PSA
+specifies for the builtin keys, i.e. between ``MBEDTLS_PSA_KEY_ID_BUILTIN_MIN``
+and ``MBEDLTS_PSA_KEY_ID_BUILTIN_MAX``. A platform can specify extra builtin key
+IDs by setting the ``PLATFORM_DEFAULT_CRYPTO_KEYS`` variable to ``OFF``,
+creating the header ``tfm_builtin_key_ids.h``, and specifying new keys and IDs.
 
 --------------------------
 Builtin key access control
 --------------------------
 
-Builtin keys by default can be used by any caller since the key handle is
-public information. TF-M must mediate access to the keys, which is done in the
-function ``tfm_plat_builtin_key_get_usage`` (part of ``crypto_keys.h``). This
-function maps the caller ID to a particular key usage, which allows granular key
-permissions. The function returns ``PSA_ERROR_NOT_PERMITTED`` if a caller does
-not have permission to use the key.
+A builtin key is accessible by all callers since the ``key_id`` associated to it
+is public information. Access to the keys must be mediated, which is done by
+matching the user requesting the ``key_id`` against the policies available for
+that user on that particular key in the policy table. If no policies are
+specified for a specific combination of user and ``key_id``, the usage flags in
+the key attributes will be all set to zeros, meaning the key will be unusable
+for any operation for that particular user.
 
 ------------------------------
 Multi-partition key derivation
@@ -122,16 +181,17 @@
 The HUK is used for key derivation by any secure partition or NS caller that
 requires keys that are bound to a particular context. For example, Protected
 Storage derives keys uniquely for each user of the service which are used to
-encrypt each user's files. In order to provide HUK derivation to every secure
+encrypt user files. In order to provide HUK derivation to every secure
 partition / NS caller, it must be ensured that no service that utilises HUK
-derivation can derive the same key as another service (simply by inputting the
-same KDF inputs).
+derivation can derive the same key as another service (simply by using the same
+inputs for the KDF APIs, i.e. accessing the same base key for derivation).
 
 This is accomplished by deriving a further "platform key" for each builtin key
-that can be used for key derivation. These platform keys are derived from the
-builtin key, using the partition ID as a KDF input, and can then be used for
-further derivation by the partition (or NS caller) with the further derived keys
-being unique for each partition even if the KDF inputs are the same.
+that has ``PSA_KEY_USAGE_DERIVE`` set in its attributes. These platform keys
+are derived from the builtin key, using the partition ID as a KDF input, and
+can then be used for safely for further derivations by the user, without risks
+to derive the same keys as other users. This is enforced directly by the
+``tfm_builtin_key_loader`` driver.
 
 .. Note::
     If the NS client ID feature is disabled, all NS callers share a partition ID
@@ -150,22 +210,35 @@
 for key derivation (``PSA_KEY_USAGE_DERIVE``), and not derive platform keys
 otherwise. For builtin keys that do not derive platform keys but are directly
 used, care must be taken with access control where multiple partitions have
-access.
+access to the same raw key material.
 
 ---------------------------------
 Mbed TLS transparent builtin keys
 ---------------------------------
 
 Mbed TLS does not natively support transparent builtin keys (transparent keys
-are keys where the key material is directly accessible to the PSA Crypto core),
+are keys where the key material is directly accessible by the PSA Crypto core),
 so some modifications had to be made. Opaque keyslots have the same basic
-structure as standard transparent keyslots, and can be passed to the functions
-usually reserved for transparent keys, though this behaviour is not defined and
-may not continue to work in future versions. Therefore, the only modification
-required currently is to force keys that have the location
-``TFM_BUILTIN_KEY_LOADER_KEY_LOCATION`` to be passed to the functions that only
-usually accept keys with the location ``PSA_KEY_LOCATION_LOCAL_STORAGE``.
+structure as standard transparent key slots, and can be passed to the functions
+usually reserved for transparent keys, though this is a private implementation
+detail of the mbed TLS library and is not specified in the driver interface.
+Therefore, the only modification required currently is to allow keys that have
+the location ``TFM_BUILTIN_KEY_LOADER_KEY_LOCATION`` to be passed to the
+functions that usually accept transparent keys only, i.e. with the location
+``PSA_KEY_LOCATION_LOCAL_STORAGE``. This is due to the fact that the standard
+assumption of the PSA Crypto core is that, if a driver that provides an
+additional location, will also provide dedicated cryptographic mechanisms to act
+on those keys, but this is not the case of the ``tfm_builtin_key_loader``, as it
+just provides a mechanism to load keys (which act as a transparent key with
+local storage, once loaded), but mbed TLS does not support such "transparent
+builtin key" concept.
+
+References
+----------
+
+.. [1] PSA cryptoprocessor driver interface: \ https://github.com/Mbed-TLS/mbedtls/blob/development/docs/proposed/psa-driver-interface.md
+.. [2] Definition of psa_key_location_t type in the PSA spec: \ https://armmbed.github.io/mbed-crypto/html/api/keys/lifetimes.html#c.psa_key_location_t
 
 --------------
 
-*Copyright (c) 2022, Arm Limited. All rights reserved.*
+*Copyright (c) 2022-2023, Arm Limited. All rights reserved.*
diff --git a/interface/include/crypto_keys/tfm_builtin_key_ids.h b/interface/include/crypto_keys/tfm_builtin_key_ids.h
index e0473a7..c254411 100644
--- a/interface/include/crypto_keys/tfm_builtin_key_ids.h
+++ b/interface/include/crypto_keys/tfm_builtin_key_ids.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
  *
@@ -9,21 +9,15 @@
 #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.
- *
+ * \note 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.
+ *       The range is specified by the limits defined through MBEDTLS_PSA_KEY_ID_BUILTIN_MIN
+ *       and MBEDTLS_PSA_KEY_ID_BUILTIN_MAX
  */
 enum tfm_builtin_key_id_t {
-    TFM_BUILTIN_KEY_ID_MIN = 0x7fff815Bu,
+    TFM_BUILTIN_KEY_ID_MIN = 0x7FFF815Bu,
     TFM_BUILTIN_KEY_ID_HUK,
     TFM_BUILTIN_KEY_ID_IAK,
 #ifdef TFM_PARTITION_DELEGATED_ATTESTATION
diff --git a/lib/ext/mbedcrypto/0001-BUILD-Update-IAR-support-in-CMakeLists.txt.patch b/lib/ext/mbedcrypto/0001-BUILD-Update-IAR-support-in-CMakeLists.txt.patch
index af1f725..8346bf5 100644
--- a/lib/ext/mbedcrypto/0001-BUILD-Update-IAR-support-in-CMakeLists.txt.patch
+++ b/lib/ext/mbedcrypto/0001-BUILD-Update-IAR-support-in-CMakeLists.txt.patch
@@ -1,4 +1,4 @@
-From 040624446b2abd6ef63d9004fdeefbb83c24d689 Mon Sep 17 00:00:00 2001
+From edd00ceca7f04206467d950bd20da8494e45ce8c Mon Sep 17 00:00:00 2001
 From: TTornblom <thomas.tornblom@iar.com>
 Date: Thu, 16 Apr 2020 13:53:38 +0200
 Subject: [PATCH 1/6] BUILD: Update IAR support in CMakeLists.txt
@@ -12,7 +12,7 @@
  1 file changed, 3 insertions(+), 1 deletion(-)
 
 diff --git a/CMakeLists.txt b/CMakeLists.txt
-index 337e4db9..ea090a69 100644
+index 337e4db95..ea090a692 100644
 --- a/CMakeLists.txt
 +++ b/CMakeLists.txt
 @@ -222,7 +222,9 @@ if(CMAKE_COMPILER_IS_CLANG)
diff --git a/lib/ext/mbedcrypto/0002-Enable-crypto-code-sharing-between-independent-binar.patch b/lib/ext/mbedcrypto/0002-Enable-crypto-code-sharing-between-independent-binar.patch
index 9d231d4..2fed3e0 100644
--- a/lib/ext/mbedcrypto/0002-Enable-crypto-code-sharing-between-independent-binar.patch
+++ b/lib/ext/mbedcrypto/0002-Enable-crypto-code-sharing-between-independent-binar.patch
@@ -1,4 +1,4 @@
-From a58629081d650ca9d7c6baab7598a52b45166a17 Mon Sep 17 00:00:00 2001
+From e4585e5285703a2769a810e90660434285e91925 Mon Sep 17 00:00:00 2001
 From: Tamas Ban <tamas.ban@arm.com>
 Date: Tue, 27 Oct 2020 08:55:37 +0000
 Subject: [PATCH 2/6] Enable crypto code sharing between independent binaries
@@ -13,7 +13,7 @@
 
 diff --git a/library/code_share.c b/library/code_share.c
 new file mode 100644
-index 00000000..2bf67fb4
+index 000000000..2bf67fb42
 --- /dev/null
 +++ b/library/code_share.c
 @@ -0,0 +1,3 @@
@@ -21,7 +21,7 @@
 + * extensive crypto code sharing was already applied on the mbedtls library.
 + */
 diff --git a/library/platform.c b/library/platform.c
-index 6151e6c4..074ecbb7 100644
+index 6151e6c49..074ecbb72 100644
 --- a/library/platform.c
 +++ b/library/platform.c
 @@ -53,8 +53,8 @@ static void platform_free_uninit( void *ptr )
@@ -36,7 +36,7 @@
  void * mbedtls_calloc( size_t nmemb, size_t size )
  {
 diff --git a/library/platform_util.c b/library/platform_util.c
-index 916a7f44..8936a9d7 100644
+index 916a7f444..8936a9d7d 100644
 --- a/library/platform_util.c
 +++ b/library/platform_util.c
 @@ -62,7 +62,7 @@
diff --git a/lib/ext/mbedcrypto/0003-Driver-wrapper-entry-points-for-CC3XX.patch b/lib/ext/mbedcrypto/0003-Driver-wrapper-entry-points-for-CC3XX.patch
index 95cc8e7..db124cb 100644
--- a/lib/ext/mbedcrypto/0003-Driver-wrapper-entry-points-for-CC3XX.patch
+++ b/lib/ext/mbedcrypto/0003-Driver-wrapper-entry-points-for-CC3XX.patch
@@ -1,4 +1,4 @@
-From 6c54899b07c7cc9615e98eba30781500f4b906c6 Mon Sep 17 00:00:00 2001
+From 89b0a57e9c29fde64572c62dbd4c3d75adde30a1 Mon Sep 17 00:00:00 2001
 From: Antonio de Angelis <Antonio.deAngelis@arm.com>
 Date: Fri, 15 Jul 2022 12:41:34 +0100
 Subject: [PATCH 3/6] Driver wrapper entry points for CC3XX
@@ -20,7 +20,7 @@
  3 files changed, 452 insertions(+), 52 deletions(-)
 
 diff --git a/include/psa/crypto_driver_contexts_composites.h b/include/psa/crypto_driver_contexts_composites.h
-index 3f1c8af4..2fdf9561 100644
+index 3f1c8af4b..2fdf9561f 100644
 --- a/include/psa/crypto_driver_contexts_composites.h
 +++ b/include/psa/crypto_driver_contexts_composites.h
 @@ -41,6 +41,9 @@
@@ -54,7 +54,7 @@
  
  #endif /* PSA_CRYPTO_DRIVER_CONTEXTS_COMPOSITES_H */
 diff --git a/include/psa/crypto_driver_contexts_primitives.h b/include/psa/crypto_driver_contexts_primitives.h
-index 2bb01ed4..2bc0bda7 100644
+index 2bb01ed43..2bc0bda70 100644
 --- a/include/psa/crypto_driver_contexts_primitives.h
 +++ b/include/psa/crypto_driver_contexts_primitives.h
 @@ -40,6 +40,9 @@
@@ -88,7 +88,7 @@
  
  #endif /* PSA_CRYPTO_DRIVER_CONTEXTS_PRIMITIVES_H */
 diff --git a/library/psa_crypto_driver_wrappers.c b/library/psa_crypto_driver_wrappers.c
-index 19819413..e822eef0 100644
+index 198194139..e822eef01 100644
 --- a/library/psa_crypto_driver_wrappers.c
 +++ b/library/psa_crypto_driver_wrappers.c
 @@ -48,6 +48,16 @@
diff --git a/lib/ext/mbedcrypto/0004-Add-TF-M-builtin-key-driver.patch b/lib/ext/mbedcrypto/0004-Add-TF-M-builtin-key-driver.patch
index 8aa37fb..bee322f 100644
--- a/lib/ext/mbedcrypto/0004-Add-TF-M-builtin-key-driver.patch
+++ b/lib/ext/mbedcrypto/0004-Add-TF-M-builtin-key-driver.patch
@@ -1,16 +1,17 @@
-From 2982172aba752f612b4e473626f787155087ddb5 Mon Sep 17 00:00:00 2001
+From 6f49a579de65abbae877ba067ee1a76671ae8e83 Mon Sep 17 00:00:00 2001
 From: Raef Coles <raef.coles@arm.com>
 Date: Tue, 19 Jul 2022 11:12:30 +0100
 Subject: [PATCH 4/6] Add TF-M builtin key driver
 
 Signed-off-by: Raef Coles <raef.coles@arm.com>
+Co-authored-by: Antonio de Angelis <antonio.deangelis@arm.com>
 ---
  library/psa_crypto.c                 | 11 +++-
- library/psa_crypto_driver_wrappers.c | 91 +++++++++++++++++++++++++++-
- 2 files changed, 100 insertions(+), 2 deletions(-)
+ library/psa_crypto_driver_wrappers.c | 97 +++++++++++++++++++++++++++-
+ 2 files changed, 106 insertions(+), 2 deletions(-)
 
 diff --git a/library/psa_crypto.c b/library/psa_crypto.c
-index cb5791fc..43c8b9e6 100644
+index cb5791fc7..43c8b9e6a 100644
 --- a/library/psa_crypto.c
 +++ b/library/psa_crypto.c
 @@ -82,6 +82,11 @@
@@ -39,7 +40,7 @@
          psa_unlock_key_slot( *p_slot );
          *p_slot = NULL;
 diff --git a/library/psa_crypto_driver_wrappers.c b/library/psa_crypto_driver_wrappers.c
-index e822eef0..08427604 100644
+index e822eef01..1fa26d562 100644
 --- a/library/psa_crypto_driver_wrappers.c
 +++ b/library/psa_crypto_driver_wrappers.c
 @@ -58,6 +58,18 @@
@@ -71,7 +72,20 @@
  
  /* END-driver id */
  
-@@ -155,6 +170,9 @@ psa_status_t psa_driver_wrapper_sign_message(
+@@ -92,6 +107,12 @@ psa_status_t psa_driver_wrapper_init( void )
+ {
+     psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+ 
++#if defined(PSA_CRYPTO_DRIVER_TFM_BUILTIN_KEY_LOADER)
++    status = tfm_builtin_key_loader_init();
++    if (status != PSA_SUCCESS)
++        return ( status );
++#endif /* PSA_CRYPTO_DRIVER_TFM_BUILTIN_KEY_LOADER */
++
+ #if defined(PSA_CRYPTO_DRIVER_CC3XX)
+     status = cc3xx_init();
+     if (status != PSA_SUCCESS)
+@@ -155,6 +176,9 @@ psa_status_t psa_driver_wrapper_sign_message(
      switch( location )
      {
          case PSA_KEY_LOCATION_LOCAL_STORAGE:
@@ -81,7 +95,7 @@
              /* Key is stored in the slot in export representation, so
               * cycle through all known transparent accelerators */
  #if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
-@@ -243,6 +261,9 @@ psa_status_t psa_driver_wrapper_verify_message(
+@@ -243,6 +267,9 @@ psa_status_t psa_driver_wrapper_verify_message(
      switch( location )
      {
          case PSA_KEY_LOCATION_LOCAL_STORAGE:
@@ -91,7 +105,7 @@
              /* Key is stored in the slot in export representation, so
               * cycle through all known transparent accelerators */
  #if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
-@@ -342,6 +363,9 @@ psa_status_t psa_driver_wrapper_sign_hash(
+@@ -342,6 +369,9 @@ psa_status_t psa_driver_wrapper_sign_hash(
      switch( location )
      {
          case PSA_KEY_LOCATION_LOCAL_STORAGE:
@@ -101,7 +115,7 @@
              /* Key is stored in the slot in export representation, so
               * cycle through all known transparent accelerators */
  #if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
-@@ -438,6 +462,9 @@ psa_status_t psa_driver_wrapper_verify_hash(
+@@ -438,6 +468,9 @@ psa_status_t psa_driver_wrapper_verify_hash(
      switch( location )
      {
          case PSA_KEY_LOCATION_LOCAL_STORAGE:
@@ -111,7 +125,7 @@
              /* Key is stored in the slot in export representation, so
               * cycle through all known transparent accelerators */
  #if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
-@@ -582,7 +609,11 @@ psa_status_t psa_driver_wrapper_get_key_buffer_size(
+@@ -582,7 +615,11 @@ psa_status_t psa_driver_wrapper_get_key_buffer_size(
              return( ( *key_buffer_size != 0 ) ?
                      PSA_SUCCESS : PSA_ERROR_NOT_SUPPORTED );
  #endif /* PSA_CRYPTO_DRIVER_TEST */
@@ -124,7 +138,7 @@
          default:
              (void)key_type;
              (void)key_bits;
-@@ -622,6 +653,9 @@ psa_status_t psa_driver_wrapper_generate_key(
+@@ -622,6 +659,9 @@ psa_status_t psa_driver_wrapper_generate_key(
      switch( location )
      {
          case PSA_KEY_LOCATION_LOCAL_STORAGE:
@@ -134,7 +148,7 @@
  #if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
              /* Transparent drivers are limited to generating asymmetric keys */
              if( PSA_KEY_TYPE_IS_ASYMMETRIC( attributes->core.type ) )
-@@ -714,6 +748,9 @@ psa_status_t psa_driver_wrapper_import_key(
+@@ -714,6 +754,9 @@ psa_status_t psa_driver_wrapper_import_key(
      switch( location )
      {
          case PSA_KEY_LOCATION_LOCAL_STORAGE:
@@ -144,7 +158,7 @@
              /* Key is stored in the slot in export representation, so
               * cycle through all known transparent accelerators */
  #if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
-@@ -800,6 +837,9 @@ psa_status_t psa_driver_wrapper_export_key(
+@@ -800,6 +843,9 @@ psa_status_t psa_driver_wrapper_export_key(
      switch( location )
      {
          case PSA_KEY_LOCATION_LOCAL_STORAGE:
@@ -154,7 +168,7 @@
              return( psa_export_key_internal( attributes,
                                               key_buffer,
                                               key_buffer_size,
-@@ -865,6 +905,9 @@ psa_status_t psa_driver_wrapper_export_public_key(
+@@ -865,6 +911,9 @@ psa_status_t psa_driver_wrapper_export_public_key(
      switch( location )
      {
          case PSA_KEY_LOCATION_LOCAL_STORAGE:
@@ -164,13 +178,13 @@
              /* Key is stored in the slot in export representation, so
               * cycle through all known transparent accelerators */
  #if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
-@@ -950,6 +993,13 @@ psa_status_t psa_driver_wrapper_get_builtin_key(
+@@ -950,6 +999,13 @@ psa_status_t psa_driver_wrapper_get_builtin_key(
  
  
  #endif /* PSA_CRYPTO_DRIVER_TEST */
 +#if defined(PSA_CRYPTO_DRIVER_TFM_BUILTIN_KEY_LOADER)
 +        case TFM_BUILTIN_KEY_LOADER_KEY_LOCATION:
-+            return( tfm_builtin_key_loader_get_key_buffer(
++            return( tfm_builtin_key_loader_get_builtin_key(
 +                        slot_number,
 +                        attributes,
 +                        key_buffer, key_buffer_size, key_buffer_length ) );
@@ -178,7 +192,7 @@
          default:
              (void) slot_number;
              (void) key_buffer;
-@@ -1035,6 +1085,9 @@ psa_status_t psa_driver_wrapper_cipher_encrypt(
+@@ -1035,6 +1091,9 @@ psa_status_t psa_driver_wrapper_cipher_encrypt(
      switch( location )
      {
          case PSA_KEY_LOCATION_LOCAL_STORAGE:
@@ -188,7 +202,7 @@
              /* Key is stored in the slot in export representation, so
               * cycle through all known transparent accelerators */
  #if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
-@@ -1139,6 +1192,9 @@ psa_status_t psa_driver_wrapper_cipher_decrypt(
+@@ -1139,6 +1198,9 @@ psa_status_t psa_driver_wrapper_cipher_decrypt(
      switch( location )
      {
          case PSA_KEY_LOCATION_LOCAL_STORAGE:
@@ -198,7 +212,7 @@
              /* Key is stored in the slot in export representation, so
               * cycle through all known transparent accelerators */
  #if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
-@@ -1228,6 +1284,9 @@ psa_status_t psa_driver_wrapper_cipher_encrypt_setup(
+@@ -1228,6 +1290,9 @@ psa_status_t psa_driver_wrapper_cipher_encrypt_setup(
      switch( location )
      {
          case PSA_KEY_LOCATION_LOCAL_STORAGE:
@@ -208,7 +222,7 @@
              /* Key is stored in the slot in export representation, so
               * cycle through all known transparent accelerators */
  #if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
-@@ -1311,6 +1370,9 @@ psa_status_t psa_driver_wrapper_cipher_decrypt_setup(
+@@ -1311,6 +1376,9 @@ psa_status_t psa_driver_wrapper_cipher_decrypt_setup(
      switch( location )
      {
          case PSA_KEY_LOCATION_LOCAL_STORAGE:
@@ -218,7 +232,7 @@
              /* Key is stored in the slot in export representation, so
               * cycle through all known transparent accelerators */
  #if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
-@@ -1794,6 +1856,9 @@ psa_status_t psa_driver_wrapper_aead_encrypt(
+@@ -1794,6 +1862,9 @@ psa_status_t psa_driver_wrapper_aead_encrypt(
      switch( location )
      {
          case PSA_KEY_LOCATION_LOCAL_STORAGE:
@@ -228,7 +242,7 @@
              /* Key is stored in the slot in export representation, so
               * cycle through all known transparent accelerators */
  
-@@ -1857,6 +1922,9 @@ psa_status_t psa_driver_wrapper_aead_decrypt(
+@@ -1857,6 +1928,9 @@ psa_status_t psa_driver_wrapper_aead_decrypt(
      switch( location )
      {
          case PSA_KEY_LOCATION_LOCAL_STORAGE:
@@ -238,7 +252,7 @@
              /* Key is stored in the slot in export representation, so
               * cycle through all known transparent accelerators */
  
-@@ -1917,6 +1985,9 @@ psa_status_t psa_driver_wrapper_aead_encrypt_setup(
+@@ -1917,6 +1991,9 @@ psa_status_t psa_driver_wrapper_aead_encrypt_setup(
      switch( location )
      {
          case PSA_KEY_LOCATION_LOCAL_STORAGE:
@@ -248,7 +262,7 @@
              /* Key is stored in the slot in export representation, so
               * cycle through all known transparent accelerators */
  
-@@ -1974,6 +2045,9 @@ psa_status_t psa_driver_wrapper_aead_decrypt_setup(
+@@ -1974,6 +2051,9 @@ psa_status_t psa_driver_wrapper_aead_decrypt_setup(
      switch( location )
      {
          case PSA_KEY_LOCATION_LOCAL_STORAGE:
@@ -258,7 +272,7 @@
              /* Key is stored in the slot in export representation, so
               * cycle through all known transparent accelerators */
  
-@@ -2371,6 +2445,9 @@ psa_status_t psa_driver_wrapper_mac_compute(
+@@ -2371,6 +2451,9 @@ psa_status_t psa_driver_wrapper_mac_compute(
      switch( location )
      {
          case PSA_KEY_LOCATION_LOCAL_STORAGE:
@@ -268,7 +282,7 @@
              /* Key is stored in the slot in export representation, so
               * cycle through all known transparent accelerators */
  #if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
-@@ -2441,6 +2518,9 @@ psa_status_t psa_driver_wrapper_mac_sign_setup(
+@@ -2441,6 +2524,9 @@ psa_status_t psa_driver_wrapper_mac_sign_setup(
      switch( location )
      {
          case PSA_KEY_LOCATION_LOCAL_STORAGE:
@@ -278,7 +292,7 @@
              /* Key is stored in the slot in export representation, so
               * cycle through all known transparent accelerators */
  #if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
-@@ -2522,6 +2602,9 @@ psa_status_t psa_driver_wrapper_mac_verify_setup(
+@@ -2522,6 +2608,9 @@ psa_status_t psa_driver_wrapper_mac_verify_setup(
      switch( location )
      {
          case PSA_KEY_LOCATION_LOCAL_STORAGE:
@@ -288,7 +302,7 @@
              /* Key is stored in the slot in export representation, so
               * cycle through all known transparent accelerators */
  #if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
-@@ -2822,6 +2905,9 @@ psa_status_t psa_driver_wrapper_asymmetric_decrypt(
+@@ -2822,6 +2911,9 @@ psa_status_t psa_driver_wrapper_asymmetric_decrypt(
      switch( location )
      {
          case PSA_KEY_LOCATION_LOCAL_STORAGE:
@@ -298,7 +312,7 @@
              /* Key is stored in the slot in export representation, so
               * cycle through all known transparent accelerators */
  #if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
-@@ -2900,6 +2986,9 @@ psa_status_t psa_driver_wrapper_key_agreement(
+@@ -2900,6 +2992,9 @@ psa_status_t psa_driver_wrapper_key_agreement(
      switch( location )
      {
          case PSA_KEY_LOCATION_LOCAL_STORAGE:
diff --git a/lib/ext/mbedcrypto/0005-CC3XX-Manually-enforce-no-software-builtin-fallback-.patch b/lib/ext/mbedcrypto/0005-CC3XX-Manually-enforce-no-software-builtin-fallback-.patch
index ae23516..335a31d 100644
--- a/lib/ext/mbedcrypto/0005-CC3XX-Manually-enforce-no-software-builtin-fallback-.patch
+++ b/lib/ext/mbedcrypto/0005-CC3XX-Manually-enforce-no-software-builtin-fallback-.patch
@@ -1,4 +1,4 @@
-From b1badf85b254f43dad2847d05dedc4d32b1fc57d Mon Sep 17 00:00:00 2001
+From 9e1fcca0e8ba9e58135c91af9c8c640b1d0b32c5 Mon Sep 17 00:00:00 2001
 From: Antonio de Angelis <Antonio.deAngelis@arm.com>
 Date: Tue, 2 Aug 2022 13:05:05 +0200
 Subject: [PATCH 5/6] CC3XX: Manually enforce no-software builtin fallback when
@@ -14,10 +14,10 @@
  1 file changed, 77 insertions(+), 69 deletions(-)
 
 diff --git a/library/psa_crypto_driver_wrappers.c b/library/psa_crypto_driver_wrappers.c
-index 08427604..7c6b618d 100644
+index 1fa26d562..80e20d440 100644
 --- a/library/psa_crypto_driver_wrappers.c
 +++ b/library/psa_crypto_driver_wrappers.c
-@@ -205,6 +205,7 @@ psa_status_t psa_driver_wrapper_sign_message(
+@@ -211,6 +211,7 @@ psa_status_t psa_driver_wrapper_sign_message(
              return( status );
  #endif /* PSA_CRYPTO_DRIVER_CC3XX */
  #endif /* PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT */
@@ -25,7 +25,7 @@
              /* Fell through, meaning no accelerator supports this operation */
              return( psa_sign_message_builtin( attributes,
                                        key_buffer,
-@@ -215,7 +216,7 @@ psa_status_t psa_driver_wrapper_sign_message(
+@@ -221,7 +222,7 @@ psa_status_t psa_driver_wrapper_sign_message(
                                        signature,
                                        signature_size,
                                        signature_length ) );
@@ -34,7 +34,7 @@
          /* Add cases for opaque driver here */
  #if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
  #if defined(PSA_CRYPTO_DRIVER_TEST)
-@@ -294,6 +295,7 @@ psa_status_t psa_driver_wrapper_verify_message(
+@@ -300,6 +301,7 @@ psa_status_t psa_driver_wrapper_verify_message(
              return( status );
  #endif /* PSA_CRYPTO_DRIVER_CC3XX */
  #endif /* PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT */
@@ -42,7 +42,7 @@
              /* Fell through, meaning no accelerator supports this operation */
              return( psa_verify_message_builtin( attributes,
                                          key_buffer,
-@@ -303,7 +305,7 @@ psa_status_t psa_driver_wrapper_verify_message(
+@@ -309,7 +311,7 @@ psa_status_t psa_driver_wrapper_verify_message(
                                          input_length,
                                          signature,
                                          signature_length ) );
@@ -51,7 +51,7 @@
          /* Add cases for opaque driver here */
  #if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
  #if defined(PSA_CRYPTO_DRIVER_TEST)
-@@ -396,6 +398,7 @@ psa_status_t psa_driver_wrapper_sign_hash(
+@@ -402,6 +404,7 @@ psa_status_t psa_driver_wrapper_sign_hash(
              return( status );
  #endif /* PSA_CRYPTO_DRIVER_CC3XX */
  #endif /* PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT */
@@ -59,7 +59,7 @@
              /* Fell through, meaning no accelerator supports this operation */
              return( psa_sign_hash_builtin( attributes,
                                             key_buffer,
-@@ -406,7 +409,7 @@ psa_status_t psa_driver_wrapper_sign_hash(
+@@ -412,7 +415,7 @@ psa_status_t psa_driver_wrapper_sign_hash(
                                             signature,
                                             signature_size,
                                             signature_length ) );
@@ -68,7 +68,7 @@
          /* Add cases for opaque driver here */
  #if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
  #if defined(PSA_CRYPTO_DRIVER_TEST)
-@@ -494,7 +497,7 @@ psa_status_t psa_driver_wrapper_verify_hash(
+@@ -500,7 +503,7 @@ psa_status_t psa_driver_wrapper_verify_hash(
                  return( status );
  #endif /* PSA_CRYPTO_DRIVER_TEST */
  #endif /* PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT */
@@ -77,7 +77,7 @@
              return( psa_verify_hash_builtin( attributes,
                                               key_buffer,
                                               key_buffer_size,
-@@ -503,7 +506,7 @@ psa_status_t psa_driver_wrapper_verify_hash(
+@@ -509,7 +512,7 @@ psa_status_t psa_driver_wrapper_verify_hash(
                                               hash_length,
                                               signature,
                                               signature_length ) );
@@ -86,7 +86,7 @@
          /* Add cases for opaque driver here */
  #if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
  #if defined(PSA_CRYPTO_DRIVER_TEST)
-@@ -937,6 +940,7 @@ psa_status_t psa_driver_wrapper_export_public_key(
+@@ -943,6 +946,7 @@ psa_status_t psa_driver_wrapper_export_public_key(
  #endif /* PSA_CRYPTO_DRIVER_CC3XX */
  
  #endif /* PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT */
@@ -94,7 +94,7 @@
              /* Fell through, meaning no accelerator supports this operation */
              return( psa_export_public_key_internal( attributes,
                                                      key_buffer,
-@@ -944,7 +948,7 @@ psa_status_t psa_driver_wrapper_export_public_key(
+@@ -950,7 +954,7 @@ psa_status_t psa_driver_wrapper_export_public_key(
                                                      data,
                                                      data_size,
                                                      data_length ) );
@@ -103,7 +103,7 @@
          /* Add cases for opaque driver here */
  #if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
  
-@@ -1123,7 +1127,7 @@ psa_status_t psa_driver_wrapper_cipher_encrypt(
+@@ -1129,7 +1133,7 @@ psa_status_t psa_driver_wrapper_cipher_encrypt(
  #endif /* PSA_CRYPTO_DRIVER_CC3XX */
  #endif /* PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT */
  
@@ -112,7 +112,7 @@
              return( mbedtls_psa_cipher_encrypt( attributes,
                                                  key_buffer,
                                                  key_buffer_size,
-@@ -1137,7 +1141,7 @@ psa_status_t psa_driver_wrapper_cipher_encrypt(
+@@ -1143,7 +1147,7 @@ psa_status_t psa_driver_wrapper_cipher_encrypt(
                                                  output_length ) );
  #else
              return( PSA_ERROR_NOT_SUPPORTED );
@@ -121,7 +121,7 @@
  
          /* Add cases for opaque driver here */
  #if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
-@@ -1226,7 +1230,7 @@ psa_status_t psa_driver_wrapper_cipher_decrypt(
+@@ -1232,7 +1236,7 @@ psa_status_t psa_driver_wrapper_cipher_decrypt(
  #endif /* PSA_CRYPTO_DRIVER_CC3XX */
  #endif /* PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT */
  
@@ -130,7 +130,7 @@
              return( mbedtls_psa_cipher_decrypt( attributes,
                                                  key_buffer,
                                                  key_buffer_size,
-@@ -1238,7 +1242,7 @@ psa_status_t psa_driver_wrapper_cipher_decrypt(
+@@ -1244,7 +1248,7 @@ psa_status_t psa_driver_wrapper_cipher_decrypt(
                                                  output_length ) );
  #else
              return( PSA_ERROR_NOT_SUPPORTED );
@@ -139,7 +139,7 @@
  
          /* Add cases for opaque driver here */
  #if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
-@@ -1315,7 +1319,7 @@ psa_status_t psa_driver_wrapper_cipher_encrypt_setup(
+@@ -1321,7 +1325,7 @@ psa_status_t psa_driver_wrapper_cipher_encrypt_setup(
              return( status );
  #endif /* PSA_CRYPTO_DRIVER_CC3XX */
  #endif /* PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT */
@@ -148,7 +148,7 @@
              /* Fell through, meaning no accelerator supports this operation */
              status = mbedtls_psa_cipher_encrypt_setup( &operation->ctx.mbedtls_ctx,
                                                         attributes,
-@@ -1327,7 +1331,7 @@ psa_status_t psa_driver_wrapper_cipher_encrypt_setup(
+@@ -1333,7 +1337,7 @@ psa_status_t psa_driver_wrapper_cipher_encrypt_setup(
  
              if( status != PSA_ERROR_NOT_SUPPORTED )
                  return( status );
@@ -157,7 +157,7 @@
              return( PSA_ERROR_NOT_SUPPORTED );
  
          /* Add cases for opaque driver here */
-@@ -1401,7 +1405,7 @@ psa_status_t psa_driver_wrapper_cipher_decrypt_setup(
+@@ -1407,7 +1411,7 @@ psa_status_t psa_driver_wrapper_cipher_decrypt_setup(
              return( status );
  #endif /* PSA_CRYPTO_DRIVER_CC3XX */
  #endif /* PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT */
@@ -166,7 +166,7 @@
              /* Fell through, meaning no accelerator supports this operation */
              status = mbedtls_psa_cipher_decrypt_setup( &operation->ctx.mbedtls_ctx,
                                                         attributes,
-@@ -1412,7 +1416,7 @@ psa_status_t psa_driver_wrapper_cipher_decrypt_setup(
+@@ -1418,7 +1422,7 @@ psa_status_t psa_driver_wrapper_cipher_decrypt_setup(
                  operation->id = PSA_CRYPTO_MBED_TLS_DRIVER_ID;
  
              return( status );
@@ -175,7 +175,7 @@
              return( PSA_ERROR_NOT_SUPPORTED );
  
          /* Add cases for opaque driver here */
-@@ -1449,12 +1453,12 @@ psa_status_t psa_driver_wrapper_cipher_set_iv(
+@@ -1455,12 +1459,12 @@ psa_status_t psa_driver_wrapper_cipher_set_iv(
  {
      switch( operation->id )
      {
@@ -190,7 +190,7 @@
  
  #if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
  #if defined(PSA_CRYPTO_DRIVER_TEST)
-@@ -1493,7 +1497,7 @@ psa_status_t psa_driver_wrapper_cipher_update(
+@@ -1499,7 +1503,7 @@ psa_status_t psa_driver_wrapper_cipher_update(
  {
      switch( operation->id )
      {
@@ -199,7 +199,7 @@
          case PSA_CRYPTO_MBED_TLS_DRIVER_ID:
              return( mbedtls_psa_cipher_update( &operation->ctx.mbedtls_ctx,
                                                 input,
-@@ -1501,7 +1505,7 @@ psa_status_t psa_driver_wrapper_cipher_update(
+@@ -1507,7 +1511,7 @@ psa_status_t psa_driver_wrapper_cipher_update(
                                                 output,
                                                 output_size,
                                                 output_length ) );
@@ -208,7 +208,7 @@
  
  #if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
  #if defined(PSA_CRYPTO_DRIVER_TEST)
-@@ -1544,13 +1548,13 @@ psa_status_t psa_driver_wrapper_cipher_finish(
+@@ -1550,13 +1554,13 @@ psa_status_t psa_driver_wrapper_cipher_finish(
  {
      switch( operation->id )
      {
@@ -224,7 +224,7 @@
  
  #if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
  #if defined(PSA_CRYPTO_DRIVER_TEST)
-@@ -1587,10 +1591,10 @@ psa_status_t psa_driver_wrapper_cipher_abort(
+@@ -1593,10 +1597,10 @@ psa_status_t psa_driver_wrapper_cipher_abort(
  
      switch( operation->id )
      {
@@ -237,7 +237,7 @@
  
  #if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
  #if defined(PSA_CRYPTO_DRIVER_TEST)
-@@ -1654,7 +1658,7 @@ psa_status_t psa_driver_wrapper_hash_compute(
+@@ -1660,7 +1664,7 @@ psa_status_t psa_driver_wrapper_hash_compute(
  #endif /* defined(PSA_CRYPTO_DRIVER_CC3XX) */
  #endif /* PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT */
      /* If software fallback is compiled in, try fallback */
@@ -246,7 +246,7 @@
      status = mbedtls_psa_hash_compute( alg, input, input_length,
                                         hash, hash_size, hash_length );
      if( status != PSA_ERROR_NOT_SUPPORTED )
-@@ -1695,7 +1699,7 @@ psa_status_t psa_driver_wrapper_hash_setup(
+@@ -1701,7 +1705,7 @@ psa_status_t psa_driver_wrapper_hash_setup(
  #endif /* PSA_CRYPTO_DRIVER_CC3XX */
  #endif /* PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT */
  
@@ -255,7 +255,7 @@
      /* If software fallback is compiled in, try fallback */
      status = mbedtls_psa_hash_setup( &operation->ctx.mbedtls_ctx, alg );
      if( status == PSA_SUCCESS )
-@@ -1703,7 +1707,7 @@ psa_status_t psa_driver_wrapper_hash_setup(
+@@ -1709,7 +1713,7 @@ psa_status_t psa_driver_wrapper_hash_setup(
  
      if( status != PSA_ERROR_NOT_SUPPORTED )
          return( status );
@@ -264,7 +264,7 @@
      /* Nothing left to try if we fall through here */
      (void) status;
      (void) operation;
-@@ -1734,12 +1738,12 @@ psa_status_t psa_driver_wrapper_hash_clone(
+@@ -1740,12 +1744,12 @@ psa_status_t psa_driver_wrapper_hash_clone(
  
  #endif /* PSA_CRYPTO_DRIVER_CC3XX */
  #endif /* PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT */
@@ -279,7 +279,7 @@
          default:
              (void) target_operation;
              return( PSA_ERROR_BAD_STATE );
-@@ -1767,11 +1771,11 @@ psa_status_t psa_driver_wrapper_hash_update(
+@@ -1773,11 +1777,11 @@ psa_status_t psa_driver_wrapper_hash_update(
                          input, input_length ) );
  #endif /* PSA_CRYPTO_DRIVER_CC3XX */
  #endif /* PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT */
@@ -293,7 +293,7 @@
          default:
              (void) input;
              (void) input_length;
-@@ -1801,11 +1805,11 @@ psa_status_t psa_driver_wrapper_hash_finish(
+@@ -1807,11 +1811,11 @@ psa_status_t psa_driver_wrapper_hash_finish(
                          hash, hash_size, hash_length ) );
  #endif /* PSA_CRYPTO_DRIVER_CC3XX */
  #endif /* PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT */
@@ -307,7 +307,7 @@
          default:
              (void) hash;
              (void) hash_size;
-@@ -1831,10 +1835,10 @@ psa_status_t psa_driver_wrapper_hash_abort(
+@@ -1837,10 +1841,10 @@ psa_status_t psa_driver_wrapper_hash_abort(
                          &operation->ctx.cc3xx_driver_ctx ) );
  #endif /* PSA_CRYPTO_DRIVER_CC3XX */
  #endif /* PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT */
@@ -320,7 +320,7 @@
          default:
              return( PSA_ERROR_BAD_STATE );
      }
-@@ -1887,7 +1891,7 @@ psa_status_t psa_driver_wrapper_aead_encrypt(
+@@ -1893,7 +1897,7 @@ psa_status_t psa_driver_wrapper_aead_encrypt(
              return( status );
  #endif /* PSA_CRYPTO_DRIVER_CC3XX */
  #endif /* PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT */
@@ -329,7 +329,7 @@
              /* Fell through, meaning no accelerator supports this operation */
              return( mbedtls_psa_aead_encrypt(
                          attributes, key_buffer, key_buffer_size,
-@@ -1896,7 +1900,7 @@ psa_status_t psa_driver_wrapper_aead_encrypt(
+@@ -1902,7 +1906,7 @@ psa_status_t psa_driver_wrapper_aead_encrypt(
                          additional_data, additional_data_length,
                          plaintext, plaintext_length,
                          ciphertext, ciphertext_size, ciphertext_length ) );
@@ -338,7 +338,7 @@
          /* Add cases for opaque driver here */
  
          default:
-@@ -1953,7 +1957,7 @@ psa_status_t psa_driver_wrapper_aead_decrypt(
+@@ -1959,7 +1963,7 @@ psa_status_t psa_driver_wrapper_aead_decrypt(
              return( status );
  #endif /* PSA_CRYPTO_DRIVER_CC3XX */
  #endif /* PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT */
@@ -347,7 +347,7 @@
              /* Fell through, meaning no accelerator supports this operation */
              return( mbedtls_psa_aead_decrypt(
                          attributes, key_buffer, key_buffer_size,
-@@ -1962,7 +1966,7 @@ psa_status_t psa_driver_wrapper_aead_decrypt(
+@@ -1968,7 +1972,7 @@ psa_status_t psa_driver_wrapper_aead_decrypt(
                          additional_data, additional_data_length,
                          ciphertext, ciphertext_length,
                          plaintext, plaintext_size, plaintext_length ) );
@@ -356,7 +356,7 @@
          /* Add cases for opaque driver here */
  
          default:
-@@ -2013,7 +2017,7 @@ psa_status_t psa_driver_wrapper_aead_encrypt_setup(
+@@ -2019,7 +2023,7 @@ psa_status_t psa_driver_wrapper_aead_encrypt_setup(
              return( status );
  #endif /* PSA_CRYPTO_DRIVER_CC3XX */
  #endif /* PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT */
@@ -365,7 +365,7 @@
              /* Fell through, meaning no accelerator supports this operation */
              operation->id = PSA_CRYPTO_MBED_TLS_DRIVER_ID;
              status = mbedtls_psa_aead_encrypt_setup(
-@@ -2022,7 +2026,7 @@ psa_status_t psa_driver_wrapper_aead_encrypt_setup(
+@@ -2028,7 +2032,7 @@ psa_status_t psa_driver_wrapper_aead_encrypt_setup(
                          alg );
  
              return( status );
@@ -374,7 +374,7 @@
          /* Add cases for opaque driver here */
  
          default:
-@@ -2075,7 +2079,7 @@ psa_status_t psa_driver_wrapper_aead_decrypt_setup(
+@@ -2081,7 +2085,7 @@ psa_status_t psa_driver_wrapper_aead_decrypt_setup(
              return( status );
  #endif /* PSA_CRYPTO_DRIVER_CC3XX */
  #endif /* PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT */
@@ -383,7 +383,7 @@
              /* Fell through, meaning no accelerator supports this operation */
              operation->id = PSA_CRYPTO_MBED_TLS_DRIVER_ID;
              status = mbedtls_psa_aead_decrypt_setup(
-@@ -2085,7 +2089,7 @@ psa_status_t psa_driver_wrapper_aead_decrypt_setup(
+@@ -2091,7 +2095,7 @@ psa_status_t psa_driver_wrapper_aead_decrypt_setup(
                          alg );
  
              return( status );
@@ -392,7 +392,7 @@
          /* Add cases for opaque driver here */
  
          default:
-@@ -2102,13 +2106,13 @@ psa_status_t psa_driver_wrapper_aead_set_nonce(
+@@ -2108,13 +2112,13 @@ psa_status_t psa_driver_wrapper_aead_set_nonce(
  {
      switch( operation->id )
      {
@@ -408,7 +408,7 @@
  
  #if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
  #if defined(PSA_CRYPTO_DRIVER_TEST)
-@@ -2143,13 +2147,13 @@ psa_status_t psa_driver_wrapper_aead_set_lengths(
+@@ -2149,13 +2153,13 @@ psa_status_t psa_driver_wrapper_aead_set_lengths(
  {
      switch( operation->id )
      {
@@ -424,7 +424,7 @@
  
  #if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
  #if defined(PSA_CRYPTO_DRIVER_TEST)
-@@ -2184,13 +2188,13 @@ psa_status_t psa_driver_wrapper_aead_update_ad(
+@@ -2190,13 +2194,13 @@ psa_status_t psa_driver_wrapper_aead_update_ad(
  {
      switch( operation->id )
      {
@@ -440,7 +440,7 @@
  
  #if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
  #if defined(PSA_CRYPTO_DRIVER_TEST)
-@@ -2228,14 +2232,14 @@ psa_status_t psa_driver_wrapper_aead_update(
+@@ -2234,14 +2238,14 @@ psa_status_t psa_driver_wrapper_aead_update(
  {
      switch( operation->id )
      {
@@ -457,7 +457,7 @@
  
  #if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
  #if defined(PSA_CRYPTO_DRIVER_TEST)
-@@ -2279,7 +2283,7 @@ psa_status_t psa_driver_wrapper_aead_finish(
+@@ -2285,7 +2289,7 @@ psa_status_t psa_driver_wrapper_aead_finish(
  {
      switch( operation->id )
      {
@@ -466,7 +466,7 @@
          case PSA_CRYPTO_MBED_TLS_DRIVER_ID:
              return( mbedtls_psa_aead_finish( &operation->ctx.mbedtls_ctx,
                                               ciphertext,
-@@ -2287,7 +2291,7 @@ psa_status_t psa_driver_wrapper_aead_finish(
+@@ -2293,7 +2297,7 @@ psa_status_t psa_driver_wrapper_aead_finish(
                                               ciphertext_length, tag,
                                               tag_size, tag_length ) );
  
@@ -475,7 +475,7 @@
  
  #if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
  #if defined(PSA_CRYPTO_DRIVER_TEST)
-@@ -2331,7 +2335,7 @@ psa_status_t psa_driver_wrapper_aead_verify(
+@@ -2337,7 +2341,7 @@ psa_status_t psa_driver_wrapper_aead_verify(
  {
      switch( operation->id )
      {
@@ -484,7 +484,7 @@
          case PSA_CRYPTO_MBED_TLS_DRIVER_ID:
              {
                  psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
-@@ -2359,7 +2363,7 @@ psa_status_t psa_driver_wrapper_aead_verify(
+@@ -2365,7 +2369,7 @@ psa_status_t psa_driver_wrapper_aead_verify(
                  return( status );
              }
  
@@ -493,7 +493,7 @@
  
  #if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
  #if defined(PSA_CRYPTO_DRIVER_TEST)
-@@ -2397,11 +2401,11 @@ psa_status_t psa_driver_wrapper_aead_abort(
+@@ -2403,11 +2407,11 @@ psa_status_t psa_driver_wrapper_aead_abort(
  {
      switch( operation->id )
      {
@@ -507,7 +507,7 @@
  
  #if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
  #if defined(PSA_CRYPTO_DRIVER_TEST)
-@@ -2467,7 +2471,7 @@ psa_status_t psa_driver_wrapper_mac_compute(
+@@ -2473,7 +2477,7 @@ psa_status_t psa_driver_wrapper_mac_compute(
              return( status );
  #endif /* PSA_CRYPTO_DRIVER_CC3XX */
  #endif /* PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT */
@@ -516,7 +516,7 @@
              /* Fell through, meaning no accelerator supports this operation */
              status = mbedtls_psa_mac_compute(
                  attributes, key_buffer, key_buffer_size, alg,
-@@ -2475,7 +2479,7 @@ psa_status_t psa_driver_wrapper_mac_compute(
+@@ -2481,7 +2485,7 @@ psa_status_t psa_driver_wrapper_mac_compute(
                  mac, mac_size, mac_length );
              if( status != PSA_ERROR_NOT_SUPPORTED )
                  return( status );
@@ -525,7 +525,7 @@
              return( PSA_ERROR_NOT_SUPPORTED );
  
          /* Add cases for opaque driver here */
-@@ -2547,7 +2551,7 @@ psa_status_t psa_driver_wrapper_mac_sign_setup(
+@@ -2553,7 +2557,7 @@ psa_status_t psa_driver_wrapper_mac_sign_setup(
              return status;
  #endif /* defined(PSA_CRYPTO_DRIVER_CC3XX) */
  #endif /* PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT */
@@ -534,7 +534,7 @@
              /* Fell through, meaning no accelerator supports this operation */
              status = mbedtls_psa_mac_sign_setup( &operation->ctx.mbedtls_ctx,
                                                   attributes,
-@@ -2558,7 +2562,7 @@ psa_status_t psa_driver_wrapper_mac_sign_setup(
+@@ -2564,7 +2568,7 @@ psa_status_t psa_driver_wrapper_mac_sign_setup(
  
              if( status != PSA_ERROR_NOT_SUPPORTED )
                  return( status );
@@ -543,7 +543,7 @@
              return( PSA_ERROR_NOT_SUPPORTED );
  
          /* Add cases for opaque driver here */
-@@ -2631,7 +2635,7 @@ psa_status_t psa_driver_wrapper_mac_verify_setup(
+@@ -2637,7 +2641,7 @@ psa_status_t psa_driver_wrapper_mac_verify_setup(
              return status;
  #endif /* defined(PSA_CRYPTO_DRIVER_CC3XX) */
  #endif /* PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT */
@@ -552,7 +552,7 @@
              /* Fell through, meaning no accelerator supports this operation */
              status = mbedtls_psa_mac_verify_setup( &operation->ctx.mbedtls_ctx,
                                                     attributes,
-@@ -2642,7 +2646,7 @@ psa_status_t psa_driver_wrapper_mac_verify_setup(
+@@ -2648,7 +2652,7 @@ psa_status_t psa_driver_wrapper_mac_verify_setup(
  
              if( status != PSA_ERROR_NOT_SUPPORTED )
                  return( status );
@@ -561,7 +561,7 @@
              return( PSA_ERROR_NOT_SUPPORTED );
  
          /* Add cases for opaque driver here */
-@@ -2679,11 +2683,11 @@ psa_status_t psa_driver_wrapper_mac_update(
+@@ -2685,11 +2689,11 @@ psa_status_t psa_driver_wrapper_mac_update(
  {
      switch( operation->id )
      {
@@ -575,7 +575,7 @@
  
  #if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
  #if defined(PSA_CRYPTO_DRIVER_TEST)
-@@ -2717,11 +2721,11 @@ psa_status_t psa_driver_wrapper_mac_sign_finish(
+@@ -2723,11 +2727,11 @@ psa_status_t psa_driver_wrapper_mac_sign_finish(
  {
      switch( operation->id )
      {
@@ -589,7 +589,7 @@
  
  #if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
  #if defined(PSA_CRYPTO_DRIVER_TEST)
-@@ -2756,11 +2760,11 @@ psa_status_t psa_driver_wrapper_mac_verify_finish(
+@@ -2762,11 +2766,11 @@ psa_status_t psa_driver_wrapper_mac_verify_finish(
  {
      switch( operation->id )
      {
@@ -603,7 +603,7 @@
  
  #if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
  #if defined(PSA_CRYPTO_DRIVER_TEST)
-@@ -2793,10 +2797,10 @@ psa_status_t psa_driver_wrapper_mac_abort(
+@@ -2799,10 +2803,10 @@ psa_status_t psa_driver_wrapper_mac_abort(
  {
      switch( operation->id )
      {
@@ -616,7 +616,7 @@
  
  #if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
  #if defined(PSA_CRYPTO_DRIVER_TEST)
-@@ -2860,10 +2864,12 @@ psa_status_t psa_driver_wrapper_asymmetric_encrypt(
+@@ -2866,10 +2870,12 @@ psa_status_t psa_driver_wrapper_asymmetric_encrypt(
                  return( status );
  #endif /* PSA_CRYPTO_DRIVER_TEST */
  #endif /* PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT */
@@ -629,7 +629,7 @@
          /* Add cases for opaque driver here */
  #if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
  #if defined(PSA_CRYPTO_DRIVER_TEST)
-@@ -2935,10 +2941,12 @@ psa_status_t psa_driver_wrapper_asymmetric_decrypt(
+@@ -2941,10 +2947,12 @@ psa_status_t psa_driver_wrapper_asymmetric_decrypt(
                  return( status );
  #endif /* PSA_CRYPTO_DRIVER_TEST */
  #endif /* PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT */
diff --git a/lib/ext/mbedcrypto/0006-Initialise-driver-wrappers-as-first-step-in-psa_cryp.patch b/lib/ext/mbedcrypto/0006-Initialise-driver-wrappers-as-first-step-in-psa_cryp.patch
index 2fb0973..ea0fd11 100644
--- a/lib/ext/mbedcrypto/0006-Initialise-driver-wrappers-as-first-step-in-psa_cryp.patch
+++ b/lib/ext/mbedcrypto/0006-Initialise-driver-wrappers-as-first-step-in-psa_cryp.patch
@@ -1,4 +1,4 @@
-From 270713ac18597abf1bd4b8eb76cd84b6cdfa2b95 Mon Sep 17 00:00:00 2001
+From b7d0156cccb681b448f015a35581e4a4713e1ca0 Mon Sep 17 00:00:00 2001
 From: Antonio de Angelis <Antonio.deAngelis@arm.com>
 Date: Tue, 23 Aug 2022 13:06:07 +0100
 Subject: [PATCH 6/6] Initialise driver wrappers as first step in
@@ -14,7 +14,7 @@
  1 file changed, 5 insertions(+), 5 deletions(-)
 
 diff --git a/library/psa_crypto.c b/library/psa_crypto.c
-index 43c8b9e6..a700a3ce 100644
+index 43c8b9e6a..a700a3ce0 100644
 --- a/library/psa_crypto.c
 +++ b/library/psa_crypto.c
 @@ -6327,6 +6327,11 @@ psa_status_t psa_crypto_init( void )
diff --git a/platform/ext/common/template/crypto_keys.c b/platform/ext/common/template/crypto_keys.c
index 255362b..43490b5 100644
--- a/platform/ext/common/template/crypto_keys.c
+++ b/platform/ext/common/template/crypto_keys.c
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2017-2022 Arm Limited. All rights reserved.
+ * Copyright (c) 2017-2023 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,96 +14,16 @@
  * limitations under the License.
  */
 
-#include "tfm_plat_crypto_keys.h"
-
-#include "tfm_builtin_key_ids.h"
-
-#include <stddef.h>
 #include <string.h>
-
-#include "region_defs.h"
-#include "cmsis_compiler.h"
+#include "tfm_plat_crypto_keys.h"
+#include "tfm_builtin_key_ids.h"
 #include "tfm_plat_otp.h"
 #include "psa_manifest/pid.h"
 #include "tfm_builtin_key_loader.h"
 
+#define NUMBER_OF_ELEMENTS_OF(x) sizeof(x)/sizeof(*x)
 #define TFM_NS_PARTITION_ID -1
 
-#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,
-                                                   int32_t owner,
-                                                   psa_key_usage_t *usage)
-{
-    *usage = 0;
-
-    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(owner) {
-#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;
-#if defined(TEST_S_ATTESTATION) || defined(TEST_NS_ATTESTATION)
-        /* So that the tests can validate created tokens */
-#ifdef TEST_S_ATTESTATION
-        case TFM_SP_SECURE_TEST_PARTITION:
-#endif /* TEST_S_ATTESTATION */
-#ifdef TEST_NS_ATTESTATION
-        case TFM_NS_PARTITION_ID:
-#endif /* TEST_NS_ATTESTATION */
-            *usage = PSA_KEY_USAGE_VERIFY_HASH;
-            break;
-#endif /* TEST_S_ATTESTATION || TEST_NS_ATTESTATION */
-#endif /* TFM_PARTITION_INITIAL_ATTESTATION */
-        default:
-            return TFM_PLAT_ERR_NOT_PERMITTED;
-        }
-        break;
-    default:
-        return TFM_PLAT_ERR_UNSUPPORTED;
-    }
-
-    return TFM_PLAT_ERR_SUCCESS;
-}
-
-enum tfm_plat_err_t tfm_plat_builtin_key_get_lifetime_and_slot(
-    struct tfm_crypto_key_id_s key_id,
-    psa_key_lifetime_t *lifetime,
-    psa_drv_slot_number_t *slot_number)
-{
-    switch (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,
@@ -175,51 +95,67 @@
 }
 #endif /* TFM_PARTITION_INITIAL_ATTESTATION */
 
-enum tfm_plat_err_t tfm_plat_load_builtin_keys(void)
-{
-    psa_status_t err;
-    mbedtls_svc_key_id_t key_id = MBEDTLS_SVC_KEY_ID_INIT;
-    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;
-
-    /* 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;
-    }
-
 #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;
-    }
-    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;
-    }
+/**
+ * @brief Table describing per-user key policy for the IAK
+ *
+ */
+static const tfm_plat_builtin_key_per_user_policy_t g_iak_per_user_policy[] = {
+    {.user = TFM_SP_INITIAL_ATTESTATION,
+#ifdef SYMMETRIC_INITIAL_ATTESTATION
+        .usage = PSA_KEY_USAGE_SIGN_HASH | PSA_KEY_USAGE_EXPORT,
+#else
+        .usage = PSA_KEY_USAGE_SIGN_HASH,
+#endif /* SYMMETRIC_INITIAL_ATTESTATION */
+    },
+#ifdef TEST_S_ATTESTATION
+    {.user = TFM_SP_SECURE_TEST_PARTITION, .usage = PSA_KEY_USAGE_VERIFY_HASH},
+#endif /* TEST_S_ATTESTATION */
+#ifdef TEST_NS_ATTESTATION
+    {.user = TFM_NS_PARTITION_ID, .usage = PSA_KEY_USAGE_VERIFY_HASH},
+#endif /* TEST_NS_ATTESTATION */
+};
 #endif /* TFM_PARTITION_INITIAL_ATTESTATION */
 
-    return TFM_PLAT_ERR_SUCCESS;
+/**
+ * @brief Table describing per-key user policies
+ *
+ */
+static const tfm_plat_builtin_key_policy_t g_builtin_keys_policy[] = {
+    {.key_id = TFM_BUILTIN_KEY_ID_HUK, .per_user_policy = 0, .usage = PSA_KEY_USAGE_DERIVE},
+#ifdef TFM_PARTITION_INITIAL_ATTESTATION
+    {.key_id = TFM_BUILTIN_KEY_ID_IAK,
+     .per_user_policy = NUMBER_OF_ELEMENTS_OF(g_iak_per_user_policy),
+     .policy_ptr = g_iak_per_user_policy},
+#endif /* TFM_PARTITION_INITIAL_ATTESTATION */
+};
+
+/**
+ * @brief Table describing the builtin-in keys (plaform keys) available in the platform. Note
+ *        that to bind the keys to the tfm_builtin_key_loader driver, the lifetime must be
+ *        explicitly set to the one associated to the driver, i.e. TFM_BUILTIN_KEY_LOADER_LIFETIME
+ */
+static const tfm_plat_builtin_key_descriptor_t g_builtin_keys_desc[] = {
+    {.key_id = TFM_BUILTIN_KEY_ID_HUK,
+     .slot_number = TFM_BUILTIN_KEY_SLOT_HUK,
+     .lifetime = TFM_BUILTIN_KEY_LOADER_LIFETIME,
+     .loader_key_func = tfm_plat_get_huk},
+#ifdef TFM_PARTITION_INITIAL_ATTESTATION
+    {.key_id = TFM_BUILTIN_KEY_ID_IAK,
+     .slot_number = TFM_BUILTIN_KEY_SLOT_IAK,
+     .lifetime = TFM_BUILTIN_KEY_LOADER_LIFETIME,
+     .loader_key_func = tfm_plat_get_iak},
+#endif /* TFM_PARTITION_INITIAL_ATTESTATION */
+};
+
+size_t tfm_plat_builtin_key_get_policy_table_ptr(const tfm_plat_builtin_key_policy_t *desc_ptr[])
+{
+    *desc_ptr = &g_builtin_keys_policy[0];
+    return NUMBER_OF_ELEMENTS_OF(g_builtin_keys_policy);
+}
+
+size_t tfm_plat_builtin_key_get_desc_table_ptr(const tfm_plat_builtin_key_descriptor_t *desc_ptr[])
+{
+    *desc_ptr = &g_builtin_keys_desc[0];
+    return NUMBER_OF_ELEMENTS_OF(g_builtin_keys_desc);
 }
diff --git a/platform/ext/target/arm/rss/common/CMakeLists.txt b/platform/ext/target/arm/rss/common/CMakeLists.txt
index 429dd84..b901894 100644
--- a/platform/ext/target/arm/rss/common/CMakeLists.txt
+++ b/platform/ext/target/arm/rss/common/CMakeLists.txt
@@ -92,7 +92,7 @@
 )
 
 
-install(FILES       platform_builtin_key_ids.h
+install(FILES       tfm_builtin_key_ids.h
         DESTINATION ${TFM_INSTALL_PATH}/interface/include)
 
 #========================= Platform Secure ====================================#
diff --git a/platform/ext/target/arm/rss/common/crypto_keys.c b/platform/ext/target/arm/rss/common/crypto_keys.c
index c6fade7..617e106 100644
--- a/platform/ext/target/arm/rss/common/crypto_keys.c
+++ b/platform/ext/target/arm/rss/common/crypto_keys.c
@@ -1,17 +1,13 @@
 /*
- * 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_plat_crypto_keys.h"
-
-#include "tfm_builtin_key_ids.h"
-
-#include <stddef.h>
 #include <string.h>
-
+#include "tfm_plat_crypto_keys.h"
+#include "tfm_builtin_key_ids.h"
 #include "region_defs.h"
 #include "cmsis_compiler.h"
 #include "tfm_plat_otp.h"
@@ -21,136 +17,9 @@
 #include "device_definition.h"
 #include "tfm_plat_otp.h"
 
+#define NUMBER_OF_ELEMENTS_OF(x) sizeof(x)/sizeof(*x)
 #define TFM_NS_PARTITION_ID -1
 
-#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,
-                                                   int32_t owner,
-                                                   psa_key_usage_t *usage)
-{
-    *usage = 0;
-
-    switch (key_id) {
-    case TFM_BUILTIN_KEY_ID_HUK:
-        switch(owner) {
-        default:
-            /* Allow access to all partitions */
-            *usage = PSA_KEY_USAGE_DERIVE;
-        }
-        break;
-    case TFM_BUILTIN_KEY_ID_IAK:
-        switch(owner) {
-        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_SP_SECURE_TEST_PARTITION
-        /* 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_SP_SECURE_TEST_PARTITION */
-        default:
-            return TFM_PLAT_ERR_NOT_PERMITTED;
-        }
-        break;
-#ifdef TFM_PARTITION_DELEGATED_ATTESTATION
-    case TFM_BUILTIN_KEY_ID_DAK_SEED:
-        switch(owner) {
-        case TFM_SP_DELEGATED_ATTESTATION:
-            *usage = PSA_KEY_USAGE_DERIVE;
-            break;
-        default:
-            return TFM_PLAT_ERR_NOT_PERMITTED;
-        }
-        break;
-#endif /* TFM_PARTITION_DELEGATED_ATTESTATION */
-    case TFM_BUILTIN_KEY_ID_HOST_S_ROTPK:
-        switch(owner) {
-        case TFM_NS_PARTITION_ID:
-            *usage = PSA_KEY_USAGE_VERIFY_HASH;
-            break;
-        default:
-            return TFM_PLAT_ERR_NOT_PERMITTED;
-        }
-        break;
-    case TFM_BUILTIN_KEY_ID_HOST_NS_ROTPK:
-        switch(owner) {
-        case TFM_NS_PARTITION_ID:
-            *usage = PSA_KEY_USAGE_VERIFY_HASH;
-            break;
-        default:
-            return TFM_PLAT_ERR_NOT_PERMITTED;
-        }
-        break;
-    case TFM_BUILTIN_KEY_ID_HOST_CCA_ROTPK:
-        switch(owner) {
-        case TFM_NS_PARTITION_ID:
-            *usage = PSA_KEY_USAGE_VERIFY_HASH;
-            break;
-        default:
-            return TFM_PLAT_ERR_NOT_PERMITTED;
-        }
-        break;
-    default:
-        return TFM_PLAT_ERR_UNSUPPORTED;
-    }
-
-    return TFM_PLAT_ERR_SUCCESS;
-}
-
-enum tfm_plat_err_t tfm_plat_builtin_key_get_lifetime_and_slot(
-    struct tfm_crypto_key_id_s key_id,
-    psa_key_lifetime_t *lifetime,
-    psa_drv_slot_number_t *slot_number)
-{
-    switch (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;
-    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);
-        break;
-#ifdef TFM_PARTITION_DELEGATED_ATTESTATION
-    case TFM_BUILTIN_KEY_ID_DAK_SEED:
-        *slot_number = TFM_BUILTIN_KEY_SLOT_DAK_SEED;
-        *lifetime = PSA_KEY_LIFETIME_FROM_PERSISTENCE_AND_LOCATION(PSA_KEY_LIFETIME_PERSISTENT,
-                                                                   TFM_BUILTIN_KEY_LOADER_KEY_LOCATION);
-        break;
-#endif /* TFM_PARTITION_DELEGATED_ATTESTATION */
-    case TFM_BUILTIN_KEY_ID_HOST_S_ROTPK:
-        *slot_number = TFM_BUILTIN_KEY_SLOT_HOST_S_ROTPK;
-        *lifetime = PSA_KEY_LIFETIME_FROM_PERSISTENCE_AND_LOCATION(PSA_KEY_LIFETIME_PERSISTENT,
-                                                                   TFM_BUILTIN_KEY_LOADER_KEY_LOCATION);
-        break;
-    case TFM_BUILTIN_KEY_ID_HOST_NS_ROTPK:
-        *slot_number = TFM_BUILTIN_KEY_SLOT_HOST_NS_ROTPK;
-        *lifetime = PSA_KEY_LIFETIME_FROM_PERSISTENCE_AND_LOCATION(PSA_KEY_LIFETIME_PERSISTENT,
-                                                                   TFM_BUILTIN_KEY_LOADER_KEY_LOCATION);
-        break;
-    case TFM_BUILTIN_KEY_ID_HOST_CCA_ROTPK:
-        *slot_number = TFM_BUILTIN_KEY_SLOT_HOST_CCA_ROTPK;
-        *lifetime = PSA_KEY_LIFETIME_FROM_PERSISTENCE_AND_LOCATION(PSA_KEY_LIFETIME_PERSISTENT,
-                                                                   TFM_BUILTIN_KEY_LOADER_KEY_LOCATION);
-        break;
-    default:
-        return PSA_ERROR_DOES_NOT_EXIST;
-    }
-
-    return PSA_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,
@@ -176,7 +45,6 @@
     return TFM_PLAT_ERR_SUCCESS;
 }
 
-
 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,
@@ -359,122 +227,111 @@
     return tfm_plat_otp_read(PLAT_OTP_ID_HOST_ROTPK_CCA, buf_len, buf);
 }
 
-enum tfm_plat_err_t tfm_plat_load_builtin_keys(void)
-{
-    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;
-    /* The KMU requires word alignment */
-    uint8_t __ALIGNED(4) buf[96];
-    size_t key_len;
-    size_t key_bits;
-    psa_algorithm_t algorithm;
-    psa_key_type_t type;
-
-    /* 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;
-    key_id.MBEDTLS_PRIVATE(owner) = 0;
-    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;
-    }
-
-    /* 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;
-    }
-    key_id.MBEDTLS_PRIVATE(key_id) = TFM_BUILTIN_KEY_ID_IAK;
-    key_id.MBEDTLS_PRIVATE(owner) = 0;
-    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;
-    }
+/**
+ * @brief Table describing per-user key policy for the IAK
+ *
+ */
+static const tfm_plat_builtin_key_per_user_policy_t g_iak_per_user_policy[] = {
+    {.user = TFM_SP_INITIAL_ATTESTATION,
+#ifdef SYMMETRIC_INITIAL_ATTESTATION
+        .usage = PSA_KEY_USAGE_SIGN_HASH | PSA_KEY_USAGE_EXPORT,
+#else
+        .usage = PSA_KEY_USAGE_SIGN_HASH,
+#endif /* SYMMETRIC_INITIAL_ATTESTATION */
+    },
+#ifdef TEST_S_ATTESTATION
+    {.user = TFM_SP_SECURE_TEST_PARTITION, .usage = PSA_KEY_USAGE_VERIFY_HASH},
+#endif /* TEST_S_ATTESTATION */
+#ifdef TEST_NS_ATTESTATION
+    {.user = TFM_NS_PARTITION_ID, .usage = PSA_KEY_USAGE_VERIFY_HASH},
+#endif /* TEST_NS_ATTESTATION */
+};
 
 #ifdef TFM_PARTITION_DELEGATED_ATTESTATION
-    /* DAK SEED */
-    plat_err = tfm_plat_get_dak_seed(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_DAK_SEED;
-    key_id.MBEDTLS_PRIVATE(owner) = 0;
-    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;
-    }
+/**
+ * @brief Table describing per-user key policy for the DAK seed
+ *
+ */
+static const tfm_plat_builtin_key_per_user_policy_t g_dak_seed_per_user_policy[] = {
+    {.user = TFM_SP_DELEGATED_ATTESTATION, .usage = PSA_KEY_USAGE_DERIVE},
+};
 #endif /* TFM_PARTITION_DELEGATED_ATTESTATION */
 
-    /* HOST S ROTPK */
-    plat_err = tfm_plat_get_host_s_rotpk(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_HOST_S_ROTPK;
-    key_id.MBEDTLS_PRIVATE(owner) = 0;
-    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;
-    }
+/**
+ * @brief Table describing per-user key policy for all the HOST RoTPK (S, NS, CCA)
+ *
+ */
+static const tfm_plat_builtin_key_per_user_policy_t g_host_rotpk_per_user_policy[] = {
+    {.user = TFM_NS_PARTITION_ID, .usage = PSA_KEY_USAGE_VERIFY_HASH},
+};
 
-    /* HOST NS ROTPK */
-    plat_err = tfm_plat_get_host_ns_rotpk(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_HOST_NS_ROTPK;
-    key_id.MBEDTLS_PRIVATE(owner) = 0;
-    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;
-    }
+/**
+ * @brief Table describing per-key user policies
+ *
+ */
+static const tfm_plat_builtin_key_policy_t g_builtin_keys_policy[] = {
+    {.key_id = TFM_BUILTIN_KEY_ID_HUK, .per_user_policy = 0, .usage = PSA_KEY_USAGE_DERIVE},
+    {.key_id = TFM_BUILTIN_KEY_ID_IAK,
+     .per_user_policy = NUMBER_OF_ELEMENTS_OF(g_iak_per_user_policy),
+     .policy_ptr = g_iak_per_user_policy},
+#ifdef TFM_PARTITION_DELEGATED_ATTESTATION
+    {.key_id = TFM_BUILTIN_KEY_ID_DAK_SEED,
+     .per_user_policy = NUMBER_OF_ELEMENTS_OF(g_dak_seed_per_user_policy),
+     .policy_ptr = g_dak_seed_per_user_policy},
+#endif /* TFM_PARTITION_DELEGATED_ATTESTATION */
+    {.key_id = TFM_BUILTIN_KEY_ID_HOST_S_ROTPK,
+     .per_user_policy = NUMBER_OF_ELEMENTS_OF(g_host_rotpk_per_user_policy),
+     .policy_ptr = g_host_rotpk_per_user_policy},
+    {.key_id = TFM_BUILTIN_KEY_ID_HOST_NS_ROTPK,
+     .per_user_policy = NUMBER_OF_ELEMENTS_OF(g_host_rotpk_per_user_policy),
+     .policy_ptr = g_host_rotpk_per_user_policy},
+    {.key_id = TFM_BUILTIN_KEY_ID_HOST_CCA_ROTPK,
+     .per_user_policy = NUMBER_OF_ELEMENTS_OF(g_host_rotpk_per_user_policy),
+     .policy_ptr = g_host_rotpk_per_user_policy},
+};
 
-    /* HOST CCA ROTPK */
-    plat_err = tfm_plat_get_host_cca_rotpk(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_HOST_CCA_ROTPK;
-    key_id.MBEDTLS_PRIVATE(owner) = 0;
-    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;
-    }
+/**
+ * @brief Table describing the builtin-in keys (plaform keys) available in the platform. Note
+ *        that to bind the keys to the tfm_builtin_key_loader driver, the lifetime must be
+ *        explicitly set to the one associated to the driver, i.e. TFM_BUILTIN_KEY_LOADER_LIFETIME
+ */
+static const tfm_plat_builtin_key_descriptor_t g_builtin_keys_desc[] = {
+    {.key_id = TFM_BUILTIN_KEY_ID_HUK,
+     .slot_number = TFM_BUILTIN_KEY_SLOT_HUK,
+     .lifetime = TFM_BUILTIN_KEY_LOADER_LIFETIME,
+     .loader_key_func = tfm_plat_get_huk},
+    {.key_id = TFM_BUILTIN_KEY_ID_IAK,
+     .slot_number = TFM_BUILTIN_KEY_SLOT_IAK,
+     .lifetime = TFM_BUILTIN_KEY_LOADER_LIFETIME,
+     .loader_key_func = tfm_plat_get_iak},
+#ifdef TFM_PARTITION_DELEGATED_ATTESTATION
+    {.key_id = TFM_BUILTIN_KEY_ID_DAK_SEED,
+     .slot_number = TFM_BUILTIN_KEY_SLOT_DAK_SEED,
+     .lifetime = TFM_BUILTIN_KEY_LOADER_LIFETIME,
+     .loader_key_func = tfm_plat_get_dak_seed},
+#endif /* TFM_PARTITION_DELEGATED_ATTESTATION */
+    {.key_id = TFM_BUILTIN_KEY_ID_HOST_S_ROTPK,
+     .slot_number = TFM_BUILTIN_KEY_SLOT_HOST_S_ROTPK,
+     .lifetime = TFM_BUILTIN_KEY_LOADER_LIFETIME,
+     .loader_key_func = tfm_plat_get_host_s_rotpk},
+    {.key_id = TFM_BUILTIN_KEY_ID_HOST_NS_ROTPK,
+     .slot_number = TFM_BUILTIN_KEY_SLOT_HOST_NS_ROTPK,
+     .lifetime = TFM_BUILTIN_KEY_LOADER_LIFETIME,
+     .loader_key_func = tfm_plat_get_host_ns_rotpk},
+    {.key_id = TFM_BUILTIN_KEY_ID_HOST_CCA_ROTPK,
+     .slot_number = TFM_BUILTIN_KEY_SLOT_HOST_CCA_ROTPK,
+     .lifetime = TFM_BUILTIN_KEY_LOADER_LIFETIME,
+     .loader_key_func = tfm_plat_get_host_cca_rotpk},
+};
 
-    return TFM_PLAT_ERR_SUCCESS;
+size_t tfm_plat_builtin_key_get_policy_table_ptr(const tfm_plat_builtin_key_policy_t *desc_ptr[])
+{
+    *desc_ptr = &g_builtin_keys_policy[0];
+    return NUMBER_OF_ELEMENTS_OF(g_builtin_keys_policy);
+}
+
+size_t tfm_plat_builtin_key_get_desc_table_ptr(const tfm_plat_builtin_key_descriptor_t *desc_ptr[])
+{
+    *desc_ptr = &g_builtin_keys_desc[0];
+    return NUMBER_OF_ELEMENTS_OF(g_builtin_keys_desc);
 }
diff --git a/platform/ext/target/arm/rss/common/platform_builtin_key_ids.h b/platform/ext/target/arm/rss/common/platform_builtin_key_ids.h
deleted file mode 100644
index b45d275..0000000
--- a/platform/ext/target/arm/rss/common/platform_builtin_key_ids.h
+++ /dev/null
@@ -1,46 +0,0 @@
-/*
- * Copyright (c) 2022, Arm Limited. All rights reserved.
- *
- * SPDX-License-Identifier: BSD-3-Clause
- *
- */
-
-#ifndef __PLATFORM_BUILTIN_KEY_IDS_H__
-#define __PLATFORM_BUILTIN_KEY_IDS_H__
-
-#ifdef __cplusplus
-extern "C" {
-#endif
-
-/**
- * \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_key_id_builtin_t {
-    TFM_BUILTIN_KEY_ID_MIN = 0x7fff815Bu,
-    TFM_BUILTIN_KEY_ID_HUK,
-    TFM_BUILTIN_KEY_ID_IAK,
-#ifdef TFM_PARTITION_DELEGATED_ATTESTATION
-    TFM_BUILTIN_KEY_ID_DAK_SEED,
-#endif /* TFM_PARTITION_DELEGATED_ATTESTATION */
-    TFM_BUILTIN_KEY_ID_PLAT_SPECIFIC_MIN = 0x7FFF816Bu,
-    TFM_BUILTIN_KEY_ID_HOST_S_ROTPK,
-    TFM_BUILTIN_KEY_ID_HOST_NS_ROTPK,
-    TFM_BUILTIN_KEY_ID_HOST_CCA_ROTPK,
-    TFM_BUILTIN_KEY_ID_MAX = 0x7FFF817Bu,
-};
-
-#ifdef __cplusplus
-}
-#endif
-#endif /* __PLATFORM_BUILTIN_KEY_IDS_H__ */
diff --git a/platform/ext/target/arm/rss/common/platform_builtin_key_loader_ids.h b/platform/ext/target/arm/rss/common/platform_builtin_key_loader_ids.h
index c2b66f0..7e4dcd9 100644
--- a/platform/ext/target/arm/rss/common/platform_builtin_key_loader_ids.h
+++ b/platform/ext/target/arm/rss/common/platform_builtin_key_loader_ids.h
@@ -1,12 +1,12 @@
 /*
- * Copyright (c) 2022, Arm Limited. All rights reserved.
+ * Copyright (c) 2022-2023, Arm Limited. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  *
  */
 
-#ifndef TFM_BUILTIN_KEY_LOADER_IDS_H
-#define TFM_BUILTIN_KEY_LOADER_IDS_H
+#ifndef __PLATFORM_BUILTIN_KEY_LOADER_IDS_H__
+#define __PLATFORM_BUILTIN_KEY_LOADER_IDS_H__
 
 #ifdef __cplusplus
 extern "C" {
@@ -30,4 +30,4 @@
 }
 #endif
 
-#endif /* TFM_BUILTIN_KEY_LOADER_IDS_H */
+#endif /* __PLATFORM_BUILTIN_KEY_LOADER_IDS_H__ */
diff --git a/platform/ext/target/arm/rss/common/tfm_builtin_key_ids.h b/platform/ext/target/arm/rss/common/tfm_builtin_key_ids.h
index b45d275..9d2cfd9 100644
--- a/platform/ext/target/arm/rss/common/tfm_builtin_key_ids.h
+++ b/platform/ext/target/arm/rss/common/tfm_builtin_key_ids.h
@@ -1,25 +1,18 @@
 /*
- * Copyright (c) 2022, Arm Limited. All rights reserved.
+ * Copyright (c) 2022-2023, Arm Limited. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  *
  */
 
-#ifndef __PLATFORM_BUILTIN_KEY_IDS_H__
-#define __PLATFORM_BUILTIN_KEY_IDS_H__
+#ifndef __TFM_BUILTIN_KEY_IDS_H__
+#define __TFM_BUILTIN_KEY_IDS_H__
 
 #ifdef __cplusplus
 extern "C" {
 #endif
 
 /**
- * \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
@@ -27,7 +20,7 @@
  *
  */
 enum tfm_key_id_builtin_t {
-    TFM_BUILTIN_KEY_ID_MIN = 0x7fff815Bu,
+    TFM_BUILTIN_KEY_ID_MIN = 0x7FFF815Bu,
     TFM_BUILTIN_KEY_ID_HUK,
     TFM_BUILTIN_KEY_ID_IAK,
 #ifdef TFM_PARTITION_DELEGATED_ATTESTATION
@@ -43,4 +36,4 @@
 #ifdef __cplusplus
 }
 #endif
-#endif /* __PLATFORM_BUILTIN_KEY_IDS_H__ */
+#endif /* __TFM_BUILTIN_KEY_IDS_H__ */
diff --git a/platform/include/tfm_plat_crypto_keys.h b/platform/include/tfm_plat_crypto_keys.h
index c692f2c..59146a1 100644
--- a/platform/include/tfm_plat_crypto_keys.h
+++ b/platform/include/tfm_plat_crypto_keys.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2017-2022, Arm Limited. All rights reserved.
+ * Copyright (c) 2017-2023, Arm Limited. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  *
@@ -8,56 +8,89 @@
 #ifndef __TFM_PLAT_CRYPTO_KEYS_H__
 #define __TFM_PLAT_CRYPTO_KEYS_H__
 
-#include "psa/crypto.h"
-
-/**
- * \note The interfaces defined in this file must be implemented for each
- *       SoC.
- */
-#include "tfm_crypto_key.h"
-
-#include "tfm_plat_defs.h"
-
-#include <stddef.h>
 #include <stdint.h>
+#include "psa/crypto.h"
+#include "tfm_plat_defs.h"
 
 #ifdef __cplusplus
 extern "C" {
 #endif
 
 /**
- * \brief Gets key usage for a given builtin key ID and owner.
+ * \brief Callback function type platform key loader functions
  *
- * \param[in]  key_id        ID of key
- * \param[in]  owner         Which owner to get the usage permissions for
- * \param[out] usage         The permissions that the given user has for the key
+ * This function pointer type defines the prototype for a builtin key loader function so that the
+ * key can be probed by the tfm_builtin_key_loader driver during the init phase. Note that the key
+ * must be readable from the secure processing element to be able to use the tfm_builtin_key_loader
  *
- * \return Returns error code specified in \ref tfm_plat_err_t
+ * \param[out] buf       Buffer to hold the retrieved key material from the platform
+ * \param[in]  buf_len   Size of the buf buffer
+ * \param[out] key_len   Actual length of the key material in bytes retrieved from the platform
+ * \param[out] key_bits  Size in bits of the key (important for those keys that are not byte-multiple)
+ * \param[out] algorithm \ref psa_algorithm_t value associated to the retrieved key material
+ * \param[out] type      \ref psa_key_type_t value associated to the retrieved key material
+ *
+ * \return Returns an error value as specified by the \ref tfm_plat_err_t type.
+ * 
  */
-enum tfm_plat_err_t tfm_plat_builtin_key_get_usage(psa_key_id_t key_id,
-                                                   int32_t owner,
-                                                   psa_key_usage_t *usage);
+typedef enum tfm_plat_err_t (*key_loader_func_ptr)
+    (uint8_t *buf, size_t buf_len, size_t *key_len, size_t *key_bits, psa_algorithm_t *algorithm, psa_key_type_t *type);
 
 /**
- * \brief Gets key lifetime and slot number for a given builtin key ID.
- *
- * \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
+ * \brief This type describes the information that each TF-M builtin key
+ *        must key in the associated descriptor table in \ref crypto_keys.c
  */
-enum tfm_plat_err_t tfm_plat_builtin_key_get_lifetime_and_slot(
-    struct tfm_crypto_key_id_s key_id,
-    psa_key_lifetime_t *lifetime,
-    psa_drv_slot_number_t *slot_number);
+typedef struct {
+    psa_key_id_t key_id;                 /*!< Key id associated to the builtin key */
+    psa_drv_slot_number_t slot_number;   /*!< Slot number for the builtin key in the platform */
+    psa_key_lifetime_t lifetime;         /*!< Lifetime (persistence + location) for the builtin key */
+    key_loader_func_ptr loader_key_func; /*!< Loader function that reads the key from the platform */
+} tfm_plat_builtin_key_descriptor_t;
 
 /**
- * \brief Load all builtin keys.
+ * \brief This function retrieves a pointer to the description table for builtin keys. Each platform
+ *        must implement this table with the details of the builtin keys available in the platform
  *
- * \return Returns error code specified in \ref tfm_plat_err_t
+ * \param[out] desc_ptr A pointer to the description table
+ *
+ * \return size_t The number of builtin keys available in the platform
  */
-enum tfm_plat_err_t tfm_plat_load_builtin_keys(void);
+size_t tfm_plat_builtin_key_get_desc_table_ptr(const tfm_plat_builtin_key_descriptor_t *desc_ptr[]);
+
+/**
+ * \brief This type maps a particular user of a builtin key (i.e. an owner) to
+ *        the allowed usage (i.e. a policy) as specified by the platform
+ */
+typedef struct {
+    int32_t user;
+    psa_key_usage_t usage;
+} tfm_plat_builtin_key_per_user_policy_t;
+
+/**
+ * \brief This type maps a particular key_id associated to a builtin key to the
+ *        allowed usage (i.e. a policy). The policy can be user specific in case
+ *        the per_user_policy field is greater than 0. In that case policy_ptr needs
+ *        to be used to access the policies for each user of the key_id which are of
+ *        type \ref tfm_platf_builtin_key_per_user_policy_t
+ */
+typedef struct {
+    psa_key_id_t key_id;
+    size_t per_user_policy;
+    union {
+        psa_key_usage_t usage;
+        const tfm_plat_builtin_key_per_user_policy_t *policy_ptr;
+    };
+} tfm_plat_builtin_key_policy_t;
+
+/**
+ * \brief This function retrieves a pointer to the policy table of the builtin keys. Each platform
+ *        must implement this table with the details of the builtin keys available in the platform
+ *
+ * \param[out] desc_ptr A pointer to the policy table
+ *
+ * \return size_t The number of builtin keys available in the platform with associated policies
+ */
+size_t tfm_plat_builtin_key_get_policy_table_ptr(const tfm_plat_builtin_key_policy_t *desc_ptr[]);
 
 #ifdef __cplusplus
 }
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 488f560..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);