Merge pull request #4105 from paul-elliott-arm/fix_test_leak

Fix for memory leak in ssl tests
diff --git a/ChangeLog b/ChangeLog
index fb231aa..184bd09 100644
--- a/ChangeLog
+++ b/ChangeLog
@@ -20,7 +20,7 @@
      warning on CMake 3.19.0. #3801
 
 New deprecations
-   * PSA_KEY_TYPE_CHACHA20 and PSA_KEY_TYPE_ARC4 have been deprecated.
+   * PSA_ALG_CHACHA20 and PSA_ALG_ARC4 have been deprecated.
      Use PSA_ALG_STREAM_CIPHER instead.
    * The functions mbedtls_cipher_auth_encrypt() and
      mbedtls_cipher_auth_decrypt() are deprecated in favour of the new
diff --git a/docs/architecture/psa-crypto-implementation-structure.md b/docs/architecture/psa-crypto-implementation-structure.md
new file mode 100644
index 0000000..025a623
--- /dev/null
+++ b/docs/architecture/psa-crypto-implementation-structure.md
@@ -0,0 +1,73 @@
+PSA Cryptograpy API implementation and PSA driver interface
+===========================================================
+
+## Introduction
+
+The [PSA Cryptography API specification](https://armmbed.github.io/mbed-crypto/psa/#application-programming-interface) defines an interface to cryptographic operations for which the Mbed TLS library provides a reference implementation. The PSA Cryptography API specification is complemented by the PSA driver interface specification which defines an interface for cryptoprocessor drivers.
+
+This document describes the high level organization of the Mbed TLS PSA Cryptography API implementation which is tightly related to the PSA driver interface.
+
+## High level organization of the Mbed TLS PSA Cryptography API implementation
+In one sentence, the Mbed TLS PSA Cryptography API implementation is made of a core and PSA drivers as defined in the PSA driver interface. The key point is that software cryptographic operations are organized as PSA drivers: they interact with the core through the PSA driver interface.
+
+### Rationale
+
+* Addressing software and hardware cryptographic implementations through the same C interface reduces the core code size and its call graph complexity. The core and its dispatching to software and hardware implementations are consequently easier to test and validate.
+* The organization of the software cryptographic implementations in drivers promotes modularization of those implementations.
+* As hardware capabilities, software cryptographic functionalities can be described by a JSON driver description file as defined in the PSA driver interface.
+* Along with JSON driver description files, the PSA driver specification defines the deliverables for a driver to be included into the Mbed TLS PSA Cryptography implementation. This provides a natural framework to integrate third party or alternative software implementations of cryptographic operations.
+
+## The Mbed TLS PSA Cryptography API implementation core
+
+The core implements all the APIs as defined in the PSA Cryptography API specification but does not perform on its own any cryptographic operation. The core relies on PSA drivers to actually
+perform the cryptographic operations. The core is responsible for:
+
+* the key store.
+* checking PSA API arguments and translating them into valid arguments for the necessary calls to the PSA driver interface.
+* dispatching the cryptographic operations to the appropriate PSA drivers.
+
+The sketch of an Mbed TLS PSA cryptographic API implementation is thus:
+```C                                                                            
+psa_status_t psa_api( ... )
+{
+    psa_status_t status;
+
+    /* Pre driver interface call processing: validation of arguments, building
+     * of arguments for the call to the driver interface, ... */
+
+    ...
+
+    /* Call to the driver interface */
+    status = psa_driver_wrapper_<entry_point>( ... );
+    if( status != PSA_SUCCESS )
+        return( status );
+
+    /* Post driver interface call processing: validation of the values returned
+     * by the driver, finalization of the values to return to the caller,
+     * clean-up in case of error ... */
+}
+```
+The code of most PSA APIs is expected to match precisely the above layout. However, it is likely that the code structure of some APIs will be more complicated with several calls to the driver interface, mainly to encompass a larger variety of hardware designs. For example, to encompass hardware accelerators that are capable of verifying a MAC and those that are only capable of computing a MAC, the psa_mac_verify() API could call first psa_driver_wrapper_mac_verify() and then fallback to psa_driver_wrapper_mac_compute().
+
+The implementations of `psa_driver_wrapper_<entry_point>` functions are generated by the build system based on the JSON driver description files of the various PSA drivers making up the Mbed TLS PSA Cryptography API implementation. The implementations are generated in a psa_crypto_driver_wrappers.c C file and the function prototypes declared in a psa_crypto_driver_wrappers.h header file.
+
+The psa_driver_wrapper_<entry_point>() functions dispatch cryptographic operations to accelerator drivers, secure element drivers as well as to the software implementations of cryptographic operations.
+
+Note that the implementation allows to build the library with only a C compiler by shipping a generated file corresponding to a pure software implementation. The driver entry points and their code in this generated file are guarded by pre-processor directives based on PSA_WANT_xyz macros (see [Conditional inclusion of cryptographic mechanism through the PSA API in Mbed TLS](psa-conditional-inclusion-c.html). That way, it is possible to compile and include in the library only the desired cryptographic operations.
+
+### Key creation
+
+Key creation implementation in Mbed TLS PSA core is articulated around three internal functions: psa_start_key_creation(), psa_finish_key_creation() and psa_fail_key_creation(). Implementations of key creation PSA APIs, namely psa_import_key(), psa_generate_key(), psa_key_derivation_output_key() and psa_copy_key() go by the following sequence:
+    1. Check the input parameters.
+    2. Call psa_start_key_creation() that allocates a key slot, prepares it with the specified key attributes, and in case of a volatile key assign it a volatile key identifier.
+    3. Generate or copy the key material into the key slot. This entails the allocation of the buffer to store the key material.
+    4. Call psa_finish_key_creation() that mostly saves persistent keys into persistent storage.
+
+In case of any error occuring at step 3 or 4, psa_fail_key_creation() is called. It wipes and cleans the slot especially the key material: reset to zero of the RAM memory that contained the key material, free the allocated buffer.
+
+
+## Mbed TLS PSA Cryptography API implementation drivers
+
+A driver of the Mbed TLS PSA Cryptography API implementation (Mbed TLS PSA driver in the following) is a driver in the sense that it is compliant with the PSA driver interface specification. But it is not an actual driver that drives some hardware. It implements cryptographic operations purely in software.
+
+An Mbed TLS PSA driver C file is named psa_crypto_<driver_name>.c and its associated header file psa_crypto_<driver_name>.h. The functions implementing a driver entry point as defined in the PSA driver interface specification are named as mbedtls_psa_<driver name>_<entry point>(). As an example, the psa_crypto_rsa.c and psa_crypto_rsa.h are the files containing the Mbed TLS PSA driver implementing RSA cryptographic operations. This RSA driver implements among other entry points the "import_key" entry point. The function implementing this entry point is named mbedtls_psa_rsa_import_key().
diff --git a/include/psa/crypto.h b/include/psa/crypto.h
index 1adeda6..271fb99 100644
--- a/include/psa/crypto.h
+++ b/include/psa/crypto.h
@@ -90,10 +90,14 @@
  *
  * \retval #PSA_SUCCESS
  * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
+ * \retval #PSA_ERROR_INSUFFICIENT_STORAGE
  * \retval #PSA_ERROR_COMMUNICATION_FAILURE
  * \retval #PSA_ERROR_HARDWARE_FAILURE
  * \retval #PSA_ERROR_CORRUPTION_DETECTED
  * \retval #PSA_ERROR_INSUFFICIENT_ENTROPY
+ * \retval #PSA_ERROR_STORAGE_FAILURE
+ * \retval #PSA_ERROR_DATA_INVALID
+ * \retval #PSA_ERROR_DATA_CORRUPT
  */
 psa_status_t psa_crypto_init(void);
 
@@ -368,6 +372,8 @@
  * \retval #PSA_ERROR_COMMUNICATION_FAILURE
  * \retval #PSA_ERROR_CORRUPTION_DETECTED
  * \retval #PSA_ERROR_STORAGE_FAILURE
+ * \retval #PSA_ERROR_DATA_CORRUPT
+ * \retval #PSA_ERROR_DATA_INVALID
  * \retval #PSA_ERROR_BAD_STATE
  *         The library has not been previously initialized by psa_crypto_init().
  *         It is implementation-dependent whether a failure to initialize
@@ -501,6 +507,8 @@
  * \retval #PSA_ERROR_INSUFFICIENT_STORAGE
  * \retval #PSA_ERROR_COMMUNICATION_FAILURE
  * \retval #PSA_ERROR_HARDWARE_FAILURE
+ * \retval #PSA_ERROR_DATA_INVALID
+ * \retval #PSA_ERROR_DATA_CORRUPT
  * \retval #PSA_ERROR_STORAGE_FAILURE
  * \retval #PSA_ERROR_CORRUPTION_DETECTED
  * \retval #PSA_ERROR_BAD_STATE
@@ -540,6 +548,10 @@
  * \retval #PSA_ERROR_COMMUNICATION_FAILURE
  *         There was an failure in communication with the cryptoprocessor.
  *         The key material may still be present in the cryptoprocessor.
+ * \retval #PSA_ERROR_DATA_INVALID
+ *         This error is typically a result of either storage corruption on a
+ *         cleartext storage backend, or an attempt to read data that was
+ *         written by an incompatible version of the library.
  * \retval #PSA_ERROR_STORAGE_FAILURE
  *         The storage is corrupted. Implementations shall make a best effort
  *         to erase key material even in this stage, however applications
@@ -625,6 +637,8 @@
  * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
  * \retval #PSA_ERROR_INSUFFICIENT_STORAGE
  * \retval #PSA_ERROR_COMMUNICATION_FAILURE
+ * \retval #PSA_ERROR_DATA_CORRUPT
+ * \retval #PSA_ERROR_DATA_INVALID
  * \retval #PSA_ERROR_STORAGE_FAILURE
  * \retval #PSA_ERROR_HARDWARE_FAILURE
  * \retval #PSA_ERROR_CORRUPTION_DETECTED
@@ -3556,6 +3570,8 @@
  * \retval #PSA_ERROR_COMMUNICATION_FAILURE
  * \retval #PSA_ERROR_HARDWARE_FAILURE
  * \retval #PSA_ERROR_CORRUPTION_DETECTED
+ * \retval #PSA_ERROR_DATA_INVALID
+ * \retval #PSA_ERROR_DATA_CORRUPT
  * \retval #PSA_ERROR_STORAGE_FAILURE
  * \retval #PSA_ERROR_BAD_STATE
  *         The library has not been previously initialized by psa_crypto_init().
@@ -3721,6 +3737,8 @@
  * \retval #PSA_ERROR_HARDWARE_FAILURE
  * \retval #PSA_ERROR_CORRUPTION_DETECTED
  * \retval #PSA_ERROR_INSUFFICIENT_STORAGE
+ * \retval #PSA_ERROR_DATA_INVALID
+ * \retval #PSA_ERROR_DATA_CORRUPT
  * \retval #PSA_ERROR_STORAGE_FAILURE
  * \retval #PSA_ERROR_BAD_STATE
  *         The library has not been previously initialized by psa_crypto_init().
diff --git a/include/psa/crypto_compat.h b/include/psa/crypto_compat.h
index 728d283..df6645a 100644
--- a/include/psa/crypto_compat.h
+++ b/include/psa/crypto_compat.h
@@ -316,6 +316,8 @@
  * \retval #PSA_ERROR_COMMUNICATION_FAILURE
  * \retval #PSA_ERROR_CORRUPTION_DETECTED
  * \retval #PSA_ERROR_STORAGE_FAILURE
+ * \retval #PSA_ERROR_DATA_INVALID
+ * \retval #PSA_ERROR_DATA_CORRUPT
  * \retval #PSA_ERROR_BAD_STATE
  *         The library has not been previously initialized by psa_crypto_init().
  *         It is implementation-dependent whether a failure to initialize
diff --git a/include/psa/crypto_extra.h b/include/psa/crypto_extra.h
index 9abce33..14b5be3 100644
--- a/include/psa/crypto_extra.h
+++ b/include/psa/crypto_extra.h
@@ -183,8 +183,10 @@
  * \retval #PSA_ERROR_NOT_PERMITTED
  *         The caller is not authorized to register the specified key slot.
  * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
+ * \retval #PSA_ERROR_INSUFFICIENT_STORAGE
  * \retval #PSA_ERROR_COMMUNICATION_FAILURE
- * \retval #PSA_ERROR_HARDWARE_FAILURE
+ * \retval #PSA_ERROR_DATA_INVALID
+ * \retval #PSA_ERROR_DATA_CORRUPT
  * \retval #PSA_ERROR_CORRUPTION_DETECTED
  * \retval #PSA_ERROR_BAD_STATE
  *         The library has not been previously initialized by psa_crypto_init().
diff --git a/include/psa/crypto_se_driver.h b/include/psa/crypto_se_driver.h
index 1fae575..f5fe029 100644
--- a/include/psa/crypto_se_driver.h
+++ b/include/psa/crypto_se_driver.h
@@ -1378,6 +1378,8 @@
  *         `methods->hal_version` is not supported by this implementation.
  * \return #PSA_ERROR_INSUFFICIENT_MEMORY
  * \return #PSA_ERROR_NOT_PERMITTED
+ * \return #PSA_ERROR_STORAGE_FAILURE
+ * \return #PSA_ERROR_DATA_CORRUPT
  */
 psa_status_t psa_register_se_driver(
     psa_key_location_t location,
diff --git a/include/psa/crypto_struct.h b/include/psa/crypto_struct.h
index 6a018e1..04ece6d 100644
--- a/include/psa/crypto_struct.h
+++ b/include/psa/crypto_struct.h
@@ -228,11 +228,11 @@
 #if defined(MBEDTLS_MD_C)
 typedef enum
 {
-    TLS12_PRF_STATE_INIT,       /* no input provided */
-    TLS12_PRF_STATE_SEED_SET,   /* seed has been set */
-    TLS12_PRF_STATE_KEY_SET,    /* key has been set */
-    TLS12_PRF_STATE_LABEL_SET,  /* label has been set */
-    TLS12_PRF_STATE_OUTPUT      /* output has been started */
+    PSA_TLS12_PRF_STATE_INIT,       /* no input provided */
+    PSA_TLS12_PRF_STATE_SEED_SET,   /* seed has been set */
+    PSA_TLS12_PRF_STATE_KEY_SET,    /* key has been set */
+    PSA_TLS12_PRF_STATE_LABEL_SET,  /* label has been set */
+    PSA_TLS12_PRF_STATE_OUTPUT      /* output has been started */
 } psa_tls12_prf_key_derivation_state_t;
 
 typedef struct psa_tls12_prf_key_derivation_s
diff --git a/include/psa/crypto_values.h b/include/psa/crypto_values.h
index 81a8575..54273ed 100644
--- a/include/psa/crypto_values.h
+++ b/include/psa/crypto_values.h
@@ -270,6 +270,46 @@
  */
 #define PSA_ERROR_INVALID_HANDLE        ((psa_status_t)-136)
 
+/** Stored data has been corrupted.
+ *
+ * This error indicates that some persistent storage has suffered corruption.
+ * It does not indicate the following situations, which have specific error
+ * codes:
+ *
+ * - A corruption of volatile memory - use #PSA_ERROR_CORRUPTION_DETECTED.
+ * - A communication error between the cryptoprocessor and its external
+ *   storage - use #PSA_ERROR_COMMUNICATION_FAILURE.
+ * - When the storage is in a valid state but is full - use
+ *   #PSA_ERROR_INSUFFICIENT_STORAGE.
+ * - When the storage fails for other reasons - use
+ *   #PSA_ERROR_STORAGE_FAILURE.
+ * - When the stored data is not valid - use #PSA_ERROR_DATA_INVALID.
+ *
+ * \note A storage corruption does not indicate that any data that was
+ * previously read is invalid. However this previously read data might no
+ * longer be readable from storage.
+ *
+ * When a storage failure occurs, it is no longer possible to ensure the
+ * global integrity of the keystore.
+ */
+#define PSA_ERROR_DATA_CORRUPT          ((psa_status_t)-152)
+
+/** Data read from storage is not valid for the implementation.
+ *
+ * This error indicates that some data read from storage does not have a valid
+ * format. It does not indicate the following situations, which have specific
+ * error codes:
+ *
+ * - When the storage or stored data is corrupted - use #PSA_ERROR_DATA_CORRUPT
+ * - When the storage fails for other reasons - use #PSA_ERROR_STORAGE_FAILURE
+ * - An invalid argument to the API - use #PSA_ERROR_INVALID_ARGUMENT
+ *
+ * This error is typically a result of either storage corruption on a
+ * cleartext storage backend, or an attempt to read data that was
+ * written by an incompatible version of the library.
+ */
+#define PSA_ERROR_DATA_INVALID          ((psa_status_t)-153)
+
 /**@}*/
 
 /** \defgroup crypto_types Key and algorithm types
diff --git a/library/CMakeLists.txt b/library/CMakeLists.txt
index b309b6e..3ecbf9a 100644
--- a/library/CMakeLists.txt
+++ b/library/CMakeLists.txt
@@ -62,6 +62,8 @@
     poly1305.c
     psa_crypto.c
     psa_crypto_driver_wrappers.c
+    psa_crypto_ecp.c
+    psa_crypto_rsa.c
     psa_crypto_se.c
     psa_crypto_slot_management.c
     psa_crypto_storage.c
diff --git a/library/Makefile b/library/Makefile
index ae33bf2..5ef2923 100644
--- a/library/Makefile
+++ b/library/Makefile
@@ -119,6 +119,8 @@
 	     poly1305.o \
 	     psa_crypto.o \
 	     psa_crypto_driver_wrappers.o \
+	     psa_crypto_ecp.o \
+	     psa_crypto_rsa.o \
 	     psa_crypto_se.o \
 	     psa_crypto_slot_management.o \
 	     psa_crypto_storage.o \
diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 39144a3..88e25e1 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -32,6 +32,8 @@
 #include "psa_crypto_core.h"
 #include "psa_crypto_invasive.h"
 #include "psa_crypto_driver_wrappers.h"
+#include "psa_crypto_ecp.h"
+#include "psa_crypto_rsa.h"
 #if defined(MBEDTLS_PSA_CRYPTO_SE_C)
 #include "psa_crypto_se.h"
 #endif
@@ -531,466 +533,6 @@
     return( PSA_SUCCESS );
 }
 
-#if defined(MBEDTLS_PSA_BUILTIN_ALG_RSA_PKCS1V15_CRYPT) || \
-    defined(MBEDTLS_PSA_BUILTIN_ALG_RSA_PKCS1V15_SIGN) || \
-    defined(MBEDTLS_PSA_BUILTIN_ALG_RSA_OAEP) || \
-    defined(MBEDTLS_PSA_BUILTIN_ALG_RSA_PSS) || \
-    defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_RSA_KEY_PAIR) || \
-    defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_RSA_PUBLIC_KEY)
-
-/* Mbed TLS doesn't support non-byte-aligned key sizes (i.e. key sizes
- * that are not a multiple of 8) well. For example, there is only
- * mbedtls_rsa_get_len(), which returns a number of bytes, and no
- * way to return the exact bit size of a key.
- * To keep things simple, reject non-byte-aligned key sizes. */
-static psa_status_t psa_check_rsa_key_byte_aligned(
-    const mbedtls_rsa_context *rsa )
-{
-    mbedtls_mpi n;
-    psa_status_t status;
-    mbedtls_mpi_init( &n );
-    status = mbedtls_to_psa_error(
-        mbedtls_rsa_export( rsa, &n, NULL, NULL, NULL, NULL ) );
-    if( status == PSA_SUCCESS )
-    {
-        if( mbedtls_mpi_bitlen( &n ) % 8 != 0 )
-            status = PSA_ERROR_NOT_SUPPORTED;
-    }
-    mbedtls_mpi_free( &n );
-    return( status );
-}
-
-/** Load the contents of a key buffer into an internal RSA representation
- *
- * \param[in] type          The type of key contained in \p data.
- * \param[in] data          The buffer from which to load the representation.
- * \param[in] data_length   The size in bytes of \p data.
- * \param[out] p_rsa        Returns a pointer to an RSA context on success.
- *                          The caller is responsible for freeing both the
- *                          contents of the context and the context itself
- *                          when done.
- */
-static psa_status_t psa_load_rsa_representation( psa_key_type_t type,
-                                                 const uint8_t *data,
-                                                 size_t data_length,
-                                                 mbedtls_rsa_context **p_rsa )
-{
-    psa_status_t status;
-    mbedtls_pk_context ctx;
-    size_t bits;
-    mbedtls_pk_init( &ctx );
-
-    /* Parse the data. */
-    if( PSA_KEY_TYPE_IS_KEY_PAIR( type ) )
-        status = mbedtls_to_psa_error(
-            mbedtls_pk_parse_key( &ctx, data, data_length, NULL, 0 ) );
-    else
-        status = mbedtls_to_psa_error(
-            mbedtls_pk_parse_public_key( &ctx, data, data_length ) );
-    if( status != PSA_SUCCESS )
-        goto exit;
-
-    /* We have something that the pkparse module recognizes. If it is a
-     * valid RSA key, store it. */
-    if( mbedtls_pk_get_type( &ctx ) != MBEDTLS_PK_RSA )
-    {
-        status = PSA_ERROR_INVALID_ARGUMENT;
-        goto exit;
-    }
-
-    /* The size of an RSA key doesn't have to be a multiple of 8. Mbed TLS
-     * supports non-byte-aligned key sizes, but not well. For example,
-     * mbedtls_rsa_get_len() returns the key size in bytes, not in bits. */
-    bits = PSA_BYTES_TO_BITS( mbedtls_rsa_get_len( mbedtls_pk_rsa( ctx ) ) );
-    if( bits > PSA_VENDOR_RSA_MAX_KEY_BITS )
-    {
-        status = PSA_ERROR_NOT_SUPPORTED;
-        goto exit;
-    }
-    status = psa_check_rsa_key_byte_aligned( mbedtls_pk_rsa( ctx ) );
-    if( status != PSA_SUCCESS )
-        goto exit;
-
-    /* Copy out the pointer to the RSA context, and reset the PK context
-     * such that pk_free doesn't free the RSA context we just grabbed. */
-    *p_rsa = mbedtls_pk_rsa( ctx );
-    ctx.pk_info = NULL;
-
-exit:
-    mbedtls_pk_free( &ctx );
-    return( status );
-}
-
-#endif /* defined(MBEDTLS_PSA_BUILTIN_ALG_RSA_PKCS1V15_CRYPT) ||
-        * defined(MBEDTLS_PSA_BUILTIN_ALG_RSA_PKCS1V15_SIGN) ||
-        * defined(MBEDTLS_PSA_BUILTIN_ALG_RSA_OAEP) ||
-        * defined(MBEDTLS_PSA_BUILTIN_ALG_RSA_PSS) ||
-        * defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_RSA_KEY_PAIR) ||
-        * defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_RSA_PUBLIC_KEY) */
-
-#if defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_RSA_KEY_PAIR) || \
-    defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_RSA_PUBLIC_KEY)
-
-/** Export an RSA key to export representation
- *
- * \param[in] type          The type of key (public/private) to export
- * \param[in] rsa           The internal RSA representation from which to export
- * \param[out] data         The buffer to export to
- * \param[in] data_size     The length of the buffer to export to
- * \param[out] data_length  The amount of bytes written to \p data
- */
-static psa_status_t psa_export_rsa_key( psa_key_type_t type,
-                                        mbedtls_rsa_context *rsa,
-                                        uint8_t *data,
-                                        size_t data_size,
-                                        size_t *data_length )
-{
-#if defined(MBEDTLS_PK_WRITE_C)
-    int ret;
-    mbedtls_pk_context pk;
-    uint8_t *pos = data + data_size;
-
-    mbedtls_pk_init( &pk );
-    pk.pk_info = &mbedtls_rsa_info;
-    pk.pk_ctx = rsa;
-
-    /* PSA Crypto API defines the format of an RSA key as a DER-encoded
-     * representation of the non-encrypted PKCS#1 RSAPrivateKey for a
-     * private key and of the RFC3279 RSAPublicKey for a public key. */
-    if( PSA_KEY_TYPE_IS_KEY_PAIR( type ) )
-        ret = mbedtls_pk_write_key_der( &pk, data, data_size );
-    else
-        ret = mbedtls_pk_write_pubkey( &pos, data, &pk );
-
-    if( ret < 0 )
-    {
-        /* Clean up in case pk_write failed halfway through. */
-        memset( data, 0, data_size );
-        return( mbedtls_to_psa_error( ret ) );
-    }
-
-    /* The mbedtls_pk_xxx functions write to the end of the buffer.
-     * Move the data to the beginning and erase remaining data
-     * at the original location. */
-    if( 2 * (size_t) ret <= data_size )
-    {
-        memcpy( data, data + data_size - ret, ret );
-        memset( data + data_size - ret, 0, ret );
-    }
-    else if( (size_t) ret < data_size )
-    {
-        memmove( data, data + data_size - ret, ret );
-        memset( data + ret, 0, data_size - ret );
-    }
-
-    *data_length = ret;
-    return( PSA_SUCCESS );
-#else
-    (void) type;
-    (void) rsa;
-    (void) data;
-    (void) data_size;
-    (void) data_length;
-    return( PSA_ERROR_NOT_SUPPORTED );
-#endif /* MBEDTLS_PK_WRITE_C */
-}
-
-/** Import an RSA key from import representation to a slot
- *
- * \param[in,out] slot      The slot where to store the export representation to
- * \param[in] data          The buffer containing the import representation
- * \param[in] data_length   The amount of bytes in \p data
- */
-static psa_status_t psa_import_rsa_key( psa_key_slot_t *slot,
-                                        const uint8_t *data,
-                                        size_t data_length )
-{
-    psa_status_t status;
-    uint8_t* output = NULL;
-    mbedtls_rsa_context *rsa = NULL;
-
-    /* Parse input */
-    status = psa_load_rsa_representation( slot->attr.type,
-                                          data,
-                                          data_length,
-                                          &rsa );
-    if( status != PSA_SUCCESS )
-        goto exit;
-
-    slot->attr.bits = (psa_key_bits_t) PSA_BYTES_TO_BITS(
-        mbedtls_rsa_get_len( rsa ) );
-
-    /* Re-export the data to PSA export format, such that we can store export
-     * representation in the key slot. Export representation in case of RSA is
-     * the smallest representation that's allowed as input, so a straight-up
-     * allocation of the same size as the input buffer will be large enough. */
-    output = mbedtls_calloc( 1, data_length );
-    if( output == NULL )
-    {
-        status = PSA_ERROR_INSUFFICIENT_MEMORY;
-        goto exit;
-    }
-
-    status = psa_export_rsa_key( slot->attr.type,
-                                 rsa,
-                                 output,
-                                 data_length,
-                                 &data_length);
-exit:
-    /* Always free the RSA object */
-    mbedtls_rsa_free( rsa );
-    mbedtls_free( rsa );
-
-    /* Free the allocated buffer only on error. */
-    if( status != PSA_SUCCESS )
-    {
-        mbedtls_free( output );
-        return( status );
-    }
-
-    /* On success, store the allocated export-formatted key. */
-    slot->data.key.data = output;
-    slot->data.key.bytes = data_length;
-
-    return( PSA_SUCCESS );
-}
-
-#endif /* defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_RSA_KEY_PAIR) ||
-        * defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_RSA_PUBLIC_KEY) */
-
-#if defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_ECC_KEY_PAIR) || \
-    defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_ECC_PUBLIC_KEY) || \
-    defined(MBEDTLS_PSA_BUILTIN_ALG_ECDSA) || \
-    defined(MBEDTLS_PSA_BUILTIN_ALG_ECDH) || \
-    defined(MBEDTLS_PSA_BUILTIN_ALG_DETERMINISTIC_ECDSA)
-/** Load the contents of a key buffer into an internal ECP representation
- *
- * \param[in] type          The type of key contained in \p data.
- * \param[in] data          The buffer from which to load the representation.
- * \param[in] data_length   The size in bytes of \p data.
- * \param[out] p_ecp        Returns a pointer to an ECP context on success.
- *                          The caller is responsible for freeing both the
- *                          contents of the context and the context itself
- *                          when done.
- */
-static psa_status_t psa_load_ecp_representation( psa_key_type_t type,
-                                                 const uint8_t *data,
-                                                 size_t data_length,
-                                                 mbedtls_ecp_keypair **p_ecp )
-{
-    mbedtls_ecp_group_id grp_id = MBEDTLS_ECP_DP_NONE;
-    psa_status_t status;
-    mbedtls_ecp_keypair *ecp = NULL;
-    size_t curve_size = data_length;
-
-    if( PSA_KEY_TYPE_IS_PUBLIC_KEY( type ) &&
-        PSA_KEY_TYPE_ECC_GET_FAMILY( type ) != PSA_ECC_FAMILY_MONTGOMERY )
-    {
-        /* A Weierstrass public key is represented as:
-         * - The byte 0x04;
-         * - `x_P` as a `ceiling(m/8)`-byte string, big-endian;
-         * - `y_P` as a `ceiling(m/8)`-byte string, big-endian.
-         * So its data length is 2m+1 where m is the curve size in bits.
-         */
-        if( ( data_length & 1 ) == 0 )
-            return( PSA_ERROR_INVALID_ARGUMENT );
-        curve_size = data_length / 2;
-
-        /* Montgomery public keys are represented in compressed format, meaning
-         * their curve_size is equal to the amount of input. */
-
-        /* Private keys are represented in uncompressed private random integer
-         * format, meaning their curve_size is equal to the amount of input. */
-    }
-
-    /* Allocate and initialize a key representation. */
-    ecp = mbedtls_calloc( 1, sizeof( mbedtls_ecp_keypair ) );
-    if( ecp == NULL )
-        return( PSA_ERROR_INSUFFICIENT_MEMORY );
-    mbedtls_ecp_keypair_init( ecp );
-
-    /* Load the group. */
-    grp_id = mbedtls_ecc_group_of_psa( PSA_KEY_TYPE_ECC_GET_FAMILY( type ),
-                                       curve_size );
-    if( grp_id == MBEDTLS_ECP_DP_NONE )
-    {
-        status = PSA_ERROR_INVALID_ARGUMENT;
-        goto exit;
-    }
-
-    status = mbedtls_to_psa_error(
-                mbedtls_ecp_group_load( &ecp->grp, grp_id ) );
-    if( status != PSA_SUCCESS )
-        goto exit;
-
-    /* Load the key material. */
-    if( PSA_KEY_TYPE_IS_PUBLIC_KEY( type ) )
-    {
-        /* Load the public value. */
-        status = mbedtls_to_psa_error(
-            mbedtls_ecp_point_read_binary( &ecp->grp, &ecp->Q,
-                                           data,
-                                           data_length ) );
-        if( status != PSA_SUCCESS )
-            goto exit;
-
-        /* Check that the point is on the curve. */
-        status = mbedtls_to_psa_error(
-            mbedtls_ecp_check_pubkey( &ecp->grp, &ecp->Q ) );
-        if( status != PSA_SUCCESS )
-            goto exit;
-    }
-    else
-    {
-        /* Load and validate the secret value. */
-        status = mbedtls_to_psa_error(
-            mbedtls_ecp_read_key( ecp->grp.id,
-                                  ecp,
-                                  data,
-                                  data_length ) );
-        if( status != PSA_SUCCESS )
-            goto exit;
-    }
-
-    *p_ecp = ecp;
-exit:
-    if( status != PSA_SUCCESS )
-    {
-        mbedtls_ecp_keypair_free( ecp );
-        mbedtls_free( ecp );
-    }
-
-    return( status );
-}
-#endif /* defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_ECC_KEY_PAIR) ||
-        * defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_ECC_PUBLIC_KEY) ||
-        * defined(MBEDTLS_PSA_BUILTIN_ALG_ECDSA) ||
-        * defined(MBEDTLS_PSA_BUILTIN_ALG_ECDH) ||
-        * defined(MBEDTLS_PSA_BUILTIN_ALG_DETERMINISTIC_ECDSA) */
-
-#if defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_ECC_KEY_PAIR) || \
-    defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_ECC_PUBLIC_KEY)
-/** Export an ECP key to export representation
- *
- * \param[in] type          The type of key (public/private) to export
- * \param[in] ecp           The internal ECP representation from which to export
- * \param[out] data         The buffer to export to
- * \param[in] data_size     The length of the buffer to export to
- * \param[out] data_length  The amount of bytes written to \p data
- */
-static psa_status_t psa_export_ecp_key( psa_key_type_t type,
-                                        mbedtls_ecp_keypair *ecp,
-                                        uint8_t *data,
-                                        size_t data_size,
-                                        size_t *data_length )
-{
-    psa_status_t status;
-
-    if( PSA_KEY_TYPE_IS_PUBLIC_KEY( type ) )
-    {
-        /* Check whether the public part is loaded */
-        if( mbedtls_ecp_is_zero( &ecp->Q ) )
-        {
-            /* Calculate the public key */
-            status = mbedtls_to_psa_error(
-                mbedtls_ecp_mul( &ecp->grp, &ecp->Q, &ecp->d, &ecp->grp.G,
-                                 mbedtls_psa_get_random, MBEDTLS_PSA_RANDOM_STATE ) );
-            if( status != PSA_SUCCESS )
-                return( status );
-        }
-
-        status = mbedtls_to_psa_error(
-                    mbedtls_ecp_point_write_binary( &ecp->grp, &ecp->Q,
-                                                    MBEDTLS_ECP_PF_UNCOMPRESSED,
-                                                    data_length,
-                                                    data,
-                                                    data_size ) );
-        if( status != PSA_SUCCESS )
-            memset( data, 0, data_size );
-
-        return( status );
-    }
-    else
-    {
-        if( data_size < PSA_BITS_TO_BYTES( ecp->grp.nbits ) )
-            return( PSA_ERROR_BUFFER_TOO_SMALL );
-
-        status = mbedtls_to_psa_error(
-                    mbedtls_ecp_write_key( ecp,
-                                           data,
-                                           PSA_BITS_TO_BYTES( ecp->grp.nbits ) ) );
-        if( status == PSA_SUCCESS )
-            *data_length = PSA_BITS_TO_BYTES( ecp->grp.nbits );
-        else
-            memset( data, 0, data_size );
-
-        return( status );
-    }
-}
-
-/** Import an ECP key from import representation to a slot
- *
- * \param[in,out] slot      The slot where to store the export representation to
- * \param[in] data          The buffer containing the import representation
- * \param[in] data_length   The amount of bytes in \p data
- */
-static psa_status_t psa_import_ecp_key( psa_key_slot_t *slot,
-                                        const uint8_t *data,
-                                        size_t data_length )
-{
-    psa_status_t status;
-    uint8_t* output = NULL;
-    mbedtls_ecp_keypair *ecp = NULL;
-
-    /* Parse input */
-    status = psa_load_ecp_representation( slot->attr.type,
-                                          data,
-                                          data_length,
-                                          &ecp );
-    if( status != PSA_SUCCESS )
-        goto exit;
-
-    if( PSA_KEY_TYPE_ECC_GET_FAMILY( slot->attr.type ) == PSA_ECC_FAMILY_MONTGOMERY)
-        slot->attr.bits = (psa_key_bits_t) ecp->grp.nbits + 1;
-    else
-        slot->attr.bits = (psa_key_bits_t) ecp->grp.nbits;
-
-    /* Re-export the data to PSA export format. There is currently no support
-     * for other input formats then the export format, so this is a 1-1
-     * copy operation. */
-    output = mbedtls_calloc( 1, data_length );
-    if( output == NULL )
-    {
-        status = PSA_ERROR_INSUFFICIENT_MEMORY;
-        goto exit;
-    }
-
-    status = psa_export_ecp_key( slot->attr.type,
-                                 ecp,
-                                 output,
-                                 data_length,
-                                 &data_length);
-exit:
-    /* Always free the PK object (will also free contained ECP context) */
-    mbedtls_ecp_keypair_free( ecp );
-    mbedtls_free( ecp );
-
-    /* Free the allocated buffer only on error. */
-    if( status != PSA_SUCCESS )
-    {
-        mbedtls_free( output );
-        return( status );
-    }
-
-    /* On success, store the allocated export-formatted key. */
-    slot->data.key.data = output;
-    slot->data.key.bytes = data_length;
-
-    return( PSA_SUCCESS );
-}
-#endif /* defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_ECC_KEY_PAIR) ||
-        * defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_ECC_PUBLIC_KEY) */
-
 /** Return the size of the key in the given slot, in bits.
  *
  * \param[in] slot      A key slot.
@@ -1017,14 +559,14 @@
 static psa_status_t psa_allocate_buffer_to_slot( psa_key_slot_t *slot,
                                                  size_t buffer_length )
 {
-    if( slot->data.key.data != NULL )
+    if( slot->key.data != NULL )
         return( PSA_ERROR_ALREADY_EXISTS );
 
-    slot->data.key.data = mbedtls_calloc( 1, buffer_length );
-    if( slot->data.key.data == NULL )
+    slot->key.data = mbedtls_calloc( 1, buffer_length );
+    if( slot->key.data == NULL )
         return( PSA_ERROR_INSUFFICIENT_MEMORY );
 
-    slot->data.key.bytes = buffer_length;
+    slot->key.bytes = buffer_length;
     return( PSA_SUCCESS );
 }
 
@@ -1037,44 +579,26 @@
     if( status != PSA_SUCCESS )
         return( status );
 
-    memcpy( slot->data.key.data, data, data_length );
+    memcpy( slot->key.data, data, data_length );
     return( PSA_SUCCESS );
 }
 
-/** Import key data into a slot.
- *
- * `slot->type` must have been set previously.
- * This function assumes that the slot does not contain any key material yet.
- * On failure, the slot content is unchanged.
- *
- * Persistent storage is not affected.
- *
- * \param[in,out] slot  The key slot to import data into.
- *                      Its `type` field must have previously been set to
- *                      the desired key type.
- *                      It must not contain any key material yet.
- * \param[in] data      Buffer containing the key material to parse and import.
- * \param data_length   Size of \p data in bytes.
- *
- * \retval #PSA_SUCCESS
- * \retval #PSA_ERROR_INVALID_ARGUMENT
- * \retval #PSA_ERROR_NOT_SUPPORTED
- * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
- */
-static psa_status_t psa_import_key_into_slot( psa_key_slot_t *slot,
-                                              const uint8_t *data,
-                                              size_t data_length )
+psa_status_t psa_import_key_into_slot(
+    const psa_key_attributes_t *attributes,
+    const uint8_t *data, size_t data_length,
+    uint8_t *key_buffer, size_t key_buffer_size,
+    size_t *key_buffer_length, size_t *bits )
 {
-    psa_status_t status = PSA_SUCCESS;
-    size_t bit_size;
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+    psa_key_type_t type = attributes->core.type;
 
     /* zero-length keys are never supported. */
     if( data_length == 0 )
         return( PSA_ERROR_NOT_SUPPORTED );
 
-    if( key_type_is_raw_bytes( slot->attr.type ) )
+    if( key_type_is_raw_bytes( type ) )
     {
-        bit_size = PSA_BYTES_TO_BITS( data_length );
+        *bits = PSA_BYTES_TO_BITS( data_length );
 
         /* Ensure that the bytes-to-bits conversion hasn't overflown. */
         if( data_length > SIZE_MAX / 8 )
@@ -1082,77 +606,49 @@
 
         /* Enforce a size limit, and in particular ensure that the bit
          * size fits in its representation type. */
-        if( bit_size > PSA_MAX_KEY_BITS )
+        if( ( *bits ) > PSA_MAX_KEY_BITS )
             return( PSA_ERROR_NOT_SUPPORTED );
 
-        status = validate_unstructured_key_bit_size( slot->attr.type, bit_size );
+        status = validate_unstructured_key_bit_size( type, *bits );
         if( status != PSA_SUCCESS )
             return( status );
 
-        /* Allocate memory for the key */
-        status = psa_copy_key_material_into_slot( slot, data, data_length );
-        if( status != PSA_SUCCESS )
-            return( status );
-
-        /* Write the actual key size to the slot.
-         * psa_start_key_creation() wrote the size declared by the
-         * caller, which may be 0 (meaning unspecified) or wrong. */
-        slot->attr.bits = (psa_key_bits_t) bit_size;
+        /* Copy the key material. */
+        memcpy( key_buffer, data, data_length );
+        *key_buffer_length = data_length;
+        (void)key_buffer_size;
 
         return( PSA_SUCCESS );
     }
-    else if( PSA_KEY_TYPE_IS_ASYMMETRIC( slot->attr.type ) )
+    else if( PSA_KEY_TYPE_IS_ASYMMETRIC( type ) )
     {
-        /* Try validation through accelerators first. */
-        bit_size = slot->attr.bits;
-        psa_key_attributes_t attributes = {
-          .core = slot->attr
-        };
-        status = psa_driver_wrapper_validate_key( &attributes,
-                                                  data,
-                                                  data_length,
-                                                  &bit_size );
-        if( status == PSA_SUCCESS )
-        {
-            /* Key has been validated successfully by an accelerator.
-             * Copy key material into slot. */
-            status = psa_copy_key_material_into_slot( slot, data, data_length );
-            if( status != PSA_SUCCESS )
-                return( status );
-
-            slot->attr.bits = (psa_key_bits_t) bit_size;
-            return( PSA_SUCCESS );
-        }
-        else if( status != PSA_ERROR_NOT_SUPPORTED )
-            return( status );
-
-        /* Key format is not supported by any accelerator, try software fallback
-         * if present. */
 #if defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_ECC_KEY_PAIR) || \
     defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_ECC_PUBLIC_KEY)
-        if( PSA_KEY_TYPE_IS_ECC( slot->attr.type ) )
+        if( PSA_KEY_TYPE_IS_ECC( type ) )
         {
-            return( psa_import_ecp_key( slot, data, data_length ) );
+            return( mbedtls_psa_ecp_import_key( attributes,
+                                                data, data_length,
+                                                key_buffer, key_buffer_size,
+                                                key_buffer_length,
+                                                bits ) );
         }
 #endif /* defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_ECC_KEY_PAIR) ||
         * defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_ECC_PUBLIC_KEY) */
 #if defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_RSA_KEY_PAIR) || \
     defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_RSA_PUBLIC_KEY)
