Merge pull request #3933 from ronald-cron-arm/psa-import-export

Psa import export
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/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..2d93397 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 )
@@ -2140,13 +1615,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 +1634,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 +1739,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 +1793,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 +1808,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 +1887,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 +2700,7 @@
         return( ret );
 
     ret = mbedtls_cipher_cmac_starts( &operation->ctx.cmac,
-                                      slot->data.key.data,
+                                      slot->key.data,
                                       key_bits );
     return( ret );
 }
@@ -3382,8 +2846,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 +3432,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 +3464,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 +3542,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 +3567,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 +3651,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 +3757,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 +3919,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 +3929,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 +4428,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 +4450,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 +4470,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 +5135,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 );
@@ -6129,8 +5609,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 +5638,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 +5694,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 +6039,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 +6063,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 +6090,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 +6138,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 );
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_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_slot_management.c b/library/psa_crypto_slot_management.c
index 39d6dbb..1188ce6 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;
             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/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/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/suites/test_suite_psa_crypto.data b/tests/suites/test_suite_psa_crypto.data
index 99e2e8c..307b2c1 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
@@ -2757,7 +2757,7 @@
 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
 
 PSA generate key: ECC, Curve25519, good
-depends_on:PSA_WANT_ALG_ECDH:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:MBEDTLS_ECP_DP_CURVE25519_ENABLED
+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
 
 PSA generate key: RSA, default e
diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index 82514d6..a65c482 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -1028,19 +1028,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 +1073,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(
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_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/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" />