-        if( PSA_KEY_TYPE_IS_RSA( slot->attr.type ) )
+        if( PSA_KEY_TYPE_IS_RSA( type ) )
         {
-            return( psa_import_rsa_key( slot, data, data_length ) );
+            return( mbedtls_psa_rsa_import_key( attributes,
+                                                data, data_length,
+                                                key_buffer, key_buffer_size,
+                                                key_buffer_length,
+                                                bits ) );
         }
 #endif /* defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_RSA_KEY_PAIR) ||
         * defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_RSA_PUBLIC_KEY) */
+    }
 
-        /* Fell through the fallback as well, so have nothing else to try. */
-        return( PSA_ERROR_NOT_SUPPORTED );
-    }
-    else
-    {
-        /* Unknown key type */
-        return( PSA_ERROR_NOT_SUPPORTED );
-    }
+    return( PSA_ERROR_NOT_SUPPORTED );
 }
 
 /** Calculate the intersection of two algorithm usage policies.
@@ -1342,23 +838,14 @@
 /** Wipe key data from a slot. Preserve metadata such as the policy. */
 static psa_status_t psa_remove_key_data_from_memory( psa_key_slot_t *slot )
 {
-#if defined(MBEDTLS_PSA_CRYPTO_SE_C)
-    if( psa_get_se_driver( slot->attr.lifetime, NULL, NULL ) &&
-        psa_key_slot_is_external( slot ) )
-    {
-        /* No key material to clean. */
-    }
-    else
-#endif /* MBEDTLS_PSA_CRYPTO_SE_C */
-    {
-        /* Data pointer will always be either a valid pointer or NULL in an
-         * initialized slot, so we can just free it. */
-        if( slot->data.key.data != NULL )
-            mbedtls_platform_zeroize( slot->data.key.data, slot->data.key.bytes);
-        mbedtls_free( slot->data.key.data );
-        slot->data.key.data = NULL;
-        slot->data.key.bytes = 0;
-    }
+    /* Data pointer will always be either a valid pointer or NULL in an
+     * initialized slot, so we can just free it. */
+    if( slot->key.data != NULL )
+        mbedtls_platform_zeroize( slot->key.data, slot->key.bytes);
+
+    mbedtls_free( slot->key.data );
+    slot->key.data = NULL;
+    slot->key.bytes = 0;
 
     return( PSA_SUCCESS );
 }
@@ -1443,7 +930,7 @@
          * three actions. */
         psa_crypto_prepare_transaction( PSA_CRYPTO_TRANSACTION_DESTROY_KEY );
         psa_crypto_transaction.key.lifetime = slot->attr.lifetime;
-        psa_crypto_transaction.key.slot = slot->data.se.slot_number;
+        psa_crypto_transaction.key.slot = psa_key_slot_get_slot_number( slot );
         psa_crypto_transaction.key.id = slot->attr.id;
         status = psa_crypto_save_transaction( );
         if( status != PSA_SUCCESS )
@@ -1460,7 +947,8 @@
             goto exit;
         }
 
-        status = psa_destroy_se_key( driver, slot->data.se.slot_number );
+        status = psa_destroy_se_key( driver,
+                                     psa_key_slot_get_slot_number( slot ) );
         if( overall_status == PSA_SUCCESS )
             overall_status = status;
     }
@@ -1616,7 +1104,8 @@
 
 #if defined(MBEDTLS_PSA_CRYPTO_SE_C)
     if( psa_key_slot_is_external( slot ) )
-        psa_set_key_slot_number( attributes, slot->data.se.slot_number );
+        psa_set_key_slot_number( attributes,
+                                 psa_key_slot_get_slot_number( slot ) );
 #endif /* MBEDTLS_PSA_CRYPTO_SE_C */
 
     switch( slot->attr.type )
@@ -1636,10 +1125,11 @@
             {
                 mbedtls_rsa_context *rsa = NULL;
 
-                status = psa_load_rsa_representation( slot->attr.type,
-                                                      slot->data.key.data,
-                                                      slot->data.key.bytes,
-                                                      &rsa );
+                status = mbedtls_psa_rsa_load_representation(
+                             slot->attr.type,
+                             slot->key.data,
+                             slot->key.bytes,
+                             &rsa );
                 if( status != PSA_SUCCESS )
                     break;
 
@@ -1679,147 +1169,35 @@
 }
 #endif /* MBEDTLS_PSA_CRYPTO_SE_C */
 
-static psa_status_t psa_internal_export_key_buffer( const psa_key_slot_t *slot,
+static psa_status_t psa_export_key_buffer_internal( const uint8_t *key_buffer,
+                                                    size_t key_buffer_size,
                                                     uint8_t *data,
                                                     size_t data_size,
                                                     size_t *data_length )
 {
-    if( slot->data.key.bytes > data_size )
+    if( key_buffer_size > data_size )
         return( PSA_ERROR_BUFFER_TOO_SMALL );
-    memcpy( data, slot->data.key.data, slot->data.key.bytes );
-    memset( data + slot->data.key.bytes, 0,
-            data_size - slot->data.key.bytes );
-    *data_length = slot->data.key.bytes;
+    memcpy( data, key_buffer, key_buffer_size );
+    memset( data + key_buffer_size, 0,
+            data_size - key_buffer_size );
+    *data_length = key_buffer_size;
     return( PSA_SUCCESS );
 }
 
-static psa_status_t psa_internal_export_key( const psa_key_slot_t *slot,
-                                             uint8_t *data,
-                                             size_t data_size,
-                                             size_t *data_length,
-                                             int export_public_key )
+psa_status_t psa_export_key_internal(
+    const psa_key_attributes_t *attributes,
+    const uint8_t *key_buffer, size_t key_buffer_size,
+    uint8_t *data, size_t data_size, size_t *data_length )
 {
-#if defined(MBEDTLS_PSA_CRYPTO_SE_C)
-    const psa_drv_se_t *drv;
-    psa_drv_se_context_t *drv_context;
-#endif /* MBEDTLS_PSA_CRYPTO_SE_C */
+    psa_key_type_t type = attributes->core.type;
 
-    *data_length = 0;
-
-    if( export_public_key && ! PSA_KEY_TYPE_IS_ASYMMETRIC( slot->attr.type ) )
-        return( PSA_ERROR_INVALID_ARGUMENT );
-
-    /* Reject a zero-length output buffer now, since this can never be a
-     * valid key representation. This way we know that data must be a valid
-     * pointer and we can do things like memset(data, ..., data_size). */
-    if( data_size == 0 )
-        return( PSA_ERROR_BUFFER_TOO_SMALL );
-
-#if defined(MBEDTLS_PSA_CRYPTO_SE_C)
-    if( psa_get_se_driver( slot->attr.lifetime, &drv, &drv_context ) )
+    if( key_type_is_raw_bytes( type ) ||
+        PSA_KEY_TYPE_IS_RSA( type )   ||
+        PSA_KEY_TYPE_IS_ECC( type )      )
     {
-        psa_drv_se_export_key_t method;
-        if( drv->key_management == NULL )
-            return( PSA_ERROR_NOT_SUPPORTED );
-        method = ( export_public_key ?
-                   drv->key_management->p_export_public :
-                   drv->key_management->p_export );
-        if( method == NULL )
-            return( PSA_ERROR_NOT_SUPPORTED );
-        return( method( drv_context,
-                        slot->data.se.slot_number,
-                        data, data_size, data_length ) );
-    }
-#endif /* MBEDTLS_PSA_CRYPTO_SE_C */
-
-    if( key_type_is_raw_bytes( slot->attr.type ) )
-    {
-        return( psa_internal_export_key_buffer( slot, data, data_size, data_length ) );
-    }
-    else if( PSA_KEY_TYPE_IS_RSA( slot->attr.type ) ||
-             PSA_KEY_TYPE_IS_ECC( slot->attr.type ) )
-    {
-        if( PSA_KEY_TYPE_IS_PUBLIC_KEY( slot->attr.type ) )
-        {
-            /* Exporting public -> public */
-            return( psa_internal_export_key_buffer( slot, data, data_size, data_length ) );
-        }
-        else if( !export_public_key )
-        {
-            /* Exporting private -> private */
-            return( psa_internal_export_key_buffer( slot, data, data_size, data_length ) );
-        }
-
-        /* Need to export the public part of a private key,
-         * so conversion is needed. Try the accelerators first. */
-        psa_status_t status = psa_driver_wrapper_export_public_key( slot,
-                                                                    data,
-                                                                    data_size,
-                                                                    data_length );
-
-        if( status != PSA_ERROR_NOT_SUPPORTED ||
-            psa_key_lifetime_is_external( slot->attr.lifetime ) )
-            return( status );
-
-        if( PSA_KEY_TYPE_IS_RSA( slot->attr.type ) )
-        {
-#if defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_RSA_KEY_PAIR) || \
-    defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_RSA_PUBLIC_KEY)
-            mbedtls_rsa_context *rsa = NULL;
-            status = psa_load_rsa_representation(
-                                    slot->attr.type,
-                                    slot->data.key.data,
-                                    slot->data.key.bytes,
-                                    &rsa );
-            if( status != PSA_SUCCESS )
-                return( status );
-
-            status = psa_export_rsa_key( PSA_KEY_TYPE_RSA_PUBLIC_KEY,
-                                         rsa,
-                                         data,
-                                         data_size,
-                                         data_length );
-
-            mbedtls_rsa_free( rsa );
-            mbedtls_free( rsa );
-
-            return( status );
-#else
-            /* We don't know how to convert a private RSA key to public. */
-            return( PSA_ERROR_NOT_SUPPORTED );
-#endif /* defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_RSA_KEY_PAIR) ||
-        * defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_RSA_PUBLIC_KEY) */
-        }
-        else
-        {
-#if defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_ECC_KEY_PAIR) || \
-    defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_ECC_PUBLIC_KEY)
-            mbedtls_ecp_keypair *ecp = NULL;
-            status = psa_load_ecp_representation(
-                                    slot->attr.type,
-                                    slot->data.key.data,
-                                    slot->data.key.bytes,
-                                    &ecp );
-            if( status != PSA_SUCCESS )
-                return( status );
-
-            status = psa_export_ecp_key( PSA_KEY_TYPE_ECC_PUBLIC_KEY(
-                                            PSA_KEY_TYPE_ECC_GET_FAMILY(
-                                                slot->attr.type ) ),
-                                         ecp,
-                                         data,
-                                         data_size,
-                                         data_length );
-
-            mbedtls_ecp_keypair_free( ecp );
-            mbedtls_free( ecp );
-            return( status );
-#else
-            /* We don't know how to convert a private ECC key to public */
-            return( PSA_ERROR_NOT_SUPPORTED );
-#endif /* defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_ECC_KEY_PAIR) ||
-        * defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_ECC_PUBLIC_KEY) */
-        }
+        return( psa_export_key_buffer_internal(
+                    key_buffer, key_buffer_size,
+                    data, data_size, data_length ) );
     }
     else
     {
@@ -1839,6 +1217,12 @@
     psa_status_t unlock_status = PSA_ERROR_CORRUPTION_DETECTED;
     psa_key_slot_t *slot;
 
+    /* Reject a zero-length output buffer now, since this can never be a
+     * valid key representation. This way we know that data must be a valid
+     * pointer and we can do things like memset(data, ..., data_size). */
+    if( data_size == 0 )
+        return( PSA_ERROR_BUFFER_TOO_SMALL );
+
     /* Set the key to empty now, so that even when there are errors, we always
      * set data_length to a value between 0 and data_size. On error, setting
      * the key to empty is a good choice because an empty key representation is
@@ -1854,12 +1238,80 @@
     if( status != PSA_SUCCESS )
         return( status );
 
-    status = psa_internal_export_key( slot, data, data_size, data_length, 0 );
+    psa_key_attributes_t attributes = {
+        .core = slot->attr
+    };
+    status = psa_driver_wrapper_export_key( &attributes,
+                 slot->key.data, slot->key.bytes,
+                 data, data_size, data_length );
+
     unlock_status = psa_unlock_key_slot( slot );
 
     return( ( status == PSA_SUCCESS ) ? unlock_status : status );
 }
 
+psa_status_t psa_export_public_key_internal(
+    const psa_key_attributes_t *attributes,
+    const uint8_t *key_buffer,
+    size_t key_buffer_size,
+    uint8_t *data,
+    size_t data_size,
+    size_t *data_length )
+{
+    psa_key_type_t type = attributes->core.type;
+
+    if( PSA_KEY_TYPE_IS_RSA( type ) || PSA_KEY_TYPE_IS_ECC( type ) )
+    {
+        if( PSA_KEY_TYPE_IS_PUBLIC_KEY( type ) )
+        {
+            /* Exporting public -> public */
+            return( psa_export_key_buffer_internal(
+                        key_buffer, key_buffer_size,
+                        data, data_size, data_length ) );
+        }
+
+        if( PSA_KEY_TYPE_IS_RSA( type ) )
+        {
+#if defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_RSA_KEY_PAIR) || \
+    defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_RSA_PUBLIC_KEY)
+            return( mbedtls_psa_rsa_export_public_key( attributes,
+                                                       key_buffer,
+                                                       key_buffer_size,
+                                                       data,
+                                                       data_size,
+                                                       data_length ) );
+#else
+            /* We don't know how to convert a private RSA key to public. */
+            return( PSA_ERROR_NOT_SUPPORTED );
+#endif /* defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_RSA_KEY_PAIR) ||
+        * defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_RSA_PUBLIC_KEY) */
+        }
+        else
+        {
+#if defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_ECC_KEY_PAIR) || \
+    defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_ECC_PUBLIC_KEY)
+            return( mbedtls_psa_ecp_export_public_key( attributes,
+                                                       key_buffer,
+                                                       key_buffer_size,
+                                                       data,
+                                                       data_size,
+                                                       data_length ) );
+#else
+            /* We don't know how to convert a private ECC key to public */
+            return( PSA_ERROR_NOT_SUPPORTED );
+#endif /* defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_ECC_KEY_PAIR) ||
+        * defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_ECC_PUBLIC_KEY) */
+        }
+    }
+    else
+    {
+        /* This shouldn't happen in the reference implementation, but
+           it is valid for a special-purpose implementation to omit
+           support for exporting certain key types. */
+        return( PSA_ERROR_NOT_SUPPORTED );
+    }
+}
+
 psa_status_t psa_export_public_key( mbedtls_svc_key_id_t key,
                                     uint8_t *data,
                                     size_t data_size,
@@ -1869,6 +1321,12 @@
     psa_status_t unlock_status = PSA_ERROR_CORRUPTION_DETECTED;
     psa_key_slot_t *slot;
 
+    /* Reject a zero-length output buffer now, since this can never be a
+     * valid key representation. This way we know that data must be a valid
+     * pointer and we can do things like memset(data, ..., data_size). */
+    if( data_size == 0 )
+        return( PSA_ERROR_BUFFER_TOO_SMALL );
+
     /* Set the key to empty now, so that even when there are errors, we always
      * set data_length to a value between 0 and data_size. On error, setting
      * the key to empty is a good choice because an empty key representation is
@@ -1880,7 +1338,20 @@
     if( status != PSA_SUCCESS )
         return( status );
 
-    status = psa_internal_export_key( slot, data, data_size, data_length, 1 );
+    if( ! PSA_KEY_TYPE_IS_ASYMMETRIC( slot->attr.type ) )
+    {
+         status = PSA_ERROR_INVALID_ARGUMENT;
+         goto exit;
+    }
+
+    psa_key_attributes_t attributes = {
+        .core = slot->attr
+    };
+    status = psa_driver_wrapper_export_public_key(
+        &attributes, slot->key.data, slot->key.bytes,
+        data, data_size, data_length );
+
+exit:
     unlock_status = psa_unlock_key_slot( slot );
 
     return( ( status == PSA_SUCCESS ) ? unlock_status : status );
@@ -2071,8 +1542,9 @@
      * we can roll back to a state where the key doesn't exist. */
     if( *p_drv != NULL )
     {
+        psa_key_slot_number_t slot_number;
         status = psa_find_se_slot_for_key( attributes, method, *p_drv,
-                                           &slot->data.se.slot_number );
+                                           &slot_number );
         if( status != PSA_SUCCESS )
             return( status );
 
@@ -2080,7 +1552,7 @@
         {
             psa_crypto_prepare_transaction( PSA_CRYPTO_TRANSACTION_CREATE_KEY );
             psa_crypto_transaction.key.lifetime = slot->attr.lifetime;
-            psa_crypto_transaction.key.slot = slot->data.se.slot_number;
+            psa_crypto_transaction.key.slot = slot_number;
             psa_crypto_transaction.key.id = slot->attr.id;
             status = psa_crypto_save_transaction( );
             if( status != PSA_SUCCESS )
@@ -2089,6 +1561,9 @@
                 return( status );
             }
         }
+
+        status = psa_copy_key_material_into_slot(
+            slot, (uint8_t *)( &slot_number ), sizeof( slot_number ) );
     }
 
     if( *p_drv == NULL && method == PSA_KEY_CREATION_REGISTER )
@@ -2121,6 +1596,13 @@
  *
  * \retval #PSA_SUCCESS
  *         The key was successfully created.
+ * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
+ * \retval #PSA_ERROR_INSUFFICIENT_STORAGE
+ * \retval #PSA_ERROR_ALREADY_EXISTS
+ * \retval #PSA_ERROR_DATA_INVALID
+ * \retval #PSA_ERROR_DATA_CORRUPT
+ * \retval #PSA_ERROR_STORAGE_FAILURE
+ *
  * \return If this function fails, the key slot is an invalid state.
  *         You must call psa_fail_key_creation() to wipe and free the slot.
  */
@@ -2140,13 +1622,15 @@
         if( driver != NULL )
         {
             psa_se_key_data_storage_t data;
+            psa_key_slot_number_t slot_number =
+                psa_key_slot_get_slot_number( slot ) ;
+
 #if defined(static_assert)
-            static_assert( sizeof( slot->data.se.slot_number ) ==
+            static_assert( sizeof( slot_number ) ==
                            sizeof( data.slot_number ),
                            "Slot number size does not match psa_se_key_data_storage_t" );
 #endif
-            memcpy( &data.slot_number, &slot->data.se.slot_number,
-                    sizeof( slot->data.se.slot_number ) );
+            memcpy( &data.slot_number, &slot_number, sizeof( slot_number ) );
             status = psa_save_persistent_key( &slot->attr,
                                               (uint8_t*) &data,
                                               sizeof( data ) );
@@ -2157,8 +1641,8 @@
             /* Key material is saved in export representation in the slot, so
              * just pass the slot buffer for storage. */
             status = psa_save_persistent_key( &slot->attr,
-                                              slot->data.key.data,
-                                              slot->data.key.bytes );
+                                              slot->key.data,
+                                              slot->key.bytes );
         }
     }
 #endif /* defined(MBEDTLS_PSA_CRYPTO_STORAGE_C) */
@@ -2262,11 +1746,11 @@
             mbedtls_mpi actual, required;
             int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
 
-            psa_status_t status = psa_load_rsa_representation(
-                                    slot->attr.type,
-                                    slot->data.key.data,
-                                    slot->data.key.bytes,
-                                    &rsa );
+            psa_status_t status = mbedtls_psa_rsa_load_representation(
+                                      slot->attr.type,
+                                      slot->key.data,
+                                      slot->key.bytes,
+                                      &rsa );
             if( status != PSA_SUCCESS )
                 return( status );
 
@@ -2316,6 +1800,7 @@
     psa_status_t status;
     psa_key_slot_t *slot = NULL;
     psa_se_drv_table_entry_t *driver = NULL;
+    size_t bits;
 
     *key = MBEDTLS_SVC_KEY_ID_INIT;
 
@@ -2330,48 +1815,34 @@
     if( status != PSA_SUCCESS )
         goto exit;
 
-#if defined(MBEDTLS_PSA_CRYPTO_SE_C)
-    if( driver != NULL )
+    /* In the case of a transparent key or an opaque key stored in local
+     * storage (thus not in the case of generating a key in a secure element
+     * or cryptoprocessor with storage), we have to allocate a buffer to
+     * hold the generated key material. */
+    if( slot->key.data == NULL )
     {
-        const psa_drv_se_t *drv = psa_get_se_driver_methods( driver );
-        /* The driver should set the number of key bits, however in
-         * case it doesn't, we initialize bits to an invalid value. */
-        size_t bits = PSA_MAX_KEY_BITS + 1;
-        if( drv->key_management == NULL ||
-            drv->key_management->p_import == NULL )
-        {
-            status = PSA_ERROR_NOT_SUPPORTED;
-            goto exit;
-        }
-        status = drv->key_management->p_import(
-            psa_get_se_driver_context( driver ),
-            slot->data.se.slot_number, attributes, data, data_length,
-            &bits );
+        status = psa_allocate_buffer_to_slot( slot, data_length );
         if( status != PSA_SUCCESS )
             goto exit;
-        if( bits > PSA_MAX_KEY_BITS )
-        {
-            status = PSA_ERROR_NOT_SUPPORTED;
-            goto exit;
-        }
-        slot->attr.bits = (psa_key_bits_t) bits;
     }
-    else
-#endif /* MBEDTLS_PSA_CRYPTO_SE_C */
-    if( psa_key_lifetime_is_external( psa_get_key_lifetime( attributes ) ) )
+
+    bits = slot->attr.bits;
+    status = psa_driver_wrapper_import_key( attributes,
+                                            data, data_length,
+                                            slot->key.data,
+                                            slot->key.bytes,
+                                            &slot->key.bytes, &bits );
+    if( status != PSA_SUCCESS )
+        goto exit;
+
+    if( slot->attr.bits == 0 )
+        slot->attr.bits = (psa_key_bits_t) bits;
+    else if( bits != slot->attr.bits )
     {
-        /* Importing a key with external lifetime through the driver wrapper
-         * interface is not yet supported. Return as if this was an invalid
-         * lifetime. */
         status = PSA_ERROR_INVALID_ARGUMENT;
         goto exit;
     }
-    else
-    {
-        status = psa_import_key_into_slot( slot, data, data_length );
-        if( status != PSA_SUCCESS )
-            goto exit;
-    }
+
     status = psa_validate_optional_attributes( slot, attributes );
     if( status != PSA_SUCCESS )
         goto exit;
@@ -2423,8 +1894,8 @@
                                            psa_key_slot_t *target )
 {
     psa_status_t status = psa_copy_key_material_into_slot( target,
-                                                           source->data.key.data,
-                                                           source->data.key.bytes );
+                                                           source->key.data,
+                                                           source->key.bytes );
     if( status != PSA_SUCCESS )
         return( status );
 
@@ -3236,7 +2707,7 @@
         return( ret );
 
     ret = mbedtls_cipher_cmac_starts( &operation->ctx.cmac,
-                                      slot->data.key.data,
+                                      slot->key.data,
                                       key_bits );
     return( ret );
 }
@@ -3382,8 +2853,8 @@
         }
 
         status = psa_hmac_setup_internal( &operation->ctx.hmac,
-                                          slot->data.key.data,
-                                          slot->data.key.bytes,
+                                          slot->key.data,
+                                          slot->key.bytes,
                                           hash_alg );
     }
     else
@@ -3968,10 +3439,10 @@
     {
         mbedtls_rsa_context *rsa = NULL;
 
-        status = psa_load_rsa_representation( slot->attr.type,
-                                              slot->data.key.data,
-                                              slot->data.key.bytes,
-                                              &rsa );
+        status = mbedtls_psa_rsa_load_representation( slot->attr.type,
+                                                      slot->key.data,
+                                                      slot->key.bytes,
+                                                      &rsa );
         if( status != PSA_SUCCESS )
             goto exit;
 
@@ -4000,10 +3471,10 @@
             )
         {
             mbedtls_ecp_keypair *ecp = NULL;
-            status = psa_load_ecp_representation( slot->attr.type,
-                                                  slot->data.key.data,
-                                                  slot->data.key.bytes,
-                                                  &ecp );
+            status = mbedtls_psa_ecp_load_representation( slot->attr.type,
+                                                          slot->key.data,
+                                                          slot->key.bytes,
+                                                          &ecp );
             if( status != PSA_SUCCESS )
                 goto exit;
             status = psa_ecdsa_sign( ecp,
@@ -4078,10 +3549,10 @@
     {
         mbedtls_rsa_context *rsa = NULL;
 
-        status = psa_load_rsa_representation( slot->attr.type,
-                                              slot->data.key.data,
-                                              slot->data.key.bytes,
-                                              &rsa );
+        status = mbedtls_psa_rsa_load_representation( slot->attr.type,
+                                                      slot->key.data,
+                                                      slot->key.bytes,
+                                                      &rsa );
         if( status != PSA_SUCCESS )
             goto exit;
 
@@ -4103,10 +3574,10 @@
         if( PSA_ALG_IS_ECDSA( alg ) )
         {
             mbedtls_ecp_keypair *ecp = NULL;
-            status = psa_load_ecp_representation( slot->attr.type,
-                                                  slot->data.key.data,
-                                                  slot->data.key.bytes,
-                                                  &ecp );
+            status = mbedtls_psa_ecp_load_representation( slot->attr.type,
+                                                          slot->key.data,
+                                                          slot->key.bytes,
+                                                          &ecp );
             if( status != PSA_SUCCESS )
                 goto exit;
             status = psa_ecdsa_verify( ecp,
@@ -4187,10 +3658,10 @@
     if( PSA_KEY_TYPE_IS_RSA( slot->attr.type ) )
     {
         mbedtls_rsa_context *rsa = NULL;
-        status = psa_load_rsa_representation( slot->attr.type,
-                                              slot->data.key.data,
-                                              slot->data.key.bytes,
-                                              &rsa );
+        status = mbedtls_psa_rsa_load_representation( slot->attr.type,
+                                                      slot->key.data,
+                                                      slot->key.bytes,
+                                                      &rsa );
         if( status != PSA_SUCCESS )
             goto rsa_exit;
 
@@ -4293,10 +3764,10 @@
     if( slot->attr.type == PSA_KEY_TYPE_RSA_KEY_PAIR )
     {
         mbedtls_rsa_context *rsa = NULL;
-        status = psa_load_rsa_representation( slot->attr.type,
-                                              slot->data.key.data,
-                                              slot->data.key.bytes,
-                                              &rsa );
+        status = mbedtls_psa_rsa_load_representation( slot->attr.type,
+                                                      slot->key.data,
+                                                      slot->key.bytes,
+                                                      &rsa );
         if( status != PSA_SUCCESS )
             goto exit;
 
@@ -4455,8 +3926,8 @@
     {
         /* Two-key Triple-DES is 3-key Triple-DES with K1=K3 */
         uint8_t keys[24];
-        memcpy( keys, slot->data.key.data, 16 );
-        memcpy( keys + 16, slot->data.key.data, 8 );
+        memcpy( keys, slot->key.data, 16 );
+        memcpy( keys + 16, slot->key.data, 8 );
         ret = mbedtls_cipher_setkey( &operation->ctx.cipher,
                                      keys,
                                      192, cipher_operation );
@@ -4465,7 +3936,7 @@
 #endif
     {
         ret = mbedtls_cipher_setkey( &operation->ctx.cipher,
-                                     slot->data.key.data,
+                                     slot->key.data,
                                      (int) key_bits, cipher_operation );
     }
     if( ret != 0 )
@@ -4964,7 +4435,7 @@
             mbedtls_ccm_init( &operation->ctx.ccm );
             status = mbedtls_to_psa_error(
                 mbedtls_ccm_setkey( &operation->ctx.ccm, cipher_id,
-                                    operation->slot->data.key.data,
+                                    operation->slot->key.data,
                                     (unsigned int) key_bits ) );
             if( status != 0 )
                 goto cleanup;
@@ -4986,7 +4457,7 @@
             mbedtls_gcm_init( &operation->ctx.gcm );
             status = mbedtls_to_psa_error(
                 mbedtls_gcm_setkey( &operation->ctx.gcm, cipher_id,
-                                    operation->slot->data.key.data,
+                                    operation->slot->key.data,
                                     (unsigned int) key_bits ) );
             if( status != 0 )
                 goto cleanup;
@@ -5006,7 +4477,7 @@
             mbedtls_chachapoly_init( &operation->ctx.chachapoly );
             status = mbedtls_to_psa_error(
                 mbedtls_chachapoly_setkey( &operation->ctx.chachapoly,
-                                           operation->slot->data.key.data ) );
+                                           operation->slot->key.data ) );
             if( status != 0 )
                 goto cleanup;
             break;
@@ -5671,7 +5142,23 @@
     if( slot->attr.type == PSA_KEY_TYPE_DES )
         psa_des_set_key_parity( data, bytes );
 #endif /* MBEDTLS_DES_C */
-    status = psa_import_key_into_slot( slot, data, bytes );
+
+    status = psa_allocate_buffer_to_slot( slot, bytes );
+    if( status != PSA_SUCCESS )
+        return( status );
+
+    slot->attr.bits = (psa_key_bits_t) bits;
+    psa_key_attributes_t attributes = {
+      .core = slot->attr
+    };
+
+    status = psa_driver_wrapper_import_key( &attributes,
+                                            data, bytes,
+                                            slot->key.data,
+                                            slot->key.bytes,
+                                            &slot->key.bytes, &bits );
+    if( bits != slot->attr.bits )
+        status = PSA_ERROR_INVALID_ARGUMENT;
 
 exit:
     mbedtls_free( data );
@@ -5876,7 +5363,7 @@
                                             const uint8_t *data,
                                             size_t data_length )
 {
-    if( prf->state != TLS12_PRF_STATE_INIT )
+    if( prf->state != PSA_TLS12_PRF_STATE_INIT )
         return( PSA_ERROR_BAD_STATE );
 
     if( data_length != 0 )
@@ -5889,7 +5376,7 @@
         prf->seed_length = data_length;
     }
 
-    prf->state = TLS12_PRF_STATE_SEED_SET;
+    prf->state = PSA_TLS12_PRF_STATE_SEED_SET;
 
     return( PSA_SUCCESS );
 }
@@ -5900,14 +5387,14 @@
                                            size_t data_length )
 {
     psa_status_t status;
-    if( prf->state != TLS12_PRF_STATE_SEED_SET )
+    if( prf->state != PSA_TLS12_PRF_STATE_SEED_SET )
         return( PSA_ERROR_BAD_STATE );
 
     status = psa_hmac_setup_internal( &prf->hmac, data, data_length, hash_alg );
     if( status != PSA_SUCCESS )
         return( status );
 
-    prf->state = TLS12_PRF_STATE_KEY_SET;
+    prf->state = PSA_TLS12_PRF_STATE_KEY_SET;
 
     return( PSA_SUCCESS );
 }
@@ -5916,7 +5403,7 @@
                                              const uint8_t *data,
                                              size_t data_length )
 {
-    if( prf->state != TLS12_PRF_STATE_KEY_SET )
+    if( prf->state != PSA_TLS12_PRF_STATE_KEY_SET )
         return( PSA_ERROR_BAD_STATE );
 
     if( data_length != 0 )
@@ -5929,7 +5416,7 @@
         prf->label_length = data_length;
     }
 
-    prf->state = TLS12_PRF_STATE_LABEL_SET;
+    prf->state = PSA_TLS12_PRF_STATE_LABEL_SET;
 
     return( PSA_SUCCESS );
 }
@@ -6129,8 +5616,8 @@
 
     status = psa_key_derivation_input_internal( operation,
                                                 step, slot->attr.type,
-                                                slot->data.key.data,
-                                                slot->data.key.bytes );
+                                                slot->key.data,
+                                                slot->key.bytes );
 
     unlock_status = psa_unlock_key_slot( slot );
 
@@ -6158,10 +5645,11 @@
     psa_ecc_family_t curve = mbedtls_ecc_group_to_psa( our_key->grp.id, &bits );
     mbedtls_ecdh_init( &ecdh );
 
-    status = psa_load_ecp_representation( PSA_KEY_TYPE_ECC_PUBLIC_KEY(curve),
-                                          peer_key,
-                                          peer_key_length,
-                                          &their_key );
+    status = mbedtls_psa_ecp_load_representation(
+                 PSA_KEY_TYPE_ECC_PUBLIC_KEY(curve),
+                 peer_key,
+                 peer_key_length,
+                 &their_key );
     if( status != PSA_SUCCESS )
         goto exit;
 
@@ -6213,11 +5701,11 @@
             if( ! PSA_KEY_TYPE_IS_ECC_KEY_PAIR( private_key->attr.type ) )
                 return( PSA_ERROR_INVALID_ARGUMENT );
             mbedtls_ecp_keypair *ecp = NULL;
-            psa_status_t status = psa_load_ecp_representation(
-                                    private_key->attr.type,
-                                    private_key->data.key.data,
-                                    private_key->data.key.bytes,
-                                    &ecp );
+            psa_status_t status = mbedtls_psa_ecp_load_representation(
+                                      private_key->attr.type,
+                                      private_key->key.data,
+                                      private_key->key.bytes,
+                                      &ecp );
             if( status != PSA_SUCCESS )
                 return( status );
             status = psa_key_agreement_ecdh( peer_key, peer_key_length,
@@ -6558,16 +6046,16 @@
         if( status != PSA_SUCCESS )
             return( status );
 
-        status = psa_generate_random( slot->data.key.data,
-                                      slot->data.key.bytes );
+        status = psa_generate_random( slot->key.data,
+                                      slot->key.bytes );
         if( status != PSA_SUCCESS )
             return( status );
 
         slot->attr.bits = (psa_key_bits_t) bits;
 #if defined(MBEDTLS_DES_C)
         if( type == PSA_KEY_TYPE_DES )
-            psa_des_set_key_parity( slot->data.key.data,
-                                    slot->data.key.bytes );
+            psa_des_set_key_parity( slot->key.data,
+                                    slot->key.bytes );
 #endif /* MBEDTLS_DES_C */
     }
     else
@@ -6582,7 +6070,7 @@
         if( bits > PSA_VENDOR_RSA_MAX_KEY_BITS )
             return( PSA_ERROR_NOT_SUPPORTED );
         /* Accept only byte-aligned keys, for the same reasons as
-         * in psa_import_rsa_key(). */
+         * in mbedtls_psa_rsa_import_key(). */
         if( bits % 8 != 0 )
             return( PSA_ERROR_NOT_SUPPORTED );
         status = psa_read_rsa_exponent( domain_parameters,
@@ -6609,11 +6097,11 @@
             return( status );
         }
 
-        status = psa_export_rsa_key( type,
-                                     &rsa,
-                                     slot->data.key.data,
-                                     bytes,
-                                     &slot->data.key.bytes );
+        status = mbedtls_psa_rsa_export_key( type,
+                                             &rsa,
+                                             slot->key.data,
+                                             bytes,
+                                             &slot->key.bytes );
         mbedtls_rsa_free( &rsa );
         if( status != PSA_SUCCESS )
             psa_remove_key_data_from_memory( slot );
@@ -6657,11 +6145,11 @@
         }
 
         status = mbedtls_to_psa_error(
-            mbedtls_ecp_write_key( &ecp, slot->data.key.data, bytes ) );
+            mbedtls_ecp_write_key( &ecp, slot->key.data, bytes ) );
 
         mbedtls_ecp_keypair_free( &ecp );
         if( status != PSA_SUCCESS ) {
-            memset( slot->data.key.data, 0, bytes );
+            memset( slot->key.data, 0, bytes );
             psa_remove_key_data_from_memory( slot );
         }
         return( status );
@@ -6771,7 +6259,7 @@
         default:
             /* We found an unsupported transaction in the storage.
              * We don't know what state the storage is in. Give up. */
-            return( PSA_ERROR_STORAGE_FAILURE );
+            return( PSA_ERROR_DATA_INVALID );
     }
 }
 #endif /* PSA_CRYPTO_STORAGE_HAS_TRANSACTIONS */
diff --git a/library/psa_crypto_core.h b/library/psa_crypto_core.h
index f61ef95..cf9d6d0 100644
--- a/library/psa_crypto_core.h
+++ b/library/psa_crypto_core.h
@@ -62,23 +62,13 @@
      */
     size_t lock_count;
 
-    union
+    /* Dynamically allocated key data buffer.
+     * Format as specified in psa_export_key(). */
+    struct key_data
     {
-        /* Dynamically allocated key data buffer.
-         * Format as specified in psa_export_key(). */
-        struct key_data
-        {
-            uint8_t *data;
-            size_t bytes;
-        } key;
-#if defined(MBEDTLS_PSA_CRYPTO_SE_C)
-        /* Any key type in a secure element */
-        struct se
-        {
-            psa_key_slot_number_t slot_number;
-        } se;
-#endif /* MBEDTLS_PSA_CRYPTO_SE_C */
-    } data;
+        uint8_t *data;
+        size_t bytes;
+    } key;
 } psa_key_slot_t;
 
 /* A mask of key attribute flags used only internally.
@@ -163,6 +153,20 @@
     slot->attr.flags &= ~mask;
 }
 
+#if defined(MBEDTLS_PSA_CRYPTO_SE_C)
+/** Get the SE slot number of a key from the key slot storing its description.
+ *
+ * \param[in]  slot  The key slot to query. This must be a key slot storing
+ *                   the description of a key of a dynamically registered
+ *                   secure element, otherwise the behaviour is undefined.
+ */
+static inline psa_key_slot_number_t psa_key_slot_get_slot_number(
+    const psa_key_slot_t *slot )
+{
+    return( *( (psa_key_slot_number_t *)( slot->key.data ) ) );
+}
+#endif
+
 /** Completely wipe a slot in memory, including its policy.
  *
  * Persistent storage is not affected.
@@ -208,4 +212,91 @@
  */
 psa_status_t mbedtls_to_psa_error( int ret );
 
+/** Import a key in binary format.
+ *
+ * \note The signature of this function is that of a PSA driver
+ *       import_key entry point. This function behaves as an import_key
+ *       entry point as defined in the PSA driver interface specification for
+ *       transparent drivers.
+ *
+ * \param[in]  attributes       The attributes for the key to import.
+ * \param[in]  data             The buffer containing the key data in import
+ *                              format.
+ * \param[in]  data_length      Size of the \p data buffer in bytes.
+ * \param[out] key_buffer       The buffer to contain the key data in output
+ *                              format upon successful return.
+ * \param[in]  key_buffer_size  Size of the \p key_buffer buffer in bytes. This
+ *                              size is greater or equal to \p data_length.
+ * \param[out] key_buffer_length  The length of the data written in \p
+ *                                key_buffer in bytes.
+ * \param[out] bits             The key size in number of bits.
+ *
+ * \retval #PSA_SUCCESS  The key was imported successfully.
+ * \retval #PSA_ERROR_INVALID_ARGUMENT
+ *         The key data is not correctly formatted.
+ * \retval #PSA_ERROR_NOT_SUPPORTED
+ * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
+ * \retval #PSA_ERROR_CORRUPTION_DETECTED
+ */
+psa_status_t psa_import_key_into_slot(
+    const psa_key_attributes_t *attributes,
+    const uint8_t *data, size_t data_length,
+    uint8_t *key_buffer, size_t key_buffer_size,
+    size_t *key_buffer_length, size_t *bits );
+
+/** Export a key in binary format
+ *
+ * \note The signature of this function is that of a PSA driver export_key
+ *       entry point. This function behaves as an export_key entry point as
+ *       defined in the PSA driver interface specification.
+ *
+ * \param[in]  attributes       The attributes for the key to export.
+ * \param[in]  key_buffer       Material or context of the key to export.
+ * \param[in]  key_buffer_size  Size of the \p key_buffer buffer in bytes.
+ * \param[out] data             Buffer where the key data is to be written.
+ * \param[in]  data_size        Size of the \p data buffer in bytes.
+ * \param[out] data_length      On success, the number of bytes written in
+ *                              \p data
+ *
+ * \retval #PSA_SUCCESS  The key was exported successfully.
+ * \retval #PSA_ERROR_NOT_SUPPORTED
+ * \retval #PSA_ERROR_COMMUNICATION_FAILURE
+ * \retval #PSA_ERROR_HARDWARE_FAILURE
+ * \retval #PSA_ERROR_CORRUPTION_DETECTED
+ * \retval #PSA_ERROR_STORAGE_FAILURE
+ * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
+ */
+psa_status_t psa_export_key_internal(
+    const psa_key_attributes_t *attributes,
+    const uint8_t *key_buffer, size_t key_buffer_size,
+    uint8_t *data, size_t data_size, size_t *data_length );
+
+/** Export a public key or the public part of a key pair in binary format.
+ *
+ * \note The signature of this function is that of a PSA driver
+ *       export_public_key entry point. This function behaves as an
+ *       export_public_key entry point as defined in the PSA driver interface
+ *       specification.
+ *
+ * \param[in]  attributes       The attributes for the key to export.
+ * \param[in]  key_buffer       Material or context of the key to export.
+ * \param[in]  key_buffer_size  Size of the \p key_buffer buffer in bytes.
+ * \param[out] data             Buffer where the key data is to be written.
+ * \param[in]  data_size        Size of the \p data buffer in bytes.
+ * \param[out] data_length      On success, the number of bytes written in
+ *                              \p data
+ *
+ * \retval #PSA_SUCCESS  The public key was exported successfully.
+ * \retval #PSA_ERROR_NOT_SUPPORTED
+ * \retval #PSA_ERROR_COMMUNICATION_FAILURE
+ * \retval #PSA_ERROR_HARDWARE_FAILURE
+ * \retval #PSA_ERROR_CORRUPTION_DETECTED
+ * \retval #PSA_ERROR_STORAGE_FAILURE
+ * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
+ */
+psa_status_t psa_export_public_key_internal(
+    const psa_key_attributes_t *attributes,
+    const uint8_t *key_buffer, size_t key_buffer_size,
+    uint8_t *data, size_t data_size, size_t *data_length );
+
 #endif /* PSA_CRYPTO_CORE_H */
diff --git a/library/psa_crypto_driver_wrappers.c b/library/psa_crypto_driver_wrappers.c
index 1243bd3..1eb9ebe 100644
--- a/library/psa_crypto_driver_wrappers.c
+++ b/library/psa_crypto_driver_wrappers.c
@@ -80,7 +80,7 @@
             return( PSA_ERROR_NOT_SUPPORTED );
         }
         return( drv->asymmetric->p_sign( drv_context,
-                                         slot->data.se.slot_number,
+                                         psa_key_slot_get_slot_number( slot ),
                                          alg,
                                          hash, hash_length,
                                          signature, signature_size,
@@ -103,8 +103,8 @@
              * cycle through all known transparent accelerators */
 #if defined(PSA_CRYPTO_DRIVER_TEST)
             status = test_transparent_signature_sign_hash( &attributes,
-                                                           slot->data.key.data,
-                                                           slot->data.key.bytes,
+                                                           slot->key.data,
+                                                           slot->key.bytes,
                                                            alg,
                                                            hash,
                                                            hash_length,
@@ -121,8 +121,8 @@
 #if defined(PSA_CRYPTO_DRIVER_TEST)
         case PSA_CRYPTO_TEST_DRIVER_LIFETIME:
             return( test_opaque_signature_sign_hash( &attributes,
-                                                     slot->data.key.data,
-                                                     slot->data.key.bytes,
+                                                     slot->key.data,
+                                                     slot->key.bytes,
                                                      alg,
                                                      hash,
                                                      hash_length,
@@ -172,7 +172,7 @@
             return( PSA_ERROR_NOT_SUPPORTED );
         }
         return( drv->asymmetric->p_verify( drv_context,
-                                           slot->data.se.slot_number,
+                                           psa_key_slot_get_slot_number( slot ),
                                            alg,
                                            hash, hash_length,
                                            signature, signature_length ) );
@@ -194,8 +194,8 @@
              * cycle through all known transparent accelerators */
 #if defined(PSA_CRYPTO_DRIVER_TEST)
             status = test_transparent_signature_verify_hash( &attributes,
-                                                             slot->data.key.data,
-                                                             slot->data.key.bytes,
+                                                             slot->key.data,
+                                                             slot->key.bytes,
                                                              alg,
                                                              hash,
                                                              hash_length,
@@ -211,8 +211,8 @@
 #if defined(PSA_CRYPTO_DRIVER_TEST)
         case PSA_CRYPTO_TEST_DRIVER_LIFETIME:
             return( test_opaque_signature_verify_hash( &attributes,
-                                                       slot->data.key.data,
-                                                       slot->data.key.bytes,
+                                                       slot->key.data,
+                                                       slot->key.bytes,
                                                        alg,
                                                        hash,
                                                        hash_length,
@@ -330,9 +330,8 @@
             return( PSA_ERROR_NOT_SUPPORTED );
         }
         return( drv->key_management->p_generate(
-            drv_context,
-            slot->data.se.slot_number, attributes,
-            NULL, 0, &pubkey_length ) );
+            drv_context, psa_key_slot_get_slot_number( slot ),
+            attributes, NULL, 0, &pubkey_length ) );
     }
 #endif /* MBEDTLS_PSA_CRYPTO_SE_C */
 
@@ -346,10 +345,10 @@
     if( status != PSA_SUCCESS )
         return( status );
 
-    slot->data.key.data = mbedtls_calloc(1, export_size);
-    if( slot->data.key.data == NULL )
+    slot->key.data = mbedtls_calloc(1, export_size);
+    if( slot->key.data == NULL )
         return( PSA_ERROR_INSUFFICIENT_MEMORY );
-    slot->data.key.bytes = export_size;
+    slot->key.bytes = export_size;
 
     switch( location )
     {
@@ -365,9 +364,9 @@
             }
 #if defined(PSA_CRYPTO_DRIVER_TEST)
             status = test_transparent_generate_key( attributes,
-                                                    slot->data.key.data,
-                                                    slot->data.key.bytes,
-                                                    &slot->data.key.bytes );
+                                                    slot->key.data,
+                                                    slot->key.bytes,
+                                                    &slot->key.bytes );
             /* Declared with fallback == true */
             if( status != PSA_ERROR_NOT_SUPPORTED )
                 break;
@@ -379,9 +378,9 @@
 #if defined(PSA_CRYPTO_DRIVER_TEST)
         case PSA_CRYPTO_TEST_DRIVER_LIFETIME:
             status = test_opaque_generate_key( attributes,
-                                               slot->data.key.data,
-                                               slot->data.key.bytes,
-                                               &slot->data.key.bytes );
+                                               slot->key.data,
+                                               slot->key.bytes,
+                                               &slot->key.bytes );
             break;
 #endif /* PSA_CRYPTO_DRIVER_TEST */
         default:
@@ -393,9 +392,9 @@
     if( status != PSA_SUCCESS )
     {
         /* free allocated buffer */
-        mbedtls_free( slot->data.key.data );
-        slot->data.key.data = NULL;
-        slot->data.key.bytes = 0;
+        mbedtls_free( slot->key.data );
+        slot->key.data = NULL;
+        slot->key.bytes = 0;
     }
 
     return( status );
@@ -410,55 +409,177 @@
 #endif /* PSA_CRYPTO_DRIVER_PRESENT */
 }
 
-psa_status_t psa_driver_wrapper_validate_key( const psa_key_attributes_t *attributes,
-                                              const uint8_t *data,
-                                              size_t data_length,
-                                              size_t *bits )
+psa_status_t psa_driver_wrapper_import_key(
+    const psa_key_attributes_t *attributes,
+    const uint8_t *data,
+    size_t data_length,
+    uint8_t *key_buffer,
+    size_t key_buffer_size,
+    size_t *key_buffer_length,
+    size_t *bits )
 {
-#if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
     psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
-    /* Try accelerators in turn */
-#if defined(PSA_CRYPTO_DRIVER_TEST)
-    status = test_transparent_validate_key( attributes,
-                                            data,
-                                            data_length,
-                                            bits );
-    /* Declared with fallback == true */
-    if( status != PSA_ERROR_NOT_SUPPORTED )
-        return( status );
-#endif /* PSA_CRYPTO_DRIVER_TEST */
+    psa_key_location_t location = PSA_KEY_LIFETIME_GET_LOCATION(
+                                      psa_get_key_lifetime( attributes ) );
 
-    return( PSA_ERROR_NOT_SUPPORTED );
-#else /* PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT */
-    (void) attributes;
-    (void) data;
-    (void) data_length;
-    (void) bits;
-    return( PSA_ERROR_NOT_SUPPORTED );
-#endif /* PSA_CRYPTO_DRIVER_PRESENT */
-}
+    /* Try dynamically-registered SE interface first */
+#if defined(MBEDTLS_PSA_CRYPTO_SE_C)
+    const psa_drv_se_t *drv;
+    psa_drv_se_context_t *drv_context;
 
-psa_status_t psa_driver_wrapper_export_public_key( const psa_key_slot_t *slot,
-                                                   uint8_t *data,
-                                                   size_t data_size,
-                                                   size_t *data_length )
-{
-#if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
-    psa_status_t status = PSA_ERROR_INVALID_ARGUMENT;
-    psa_key_location_t location = PSA_KEY_LIFETIME_GET_LOCATION(slot->attr.lifetime);
-    psa_key_attributes_t attributes = {
-      .core = slot->attr
-    };
+    if( psa_get_se_driver( attributes->core.lifetime, &drv, &drv_context ) )
+    {
+        if( drv->key_management == NULL ||
+            drv->key_management->p_import == NULL )
+            return( PSA_ERROR_NOT_SUPPORTED );
+
+        /* The driver should set the number of key bits, however in
+         * case it doesn't, we initialize bits to an invalid value. */
+        *bits = PSA_MAX_KEY_BITS + 1;
+        status = drv->key_management->p_import(
+            drv_context,
+            *( (psa_key_slot_number_t *)key_buffer ),
+            attributes, data, data_length, bits );
+
+        if( status != PSA_SUCCESS )
+            return( status );
+
+        if( (*bits) > PSA_MAX_KEY_BITS )
+            return( PSA_ERROR_NOT_SUPPORTED );
+
+        return( PSA_SUCCESS );
+    }
+#endif /* PSA_CRYPTO_SE_C */
 
     switch( location )
     {
         case PSA_KEY_LOCATION_LOCAL_STORAGE:
             /* Key is stored in the slot in export representation, so
              * cycle through all known transparent accelerators */
+#if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
 #if defined(PSA_CRYPTO_DRIVER_TEST)
-            status = test_transparent_export_public_key( &attributes,
-                                                         slot->data.key.data,
-                                                         slot->data.key.bytes,
+            status = test_transparent_import_key( attributes,
+                                                  data, data_length,
+                                                  key_buffer, key_buffer_size,
+                                                  key_buffer_length, bits );
+            /* Declared with fallback == true */
+            if( status != PSA_ERROR_NOT_SUPPORTED )
+                return( status );
+#endif /* PSA_CRYPTO_DRIVER_TEST */
+#endif /* PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT */
+            /* Fell through, meaning no accelerator supports this operation */
+            return( psa_import_key_into_slot( attributes,
+                                              data, data_length,
+                                              key_buffer, key_buffer_size,
+                                              key_buffer_length, bits ) );
+
+        default:
+            /* Importing a key with external storage in not yet supported.
+             * Return in error indicating that the lifetime is not valid. */
+            (void)status;
+            return( PSA_ERROR_INVALID_ARGUMENT );
+    }
+
+}
+
+psa_status_t psa_driver_wrapper_export_key(
+    const psa_key_attributes_t *attributes,
+    const uint8_t *key_buffer, size_t key_buffer_size,
+    uint8_t *data, size_t data_size, size_t *data_length )
+
+{
+    psa_status_t status = PSA_ERROR_INVALID_ARGUMENT;
+    psa_key_location_t location = PSA_KEY_LIFETIME_GET_LOCATION(
+                                      psa_get_key_lifetime( attributes ) );
+
+    /* Try dynamically-registered SE interface first */
+#if defined(MBEDTLS_PSA_CRYPTO_SE_C)
+    const psa_drv_se_t *drv;
+    psa_drv_se_context_t *drv_context;
+
+    if( psa_get_se_driver( attributes->core.lifetime, &drv, &drv_context ) )
+    {
+        if( ( drv->key_management == NULL   ) ||
+            ( drv->key_management->p_export == NULL ) )
+        {
+            return( PSA_ERROR_NOT_SUPPORTED );
+        }
+
+        return( drv->key_management->p_export(
+                     drv_context,
+                     *( (psa_key_slot_number_t *)key_buffer ),
+                     data, data_size, data_length ) );
+    }
+#endif /* PSA_CRYPTO_SE_C */
+
+    switch( location )
+    {
+        case PSA_KEY_LOCATION_LOCAL_STORAGE:
+            return( psa_export_key_internal( attributes,
+                                             key_buffer,
+                                             key_buffer_size,
+                                             data,
+                                             data_size,
+                                             data_length ) );
+
+        /* Add cases for opaque driver here */
+#if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
+#if defined(PSA_CRYPTO_DRIVER_TEST)
+        case PSA_CRYPTO_TEST_DRIVER_LIFETIME:
+            return( test_opaque_export_key( attributes,
+                                            key_buffer,
+                                            key_buffer_size,
+                                            data,
+                                            data_size,
+                                            data_length ) );
+#endif /* PSA_CRYPTO_DRIVER_TEST */
+#endif /* PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT */
+        default:
+            /* Key is declared with a lifetime not known to us */
+            return( status );
+    }
+}
+
+psa_status_t psa_driver_wrapper_export_public_key(
+    const psa_key_attributes_t *attributes,
+    const uint8_t *key_buffer, size_t key_buffer_size,
+    uint8_t *data, size_t data_size, size_t *data_length )
+
+{
+    psa_status_t status = PSA_ERROR_INVALID_ARGUMENT;
+    psa_key_location_t location = PSA_KEY_LIFETIME_GET_LOCATION(
+                                      psa_get_key_lifetime( attributes ) );
+
+    /* Try dynamically-registered SE interface first */
+#if defined(MBEDTLS_PSA_CRYPTO_SE_C)
+    const psa_drv_se_t *drv;
+    psa_drv_se_context_t *drv_context;
+
+    if( psa_get_se_driver( attributes->core.lifetime, &drv, &drv_context ) )
+    {
+        if( ( drv->key_management == NULL ) ||
+            ( drv->key_management->p_export_public == NULL ) )
+        {
+            return( PSA_ERROR_NOT_SUPPORTED );
+        }
+
+        return( drv->key_management->p_export_public(
+                    drv_context,
+                    *( (psa_key_slot_number_t *)key_buffer ),
+                    data, data_size, data_length ) );
+    }
+#endif /* MBEDTLS_PSA_CRYPTO_SE_C */
+
+    switch( location )
+    {
+        case PSA_KEY_LOCATION_LOCAL_STORAGE:
+            /* Key is stored in the slot in export representation, so
+             * cycle through all known transparent accelerators */
+#if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
+#if defined(PSA_CRYPTO_DRIVER_TEST)
+            status = test_transparent_export_public_key( attributes,
+                                                         key_buffer,
+                                                         key_buffer_size,
                                                          data,
                                                          data_size,
                                                          data_length );
@@ -466,29 +587,31 @@
             if( status != PSA_ERROR_NOT_SUPPORTED )
                 return( status );
 #endif /* PSA_CRYPTO_DRIVER_TEST */
+#endif /* PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT */
             /* Fell through, meaning no accelerator supports this operation */
-            return( PSA_ERROR_NOT_SUPPORTED );
+            return( psa_export_public_key_internal( attributes,
+                                                    key_buffer,
+                                                    key_buffer_size,
+                                                    data,
+                                                    data_size,
+                                                    data_length ) );
+
         /* Add cases for opaque driver here */
+#if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
 #if defined(PSA_CRYPTO_DRIVER_TEST)
         case PSA_CRYPTO_TEST_DRIVER_LIFETIME:
-            return( test_opaque_export_public_key( &attributes,
-                                                   slot->data.key.data,
-                                                   slot->data.key.bytes,
+            return( test_opaque_export_public_key( attributes,
+                                                   key_buffer,
+                                                   key_buffer_size,
                                                    data,
                                                    data_size,
                                                    data_length ) );
 #endif /* PSA_CRYPTO_DRIVER_TEST */
+#endif /* PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT */
         default:
             /* Key is declared with a lifetime not known to us */
             return( status );
     }
-#else /* PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT */
-    (void) slot;
-    (void) data;
-    (void) data_size;
-    (void) data_length;
-    return( PSA_ERROR_NOT_SUPPORTED );
-#endif /* PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT */
 }
 
 /*
@@ -517,8 +640,8 @@
              * cycle through all known transparent accelerators */
 #if defined(PSA_CRYPTO_DRIVER_TEST)
             status = test_transparent_cipher_encrypt( &attributes,
-                                                      slot->data.key.data,
-                                                      slot->data.key.bytes,
+                                                      slot->key.data,
+                                                      slot->key.bytes,
                                                       alg,
                                                       input,
                                                       input_length,
@@ -535,8 +658,8 @@
 #if defined(PSA_CRYPTO_DRIVER_TEST)
         case PSA_CRYPTO_TEST_DRIVER_LIFETIME:
             return( test_opaque_cipher_encrypt( &attributes,
-                                                slot->data.key.data,
-                                                slot->data.key.bytes,
+                                                slot->key.data,
+                                                slot->key.bytes,
                                                 alg,
                                                 input,
                                                 input_length,
@@ -584,8 +707,8 @@
              * cycle through all known transparent accelerators */
 #if defined(PSA_CRYPTO_DRIVER_TEST)
             status = test_transparent_cipher_decrypt( &attributes,
-                                                      slot->data.key.data,
-                                                      slot->data.key.bytes,
+                                                      slot->key.data,
+                                                      slot->key.bytes,
                                                       alg,
                                                       input,
                                                       input_length,
@@ -602,8 +725,8 @@
 #if defined(PSA_CRYPTO_DRIVER_TEST)
         case PSA_CRYPTO_TEST_DRIVER_LIFETIME:
             return( test_opaque_cipher_decrypt( &attributes,
-                                                slot->data.key.data,
-                                                slot->data.key.bytes,
+                                                slot->key.data,
+                                                slot->key.bytes,
                                                 alg,
                                                 input,
                                                 input_length,
@@ -652,8 +775,8 @@
 
             status = test_transparent_cipher_encrypt_setup( operation->ctx,
                                                             &attributes,
-                                                            slot->data.key.data,
-                                                            slot->data.key.bytes,
+                                                            slot->key.data,
+                                                            slot->key.bytes,
                                                             alg );
             /* Declared with fallback == true */
             if( status == PSA_SUCCESS )
@@ -680,8 +803,8 @@
 
             status = test_opaque_cipher_encrypt_setup( operation->ctx,
                                                        &attributes,
-                                                       slot->data.key.data,
-                                                       slot->data.key.bytes,
+                                                       slot->key.data,
+                                                       slot->key.bytes,
                                                        alg );
             if( status == PSA_SUCCESS )
                 operation->id = PSA_CRYPTO_OPAQUE_TEST_DRIVER_ID;
@@ -733,8 +856,8 @@
 
             status = test_transparent_cipher_decrypt_setup( operation->ctx,
                                                             &attributes,
-                                                            slot->data.key.data,
-                                                            slot->data.key.bytes,
+                                                            slot->key.data,
+                                                            slot->key.bytes,
                                                             alg );
             /* Declared with fallback == true */
             if( status == PSA_SUCCESS )
@@ -761,8 +884,8 @@
 
             status = test_opaque_cipher_decrypt_setup( operation->ctx,
                                                        &attributes,
-                                                       slot->data.key.data,
-                                                       slot->data.key.bytes,
+                                                       slot->key.data,
+                                                       slot->key.bytes,
                                                        alg );
             if( status == PSA_SUCCESS )
                 operation->id = PSA_CRYPTO_OPAQUE_TEST_DRIVER_ID;
diff --git a/library/psa_crypto_driver_wrappers.h b/library/psa_crypto_driver_wrappers.h
index 6b51437..27d8b64 100644
--- a/library/psa_crypto_driver_wrappers.h
+++ b/library/psa_crypto_driver_wrappers.h
@@ -47,18 +47,25 @@
  * Key handling functions
  */
 
-psa_status_t psa_driver_wrapper_generate_key( const psa_key_attributes_t *attributes,
-                                              psa_key_slot_t *slot );
+psa_status_t psa_driver_wrapper_import_key(
+    const psa_key_attributes_t *attributes,
+    const uint8_t *data, size_t data_length,
+    uint8_t *key_buffer, size_t key_buffer_size,
+    size_t *key_buffer_length, size_t *bits );
 
-psa_status_t psa_driver_wrapper_validate_key( const psa_key_attributes_t *attributes,
-                                              const uint8_t *data,
-                                              size_t data_length,
-                                              size_t *bits );
+psa_status_t psa_driver_wrapper_export_key(
+    const psa_key_attributes_t *attributes,
+    const uint8_t *key_buffer, size_t key_buffer_size,
+    uint8_t *data, size_t data_size, size_t *data_length );
 
-psa_status_t psa_driver_wrapper_export_public_key( const psa_key_slot_t *slot,
-                                                   uint8_t *data,
-                                                   size_t data_size,
-                                                   size_t *data_length );
+psa_status_t psa_driver_wrapper_export_public_key(
+    const psa_key_attributes_t *attributes,
+    const uint8_t *key_buffer, size_t key_buffer_size,
+    uint8_t *data, size_t data_size, size_t *data_length );
+
+psa_status_t psa_driver_wrapper_generate_key(
+    const psa_key_attributes_t *attributes,
+    psa_key_slot_t *slot );
 
 /*
  * Cipher functions
diff --git a/library/psa_crypto_ecp.c b/library/psa_crypto_ecp.c
new file mode 100644
index 0000000..bfc9674
--- /dev/null
+++ b/library/psa_crypto_ecp.c
@@ -0,0 +1,331 @@
+/*
+ *  PSA ECP layer on top of Mbed TLS crypto
+ */
+/*
+ *  Copyright The Mbed TLS Contributors
+ *  SPDX-License-Identifier: Apache-2.0
+ *
+ *  Licensed under the Apache License, Version 2.0 (the "License"); you may
+ *  not use this file except in compliance with the License.
+ *  You may obtain a copy of the License at
+ *
+ *  http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *  Unless required by applicable law or agreed to in writing, software
+ *  distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ *  WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *  See the License for the specific language governing permissions and
+ *  limitations under the License.
+ */
+
+#include "common.h"
+
+#if defined(MBEDTLS_PSA_CRYPTO_C)
+
+#include <psa/crypto.h>
+#include "psa_crypto_core.h"
+#include "psa_crypto_ecp.h"
+#include "psa_crypto_random_impl.h"
+
+#include <stdlib.h>
+#include <string.h>
+#include "mbedtls/platform.h"
+#if !defined(MBEDTLS_PLATFORM_C)
+#define mbedtls_calloc calloc
+#define mbedtls_free   free
+#endif
+
+#include <mbedtls/ecp.h>
+#include <mbedtls/error.h>
+
+#if ( defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_ECC_KEY_PAIR) ||  \
+      ( defined(PSA_CRYPTO_DRIVER_TEST) &&                   \
+        defined(MBEDTLS_PSA_ACCEL_KEY_TYPE_ECC_KEY_PAIR) ) )
+#define BUILTIN_KEY_TYPE_ECC_KEY_PAIR    1
+#endif
+
+#if ( defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_ECC_PUBLIC_KEY) ||  \
+      ( defined(PSA_CRYPTO_DRIVER_TEST) &&                   \
+        defined(MBEDTLS_PSA_ACCEL_KEY_TYPE_ECC_PUBLIC_KEY) ) )
+#define BUILTIN_KEY_TYPE_ECC_PUBLIC_KEY  1
+#endif
+
+#if defined(BUILTIN_KEY_TYPE_ECC_KEY_PAIR) || \
+    defined(BUILTIN_KEY_TYPE_ECC_PUBLIC_KEY) || \
+    defined(MBEDTLS_PSA_BUILTIN_ALG_ECDSA) || \
+    defined(MBEDTLS_PSA_BUILTIN_ALG_ECDH) || \
+    defined(MBEDTLS_PSA_BUILTIN_ALG_DETERMINISTIC_ECDSA)
+psa_status_t mbedtls_psa_ecp_load_representation(
+    psa_key_type_t type, const uint8_t *data, size_t data_length,
+    mbedtls_ecp_keypair **p_ecp )
+{
+    mbedtls_ecp_group_id grp_id = MBEDTLS_ECP_DP_NONE;
+    psa_status_t status;
+    mbedtls_ecp_keypair *ecp = NULL;
+    size_t curve_size = data_length;
+
+    if( PSA_KEY_TYPE_IS_PUBLIC_KEY( type ) &&
+        PSA_KEY_TYPE_ECC_GET_FAMILY( type ) != PSA_ECC_FAMILY_MONTGOMERY )
+    {
+        /* A Weierstrass public key is represented as:
+         * - The byte 0x04;
+         * - `x_P` as a `ceiling(m/8)`-byte string, big-endian;
+         * - `y_P` as a `ceiling(m/8)`-byte string, big-endian.
+         * So its data length is 2m+1 where m is the curve size in bits.
+         */
+        if( ( data_length & 1 ) == 0 )
+            return( PSA_ERROR_INVALID_ARGUMENT );
+        curve_size = data_length / 2;
+
+        /* Montgomery public keys are represented in compressed format, meaning
+         * their curve_size is equal to the amount of input. */
+
+        /* Private keys are represented in uncompressed private random integer
+         * format, meaning their curve_size is equal to the amount of input. */
+    }
+
+    /* Allocate and initialize a key representation. */
+    ecp = mbedtls_calloc( 1, sizeof( mbedtls_ecp_keypair ) );
+    if( ecp == NULL )
+        return( PSA_ERROR_INSUFFICIENT_MEMORY );
+    mbedtls_ecp_keypair_init( ecp );
+
+    /* Load the group. */
+    grp_id = mbedtls_ecc_group_of_psa( PSA_KEY_TYPE_ECC_GET_FAMILY( type ),
+                                       curve_size );
+    if( grp_id == MBEDTLS_ECP_DP_NONE )
+    {
+        status = PSA_ERROR_INVALID_ARGUMENT;
+        goto exit;
+    }
+
+    status = mbedtls_to_psa_error(
+                mbedtls_ecp_group_load( &ecp->grp, grp_id ) );
+    if( status != PSA_SUCCESS )
+        goto exit;
+
+    /* Load the key material. */
+    if( PSA_KEY_TYPE_IS_PUBLIC_KEY( type ) )
+    {
+        /* Load the public value. */
+        status = mbedtls_to_psa_error(
+            mbedtls_ecp_point_read_binary( &ecp->grp, &ecp->Q,
+                                           data,
+                                           data_length ) );
+        if( status != PSA_SUCCESS )
+            goto exit;
+
+        /* Check that the point is on the curve. */
+        status = mbedtls_to_psa_error(
+            mbedtls_ecp_check_pubkey( &ecp->grp, &ecp->Q ) );
+        if( status != PSA_SUCCESS )
+            goto exit;
+    }
+    else
+    {
+        /* Load and validate the secret value. */
+        status = mbedtls_to_psa_error(
+            mbedtls_ecp_read_key( ecp->grp.id,
+                                  ecp,
+                                  data,
+                                  data_length ) );
+        if( status != PSA_SUCCESS )
+            goto exit;
+    }
+
+    *p_ecp = ecp;
+exit:
+    if( status != PSA_SUCCESS )
+    {
+        mbedtls_ecp_keypair_free( ecp );
+        mbedtls_free( ecp );
+    }
+
+    return( status );
+}
+#endif /* defined(BUILTIN_KEY_TYPE_ECC_KEY_PAIR) ||
+        * defined(BUILTIN_KEY_TYPE_ECC_PUBLIC_KEY) ||
+        * defined(MBEDTLS_PSA_BUILTIN_ALG_ECDSA) ||
+        * defined(MBEDTLS_PSA_BUILTIN_ALG_ECDH) ||
+        * defined(MBEDTLS_PSA_BUILTIN_ALG_DETERMINISTIC_ECDSA) */
+
+#if defined(BUILTIN_KEY_TYPE_ECC_KEY_PAIR) || \
+    defined(BUILTIN_KEY_TYPE_ECC_PUBLIC_KEY)
+
+static psa_status_t ecp_import_key(
+    const psa_key_attributes_t *attributes,
+    const uint8_t *data, size_t data_length,
+    uint8_t *key_buffer, size_t key_buffer_size,
+    size_t *key_buffer_length, size_t *bits )
+{
+    psa_status_t status;
+    mbedtls_ecp_keypair *ecp = NULL;
+
+    /* Parse input */
+    status = mbedtls_psa_ecp_load_representation( attributes->core.type,
+                                                  data,
+                                                  data_length,
+                                                  &ecp );
+    if( status != PSA_SUCCESS )
+        goto exit;
+
+    if( PSA_KEY_TYPE_ECC_GET_FAMILY( attributes->core.type ) ==
+        PSA_ECC_FAMILY_MONTGOMERY )
+        *bits = ecp->grp.nbits + 1;
+    else
+        *bits = ecp->grp.nbits;
+
+    /* Re-export the data to PSA export format. There is currently no support
+     * for other input formats then the export format, so this is a 1-1
+     * copy operation. */
+    status = mbedtls_psa_ecp_export_key( attributes->core.type,
+                                         ecp,
+                                         key_buffer,
+                                         key_buffer_size,
+                                         key_buffer_length );
+exit:
+    /* Always free the PK object (will also free contained ECP context) */
+    mbedtls_ecp_keypair_free( ecp );
+    mbedtls_free( ecp );
+
+    return( status );
+}
+
+psa_status_t mbedtls_psa_ecp_export_key( psa_key_type_t type,
+                                         mbedtls_ecp_keypair *ecp,
+                                         uint8_t *data,
+                                         size_t data_size,
+                                         size_t *data_length )
+{
+    psa_status_t status;
+
+    if( PSA_KEY_TYPE_IS_PUBLIC_KEY( type ) )
+    {
+        /* Check whether the public part is loaded */
+        if( mbedtls_ecp_is_zero( &ecp->Q ) )
+        {
+            /* Calculate the public key */
+            status = mbedtls_to_psa_error(
+                mbedtls_ecp_mul( &ecp->grp, &ecp->Q, &ecp->d, &ecp->grp.G,
+                                 mbedtls_psa_get_random,
+                                 MBEDTLS_PSA_RANDOM_STATE ) );
+            if( status != PSA_SUCCESS )
+                return( status );
+        }
+
+        status = mbedtls_to_psa_error(
+                    mbedtls_ecp_point_write_binary( &ecp->grp, &ecp->Q,
+                                                    MBEDTLS_ECP_PF_UNCOMPRESSED,
+                                                    data_length,
+                                                    data,
+                                                    data_size ) );
+        if( status != PSA_SUCCESS )
+            memset( data, 0, data_size );
+
+        return( status );
+    }
+    else
+    {
+        if( data_size < PSA_BITS_TO_BYTES( ecp->grp.nbits ) )
+            return( PSA_ERROR_BUFFER_TOO_SMALL );
+
+        status = mbedtls_to_psa_error(
+                    mbedtls_ecp_write_key( ecp,
+                                           data,
+                                           PSA_BITS_TO_BYTES( ecp->grp.nbits ) ) );
+        if( status == PSA_SUCCESS )
+            *data_length = PSA_BITS_TO_BYTES( ecp->grp.nbits );
+        else
+            memset( data, 0, data_size );
+
+        return( status );
+    }
+}
+
+static psa_status_t ecp_export_public_key(
+    const psa_key_attributes_t *attributes,
+    const uint8_t *key_buffer, size_t key_buffer_size,
+    uint8_t *data, size_t data_size, size_t *data_length )
+{
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+    mbedtls_ecp_keypair *ecp = NULL;
+
+    status = mbedtls_psa_ecp_load_representation(
+                 attributes->core.type, key_buffer, key_buffer_size, &ecp );
+    if( status != PSA_SUCCESS )
+        return( status );
+
+    status = mbedtls_psa_ecp_export_key(
+                 PSA_KEY_TYPE_ECC_PUBLIC_KEY(
+                     PSA_KEY_TYPE_ECC_GET_FAMILY( attributes->core.type ) ),
+                 ecp, data, data_size, data_length );
+
+    mbedtls_ecp_keypair_free( ecp );
+    mbedtls_free( ecp );
+
+    return( status );
+}
+#endif /* defined(BUILTIN_KEY_TYPE_ECC_KEY_PAIR) ||
+        * defined(BUILTIN_KEY_TYPE_ECC_PUBLIC_KEY) */
+
+#if defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_ECC_KEY_PAIR) || \
+    defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_ECC_PUBLIC_KEY)
+
+psa_status_t mbedtls_psa_ecp_import_key(
+    const psa_key_attributes_t *attributes,
+    const uint8_t *data, size_t data_length,
+    uint8_t *key_buffer, size_t key_buffer_size,
+    size_t *key_buffer_length, size_t *bits )
+{
+    return( ecp_import_key( attributes, data, data_length,
+                            key_buffer, key_buffer_size,
+                            key_buffer_length, bits ) );
+}
+
+psa_status_t mbedtls_psa_ecp_export_public_key(
+    const psa_key_attributes_t *attributes,
+    const uint8_t *key_buffer, size_t key_buffer_size,
+    uint8_t *data, size_t data_size, size_t *data_length )
+{
+    return( ecp_export_public_key( attributes, key_buffer, key_buffer_size,
+                                   data, data_size, data_length ) );
+}
+
+#endif /* defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_ECC_KEY_PAIR) ||
+        * defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_ECC_PUBLIC_KEY) */
+
+/*
+ * BEYOND THIS POINT, TEST DRIVER ENTRY POINTS ONLY.
+ */
+
+#if defined(PSA_CRYPTO_DRIVER_TEST)
+
+#if defined(MBEDTLS_PSA_ACCEL_KEY_TYPE_ECC_KEY_PAIR) || \
+    defined(MBEDTLS_PSA_ACCEL_KEY_TYPE_ECC_PUBLIC_KEY)
+
+psa_status_t mbedtls_transparent_test_driver_ecp_import_key(
+    const psa_key_attributes_t *attributes,
+    const uint8_t *data, size_t data_length,
+    uint8_t *key_buffer, size_t key_buffer_size,
+    size_t *key_buffer_length, size_t *bits )
+{
+    return( ecp_import_key( attributes, data, data_length,
+                            key_buffer, key_buffer_size,
+                            key_buffer_length, bits ) );
+}
+
+psa_status_t mbedtls_transparent_test_driver_ecp_export_public_key(
+    const psa_key_attributes_t *attributes,
+    const uint8_t *key_buffer, size_t key_buffer_size,
+    uint8_t *data, size_t data_size, size_t *data_length )
+{
+    return( ecp_export_public_key( attributes, key_buffer, key_buffer_size,
+                                   data, data_size, data_length ) );
+}
+
+#endif /* defined(MBEDTLS_PSA_ACCEL_KEY_TYPE_ECC_KEY_PAIR) ||
+          defined(MBEDTLS_PSA_ACCEL_KEY_TYPE_ECC_PUBLIC_KEY) */
+
+#endif /* PSA_CRYPTO_DRIVER_TEST */
+
+#endif /* MBEDTLS_PSA_CRYPTO_C */
diff --git a/library/psa_crypto_ecp.h b/library/psa_crypto_ecp.h
new file mode 100644
index 0000000..695f07e
--- /dev/null
+++ b/library/psa_crypto_ecp.h
@@ -0,0 +1,135 @@
+/*
+ *  PSA ECP layer on top of Mbed TLS crypto
+ */
+/*
+ *  Copyright The Mbed TLS Contributors
+ *  SPDX-License-Identifier: Apache-2.0
+ *
+ *  Licensed under the Apache License, Version 2.0 (the "License"); you may
+ *  not use this file except in compliance with the License.
+ *  You may obtain a copy of the License at
+ *
+ *  http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *  Unless required by applicable law or agreed to in writing, software
+ *  distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ *  WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *  See the License for the specific language governing permissions and
+ *  limitations under the License.
+ */
+
+#ifndef PSA_CRYPTO_ECP_H
+#define PSA_CRYPTO_ECP_H
+
+#include <psa/crypto.h>
+#include <mbedtls/ecp.h>
+
+/** Load the contents of a key buffer into an internal ECP representation
+ *
+ * \param[in] type          The type of key contained in \p data.
+ * \param[in] data          The buffer from which to load the representation.
+ * \param[in] data_length   The size in bytes of \p data.
+ * \param[out] p_ecp        Returns a pointer to an ECP context on success.
+ *                          The caller is responsible for freeing both the
+ *                          contents of the context and the context itself
+ *                          when done.
+ */
+psa_status_t mbedtls_psa_ecp_load_representation( psa_key_type_t type,
+                                                  const uint8_t *data,
+                                                  size_t data_length,
+                                                  mbedtls_ecp_keypair **p_ecp );
+
+/** Import an ECP key in binary format.
+ *
+ * \note The signature of this function is that of a PSA driver
+ *       import_key entry point. This function behaves as an import_key
+ *       entry point as defined in the PSA driver interface specification for
+ *       transparent drivers.
+ *
+ * \param[in]  attributes       The attributes for the key to import.
+ * \param[in]  data             The buffer containing the key data in import
+ *                              format.
+ * \param[in]  data_length      Size of the \p data buffer in bytes.
+ * \param[out] key_buffer       The buffer containing the key data in output
+ *                              format.
+ * \param[in]  key_buffer_size  Size of the \p key_buffer buffer in bytes. This
+ *                              size is greater or equal to \p data_length.
+ * \param[out] key_buffer_length  The length of the data written in \p
+ *                                key_buffer in bytes.
+ * \param[out] bits             The key size in number of bits.
+ *
+ * \retval #PSA_SUCCESS  The ECP key was imported successfully.
+ * \retval #PSA_ERROR_INVALID_ARGUMENT
+ *         The key data is not correctly formatted.
+ * \retval #PSA_ERROR_NOT_SUPPORTED
+ * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
+ * \retval #PSA_ERROR_CORRUPTION_DETECTED
+ */
+psa_status_t mbedtls_psa_ecp_import_key(
+    const psa_key_attributes_t *attributes,
+    const uint8_t *data, size_t data_length,
+    uint8_t *key_buffer, size_t key_buffer_size,
+    size_t *key_buffer_length, size_t *bits );
+
+/** Export an ECP key to export representation
+ *
+ * \param[in] type          The type of key (public/private) to export
+ * \param[in] ecp           The internal ECP representation from which to export
+ * \param[out] data         The buffer to export to
+ * \param[in] data_size     The length of the buffer to export to
+ * \param[out] data_length  The amount of bytes written to \p data
+ */
+psa_status_t mbedtls_psa_ecp_export_key( psa_key_type_t type,
+                                         mbedtls_ecp_keypair *ecp,
+                                         uint8_t *data,
+                                         size_t data_size,
+                                         size_t *data_length );
+
+/** Export an ECP public key or the public part of an ECP key pair in binary
+ *  format.
+ *
+ * \note The signature of this function is that of a PSA driver
+ *       export_public_key entry point. This function behaves as an
+ *       export_public_key entry point as defined in the PSA driver interface
+ *       specification.
+ *
+ * \param[in]  attributes       The attributes for the key to export.
+ * \param[in]  key_buffer       Material or context of the key to export.
+ * \param[in]  key_buffer_size  Size of the \p key_buffer buffer in bytes.
+ * \param[out] data             Buffer where the key data is to be written.
+ * \param[in]  data_size        Size of the \p data buffer in bytes.
+ * \param[out] data_length      On success, the number of bytes written in
+ *                              \p data
+ *
+ * \retval #PSA_SUCCESS  The ECP public key was exported successfully.
+ * \retval #PSA_ERROR_NOT_SUPPORTED
+ * \retval #PSA_ERROR_COMMUNICATION_FAILURE
+ * \retval #PSA_ERROR_HARDWARE_FAILURE
+ * \retval #PSA_ERROR_CORRUPTION_DETECTED
+ * \retval #PSA_ERROR_STORAGE_FAILURE
+ * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
+ */
+psa_status_t mbedtls_psa_ecp_export_public_key(
+    const psa_key_attributes_t *attributes,
+    const uint8_t *key_buffer, size_t key_buffer_size,
+    uint8_t *data, size_t data_size, size_t *data_length );
+
+/*
+ * BEYOND THIS POINT, TEST DRIVER ENTRY POINTS ONLY.
+ */
+
+#if defined(PSA_CRYPTO_DRIVER_TEST)
+psa_status_t mbedtls_transparent_test_driver_ecp_import_key(
+    const psa_key_attributes_t *attributes,
+    const uint8_t *data, size_t data_length,
+    uint8_t *key_buffer, size_t key_buffer_size,
+    size_t *key_buffer_length, size_t *bits );
+
+psa_status_t mbedtls_transparent_test_driver_ecp_export_public_key(
+    const psa_key_attributes_t *attributes,
+    const uint8_t *key_buffer, size_t key_buffer_size,
+    uint8_t *data, size_t data_size, size_t *data_length );
+
+#endif /* PSA_CRYPTO_DRIVER_TEST */
+
+#endif /* PSA_CRYPTO_ECP_H */
diff --git a/library/psa_crypto_its.h b/library/psa_crypto_its.h
index 11703a0..3a3f49a 100644
--- a/library/psa_crypto_its.h
+++ b/library/psa_crypto_its.h
@@ -72,12 +72,12 @@
  *
  * \return      A status indicating the success/failure of the operation
  *
- * \retval     #PSA_SUCCESS                      The operation completed successfully
- * \retval     #PSA_ERROR_NOT_PERMITTED          The operation failed because the provided `uid` value was already created with PSA_STORAGE_WRITE_ONCE_FLAG
- * \retval     #PSA_ERROR_NOT_SUPPORTED          The operation failed because one or more of the flags provided in `create_flags` is not supported or is not valid
- * \retval     #PSA_ERROR_INSUFFICIENT_STORAGE   The operation failed because there was insufficient space on the storage medium
- * \retval     #PSA_ERROR_STORAGE_FAILURE        The operation failed because the physical storage has failed (Fatal error)
- * \retval     #PSA_ERROR_INVALID_ARGUMENT       The operation failed because one of the provided pointers(`p_data`)
+ * \retval      #PSA_SUCCESS                     The operation completed successfully
+ * \retval      #PSA_ERROR_NOT_PERMITTED         The operation failed because the provided `uid` value was already created with PSA_STORAGE_WRITE_ONCE_FLAG
+ * \retval      #PSA_ERROR_NOT_SUPPORTED         The operation failed because one or more of the flags provided in `create_flags` is not supported or is not valid
+ * \retval      #PSA_ERROR_INSUFFICIENT_STORAGE  The operation failed because there was insufficient space on the storage medium
+ * \retval      #PSA_ERROR_STORAGE_FAILURE       The operation failed because the physical storage has failed (Fatal error)
+ * \retval      #PSA_ERROR_INVALID_ARGUMENT      The operation failed because one of the provided pointers(`p_data`)
  *                                               is invalid, for example is `NULL` or references memory the caller cannot access
  */
 psa_status_t psa_its_set(psa_storage_uid_t uid,
@@ -97,11 +97,11 @@
  *
  * \return      A status indicating the success/failure of the operation
  *
- * \retval     #PSA_SUCCESS                  The operation completed successfully
- * \retval     #PSA_ERROR_DOES_NOT_EXIST     The operation failed because the provided `uid` value was not found in the storage
- * \retval     #PSA_ERROR_INVALID_SIZE       The operation failed because the data associated with provided uid is larger than `data_size`
- * \retval     #PSA_ERROR_STORAGE_FAILURE    The operation failed because the physical storage has failed (Fatal error)
- * \retval     #PSA_ERROR_INVALID_ARGUMENT   The operation failed because one of the provided pointers(`p_data`, `p_data_length`)
+ * \retval      #PSA_SUCCESS                 The operation completed successfully
+ * \retval      #PSA_ERROR_DOES_NOT_EXIST    The operation failed because the provided `uid` value was not found in the storage
+ * \retval      #PSA_ERROR_STORAGE_FAILURE   The operation failed because the physical storage has failed (Fatal error)
+ * \retval      #PSA_ERROR_DATA_CORRUPT      The operation failed because stored data has been corrupted
+ * \retval      #PSA_ERROR_INVALID_ARGUMENT  The operation failed because one of the provided pointers(`p_data`, `p_data_length`)
  *                                           is invalid. For example is `NULL` or references memory the caller cannot access.
  *                                           In addition, this can also happen if an invalid offset was provided.
  */
@@ -119,10 +119,10 @@
  *
  * \return      A status indicating the success/failure of the operation
  *
- * \retval     #PSA_SUCCESS                  The operation completed successfully
- * \retval     #PSA_ERROR_DOES_NOT_EXIST     The operation failed because the provided uid value was not found in the storage
- * \retval     #PSA_ERROR_STORAGE_FAILURE    The operation failed because the physical storage has failed (Fatal error)
- * \retval     #PSA_ERROR_INVALID_ARGUMENT   The operation failed because one of the provided pointers(`p_info`)
+ * \retval      #PSA_SUCCESS                 The operation completed successfully
+ * \retval      #PSA_ERROR_DOES_NOT_EXIST    The operation failed because the provided uid value was not found in the storage
+ * \retval      #PSA_ERROR_DATA_CORRUPT      The operation failed because stored data has been corrupted
+ * \retval      #PSA_ERROR_INVALID_ARGUMENT  The operation failed because one of the provided pointers(`p_info`)
  *                                           is invalid, for example is `NULL` or references memory the caller cannot access
  */
 psa_status_t psa_its_get_info(psa_storage_uid_t uid,
@@ -135,10 +135,10 @@
  *
  * \return  A status indicating the success/failure of the operation
  *
- * \retval     #PSA_SUCCESS                  The operation completed successfully
- * \retval     #PSA_ERROR_DOES_NOT_EXIST     The operation failed because the provided key value was not found in the storage
- * \retval     #PSA_ERROR_NOT_PERMITTED      The operation failed because the provided key value was created with PSA_STORAGE_WRITE_ONCE_FLAG
- * \retval     #PSA_ERROR_STORAGE_FAILURE    The operation failed because the physical storage has failed (Fatal error)
+ * \retval      #PSA_SUCCESS                  The operation completed successfully
+ * \retval      #PSA_ERROR_DOES_NOT_EXIST     The operation failed because the provided key value was not found in the storage
+ * \retval      #PSA_ERROR_NOT_PERMITTED      The operation failed because the provided key value was created with PSA_STORAGE_WRITE_ONCE_FLAG
+ * \retval      #PSA_ERROR_STORAGE_FAILURE    The operation failed because the physical storage has failed (Fatal error)
  */
 psa_status_t psa_its_remove(psa_storage_uid_t uid);
 
diff --git a/library/psa_crypto_rsa.c b/library/psa_crypto_rsa.c
new file mode 100644
index 0000000..aae1d05
--- /dev/null
+++ b/library/psa_crypto_rsa.c
@@ -0,0 +1,321 @@
+/*
+ *  PSA RSA layer on top of Mbed TLS crypto
+ */
+/*
+ *  Copyright The Mbed TLS Contributors
+ *  SPDX-License-Identifier: Apache-2.0
+ *
+ *  Licensed under the Apache License, Version 2.0 (the "License"); you may
+ *  not use this file except in compliance with the License.
+ *  You may obtain a copy of the License at
+ *
+ *  http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *  Unless required by applicable law or agreed to in writing, software
+ *  distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ *  WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *  See the License for the specific language governing permissions and
+ *  limitations under the License.
+ */
+
+#include "common.h"
+
+#if defined(MBEDTLS_PSA_CRYPTO_C)
+
+#include <psa/crypto.h>
+#include "psa_crypto_core.h"
+#include "psa_crypto_rsa.h"
+
+#include <stdlib.h>
+#include <string.h>
+#include "mbedtls/platform.h"
+#if !defined(MBEDTLS_PLATFORM_C)
+#define mbedtls_calloc calloc
+#define mbedtls_free   free
+#endif
+
+#include <mbedtls/rsa.h>
+#include <mbedtls/error.h>
+#include <mbedtls/pk.h>
+#include <mbedtls/pk_internal.h>
+
+#if ( defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_RSA_KEY_PAIR) ||  \
+      ( defined(PSA_CRYPTO_DRIVER_TEST) &&                   \
+        defined(MBEDTLS_PSA_ACCEL_KEY_TYPE_RSA_KEY_PAIR) ) )
+#define BUILTIN_KEY_TYPE_RSA_KEY_PAIR    1
+#endif
+
+#if ( defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_RSA_PUBLIC_KEY) ||  \
+      ( defined(PSA_CRYPTO_DRIVER_TEST) &&                   \
+        defined(MBEDTLS_PSA_ACCEL_KEY_TYPE_RSA_PUBLIC_KEY) ) )
+#define BUILTIN_KEY_TYPE_RSA_PUBLIC_KEY  1
+#endif
+
+#if defined(MBEDTLS_PSA_BUILTIN_ALG_RSA_PKCS1V15_CRYPT) || \
+    defined(MBEDTLS_PSA_BUILTIN_ALG_RSA_PKCS1V15_SIGN) || \
+    defined(MBEDTLS_PSA_BUILTIN_ALG_RSA_OAEP) || \
+    defined(MBEDTLS_PSA_BUILTIN_ALG_RSA_PSS) || \
+    defined(BUILTIN_KEY_TYPE_RSA_KEY_PAIR) || \
+    defined(BUILTIN_KEY_TYPE_RSA_PUBLIC_KEY)
+
+/* Mbed TLS doesn't support non-byte-aligned key sizes (i.e. key sizes
+ * that are not a multiple of 8) well. For example, there is only
+ * mbedtls_rsa_get_len(), which returns a number of bytes, and no
+ * way to return the exact bit size of a key.
+ * To keep things simple, reject non-byte-aligned key sizes. */
+static psa_status_t psa_check_rsa_key_byte_aligned(
+    const mbedtls_rsa_context *rsa )
+{
+    mbedtls_mpi n;
+    psa_status_t status;
+    mbedtls_mpi_init( &n );
+    status = mbedtls_to_psa_error(
+        mbedtls_rsa_export( rsa, &n, NULL, NULL, NULL, NULL ) );
+    if( status == PSA_SUCCESS )
+    {
+        if( mbedtls_mpi_bitlen( &n ) % 8 != 0 )
+            status = PSA_ERROR_NOT_SUPPORTED;
+    }
+    mbedtls_mpi_free( &n );
+    return( status );
+}
+
+psa_status_t mbedtls_psa_rsa_load_representation(
+    psa_key_type_t type, const uint8_t *data, size_t data_length,
+    mbedtls_rsa_context **p_rsa )
+{
+    psa_status_t status;
+    mbedtls_pk_context ctx;
+    size_t bits;
+    mbedtls_pk_init( &ctx );
+
+    /* Parse the data. */
+    if( PSA_KEY_TYPE_IS_KEY_PAIR( type ) )
+        status = mbedtls_to_psa_error(
+            mbedtls_pk_parse_key( &ctx, data, data_length, NULL, 0 ) );
+    else
+        status = mbedtls_to_psa_error(
+            mbedtls_pk_parse_public_key( &ctx, data, data_length ) );
+    if( status != PSA_SUCCESS )
+        goto exit;
+
+    /* We have something that the pkparse module recognizes. If it is a
+     * valid RSA key, store it. */
+    if( mbedtls_pk_get_type( &ctx ) != MBEDTLS_PK_RSA )
+    {
+        status = PSA_ERROR_INVALID_ARGUMENT;
+        goto exit;
+    }
+
+    /* The size of an RSA key doesn't have to be a multiple of 8. Mbed TLS
+     * supports non-byte-aligned key sizes, but not well. For example,
+     * mbedtls_rsa_get_len() returns the key size in bytes, not in bits. */
+    bits = PSA_BYTES_TO_BITS( mbedtls_rsa_get_len( mbedtls_pk_rsa( ctx ) ) );
+    if( bits > PSA_VENDOR_RSA_MAX_KEY_BITS )
+    {
+        status = PSA_ERROR_NOT_SUPPORTED;
+        goto exit;
+    }
+    status = psa_check_rsa_key_byte_aligned( mbedtls_pk_rsa( ctx ) );
+    if( status != PSA_SUCCESS )
+        goto exit;
+
+    /* Copy out the pointer to the RSA context, and reset the PK context
+     * such that pk_free doesn't free the RSA context we just grabbed. */
+    *p_rsa = mbedtls_pk_rsa( ctx );
+    ctx.pk_info = NULL;
+
+exit:
+    mbedtls_pk_free( &ctx );
+    return( status );
+}
+#endif /* defined(MBEDTLS_PSA_BUILTIN_ALG_RSA_PKCS1V15_CRYPT) ||
+        * defined(MBEDTLS_PSA_BUILTIN_ALG_RSA_PKCS1V15_SIGN) ||
+        * defined(MBEDTLS_PSA_BUILTIN_ALG_RSA_OAEP) ||
+        * defined(MBEDTLS_PSA_BUILTIN_ALG_RSA_PSS) ||
+        * defined(BUILTIN_KEY_TYPE_RSA_KEY_PAIR) ||
+        * defined(BUILTIN_KEY_TYPE_RSA_PUBLIC_KEY) */
+
+#if defined(BUILTIN_KEY_TYPE_RSA_KEY_PAIR) || \
+    defined(BUILTIN_KEY_TYPE_RSA_PUBLIC_KEY)
+
+static psa_status_t rsa_import_key(
+    const psa_key_attributes_t *attributes,
+    const uint8_t *data, size_t data_length,
+    uint8_t *key_buffer, size_t key_buffer_size,
+    size_t *key_buffer_length, size_t *bits )
+{
+    psa_status_t status;
+    mbedtls_rsa_context *rsa = NULL;
+
+    /* Parse input */
+    status = mbedtls_psa_rsa_load_representation( attributes->core.type,
+                                                  data,
+                                                  data_length,
+                                                  &rsa );
+    if( status != PSA_SUCCESS )
+        goto exit;
+
+    *bits = (psa_key_bits_t) PSA_BYTES_TO_BITS( mbedtls_rsa_get_len( rsa ) );
+
+    /* Re-export the data to PSA export format, such that we can store export
+     * representation in the key slot. Export representation in case of RSA is
+     * the smallest representation that's allowed as input, so a straight-up
+     * allocation of the same size as the input buffer will be large enough. */
+    status = mbedtls_psa_rsa_export_key( attributes->core.type,
+                                         rsa,
+                                         key_buffer,
+                                         key_buffer_size,
+                                         key_buffer_length );
+exit:
+    /* Always free the RSA object */
+    mbedtls_rsa_free( rsa );
+    mbedtls_free( rsa );
+
+    return( status );
+}
+
+psa_status_t mbedtls_psa_rsa_export_key( psa_key_type_t type,
+                                         mbedtls_rsa_context *rsa,
+                                         uint8_t *data,
+                                         size_t data_size,
+                                         size_t *data_length )
+{
+#if defined(MBEDTLS_PK_WRITE_C)
+    int ret;
+    mbedtls_pk_context pk;
+    uint8_t *pos = data + data_size;
+
+    mbedtls_pk_init( &pk );
+    pk.pk_info = &mbedtls_rsa_info;
+    pk.pk_ctx = rsa;
+
+    /* PSA Crypto API defines the format of an RSA key as a DER-encoded
+     * representation of the non-encrypted PKCS#1 RSAPrivateKey for a
+     * private key and of the RFC3279 RSAPublicKey for a public key. */
+    if( PSA_KEY_TYPE_IS_KEY_PAIR( type ) )
+        ret = mbedtls_pk_write_key_der( &pk, data, data_size );
+    else
+        ret = mbedtls_pk_write_pubkey( &pos, data, &pk );
+
+    if( ret < 0 )
+    {
+        /* Clean up in case pk_write failed halfway through. */
+        memset( data, 0, data_size );
+        return( mbedtls_to_psa_error( ret ) );
+    }
+
+    /* The mbedtls_pk_xxx functions write to the end of the buffer.
+     * Move the data to the beginning and erase remaining data
+     * at the original location. */
+    if( 2 * (size_t) ret <= data_size )
+    {
+        memcpy( data, data + data_size - ret, ret );
+        memset( data + data_size - ret, 0, ret );
+    }
+    else if( (size_t) ret < data_size )
+    {
+        memmove( data, data + data_size - ret, ret );
+        memset( data + ret, 0, data_size - ret );
+    }
+
+    *data_length = ret;
+    return( PSA_SUCCESS );
+#else
+    (void) type;
+    (void) rsa;
+    (void) data;
+    (void) data_size;
+    (void) data_length;
+    return( PSA_ERROR_NOT_SUPPORTED );
+#endif /* MBEDTLS_PK_WRITE_C */
+}
+
+static psa_status_t rsa_export_public_key(
+    const psa_key_attributes_t *attributes,
+    const uint8_t *key_buffer, size_t key_buffer_size,
+    uint8_t *data, size_t data_size, size_t *data_length )
+{
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+    mbedtls_rsa_context *rsa = NULL;
+
+    status = mbedtls_psa_rsa_load_representation(
+                 attributes->core.type, key_buffer, key_buffer_size, &rsa );
+    if( status != PSA_SUCCESS )
+        return( status );
+
+    status = mbedtls_psa_rsa_export_key( PSA_KEY_TYPE_RSA_PUBLIC_KEY,
+                                         rsa,
+                                         data,
+                                         data_size,
+                                         data_length );
+
+    mbedtls_rsa_free( rsa );
+    mbedtls_free( rsa );
+
+    return( status );
+}
+#endif /* defined(BUILTIN_KEY_TYPE_RSA_KEY_PAIR) ||
+        * defined(BUILTIN_KEY_TYPE_RSA_PUBLIC_KEY) */
+
+#if defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_RSA_KEY_PAIR) || \
+    defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_RSA_PUBLIC_KEY)
+
+psa_status_t mbedtls_psa_rsa_import_key(
+    const psa_key_attributes_t *attributes,
+    const uint8_t *data, size_t data_length,
+    uint8_t *key_buffer, size_t key_buffer_size,
+    size_t *key_buffer_length, size_t *bits )
+{
+    return( rsa_import_key( attributes, data, data_length,
+                            key_buffer, key_buffer_size,
+                            key_buffer_length, bits ) );
+}
+
+psa_status_t mbedtls_psa_rsa_export_public_key(
+    const psa_key_attributes_t *attributes,
+    const uint8_t *key_buffer, size_t key_buffer_size,
+    uint8_t *data, size_t data_size, size_t *data_length )
+{
+    return( rsa_export_public_key( attributes, key_buffer, key_buffer_size,
+                                   data, data_size, data_length ) );
+}
+
+#endif /* defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_RSA_KEY_PAIR) ||
+        * defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_RSA_PUBLIC_KEY) */
+
+/*
+ * BEYOND THIS POINT, TEST DRIVER ENTRY POINTS ONLY.
+ */
+
+#if defined(PSA_CRYPTO_DRIVER_TEST)
+
+#if defined(MBEDTLS_PSA_ACCEL_KEY_TYPE_RSA_KEY_PAIR) || \
+    defined(MBEDTLS_PSA_ACCEL_KEY_TYPE_RSA_PUBLIC_KEY)
+
+psa_status_t mbedtls_transparent_test_driver_rsa_import_key(
+    const psa_key_attributes_t *attributes,
+    const uint8_t *data, size_t data_length,
+    uint8_t *key_buffer, size_t key_buffer_size,
+    size_t *key_buffer_length, size_t *bits )
+{
+    return( rsa_import_key( attributes, data, data_length,
+                            key_buffer, key_buffer_size,
+                            key_buffer_length, bits ) );
+}
+
+psa_status_t mbedtls_transparent_test_driver_rsa_export_public_key(
+    const psa_key_attributes_t *attributes,
+    const uint8_t *key_buffer, size_t key_buffer_size,
+    uint8_t *data, size_t data_size, size_t *data_length )
+{
+    return( rsa_export_public_key( attributes, key_buffer, key_buffer_size,
+                                   data, data_size, data_length ) );
+}
+
+#endif /* defined(MBEDTLS_PSA_ACCEL_KEY_TYPE_RSA_KEY_PAIR) ||
+          defined(MBEDTLS_PSA_ACCEL_KEY_TYPE_RSA_PUBLIC_KEY) */
+
+#endif /* PSA_CRYPTO_DRIVER_TEST */
+
+#endif /* MBEDTLS_PSA_CRYPTO_C */
diff --git a/library/psa_crypto_rsa.h b/library/psa_crypto_rsa.h
new file mode 100644
index 0000000..64013df
--- /dev/null
+++ b/library/psa_crypto_rsa.h
@@ -0,0 +1,136 @@
+/*
+ *  PSA RSA layer on top of Mbed TLS crypto
+ */
+/*
+ *  Copyright The Mbed TLS Contributors
+ *  SPDX-License-Identifier: Apache-2.0
+ *
+ *  Licensed under the Apache License, Version 2.0 (the "License"); you may
+ *  not use this file except in compliance with the License.
+ *  You may obtain a copy of the License at
+ *
+ *  http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *  Unless required by applicable law or agreed to in writing, software
+ *  distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ *  WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *  See the License for the specific language governing permissions and
+ *  limitations under the License.
+ */
+
+#ifndef PSA_CRYPTO_RSA_H
+#define PSA_CRYPTO_RSA_H
+
+#include <psa/crypto.h>
+#include <mbedtls/rsa.h>
+
+/** Load the contents of a key buffer into an internal RSA representation
+ *
+ * \param[in] type          The type of key contained in \p data.
+ * \param[in] data          The buffer from which to load the representation.
+ * \param[in] data_length   The size in bytes of \p data.
+ * \param[out] p_rsa        Returns a pointer to an RSA context on success.
+ *                          The caller is responsible for freeing both the
+ *                          contents of the context and the context itself
+ *                          when done.
+ */
+psa_status_t mbedtls_psa_rsa_load_representation( psa_key_type_t type,
+                                                  const uint8_t *data,
+                                                  size_t data_length,
+                                                  mbedtls_rsa_context **p_rsa );
+
+/** Import an RSA key in binary format.
+ *
+ * \note The signature of this function is that of a PSA driver
+ *       import_key entry point. This function behaves as an import_key
+ *       entry point as defined in the PSA driver interface specification for
+ *       transparent drivers.
+ *
+ * \param[in]  attributes       The attributes for the key to import.
+ * \param[in]  data             The buffer containing the key data in import
+ *                              format.
+ * \param[in]  data_length      Size of the \p data buffer in bytes.
+ * \param[out] key_buffer       The buffer containing the key data in output
+ *                              format.
+ * \param[in]  key_buffer_size  Size of the \p key_buffer buffer in bytes. This
+ *                              size is greater or equal to \p data_length.
+ * \param[out] key_buffer_length  The length of the data written in \p
+ *                                key_buffer in bytes.
+ * \param[out] bits             The key size in number of bits.
+ *
+ * \retval #PSA_SUCCESS  The RSA key was imported successfully.
+ * \retval #PSA_ERROR_INVALID_ARGUMENT
+ *         The key data is not correctly formatted.
+ * \retval #PSA_ERROR_NOT_SUPPORTED
+ * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
+ * \retval #PSA_ERROR_CORRUPTION_DETECTED
+ */
+psa_status_t mbedtls_psa_rsa_import_key(
+    const psa_key_attributes_t *attributes,
+    const uint8_t *data, size_t data_length,
+    uint8_t *key_buffer, size_t key_buffer_size,
+    size_t *key_buffer_length, size_t *bits );
+
+/** Export an RSA key to export representation
+ *
+ * \param[in] type          The type of key (public/private) to export
+ * \param[in] rsa           The internal RSA representation from which to export
+ * \param[out] data         The buffer to export to
+ * \param[in] data_size     The length of the buffer to export to
+ * \param[out] data_length  The amount of bytes written to \p data
+ */
+psa_status_t mbedtls_psa_rsa_export_key( psa_key_type_t type,
+                                         mbedtls_rsa_context *rsa,
+                                         uint8_t *data,
+                                         size_t data_size,
+                                         size_t *data_length );
+
+/** Export a public RSA key or the public part of an RSA key pair in binary
+ *  format.
+ *
+ * \note The signature of this function is that of a PSA driver
+ *       export_public_key entry point. This function behaves as an
+ *       export_public_key entry point as defined in the PSA driver interface
+ *       specification.
+ *
+ * \param[in]  attributes       The attributes for the key to export.
+ * \param[in]  key_buffer       Material or context of the key to export.
+ * \param[in]  key_buffer_size  Size of the \p key_buffer buffer in bytes.
+ * \param[out] data             Buffer where the key data is to be written.
+ * \param[in]  data_size        Size of the \p data buffer in bytes.
+ * \param[out] data_length      On success, the number of bytes written in
+ *                              \p data.
+ *
+ * \retval #PSA_SUCCESS  The RSA public key was exported successfully.
+ * \retval #PSA_ERROR_NOT_SUPPORTED
+ * \retval #PSA_ERROR_COMMUNICATION_FAILURE
+ * \retval #PSA_ERROR_HARDWARE_FAILURE
+ * \retval #PSA_ERROR_CORRUPTION_DETECTED
+ * \retval #PSA_ERROR_STORAGE_FAILURE
+ * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
+ */
+psa_status_t mbedtls_psa_rsa_export_public_key(
+    const psa_key_attributes_t *attributes,
+    const uint8_t *key_buffer, size_t key_buffer_size,
+    uint8_t *data, size_t data_size, size_t *data_length );
+
+/*
+ * BEYOND THIS POINT, TEST DRIVER ENTRY POINTS ONLY.
+ */
+
+#if defined(PSA_CRYPTO_DRIVER_TEST)
+
+psa_status_t mbedtls_transparent_test_driver_rsa_import_key(
+    const psa_key_attributes_t *attributes,
+    const uint8_t *data, size_t data_length,
+    uint8_t *key_buffer, size_t key_buffer_size,
+    size_t *key_buffer_length, size_t *bits );
+
+psa_status_t mbedtls_transparent_test_driver_rsa_export_public_key(
+    const psa_key_attributes_t *attributes,
+    const uint8_t *key_buffer, size_t key_buffer_size,
+    uint8_t *data, size_t data_size, size_t *data_length );
+
+#endif /* PSA_CRYPTO_DRIVER_TEST */
+
+#endif /* PSA_CRYPTO_RSA_H */
diff --git a/library/psa_crypto_se.h b/library/psa_crypto_se.h
index 67fadf8..7104261 100644
--- a/library/psa_crypto_se.h
+++ b/library/psa_crypto_se.h
@@ -155,6 +155,13 @@
  *
  * \param driver        The driver table entry containing the persistent
  *                      data to load from storage.
+ *
+ * \return #PSA_SUCCESS
+ * \return #PSA_ERROR_NOT_SUPPORTED
+ * \return #PSA_ERROR_DOES_NOT_EXIST
+ * \return #PSA_ERROR_STORAGE_FAILURE
+ * \return #PSA_ERROR_DATA_CORRUPT
+ * \return #PSA_ERROR_INVALID_ARGUMENT
  */
 psa_status_t psa_load_se_persistent_data(
     const psa_se_drv_table_entry_t *driver );
@@ -163,6 +170,14 @@
  *
  * \param[in] driver    The driver table entry containing the persistent
  *                      data to save to storage.
+ *
+ * \return #PSA_SUCCESS
+ * \return #PSA_ERROR_NOT_SUPPORTED
+ * \return #PSA_ERROR_NOT_PERMITTED
+ * \return #PSA_ERROR_NOT_SUPPORTED
+ * \return #PSA_ERROR_INSUFFICIENT_STORAGE
+ * \return #PSA_ERROR_STORAGE_FAILURE
+ * \return #PSA_ERROR_INVALID_ARGUMENT
  */
 psa_status_t psa_save_se_persistent_data(
     const psa_se_drv_table_entry_t *driver );
diff --git a/library/psa_crypto_slot_management.c b/library/psa_crypto_slot_management.c
index 39d6dbb..6dca0ef 100644
--- a/library/psa_crypto_slot_management.c
+++ b/library/psa_crypto_slot_management.c
@@ -255,16 +255,15 @@
     if( psa_get_se_driver( slot->attr.lifetime, &drv, &drv_context ) )
     {
         psa_se_key_data_storage_t *data;
+
         if( key_data_length != sizeof( *data ) )
         {
-            status = PSA_ERROR_STORAGE_FAILURE;
+            status = PSA_ERROR_DATA_INVALID;
             goto exit;
         }
         data = (psa_se_key_data_storage_t *) key_data;
-        memcpy( &slot->data.se.slot_number, &data->slot_number,
-                sizeof( slot->data.se.slot_number ) );
-
-        status = PSA_SUCCESS;
+        status = psa_copy_key_material_into_slot(
+                     slot, data->slot_number, sizeof( data->slot_number ) );
         goto exit;
     }
 #endif /* MBEDTLS_PSA_CRYPTO_SE_C */
diff --git a/library/psa_crypto_storage.c b/library/psa_crypto_storage.c
index 10a1ad3..773d3aa 100644
--- a/library/psa_crypto_storage.c
+++ b/library/psa_crypto_storage.c
@@ -91,6 +91,8 @@
  * \param data_size         Size of the \c data buffer in bytes.
  *
  * \retval #PSA_SUCCESS
+ * \retval #PSA_ERROR_DATA_INVALID
+ * \retval #PSA_ERROR_DATA_CORRUPT
  * \retval #PSA_ERROR_STORAGE_FAILURE
  * \retval #PSA_ERROR_DOES_NOT_EXIST
  */
@@ -108,7 +110,7 @@
 
     status = psa_its_get( data_identifier, 0, (uint32_t) data_size, data, &data_length );
     if( data_size  != data_length )
-        return( PSA_ERROR_STORAGE_FAILURE );
+        return( PSA_ERROR_DATA_INVALID );
 
     return( status );
 }
@@ -139,8 +141,9 @@
  *
  * \retval #PSA_SUCCESS
  * \retval #PSA_ERROR_INSUFFICIENT_STORAGE
- * \retval #PSA_ERROR_STORAGE_FAILURE
  * \retval #PSA_ERROR_ALREADY_EXISTS
+ * \retval #PSA_ERROR_STORAGE_FAILURE
+ * \retval #PSA_ERROR_DATA_INVALID
  */
 static psa_status_t psa_crypto_storage_store( const mbedtls_svc_key_id_t key,
                                               const uint8_t *data,
@@ -156,7 +159,7 @@
     status = psa_its_set( data_identifier, (uint32_t) data_length, data, 0 );
     if( status != PSA_SUCCESS )
     {
-        return( PSA_ERROR_STORAGE_FAILURE );
+        return( PSA_ERROR_DATA_INVALID );
     }
 
     status = psa_its_get_info( data_identifier, &data_identifier_info );
@@ -167,7 +170,7 @@
 
     if( data_identifier_info.size != data_length )
     {
-        status = PSA_ERROR_STORAGE_FAILURE;
+        status = PSA_ERROR_DATA_INVALID;
         goto exit;
     }
 
@@ -194,11 +197,11 @@
         return( PSA_SUCCESS );
 
     if( psa_its_remove( data_identifier ) != PSA_SUCCESS )
-        return( PSA_ERROR_STORAGE_FAILURE );
+        return( PSA_ERROR_DATA_INVALID );
 
     ret = psa_its_get_info( data_identifier, &data_identifier_info );
     if( ret != PSA_ERROR_DOES_NOT_EXIST )
-        return( PSA_ERROR_STORAGE_FAILURE );
+        return( PSA_ERROR_DATA_INVALID );
 
     return( PSA_SUCCESS );
 }
@@ -212,6 +215,8 @@
  *
  * \retval #PSA_SUCCESS
  * \retval #PSA_ERROR_STORAGE_FAILURE
+ * \retval #PSA_ERROR_DOES_NOT_EXIST
+ * \retval #PSA_ERROR_DATA_CORRUPT
  */
 static psa_status_t psa_crypto_storage_get_data_length(
     const mbedtls_svc_key_id_t key,
@@ -313,7 +318,7 @@
 {
     if( memcmp( data, PSA_KEY_STORAGE_MAGIC_HEADER,
                 PSA_KEY_STORAGE_MAGIC_HEADER_LENGTH ) != 0 )
-        return( PSA_ERROR_STORAGE_FAILURE );
+        return( PSA_ERROR_DATA_INVALID );
     return( PSA_SUCCESS );
 }
 
@@ -329,7 +334,7 @@
     uint32_t version;
 
     if( storage_data_length < sizeof(*storage_format) )
-        return( PSA_ERROR_STORAGE_FAILURE );
+        return( PSA_ERROR_DATA_INVALID );
 
     status = check_magic_header( storage_data );
     if( status != PSA_SUCCESS )
@@ -337,12 +342,12 @@
 
     GET_UINT32_LE( version, storage_format->version, 0 );
     if( version != 0 )
-        return( PSA_ERROR_STORAGE_FAILURE );
+        return( PSA_ERROR_DATA_INVALID );
 
     GET_UINT32_LE( *key_data_length, storage_format->data_len, 0 );
     if( *key_data_length > ( storage_data_length - sizeof(*storage_format) ) ||
         *key_data_length > PSA_CRYPTO_MAX_STORAGE_SIZE )
-        return( PSA_ERROR_STORAGE_FAILURE );
+        return( PSA_ERROR_DATA_INVALID );
 
     if( *key_data_length == 0 )
     {
@@ -479,7 +484,7 @@
     if( status != PSA_SUCCESS )
         return( status );
     if( length != sizeof( psa_crypto_transaction ) )
-        return( PSA_ERROR_STORAGE_FAILURE );
+        return( PSA_ERROR_DATA_INVALID );
     return( PSA_SUCCESS );
 }
 
diff --git a/library/psa_crypto_storage.h b/library/psa_crypto_storage.h
index 1d597c0..8461691 100644
--- a/library/psa_crypto_storage.h
+++ b/library/psa_crypto_storage.h
@@ -102,6 +102,8 @@
  * \retval #PSA_ERROR_INSUFFICIENT_STORAGE
  * \retval #PSA_ERROR_STORAGE_FAILURE
  * \retval #PSA_ERROR_ALREADY_EXISTS
+ * \retval #PSA_ERROR_DATA_INVALID
+ * \retval #PSA_ERROR_DATA_CORRUPT
  */
 psa_status_t psa_save_persistent_key( const psa_core_key_attributes_t *attr,
                                       const uint8_t *data,
@@ -129,7 +131,8 @@
  *
  * \retval #PSA_SUCCESS
  * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
- * \retval #PSA_ERROR_STORAGE_FAILURE
+ * \retval #PSA_ERROR_DATA_INVALID
+ * \retval #PSA_ERROR_DATA_CORRUPT
  * \retval #PSA_ERROR_DOES_NOT_EXIST
  */
 psa_status_t psa_load_persistent_key( psa_core_key_attributes_t *attr,
@@ -145,7 +148,7 @@
  * \retval #PSA_SUCCESS
  *         The key was successfully removed,
  *         or the key did not exist.
- * \retval #PSA_ERROR_STORAGE_FAILURE
+ * \retval #PSA_ERROR_DATA_INVALID
  */
 psa_status_t psa_destroy_persistent_key( const mbedtls_svc_key_id_t key );
 
@@ -188,9 +191,8 @@
  *                             with the loaded key metadata.
  *
  * \retval #PSA_SUCCESS
- * \retval #PSA_ERROR_INSUFFICIENT_STORAGE
  * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
- * \retval #PSA_ERROR_STORAGE_FAILURE
+ * \retval #PSA_ERROR_DATA_INVALID
  */
 psa_status_t psa_parse_key_data_from_storage( const uint8_t *storage_data,
                                               size_t storage_data_length,
@@ -324,6 +326,7 @@
  * atomically update the transaction state.
  *
  * \retval #PSA_SUCCESS
+ * \retval #PSA_ERROR_DATA_CORRUPT
  * \retval #PSA_ERROR_INSUFFICIENT_STORAGE
  * \retval #PSA_ERROR_STORAGE_FAILURE
  */
@@ -340,6 +343,8 @@
  * \retval #PSA_ERROR_DOES_NOT_EXIST
  *         There is no ongoing transaction.
  * \retval #PSA_ERROR_STORAGE_FAILURE
+ * \retval #PSA_ERROR_DATA_INVALID
+ * \retval #PSA_ERROR_DATA_CORRUPT
  */
 psa_status_t psa_crypto_load_transaction( void );
 
diff --git a/programs/psa/psa_constant_names_generated.c b/programs/psa/psa_constant_names_generated.c
index a6eee21..ba7a625 100644
--- a/programs/psa/psa_constant_names_generated.c
+++ b/programs/psa/psa_constant_names_generated.c
@@ -8,6 +8,8 @@
     case PSA_ERROR_BUFFER_TOO_SMALL: return "PSA_ERROR_BUFFER_TOO_SMALL";
     case PSA_ERROR_COMMUNICATION_FAILURE: return "PSA_ERROR_COMMUNICATION_FAILURE";
     case PSA_ERROR_CORRUPTION_DETECTED: return "PSA_ERROR_CORRUPTION_DETECTED";
+    case PSA_ERROR_DATA_CORRUPT: return "PSA_ERROR_DATA_CORRUPT";
+    case PSA_ERROR_DATA_INVALID: return "PSA_ERROR_DATA_INVALID";
     case PSA_ERROR_DOES_NOT_EXIST: return "PSA_ERROR_DOES_NOT_EXIST";
     case PSA_ERROR_GENERIC_ERROR: return "PSA_ERROR_GENERIC_ERROR";
     case PSA_ERROR_HARDWARE_FAILURE: return "PSA_ERROR_HARDWARE_FAILURE";
diff --git a/tests/include/test/drivers/key_management.h b/tests/include/test/drivers/key_management.h
index 90f8c58..b30baa2 100644
--- a/tests/include/test/drivers/key_management.h
+++ b/tests/include/test/drivers/key_management.h
@@ -41,7 +41,7 @@
     unsigned long hits;
 } test_driver_key_management_hooks_t;
 
-#define TEST_DRIVER_KEY_MANAGEMENT_INIT { NULL, 0, PSA_ERROR_NOT_SUPPORTED, 0 }
+#define TEST_DRIVER_KEY_MANAGEMENT_INIT { NULL, 0, PSA_SUCCESS, 0 }
 static inline test_driver_key_management_hooks_t test_driver_key_management_hooks_init( void )
 {
     const test_driver_key_management_hooks_t v = TEST_DRIVER_KEY_MANAGEMENT_INIT;
@@ -58,11 +58,10 @@
     const psa_key_attributes_t *attributes,
     uint8_t *key, size_t key_size, size_t *key_length );
 
-psa_status_t test_transparent_validate_key(
+psa_status_t test_opaque_export_key(
     const psa_key_attributes_t *attributes,
-    const uint8_t *data,
-    size_t data_length,
-    size_t *bits);
+    const uint8_t *key, size_t key_length,
+    uint8_t *data, size_t data_size, size_t *data_length );
 
 psa_status_t test_transparent_export_public_key(
     const psa_key_attributes_t *attributes,
@@ -74,5 +73,14 @@
     const uint8_t *key, size_t key_length,
     uint8_t *data, size_t data_size, size_t *data_length );
 
+psa_status_t test_transparent_import_key(
+    const psa_key_attributes_t *attributes,
+    const uint8_t *data,
+    size_t data_length,
+    uint8_t *key_buffer,
+    size_t key_buffer_size,
+    size_t *key_buffer_length,
+    size_t *bits);
+
 #endif /* PSA_CRYPTO_DRIVER_TEST */
 #endif /* PSA_CRYPTO_TEST_DRIVERS_KEY_MANAGEMENT_H */
diff --git a/tests/include/test/helpers.h b/tests/include/test/helpers.h
index ce8a1e2..8b56d3a 100644
--- a/tests/include/test/helpers.h
+++ b/tests/include/test/helpers.h
@@ -49,10 +49,76 @@
 #include <stddef.h>
 #include <stdint.h>
 
+typedef enum
+{
+    MBEDTLS_TEST_RESULT_SUCCESS = 0,
+    MBEDTLS_TEST_RESULT_FAILED,
+    MBEDTLS_TEST_RESULT_SKIPPED
+} mbedtls_test_result_t;
+
+typedef struct
+{
+    mbedtls_test_result_t result;
+    const char *test;
+    const char *filename;
+    int line_no;
+    unsigned long step;
+}
+mbedtls_test_info_t;
+extern mbedtls_test_info_t mbedtls_test_info;
+
 int mbedtls_test_platform_setup( void );
 void mbedtls_test_platform_teardown( void );
 
 /**
+ * \brief           Record the current test case as a failure.
+ *
+ *                  This function can be called directly however it is usually
+ *                  called via macros such as TEST_ASSERT, TEST_EQUAL,
+ *                  PSA_ASSERT, etc...
+ *
+ * \note            If the test case was already marked as failed, calling
+ *                  `mbedtls_test_fail( )` again will not overwrite any
+ *                  previous information about the failure.
+ *
+ * \param test      Description of the failure or assertion that failed. This
+ *                  MUST be a string literal.
+ * \param line_no   Line number where the failure originated.
+ * \param filename  Filename where the failure originated.
+ */
+void mbedtls_test_fail( const char *test, int line_no, const char* filename );
+
+/**
+ * \brief           Record the current test case as skipped.
+ *
+ *                  This function can be called directly however it is usually
+ *                  called via the TEST_ASSUME macro.
+ *
+ * \param test      Description of the assumption that caused the test case to
+ *                  be skipped. This MUST be a string literal.
+ * \param line_no   Line number where the test case was skipped.
+ * \param filename  Filename where the test case was skipped.
+ */
+void mbedtls_test_skip( const char *test, int line_no, const char* filename );
+
+/**
+ * \brief       Set the test step number for failure reports.
+ *
+ *              Call this function to display "step NNN" in addition to the
+ *              line number and file name if a test fails. Typically the "step
+ *              number" is the index of a for loop but it can be whatever you
+ *              want.
+ *
+ * \param step  The step number to report.
+ */
+void mbedtls_test_set_step( unsigned long step );
+
+/**
+ * \brief       Reset mbedtls_test_info to a ready/starting state.
+ */
+void mbedtls_test_info_reset( void );
+
+/**
  * \brief          This function decodes the hexadecimal representation of
  *                 data.
  *
diff --git a/tests/scripts/all.sh b/tests/scripts/all.sh
index 676c804..a4c7a74 100755
--- a/tests/scripts/all.sh
+++ b/tests/scripts/all.sh
@@ -1360,7 +1360,7 @@
     scripts/config.py set MBEDTLS_PSA_CRYPTO_DRIVERS
     scripts/config.py unset MBEDTLS_USE_PSA_CRYPTO
     # Need to define the correct symbol and include the test driver header path in order to build with the test driver
-    make CC=gcc CFLAGS="$ASAN_CFLAGS -DPSA_CRYPTO_DRIVER_TEST -I../tests/include -O2" LDFLAGS="$ASAN_CFLAGS"
+    make CC=gcc CFLAGS="$ASAN_CFLAGS -DPSA_CRYPTO_DRIVER_TEST -DMBEDTLS_PSA_ACCEL_KEY_TYPE_ECC_KEY_PAIR -I../tests/include -O2" LDFLAGS="$ASAN_CFLAGS"
 
     msg "test: full + MBEDTLS_PSA_CRYPTO_CONFIG"
     make test
@@ -2128,7 +2128,7 @@
     msg "build: MBEDTLS_PSA_CRYPTO_DRIVERS w/ driver hooks"
     scripts/config.py set MBEDTLS_PSA_CRYPTO_DRIVERS
     # Need to define the correct symbol and include the test driver header path in order to build with the test driver
-    make CC=gcc CFLAGS="$ASAN_CFLAGS -DPSA_CRYPTO_DRIVER_TEST -I../tests/include -O2" LDFLAGS="$ASAN_CFLAGS"
+    make CC=gcc CFLAGS="$ASAN_CFLAGS -DPSA_CRYPTO_DRIVER_TEST -DMBEDTLS_PSA_ACCEL_KEY_TYPE_RSA_KEY_PAIR -DMBEDTLS_PSA_ACCEL_KEY_TYPE_ECC_KEY_PAIR -I../tests/include -O2" LDFLAGS="$ASAN_CFLAGS"
 
     msg "test: MBEDTLS_PSA_CRYPTO_DRIVERS, signature"
     make test
diff --git a/tests/src/drivers/key_management.c b/tests/src/drivers/key_management.c
index 00d2b45..98990d1 100644
--- a/tests/src/drivers/key_management.c
+++ b/tests/src/drivers/key_management.c
@@ -27,6 +27,8 @@
 #if defined(MBEDTLS_PSA_CRYPTO_DRIVERS) && defined(PSA_CRYPTO_DRIVER_TEST)
 #include "psa/crypto.h"
 #include "psa_crypto_core.h"
+#include "psa_crypto_ecp.h"
+#include "psa_crypto_rsa.h"
 #include "mbedtls/ecp.h"
 #include "mbedtls/error.h"
 
@@ -137,121 +139,78 @@
     return( PSA_ERROR_NOT_SUPPORTED );
 }
 
-psa_status_t test_transparent_validate_key(
+psa_status_t test_transparent_import_key(
     const psa_key_attributes_t *attributes,
     const uint8_t *data,
     size_t data_length,
-    size_t *bits )
+    uint8_t *key_buffer,
+    size_t key_buffer_size,
+    size_t *key_buffer_length,
+    size_t *bits)
 {
     ++test_driver_key_management_hooks.hits;
 
     if( test_driver_key_management_hooks.forced_status != PSA_SUCCESS )
         return( test_driver_key_management_hooks.forced_status );
 
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+    psa_key_type_t type = psa_get_key_type( attributes );
+
 #if defined(MBEDTLS_PSA_ACCEL_KEY_TYPE_ECC_KEY_PAIR) || \
     defined(MBEDTLS_PSA_ACCEL_KEY_TYPE_ECC_PUBLIC_KEY)
-    psa_key_type_t type = psa_get_key_type( attributes );
-    if ( PSA_KEY_TYPE_IS_ECC( type ) )
+    if( PSA_KEY_TYPE_IS_ECC( type ) )
     {
-        // Code mostly copied from psa_load_ecp_representation
-        psa_ecc_family_t curve = PSA_KEY_TYPE_ECC_GET_FAMILY( type );
-        mbedtls_ecp_group_id grp_id;
-        mbedtls_ecp_keypair ecp;
-        psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
-
-        if( psa_get_key_bits( attributes ) == 0 )
-        {
-            // Attempt auto-detect of curve bit size
-            size_t curve_size = data_length;
-
-            if( PSA_KEY_TYPE_IS_PUBLIC_KEY( type ) &&
-                PSA_KEY_TYPE_ECC_GET_FAMILY( type ) != PSA_ECC_FAMILY_MONTGOMERY )
-            {
-                /* A Weierstrass public key is represented as:
-                 * - The byte 0x04;
-                 * - `x_P` as a `ceiling(m/8)`-byte string, big-endian;
-                 * - `y_P` as a `ceiling(m/8)`-byte string, big-endian.
-                 * So its data length is 2m+1 where m is the curve size in bits.
-                 */
-                if( ( data_length & 1 ) == 0 )
-                    return( PSA_ERROR_INVALID_ARGUMENT );
-                curve_size = data_length / 2;
-
-                /* Montgomery public keys are represented in compressed format, meaning
-                 * their curve_size is equal to the amount of input. */
-
-                /* Private keys are represented in uncompressed private random integer
-                 * format, meaning their curve_size is equal to the amount of input. */
-            }
-
-            grp_id = mbedtls_ecc_group_of_psa( curve, curve_size );
-        }
-        else
-        {
-            grp_id = mbedtls_ecc_group_of_psa( curve,
-                PSA_BITS_TO_BYTES( psa_get_key_bits( attributes ) ) );
-        }
-
-        const mbedtls_ecp_curve_info *curve_info =
-            mbedtls_ecp_curve_info_from_grp_id( grp_id );
-
-        if( attributes->domain_parameters_size != 0 )
-            return( PSA_ERROR_NOT_SUPPORTED );
-        if( grp_id == MBEDTLS_ECP_DP_NONE || curve_info == NULL )
-            return( PSA_ERROR_NOT_SUPPORTED );
-
-        *bits = curve_info->bit_size;
-
-        mbedtls_ecp_keypair_init( &ecp );
-
-        status = mbedtls_to_psa_error(
-                    mbedtls_ecp_group_load( &ecp.grp, grp_id ) );
-        if( status != PSA_SUCCESS )
-            goto ecp_exit;
-
-        /* Load the key material. */
-        if( PSA_KEY_TYPE_IS_PUBLIC_KEY( type ) )
-        {
-            /* Load the public value. */
-            status = mbedtls_to_psa_error(
-                mbedtls_ecp_point_read_binary( &ecp.grp, &ecp.Q,
-                                               data,
-                                               data_length ) );
-            if( status != PSA_SUCCESS )
-                goto ecp_exit;
-
-            /* Check that the point is on the curve. */
-            status = mbedtls_to_psa_error(
-                mbedtls_ecp_check_pubkey( &ecp.grp, &ecp.Q ) );
-        }
-        else
-        {
-            /* Load and validate the secret value. */
-            status = mbedtls_to_psa_error(
-                mbedtls_ecp_read_key( ecp.grp.id,
-                                      &ecp,
-                                      data,
-                                      data_length ) );
-        }
-
-ecp_exit:
-        mbedtls_ecp_keypair_free( &ecp );
-        return( status );
+        status = mbedtls_transparent_test_driver_ecp_import_key(
+                     attributes,
+                     data, data_length,
+                     key_buffer, key_buffer_size,
+                     key_buffer_length, bits );
     }
-    return( PSA_ERROR_NOT_SUPPORTED );
-#else
+    else
+#endif
+#if defined(MBEDTLS_PSA_ACCEL_KEY_TYPE_RSA_KEY_PAIR) || \
+    defined(MBEDTLS_PSA_ACCEL_KEY_TYPE_RSA_PUBLIC_KEY)
+    if( PSA_KEY_TYPE_IS_RSA( type ) )
+    {
+        status = mbedtls_transparent_test_driver_rsa_import_key(
+                     attributes,
+                     data, data_length,
+                     key_buffer, key_buffer_size,
+                     key_buffer_length, bits );
+    }
+    else
+#endif
+    {
+        status = PSA_ERROR_NOT_SUPPORTED;
+        (void)data;
+        (void)data_length;
+        (void)key_buffer;
+        (void)key_buffer_size;
+        (void)key_buffer_length;
+        (void)bits;
+        (void)type;
+    }
+
+    return( status );
+}
+
+psa_status_t test_opaque_export_key(
+    const psa_key_attributes_t *attributes,
+    const uint8_t *key, size_t key_length,
+    uint8_t *data, size_t data_size, size_t *data_length )
+{
     (void) attributes;
+    (void) key;
+    (void) key_length;
     (void) data;
+    (void) data_size;
     (void) data_length;
-    (void) bits;
     return( PSA_ERROR_NOT_SUPPORTED );
-#endif /* MBEDTLS_PSA_ACCEL_KEY_TYPE_ECC_KEY_PAIR ||
-        * MBEDTLS_PSA_ACCEL_KEY_TYPE_ECC_PUBLIC_KEY */
 }
 
 psa_status_t test_transparent_export_public_key(
     const psa_key_attributes_t *attributes,
-    const uint8_t *key, size_t key_length,
+    const uint8_t *key_buffer, size_t key_buffer_size,
     uint8_t *data, size_t data_size, size_t *data_length )
 {
     ++test_driver_key_management_hooks.hits;
@@ -269,73 +228,39 @@
         return( PSA_SUCCESS );
     }
 
-    if( key == NULL || key_length == 0 )
-        return( PSA_ERROR_INVALID_ARGUMENT );
-
-    psa_key_type_t keytype = psa_get_key_type( attributes );
-    (void) keytype;
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+    psa_key_type_t key_type = psa_get_key_type( attributes );
 
 #if defined(MBEDTLS_PSA_ACCEL_KEY_TYPE_ECC_KEY_PAIR) || \
     defined(MBEDTLS_PSA_ACCEL_KEY_TYPE_ECC_PUBLIC_KEY)
-    if( PSA_KEY_TYPE_IS_ECC( keytype ) )
+    if( PSA_KEY_TYPE_IS_ECC( key_type ) )
     {
-        if( !PSA_KEY_TYPE_IS_KEY_PAIR( keytype ) )
-            return( PSA_ERROR_INVALID_ARGUMENT );
-
-        /* Mostly copied from psa_crypto.c */
-        mbedtls_ecp_group_id grp_id = MBEDTLS_ECP_DP_NONE;
-        psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
-        mbedtls_ecp_keypair ecp;
-        mbedtls_test_rnd_pseudo_info rnd_info;
-        memset( &rnd_info, 0x5A, sizeof( mbedtls_test_rnd_pseudo_info ) );
-
-        if( attributes->domain_parameters_size != 0 )
-            return( PSA_ERROR_NOT_SUPPORTED );
-
-        grp_id = mbedtls_ecc_group_of_psa( PSA_KEY_TYPE_ECC_GET_FAMILY( keytype ),
-                                           PSA_BITS_TO_BYTES( psa_get_key_bits( attributes ) ) );
-        if( grp_id == MBEDTLS_ECP_DP_NONE )
-            return( PSA_ERROR_NOT_SUPPORTED );
-
-        mbedtls_ecp_keypair_init( &ecp );
-
-        status = mbedtls_to_psa_error(
-                    mbedtls_ecp_group_load( &ecp.grp, grp_id ) );
-        if( status != PSA_SUCCESS )
-            goto ecp_exit;
-
-        status = mbedtls_to_psa_error(
-            mbedtls_ecp_read_key( ecp.grp.id,
-                                  &ecp,
-                                  key,
-                                  key_length ) );
-        if( status != PSA_SUCCESS )
-            goto ecp_exit;
-
-        /* Calculate the public key */
-        status = mbedtls_to_psa_error(
-            mbedtls_ecp_mul( &ecp.grp, &ecp.Q, &ecp.d, &ecp.grp.G,
-                             &mbedtls_test_rnd_pseudo_rand,
-                             &rnd_info ) );
-        if( status != PSA_SUCCESS )
-            goto ecp_exit;
-
-        status = mbedtls_to_psa_error(
-                    mbedtls_ecp_point_write_binary( &ecp.grp, &ecp.Q,
-                                                    MBEDTLS_ECP_PF_UNCOMPRESSED,
-                                                    data_length,
-                                                    data,
-                                                    data_size ) );
-        if( status != PSA_SUCCESS )
-            memset( data, 0, data_size );
-ecp_exit:
-        mbedtls_ecp_keypair_free( &ecp );
-        return( status );
+        status = mbedtls_transparent_test_driver_ecp_export_public_key(
+                      attributes,
+                      key_buffer, key_buffer_size,
+                      data, data_size, data_length );
     }
-#endif /* MBEDTLS_PSA_ACCEL_KEY_TYPE_ECC_KEY_PAIR ||
-        * MBEDTLS_PSA_ACCEL_KEY_TYPE_ECC_PUBLIC_KEY */
+    else
+#endif
+#if defined(MBEDTLS_PSA_ACCEL_KEY_TYPE_RSA_KEY_PAIR) || \
+    defined(MBEDTLS_PSA_ACCEL_KEY_TYPE_RSA_PUBLIC_KEY)
+    if( PSA_KEY_TYPE_IS_RSA( key_type ) )
+    {
+        status = mbedtls_transparent_test_driver_rsa_export_public_key(
+                      attributes,
+                      key_buffer, key_buffer_size,
+                      data, data_size, data_length );
+    }
+    else
+#endif
+    {
+        status = PSA_ERROR_NOT_SUPPORTED;
+        (void)key_buffer;
+        (void)key_buffer_size;
+        (void)key_type;
+    }
 
-    return( PSA_ERROR_NOT_SUPPORTED );
+    return( status );
 }
 
 psa_status_t test_opaque_export_public_key(
diff --git a/tests/src/helpers.c b/tests/src/helpers.c
index a18f1d4..e323275 100644
--- a/tests/src/helpers.c
+++ b/tests/src/helpers.c
@@ -44,6 +44,8 @@
 static mbedtls_platform_context platform_ctx;
 #endif
 
+mbedtls_test_info_t mbedtls_test_info;
+
 /*----------------------------------------------------------------------------*/
 /* Helper Functions */
 
@@ -77,6 +79,42 @@
     return( 0 );
 }
 
+void mbedtls_test_fail( const char *test, int line_no, const char* filename )
+{
+    if( mbedtls_test_info.result == MBEDTLS_TEST_RESULT_FAILED )
+    {
+        /* We've already recorded the test as having failed. Don't
+         * overwrite any previous information about the failure. */
+        return;
+    }
+    mbedtls_test_info.result = MBEDTLS_TEST_RESULT_FAILED;
+    mbedtls_test_info.test = test;
+    mbedtls_test_info.line_no = line_no;
+    mbedtls_test_info.filename = filename;
+}
+
+void mbedtls_test_skip( const char *test, int line_no, const char* filename )
+{
+    mbedtls_test_info.result = MBEDTLS_TEST_RESULT_SKIPPED;
+    mbedtls_test_info.test = test;
+    mbedtls_test_info.line_no = line_no;
+    mbedtls_test_info.filename = filename;
+}
+
+void mbedtls_test_set_step( unsigned long step )
+{
+    mbedtls_test_info.step = step;
+}
+
+void mbedtls_test_info_reset( void )
+{
+    mbedtls_test_info.result = MBEDTLS_TEST_RESULT_SUCCESS;
+    mbedtls_test_info.step = (unsigned long)( -1 );
+    mbedtls_test_info.test = 0;
+    mbedtls_test_info.line_no = 0;
+    mbedtls_test_info.filename = 0;
+}
+
 int mbedtls_test_unhexify( unsigned char *obuf,
                            size_t obufmax,
                            const char *ibuf,
diff --git a/tests/suites/helpers.function b/tests/suites/helpers.function
index 1dc6721..9762d41 100644
--- a/tests/suites/helpers.function
+++ b/tests/suites/helpers.function
@@ -108,7 +108,7 @@
     do {                                                    \
        if( ! (TEST) )                                       \
        {                                                    \
-          test_fail( #TEST, __LINE__, __FILE__ );           \
+          mbedtls_test_fail( #TEST, __LINE__, __FILE__ );   \
           goto exit;                                        \
        }                                                    \
     } while( 0 )
@@ -201,13 +201,13 @@
  *
  * \param   TEST    The test expression to be tested.
  */
-#define TEST_ASSUME( TEST )                         \
-    do {                                            \
-        if( ! (TEST) )                              \
-        {                                           \
-            test_skip( #TEST, __LINE__, __FILE__ ); \
-            goto exit;                              \
-        }                                           \
+#define TEST_ASSUME( TEST )                                 \
+    do {                                                    \
+        if( ! (TEST) )                                      \
+        {                                                   \
+            mbedtls_test_skip( #TEST, __LINE__, __FILE__ ); \
+            goto exit;                                      \
+        }                                                   \
     } while( 0 )
 
 #if defined(MBEDTLS_CHECK_PARAMS) && !defined(MBEDTLS_PARAM_FAILED_ALT)
@@ -237,7 +237,7 @@
         if( ( ( TEST ) != ( PARAM_ERR_VALUE ) ) ||                      \
             ( mbedtls_test_param_failed_check_expected_call( ) != 0 ) ) \
         {                                                               \
-            test_fail( #TEST, __LINE__, __FILE__ );                     \
+            mbedtls_test_fail( #TEST, __LINE__, __FILE__ );             \
             goto exit;                                                  \
         }                                                               \
         mbedtls_test_param_failed_check_expected_call( );               \
@@ -270,7 +270,7 @@
         if( setjmp(  mbedtls_test_param_failed_get_state_buf( ) ) == 0 ) \
         {                                                                \
             TEST;                                                        \
-            test_fail( #TEST, __LINE__, __FILE__ );                      \
+            mbedtls_test_fail( #TEST, __LINE__, __FILE__ );              \
             goto exit;                                                   \
         }                                                                \
         mbedtls_test_param_failed_reset_state( );                        \
@@ -346,24 +346,6 @@
 /*----------------------------------------------------------------------------*/
 /* Global variables */
 
-typedef enum
-{
-    TEST_RESULT_SUCCESS = 0,
-    TEST_RESULT_FAILED,
-    TEST_RESULT_SKIPPED
-} test_result_t;
-
-typedef struct
-{
-    test_result_t result;
-    const char *test;
-    const char *filename;
-    int line_no;
-    unsigned long step;
-}
-test_info_t;
-static test_info_t test_info;
-
 #if defined(MBEDTLS_CHECK_PARAMS)
 jmp_buf jmp_tmp;
 #endif
@@ -386,41 +368,6 @@
 /*----------------------------------------------------------------------------*/
 /* Helper Functions */
 
-/** Set the test step number for failure reports.
- *
- * Call this function to display "step NNN" in addition to the line number
- * and file name if a test fails. Typically the "step number" is the index
- * of a for loop but it can be whatever you want.
- *
- * \param step  The step number to report.
- */
-void test_set_step( unsigned long step )
-{
-    test_info.step = step;
-}
-
-void test_fail( const char *test, int line_no, const char* filename )
-{
-    if( test_info.result == TEST_RESULT_FAILED )
-    {
-        /* We've already recorded the test as having failed. Don't
-         * overwrite any previous information about the failure. */
-        return;
-    }
-    test_info.result = TEST_RESULT_FAILED;
-    test_info.test = test;
-    test_info.line_no = line_no;
-    test_info.filename = filename;
-}
-
-void test_skip( const char *test, int line_no, const char* filename )
-{
-    test_info.result = TEST_RESULT_SKIPPED;
-    test_info.test = test;
-    test_info.line_no = line_no;
-    test_info.filename = filename;
-}
-
 #if defined(MBEDTLS_PSA_CRYPTO_C)
 /** Check that no PSA Crypto key slots are in use.
  *
@@ -435,7 +382,7 @@
         return 0;
     else
     {
-        test_fail( msg, line_no, filename );
+        mbedtls_test_fail( msg, line_no, filename );
         return 1;
     }
 }
diff --git a/tests/suites/host_test.function b/tests/suites/host_test.function
index 872a3a4..3138c33 100644
--- a/tests/suites/host_test.function
+++ b/tests/suites/host_test.function
@@ -428,15 +428,15 @@
  * \param unmet_dependencies         The array of unmet dependencies.
  * \param missing_unmet_dependencies Non-zero if there was a problem tracking
  *                                   all unmet dependencies, 0 otherwise.
- * \param ret           The test dispatch status (DISPATCH_xxx).
- * \param test_info     A pointer to the test info structure.
+ * \param ret                   The test dispatch status (DISPATCH_xxx).
+ * \param mbedtls_test_info     A pointer to the test info structure.
  */
 static void write_outcome_result( FILE *outcome_file,
                                   size_t unmet_dep_count,
                                   int unmet_dependencies[],
                                   int missing_unmet_dependencies,
                                   int ret,
-                                  const test_info_t *info )
+                                  const mbedtls_test_info_t *info )
 {
     if( outcome_file == NULL )
         return;
@@ -462,10 +462,10 @@
             }
             switch( info->result )
             {
-                case TEST_RESULT_SUCCESS:
+                case MBEDTLS_TEST_RESULT_SUCCESS:
                     mbedtls_fprintf( outcome_file, "PASS;" );
                     break;
-                case TEST_RESULT_SKIPPED:
+                case MBEDTLS_TEST_RESULT_SKIPPED:
                     mbedtls_fprintf( outcome_file, "SKIP;Runtime skip" );
                     break;
                 default:
@@ -601,7 +601,7 @@
     }
 
     /* Initialize the struct that holds information about the last test */
-    memset( &test_info, 0, sizeof( test_info ) );
+    mbedtls_test_info_reset( );
 
     /* Now begin to execute the tests in the testfiles */
     for ( testfile_index = 0;
@@ -638,7 +638,8 @@
             if( ( ret = get_line( file, buf, sizeof(buf) ) ) != 0 )
                 break;
             mbedtls_fprintf( stdout, "%s%.66s",
-                    test_info.result == TEST_RESULT_FAILED ? "\n" : "", buf );
+                    mbedtls_test_info.result == MBEDTLS_TEST_RESULT_FAILED ?
+                    "\n" : "", buf );
             mbedtls_fprintf( stdout, " " );
             for( i = strlen( buf ) + 1; i < 67; i++ )
                 mbedtls_fprintf( stdout, "." );
@@ -682,8 +683,7 @@
             // If there are no unmet dependencies execute the test
             if( unmet_dep_count == 0 )
             {
-                test_info.result = TEST_RESULT_SUCCESS;
-                test_info.step = (unsigned long)( -1 );
+                mbedtls_test_info_reset( );
 
 #if defined(__unix__) || (defined(__APPLE__) && defined(__MACH__))
                 /* Suppress all output from the library unless we're verbose
@@ -723,7 +723,7 @@
             write_outcome_result( outcome_file,
                                   unmet_dep_count, unmet_dependencies,
                                   missing_unmet_dependencies,
-                                  ret, &test_info );
+                                  ret, &mbedtls_test_info );
             if( unmet_dep_count > 0 || ret == DISPATCH_UNSUPPORTED_SUITE )
             {
                 total_skipped++;
@@ -753,11 +753,11 @@
             }
             else if( ret == DISPATCH_TEST_SUCCESS )
             {
-                if( test_info.result == TEST_RESULT_SUCCESS )
+                if( mbedtls_test_info.result == MBEDTLS_TEST_RESULT_SUCCESS )
                 {
                     mbedtls_fprintf( stdout, "PASS\n" );
                 }
-                else if( test_info.result == TEST_RESULT_SKIPPED )
+                else if( mbedtls_test_info.result == MBEDTLS_TEST_RESULT_SKIPPED )
                 {
                     mbedtls_fprintf( stdout, "----\n" );
                     total_skipped++;
@@ -767,14 +767,15 @@
                     total_errors++;
                     mbedtls_fprintf( stdout, "FAILED\n" );
                     mbedtls_fprintf( stdout, "  %s\n  at ",
-                                     test_info.test );
-                    if( test_info.step != (unsigned long)( -1 ) )
+                                     mbedtls_test_info.test );
+                    if( mbedtls_test_info.step != (unsigned long)( -1 ) )
                     {
                         mbedtls_fprintf( stdout, "step %lu, ",
-                                         test_info.step );
+                                         mbedtls_test_info.step );
                     }
                     mbedtls_fprintf( stdout, "line %d, %s",
-                                     test_info.line_no, test_info.filename );
+                                     mbedtls_test_info.line_no,
+                                     mbedtls_test_info.filename );
                 }
                 fflush( stdout );
             }
diff --git a/tests/suites/main_test.function b/tests/suites/main_test.function
index 98dab3e..57395ae 100644
--- a/tests/suites/main_test.function
+++ b/tests/suites/main_test.function
@@ -179,10 +179,9 @@
     {
         /* Unexpected parameter validation error */
         mbedtls_test_param_failed_get_location_record( &location_record );
-        test_fail( location_record.failure_condition,
-                   location_record.line,
-                   location_record.file );
-        test_info.result = TEST_RESULT_FAILED;
+        mbedtls_test_fail( location_record.failure_condition,
+                           location_record.line,
+                           location_record.file );
     }
 
     mbedtls_test_param_failed_reset_state( );
diff --git a/tests/suites/target_test.function b/tests/suites/target_test.function
index 8354b96..637a79d 100644
--- a/tests/suites/target_test.function
+++ b/tests/suites/target_test.function
@@ -384,8 +384,7 @@
     while ( 1 )
     {
         ret = 0;
-        test_info.result = TEST_RESULT_SUCCESS;
-        test_info.step = (unsigned long)( -1 );
+        mbedtls_test_info_reset( );
         data_len = 0;
 
         data = receive_data( &data_len );
@@ -443,7 +442,7 @@
         if ( ret )
             send_failure( ret );
         else
-            send_status( test_info.result );
+            send_status( mbedtls_test_info.result );
     }
     return( 0 );
 }
diff --git a/tests/suites/test_suite_asn1parse.function b/tests/suites/test_suite_asn1parse.function
index 990f343..47a4340 100644
--- a/tests/suites/test_suite_asn1parse.function
+++ b/tests/suites/test_suite_asn1parse.function
@@ -130,7 +130,7 @@
     size_t parsed_length;
     int ret;
 
-    test_set_step( buffer_size );
+    mbedtls_test_set_step( buffer_size );
     /* Allocate a new buffer of exactly the length to parse each time.
      * This gives memory sanitizers a chance to catch buffer overreads. */
     if( buffer_size == 0 )
@@ -198,7 +198,7 @@
 
     TEST_ASSERT( content > state->input_start );
     offset = content - state->input_start;
-    test_set_step( offset );
+    mbedtls_test_set_step( offset );
 
     if( *rest == 0 )
         return( RET_TRAVERSE_STOP );
@@ -252,7 +252,7 @@
      */
     for( buffer_size = 1; buffer_size <= input->len + 1; buffer_size++ )
     {
-        test_set_step( buffer_size );
+        mbedtls_test_set_step( buffer_size );
         /* Allocate a new buffer of exactly the length to parse each time.
          * This gives memory sanitizers a chance to catch buffer overreads. */
         ASSERT_ALLOC( buf, buffer_size );
@@ -594,6 +594,7 @@
     unsigned char *p = input->x;
     const char *rest = description;
     unsigned long n;
+    unsigned int step = 0;
 
     TEST_EQUAL( mbedtls_asn1_get_sequence_of( &p, input->x + input->len,
                                               &head, tag ),
@@ -614,7 +615,7 @@
             cur = &head;
             while( *rest )
             {
-                ++test_info.step;
+                mbedtls_test_set_step( step );
                 TEST_ASSERT( cur != NULL );
                 TEST_EQUAL( cur->buf.tag, tag );
                 n = strtoul( rest, (char **) &rest, 0 );
@@ -625,6 +626,7 @@
                 if( *rest )
                     ++rest;
                 cur = cur->next;
+                ++step;
             }
             TEST_ASSERT( cur == NULL );
         }
diff --git a/tests/suites/test_suite_asn1write.function b/tests/suites/test_suite_asn1write.function
index 21465c7..8824739 100644
--- a/tests/suites/test_suite_asn1write.function
+++ b/tests/suites/test_suite_asn1write.function
@@ -15,7 +15,7 @@
 
 int generic_write_start_step( generic_write_data_t *data )
 {
-    test_set_step( data->size );
+    mbedtls_test_set_step( data->size );
     ASSERT_ALLOC( data->output, data->size == 0 ? 1 : data->size );
     data->end = data->output + data->size;
     data->p = data->end;
diff --git a/tests/suites/test_suite_psa_crypto.data b/tests/suites/test_suite_psa_crypto.data
index 99e2e8c..6af4120 100644
--- a/tests/suites/test_suite_psa_crypto.data
+++ b/tests/suites/test_suite_psa_crypto.data
@@ -244,7 +244,7 @@
 
 PSA import/export-public: cannot export-public a symmetric key
 depends_on:MBEDTLS_PK_PARSE_C:MBEDTLS_PK_WRITE_C
-import_export_public_key:"2b7e151628aed2a6abf7158809cf4f3c":PSA_KEY_TYPE_AES:PSA_ALG_CBC_NO_PADDING:0:PSA_ERROR_INVALID_ARGUMENT:""
+import_export_public_key:"2b7e151628aed2a6abf7158809cf4f3c":PSA_KEY_TYPE_AES:PSA_ALG_CBC_NO_PADDING:0:PSA_ERROR_INVALID_ARGUMENT:"2b7e151628aed2a6abf7158809cf4f3c"
 
 PSA import/export EC secp256r1 public key: good
 depends_on:PSA_WANT_ALG_ECDSA:PSA_WANT_KEY_TYPE_ECC_PUBLIC_KEY:MBEDTLS_PK_PARSE_C:MBEDTLS_PK_WRITE_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED
@@ -357,11 +357,9 @@
 import_with_data:"aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa":PSA_KEY_TYPE_AES:128:PSA_ERROR_INVALID_ARGUMENT
 
 PSA import large key: raw, 65528 bits (ok)
-depends_on:HAVE_RAM_AVAILABLE_128K
 import_large_key:PSA_KEY_TYPE_RAW_DATA:8191:PSA_SUCCESS
 
 PSA import large key: raw, 65536 bits (not supported)
-depends_on:HAVE_RAM_AVAILABLE_128K
 import_large_key:PSA_KEY_TYPE_RAW_DATA:8192:PSA_ERROR_NOT_SUPPORTED
 
 PSA import RSA key pair: maximum size exceeded
@@ -2376,11 +2374,11 @@
 derive_output:PSA_ALG_TLS12_PSK_TO_MS(PSA_ALG_SHA_384):PSA_KEY_DERIVATION_INPUT_SEED:"5bed47716a11a49a6268a8350b085929116ad9ccc8181f09a05b07a7741576d65bed47718dfd82f2d3f57544afe52decae6819b970dc716ada72ae0dd3072e9a":PSA_KEY_DERIVATION_INPUT_SECRET:"01020304":PSA_KEY_DERIVATION_INPUT_LABEL:"6d617374657220736563726574":48:"f5a61fbdd2ec415762abb8042a6c16645a53d2edb6dec8c8":"5ca71689301f9f4d875128c87608b75250b20a9550e4fe18"
 
 PSA key derivation: HKDF SHA-256, request maximum capacity
-depends_on:PSA_WANT_ALG_HKDF:PSA_WANT_ALG_SHA_256:HAVE_RAM_AVAILABLE_128K
+depends_on:PSA_WANT_ALG_HKDF:PSA_WANT_ALG_SHA_256
 derive_output:PSA_ALG_HKDF(PSA_ALG_SHA_256):PSA_KEY_DERIVATION_INPUT_SALT:"000102030405060708090a0b0c":PSA_KEY_DERIVATION_INPUT_SECRET:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_KEY_DERIVATION_INPUT_INFO:"f0f1f2f3f4f5f6f7f8f9":255 * 32:"3cb25f25faacd57a90434f64d0362f2a2d2d0a90cf1a5a4c5db02d56ecc4c5bf34007208d5b887185865":""
 
 PSA key derivation: HKDF SHA-1, request maximum capacity
-depends_on:PSA_WANT_ALG_HKDF:PSA_WANT_ALG_SHA_1:HAVE_RAM_AVAILABLE_128K
+depends_on:PSA_WANT_ALG_HKDF:PSA_WANT_ALG_SHA_1
 derive_output:PSA_ALG_HKDF(PSA_ALG_SHA_1):PSA_KEY_DERIVATION_INPUT_SALT:"":PSA_KEY_DERIVATION_INPUT_SECRET:"0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c0c":PSA_KEY_DERIVATION_INPUT_INFO:"":255 * 20:"2c91117204d745f3500d636a62f64f0ab3bae548aa53d423b0d1f27ebba6f5e5673a081d70cce7acfc48":""
 
 PSA key derivation: HKDF SHA-256, request too much capacity
@@ -2408,19 +2406,19 @@
 derive_output:PSA_ALG_HKDF(PSA_ALG_SHA_256):PSA_KEY_DERIVATION_INPUT_SALT:"000102030405060708090a0b0c":PSA_KEY_DERIVATION_INPUT_SECRET:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_KEY_DERIVATION_INPUT_INFO:"f0f1f2f3f4f5f6f7f8f9":42:"3cb25f25faacd57a90434f64d0362f2a2d2d0a90cf1a5a4c5db02d56ecc4c5bf34007208d5b887185865ff":"ff"
 
 PSA key derivation: HKDF SHA-256, read maximum capacity minus 1
-depends_on:PSA_WANT_ALG_HKDF:PSA_WANT_ALG_SHA_256:HAVE_RAM_AVAILABLE_128K
+depends_on:PSA_WANT_ALG_HKDF:PSA_WANT_ALG_SHA_256
 derive_full:PSA_ALG_HKDF(PSA_ALG_SHA_256):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":255 * 32 - 1
 
 PSA key derivation: HKDF SHA-256, read maximum capacity
-depends_on:PSA_WANT_ALG_HKDF:PSA_WANT_ALG_SHA_256:HAVE_RAM_AVAILABLE_128K
+depends_on:PSA_WANT_ALG_HKDF:PSA_WANT_ALG_SHA_256
 derive_full:PSA_ALG_HKDF(PSA_ALG_SHA_256):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":255 * 32
 
 PSA key derivation: TLS 1.2 PRF SHA-256, read maximum capacity minus 1
-depends_on:PSA_WANT_ALG_SHA_256:PSA_WANT_ALG_TLS12_PRF:HAVE_RAM_AVAILABLE_128K
+depends_on:PSA_WANT_ALG_SHA_256:PSA_WANT_ALG_TLS12_PRF
 derive_full:PSA_ALG_TLS12_PRF(PSA_ALG_SHA_256):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":255 * 32 - 1
 
 PSA key derivation: TLS 1.2 PRF SHA-256, read maximum capacity
-depends_on:PSA_WANT_ALG_SHA_256:PSA_WANT_ALG_TLS12_PRF:HAVE_RAM_AVAILABLE_128K
+depends_on:PSA_WANT_ALG_SHA_256:PSA_WANT_ALG_TLS12_PRF
 derive_full:PSA_ALG_TLS12_PRF(PSA_ALG_SHA_256):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":255 * 32
 
 PSA key derivation: HKDF SHA-256, exercise AES128-CTR
@@ -2493,39 +2491,39 @@
 
 PSA key derivation: invalid type (0)
 depends_on:PSA_WANT_ALG_HKDF:PSA_WANT_ALG_SHA_256
-derive_key:PSA_ALG_HKDF(PSA_ALG_SHA_256):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":PSA_KEY_TYPE_NONE:128:PSA_ERROR_INVALID_ARGUMENT
+derive_key:PSA_ALG_HKDF(PSA_ALG_SHA_256):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":PSA_KEY_TYPE_NONE:128:PSA_ERROR_INVALID_ARGUMENT:0
 
 PSA key derivation: invalid type (PSA_KEY_TYPE_CATEGORY_MASK)
 depends_on:PSA_WANT_ALG_HKDF:PSA_WANT_ALG_SHA_256
-derive_key:PSA_ALG_HKDF(PSA_ALG_SHA_256):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":PSA_KEY_TYPE_CATEGORY_MASK:128:PSA_ERROR_INVALID_ARGUMENT
+derive_key:PSA_ALG_HKDF(PSA_ALG_SHA_256):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":PSA_KEY_TYPE_CATEGORY_MASK:128:PSA_ERROR_INVALID_ARGUMENT:0
 
 PSA key derivation: invalid length (0)
 depends_on:PSA_WANT_ALG_HKDF:PSA_WANT_ALG_SHA_256
 # The spec allows either INVALID_ARGUMENT or NOT_SUPPORTED
-derive_key:PSA_ALG_HKDF(PSA_ALG_SHA_256):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":PSA_KEY_TYPE_RAW_DATA:0:PSA_ERROR_INVALID_ARGUMENT
+derive_key:PSA_ALG_HKDF(PSA_ALG_SHA_256):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":PSA_KEY_TYPE_RAW_DATA:0:PSA_ERROR_INVALID_ARGUMENT:0
 
 PSA key derivation: invalid length (7 bits)
 depends_on:PSA_WANT_ALG_HKDF:PSA_WANT_ALG_SHA_256
-derive_key:PSA_ALG_HKDF(PSA_ALG_SHA_256):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":PSA_KEY_TYPE_RAW_DATA:7:PSA_ERROR_INVALID_ARGUMENT
+derive_key:PSA_ALG_HKDF(PSA_ALG_SHA_256):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":PSA_KEY_TYPE_RAW_DATA:7:PSA_ERROR_INVALID_ARGUMENT:0
 
 PSA key derivation: raw data, 8 bits
 depends_on:PSA_WANT_ALG_HKDF:PSA_WANT_ALG_SHA_256
-derive_key:PSA_ALG_HKDF(PSA_ALG_SHA_256):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":PSA_KEY_TYPE_RAW_DATA:8:PSA_SUCCESS
+derive_key:PSA_ALG_HKDF(PSA_ALG_SHA_256):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":PSA_KEY_TYPE_RAW_DATA:8:PSA_SUCCESS:0
 
 PSA key derivation: invalid length (9 bits)
 depends_on:PSA_WANT_ALG_HKDF:PSA_WANT_ALG_SHA_256
-derive_key:PSA_ALG_HKDF(PSA_ALG_SHA_256):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":PSA_KEY_TYPE_RAW_DATA:9:PSA_ERROR_INVALID_ARGUMENT
+derive_key:PSA_ALG_HKDF(PSA_ALG_SHA_256):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":PSA_KEY_TYPE_RAW_DATA:9:PSA_ERROR_INVALID_ARGUMENT:0
 
 # This test assumes that PSA_MAX_KEY_BITS (currently 65536-8 bits = 8191 bytes
 # and not expected to be raised any time soon) is less than the maximum
 # output from HKDF-SHA512 (255*64 = 16320 bytes).
 PSA key derivation: largest possible key
-depends_on:PSA_WANT_ALG_HKDF:PSA_WANT_ALG_SHA_512:HAVE_RAM_AVAILABLE_128K
-derive_key:PSA_ALG_HKDF(PSA_ALG_SHA_512):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":PSA_KEY_TYPE_RAW_DATA:PSA_MAX_KEY_BITS:PSA_SUCCESS
+depends_on:PSA_WANT_ALG_HKDF:PSA_WANT_ALG_SHA_512
+derive_key:PSA_ALG_HKDF(PSA_ALG_SHA_512):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":PSA_KEY_TYPE_RAW_DATA:PSA_MAX_KEY_BITS:PSA_SUCCESS:1
 
 PSA key derivation: key too large
 depends_on:MBEDTLS_SHA512_C
-derive_key:PSA_ALG_HKDF(PSA_ALG_SHA_512):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":PSA_KEY_TYPE_RAW_DATA:PSA_MAX_KEY_BITS + 1:PSA_ERROR_NOT_SUPPORTED
+derive_key:PSA_ALG_HKDF(PSA_ALG_SHA_512):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":PSA_KEY_TYPE_RAW_DATA:PSA_MAX_KEY_BITS + 1:PSA_ERROR_NOT_SUPPORTED:0
 
 PSA key agreement setup: ECDH + HKDF-SHA-256: good
 depends_on:PSA_WANT_ALG_ECDH:PSA_WANT_ALG_HKDF:PSA_WANT_ALG_SHA_256:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:MBEDTLS_PK_PARSE_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED
@@ -2651,114 +2649,113 @@
 generate_random:2 * MBEDTLS_CTR_DRBG_MAX_REQUEST + 1
 
 PSA generate key: bad type (0)
-generate_key:PSA_KEY_TYPE_NONE:128:PSA_KEY_USAGE_EXPORT:0:PSA_ERROR_NOT_SUPPORTED
+generate_key:PSA_KEY_TYPE_NONE:128:PSA_KEY_USAGE_EXPORT:0:PSA_ERROR_NOT_SUPPORTED:0
 
 PSA generate key: bad type (PSA_KEY_TYPE_CATEGORY_MASK)
-generate_key:PSA_KEY_TYPE_CATEGORY_MASK:128:PSA_KEY_USAGE_EXPORT:0:PSA_ERROR_NOT_SUPPORTED
+generate_key:PSA_KEY_TYPE_CATEGORY_MASK:128:PSA_KEY_USAGE_EXPORT:0:PSA_ERROR_NOT_SUPPORTED:0
 
 PSA generate key: bad type (RSA public key)
-generate_key:PSA_KEY_TYPE_RSA_PUBLIC_KEY:512:PSA_KEY_USAGE_EXPORT:0:PSA_ERROR_NOT_SUPPORTED
+generate_key:PSA_KEY_TYPE_RSA_PUBLIC_KEY:512:PSA_KEY_USAGE_EXPORT:0:PSA_ERROR_NOT_SUPPORTED:0
 
 PSA generate key: raw data, 0 bits: invalid argument
 # The spec allows either INVALID_ARGUMENT or NOT_SUPPORTED
-generate_key:PSA_KEY_TYPE_RAW_DATA:0:PSA_KEY_USAGE_EXPORT:0:PSA_ERROR_INVALID_ARGUMENT
+generate_key:PSA_KEY_TYPE_RAW_DATA:0:PSA_KEY_USAGE_EXPORT:0:PSA_ERROR_INVALID_ARGUMENT:0
 
 PSA generate key: raw data, 7 bits: invalid argument
-generate_key:PSA_KEY_TYPE_RAW_DATA:7:PSA_KEY_USAGE_EXPORT:0:PSA_ERROR_INVALID_ARGUMENT
+generate_key:PSA_KEY_TYPE_RAW_DATA:7:PSA_KEY_USAGE_EXPORT:0:PSA_ERROR_INVALID_ARGUMENT:0
 
 PSA generate key: raw data, 8 bits
-generate_key:PSA_KEY_TYPE_RAW_DATA:8:PSA_KEY_USAGE_EXPORT:0:PSA_SUCCESS
+generate_key:PSA_KEY_TYPE_RAW_DATA:8:PSA_KEY_USAGE_EXPORT:0:PSA_SUCCESS:0
 
 PSA generate key: raw data, 9 bits: invalid argument
-generate_key:PSA_KEY_TYPE_RAW_DATA:9:PSA_KEY_USAGE_EXPORT:0:PSA_ERROR_INVALID_ARGUMENT
+generate_key:PSA_KEY_TYPE_RAW_DATA:9:PSA_KEY_USAGE_EXPORT:0:PSA_ERROR_INVALID_ARGUMENT:0
 
 PSA generate key: raw data, (MBEDTLS_CTR_DRBG_MAX_REQUEST + 1) * 8 bits
-generate_key:PSA_KEY_TYPE_RAW_DATA:(MBEDTLS_CTR_DRBG_MAX_REQUEST + 1) * 8:PSA_KEY_USAGE_EXPORT:0:PSA_SUCCESS
+generate_key:PSA_KEY_TYPE_RAW_DATA:(MBEDTLS_CTR_DRBG_MAX_REQUEST + 1) * 8:PSA_KEY_USAGE_EXPORT:0:PSA_SUCCESS:0
 
 PSA generate key: raw data, (2 * MBEDTLS_CTR_DRBG_MAX_REQUEST + 1) * 8 bits
-generate_key:PSA_KEY_TYPE_RAW_DATA:(2 * MBEDTLS_CTR_DRBG_MAX_REQUEST + 1) * 8:PSA_KEY_USAGE_EXPORT:0:PSA_SUCCESS
+generate_key:PSA_KEY_TYPE_RAW_DATA:(2 * MBEDTLS_CTR_DRBG_MAX_REQUEST + 1) * 8:PSA_KEY_USAGE_EXPORT:0:PSA_SUCCESS:0
 
-PSA generate key: raw data, 65528 bits (ok)
-depends_on:HAVE_RAM_AVAILABLE_128K
-generate_key:PSA_KEY_TYPE_RAW_DATA:65528:PSA_KEY_USAGE_EXPORT:0:PSA_SUCCESS
+PSA generate key: raw data, 65528 bits (large key, ok if it fits)
+generate_key:PSA_KEY_TYPE_RAW_DATA:65528:PSA_KEY_USAGE_EXPORT:0:PSA_SUCCESS:1
 
 PSA generate key: raw data, 65536 bits (not supported)
-generate_key:PSA_KEY_TYPE_RAW_DATA:65536:PSA_KEY_USAGE_EXPORT:0:PSA_ERROR_NOT_SUPPORTED
+generate_key:PSA_KEY_TYPE_RAW_DATA:65536:PSA_KEY_USAGE_EXPORT:0:PSA_ERROR_NOT_SUPPORTED:0
 
 PSA generate key: AES, 128 bits, CTR
 depends_on:MBEDTLS_AES_C:MBEDTLS_CIPHER_MODE_CTR
-generate_key:PSA_KEY_TYPE_AES:128:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT:PSA_ALG_CTR:PSA_SUCCESS
+generate_key:PSA_KEY_TYPE_AES:128:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT:PSA_ALG_CTR:PSA_SUCCESS:0
 
 PSA generate key: AES, 128 bits, GCM
 depends_on:MBEDTLS_AES_C:MBEDTLS_GCM_C
-generate_key:PSA_KEY_TYPE_AES:128:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT:PSA_ALG_GCM:PSA_SUCCESS
+generate_key:PSA_KEY_TYPE_AES:128:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT:PSA_ALG_GCM:PSA_SUCCESS:0
 
 PSA generate key: DES, 64 bits, CBC-nopad
 depends_on:MBEDTLS_DES_C:MBEDTLS_CIPHER_MODE_CBC
-generate_key:PSA_KEY_TYPE_DES:64:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT:PSA_ALG_CBC_NO_PADDING:PSA_SUCCESS
+generate_key:PSA_KEY_TYPE_DES:64:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT:PSA_ALG_CBC_NO_PADDING:PSA_SUCCESS:0
 
 PSA generate key: DES, 128 bits, CBC-nopad
 depends_on:MBEDTLS_DES_C:MBEDTLS_CIPHER_MODE_CBC
-generate_key:PSA_KEY_TYPE_DES:128:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT:PSA_ALG_CBC_NO_PADDING:PSA_SUCCESS
+generate_key:PSA_KEY_TYPE_DES:128:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT:PSA_ALG_CBC_NO_PADDING:PSA_SUCCESS:0
 
 PSA generate key: DES, 192 bits, CBC-nopad
 depends_on:MBEDTLS_DES_C:MBEDTLS_CIPHER_MODE_CBC
-generate_key:PSA_KEY_TYPE_DES:192:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT:PSA_ALG_CBC_NO_PADDING:PSA_SUCCESS
+generate_key:PSA_KEY_TYPE_DES:192:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT:PSA_ALG_CBC_NO_PADDING:PSA_SUCCESS:0
 
 PSA generate key: invalid key size: AES, 64 bits
 depends_on:MBEDTLS_AES_C
-generate_key:PSA_KEY_TYPE_AES:64:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT:PSA_ALG_CTR:PSA_ERROR_INVALID_ARGUMENT
+generate_key:PSA_KEY_TYPE_AES:64:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT:PSA_ALG_CTR:PSA_ERROR_INVALID_ARGUMENT:0
 
 PSA generate key: RSA, 512 bits, good, sign (PKCS#1 v1.5)
 depends_on:PSA_WANT_ALG_RSA_PKCS1V15_SIGN:PSA_WANT_KEY_TYPE_RSA_KEY_PAIR:MBEDTLS_GENPRIME
-generate_key:PSA_KEY_TYPE_RSA_KEY_PAIR:512:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_SIGN_HASH | PSA_KEY_USAGE_VERIFY_HASH:PSA_ALG_RSA_PKCS1V15_SIGN_RAW:PSA_SUCCESS
+generate_key:PSA_KEY_TYPE_RSA_KEY_PAIR:512:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_SIGN_HASH | PSA_KEY_USAGE_VERIFY_HASH:PSA_ALG_RSA_PKCS1V15_SIGN_RAW:PSA_SUCCESS:0
 
 PSA generate key: RSA, 1016 bits, good, sign (PKCS#1 v1.5)
 depends_on:PSA_WANT_ALG_RSA_PKCS1V15_SIGN:PSA_WANT_KEY_TYPE_RSA_KEY_PAIR:MBEDTLS_GENPRIME:MBEDTLS_MD_C
-generate_key:PSA_KEY_TYPE_RSA_KEY_PAIR:1016:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_SIGN_HASH | PSA_KEY_USAGE_VERIFY_HASH:PSA_ALG_RSA_PKCS1V15_SIGN_RAW:PSA_SUCCESS
+generate_key:PSA_KEY_TYPE_RSA_KEY_PAIR:1016:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_SIGN_HASH | PSA_KEY_USAGE_VERIFY_HASH:PSA_ALG_RSA_PKCS1V15_SIGN_RAW:PSA_SUCCESS:0
 
 PSA generate key: RSA, 1024 bits, good, sign (PSS SHA-256)
 depends_on:PSA_WANT_ALG_RSA_PSS:PSA_WANT_ALG_SHA_256:PSA_WANT_KEY_TYPE_RSA_KEY_PAIR:MBEDTLS_GENPRIME:MBEDTLS_MD_C
-generate_key:PSA_KEY_TYPE_RSA_KEY_PAIR:1024:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_SIGN_HASH | PSA_KEY_USAGE_VERIFY_HASH:PSA_ALG_RSA_PSS(PSA_ALG_SHA_256):PSA_SUCCESS
+generate_key:PSA_KEY_TYPE_RSA_KEY_PAIR:1024:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_SIGN_HASH | PSA_KEY_USAGE_VERIFY_HASH:PSA_ALG_RSA_PSS(PSA_ALG_SHA_256):PSA_SUCCESS:0
 
 PSA generate key: RSA, 512 bits, good, encrypt (PKCS#1 v1.5)
 depends_on:PSA_WANT_ALG_RSA_PKCS1V15_CRYPT:PSA_WANT_KEY_TYPE_RSA_KEY_PAIR:MBEDTLS_GENPRIME
-generate_key:PSA_KEY_TYPE_RSA_KEY_PAIR:512:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT:PSA_ALG_RSA_PKCS1V15_CRYPT:PSA_SUCCESS
+generate_key:PSA_KEY_TYPE_RSA_KEY_PAIR:512:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT:PSA_ALG_RSA_PKCS1V15_CRYPT:PSA_SUCCESS:0
 
 PSA generate key: RSA, 1024 bits, good, encrypt (OAEP SHA-256)
 depends_on:PSA_WANT_ALG_RSA_OAEP:PSA_WANT_ALG_SHA_256:PSA_WANT_KEY_TYPE_RSA_KEY_PAIR:MBEDTLS_MD_C
-generate_key:PSA_KEY_TYPE_RSA_KEY_PAIR:1024:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT:PSA_ALG_RSA_OAEP(PSA_ALG_SHA_256):PSA_SUCCESS
+generate_key:PSA_KEY_TYPE_RSA_KEY_PAIR:1024:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT:PSA_ALG_RSA_OAEP(PSA_ALG_SHA_256):PSA_SUCCESS:0
 
 PSA generate key: RSA, 0 bits: invalid
 depends_on:PSA_WANT_ALG_RSA_PKCS1V15_CRYPT:PSA_WANT_KEY_TYPE_RSA_KEY_PAIR:MBEDTLS_GENPRIME
 # The spec allows either INVALID_ARGUMENT or NOT_SUPPORTED
-generate_key:PSA_KEY_TYPE_RSA_KEY_PAIR:0:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT:PSA_ALG_RSA_PKCS1V15_CRYPT:PSA_ERROR_INVALID_ARGUMENT
+generate_key:PSA_KEY_TYPE_RSA_KEY_PAIR:0:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT:PSA_ALG_RSA_PKCS1V15_CRYPT:PSA_ERROR_INVALID_ARGUMENT:0
 
 PSA generate key: RSA, 1022 bits: not supported
 depends_on:MBEDTLS_RSA_C:MBEDTLS_GENPRIME
-generate_key:PSA_KEY_TYPE_RSA_KEY_PAIR:1022:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT:PSA_ALG_RSA_PKCS1V15_CRYPT:PSA_ERROR_NOT_SUPPORTED
+generate_key:PSA_KEY_TYPE_RSA_KEY_PAIR:1022:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT:PSA_ALG_RSA_PKCS1V15_CRYPT:PSA_ERROR_NOT_SUPPORTED:0
 
 PSA generate key: RSA, 1023 bits: not supported
 depends_on:MBEDTLS_RSA_C:MBEDTLS_GENPRIME
-generate_key:PSA_KEY_TYPE_RSA_KEY_PAIR:1023:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT:PSA_ALG_RSA_PKCS1V15_CRYPT:PSA_ERROR_NOT_SUPPORTED
+generate_key:PSA_KEY_TYPE_RSA_KEY_PAIR:1023:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT:PSA_ALG_RSA_PKCS1V15_CRYPT:PSA_ERROR_NOT_SUPPORTED:0
 
 PSA generate key: RSA, maximum size exceeded
 depends_on:MBEDTLS_RSA_C:MBEDTLS_GENPRIME
-generate_key:PSA_KEY_TYPE_RSA_KEY_PAIR:PSA_VENDOR_RSA_MAX_KEY_BITS+1:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT:PSA_ALG_RSA_PKCS1V15_CRYPT:PSA_ERROR_NOT_SUPPORTED
+generate_key:PSA_KEY_TYPE_RSA_KEY_PAIR:PSA_VENDOR_RSA_MAX_KEY_BITS+1:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT:PSA_ALG_RSA_PKCS1V15_CRYPT:PSA_ERROR_NOT_SUPPORTED:0
 
 PSA generate key: ECC, SECP256R1, good
 depends_on:PSA_WANT_ALG_ECDSA:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:MBEDTLS_ECP_DP_SECP256R1_ENABLED
-generate_key:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1):256:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_SIGN_HASH | PSA_KEY_USAGE_VERIFY_HASH:PSA_ALG_ECDSA_ANY:PSA_SUCCESS
+generate_key:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1):256:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_SIGN_HASH | PSA_KEY_USAGE_VERIFY_HASH:PSA_ALG_ECDSA_ANY:PSA_SUCCESS:0
 
 PSA generate key: ECC, SECP256R1, incorrect bit size
 depends_on:MBEDTLS_ECP_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED:MBEDTLS_ECDSA_C
 # INVALID_ARGUMENT would make more sense, but our code as currently structured
 # doesn't fully relate the curve with its size.
-generate_key:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1):128:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_SIGN_HASH | PSA_KEY_USAGE_VERIFY_HASH:PSA_ALG_ECDSA_ANY:PSA_ERROR_NOT_SUPPORTED
+generate_key:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1):128:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_SIGN_HASH | PSA_KEY_USAGE_VERIFY_HASH:PSA_ALG_ECDSA_ANY:PSA_ERROR_NOT_SUPPORTED:0
 
 PSA generate key: ECC, Curve25519, good
-depends_on:PSA_WANT_ALG_ECDH:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:MBEDTLS_ECP_DP_CURVE25519_ENABLED
-generate_key:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_MONTGOMERY):255:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_DERIVE:PSA_ALG_ECDH:PSA_SUCCESS
+depends_on:PSA_WANT_ALG_ECDH:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:MBEDTLS_ECP_DP_CURVE25519_ENABLED:!MBEDTLS_PSA_ACCEL_KEY_TYPE_ECC_KEY_PAIR
+generate_key:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_MONTGOMERY):255:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_DERIVE:PSA_ALG_ECDH:PSA_SUCCESS:0
 
 PSA generate key: RSA, default e
 generate_key_rsa:512:"":PSA_SUCCESS
diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index 82514d6..ef35d57 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -9,13 +9,6 @@
  * uses mbedtls_ctr_drbg internally. */
 #include "mbedtls/ctr_drbg.h"
 
-/* Tests that require more than 128kB of RAM plus change have this symbol
- * as a dependency. Currently we always define this symbol, so the tests
- * are always executed. In the future we should make this conditional
- * so that tests that require a lot of memory are skipped on constrained
- * platforms. */
-#define HAVE_RAM_AVAILABLE_128K
-
 #include "psa/crypto.h"
 #include "psa_crypto_slot_management.h"
 
@@ -539,7 +532,8 @@
         hash_alg = KNOWN_SUPPORTED_HASH_ALG;
         alg ^= PSA_ALG_ANY_HASH ^ hash_alg;
 #else
-        test_fail( "No hash algorithm for hash-and-sign testing", __LINE__, __FILE__ );
+        mbedtls_test_fail( "No hash algorithm for hash-and-sign testing",
+                            __LINE__, __FILE__ );
         return( 1 );
 #endif
     }
@@ -999,7 +993,7 @@
             mbedtls_snprintf( message, sizeof( message ),
                               "No sanity check for public key type=0x%08lx",
                               (unsigned long) type );
-            test_fail( message, __LINE__, __FILE__ );
+            mbedtls_test_fail( message, __LINE__, __FILE__ );
             (void) p;
             (void) end;
             return( 0 );
@@ -1028,19 +1022,21 @@
 
     PSA_ASSERT( psa_get_key_attributes( key, &attributes ) );
 
+    exported_size = PSA_EXPORT_KEY_OUTPUT_SIZE(
+                        psa_get_key_type( &attributes ),
+                        psa_get_key_bits( &attributes ) );
+    ASSERT_ALLOC( exported, exported_size );
+
     if( ( usage & PSA_KEY_USAGE_EXPORT ) == 0 &&
         ! PSA_KEY_TYPE_IS_PUBLIC_KEY( psa_get_key_type( &attributes ) ) )
     {
-        TEST_EQUAL( psa_export_key( key, NULL, 0, &exported_length ),
+        TEST_EQUAL( psa_export_key( key, exported,
+                                    exported_size, &exported_length ),
                     PSA_ERROR_NOT_PERMITTED );
         ok = 1;
         goto exit;
     }
 
-    exported_size = PSA_EXPORT_KEY_OUTPUT_SIZE( psa_get_key_type( &attributes ),
-                                                psa_get_key_bits( &attributes ) );
-    ASSERT_ALLOC( exported, exported_size );
-
     PSA_ASSERT( psa_export_key( key,
                                 exported, exported_size,
                                 &exported_length ) );
@@ -1071,9 +1067,16 @@
     PSA_ASSERT( psa_get_key_attributes( key, &attributes ) );
     if( ! PSA_KEY_TYPE_IS_ASYMMETRIC( psa_get_key_type( &attributes ) ) )
     {
-        TEST_EQUAL( psa_export_public_key( key, NULL, 0, &exported_length ),
+        exported_size = PSA_EXPORT_KEY_OUTPUT_SIZE(
+                            psa_get_key_type( &attributes ),
+                            psa_get_key_bits( &attributes ) );
+        ASSERT_ALLOC( exported, exported_size );
+
+        TEST_EQUAL( psa_export_public_key( key, exported,
+                                           exported_size, &exported_length ),
                     PSA_ERROR_INVALID_ARGUMENT );
-        return( 1 );
+        ok = 1;
+        goto exit;
     }
 
     public_type = PSA_KEY_TYPE_PUBLIC_KEY_OF_KEY_PAIR(
@@ -1113,8 +1116,8 @@
  * asymmetric, also check \p psa_export_public_key.
  *
  * If the key fails the tests, this function calls the test framework's
- * `test_fail` function and returns false. Otherwise this function returns
- * true. Therefore it should be used as follows:
+ * `mbedtls_test_fail` function and returns false. Otherwise this function
+ * returns true. Therefore it should be used as follows:
  * ```
  * if( ! exercise_key( ... ) ) goto exit;
  * ```
@@ -1160,7 +1163,7 @@
         mbedtls_snprintf( message, sizeof( message ),
                           "No code to exercise alg=0x%08lx",
                           (unsigned long) alg );
-        test_fail( message, __LINE__, __FILE__ );
+        mbedtls_test_fail( message, __LINE__, __FILE__ );
         ok = 0;
     }
 
@@ -1579,9 +1582,9 @@
     size_t buffer_size = byte_size + 1;
     size_t n;
 
-    /* It would be better to skip the test than fail it if the allocation
-     * fails, but the test framework doesn't support this yet. */
-    ASSERT_ALLOC( buffer, buffer_size );
+    /* Skip the test case if the target running the test cannot
+     * accomodate large keys due to heap size constraints */
+    ASSERT_ALLOC_WEAK( buffer, buffer_size );
     memset( buffer, 'K', byte_size );
 
     PSA_ASSERT( psa_crypto_init( ) );
@@ -1590,6 +1593,7 @@
     psa_set_key_usage_flags( &attributes, PSA_KEY_USAGE_EXPORT );
     psa_set_key_type( &attributes, type );
     status = psa_import_key( &attributes, buffer, byte_size, &key );
+    TEST_ASSUME( status != PSA_ERROR_INSUFFICIENT_MEMORY );
     TEST_EQUAL( status, expected_status );
 
     if( status == PSA_SUCCESS )
@@ -2674,7 +2678,7 @@
     /* Compare with corrupted value */
     for( i = 0; i < output_length; i++ )
     {
-        test_set_step( i );
+        mbedtls_test_set_step( i );
         output[i] ^= 1;
         TEST_EQUAL( psa_hash_compare( alg, input->x, input->len,
                                       output, output_length ),
@@ -3149,7 +3153,7 @@
             ( output_size >= expected_mac->len ? PSA_SUCCESS :
               PSA_ERROR_BUFFER_TOO_SMALL );
 
-        test_set_step( output_size );
+        mbedtls_test_set_step( output_size );
         ASSERT_ALLOC( actual_mac, output_size );
 
         /* Calculate the MAC. */
@@ -3235,7 +3239,7 @@
     /* Test changing one byte. */
     for( size_t i = 0; i < expected_mac->len; i++ )
     {
-        test_set_step( i );
+        mbedtls_test_set_step( i );
         perturbed_mac[i] ^= 1;
         PSA_ASSERT( psa_mac_verify_setup( &operation, key, alg ) );
         PSA_ASSERT( psa_mac_update( &operation,
@@ -5282,7 +5286,8 @@
 void derive_key( int alg_arg,
                  data_t *key_data, data_t *input1, data_t *input2,
                  int type_arg, int bits_arg,
-                 int expected_status_arg )
+                 int expected_status_arg,
+                 int is_large_output )
 {
     mbedtls_svc_key_id_t base_key = MBEDTLS_SVC_KEY_ID_INIT;
     mbedtls_svc_key_id_t derived_key = MBEDTLS_SVC_KEY_ID_INIT;
@@ -5311,9 +5316,14 @@
     psa_set_key_algorithm( &derived_attributes, 0 );
     psa_set_key_type( &derived_attributes, type );
     psa_set_key_bits( &derived_attributes, bits );
-    TEST_EQUAL( psa_key_derivation_output_key( &derived_attributes, &operation,
-                                               &derived_key ),
-                expected_status );
+
+    psa_status_t status =
+      psa_key_derivation_output_key( &derived_attributes,
+                                     &operation,
+                                     &derived_key );
+    if( is_large_output > 0 )
+      TEST_ASSUME( status != PSA_ERROR_INSUFFICIENT_MEMORY );
+    TEST_EQUAL( status, expected_status );
 
 exit:
     psa_key_derivation_abort( &operation );
@@ -5533,7 +5543,6 @@
 void generate_random( int bytes_arg )
 {
     size_t bytes = bytes_arg;
-    const unsigned char trail[] = "don't overwrite me";
     unsigned char *output = NULL;
     unsigned char *changed = NULL;
     size_t i;
@@ -5541,9 +5550,8 @@
 
     TEST_ASSERT( bytes_arg >= 0 );
 
-    ASSERT_ALLOC( output, bytes + sizeof( trail ) );
+    ASSERT_ALLOC( output, bytes );
     ASSERT_ALLOC( changed, bytes );
-    memcpy( output + bytes, trail, sizeof( trail ) );
 
     PSA_ASSERT( psa_crypto_init( ) );
 
@@ -5556,10 +5564,6 @@
             memset( output, 0, bytes );
         PSA_ASSERT( psa_generate_random( output, bytes ) );
 
-        /* Check that no more than bytes have been overwritten */
-        ASSERT_COMPARE( output + bytes, sizeof( trail ),
-                        trail, sizeof( trail ) );
-
         for( i = 0; i < bytes; i++ )
         {
             if( output[i] != 0 )
@@ -5587,7 +5591,8 @@
                    int bits_arg,
                    int usage_arg,
                    int alg_arg,
-                   int expected_status_arg )
+                   int expected_status_arg,
+                   int is_large_key )
 {
     mbedtls_svc_key_id_t key = MBEDTLS_SVC_KEY_ID_INIT;
     psa_key_type_t type = type_arg;
@@ -5606,7 +5611,11 @@
     psa_set_key_bits( &attributes, bits );
 
     /* Generate a key */
-    TEST_EQUAL( psa_generate_key( &attributes, &key ), expected_status );
+    psa_status_t status = psa_generate_key( &attributes, &key );
+
+    if( is_large_key > 0 )
+      TEST_ASSUME( status != PSA_ERROR_INSUFFICIENT_MEMORY );
+    TEST_EQUAL( status , expected_status );
     if( expected_status != PSA_SUCCESS )
         goto exit;
 
diff --git a/tests/suites/test_suite_psa_crypto_driver_wrappers.data b/tests/suites/test_suite_psa_crypto_driver_wrappers.data
index a85861c..1fd449f 100644
--- a/tests/suites/test_suite_psa_crypto_driver_wrappers.data
+++ b/tests/suites/test_suite_psa_crypto_driver_wrappers.data
@@ -35,6 +35,7 @@
 generate_key:PSA_SUCCESS:"":PSA_SUCCESS
 
 generate_key through transparent driver: fallback
+depends_on:MBEDTLS_PSA_BUILTIN_KEY_TYPE_ECC_KEY_PAIR
 generate_key:PSA_ERROR_NOT_SUPPORTED:"":PSA_SUCCESS
 
 generate_key through transparent driver: error
diff --git a/tests/suites/test_suite_psa_crypto_persistent_key.data b/tests/suites/test_suite_psa_crypto_persistent_key.data
index b25d1e0..fd71dff 100644
--- a/tests/suites/test_suite_psa_crypto_persistent_key.data
+++ b/tests/suites/test_suite_psa_crypto_persistent_key.data
@@ -11,16 +11,22 @@
 parse_storage_data_check:"505341004b45590000000000010000000024800000030000000250050000000010000000404142434445464748494a4b4c4d4e4f":"404142434445464748494a4b4c4d4e4f":PSA_KEY_LIFETIME_PERSISTENT:PSA_KEY_TYPE_AES:128:PSA_KEY_USAGE_DECRYPT | PSA_KEY_USAGE_ENCRYPT:PSA_ALG_GCM:0:PSA_SUCCESS
 
 Parse storage: wrong version
-parse_storage_data_check:"505341004b455900ffffffff0100000001700004010000000000001200000010620200003082025e02010002818100af057d396ee84fb75fdbb5c2b13c7fe5a654aa8aa2470b541ee1feb0b12d25c79711531249e1129628042dbbb6c120d1443524ef4c0e6e1d8956eeb2077af12349ddeee54483bc06c2c61948cd02b202e796aebd94d3a7cbf859c2c1819c324cb82b9cd34ede263a2abffe4733f077869e8660f7d6834da53d690ef7985f6bc3020301000102818100874bf0ffc2f2a71d14671ddd0171c954d7fdbf50281e4f6d99ea0e1ebcf82faa58e7b595ffb293d1abe17f110b37c48cc0f36c37e84d876621d327f64bbe08457d3ec4098ba2fa0a319fba411c2841ed7be83196a8cdf9daa5d00694bc335fc4c32217fe0488bce9cb7202e59468b1ead119000477db2ca797fac19eda3f58c1024100e2ab760841bb9d30a81d222de1eb7381d82214407f1b975cbbfe4e1a9467fd98adbd78f607836ca5be1928b9d160d97fd45c12d6b52e2c9871a174c66b488113024100c5ab27602159ae7d6f20c3c2ee851e46dc112e689e28d5fcbbf990a99ef8a90b8bb44fd36467e7fc1789ceb663abda338652c3c73f111774902e840565927091024100b6cdbd354f7df579a63b48b3643e353b84898777b48b15f94e0bfc0567a6ae5911d57ad6409cf7647bf96264e9bd87eb95e263b7110b9a1f9f94acced0fafa4d024071195eec37e8d257decfc672b07ae639f10cbb9b0c739d0c809968d644a94e3fd6ed9287077a14583f379058f76a8aecd43c62dc8c0f41766650d725275ac4a1024100bb32d133edc2e048d463388b7be9cb4be29f4b6250be603e70e3647501c97ddde20a4e71be95fd5e71784e25aca4baf25be5738aae59bbfe1c997781447a2b24":"3082025e02010002818100af057d396ee84fb75fdbb5c2b13c7fe5a654aa8aa2470b541ee1feb0b12d25c79711531249e1129628042dbbb6c120d1443524ef4c0e6e1d8956eeb2077af12349ddeee54483bc06c2c61948cd02b202e796aebd94d3a7cbf859c2c1819c324cb82b9cd34ede263a2abffe4733f077869e8660f7d6834da53d690ef7985f6bc3020301000102818100874bf0ffc2f2a71d14671ddd0171c954d7fdbf50281e4f6d99ea0e1ebcf82faa58e7b595ffb293d1abe17f110b37c48cc0f36c37e84d876621d327f64bbe08457d3ec4098ba2fa0a319fba411c2841ed7be83196a8cdf9daa5d00694bc335fc4c32217fe0488bce9cb7202e59468b1ead119000477db2ca797fac19eda3f58c1024100e2ab760841bb9d30a81d222de1eb7381d82214407f1b975cbbfe4e1a9467fd98adbd78f607836ca5be1928b9d160d97fd45c12d6b52e2c9871a174c66b488113024100c5ab27602159ae7d6f20c3c2ee851e46dc112e689e28d5fcbbf990a99ef8a90b8bb44fd36467e7fc1789ceb663abda338652c3c73f111774902e840565927091024100b6cdbd354f7df579a63b48b3643e353b84898777b48b15f94e0bfc0567a6ae5911d57ad6409cf7647bf96264e9bd87eb95e263b7110b9a1f9f94acced0fafa4d024071195eec37e8d257decfc672b07ae639f10cbb9b0c739d0c809968d644a94e3fd6ed9287077a14583f379058f76a8aecd43c62dc8c0f41766650d725275ac4a1024100bb32d133edc2e048d463388b7be9cb4be29f4b6250be603e70e3647501c97ddde20a4e71be95fd5e71784e25aca4baf25be5738aae59bbfe1c997781447a2b24":PSA_KEY_LIFETIME_PERSISTENT:PSA_KEY_TYPE_RSA_KEY_PAIR:1024:PSA_KEY_USAGE_EXPORT:PSA_ALG_CATEGORY_ASYMMETRIC_ENCRYPTION:PSA_ALG_CATEGORY_SIGN:PSA_ERROR_STORAGE_FAILURE
+parse_storage_data_check:"505341004b455900ffffffff0100000001700004010000000000001200000010620200003082025e02010002818100af057d396ee84fb75fdbb5c2b13c7fe5a654aa8aa2470b541ee1feb0b12d25c79711531249e1129628042dbbb6c120d1443524ef4c0e6e1d8956eeb2077af12349ddeee54483bc06c2c61948cd02b202e796aebd94d3a7cbf859c2c1819c324cb82b9cd34ede263a2abffe4733f077869e8660f7d6834da53d690ef7985f6bc3020301000102818100874bf0ffc2f2a71d14671ddd0171c954d7fdbf50281e4f6d99ea0e1ebcf82faa58e7b595ffb293d1abe17f110b37c48cc0f36c37e84d876621d327f64bbe08457d3ec4098ba2fa0a319fba411c2841ed7be83196a8cdf9daa5d00694bc335fc4c32217fe0488bce9cb7202e59468b1ead119000477db2ca797fac19eda3f58c1024100e2ab760841bb9d30a81d222de1eb7381d82214407f1b975cbbfe4e1a9467fd98adbd78f607836ca5be1928b9d160d97fd45c12d6b52e2c9871a174c66b488113024100c5ab27602159ae7d6f20c3c2ee851e46dc112e689e28d5fcbbf990a99ef8a90b8bb44fd36467e7fc1789ceb663abda338652c3c73f111774902e840565927091024100b6cdbd354f7df579a63b48b3643e353b84898777b48b15f94e0bfc0567a6ae5911d57ad6409cf7647bf96264e9bd87eb95e263b7110b9a1f9f94acced0fafa4d024071195eec37e8d257decfc672b07ae639f10cbb9b0c739d0c809968d644a94e3fd6ed9287077a14583f379058f76a8aecd43c62dc8c0f41766650d725275ac4a1024100bb32d133edc2e048d463388b7be9cb4be29f4b6250be603e70e3647501c97ddde20a4e71be95fd5e71784e25aca4baf25be5738aae59bbfe1c997781447a2b24":"3082025e02010002818100af057d396ee84fb75fdbb5c2b13c7fe5a654aa8aa2470b541ee1feb0b12d25c79711531249e1129628042dbbb6c120d1443524ef4c0e6e1d8956eeb2077af12349ddeee54483bc06c2c61948cd02b202e796aebd94d3a7cbf859c2c1819c324cb82b9cd34ede263a2abffe4733f077869e8660f7d6834da53d690ef7985f6bc3020301000102818100874bf0ffc2f2a71d14671ddd0171c954d7fdbf50281e4f6d99ea0e1ebcf82faa58e7b595ffb293d1abe17f110b37c48cc0f36c37e84d876621d327f64bbe08457d3ec4098ba2fa0a319fba411c2841ed7be83196a8cdf9daa5d00694bc335fc4c32217fe0488bce9cb7202e59468b1ead119000477db2ca797fac19eda3f58c1024100e2ab760841bb9d30a81d222de1eb7381d82214407f1b975cbbfe4e1a9467fd98adbd78f607836ca5be1928b9d160d97fd45c12d6b52e2c9871a174c66b488113024100c5ab27602159ae7d6f20c3c2ee851e46dc112e689e28d5fcbbf990a99ef8a90b8bb44fd36467e7fc1789ceb663abda338652c3c73f111774902e840565927091024100b6cdbd354f7df579a63b48b3643e353b84898777b48b15f94e0bfc0567a6ae5911d57ad6409cf7647bf96264e9bd87eb95e263b7110b9a1f9f94acced0fafa4d024071195eec37e8d257decfc672b07ae639f10cbb9b0c739d0c809968d644a94e3fd6ed9287077a14583f379058f76a8aecd43c62dc8c0f41766650d725275ac4a1024100bb32d133edc2e048d463388b7be9cb4be29f4b6250be603e70e3647501c97ddde20a4e71be95fd5e71784e25aca4baf25be5738aae59bbfe1c997781447a2b24":PSA_KEY_LIFETIME_PERSISTENT:PSA_KEY_TYPE_RSA_KEY_PAIR:1024:PSA_KEY_USAGE_EXPORT:PSA_ALG_CATEGORY_ASYMMETRIC_ENCRYPTION:PSA_ALG_CATEGORY_SIGN:PSA_ERROR_DATA_INVALID
 
 Parse storage: data too big
-parse_storage_data_check:"505341004b455900000000000100000001700000010000000000001200000010ffffffff3082025e02010002818100af057d396ee84fb75fdbb5c2b13c7fe5a654aa8aa2470b541ee1feb0b12d25c79711531249e1129628042dbbb6c120d1443524ef4c0e6e1d8956eeb2077af12349ddeee54483bc06c2c61948cd02b202e796aebd94d3a7cbf859c2c1819c324cb82b9cd34ede263a2abffe4733f077869e8660f7d6834da53d690ef7985f6bc3020301000102818100874bf0ffc2f2a71d14671ddd0171c954d7fdbf50281e4f6d99ea0e1ebcf82faa58e7b595ffb293d1abe17f110b37c48cc0f36c37e84d876621d327f64bbe08457d3ec4098ba2fa0a319fba411c2841ed7be83196a8cdf9daa5d00694bc335fc4c32217fe0488bce9cb7202e59468b1ead119000477db2ca797fac19eda3f58c1024100e2ab760841bb9d30a81d222de1eb7381d82214407f1b975cbbfe4e1a9467fd98adbd78f607836ca5be1928b9d160d97fd45c12d6b52e2c9871a174c66b488113024100c5ab27602159ae7d6f20c3c2ee851e46dc112e689e28d5fcbbf990a99ef8a90b8bb44fd36467e7fc1789ceb663abda338652c3c73f111774902e840565927091024100b6cdbd354f7df579a63b48b3643e353b84898777b48b15f94e0bfc0567a6ae5911d57ad6409cf7647bf96264e9bd87eb95e263b7110b9a1f9f94acced0fafa4d024071195eec37e8d257decfc672b07ae639f10cbb9b0c739d0c809968d644a94e3fd6ed9287077a14583f379058f76a8aecd43c62dc8c0f41766650d725275ac4a1024100bb32d133edc2e048d463388b7be9cb4be29f4b6250be603e70e3647501c97ddde20a4e71be95fd5e71784e25aca4baf25be5738aae59bbfe1c997781447a2b24":"":PSA_KEY_LIFETIME_PERSISTENT:PSA_KEY_TYPE_RSA_KEY_PAIR:0:PSA_KEY_USAGE_EXPORT:PSA_ALG_CATEGORY_ASYMMETRIC_ENCRYPTION:PSA_ALG_CATEGORY_SIGN:PSA_ERROR_STORAGE_FAILURE
+parse_storage_data_check:"505341004b455900000000000100000001700000010000000000001200000010ffffffff3082025e02010002818100af057d396ee84fb75fdbb5c2b13c7fe5a654aa8aa2470b541ee1feb0b12d25c79711531249e1129628042dbbb6c120d1443524ef4c0e6e1d8956eeb2077af12349ddeee54483bc06c2c61948cd02b202e796aebd94d3a7cbf859c2c1819c324cb82b9cd34ede263a2abffe4733f077869e8660f7d6834da53d690ef7985f6bc3020301000102818100874bf0ffc2f2a71d14671ddd0171c954d7fdbf50281e4f6d99ea0e1ebcf82faa58e7b595ffb293d1abe17f110b37c48cc0f36c37e84d876621d327f64bbe08457d3ec4098ba2fa0a319fba411c2841ed7be83196a8cdf9daa5d00694bc335fc4c32217fe0488bce9cb7202e59468b1ead119000477db2ca797fac19eda3f58c1024100e2ab760841bb9d30a81d222de1eb7381d82214407f1b975cbbfe4e1a9467fd98adbd78f607836ca5be1928b9d160d97fd45c12d6b52e2c9871a174c66b488113024100c5ab27602159ae7d6f20c3c2ee851e46dc112e689e28d5fcbbf990a99ef8a90b8bb44fd36467e7fc1789ceb663abda338652c3c73f111774902e840565927091024100b6cdbd354f7df579a63b48b3643e353b84898777b48b15f94e0bfc0567a6ae5911d57ad6409cf7647bf96264e9bd87eb95e263b7110b9a1f9f94acced0fafa4d024071195eec37e8d257decfc672b07ae639f10cbb9b0c739d0c809968d644a94e3fd6ed9287077a14583f379058f76a8aecd43c62dc8c0f41766650d725275ac4a1024100bb32d133edc2e048d463388b7be9cb4be29f4b6250be603e70e3647501c97ddde20a4e71be95fd5e71784e25aca4baf25be5738aae59bbfe1c997781447a2b24":"":PSA_KEY_LIFETIME_PERSISTENT:PSA_KEY_TYPE_RSA_KEY_PAIR:0:PSA_KEY_USAGE_EXPORT:PSA_ALG_CATEGORY_ASYMMETRIC_ENCRYPTION:PSA_ALG_CATEGORY_SIGN:PSA_ERROR_DATA_INVALID
 
 Parse storage: bad magic
-parse_storage_data_check:"645341004b455900000000000100000001700004010000000000001200000010620200003082025e02010002818100af057d396ee84fb75fdbb5c2b13c7fe5a654aa8aa2470b541ee1feb0b12d25c79711531249e1129628042dbbb6c120d1443524ef4c0e6e1d8956eeb2077af12349ddeee54483bc06c2c61948cd02b202e796aebd94d3a7cbf859c2c1819c324cb82b9cd34ede263a2abffe4733f077869e8660f7d6834da53d690ef7985f6bc3020301000102818100874bf0ffc2f2a71d14671ddd0171c954d7fdbf50281e4f6d99ea0e1ebcf82faa58e7b595ffb293d1abe17f110b37c48cc0f36c37e84d876621d327f64bbe08457d3ec4098ba2fa0a319fba411c2841ed7be83196a8cdf9daa5d00694bc335fc4c32217fe0488bce9cb7202e59468b1ead119000477db2ca797fac19eda3f58c1024100e2ab760841bb9d30a81d222de1eb7381d82214407f1b975cbbfe4e1a9467fd98adbd78f607836ca5be1928b9d160d97fd45c12d6b52e2c9871a174c66b488113024100c5ab27602159ae7d6f20c3c2ee851e46dc112e689e28d5fcbbf990a99ef8a90b8bb44fd36467e7fc1789ceb663abda338652c3c73f111774902e840565927091024100b6cdbd354f7df579a63b48b3643e353b84898777b48b15f94e0bfc0567a6ae5911d57ad6409cf7647bf96264e9bd87eb95e263b7110b9a1f9f94acced0fafa4d024071195eec37e8d257decfc672b07ae639f10cbb9b0c739d0c809968d644a94e3fd6ed9287077a14583f379058f76a8aecd43c62dc8c0f41766650d725275ac4a1024100bb32d133edc2e048d463388b7be9cb4be29f4b6250be603e70e3647501c97ddde20a4e71be95fd5e71784e25aca4baf25be5738aae59bbfe1c997781447a2b24":"3082025e02010002818100af057d396ee84fb75fdbb5c2b13c7fe5a654aa8aa2470b541ee1feb0b12d25c79711531249e1129628042dbbb6c120d1443524ef4c0e6e1d8956eeb2077af12349ddeee54483bc06c2c61948cd02b202e796aebd94d3a7cbf859c2c1819c324cb82b9cd34ede263a2abffe4733f077869e8660f7d6834da53d690ef7985f6bc3020301000102818100874bf0ffc2f2a71d14671ddd0171c954d7fdbf50281e4f6d99ea0e1ebcf82faa58e7b595ffb293d1abe17f110b37c48cc0f36c37e84d876621d327f64bbe08457d3ec4098ba2fa0a319fba411c2841ed7be83196a8cdf9daa5d00694bc335fc4c32217fe0488bce9cb7202e59468b1ead119000477db2ca797fac19eda3f58c1024100e2ab760841bb9d30a81d222de1eb7381d82214407f1b975cbbfe4e1a9467fd98adbd78f607836ca5be1928b9d160d97fd45c12d6b52e2c9871a174c66b488113024100c5ab27602159ae7d6f20c3c2ee851e46dc112e689e28d5fcbbf990a99ef8a90b8bb44fd36467e7fc1789ceb663abda338652c3c73f111774902e840565927091024100b6cdbd354f7df579a63b48b3643e353b84898777b48b15f94e0bfc0567a6ae5911d57ad6409cf7647bf96264e9bd87eb95e263b7110b9a1f9f94acced0fafa4d024071195eec37e8d257decfc672b07ae639f10cbb9b0c739d0c809968d644a94e3fd6ed9287077a14583f379058f76a8aecd43c62dc8c0f41766650d725275ac4a1024100bb32d133edc2e048d463388b7be9cb4be29f4b6250be603e70e3647501c97ddde20a4e71be95fd5e71784e25aca4baf25be5738aae59bbfe1c997781447a2b24":PSA_KEY_LIFETIME_PERSISTENT:PSA_KEY_TYPE_RSA_KEY_PAIR:1024:PSA_KEY_USAGE_EXPORT:PSA_ALG_CATEGORY_ASYMMETRIC_ENCRYPTION:PSA_ALG_CATEGORY_SIGN:PSA_ERROR_STORAGE_FAILURE
+parse_storage_data_check:"645341004b455900000000000100000001700004010000000000001200000010620200003082025e02010002818100af057d396ee84fb75fdbb5c2b13c7fe5a654aa8aa2470b541ee1feb0b12d25c79711531249e1129628042dbbb6c120d1443524ef4c0e6e1d8956eeb2077af12349ddeee54483bc06c2c61948cd02b202e796aebd94d3a7cbf859c2c1819c324cb82b9cd34ede263a2abffe4733f077869e8660f7d6834da53d690ef7985f6bc3020301000102818100874bf0ffc2f2a71d14671ddd0171c954d7fdbf50281e4f6d99ea0e1ebcf82faa58e7b595ffb293d1abe17f110b37c48cc0f36c37e84d876621d327f64bbe08457d3ec4098ba2fa0a319fba411c2841ed7be83196a8cdf9daa5d00694bc335fc4c32217fe0488bce9cb7202e59468b1ead119000477db2ca797fac19eda3f58c1024100e2ab760841bb9d30a81d222de1eb7381d82214407f1b975cbbfe4e1a9467fd98adbd78f607836ca5be1928b9d160d97fd45c12d6b52e2c9871a174c66b488113024100c5ab27602159ae7d6f20c3c2ee851e46dc112e689e28d5fcbbf990a99ef8a90b8bb44fd36467e7fc1789ceb663abda338652c3c73f111774902e840565927091024100b6cdbd354f7df579a63b48b3643e353b84898777b48b15f94e0bfc0567a6ae5911d57ad6409cf7647bf96264e9bd87eb95e263b7110b9a1f9f94acced0fafa4d024071195eec37e8d257decfc672b07ae639f10cbb9b0c739d0c809968d644a94e3fd6ed9287077a14583f379058f76a8aecd43c62dc8c0f41766650d725275ac4a1024100bb32d133edc2e048d463388b7be9cb4be29f4b6250be603e70e3647501c97ddde20a4e71be95fd5e71784e25aca4baf25be5738aae59bbfe1c997781447a2b24":"3082025e02010002818100af057d396ee84fb75fdbb5c2b13c7fe5a654aa8aa2470b541ee1feb0b12d25c79711531249e1129628042dbbb6c120d1443524ef4c0e6e1d8956eeb2077af12349ddeee54483bc06c2c61948cd02b202e796aebd94d3a7cbf859c2c1819c324cb82b9cd34ede263a2abffe4733f077869e8660f7d6834da53d690ef7985f6bc3020301000102818100874bf0ffc2f2a71d14671ddd0171c954d7fdbf50281e4f6d99ea0e1ebcf82faa58e7b595ffb293d1abe17f110b37c48cc0f36c37e84d876621d327f64bbe08457d3ec4098ba2fa0a319fba411c2841ed7be83196a8cdf9daa5d00694bc335fc4c32217fe0488bce9cb7202e59468b1ead119000477db2ca797fac19eda3f58c1024100e2ab760841bb9d30a81d222de1eb7381d82214407f1b975cbbfe4e1a9467fd98adbd78f607836ca5be1928b9d160d97fd45c12d6b52e2c9871a174c66b488113024100c5ab27602159ae7d6f20c3c2ee851e46dc112e689e28d5fcbbf990a99ef8a90b8bb44fd36467e7fc1789ceb663abda338652c3c73f111774902e840565927091024100b6cdbd354f7df579a63b48b3643e353b84898777b48b15f94e0bfc0567a6ae5911d57ad6409cf7647bf96264e9bd87eb95e263b7110b9a1f9f94acced0fafa4d024071195eec37e8d257decfc672b07ae639f10cbb9b0c739d0c809968d644a94e3fd6ed9287077a14583f379058f76a8aecd43c62dc8c0f41766650d725275ac4a1024100bb32d133edc2e048d463388b7be9cb4be29f4b6250be603e70e3647501c97ddde20a4e71be95fd5e71784e25aca4baf25be5738aae59bbfe1c997781447a2b24":PSA_KEY_LIFETIME_PERSISTENT:PSA_KEY_TYPE_RSA_KEY_PAIR:1024:PSA_KEY_USAGE_EXPORT:PSA_ALG_CATEGORY_ASYMMETRIC_ENCRYPTION:PSA_ALG_CATEGORY_SIGN:PSA_ERROR_DATA_INVALID
 
 Parse storage: truncated magic
-parse_storage_data_check:"505341004b4559":"":PSA_KEY_LIFETIME_PERSISTENT:PSA_KEY_TYPE_RSA_KEY_PAIR:0:PSA_KEY_USAGE_EXPORT:PSA_ALG_CATEGORY_ASYMMETRIC_ENCRYPTION:PSA_ALG_CATEGORY_SIGN:PSA_ERROR_STORAGE_FAILURE
+parse_storage_data_check:"505341004b4559":"":PSA_KEY_LIFETIME_PERSISTENT:PSA_KEY_TYPE_RSA_KEY_PAIR:0:PSA_KEY_USAGE_EXPORT:PSA_ALG_CATEGORY_ASYMMETRIC_ENCRYPTION:PSA_ALG_CATEGORY_SIGN:PSA_ERROR_DATA_INVALID
+
+Parse storage: truncated header
+parse_storage_data_check:"505341004b455900000000000100000001700000010000000000001200000010620200":"":PSA_KEY_LIFETIME_PERSISTENT:PSA_KEY_TYPE_RSA_KEY_PAIR:0:PSA_KEY_USAGE_EXPORT:PSA_ALG_CATEGORY_ASYMMETRIC_ENCRYPTION:PSA_ALG_CATEGORY_SIGN:PSA_ERROR_DATA_INVALID
+
+Parse storage: truncated key
+parse_storage_data_check:"505341004b455900000000000100000001700000010000000000001200000010620200003082025e02010002818100af057d396ee84fb75fdbb5c2b13c7fe5a654aa8aa2470b541ee1feb0b12d25c79711531249e1129628042dbbb6c120d1443524ef4c0e6e1d8956eeb2077af12349ddeee54483bc06c2c61948cd02b202e796aebd94d3a7cbf859c2c1819c324cb82b9cd34ede263a2abffe4733f077869e8660f7d6834da53d690ef7985f6bc3020301000102818100874bf0ffc2f2a71d14671ddd0171c954d7fdbf50281e4f6d99ea0e1ebcf82faa58e7b595ffb293d1abe17f110b37c48cc0f36c37e84d876621d327f64bbe08457d3ec4098ba2fa0a319fba411c2841ed7be83196a8cdf9daa5d00694bc335fc4c32217fe0488bce9cb7202e59468b1ead119000477db2ca797fac19eda3f58c1024100e2ab760841bb9d30a81d222de1eb7381d82214407f1b975cbbfe4e1a9467fd98adbd78f607836ca5be1928b9d160d97fd45c12d6b52e2c9871a174c66b488113024100c5ab27602159ae7d6f20c3c2ee851e46dc112e689e28d5fcbbf990a99ef8a90b8bb44fd36467e7fc1789ceb663abda338652c3c73f111774902e840565927091024100b6cdbd354f7df579a63b48b3643e353b84898777b48b15f94e0bfc0567a6ae5911d57ad6409cf7647bf96264e9bd87eb95e263b7110b9a1f9f94acced0fafa4d024071195eec37e8d257decfc672b07ae639f10cbb9b0c739d0c809968d644a94e3fd6ed9287077a14583f379058f76a8aecd43c62dc8c0f41766650d725275ac4a1024100bb32d133edc2e048d463388b7be9cb4be29f4b6250be603e70e3647501c97ddde20a4e71be95fd5e71784e25aca4baf25be5738aae59bbfe1c997781447a2b":"":PSA_KEY_LIFETIME_PERSISTENT:PSA_KEY_TYPE_RSA_KEY_PAIR:0:PSA_KEY_USAGE_EXPORT:PSA_ALG_CATEGORY_ASYMMETRIC_ENCRYPTION:PSA_ALG_CATEGORY_SIGN:PSA_ERROR_DATA_INVALID
 
 # Not specific to files, but only run this test in an environment where the maximum size could be reached.
 Save maximum-size persistent raw key
diff --git a/tests/suites/test_suite_psa_crypto_se_driver_hal.function b/tests/suites/test_suite_psa_crypto_se_driver_hal.function
index cf6ca5e..fa516c5 100644
--- a/tests/suites/test_suite_psa_crypto_se_driver_hal.function
+++ b/tests/suites/test_suite_psa_crypto_se_driver_hal.function
@@ -45,7 +45,7 @@
     do {                                                    \
        if( ! (TEST) )                                       \
        {                                                    \
-          test_fail( #TEST, __LINE__, __FILE__ );           \
+          mbedtls_test_fail( #TEST, __LINE__, __FILE__ );   \
           return( PSA_ERROR_DETECTED_BY_DRIVER );           \
        }                                                    \
     } while( 0 )
@@ -61,7 +61,7 @@
     do {                                                    \
        if( ! (TEST) )                                       \
        {                                                    \
-          test_fail( #TEST, __LINE__, __FILE__ );           \
+          mbedtls_test_fail( #TEST, __LINE__, __FILE__ );   \
           status = PSA_ERROR_DETECTED_BY_DRIVER;            \
           goto exit;                                        \
        }                                                    \
@@ -72,10 +72,10 @@
  * Run the code \p expr. If this returns \p expected_status,
  * do nothing. If this returns #PSA_ERROR_DETECTED_BY_DRIVER,
  * jump directly to the `exit` label. If this returns any other
- * status, call test_fail() then jump to `exit`.
+ * status, call mbedtls_test_fail() then jump to `exit`.
  *
  * The special case for #PSA_ERROR_DETECTED_BY_DRIVER is because in this
- * case, the test driver code is expected to have called test_fail()
+ * case, the test driver code is expected to have called mbedtls_test_fail()
  * already, so we make sure not to overwrite the failure information.
  */
 #define PSA_ASSERT_VIA_DRIVER( expr, expected_status )                  \
@@ -85,7 +85,7 @@
             goto exit;                                                  \
         if( PSA_ASSERT_VIA_DRIVER_status != ( expected_status ) )       \
         {                                                               \
-            test_fail( #expr, __LINE__, __FILE__ );                     \
+            mbedtls_test_fail( #expr, __LINE__, __FILE__ );                     \
             goto exit;                                                  \
         }                                                               \
     } while( 0 )
diff --git a/tests/suites/test_suite_psa_crypto_slot_management.function b/tests/suites/test_suite_psa_crypto_slot_management.function
index b0c660b..d14dfbb 100644
--- a/tests/suites/test_suite_psa_crypto_slot_management.function
+++ b/tests/suites/test_suite_psa_crypto_slot_management.function
@@ -343,9 +343,9 @@
                         psa_get_key_type( &read_attributes ) );
             TEST_EQUAL( psa_get_key_bits( &attributes ),
                         psa_get_key_bits( &read_attributes ) );
+            ASSERT_ALLOC( reexported, key_data->len );
             if( usage_flags & PSA_KEY_USAGE_EXPORT )
             {
-                ASSERT_ALLOC( reexported, key_data->len );
                 PSA_ASSERT( psa_export_key( id, reexported, key_data->len,
                                             &reexported_length ) );
                 ASSERT_COMPARE( key_data->x, key_data->len,
@@ -353,7 +353,8 @@
             }
             else
             {
-                TEST_EQUAL( psa_export_key( id, NULL, 0, &reexported_length ),
+                TEST_EQUAL( psa_export_key( id, reexported,
+                                key_data->len, &reexported_length ),
                             PSA_ERROR_NOT_PERMITTED );
             }
             PSA_ASSERT( psa_close_key( handle ) );
diff --git a/tests/suites/test_suite_psa_its.data b/tests/suites/test_suite_psa_its.data
index 63ca129..9057a1a 100644
--- a/tests/suites/test_suite_psa_its.data
+++ b/tests/suites/test_suite_psa_its.data
@@ -63,3 +63,9 @@
 
 Get 1 byte of 10 at -1: out of range
 get_at:0:"40414243444546474849":-1:1:PSA_ERROR_INVALID_ARGUMENT
+
+Overwrite ITS header magic
+get_fail:0:"40414243444546474849":1:0:PSA_ERROR_DATA_CORRUPT
+
+Truncate ITS header
+get_fail:0:"40414243444546474849":0:1:PSA_ERROR_DATA_CORRUPT
diff --git a/tests/suites/test_suite_psa_its.function b/tests/suites/test_suite_psa_its.function
index a7ce7b1..330846a 100644
--- a/tests/suites/test_suite_psa_its.function
+++ b/tests/suites/test_suite_psa_its.function
@@ -231,3 +231,56 @@
     cleanup( );
 }
 /* END_CASE */
+
+/* BEGIN_CASE */
+void get_fail( int uid_arg, data_t *data,
+               int overwrite_magic, int cut_header,
+               int expected_status )
+{
+    psa_storage_uid_t uid = uid_arg;
+    unsigned char *buffer = NULL;
+    psa_status_t status;
+    size_t n;
+    size_t ret_len = 0;
+    char filename[PSA_ITS_STORAGE_FILENAME_LENGTH];
+    FILE *stream = NULL;
+    char bad_char = 'X';
+
+    PSA_ASSERT( psa_its_set_wrap( uid, data->len, data->x, 0 ) );
+
+    psa_its_fill_filename( uid, filename );
+    stream = fopen( filename, "rb+" );
+    TEST_ASSERT( NULL != stream );
+    if( 0 != overwrite_magic )
+    {
+        /* Overwrite the 1st byte of the file, the ITS magic number */
+        TEST_ASSERT( fseek( stream, 0, SEEK_SET ) == 0 );
+        n = fwrite( &bad_char, 1, 1, stream );
+        TEST_ASSERT( 1 == n );
+    }
+    if( 0 != cut_header )
+    {
+        /* Reopen file and truncate it to 0 byte by specifying the 'w' flag */
+        stream = freopen( filename, "wb", stream );
+        TEST_ASSERT( NULL != stream );
+    }
+    fclose( stream );
+    stream = NULL;
+
+    status = psa_its_get( uid, 0, 0, buffer, &ret_len );
+    TEST_ASSERT( status == (psa_status_t) expected_status );
+    TEST_ASSERT( 0 == ret_len );
+    PSA_ASSERT( psa_its_remove( uid ) );
+
+    /* Check if the file is really deleted. */
+    stream = fopen( filename, "rb" );
+    TEST_ASSERT( NULL == stream );
+
+exit:
+    if( stream != NULL )
+        fclose( stream );
+
+    mbedtls_free( buffer );
+    cleanup( );
+}
+/* END_CASE */
diff --git a/tests/suites/test_suite_ssl.function b/tests/suites/test_suite_ssl.function
index d021cbd..5c97d90 100644
--- a/tests/suites/test_suite_ssl.function
+++ b/tests/suites/test_suite_ssl.function
@@ -3613,7 +3613,7 @@
      */
     for( i = block_size; i < buflen; i++ )
     {
-        test_set_step( i );
+        mbedtls_test_set_step( i );
 
         /* Restore correct pre-encryption record */
         rec = rec_save;
@@ -3650,7 +3650,7 @@
      */
     for( i = padlen; i <= pad_max_len; i++ )
     {
-        test_set_step( i );
+        mbedtls_test_set_step( i );
 
         /* Restore correct pre-encryption record */
         rec = rec_save;
@@ -4470,7 +4470,7 @@
      */
     for( max_in_len = 0; max_in_len <= 255 + block_size; max_in_len++ )
     {
-        test_set_step( max_in_len * 10000 );
+        mbedtls_test_set_step( max_in_len * 10000 );
 
         /* Use allocated in buffer to catch overreads */
         ASSERT_ALLOC( data, max_in_len );
@@ -4478,7 +4478,7 @@
         min_in_len = max_in_len > 255 ? max_in_len - 255 : 0;
         for( in_len = min_in_len; in_len <= max_in_len; in_len++ )
         {
-            test_set_step( max_in_len * 10000 + in_len );
+            mbedtls_test_set_step( max_in_len * 10000 + in_len );
 
             /* Set up dummy data and add_data */
             rec_num++;
@@ -4535,7 +4535,7 @@
 
     for( secret = offset_min; secret <= (size_t) offset_max; secret++ )
     {
-        test_set_step( (int) secret );
+        mbedtls_test_set_step( (int) secret );
 
         TEST_CF_SECRET( &secret, sizeof( secret ) );
         mbedtls_ssl_cf_memcpy_offset( dst, src, secret,
diff --git a/visualc/VS2010/mbedTLS.vcxproj b/visualc/VS2010/mbedTLS.vcxproj
index 78832eb..42cca89 100644
--- a/visualc/VS2010/mbedTLS.vcxproj
+++ b/visualc/VS2010/mbedTLS.vcxproj
@@ -248,9 +248,11 @@
     <ClInclude Include="..\..\library\common.h" />

     <ClInclude Include="..\..\library\psa_crypto_core.h" />

     <ClInclude Include="..\..\library\psa_crypto_driver_wrappers.h" />

+    <ClInclude Include="..\..\library\psa_crypto_ecp.h" />

     <ClInclude Include="..\..\library\psa_crypto_invasive.h" />

     <ClInclude Include="..\..\library\psa_crypto_its.h" />

     <ClInclude Include="..\..\library\psa_crypto_random_impl.h" />

+    <ClInclude Include="..\..\library\psa_crypto_rsa.h" />

     <ClInclude Include="..\..\library\psa_crypto_se.h" />

     <ClInclude Include="..\..\library\psa_crypto_service_integration.h" />

     <ClInclude Include="..\..\library\psa_crypto_slot_management.h" />

@@ -318,6 +320,8 @@
     <ClCompile Include="..\..\library\poly1305.c" />

     <ClCompile Include="..\..\library\psa_crypto.c" />

     <ClCompile Include="..\..\library\psa_crypto_driver_wrappers.c" />

+    <ClCompile Include="..\..\library\psa_crypto_ecp.c" />

+    <ClCompile Include="..\..\library\psa_crypto_rsa.c" />

     <ClCompile Include="..\..\library\psa_crypto_se.c" />

     <ClCompile Include="..\..\library\psa_crypto_slot_management.c" />

     <ClCompile Include="..\..\library\psa_crypto_storage.c" />