Merge pull request #4141 from gilles-peskine-arm/psa-test-functions-lift

Lift some PSA test functions to a separate module
diff --git a/library/md.c b/library/md.c
index de77b16..a10a835 100644
--- a/library/md.c
+++ b/library/md.c
@@ -1,5 +1,5 @@
 /**
- * \file mbedtls_md.c
+ * \file md.c
  *
  * \brief Generic message digest wrapper for mbed TLS
  *
diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt
index cefc5c9..e2bc420 100644
--- a/tests/CMakeLists.txt
+++ b/tests/CMakeLists.txt
@@ -67,6 +67,10 @@
     set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} -Wno-unused-function")
 endif(CMAKE_COMPILER_IS_GNUCC OR CMAKE_COMPILER_IS_CLANG)
 
+if(CMAKE_COMPILER_IS_CLANG)
+    set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} -Wdocumentation -Wno-documentation-deprecated-sync -Wunreachable-code")
+endif(CMAKE_COMPILER_IS_CLANG)
+
 if(MSVC)
     # If a warning level has been defined, suppress all warnings for test code
     set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} /W0")
diff --git a/tests/include/test/asn1_helpers.h b/tests/include/test/asn1_helpers.h
new file mode 100644
index 0000000..91ae260
--- /dev/null
+++ b/tests/include/test/asn1_helpers.h
@@ -0,0 +1,50 @@
+/** Helper functions for tests that manipulate ASN.1 data.
+ */
+/*
+ *  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 ASN1_HELPERS_H
+#define ASN1_HELPERS_H
+
+#include "test/helpers.h"
+
+/** Skip past an INTEGER in an ASN.1 buffer.
+ *
+ * Mark the current test case as failed in any of the following conditions:
+ * - The buffer does not start with an ASN.1 INTEGER.
+ * - The integer's size or parity does not match the constraints expressed
+ *   through \p min_bits, \p max_bits and \p must_be_odd.
+ *
+ * \param p             Upon entry, `*p` points to the first byte of the
+ *                      buffer to parse.
+ *                      On successful return, `*p` points to the first byte
+ *                      after the parsed INTEGER.
+ *                      On failure, `*p` is unspecified.
+ * \param end           The end of the ASN.1 buffer.
+ * \param min_bits      Fail the test case if the integer does not have at
+ *                      least this many significant bits.
+ * \param max_bits      Fail the test case if the integer has more than
+ *                      this many significant bits.
+ * \param must_be_odd   Fail the test case if the integer is even.
+ *
+ * \return              \c 0 if the test failed, otherwise 1.
+ */
+int mbedtls_test_asn1_skip_integer( unsigned char **p, const unsigned char *end,
+                                    size_t min_bits, size_t max_bits,
+                                    int must_be_odd );
+
+#endif /* ASN1_HELPERS_H */
diff --git a/tests/include/test/macros.h b/tests/include/test/macros.h
index 6930a5d..450bc2c 100644
--- a/tests/include/test/macros.h
+++ b/tests/include/test/macros.h
@@ -321,40 +321,45 @@
     mbedtls_exit( 1 );                                              \
 }
 
+/** \def ARRAY_LENGTH
+ * Return the number of elements of a static or stack array.
+ *
+ * \param array         A value of array (not pointer) type.
+ *
+ * \return The number of elements of the array.
+ */
+/* A correct implementation of ARRAY_LENGTH, but which silently gives
+ * a nonsensical result if called with a pointer rather than an array. */
+#define ARRAY_LENGTH_UNSAFE( array )            \
+    ( sizeof( array ) / sizeof( *( array ) ) )
+
 #if defined(__GNUC__)
 /* Test if arg and &(arg)[0] have the same type. This is true if arg is
  * an array but not if it's a pointer. */
 #define IS_ARRAY_NOT_POINTER( arg )                                     \
     ( ! __builtin_types_compatible_p( __typeof__( arg ),                \
                                       __typeof__( &( arg )[0] ) ) )
-#else
-/* On platforms where we don't know how to implement this check,
- * omit it. Oh well, a non-portable check is better than nothing. */
-#define IS_ARRAY_NOT_POINTER( arg ) 1
-#endif
-
 /* A compile-time constant with the value 0. If `const_expr` is not a
  * compile-time constant with a nonzero value, cause a compile-time error. */
 #define STATIC_ASSERT_EXPR( const_expr )                                \
     ( 0 && sizeof( struct { unsigned int STATIC_ASSERT : 1 - 2 * ! ( const_expr ); } ) )
+
 /* Return the scalar value `value` (possibly promoted). This is a compile-time
  * constant if `value` is. `condition` must be a compile-time constant.
  * If `condition` is false, arrange to cause a compile-time error. */
 #define STATIC_ASSERT_THEN_RETURN( condition, value )   \
     ( STATIC_ASSERT_EXPR( condition ) ? 0 : ( value ) )
 
-#define ARRAY_LENGTH_UNSAFE( array )            \
-    ( sizeof( array ) / sizeof( *( array ) ) )
-/** Return the number of elements of a static or stack array.
- *
- * \param array         A value of array (not pointer) type.
- *
- * \return The number of elements of the array.
- */
 #define ARRAY_LENGTH( array )                                           \
     ( STATIC_ASSERT_THEN_RETURN( IS_ARRAY_NOT_POINTER( array ),         \
                                  ARRAY_LENGTH_UNSAFE( array ) ) )
 
+#else
+/* If we aren't sure the compiler supports our non-standard tricks,
+ * fall back to the unsafe implementation. */
+#define ARRAY_LENGTH( array ) ARRAY_LENGTH_UNSAFE( array )
+#endif
+
 /** Return the smaller of two values.
  *
  * \param x         An integer-valued expression without side effects.
diff --git a/tests/include/test/psa_crypto_helpers.h b/tests/include/test/psa_crypto_helpers.h
index 30bb20f..9881eae 100644
--- a/tests/include/test/psa_crypto_helpers.h
+++ b/tests/include/test/psa_crypto_helpers.h
@@ -34,6 +34,59 @@
 #include "mbedtls/psa_util.h"
 #endif
 
+#if defined(MBEDTLS_PSA_CRYPTO_STORAGE_C)
+
+/* Internal function for #TEST_USES_KEY_ID. Return 1 on success, 0 on failure. */
+int mbedtls_test_uses_key_id( mbedtls_svc_key_id_t key_id );
+
+/** Destroy persistent keys recorded with #TEST_USES_KEY_ID.
+ */
+void mbedtls_test_psa_purge_key_storage( void );
+
+/** Purge the in-memory cache of persistent keys recorded with
+ * #TEST_USES_KEY_ID.
+ *
+ * Call this function before calling PSA_DONE() if it's ok for
+ * persistent keys to still exist at this point.
+ */
+void mbedtls_test_psa_purge_key_cache( void );
+
+/** \def TEST_USES_KEY_ID
+ *
+ * Call this macro in a test function before potentially creating a
+ * persistent key. Test functions that use this mechanism must call
+ * mbedtls_test_psa_purge_key_storage() in their cleanup code.
+ *
+ * This macro records a persistent key identifier as potentially used in the
+ * current test case. Recorded key identifiers will be cleaned up at the end
+ * of the test case, even on failure.
+ *
+ * This macro has no effect on volatile keys. Therefore, it is safe to call
+ * this macro in a test function that creates either volatile or persistent
+ * keys depending on the test data.
+ *
+ * This macro currently has no effect on special identifiers
+ * used to store implementation-specific files.
+ *
+ * Calling this macro multiple times on the same key identifier in the same
+ * test case has no effect.
+ *
+ * This macro can fail the test case if there isn't enough memory to
+ * record the key id.
+ *
+ * \param key_id    The PSA key identifier to record.
+ */
+#define TEST_USES_KEY_ID( key_id )                      \
+    TEST_ASSERT( mbedtls_test_uses_key_id( key_id ) )
+
+#else /* MBEDTLS_PSA_CRYPTO_STORAGE_C */
+
+#define TEST_USES_KEY_ID( key_id ) ( (void) ( key_id ) )
+#define mbedtls_test_psa_purge_key_storage( ) ( (void) 0 )
+#define mbedtls_test_psa_purge_key_cache( ) ( (void) 0 )
+
+#endif /* MBEDTLS_PSA_CRYPTO_STORAGE_C */
+
 #define PSA_INIT( ) PSA_ASSERT( psa_crypto_init( ) )
 
 /** Check for things that have not been cleaned up properly in the
@@ -60,13 +113,36 @@
     }                                                                   \
     while( 0 )
 
-/** Shut down the PSA Crypto subsystem. Expect a clean shutdown, with no slots
- * in use.
+/** Shut down the PSA Crypto subsystem and destroy persistent keys.
+ * Expect a clean shutdown, with no slots in use.
+ *
+ * If some key slots are still in use, record the test case as failed,
+ * but continue executing. This macro is suitable (and primarily intended)
+ * for use in the cleanup section of test functions.
+ *
+ * \note Persistent keys must be recorded with #TEST_USES_KEY_ID before
+ *       creating them.
  */
 #define PSA_DONE( )                                                     \
     do                                                                  \
     {                                                                   \
         test_fail_if_psa_leaking( __LINE__, __FILE__ );                 \
+        mbedtls_test_psa_purge_key_storage( );                          \
+        mbedtls_psa_crypto_free( );                                     \
+    }                                                                   \
+    while( 0 )
+
+/** Shut down the PSA Crypto subsystem, allowing persistent keys to survive.
+ * Expect a clean shutdown, with no slots in use.
+ *
+ * If some key slots are still in use, record the test case as failed and
+ * jump to the `exit` label.
+ */
+#define PSA_SESSION_DONE( )                                             \
+    do                                                                  \
+    {                                                                   \
+        mbedtls_test_psa_purge_key_cache( );                            \
+        ASSERT_PSA_PRISTINE( );                                         \
         mbedtls_psa_crypto_free( );                                     \
     }                                                                   \
     while( 0 )
diff --git a/tests/include/test/psa_exercise_key.h b/tests/include/test/psa_exercise_key.h
new file mode 100644
index 0000000..57eae58
--- /dev/null
+++ b/tests/include/test/psa_exercise_key.h
@@ -0,0 +1,241 @@
+/** Code to exercise a PSA key object, i.e. validate that it seems well-formed
+ * and can do what it is supposed to do.
+ */
+/*
+ *  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_EXERCISE_KEY_H
+#define PSA_EXERCISE_KEY_H
+
+#include "test/helpers.h"
+#include "test/psa_crypto_helpers.h"
+
+#include <psa/crypto.h>
+
+/** \def KNOWN_SUPPORTED_HASH_ALG
+ *
+ * A hash algorithm that is known to be supported.
+ *
+ * This is used in some smoke tests.
+ */
+#if defined(PSA_WANT_ALG_MD2)
+#define KNOWN_SUPPORTED_HASH_ALG PSA_ALG_MD2
+#elif defined(PSA_WANT_ALG_MD4)
+#define KNOWN_SUPPORTED_HASH_ALG PSA_ALG_MD4
+#elif defined(PSA_WANT_ALG_MD5)
+#define KNOWN_SUPPORTED_HASH_ALG PSA_ALG_MD5
+/* MBEDTLS_RIPEMD160_C omitted. This is necessary for the sake of
+ * exercise_signature_key() because Mbed TLS doesn't support RIPEMD160
+ * in RSA PKCS#1v1.5 signatures. A RIPEMD160-only configuration would be
+ * implausible anyway. */
+#elif defined(PSA_WANT_ALG_SHA_1)
+#define KNOWN_SUPPORTED_HASH_ALG PSA_ALG_SHA_1
+#elif defined(PSA_WANT_ALG_SHA_256)
+#define KNOWN_SUPPORTED_HASH_ALG PSA_ALG_SHA_256
+#elif defined(PSA_WANT_ALG_SHA_384)
+#define KNOWN_SUPPORTED_HASH_ALG PSA_ALG_SHA_384
+#elif defined(PSA_WANT_ALG_SHA_512)
+#define KNOWN_SUPPORTED_HASH_ALG PSA_ALG_SHA_512
+#elif defined(PSA_WANT_ALG_SHA3_256)
+#define KNOWN_SUPPORTED_HASH_ALG PSA_ALG_SHA3_256
+#else
+#undef KNOWN_SUPPORTED_HASH_ALG
+#endif
+
+/** \def KNOWN_SUPPORTED_BLOCK_CIPHER
+ *
+ * A block cipher that is known to be supported.
+ *
+ * For simplicity's sake, stick to block ciphers with 16-byte blocks.
+ */
+#if defined(MBEDTLS_AES_C)
+#define KNOWN_SUPPORTED_BLOCK_CIPHER PSA_KEY_TYPE_AES
+#elif defined(MBEDTLS_ARIA_C)
+#define KNOWN_SUPPORTED_BLOCK_CIPHER PSA_KEY_TYPE_ARIA
+#elif defined(MBEDTLS_CAMELLIA_C)
+#define KNOWN_SUPPORTED_BLOCK_CIPHER PSA_KEY_TYPE_CAMELLIA
+#undef KNOWN_SUPPORTED_BLOCK_CIPHER
+#endif
+
+/** \def KNOWN_SUPPORTED_MAC_ALG
+ *
+ * A MAC mode that is known to be supported.
+ *
+ * It must either be HMAC with #KNOWN_SUPPORTED_HASH_ALG or
+ * a block cipher-based MAC with #KNOWN_SUPPORTED_BLOCK_CIPHER.
+ *
+ * This is used in some smoke tests.
+ */
+#if defined(KNOWN_SUPPORTED_HASH_ALG) && defined(PSA_WANT_ALG_HMAC)
+#define KNOWN_SUPPORTED_MAC_ALG ( PSA_ALG_HMAC( KNOWN_SUPPORTED_HASH_ALG ) )
+#define KNOWN_SUPPORTED_MAC_KEY_TYPE PSA_KEY_TYPE_HMAC
+#elif defined(KNOWN_SUPPORTED_BLOCK_CIPHER) && defined(MBEDTLS_CMAC_C)
+#define KNOWN_SUPPORTED_MAC_ALG PSA_ALG_CMAC
+#define KNOWN_SUPPORTED_MAC_KEY_TYPE KNOWN_SUPPORTED_BLOCK_CIPHER
+#else
+#undef KNOWN_SUPPORTED_MAC_ALG
+#undef KNOWN_SUPPORTED_MAC_KEY_TYPE
+#endif
+
+/** \def KNOWN_SUPPORTED_BLOCK_CIPHER_ALG
+ *
+ * A cipher algorithm and key type that are known to be supported.
+ *
+ * This is used in some smoke tests.
+ */
+#if defined(KNOWN_SUPPORTED_BLOCK_CIPHER) && defined(MBEDTLS_CIPHER_MODE_CTR)
+#define KNOWN_SUPPORTED_BLOCK_CIPHER_ALG PSA_ALG_CTR
+#elif defined(KNOWN_SUPPORTED_BLOCK_CIPHER) && defined(MBEDTLS_CIPHER_MODE_CBC)
+#define KNOWN_SUPPORTED_BLOCK_CIPHER_ALG PSA_ALG_CBC_NO_PADDING
+#elif defined(KNOWN_SUPPORTED_BLOCK_CIPHER) && defined(MBEDTLS_CIPHER_MODE_CFB)
+#define KNOWN_SUPPORTED_BLOCK_CIPHER_ALG PSA_ALG_CFB
+#elif defined(KNOWN_SUPPORTED_BLOCK_CIPHER) && defined(MBEDTLS_CIPHER_MODE_OFB)
+#define KNOWN_SUPPORTED_BLOCK_CIPHER_ALG PSA_ALG_OFB
+#else
+#undef KNOWN_SUPPORTED_BLOCK_CIPHER_ALG
+#endif
+#if defined(KNOWN_SUPPORTED_BLOCK_CIPHER_ALG)
+#define KNOWN_SUPPORTED_CIPHER_ALG KNOWN_SUPPORTED_BLOCK_CIPHER_ALG
+#define KNOWN_SUPPORTED_CIPHER_KEY_TYPE KNOWN_SUPPORTED_BLOCK_CIPHER
+#elif defined(MBEDTLS_RC4_C)
+#define KNOWN_SUPPORTED_CIPHER_ALG PSA_ALG_RC4
+#define KNOWN_SUPPORTED_CIPHER_KEY_TYPE PSA_KEY_TYPE_RC4
+#else
+#undef KNOWN_SUPPORTED_CIPHER_ALG
+#undef KNOWN_SUPPORTED_CIPHER_KEY_TYPE
+#endif
+
+/** Convenience function to set up a key derivation.
+ *
+ * In case of failure, mark the current test case as failed.
+ *
+ * The inputs \p input1 and \p input2 are, in order:
+ * - HKDF: salt, info.
+ * - TKS 1.2 PRF, TLS 1.2 PSK-to-MS: seed, label.
+ *
+ * \param operation         The operation object to use.
+ *                          It must be in the initialized state.
+ * \param key               The key to use.
+ * \param alg               The algorithm to use.
+ * \param input1            The first input to pass.
+ * \param input1_length     The length of \p input1 in bytes.
+ * \param input2            The first input to pass.
+ * \param input2_length     The length of \p input2 in bytes.
+ * \param capacity          The capacity to set.
+ *
+ * \return                  \c 1 on success, \c 0 on failure.
+ */
+int mbedtls_test_psa_setup_key_derivation_wrap(
+    psa_key_derivation_operation_t* operation,
+    mbedtls_svc_key_id_t key,
+    psa_algorithm_t alg,
+    const unsigned char* input1, size_t input1_length,
+    const unsigned char* input2, size_t input2_length,
+    size_t capacity );
+
+/** Perform a key agreement using the given key pair against its public key
+ * using psa_raw_key_agreement().
+ *
+ * The result is discarded. The purpose of this function is to smoke-test a key.
+ *
+ * In case of failure, mark the current test case as failed.
+ *
+ * \param alg               A key agreement algorithm compatible with \p key.
+ * \param key               A key that allows key agreement with \p alg.
+ *
+ * \return                  \c 1 on success, \c 0 on failure.
+ */
+psa_status_t mbedtls_test_psa_raw_key_agreement_with_self(
+    psa_algorithm_t alg,
+    mbedtls_svc_key_id_t key );
+
+/** Perform a key agreement using the given key pair against its public key
+ * using psa_key_derivation_raw_key().
+ *
+ * The result is discarded. The purpose of this function is to smoke-test a key.
+ *
+ * In case of failure, mark the current test case as failed.
+ *
+ * \param operation         An operation that has been set up for a key
+ *                          agreement algorithm that is compatible with
+ *                          \p key.
+ * \param key               A key pair object that is suitable for a key
+ *                          agreement with \p operation.
+ *
+ * \return                  \c 1 on success, \c 0 on failure.
+ */
+psa_status_t mbedtls_test_psa_key_agreement_with_self(
+    psa_key_derivation_operation_t *operation,
+    mbedtls_svc_key_id_t key );
+
+/** Perform sanity checks on the given key representation.
+ *
+ * If any of the checks fail, mark the current test case as failed.
+ *
+ * The checks depend on the key type.
+ * - All types: check the export size against maximum-size macros.
+ * - DES: parity bits.
+ * - RSA: check the ASN.1 structure and the size and parity of the integers.
+ * - ECC private or public key: exact representation length.
+ * - Montgomery public key: first byte.
+ *
+ * \param type              The key type.
+ * \param bits              The key size in bits.
+ * \param exported          A buffer containing the key representation.
+ * \param exported_length   The length of \p exported in bytes.
+ *
+ * \return                  \c 1 if all checks passed, \c 0 on failure.
+ */
+int mbedtls_test_psa_exported_key_sanity_check(
+    psa_key_type_t type, size_t bits,
+    const uint8_t *exported, size_t exported_length );
+
+/** Do smoke tests on a key.
+ *
+ * Perform one of each operation indicated by \p alg (decrypt/encrypt,
+ * sign/verify, or derivation) that is permitted according to \p usage.
+ * \p usage and \p alg should correspond to the expected policy on the
+ * key.
+ *
+ * Export the key if permitted by \p usage, and check that the output
+ * looks sensible. If \p usage forbids export, check that
+ * \p psa_export_key correctly rejects the attempt. If the key is
+ * asymmetric, also check \p psa_export_public_key.
+ *
+ * If the key fails the tests, this function calls the test framework's
+ * `mbedtls_test_fail` function and returns false. Otherwise this function
+ * returns true. Therefore it should be used as follows:
+ * ```
+ * if( ! exercise_key( ... ) ) goto exit;
+ * ```
+ *
+ * \param key       The key to exercise. It should be capable of performing
+ *                  \p alg.
+ * \param usage     The usage flags to assume.
+ * \param alg       The algorithm to exercise.
+ *
+ * \retval 0 The key failed the smoke tests.
+ * \retval 1 The key passed the smoke tests.
+ */
+int mbedtls_test_psa_exercise_key( mbedtls_svc_key_id_t key,
+                                   psa_key_usage_t usage,
+                                   psa_algorithm_t alg );
+
+psa_key_usage_t mbedtls_test_psa_usage_to_exercise( psa_key_type_t type,
+                                                    psa_algorithm_t alg );
+
+#endif /* PSA_EXERCISE_KEY_H */
diff --git a/tests/src/asn1_helpers.c b/tests/src/asn1_helpers.c
new file mode 100644
index 0000000..79aa166
--- /dev/null
+++ b/tests/src/asn1_helpers.c
@@ -0,0 +1,74 @@
+/** \file asn1_helpers.c
+ *
+ * \brief Helper functions for tests that manipulate ASN.1 data.
+ */
+
+/*
+ *  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 <test/helpers.h>
+#include <test/macros.h>
+
+#if defined(MBEDTLS_ASN1_PARSE_C)
+
+#include <mbedtls/asn1.h>
+
+int mbedtls_test_asn1_skip_integer( unsigned char **p, const unsigned char *end,
+                                    size_t min_bits, size_t max_bits,
+                                    int must_be_odd )
+{
+    size_t len;
+    size_t actual_bits;
+    unsigned char msb;
+    TEST_EQUAL( mbedtls_asn1_get_tag( p, end, &len,
+                                      MBEDTLS_ASN1_INTEGER ),
+                0 );
+
+    /* Check if the retrieved length doesn't extend the actual buffer's size.
+     * It is assumed here, that end >= p, which validates casting to size_t. */
+    TEST_ASSERT( len <= (size_t)( end - *p) );
+
+    /* Tolerate a slight departure from DER encoding:
+     * - 0 may be represented by an empty string or a 1-byte string.
+     * - The sign bit may be used as a value bit. */
+    if( ( len == 1 && ( *p )[0] == 0 ) ||
+        ( len > 1 && ( *p )[0] == 0 && ( ( *p )[1] & 0x80 ) != 0 ) )
+    {
+        ++( *p );
+        --len;
+    }
+    if( min_bits == 0 && len == 0 )
+        return( 1 );
+    msb = ( *p )[0];
+    TEST_ASSERT( msb != 0 );
+    actual_bits = 8 * ( len - 1 );
+    while( msb != 0 )
+    {
+        msb >>= 1;
+        ++actual_bits;
+    }
+    TEST_ASSERT( actual_bits >= min_bits );
+    TEST_ASSERT( actual_bits <= max_bits );
+    if( must_be_odd )
+        TEST_ASSERT( ( ( *p )[len-1] & 1 ) != 0 );
+    *p += len;
+    return( 1 );
+exit:
+    return( 0 );
+}
+
+#endif /* MBEDTLS_ASN1_PARSE_C */
diff --git a/tests/src/fake_external_rng_for_test.c b/tests/src/fake_external_rng_for_test.c
index 98b3fe0..9c2195b 100644
--- a/tests/src/fake_external_rng_for_test.c
+++ b/tests/src/fake_external_rng_for_test.c
@@ -1,4 +1,4 @@
-/** \file psa_crypto_helpers.c
+/** \file fake_external_rng_for_test.c
  *
  * \brief Helper functions to test PSA crypto functionality.
  */
diff --git a/tests/src/psa_crypto_helpers.c b/tests/src/psa_crypto_helpers.c
index cb79a22..f2222cb 100644
--- a/tests/src/psa_crypto_helpers.c
+++ b/tests/src/psa_crypto_helpers.c
@@ -28,6 +28,51 @@
 
 #include <psa/crypto.h>
 
+#if defined(MBEDTLS_PSA_CRYPTO_STORAGE_C)
+
+#include <psa_crypto_storage.h>
+
+static mbedtls_svc_key_id_t key_ids_used_in_test[9];
+static size_t num_key_ids_used;
+
+int mbedtls_test_uses_key_id( mbedtls_svc_key_id_t key_id )
+{
+    size_t i;
+    if( MBEDTLS_SVC_KEY_ID_GET_KEY_ID( key_id ) >
+        PSA_MAX_PERSISTENT_KEY_IDENTIFIER )
+    {
+        /* Don't touch key id values that designate non-key files. */
+        return( 1 );
+    }
+    for( i = 0; i < num_key_ids_used ; i++ )
+    {
+        if( mbedtls_svc_key_id_equal( key_id, key_ids_used_in_test[i] ) )
+            return( 1 );
+    }
+    if( num_key_ids_used == ARRAY_LENGTH( key_ids_used_in_test ) )
+        return( 0 );
+    key_ids_used_in_test[num_key_ids_used] = key_id;
+    ++num_key_ids_used;
+    return( 1 );
+}
+
+void mbedtls_test_psa_purge_key_storage( void )
+{
+    size_t i;
+    for( i = 0; i < num_key_ids_used; i++ )
+        psa_destroy_persistent_key( key_ids_used_in_test[i] );
+    num_key_ids_used = 0;
+}
+
+void mbedtls_test_psa_purge_key_cache( void )
+{
+    size_t i;
+    for( i = 0; i < num_key_ids_used; i++ )
+        psa_purge_key( key_ids_used_in_test[i] );
+}
+
+#endif /* MBEDTLS_PSA_CRYPTO_STORAGE_C */
+
 const char *mbedtls_test_helper_is_psa_leaking( void )
 {
     mbedtls_psa_stats_t stats;
diff --git a/tests/src/psa_exercise_key.c b/tests/src/psa_exercise_key.c
new file mode 100644
index 0000000..9f80d7b
--- /dev/null
+++ b/tests/src/psa_exercise_key.c
@@ -0,0 +1,875 @@
+/** Code to exercise a PSA key object, i.e. validate that it seems well-formed
+ * and can do what it is supposed to do.
+ */
+
+/*
+ *  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 <test/helpers.h>
+#include <test/macros.h>
+#include <test/psa_exercise_key.h>
+
+#if defined(MBEDTLS_PSA_CRYPTO_C)
+
+#include <mbedtls/asn1.h>
+#include <psa/crypto.h>
+
+#include <test/asn1_helpers.h>
+#include <test/psa_crypto_helpers.h>
+
+#if defined(MBEDTLS_PSA_CRYPTO_SE_C)
+static int lifetime_is_dynamic_secure_element( psa_key_lifetime_t lifetime )
+{
+    return( PSA_KEY_LIFETIME_GET_LOCATION( lifetime ) !=
+            PSA_KEY_LOCATION_LOCAL_STORAGE );
+}
+#endif
+
+static int check_key_attributes_sanity( mbedtls_svc_key_id_t key )
+{
+    int ok = 0;
+    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
+    psa_key_lifetime_t lifetime;
+    mbedtls_svc_key_id_t id;
+    psa_key_type_t type;
+    size_t bits;
+
+    PSA_ASSERT( psa_get_key_attributes( key, &attributes ) );
+    lifetime = psa_get_key_lifetime( &attributes );
+    id = psa_get_key_id( &attributes );
+    type = psa_get_key_type( &attributes );
+    bits = psa_get_key_bits( &attributes );
+
+    /* Persistence */
+    if( PSA_KEY_LIFETIME_IS_VOLATILE( lifetime ) )
+    {
+        TEST_ASSERT(
+            ( PSA_KEY_ID_VOLATILE_MIN <=
+              MBEDTLS_SVC_KEY_ID_GET_KEY_ID( id ) ) &&
+            ( MBEDTLS_SVC_KEY_ID_GET_KEY_ID( id ) <=
+              PSA_KEY_ID_VOLATILE_MAX ) );
+    }
+    else
+    {
+        TEST_ASSERT(
+            ( PSA_KEY_ID_USER_MIN <= MBEDTLS_SVC_KEY_ID_GET_KEY_ID( id ) ) &&
+            ( MBEDTLS_SVC_KEY_ID_GET_KEY_ID( id ) <= PSA_KEY_ID_USER_MAX ) );
+    }
+#if defined(MBEDTLS_PSA_CRYPTO_SE_C)
+    /* randomly-generated 64-bit constant, should never appear in test data */
+    psa_key_slot_number_t slot_number = 0xec94d4a5058a1a21;
+    psa_status_t status = psa_get_key_slot_number( &attributes, &slot_number );
+    if( lifetime_is_dynamic_secure_element( lifetime ) )
+    {
+        /* Mbed Crypto currently always exposes the slot number to
+         * applications. This is not mandated by the PSA specification
+         * and may change in future versions. */
+        TEST_EQUAL( status, 0 );
+        TEST_ASSERT( slot_number != 0xec94d4a5058a1a21 );
+    }
+    else
+    {
+        TEST_EQUAL( status, PSA_ERROR_INVALID_ARGUMENT );
+    }
+#endif
+
+    /* Type and size */
+    TEST_ASSERT( type != 0 );
+    TEST_ASSERT( bits != 0 );
+    TEST_ASSERT( bits <= PSA_MAX_KEY_BITS );
+    if( PSA_KEY_TYPE_IS_UNSTRUCTURED( type ) )
+        TEST_ASSERT( bits % 8 == 0 );
+
+    /* MAX macros concerning specific key types */
+    if( PSA_KEY_TYPE_IS_ECC( type ) )
+        TEST_ASSERT( bits <= PSA_VENDOR_ECC_MAX_CURVE_BITS );
+    else if( PSA_KEY_TYPE_IS_RSA( type ) )
+        TEST_ASSERT( bits <= PSA_VENDOR_RSA_MAX_KEY_BITS );
+    TEST_ASSERT( PSA_BLOCK_CIPHER_BLOCK_LENGTH( type ) <= PSA_BLOCK_CIPHER_BLOCK_MAX_SIZE );
+
+    ok = 1;
+
+exit:
+    /*
+     * Key attributes may have been returned by psa_get_key_attributes()
+     * thus reset them as required.
+     */
+    psa_reset_key_attributes( &attributes );
+
+    return( ok );
+}
+
+static int exercise_mac_key( mbedtls_svc_key_id_t key,
+                             psa_key_usage_t usage,
+                             psa_algorithm_t alg )
+{
+    psa_mac_operation_t operation = PSA_MAC_OPERATION_INIT;
+    const unsigned char input[] = "foo";
+    unsigned char mac[PSA_MAC_MAX_SIZE] = {0};
+    size_t mac_length = sizeof( mac );
+
+    if( usage & PSA_KEY_USAGE_SIGN_HASH )
+    {
+        PSA_ASSERT( psa_mac_sign_setup( &operation, key, alg ) );
+        PSA_ASSERT( psa_mac_update( &operation,
+                                    input, sizeof( input ) ) );
+        PSA_ASSERT( psa_mac_sign_finish( &operation,
+                                         mac, sizeof( mac ),
+                                         &mac_length ) );
+    }
+
+    if( usage & PSA_KEY_USAGE_VERIFY_HASH )
+    {
+        psa_status_t verify_status =
+            ( usage & PSA_KEY_USAGE_SIGN_HASH ?
+              PSA_SUCCESS :
+              PSA_ERROR_INVALID_SIGNATURE );
+        PSA_ASSERT( psa_mac_verify_setup( &operation, key, alg ) );
+        PSA_ASSERT( psa_mac_update( &operation,
+                                    input, sizeof( input ) ) );
+        TEST_EQUAL( psa_mac_verify_finish( &operation, mac, mac_length ),
+                    verify_status );
+    }
+
+    return( 1 );
+
+exit:
+    psa_mac_abort( &operation );
+    return( 0 );
+}
+
+static int exercise_cipher_key( mbedtls_svc_key_id_t key,
+                                psa_key_usage_t usage,
+                                psa_algorithm_t alg )
+{
+    psa_cipher_operation_t operation = PSA_CIPHER_OPERATION_INIT;
+    unsigned char iv[16] = {0};
+    size_t iv_length = sizeof( iv );
+    const unsigned char plaintext[16] = "Hello, world...";
+    unsigned char ciphertext[32] = "(wabblewebblewibblewobblewubble)";
+    size_t ciphertext_length = sizeof( ciphertext );
+    unsigned char decrypted[sizeof( ciphertext )];
+    size_t part_length;
+
+    if( usage & PSA_KEY_USAGE_ENCRYPT )
+    {
+        PSA_ASSERT( psa_cipher_encrypt_setup( &operation, key, alg ) );
+        PSA_ASSERT( psa_cipher_generate_iv( &operation,
+                                            iv, sizeof( iv ),
+                                            &iv_length ) );
+        PSA_ASSERT( psa_cipher_update( &operation,
+                                       plaintext, sizeof( plaintext ),
+                                       ciphertext, sizeof( ciphertext ),
+                                       &ciphertext_length ) );
+        PSA_ASSERT( psa_cipher_finish( &operation,
+                                       ciphertext + ciphertext_length,
+                                       sizeof( ciphertext ) - ciphertext_length,
+                                       &part_length ) );
+        ciphertext_length += part_length;
+    }
+
+    if( usage & PSA_KEY_USAGE_DECRYPT )
+    {
+        psa_status_t status;
+        int maybe_invalid_padding = 0;
+        if( ! ( usage & PSA_KEY_USAGE_ENCRYPT ) )
+        {
+            psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
+            PSA_ASSERT( psa_get_key_attributes( key, &attributes ) );
+            /* This should be PSA_CIPHER_GET_IV_SIZE but the API doesn't
+             * have this macro yet. */
+            iv_length = PSA_BLOCK_CIPHER_BLOCK_LENGTH(
+                psa_get_key_type( &attributes ) );
+            maybe_invalid_padding = ! PSA_ALG_IS_STREAM_CIPHER( alg );
+            psa_reset_key_attributes( &attributes );
+        }
+        PSA_ASSERT( psa_cipher_decrypt_setup( &operation, key, alg ) );
+        PSA_ASSERT( psa_cipher_set_iv( &operation,
+                                       iv, iv_length ) );
+        PSA_ASSERT( psa_cipher_update( &operation,
+                                       ciphertext, ciphertext_length,
+                                       decrypted, sizeof( decrypted ),
+                                       &part_length ) );
+        status = psa_cipher_finish( &operation,
+                                    decrypted + part_length,
+                                    sizeof( decrypted ) - part_length,
+                                    &part_length );
+        /* For a stream cipher, all inputs are valid. For a block cipher,
+         * if the input is some aribtrary data rather than an actual
+         ciphertext, a padding error is likely.  */
+        if( maybe_invalid_padding )
+            TEST_ASSERT( status == PSA_SUCCESS ||
+                         status == PSA_ERROR_INVALID_PADDING );
+        else
+            PSA_ASSERT( status );
+    }
+
+    return( 1 );
+
+exit:
+    psa_cipher_abort( &operation );
+    return( 0 );
+}
+
+static int exercise_aead_key( mbedtls_svc_key_id_t key,
+                              psa_key_usage_t usage,
+                              psa_algorithm_t alg )
+{
+    unsigned char nonce[16] = {0};
+    size_t nonce_length = sizeof( nonce );
+    unsigned char plaintext[16] = "Hello, world...";
+    unsigned char ciphertext[48] = "(wabblewebblewibblewobblewubble)";
+    size_t ciphertext_length = sizeof( ciphertext );
+    size_t plaintext_length = sizeof( ciphertext );
+
+    /* Default IV length for AES-GCM is 12 bytes */
+    if( PSA_ALG_AEAD_WITH_SHORTENED_TAG( alg, 0 ) ==
+        PSA_ALG_AEAD_WITH_SHORTENED_TAG( PSA_ALG_GCM, 0 ) )
+    {
+        nonce_length = 12;
+    }
+
+    if( usage & PSA_KEY_USAGE_ENCRYPT )
+    {
+        PSA_ASSERT( psa_aead_encrypt( key, alg,
+                                      nonce, nonce_length,
+                                      NULL, 0,
+                                      plaintext, sizeof( plaintext ),
+                                      ciphertext, sizeof( ciphertext ),
+                                      &ciphertext_length ) );
+    }
+
+    if( usage & PSA_KEY_USAGE_DECRYPT )
+    {
+        psa_status_t verify_status =
+            ( usage & PSA_KEY_USAGE_ENCRYPT ?
+              PSA_SUCCESS :
+              PSA_ERROR_INVALID_SIGNATURE );
+        TEST_EQUAL( psa_aead_decrypt( key, alg,
+                                      nonce, nonce_length,
+                                      NULL, 0,
+                                      ciphertext, ciphertext_length,
+                                      plaintext, sizeof( plaintext ),
+                                      &plaintext_length ),
+                    verify_status );
+    }
+
+    return( 1 );
+
+exit:
+    return( 0 );
+}
+
+static int exercise_signature_key( mbedtls_svc_key_id_t key,
+                                   psa_key_usage_t usage,
+                                   psa_algorithm_t alg )
+{
+    unsigned char payload[PSA_HASH_MAX_SIZE] = {1};
+    size_t payload_length = 16;
+    unsigned char signature[PSA_SIGNATURE_MAX_SIZE] = {0};
+    size_t signature_length = sizeof( signature );
+    psa_algorithm_t hash_alg = PSA_ALG_SIGN_GET_HASH( alg );
+
+    /* If the policy allows signing with any hash, just pick one. */
+    if( PSA_ALG_IS_HASH_AND_SIGN( alg ) && hash_alg == PSA_ALG_ANY_HASH )
+    {
+#if defined(KNOWN_SUPPORTED_HASH_ALG)
+        hash_alg = KNOWN_SUPPORTED_HASH_ALG;
+        alg ^= PSA_ALG_ANY_HASH ^ hash_alg;
+#else
+        TEST_ASSERT( ! "No hash algorithm for hash-and-sign testing" );
+#endif
+    }
+
+    if( usage & PSA_KEY_USAGE_SIGN_HASH )
+    {
+        /* Some algorithms require the payload to have the size of
+         * the hash encoded in the algorithm. Use this input size
+         * even for algorithms that allow other input sizes. */
+        if( hash_alg != 0 )
+            payload_length = PSA_HASH_LENGTH( hash_alg );
+        PSA_ASSERT( psa_sign_hash( key, alg,
+                                   payload, payload_length,
+                                   signature, sizeof( signature ),
+                                   &signature_length ) );
+    }
+
+    if( usage & PSA_KEY_USAGE_VERIFY_HASH )
+    {
+        psa_status_t verify_status =
+            ( usage & PSA_KEY_USAGE_SIGN_HASH ?
+              PSA_SUCCESS :
+              PSA_ERROR_INVALID_SIGNATURE );
+        TEST_EQUAL( psa_verify_hash( key, alg,
+                                     payload, payload_length,
+                                     signature, signature_length ),
+                    verify_status );
+    }
+
+    return( 1 );
+
+exit:
+    return( 0 );
+}
+
+static int exercise_asymmetric_encryption_key( mbedtls_svc_key_id_t key,
+                                               psa_key_usage_t usage,
+                                               psa_algorithm_t alg )
+{
+    unsigned char plaintext[256] = "Hello, world...";
+    unsigned char ciphertext[256] = "(wabblewebblewibblewobblewubble)";
+    size_t ciphertext_length = sizeof( ciphertext );
+    size_t plaintext_length = 16;
+
+    if( usage & PSA_KEY_USAGE_ENCRYPT )
+    {
+        PSA_ASSERT( psa_asymmetric_encrypt( key, alg,
+                                            plaintext, plaintext_length,
+                                            NULL, 0,
+                                            ciphertext, sizeof( ciphertext ),
+                                            &ciphertext_length ) );
+    }
+
+    if( usage & PSA_KEY_USAGE_DECRYPT )
+    {
+        psa_status_t status =
+            psa_asymmetric_decrypt( key, alg,
+                                    ciphertext, ciphertext_length,
+                                    NULL, 0,
+                                    plaintext, sizeof( plaintext ),
+                                    &plaintext_length );
+        TEST_ASSERT( status == PSA_SUCCESS ||
+                     ( ( usage & PSA_KEY_USAGE_ENCRYPT ) == 0 &&
+                       ( status == PSA_ERROR_INVALID_ARGUMENT ||
+                         status == PSA_ERROR_INVALID_PADDING ) ) );
+    }
+
+    return( 1 );
+
+exit:
+    return( 0 );
+}
+
+int mbedtls_test_psa_setup_key_derivation_wrap(
+    psa_key_derivation_operation_t* operation,
+    mbedtls_svc_key_id_t key,
+    psa_algorithm_t alg,
+    const unsigned char* input1, size_t input1_length,
+    const unsigned char* input2, size_t input2_length,
+    size_t capacity )
+{
+    PSA_ASSERT( psa_key_derivation_setup( operation, alg ) );
+    if( PSA_ALG_IS_HKDF( alg ) )
+    {
+        PSA_ASSERT( psa_key_derivation_input_bytes( operation,
+                                                    PSA_KEY_DERIVATION_INPUT_SALT,
+                                                    input1, input1_length ) );
+        PSA_ASSERT( psa_key_derivation_input_key( operation,
+                                                  PSA_KEY_DERIVATION_INPUT_SECRET,
+                                                  key ) );
+        PSA_ASSERT( psa_key_derivation_input_bytes( operation,
+                                                    PSA_KEY_DERIVATION_INPUT_INFO,
+                                                    input2,
+                                                    input2_length ) );
+    }
+    else if( PSA_ALG_IS_TLS12_PRF( alg ) ||
+             PSA_ALG_IS_TLS12_PSK_TO_MS( alg ) )
+    {
+        PSA_ASSERT( psa_key_derivation_input_bytes( operation,
+                                                    PSA_KEY_DERIVATION_INPUT_SEED,
+                                                    input1, input1_length ) );
+        PSA_ASSERT( psa_key_derivation_input_key( operation,
+                                                  PSA_KEY_DERIVATION_INPUT_SECRET,
+                                                  key ) );
+        PSA_ASSERT( psa_key_derivation_input_bytes( operation,
+                                                    PSA_KEY_DERIVATION_INPUT_LABEL,
+                                                    input2, input2_length ) );
+    }
+    else
+    {
+        TEST_ASSERT( ! "Key derivation algorithm not supported" );
+    }
+
+    if( capacity != SIZE_MAX )
+        PSA_ASSERT( psa_key_derivation_set_capacity( operation, capacity ) );
+
+    return( 1 );
+
+exit:
+    return( 0 );
+}
+
+
+static int exercise_key_derivation_key( mbedtls_svc_key_id_t key,
+                                        psa_key_usage_t usage,
+                                        psa_algorithm_t alg )
+{
+    psa_key_derivation_operation_t operation = PSA_KEY_DERIVATION_OPERATION_INIT;
+    unsigned char input1[] = "Input 1";
+    size_t input1_length = sizeof( input1 );
+    unsigned char input2[] = "Input 2";
+    size_t input2_length = sizeof( input2 );
+    unsigned char output[1];
+    size_t capacity = sizeof( output );
+
+    if( usage & PSA_KEY_USAGE_DERIVE )
+    {
+        if( !mbedtls_test_psa_setup_key_derivation_wrap( &operation, key, alg,
+                                                         input1, input1_length,
+                                                         input2, input2_length,
+                                                         capacity ) )
+            goto exit;
+
+        PSA_ASSERT( psa_key_derivation_output_bytes( &operation,
+                                                     output,
+                                                     capacity ) );
+        PSA_ASSERT( psa_key_derivation_abort( &operation ) );
+    }
+
+    return( 1 );
+
+exit:
+    return( 0 );
+}
+
+/* We need two keys to exercise key agreement. Exercise the
+ * private key against its own public key. */
+psa_status_t mbedtls_test_psa_key_agreement_with_self(
+    psa_key_derivation_operation_t *operation,
+    mbedtls_svc_key_id_t key )
+{
+    psa_key_type_t private_key_type;
+    psa_key_type_t public_key_type;
+    size_t key_bits;
+    uint8_t *public_key = NULL;
+    size_t public_key_length;
+    /* Return GENERIC_ERROR if something other than the final call to
+     * psa_key_derivation_key_agreement fails. This isn't fully satisfactory,
+     * but it's good enough: callers will report it as a failed test anyway. */
+    psa_status_t status = PSA_ERROR_GENERIC_ERROR;
+    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
+
+    PSA_ASSERT( psa_get_key_attributes( key, &attributes ) );
+    private_key_type = psa_get_key_type( &attributes );
+    key_bits = psa_get_key_bits( &attributes );
+    public_key_type = PSA_KEY_TYPE_PUBLIC_KEY_OF_KEY_PAIR( private_key_type );
+    public_key_length = PSA_EXPORT_KEY_OUTPUT_SIZE( public_key_type, key_bits );
+    ASSERT_ALLOC( public_key, public_key_length );
+    PSA_ASSERT( psa_export_public_key( key, public_key, public_key_length,
+                                       &public_key_length ) );
+
+    status = psa_key_derivation_key_agreement(
+        operation, PSA_KEY_DERIVATION_INPUT_SECRET, key,
+        public_key, public_key_length );
+exit:
+    /*
+     * Key attributes may have been returned by psa_get_key_attributes()
+     * thus reset them as required.
+     */
+    psa_reset_key_attributes( &attributes );
+
+    mbedtls_free( public_key );
+    return( status );
+}
+
+/* We need two keys to exercise key agreement. Exercise the
+ * private key against its own public key. */
+psa_status_t mbedtls_test_psa_raw_key_agreement_with_self(
+    psa_algorithm_t alg,
+    mbedtls_svc_key_id_t key )
+{
+    psa_key_type_t private_key_type;
+    psa_key_type_t public_key_type;
+    size_t key_bits;
+    uint8_t *public_key = NULL;
+    size_t public_key_length;
+    uint8_t output[1024];
+    size_t output_length;
+    /* Return GENERIC_ERROR if something other than the final call to
+     * psa_key_derivation_key_agreement fails. This isn't fully satisfactory,
+     * but it's good enough: callers will report it as a failed test anyway. */
+    psa_status_t status = PSA_ERROR_GENERIC_ERROR;
+    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
+
+    PSA_ASSERT( psa_get_key_attributes( key, &attributes ) );
+    private_key_type = psa_get_key_type( &attributes );
+    key_bits = psa_get_key_bits( &attributes );
+    public_key_type = PSA_KEY_TYPE_PUBLIC_KEY_OF_KEY_PAIR( private_key_type );
+    public_key_length = PSA_EXPORT_KEY_OUTPUT_SIZE( public_key_type, key_bits );
+    ASSERT_ALLOC( public_key, public_key_length );
+    PSA_ASSERT( psa_export_public_key( key,
+                                       public_key, public_key_length,
+                                       &public_key_length ) );
+
+    status = psa_raw_key_agreement( alg, key,
+                                    public_key, public_key_length,
+                                    output, sizeof( output ), &output_length );
+exit:
+    /*
+     * Key attributes may have been returned by psa_get_key_attributes()
+     * thus reset them as required.
+     */
+    psa_reset_key_attributes( &attributes );
+
+    mbedtls_free( public_key );
+    return( status );
+}
+
+static int exercise_raw_key_agreement_key( mbedtls_svc_key_id_t key,
+                                           psa_key_usage_t usage,
+                                           psa_algorithm_t alg )
+{
+    int ok = 0;
+
+    if( usage & PSA_KEY_USAGE_DERIVE )
+    {
+        /* We need two keys to exercise key agreement. Exercise the
+         * private key against its own public key. */
+        PSA_ASSERT( mbedtls_test_psa_raw_key_agreement_with_self( alg, key ) );
+    }
+    ok = 1;
+
+exit:
+    return( ok );
+}
+
+static int exercise_key_agreement_key( mbedtls_svc_key_id_t key,
+                                       psa_key_usage_t usage,
+                                       psa_algorithm_t alg )
+{
+    psa_key_derivation_operation_t operation = PSA_KEY_DERIVATION_OPERATION_INIT;
+    unsigned char output[1];
+    int ok = 0;
+
+    if( usage & PSA_KEY_USAGE_DERIVE )
+    {
+        /* We need two keys to exercise key agreement. Exercise the
+         * private key against its own public key. */
+        PSA_ASSERT( psa_key_derivation_setup( &operation, alg ) );
+        PSA_ASSERT( mbedtls_test_psa_key_agreement_with_self( &operation, key ) );
+        PSA_ASSERT( psa_key_derivation_output_bytes( &operation,
+                                                     output,
+                                                     sizeof( output ) ) );
+        PSA_ASSERT( psa_key_derivation_abort( &operation ) );
+    }
+    ok = 1;
+
+exit:
+    return( ok );
+}
+
+int mbedtls_test_psa_exported_key_sanity_check(
+    psa_key_type_t type, size_t bits,
+    const uint8_t *exported, size_t exported_length )
+{
+    TEST_ASSERT( exported_length <= PSA_EXPORT_KEY_OUTPUT_SIZE( type, bits ) );
+
+    if( PSA_KEY_TYPE_IS_UNSTRUCTURED( type ) )
+        TEST_EQUAL( exported_length, PSA_BITS_TO_BYTES( bits ) );
+    else
+
+#if defined(MBEDTLS_RSA_C) && defined(MBEDTLS_PK_PARSE_C)
+    if( type == PSA_KEY_TYPE_RSA_KEY_PAIR )
+    {
+        uint8_t *p = (uint8_t*) exported;
+        const uint8_t *end = exported + exported_length;
+        size_t len;
+        /*   RSAPrivateKey ::= SEQUENCE {
+         *       version             INTEGER,  -- must be 0
+         *       modulus             INTEGER,  -- n
+         *       publicExponent      INTEGER,  -- e
+         *       privateExponent     INTEGER,  -- d
+         *       prime1              INTEGER,  -- p
+         *       prime2              INTEGER,  -- q
+         *       exponent1           INTEGER,  -- d mod (p-1)
+         *       exponent2           INTEGER,  -- d mod (q-1)
+         *       coefficient         INTEGER,  -- (inverse of q) mod p
+         *   }
+         */
+        TEST_EQUAL( mbedtls_asn1_get_tag( &p, end, &len,
+                                          MBEDTLS_ASN1_SEQUENCE |
+                                          MBEDTLS_ASN1_CONSTRUCTED ), 0 );
+        TEST_EQUAL( p + len, end );
+        if( ! mbedtls_test_asn1_skip_integer( &p, end, 0, 0, 0 ) )
+            goto exit;
+        if( ! mbedtls_test_asn1_skip_integer( &p, end, bits, bits, 1 ) )
+            goto exit;
+        if( ! mbedtls_test_asn1_skip_integer( &p, end, 2, bits, 1 ) )
+            goto exit;
+        /* Require d to be at least half the size of n. */
+        if( ! mbedtls_test_asn1_skip_integer( &p, end, bits / 2, bits, 1 ) )
+            goto exit;
+        /* Require p and q to be at most half the size of n, rounded up. */
+        if( ! mbedtls_test_asn1_skip_integer( &p, end, bits / 2, bits / 2 + 1, 1 ) )
+            goto exit;
+        if( ! mbedtls_test_asn1_skip_integer( &p, end, bits / 2, bits / 2 + 1, 1 ) )
+            goto exit;
+        if( ! mbedtls_test_asn1_skip_integer( &p, end, 1, bits / 2 + 1, 0 ) )
+            goto exit;
+        if( ! mbedtls_test_asn1_skip_integer( &p, end, 1, bits / 2 + 1, 0 ) )
+            goto exit;
+        if( ! mbedtls_test_asn1_skip_integer( &p, end, 1, bits / 2 + 1, 0 ) )
+            goto exit;
+        TEST_EQUAL( p, end );
+    }
+    else
+#endif /* MBEDTLS_RSA_C */
+
+#if defined(MBEDTLS_ECP_C)
+    if( PSA_KEY_TYPE_IS_ECC_KEY_PAIR( type ) )
+    {
+        /* Just the secret value */
+        TEST_EQUAL( exported_length, PSA_BITS_TO_BYTES( bits ) );
+    }
+    else
+#endif /* MBEDTLS_ECP_C */
+
+#if defined(MBEDTLS_RSA_C)
+    if( type == PSA_KEY_TYPE_RSA_PUBLIC_KEY )
+    {
+        uint8_t *p = (uint8_t*) exported;
+        const uint8_t *end = exported + exported_length;
+        size_t len;
+        /*   RSAPublicKey ::= SEQUENCE {
+         *      modulus            INTEGER,    -- n
+         *      publicExponent     INTEGER  }  -- e
+         */
+        TEST_EQUAL( mbedtls_asn1_get_tag( &p, end, &len,
+                                          MBEDTLS_ASN1_SEQUENCE |
+                                          MBEDTLS_ASN1_CONSTRUCTED ),
+                    0 );
+        TEST_EQUAL( p + len, end );
+        if( ! mbedtls_test_asn1_skip_integer( &p, end, bits, bits, 1 ) )
+            goto exit;
+        if( ! mbedtls_test_asn1_skip_integer( &p, end, 2, bits, 1 ) )
+            goto exit;
+        TEST_EQUAL( p, end );
+    }
+    else
+#endif /* MBEDTLS_RSA_C */
+
+#if defined(MBEDTLS_ECP_C)
+    if( PSA_KEY_TYPE_IS_ECC_PUBLIC_KEY( type ) )
+    {
+        if( PSA_KEY_TYPE_ECC_GET_FAMILY( type ) == PSA_ECC_FAMILY_MONTGOMERY )
+        {
+            /* The representation of an ECC Montgomery public key is
+             * the raw compressed point */
+             TEST_EQUAL( PSA_BITS_TO_BYTES( bits ), exported_length );
+        }
+        else
+        {
+            /* The representation of an ECC Weierstrass public key is:
+             *      - 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;
+             *      - where m is the bit size associated with the curve.
+             */
+            TEST_EQUAL( 1 + 2 * PSA_BITS_TO_BYTES( bits ), exported_length );
+            TEST_EQUAL( exported[0], 4 );
+        }
+    }
+    else
+#endif /* MBEDTLS_ECP_C */
+
+    {
+        TEST_ASSERT( ! "Sanity check not implemented for this key type" );
+    }
+
+#if defined(MBEDTLS_DES_C)
+    if( type == PSA_KEY_TYPE_DES )
+    {
+        /* Check the parity bits. */
+        unsigned i;
+        for( i = 0; i < bits / 8; i++ )
+        {
+            unsigned bit_count = 0;
+            unsigned m;
+            for( m = 1; m <= 0x100; m <<= 1 )
+            {
+                if( exported[i] & m )
+                    ++bit_count;
+            }
+            TEST_ASSERT( bit_count % 2 != 0 );
+        }
+    }
+#endif
+
+    return( 1 );
+
+exit:
+    return( 0 );
+}
+
+static int exercise_export_key( mbedtls_svc_key_id_t key,
+                                psa_key_usage_t usage )
+{
+    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
+    uint8_t *exported = NULL;
+    size_t exported_size = 0;
+    size_t exported_length = 0;
+    int ok = 0;
+
+    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, exported,
+                                    exported_size, &exported_length ),
+                    PSA_ERROR_NOT_PERMITTED );
+        ok = 1;
+        goto exit;
+    }
+
+    PSA_ASSERT( psa_export_key( key,
+                                exported, exported_size,
+                                &exported_length ) );
+    ok = mbedtls_test_psa_exported_key_sanity_check(
+        psa_get_key_type( &attributes ), psa_get_key_bits( &attributes ),
+        exported, exported_length );
+
+exit:
+    /*
+     * Key attributes may have been returned by psa_get_key_attributes()
+     * thus reset them as required.
+     */
+    psa_reset_key_attributes( &attributes );
+
+    mbedtls_free( exported );
+    return( ok );
+}
+
+static int exercise_export_public_key( mbedtls_svc_key_id_t key )
+{
+    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
+    psa_key_type_t public_type;
+    uint8_t *exported = NULL;
+    size_t exported_size = 0;
+    size_t exported_length = 0;
+    int ok = 0;
+
+    PSA_ASSERT( psa_get_key_attributes( key, &attributes ) );
+    if( ! PSA_KEY_TYPE_IS_ASYMMETRIC( psa_get_key_type( &attributes ) ) )
+    {
+        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 );
+        ok = 1;
+        goto exit;
+    }
+
+    public_type = PSA_KEY_TYPE_PUBLIC_KEY_OF_KEY_PAIR(
+        psa_get_key_type( &attributes ) );
+    exported_size = PSA_EXPORT_KEY_OUTPUT_SIZE( public_type,
+                                                psa_get_key_bits( &attributes ) );
+    ASSERT_ALLOC( exported, exported_size );
+
+    PSA_ASSERT( psa_export_public_key( key,
+                                       exported, exported_size,
+                                       &exported_length ) );
+    ok = mbedtls_test_psa_exported_key_sanity_check(
+        public_type, psa_get_key_bits( &attributes ),
+        exported, exported_length );
+
+exit:
+    /*
+     * Key attributes may have been returned by psa_get_key_attributes()
+     * thus reset them as required.
+     */
+    psa_reset_key_attributes( &attributes );
+
+    mbedtls_free( exported );
+    return( ok );
+}
+
+int mbedtls_test_psa_exercise_key( mbedtls_svc_key_id_t key,
+                                   psa_key_usage_t usage,
+                                   psa_algorithm_t alg )
+{
+    int ok = 0;
+
+    if( ! check_key_attributes_sanity( key ) )
+        return( 0 );
+
+    if( alg == 0 )
+        ok = 1; /* If no algorihm, do nothing (used for raw data "keys"). */
+    else if( PSA_ALG_IS_MAC( alg ) )
+        ok = exercise_mac_key( key, usage, alg );
+    else if( PSA_ALG_IS_CIPHER( alg ) )
+        ok = exercise_cipher_key( key, usage, alg );
+    else if( PSA_ALG_IS_AEAD( alg ) )
+        ok = exercise_aead_key( key, usage, alg );
+    else if( PSA_ALG_IS_SIGN( alg ) )
+        ok = exercise_signature_key( key, usage, alg );
+    else if( PSA_ALG_IS_ASYMMETRIC_ENCRYPTION( alg ) )
+        ok = exercise_asymmetric_encryption_key( key, usage, alg );
+    else if( PSA_ALG_IS_KEY_DERIVATION( alg ) )
+        ok = exercise_key_derivation_key( key, usage, alg );
+    else if( PSA_ALG_IS_RAW_KEY_AGREEMENT( alg ) )
+        ok = exercise_raw_key_agreement_key( key, usage, alg );
+    else if( PSA_ALG_IS_KEY_AGREEMENT( alg ) )
+        ok = exercise_key_agreement_key( key, usage, alg );
+    else
+        TEST_ASSERT( ! "No code to exercise this category of algorithm" );
+
+    ok = ok && exercise_export_key( key, usage );
+    ok = ok && exercise_export_public_key( key );
+
+exit:
+    return( ok );
+}
+
+psa_key_usage_t mbedtls_test_psa_usage_to_exercise( psa_key_type_t type,
+                                                    psa_algorithm_t alg )
+{
+    if( PSA_ALG_IS_MAC( alg ) || PSA_ALG_IS_SIGN( alg ) )
+    {
+        return( PSA_KEY_TYPE_IS_PUBLIC_KEY( type ) ?
+                PSA_KEY_USAGE_VERIFY_HASH :
+                PSA_KEY_USAGE_SIGN_HASH | PSA_KEY_USAGE_VERIFY_HASH );
+    }
+    else if( PSA_ALG_IS_CIPHER( alg ) || PSA_ALG_IS_AEAD( alg ) ||
+             PSA_ALG_IS_ASYMMETRIC_ENCRYPTION( alg ) )
+    {
+        return( PSA_KEY_TYPE_IS_PUBLIC_KEY( type ) ?
+                PSA_KEY_USAGE_ENCRYPT :
+                PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT );
+    }
+    else if( PSA_ALG_IS_KEY_DERIVATION( alg ) ||
+             PSA_ALG_IS_KEY_AGREEMENT( alg ) )
+    {
+        return( PSA_KEY_USAGE_DERIVE );
+    }
+    else
+    {
+        return( 0 );
+    }
+
+}
+
+#endif /* MBEDTLS_PSA_CRYPTO_C */
diff --git a/tests/suites/host_test.function b/tests/suites/host_test.function
index f4f4f45..d83d751 100644
--- a/tests/suites/host_test.function
+++ b/tests/suites/host_test.function
@@ -360,8 +360,6 @@
 /**
  * \brief       Tests snprintf implementation.
  *
- * \param none
- *
  * \return      0 for success else 1
  */
 static int run_test_snprintf( void )
@@ -428,8 +426,8 @@
  * \param unmet_dependencies         The array of unmet dependencies.
  * \param missing_unmet_dependencies Non-zero if there was a problem tracking
  *                                   all unmet dependencies, 0 otherwise.
- * \param ret                   The test dispatch status (DISPATCH_xxx).
- * \param mbedtls_test_info     A pointer to the test info structure.
+ * \param ret                        The test dispatch status (DISPATCH_xxx).
+ * \param info                       A pointer to the test info structure.
  */
 static void write_outcome_result( FILE *outcome_file,
                                   size_t unmet_dep_count,
diff --git a/tests/suites/main_test.function b/tests/suites/main_test.function
index aa408ea..36a7d23 100644
--- a/tests/suites/main_test.function
+++ b/tests/suites/main_test.function
@@ -105,7 +105,7 @@
  *              Identifiers and check code is generated by script:
  *              $generator_script
  *
- * \param exp_id    Dependency identifier.
+ * \param dep_id    Dependency identifier.
  *
  * \return       DEPENDENCY_SUPPORTED if set else DEPENDENCY_NOT_SUPPORTED
  */
@@ -129,13 +129,17 @@
 /**
  * \brief       Function pointer type for test function wrappers.
  *
+ * A test function wrapper decodes the parameters and passes them to the
+ * underlying test function. Both the wrapper and the underlying function
+ * return void. Test wrappers assume that they are passed a suitable
+ * parameter array and do not perform any error detection.
  *
- * \param void **   Pointer to void pointers. Represents an array of test
- *                  function parameters.
- *
- * \return       void
+ * \param param_array   The array of parameters. Each element is a `void *`
+ *                      which the wrapper casts to the correct type and
+ *                      dereferences. Each wrapper function hard-codes the
+ *                      number and types of the parameters.
  */
-typedef void (*TestWrapper_t)( void ** );
+typedef void (*TestWrapper_t)( void **param_array );
 
 
 /**
@@ -158,8 +162,8 @@
  *               parameter failure callback, to be used. Calls to setjmp()
  *               can invalidate the state of any local auto variables.
  *
- * \param fp     Function pointer to the test function
- * \param params Parameters to pass
+ * \param fp     Function pointer to the test function.
+ * \param params Parameters to pass to the #TestWrapper_t wrapper function.
  *
  */
 void execute_function_ptr(TestWrapper_t fp, void **params)
@@ -197,7 +201,9 @@
 /**
  * \brief        Dispatches test functions based on function index.
  *
- * \param exp_id    Test function index.
+ * \param func_idx    Test function index.
+ * \param params      The array of parameters to pass to the test function.
+ *                    It will be decoded by the #TestWrapper_t wrapper function.
  *
  * \return       DISPATCH_TEST_SUCCESS if found
  *               DISPATCH_TEST_FN_NOT_FOUND if not found
@@ -226,9 +232,10 @@
 
 
 /**
- * \brief       Checks if test function is supported
+ * \brief       Checks if test function is supported in this build-time
+ *              configuration.
  *
- * \param exp_id    Test function index.
+ * \param func_idx    Test function index.
  *
  * \return       DISPATCH_TEST_SUCCESS if found
  *               DISPATCH_TEST_FN_NOT_FOUND if not found
diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index 93f41b5..7ae6725 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -12,110 +12,13 @@
 #include "psa/crypto.h"
 #include "psa_crypto_slot_management.h"
 
+#include "test/asn1_helpers.h"
 #include "test/psa_crypto_helpers.h"
+#include "test/psa_exercise_key.h"
 
 /** An invalid export length that will never be set by psa_export_key(). */
 static const size_t INVALID_EXPORT_LENGTH = ~0U;
 
-/* A hash algorithm that is known to be supported.
- *
- * This is used in some smoke tests.
- */
-#if defined(PSA_WANT_ALG_MD2)
-#define KNOWN_SUPPORTED_HASH_ALG PSA_ALG_MD2
-#elif defined(PSA_WANT_ALG_MD4)
-#define KNOWN_SUPPORTED_HASH_ALG PSA_ALG_MD4
-#elif defined(PSA_WANT_ALG_MD5)
-#define KNOWN_SUPPORTED_HASH_ALG PSA_ALG_MD5
-/* MBEDTLS_RIPEMD160_C omitted. This is necessary for the sake of
- * exercise_signature_key() because Mbed TLS doesn't support RIPEMD160
- * in RSA PKCS#1v1.5 signatures. A RIPEMD160-only configuration would be
- * implausible anyway. */
-#elif defined(PSA_WANT_ALG_SHA_1)
-#define KNOWN_SUPPORTED_HASH_ALG PSA_ALG_SHA_1
-#elif defined(PSA_WANT_ALG_SHA_256)
-#define KNOWN_SUPPORTED_HASH_ALG PSA_ALG_SHA_256
-#elif defined(PSA_WANT_ALG_SHA_384)
-#define KNOWN_SUPPORTED_HASH_ALG PSA_ALG_SHA_384
-#elif defined(PSA_WANT_ALG_SHA_512)
-#define KNOWN_SUPPORTED_HASH_ALG PSA_ALG_SHA_512
-#elif defined(PSA_WANT_ALG_SHA3_256)
-#define KNOWN_SUPPORTED_HASH_ALG PSA_ALG_SHA3_256
-#else
-#undef KNOWN_SUPPORTED_HASH_ALG
-#endif
-
-/* A block cipher that is known to be supported.
- *
- * For simplicity's sake, stick to block ciphers with 16-byte blocks.
- */
-#if defined(MBEDTLS_AES_C)
-#define KNOWN_SUPPORTED_BLOCK_CIPHER PSA_KEY_TYPE_AES
-#elif defined(MBEDTLS_ARIA_C)
-#define KNOWN_SUPPORTED_BLOCK_CIPHER PSA_KEY_TYPE_ARIA
-#elif defined(MBEDTLS_CAMELLIA_C)
-#define KNOWN_SUPPORTED_BLOCK_CIPHER PSA_KEY_TYPE_CAMELLIA
-#undef KNOWN_SUPPORTED_BLOCK_CIPHER
-#endif
-
-/* A MAC mode that is known to be supported.
- *
- * It must either be HMAC with #KNOWN_SUPPORTED_HASH_ALG or
- * a block cipher-based MAC with #KNOWN_SUPPORTED_BLOCK_CIPHER.
- *
- * This is used in some smoke tests.
- */
-#if defined(KNOWN_SUPPORTED_HASH_ALG) && defined(PSA_WANT_ALG_HMAC)
-#define KNOWN_SUPPORTED_MAC_ALG ( PSA_ALG_HMAC( KNOWN_SUPPORTED_HASH_ALG ) )
-#define KNOWN_SUPPORTED_MAC_KEY_TYPE PSA_KEY_TYPE_HMAC
-#elif defined(KNOWN_SUPPORTED_BLOCK_CIPHER) && defined(MBEDTLS_CMAC_C)
-#define KNOWN_SUPPORTED_MAC_ALG PSA_ALG_CMAC
-#define KNOWN_SUPPORTED_MAC_KEY_TYPE KNOWN_SUPPORTED_BLOCK_CIPHER
-#else
-#undef KNOWN_SUPPORTED_MAC_ALG
-#undef KNOWN_SUPPORTED_MAC_KEY_TYPE
-#endif
-
-/* A cipher algorithm and key type that are known to be supported.
- *
- * This is used in some smoke tests.
- */
-#if defined(KNOWN_SUPPORTED_BLOCK_CIPHER) && defined(MBEDTLS_CIPHER_MODE_CTR)
-#define KNOWN_SUPPORTED_BLOCK_CIPHER_ALG PSA_ALG_CTR
-#elif defined(KNOWN_SUPPORTED_BLOCK_CIPHER) && defined(MBEDTLS_CIPHER_MODE_CBC)
-#define KNOWN_SUPPORTED_BLOCK_CIPHER_ALG PSA_ALG_CBC_NO_PADDING
-#elif defined(KNOWN_SUPPORTED_BLOCK_CIPHER) && defined(MBEDTLS_CIPHER_MODE_CFB)
-#define KNOWN_SUPPORTED_BLOCK_CIPHER_ALG PSA_ALG_CFB
-#elif defined(KNOWN_SUPPORTED_BLOCK_CIPHER) && defined(MBEDTLS_CIPHER_MODE_OFB)
-#define KNOWN_SUPPORTED_BLOCK_CIPHER_ALG PSA_ALG_OFB
-#else
-#undef KNOWN_SUPPORTED_BLOCK_CIPHER_ALG
-#endif
-#if defined(KNOWN_SUPPORTED_BLOCK_CIPHER_ALG)
-#define KNOWN_SUPPORTED_CIPHER_ALG KNOWN_SUPPORTED_BLOCK_CIPHER_ALG
-#define KNOWN_SUPPORTED_CIPHER_KEY_TYPE KNOWN_SUPPORTED_BLOCK_CIPHER
-#elif defined(MBEDTLS_RC4_C)
-#define KNOWN_SUPPORTED_CIPHER_ALG PSA_ALG_RC4
-#define KNOWN_SUPPORTED_CIPHER_KEY_TYPE PSA_KEY_TYPE_RC4
-#else
-#undef KNOWN_SUPPORTED_CIPHER_ALG
-#undef KNOWN_SUPPORTED_CIPHER_KEY_TYPE
-#endif
-
-#if defined(MBEDTLS_PSA_CRYPTO_SE_C)
-int lifetime_is_dynamic_secure_element( psa_key_lifetime_t lifetime )
-{
-    return( PSA_KEY_LIFETIME_GET_LOCATION( lifetime ) !=
-            PSA_KEY_LOCATION_LOCAL_STORAGE );
-}
-#else
-int lifetime_is_secure_element( psa_key_lifetime_t lifetime )
-{
-    (void) lifetime;
-    return( 0 );
-}
-#endif
-
 /** Test if a buffer contains a constant byte value.
  *
  * `mem_is_char(buffer, c, size)` is true after `memset(buffer, c, size)`.
@@ -222,80 +125,6 @@
     return( len );
 }
 
-int check_key_attributes_sanity( mbedtls_svc_key_id_t key )
-{
-    int ok = 0;
-    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
-    psa_key_lifetime_t lifetime;
-    mbedtls_svc_key_id_t id;
-    psa_key_type_t type;
-    psa_key_type_t bits;
-
-    PSA_ASSERT( psa_get_key_attributes( key, &attributes ) );
-    lifetime = psa_get_key_lifetime( &attributes );
-    id = psa_get_key_id( &attributes );
-    type = psa_get_key_type( &attributes );
-    bits = psa_get_key_bits( &attributes );
-
-    /* Persistence */
-    if( PSA_KEY_LIFETIME_IS_VOLATILE( lifetime ) )
-    {
-        TEST_ASSERT(
-            ( PSA_KEY_ID_VOLATILE_MIN <=
-              MBEDTLS_SVC_KEY_ID_GET_KEY_ID( id ) ) &&
-            ( MBEDTLS_SVC_KEY_ID_GET_KEY_ID( id ) <=
-              PSA_KEY_ID_VOLATILE_MAX ) );
-    }
-    else
-    {
-        TEST_ASSERT(
-            ( PSA_KEY_ID_USER_MIN <= MBEDTLS_SVC_KEY_ID_GET_KEY_ID( id ) ) &&
-            ( MBEDTLS_SVC_KEY_ID_GET_KEY_ID( id ) <= PSA_KEY_ID_USER_MAX ) );
-    }
-#if defined(MBEDTLS_PSA_CRYPTO_SE_C)
-    /* randomly-generated 64-bit constant, should never appear in test data */
-    psa_key_slot_number_t slot_number = 0xec94d4a5058a1a21;
-    psa_status_t status = psa_get_key_slot_number( &attributes, &slot_number );
-    if( lifetime_is_dynamic_secure_element( lifetime ) )
-    {
-        /* Mbed Crypto currently always exposes the slot number to
-         * applications. This is not mandated by the PSA specification
-         * and may change in future versions. */
-        TEST_EQUAL( status, 0 );
-        TEST_ASSERT( slot_number != 0xec94d4a5058a1a21 );
-    }
-    else
-    {
-        TEST_EQUAL( status, PSA_ERROR_INVALID_ARGUMENT );
-    }
-#endif
-
-    /* Type and size */
-    TEST_ASSERT( type != 0 );
-    TEST_ASSERT( bits != 0 );
-    TEST_ASSERT( bits <= PSA_MAX_KEY_BITS );
-    if( PSA_KEY_TYPE_IS_UNSTRUCTURED( type ) )
-        TEST_ASSERT( bits % 8 == 0 );
-
-    /* MAX macros concerning specific key types */
-    if( PSA_KEY_TYPE_IS_ECC( type ) )
-        TEST_ASSERT( bits <= PSA_VENDOR_ECC_MAX_CURVE_BITS );
-    else if( PSA_KEY_TYPE_IS_RSA( type ) )
-        TEST_ASSERT( bits <= PSA_VENDOR_RSA_MAX_KEY_BITS );
-    TEST_ASSERT( PSA_BLOCK_CIPHER_BLOCK_LENGTH( type ) <= PSA_BLOCK_CIPHER_BLOCK_MAX_SIZE );
-
-    ok = 1;
-
-exit:
-    /*
-     * Key attributes may have been returned by psa_get_key_attributes()
-     * thus reset them as required.
-     */
-    psa_reset_key_attributes( &attributes );
-
-    return( ok );
-}
-
 int exercise_mac_setup( psa_key_type_t key_type,
                         const unsigned char *key_bytes,
                         size_t key_length,
@@ -363,853 +192,6 @@
     return( 0 );
 }
 
-static int exercise_mac_key( mbedtls_svc_key_id_t key,
-                             psa_key_usage_t usage,
-                             psa_algorithm_t alg )
-{
-    psa_mac_operation_t operation = PSA_MAC_OPERATION_INIT;
-    const unsigned char input[] = "foo";
-    unsigned char mac[PSA_MAC_MAX_SIZE] = {0};
-    size_t mac_length = sizeof( mac );
-
-    if( usage & PSA_KEY_USAGE_SIGN_HASH )
-    {
-        PSA_ASSERT( psa_mac_sign_setup( &operation, key, alg ) );
-        PSA_ASSERT( psa_mac_update( &operation,
-                                    input, sizeof( input ) ) );
-        PSA_ASSERT( psa_mac_sign_finish( &operation,
-                                         mac, sizeof( mac ),
-                                         &mac_length ) );
-    }
-
-    if( usage & PSA_KEY_USAGE_VERIFY_HASH )
-    {
-        psa_status_t verify_status =
-            ( usage & PSA_KEY_USAGE_SIGN_HASH ?
-              PSA_SUCCESS :
-              PSA_ERROR_INVALID_SIGNATURE );
-        PSA_ASSERT( psa_mac_verify_setup( &operation, key, alg ) );
-        PSA_ASSERT( psa_mac_update( &operation,
-                                    input, sizeof( input ) ) );
-        TEST_EQUAL( psa_mac_verify_finish( &operation, mac, mac_length ),
-                    verify_status );
-    }
-
-    return( 1 );
-
-exit:
-    psa_mac_abort( &operation );
-    return( 0 );
-}
-
-static int exercise_cipher_key( mbedtls_svc_key_id_t key,
-                                psa_key_usage_t usage,
-                                psa_algorithm_t alg )
-{
-    psa_cipher_operation_t operation = PSA_CIPHER_OPERATION_INIT;
-    unsigned char iv[16] = {0};
-    size_t iv_length = sizeof( iv );
-    const unsigned char plaintext[16] = "Hello, world...";
-    unsigned char ciphertext[32] = "(wabblewebblewibblewobblewubble)";
-    size_t ciphertext_length = sizeof( ciphertext );
-    unsigned char decrypted[sizeof( ciphertext )];
-    size_t part_length;
-
-    if( usage & PSA_KEY_USAGE_ENCRYPT )
-    {
-        PSA_ASSERT( psa_cipher_encrypt_setup( &operation, key, alg ) );
-        PSA_ASSERT( psa_cipher_generate_iv( &operation,
-                                            iv, sizeof( iv ),
-                                            &iv_length ) );
-        PSA_ASSERT( psa_cipher_update( &operation,
-                                       plaintext, sizeof( plaintext ),
-                                       ciphertext, sizeof( ciphertext ),
-                                       &ciphertext_length ) );
-        PSA_ASSERT( psa_cipher_finish( &operation,
-                                       ciphertext + ciphertext_length,
-                                       sizeof( ciphertext ) - ciphertext_length,
-                                       &part_length ) );
-        ciphertext_length += part_length;
-    }
-
-    if( usage & PSA_KEY_USAGE_DECRYPT )
-    {
-        psa_status_t status;
-        int maybe_invalid_padding = 0;
-        if( ! ( usage & PSA_KEY_USAGE_ENCRYPT ) )
-        {
-            psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
-            PSA_ASSERT( psa_get_key_attributes( key, &attributes ) );
-            /* This should be PSA_CIPHER_GET_IV_SIZE but the API doesn't
-             * have this macro yet. */
-            iv_length = PSA_BLOCK_CIPHER_BLOCK_LENGTH(
-                psa_get_key_type( &attributes ) );
-            maybe_invalid_padding = ! PSA_ALG_IS_STREAM_CIPHER( alg );
-            psa_reset_key_attributes( &attributes );
-        }
-        PSA_ASSERT( psa_cipher_decrypt_setup( &operation, key, alg ) );
-        PSA_ASSERT( psa_cipher_set_iv( &operation,
-                                       iv, iv_length ) );
-        PSA_ASSERT( psa_cipher_update( &operation,
-                                       ciphertext, ciphertext_length,
-                                       decrypted, sizeof( decrypted ),
-                                       &part_length ) );
-        status = psa_cipher_finish( &operation,
-                                    decrypted + part_length,
-                                    sizeof( decrypted ) - part_length,
-                                    &part_length );
-        /* For a stream cipher, all inputs are valid. For a block cipher,
-         * if the input is some aribtrary data rather than an actual
-         ciphertext, a padding error is likely.  */
-        if( maybe_invalid_padding )
-            TEST_ASSERT( status == PSA_SUCCESS ||
-                         status == PSA_ERROR_INVALID_PADDING );
-        else
-            PSA_ASSERT( status );
-    }
-
-    return( 1 );
-
-exit:
-    psa_cipher_abort( &operation );
-    return( 0 );
-}
-
-static int exercise_aead_key( mbedtls_svc_key_id_t key,
-                              psa_key_usage_t usage,
-                              psa_algorithm_t alg )
-{
-    unsigned char nonce[16] = {0};
-    size_t nonce_length = sizeof( nonce );
-    unsigned char plaintext[16] = "Hello, world...";
-    unsigned char ciphertext[48] = "(wabblewebblewibblewobblewubble)";
-    size_t ciphertext_length = sizeof( ciphertext );
-    size_t plaintext_length = sizeof( ciphertext );
-
-    /* Default IV length for AES-GCM is 12 bytes */
-    if( PSA_ALG_AEAD_WITH_SHORTENED_TAG( alg, 0 ) ==
-        PSA_ALG_AEAD_WITH_SHORTENED_TAG( PSA_ALG_GCM, 0 ) )
-    {
-        nonce_length = 12;
-    }
-
-    if( usage & PSA_KEY_USAGE_ENCRYPT )
-    {
-        PSA_ASSERT( psa_aead_encrypt( key, alg,
-                                      nonce, nonce_length,
-                                      NULL, 0,
-                                      plaintext, sizeof( plaintext ),
-                                      ciphertext, sizeof( ciphertext ),
-                                      &ciphertext_length ) );
-    }
-
-    if( usage & PSA_KEY_USAGE_DECRYPT )
-    {
-        psa_status_t verify_status =
-            ( usage & PSA_KEY_USAGE_ENCRYPT ?
-              PSA_SUCCESS :
-              PSA_ERROR_INVALID_SIGNATURE );
-        TEST_EQUAL( psa_aead_decrypt( key, alg,
-                                      nonce, nonce_length,
-                                      NULL, 0,
-                                      ciphertext, ciphertext_length,
-                                      plaintext, sizeof( plaintext ),
-                                      &plaintext_length ),
-                    verify_status );
-    }
-
-    return( 1 );
-
-exit:
-    return( 0 );
-}
-
-static int exercise_signature_key( mbedtls_svc_key_id_t key,
-                                   psa_key_usage_t usage,
-                                   psa_algorithm_t alg )
-{
-    unsigned char payload[PSA_HASH_MAX_SIZE] = {1};
-    size_t payload_length = 16;
-    unsigned char signature[PSA_SIGNATURE_MAX_SIZE] = {0};
-    size_t signature_length = sizeof( signature );
-    psa_algorithm_t hash_alg = PSA_ALG_SIGN_GET_HASH( alg );
-
-    /* If the policy allows signing with any hash, just pick one. */
-    if( PSA_ALG_IS_HASH_AND_SIGN( alg ) && hash_alg == PSA_ALG_ANY_HASH )
-    {
-#if defined(KNOWN_SUPPORTED_HASH_ALG)
-        hash_alg = KNOWN_SUPPORTED_HASH_ALG;
-        alg ^= PSA_ALG_ANY_HASH ^ hash_alg;
-#else
-        mbedtls_test_fail( "No hash algorithm for hash-and-sign testing",
-                            __LINE__, __FILE__ );
-        return( 1 );
-#endif
-    }
-
-    if( usage & PSA_KEY_USAGE_SIGN_HASH )
-    {
-        /* Some algorithms require the payload to have the size of
-         * the hash encoded in the algorithm. Use this input size
-         * even for algorithms that allow other input sizes. */
-        if( hash_alg != 0 )
-            payload_length = PSA_HASH_LENGTH( hash_alg );
-        PSA_ASSERT( psa_sign_hash( key, alg,
-                                   payload, payload_length,
-                                   signature, sizeof( signature ),
-                                   &signature_length ) );
-    }
-
-    if( usage & PSA_KEY_USAGE_VERIFY_HASH )
-    {
-        psa_status_t verify_status =
-            ( usage & PSA_KEY_USAGE_SIGN_HASH ?
-              PSA_SUCCESS :
-              PSA_ERROR_INVALID_SIGNATURE );
-        TEST_EQUAL( psa_verify_hash( key, alg,
-                                     payload, payload_length,
-                                     signature, signature_length ),
-                    verify_status );
-    }
-
-    return( 1 );
-
-exit:
-    return( 0 );
-}
-
-static int exercise_asymmetric_encryption_key( mbedtls_svc_key_id_t key,
-                                               psa_key_usage_t usage,
-                                               psa_algorithm_t alg )
-{
-    unsigned char plaintext[256] = "Hello, world...";
-    unsigned char ciphertext[256] = "(wabblewebblewibblewobblewubble)";
-    size_t ciphertext_length = sizeof( ciphertext );
-    size_t plaintext_length = 16;
-
-    if( usage & PSA_KEY_USAGE_ENCRYPT )
-    {
-        PSA_ASSERT( psa_asymmetric_encrypt( key, alg,
-                                            plaintext, plaintext_length,
-                                            NULL, 0,
-                                            ciphertext, sizeof( ciphertext ),
-                                            &ciphertext_length ) );
-    }
-
-    if( usage & PSA_KEY_USAGE_DECRYPT )
-    {
-        psa_status_t status =
-            psa_asymmetric_decrypt( key, alg,
-                                    ciphertext, ciphertext_length,
-                                    NULL, 0,
-                                    plaintext, sizeof( plaintext ),
-                                    &plaintext_length );
-        TEST_ASSERT( status == PSA_SUCCESS ||
-                     ( ( usage & PSA_KEY_USAGE_ENCRYPT ) == 0 &&
-                       ( status == PSA_ERROR_INVALID_ARGUMENT ||
-                         status == PSA_ERROR_INVALID_PADDING ) ) );
-    }
-
-    return( 1 );
-
-exit:
-    return( 0 );
-}
-
-static int setup_key_derivation_wrap( psa_key_derivation_operation_t* operation,
-                                      mbedtls_svc_key_id_t key,
-                                      psa_algorithm_t alg,
-                                      unsigned char* input1, size_t input1_length,
-                                      unsigned char* input2, size_t input2_length,
-                                      size_t capacity )
-{
-    PSA_ASSERT( psa_key_derivation_setup( operation, alg ) );
-    if( PSA_ALG_IS_HKDF( alg ) )
-    {
-        PSA_ASSERT( psa_key_derivation_input_bytes( operation,
-                                                    PSA_KEY_DERIVATION_INPUT_SALT,
-                                                    input1, input1_length ) );
-        PSA_ASSERT( psa_key_derivation_input_key( operation,
-                                                  PSA_KEY_DERIVATION_INPUT_SECRET,
-                                                  key ) );
-        PSA_ASSERT( psa_key_derivation_input_bytes( operation,
-                                                    PSA_KEY_DERIVATION_INPUT_INFO,
-                                                    input2,
-                                                    input2_length ) );
-    }
-    else if( PSA_ALG_IS_TLS12_PRF( alg ) ||
-             PSA_ALG_IS_TLS12_PSK_TO_MS( alg ) )
-    {
-        PSA_ASSERT( psa_key_derivation_input_bytes( operation,
-                                                    PSA_KEY_DERIVATION_INPUT_SEED,
-                                                    input1, input1_length ) );
-        PSA_ASSERT( psa_key_derivation_input_key( operation,
-                                                  PSA_KEY_DERIVATION_INPUT_SECRET,
-                                                  key ) );
-        PSA_ASSERT( psa_key_derivation_input_bytes( operation,
-                                                    PSA_KEY_DERIVATION_INPUT_LABEL,
-                                                    input2, input2_length ) );
-    }
-    else
-    {
-        TEST_ASSERT( ! "Key derivation algorithm not supported" );
-    }
-
-    if( capacity != SIZE_MAX )
-        PSA_ASSERT( psa_key_derivation_set_capacity( operation, capacity ) );
-
-    return( 1 );
-
-exit:
-    return( 0 );
-}
-
-
-static int exercise_key_derivation_key( mbedtls_svc_key_id_t key,
-                                        psa_key_usage_t usage,
-                                        psa_algorithm_t alg )
-{
-    psa_key_derivation_operation_t operation = PSA_KEY_DERIVATION_OPERATION_INIT;
-    unsigned char input1[] = "Input 1";
-    size_t input1_length = sizeof( input1 );
-    unsigned char input2[] = "Input 2";
-    size_t input2_length = sizeof( input2 );
-    unsigned char output[1];
-    size_t capacity = sizeof( output );
-
-    if( usage & PSA_KEY_USAGE_DERIVE )
-    {
-        if( !setup_key_derivation_wrap( &operation, key, alg,
-                                        input1, input1_length,
-                                        input2, input2_length, capacity ) )
-            goto exit;
-
-        PSA_ASSERT( psa_key_derivation_output_bytes( &operation,
-                                                     output,
-                                                     capacity ) );
-        PSA_ASSERT( psa_key_derivation_abort( &operation ) );
-    }
-
-    return( 1 );
-
-exit:
-    return( 0 );
-}
-
-/* We need two keys to exercise key agreement. Exercise the
- * private key against its own public key. */
-static psa_status_t key_agreement_with_self(
-    psa_key_derivation_operation_t *operation,
-    mbedtls_svc_key_id_t key )
-{
-    psa_key_type_t private_key_type;
-    psa_key_type_t public_key_type;
-    size_t key_bits;
-    uint8_t *public_key = NULL;
-    size_t public_key_length;
-    /* Return GENERIC_ERROR if something other than the final call to
-     * psa_key_derivation_key_agreement fails. This isn't fully satisfactory,
-     * but it's good enough: callers will report it as a failed test anyway. */
-    psa_status_t status = PSA_ERROR_GENERIC_ERROR;
-    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
-
-    PSA_ASSERT( psa_get_key_attributes( key, &attributes ) );
-    private_key_type = psa_get_key_type( &attributes );
-    key_bits = psa_get_key_bits( &attributes );
-    public_key_type = PSA_KEY_TYPE_PUBLIC_KEY_OF_KEY_PAIR( private_key_type );
-    public_key_length = PSA_EXPORT_KEY_OUTPUT_SIZE( public_key_type, key_bits );
-    ASSERT_ALLOC( public_key, public_key_length );
-    PSA_ASSERT( psa_export_public_key( key, public_key, public_key_length,
-                                       &public_key_length ) );
-
-    status = psa_key_derivation_key_agreement(
-        operation, PSA_KEY_DERIVATION_INPUT_SECRET, key,
-        public_key, public_key_length );
-exit:
-    /*
-     * Key attributes may have been returned by psa_get_key_attributes()
-     * thus reset them as required.
-     */
-    psa_reset_key_attributes( &attributes );
-
-    mbedtls_free( public_key );
-    return( status );
-}
-
-/* We need two keys to exercise key agreement. Exercise the
- * private key against its own public key. */
-static psa_status_t raw_key_agreement_with_self( psa_algorithm_t alg,
-                                                 mbedtls_svc_key_id_t key )
-{
-    psa_key_type_t private_key_type;
-    psa_key_type_t public_key_type;
-    size_t key_bits;
-    uint8_t *public_key = NULL;
-    size_t public_key_length;
-    uint8_t output[1024];
-    size_t output_length;
-    /* Return GENERIC_ERROR if something other than the final call to
-     * psa_key_derivation_key_agreement fails. This isn't fully satisfactory,
-     * but it's good enough: callers will report it as a failed test anyway. */
-    psa_status_t status = PSA_ERROR_GENERIC_ERROR;
-    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
-
-    PSA_ASSERT( psa_get_key_attributes( key, &attributes ) );
-    private_key_type = psa_get_key_type( &attributes );
-    key_bits = psa_get_key_bits( &attributes );
-    public_key_type = PSA_KEY_TYPE_PUBLIC_KEY_OF_KEY_PAIR( private_key_type );
-    public_key_length = PSA_EXPORT_KEY_OUTPUT_SIZE( public_key_type, key_bits );
-    ASSERT_ALLOC( public_key, public_key_length );
-    PSA_ASSERT( psa_export_public_key( key,
-                                       public_key, public_key_length,
-                                       &public_key_length ) );
-
-    status = psa_raw_key_agreement( alg, key,
-                                    public_key, public_key_length,
-                                    output, sizeof( output ), &output_length );
-exit:
-    /*
-     * Key attributes may have been returned by psa_get_key_attributes()
-     * thus reset them as required.
-     */
-    psa_reset_key_attributes( &attributes );
-
-    mbedtls_free( public_key );
-    return( status );
-}
-
-static int exercise_raw_key_agreement_key( mbedtls_svc_key_id_t key,
-                                           psa_key_usage_t usage,
-                                           psa_algorithm_t alg )
-{
-    int ok = 0;
-
-    if( usage & PSA_KEY_USAGE_DERIVE )
-    {
-        /* We need two keys to exercise key agreement. Exercise the
-         * private key against its own public key. */
-        PSA_ASSERT( raw_key_agreement_with_self( alg, key ) );
-    }
-    ok = 1;
-
-exit:
-    return( ok );
-}
-
-static int exercise_key_agreement_key( mbedtls_svc_key_id_t key,
-                                       psa_key_usage_t usage,
-                                       psa_algorithm_t alg )
-{
-    psa_key_derivation_operation_t operation = PSA_KEY_DERIVATION_OPERATION_INIT;
-    unsigned char output[1];
-    int ok = 0;
-
-    if( usage & PSA_KEY_USAGE_DERIVE )
-    {
-        /* We need two keys to exercise key agreement. Exercise the
-         * private key against its own public key. */
-        PSA_ASSERT( psa_key_derivation_setup( &operation, alg ) );
-        PSA_ASSERT( key_agreement_with_self( &operation, key ) );
-        PSA_ASSERT( psa_key_derivation_output_bytes( &operation,
-                                                     output,
-                                                     sizeof( output ) ) );
-        PSA_ASSERT( psa_key_derivation_abort( &operation ) );
-    }
-    ok = 1;
-
-exit:
-    return( ok );
-}
-
-int asn1_skip_integer( unsigned char **p, const unsigned char *end,
-                       size_t min_bits, size_t max_bits,
-                       int must_be_odd )
-{
-    size_t len;
-    size_t actual_bits;
-    unsigned char msb;
-    TEST_EQUAL( mbedtls_asn1_get_tag( p, end, &len,
-                                      MBEDTLS_ASN1_INTEGER ),
-                0 );
-
-    /* Check if the retrieved length doesn't extend the actual buffer's size.
-     * It is assumed here, that end >= p, which validates casting to size_t. */
-    TEST_ASSERT( len <= (size_t)( end - *p) );
-
-    /* Tolerate a slight departure from DER encoding:
-     * - 0 may be represented by an empty string or a 1-byte string.
-     * - The sign bit may be used as a value bit. */
-    if( ( len == 1 && ( *p )[0] == 0 ) ||
-        ( len > 1 && ( *p )[0] == 0 && ( ( *p )[1] & 0x80 ) != 0 ) )
-    {
-        ++( *p );
-        --len;
-    }
-    if( min_bits == 0 && len == 0 )
-        return( 1 );
-    msb = ( *p )[0];
-    TEST_ASSERT( msb != 0 );
-    actual_bits = 8 * ( len - 1 );
-    while( msb != 0 )
-    {
-        msb >>= 1;
-        ++actual_bits;
-    }
-    TEST_ASSERT( actual_bits >= min_bits );
-    TEST_ASSERT( actual_bits <= max_bits );
-    if( must_be_odd )
-        TEST_ASSERT( ( ( *p )[len-1] & 1 ) != 0 );
-    *p += len;
-    return( 1 );
-exit:
-    return( 0 );
-}
-
-static int exported_key_sanity_check( psa_key_type_t type, size_t bits,
-                                      uint8_t *exported, size_t exported_length )
-{
-    if( PSA_KEY_TYPE_IS_UNSTRUCTURED( type ) )
-        TEST_EQUAL( exported_length, ( bits + 7 ) / 8 );
-    else
-        TEST_ASSERT( exported_length <= PSA_EXPORT_KEY_OUTPUT_SIZE( type, bits ) );
-
-#if defined(MBEDTLS_DES_C)
-    if( type == PSA_KEY_TYPE_DES )
-    {
-        /* Check the parity bits. */
-        unsigned i;
-        for( i = 0; i < bits / 8; i++ )
-        {
-            unsigned bit_count = 0;
-            unsigned m;
-            for( m = 1; m <= 0x100; m <<= 1 )
-            {
-                if( exported[i] & m )
-                    ++bit_count;
-            }
-            TEST_ASSERT( bit_count % 2 != 0 );
-        }
-    }
-    else
-#endif
-
-#if defined(MBEDTLS_RSA_C) && defined(MBEDTLS_PK_PARSE_C)
-    if( type == PSA_KEY_TYPE_RSA_KEY_PAIR )
-    {
-        uint8_t *p = exported;
-        uint8_t *end = exported + exported_length;
-        size_t len;
-        /*   RSAPrivateKey ::= SEQUENCE {
-         *       version             INTEGER,  -- must be 0
-         *       modulus             INTEGER,  -- n
-         *       publicExponent      INTEGER,  -- e
-         *       privateExponent     INTEGER,  -- d
-         *       prime1              INTEGER,  -- p
-         *       prime2              INTEGER,  -- q
-         *       exponent1           INTEGER,  -- d mod (p-1)
-         *       exponent2           INTEGER,  -- d mod (q-1)
-         *       coefficient         INTEGER,  -- (inverse of q) mod p
-         *   }
-         */
-        TEST_EQUAL( mbedtls_asn1_get_tag( &p, end, &len,
-                                          MBEDTLS_ASN1_SEQUENCE |
-                                          MBEDTLS_ASN1_CONSTRUCTED ), 0 );
-        TEST_EQUAL( p + len, end );
-        if( ! asn1_skip_integer( &p, end, 0, 0, 0 ) )
-            goto exit;
-        if( ! asn1_skip_integer( &p, end, bits, bits, 1 ) )
-            goto exit;
-        if( ! asn1_skip_integer( &p, end, 2, bits, 1 ) )
-            goto exit;
-        /* Require d to be at least half the size of n. */
-        if( ! asn1_skip_integer( &p, end, bits / 2, bits, 1 ) )
-            goto exit;
-        /* Require p and q to be at most half the size of n, rounded up. */
-        if( ! asn1_skip_integer( &p, end, bits / 2, bits / 2 + 1, 1 ) )
-            goto exit;
-        if( ! asn1_skip_integer( &p, end, bits / 2, bits / 2 + 1, 1 ) )
-            goto exit;
-        if( ! asn1_skip_integer( &p, end, 1, bits / 2 + 1, 0 ) )
-            goto exit;
-        if( ! asn1_skip_integer( &p, end, 1, bits / 2 + 1, 0 ) )
-            goto exit;
-        if( ! asn1_skip_integer( &p, end, 1, bits / 2 + 1, 0 ) )
-            goto exit;
-        TEST_EQUAL( p, end );
-    }
-    else
-#endif /* MBEDTLS_RSA_C */
-
-#if defined(MBEDTLS_ECP_C)
-    if( PSA_KEY_TYPE_IS_ECC_KEY_PAIR( type ) )
-    {
-        /* Just the secret value */
-        TEST_EQUAL( exported_length, PSA_BITS_TO_BYTES( bits ) );
-    }
-    else
-#endif /* MBEDTLS_ECP_C */
-
-    if( PSA_KEY_TYPE_IS_PUBLIC_KEY( type ) )
-    {
-        uint8_t *p = exported;
-        uint8_t *end = exported + exported_length;
-#if defined(MBEDTLS_RSA_C)
-        if( type == PSA_KEY_TYPE_RSA_PUBLIC_KEY )
-        {
-            size_t len;
-            /*   RSAPublicKey ::= SEQUENCE {
-             *      modulus            INTEGER,    -- n
-             *      publicExponent     INTEGER  }  -- e
-             */
-            TEST_EQUAL( mbedtls_asn1_get_tag( &p, end, &len,
-                                              MBEDTLS_ASN1_SEQUENCE |
-                                              MBEDTLS_ASN1_CONSTRUCTED ),
-                        0 );
-            TEST_EQUAL( p + len, end );
-            if( ! asn1_skip_integer( &p, end, bits, bits, 1 ) )
-                goto exit;
-            if( ! asn1_skip_integer( &p, end, 2, bits, 1 ) )
-                goto exit;
-            TEST_EQUAL( p, end );
-        }
-        else
-#endif /* MBEDTLS_RSA_C */
-#if defined(MBEDTLS_ECP_C)
-        if( PSA_KEY_TYPE_IS_ECC_PUBLIC_KEY( type ) )
-        {
-            if( PSA_KEY_TYPE_ECC_GET_FAMILY( type ) == PSA_ECC_FAMILY_MONTGOMERY )
-            {
-                /* The representation of an ECC Montgomery public key is
-                 * the raw compressed point */
-                 TEST_EQUAL( p + PSA_BITS_TO_BYTES( bits ), end );
-            }
-            else
-            {
-                /* The representation of an ECC Weierstrass public key is:
-                 *      - 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;
-                 *      - where m is the bit size associated with the curve.
-                 */
-                TEST_EQUAL( p + 1 + 2 * PSA_BITS_TO_BYTES( bits ), end );
-                TEST_EQUAL( p[0], 4 );
-            }
-        }
-        else
-#endif /* MBEDTLS_ECP_C */
-        {
-            char message[47];
-            mbedtls_snprintf( message, sizeof( message ),
-                              "No sanity check for public key type=0x%08lx",
-                              (unsigned long) type );
-            mbedtls_test_fail( message, __LINE__, __FILE__ );
-            (void) p;
-            (void) end;
-            return( 0 );
-        }
-    }
-    else
-
-    {
-        /* No sanity checks for other types */
-    }
-
-    return( 1 );
-
-exit:
-    return( 0 );
-}
-
-static int exercise_export_key( mbedtls_svc_key_id_t key,
-                                psa_key_usage_t usage )
-{
-    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
-    uint8_t *exported = NULL;
-    size_t exported_size = 0;
-    size_t exported_length = 0;
-    int ok = 0;
-
-    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, exported,
-                                    exported_size, &exported_length ),
-                    PSA_ERROR_NOT_PERMITTED );
-        ok = 1;
-        goto exit;
-    }
-
-    PSA_ASSERT( psa_export_key( key,
-                                exported, exported_size,
-                                &exported_length ) );
-    ok = exported_key_sanity_check( psa_get_key_type( &attributes ),
-                                    psa_get_key_bits( &attributes ),
-                                    exported, exported_length );
-
-exit:
-    /*
-     * Key attributes may have been returned by psa_get_key_attributes()
-     * thus reset them as required.
-     */
-    psa_reset_key_attributes( &attributes );
-
-    mbedtls_free( exported );
-    return( ok );
-}
-
-static int exercise_export_public_key( mbedtls_svc_key_id_t key )
-{
-    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
-    psa_key_type_t public_type;
-    uint8_t *exported = NULL;
-    size_t exported_size = 0;
-    size_t exported_length = 0;
-    int ok = 0;
-
-    PSA_ASSERT( psa_get_key_attributes( key, &attributes ) );
-    if( ! PSA_KEY_TYPE_IS_ASYMMETRIC( psa_get_key_type( &attributes ) ) )
-    {
-        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 );
-        ok = 1;
-        goto exit;
-    }
-
-    public_type = PSA_KEY_TYPE_PUBLIC_KEY_OF_KEY_PAIR(
-        psa_get_key_type( &attributes ) );
-    exported_size = PSA_EXPORT_KEY_OUTPUT_SIZE( public_type,
-                                                psa_get_key_bits( &attributes ) );
-    ASSERT_ALLOC( exported, exported_size );
-
-    PSA_ASSERT( psa_export_public_key( key,
-                                       exported, exported_size,
-                                       &exported_length ) );
-    ok = exported_key_sanity_check( public_type,
-                                    psa_get_key_bits( &attributes ),
-                                    exported, exported_length );
-
-exit:
-    /*
-     * Key attributes may have been returned by psa_get_key_attributes()
-     * thus reset them as required.
-     */
-    psa_reset_key_attributes( &attributes );
-
-    mbedtls_free( exported );
-    return( ok );
-}
-
-/** Do smoke tests on a key.
- *
- * Perform one of each operation indicated by \p alg (decrypt/encrypt,
- * sign/verify, or derivation) that is permitted according to \p usage.
- * \p usage and \p alg should correspond to the expected policy on the
- * key.
- *
- * Export the key if permitted by \p usage, and check that the output
- * looks sensible. If \p usage forbids export, check that
- * \p psa_export_key correctly rejects the attempt. If the key is
- * asymmetric, also check \p psa_export_public_key.
- *
- * If the key fails the tests, this function calls the test framework's
- * `mbedtls_test_fail` function and returns false. Otherwise this function
- * returns true. Therefore it should be used as follows:
- * ```
- * if( ! exercise_key( ... ) ) goto exit;
- * ```
- *
- * \param key       The key to exercise. It should be capable of performing
- *                  \p alg.
- * \param usage     The usage flags to assume.
- * \param alg       The algorithm to exercise.
- *
- * \retval 0 The key failed the smoke tests.
- * \retval 1 The key passed the smoke tests.
- */
-static int exercise_key( mbedtls_svc_key_id_t key,
-                         psa_key_usage_t usage,
-                         psa_algorithm_t alg )
-{
-    int ok;
-
-    if( ! check_key_attributes_sanity( key ) )
-        return( 0 );
-
-    if( alg == 0 )
-        ok = 1; /* If no algorihm, do nothing (used for raw data "keys"). */
-    else if( PSA_ALG_IS_MAC( alg ) )
-        ok = exercise_mac_key( key, usage, alg );
-    else if( PSA_ALG_IS_CIPHER( alg ) )
-        ok = exercise_cipher_key( key, usage, alg );
-    else if( PSA_ALG_IS_AEAD( alg ) )
-        ok = exercise_aead_key( key, usage, alg );
-    else if( PSA_ALG_IS_SIGN( alg ) )
-        ok = exercise_signature_key( key, usage, alg );
-    else if( PSA_ALG_IS_ASYMMETRIC_ENCRYPTION( alg ) )
-        ok = exercise_asymmetric_encryption_key( key, usage, alg );
-    else if( PSA_ALG_IS_KEY_DERIVATION( alg ) )
-        ok = exercise_key_derivation_key( key, usage, alg );
-    else if( PSA_ALG_IS_RAW_KEY_AGREEMENT( alg ) )
-        ok = exercise_raw_key_agreement_key( key, usage, alg );
-    else if( PSA_ALG_IS_KEY_AGREEMENT( alg ) )
-        ok = exercise_key_agreement_key( key, usage, alg );
-    else
-    {
-        char message[40];
-        mbedtls_snprintf( message, sizeof( message ),
-                          "No code to exercise alg=0x%08lx",
-                          (unsigned long) alg );
-        mbedtls_test_fail( message, __LINE__, __FILE__ );
-        ok = 0;
-    }
-
-    ok = ok && exercise_export_key( key, usage );
-    ok = ok && exercise_export_public_key( key );
-
-    return( ok );
-}
-
-static psa_key_usage_t usage_to_exercise( psa_key_type_t type,
-                                          psa_algorithm_t alg )
-{
-    if( PSA_ALG_IS_MAC( alg ) || PSA_ALG_IS_SIGN( alg ) )
-    {
-        return( PSA_KEY_TYPE_IS_PUBLIC_KEY( type ) ?
-                PSA_KEY_USAGE_VERIFY_HASH :
-                PSA_KEY_USAGE_SIGN_HASH | PSA_KEY_USAGE_VERIFY_HASH );
-    }
-    else if( PSA_ALG_IS_CIPHER( alg ) || PSA_ALG_IS_AEAD( alg ) ||
-             PSA_ALG_IS_ASYMMETRIC_ENCRYPTION( alg ) )
-    {
-        return( PSA_KEY_TYPE_IS_PUBLIC_KEY( type ) ?
-                PSA_KEY_USAGE_ENCRYPT :
-                PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT );
-    }
-    else if( PSA_ALG_IS_KEY_DERIVATION( alg ) ||
-             PSA_ALG_IS_KEY_AGREEMENT( alg ) )
-    {
-        return( PSA_KEY_USAGE_DERIVE );
-    }
-    else
-    {
-        return( 0 );
-    }
-
-}
-
 static int test_operations_on_invalid_key( mbedtls_svc_key_id_t key )
 {
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
@@ -1609,7 +591,11 @@
         goto destroy;
     }
 
-    if( ! exercise_export_key( key, usage_arg ) )
+    /* Run sanity checks on the exported key. For non-canonical inputs,
+     * this validates the canonical representations. For canonical inputs,
+     * this doesn't directly validate the implementation, but it still helps
+     * by cross-validating the test data with the sanity check code. */
+    if( ! mbedtls_test_psa_exercise_key( key, usage_arg, 0 ) )
         goto exit;
 
     if( canonical_input )
@@ -1715,7 +701,7 @@
     psa_key_type_t type = type_arg;
     size_t bits = bits_arg;
     psa_algorithm_t alg = alg_arg;
-    psa_key_usage_t usage = usage_to_exercise( type, alg );
+    psa_key_usage_t usage = mbedtls_test_psa_usage_to_exercise( type, alg );
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
     psa_key_attributes_t got_attributes = PSA_KEY_ATTRIBUTES_INIT;
 
@@ -1734,7 +720,7 @@
     TEST_EQUAL( psa_get_key_bits( &got_attributes ), bits );
 
     /* Do something with the key according to its type and permitted usage. */
-    if( ! exercise_key( key, usage, alg ) )
+    if( ! mbedtls_test_psa_exercise_key( key, usage, alg ) )
         goto exit;
 
     PSA_ASSERT( psa_destroy_key( key ) );
@@ -2187,7 +1173,7 @@
                                 &key ) );
 
     PSA_ASSERT( psa_key_derivation_setup( &operation, exercise_alg ) );
-    status = key_agreement_with_self( &operation, key );
+    status = mbedtls_test_psa_key_agreement_with_self( &operation, key );
 
     TEST_EQUAL( status, expected_status );
 
@@ -2224,9 +1210,9 @@
     TEST_EQUAL( psa_get_key_algorithm( &got_attributes ), alg );
     TEST_EQUAL( psa_get_key_enrollment_algorithm( &got_attributes ), alg2 );
 
-    if( ! exercise_key( key, usage, alg ) )
+    if( ! mbedtls_test_psa_exercise_key( key, usage, alg ) )
         goto exit;
-    if( ! exercise_key( key, usage, alg2 ) )
+    if( ! mbedtls_test_psa_exercise_key( key, usage, alg2 ) )
         goto exit;
 
 exit:
@@ -2265,7 +1251,7 @@
     PSA_ASSERT( psa_import_key( &attributes, key_data->x, key_data->len,
                                 &key ) );
 
-    status = raw_key_agreement_with_self( exercise_alg, key );
+    status = mbedtls_test_psa_raw_key_agreement_with_self( exercise_alg, key );
 
     TEST_EQUAL( status, expected_status );
 
@@ -2348,9 +1334,9 @@
         ASSERT_COMPARE( material->x, material->len,
                         export_buffer, length );
     }
-    if( ! exercise_key( target_key, expected_usage, expected_alg ) )
+    if( ! mbedtls_test_psa_exercise_key( target_key, expected_usage, expected_alg ) )
         goto exit;
-    if( ! exercise_key( target_key, expected_usage, expected_alg2 ) )
+    if( ! mbedtls_test_psa_exercise_key( target_key, expected_usage, expected_alg2 ) )
         goto exit;
 
     PSA_ASSERT( psa_destroy_key( target_key ) );
@@ -4740,7 +3726,8 @@
             {
                 // When taking a private key as secret input, use key agreement
                 // to add the shared secret to the derivation
-                TEST_EQUAL( key_agreement_with_self( &operation, keys[i] ),
+                TEST_EQUAL( mbedtls_test_psa_key_agreement_with_self(
+                                &operation, keys[i] ),
                             expected_statuses[i] );
             }
             else
@@ -4815,10 +3802,10 @@
                                 &key ) );
 
     /* valid key derivation */
-    if( !setup_key_derivation_wrap( &operation, key, alg,
-                                    input1, input1_length,
-                                    input2, input2_length,
-                                    capacity ) )
+    if( !mbedtls_test_psa_setup_key_derivation_wrap( &operation, key, alg,
+                                                     input1, input1_length,
+                                                     input2, input2_length,
+                                                     capacity ) )
         goto exit;
 
     /* state of operation shouldn't allow additional generation */
@@ -5007,10 +3994,10 @@
     PSA_ASSERT( psa_import_key( &attributes, key_data->x, key_data->len,
                                 &key ) );
 
-    if( !setup_key_derivation_wrap( &operation, key, alg,
-                                    input1->x, input1->len,
-                                    input2->x, input2->len,
-                                    requested_capacity ) )
+    if( !mbedtls_test_psa_setup_key_derivation_wrap( &operation, key, alg,
+                                                     input1->x, input1->len,
+                                                     input2->x, input2->len,
+                                                     requested_capacity ) )
         goto exit;
 
     PSA_ASSERT( psa_key_derivation_get_capacity( &operation,
@@ -5076,9 +4063,10 @@
                                 &base_key ) );
 
     /* Derive a key. */
-    if ( setup_key_derivation_wrap( &operation, base_key, alg,
-                                    input1->x, input1->len,
-                                    input2->x, input2->len, capacity ) )
+    if ( mbedtls_test_psa_setup_key_derivation_wrap( &operation, base_key, alg,
+                                                     input1->x, input1->len,
+                                                     input2->x, input2->len,
+                                                     capacity ) )
         goto exit;
 
     psa_set_key_usage_flags( &attributes, derived_usage );
@@ -5094,7 +4082,7 @@
     TEST_EQUAL( psa_get_key_bits( &got_attributes ), derived_bits );
 
     /* Exercise the derived key. */
-    if( ! exercise_key( derived_key, derived_usage, derived_alg ) )
+    if( ! mbedtls_test_psa_exercise_key( derived_key, derived_usage, derived_alg ) )
         goto exit;
 
 exit:
@@ -5143,9 +4131,10 @@
                                 &base_key ) );
 
     /* Derive some material and output it. */
-    if( !setup_key_derivation_wrap( &operation, base_key, alg,
-                                    input1->x, input1->len,
-                                    input2->x, input2->len, capacity ) )
+    if( !mbedtls_test_psa_setup_key_derivation_wrap( &operation, base_key, alg,
+                                                     input1->x, input1->len,
+                                                     input2->x, input2->len,
+                                                     capacity ) )
         goto exit;
 
     PSA_ASSERT( psa_key_derivation_output_bytes( &operation,
@@ -5154,9 +4143,10 @@
     PSA_ASSERT( psa_key_derivation_abort( &operation ) );
 
     /* Derive the same output again, but this time store it in key objects. */
-    if( !setup_key_derivation_wrap( &operation, base_key, alg,
-                                    input1->x, input1->len,
-                                    input2->x, input2->len, capacity ) )
+    if( !mbedtls_test_psa_setup_key_derivation_wrap( &operation, base_key, alg,
+                                                     input1->x, input1->len,
+                                                     input2->x, input2->len,
+                                                     capacity ) )
         goto exit;
 
     psa_set_key_usage_flags( &derived_attributes, PSA_KEY_USAGE_EXPORT );
@@ -5217,9 +4207,10 @@
     PSA_ASSERT( psa_import_key( &base_attributes, key_data->x, key_data->len,
                                 &base_key ) );
 
-    if( !setup_key_derivation_wrap( &operation, base_key, alg,
-                                    input1->x, input1->len,
-                                    input2->x, input2->len, SIZE_MAX ) )
+    if( !mbedtls_test_psa_setup_key_derivation_wrap( &operation, base_key, alg,
+                                                     input1->x, input1->len,
+                                                     input2->x, input2->len,
+                                                     SIZE_MAX ) )
         goto exit;
 
     psa_set_key_usage_flags( &derived_attributes, PSA_KEY_USAGE_EXPORT );
@@ -5535,7 +4526,7 @@
     TEST_EQUAL( psa_get_key_bits( &got_attributes ), bits );
 
     /* Do something with the key according to its type and permitted usage. */
-    if( ! exercise_key( key, usage, alg ) )
+    if( ! mbedtls_test_psa_exercise_key( key, usage, alg ) )
         goto exit;
 
 exit:
@@ -5607,7 +4598,7 @@
         ASSERT_COMPARE( e_read_buffer, e_read_length, e_arg->x, e_arg->len );
 
     /* Do something with the key according to its type and permitted usage. */
-    if( ! exercise_key( key, usage, alg ) )
+    if( ! mbedtls_test_psa_exercise_key( key, usage, alg ) )
         goto exit;
 
     /* Export the key and check the public exponent. */
@@ -5625,7 +4616,7 @@
         TEST_EQUAL( 0, mbedtls_asn1_get_tag( &p, end, &len,
                                              MBEDTLS_ASN1_SEQUENCE |
                                              MBEDTLS_ASN1_CONSTRUCTED ) );
-        TEST_ASSERT( asn1_skip_integer( &p, end, bits, bits, 1 ) );
+        TEST_ASSERT( mbedtls_test_asn1_skip_integer( &p, end, bits, bits, 1 ) );
         TEST_EQUAL( 0, mbedtls_asn1_get_tag( &p, end, &len,
                                              MBEDTLS_ASN1_INTEGER ) );
         if( len >= 1 && p[0] == 0 )
@@ -5783,7 +4774,7 @@
     }
 
     /* Do something with the key according to its type and permitted usage. */
-    if( ! exercise_key( key, usage_flags, alg ) )
+    if( ! mbedtls_test_psa_exercise_key( key, usage_flags, alg ) )
         goto exit;
 
 exit:
diff --git a/tests/suites/test_suite_psa_crypto_se_driver_hal.function b/tests/suites/test_suite_psa_crypto_se_driver_hal.function
index fa516c5..79d658f 100644
--- a/tests/suites/test_suite_psa_crypto_se_driver_hal.function
+++ b/tests/suites/test_suite_psa_crypto_se_driver_hal.function
@@ -767,41 +767,12 @@
     return( ok );
 }
 
-static mbedtls_svc_key_id_t key_ids_used_in_test[10];
-static size_t num_key_ids_used;
-
-/* Record a key id as potentially used in a test case. */
-static int test_uses_key_id( mbedtls_svc_key_id_t key_id )
-{
-    size_t i;
-
-    for( i = 0; i < num_key_ids_used ; i++ )
-    {
-        if( mbedtls_svc_key_id_equal( key_id, key_ids_used_in_test[i] ) )
-            return( 1 );
-    }
-
-    if( num_key_ids_used >= ARRAY_LENGTH( key_ids_used_in_test ) )
-        return( 0 );
-
-    key_ids_used_in_test[num_key_ids_used] = key_id;
-    ++num_key_ids_used;
-
-    return( 1 );
-}
-
-#define TEST_USES_KEY_ID( key_id )                       \
-    TEST_ASSERT( test_uses_key_id( key_id ) )
-
 static void psa_purge_storage( void )
 {
-    size_t i;
+    /* The generic code in mbedtls_test_psa_purge_key_storage()
+     * (which is called by PSA_DONE()) doesn't take care of things that are
+     * specific to dynamic secure elements. */
     psa_key_location_t location;
-
-    for( i = 0; i < num_key_ids_used; i++ )
-        psa_destroy_persistent_key( key_ids_used_in_test[i] );
-    num_key_ids_used = 0;
-
     /* Purge the transaction file. */
     psa_crypto_stop_transaction( );
     /* Purge driver persistent data. */
@@ -1496,6 +1467,7 @@
             ( validate > 0 ? PSA_SUCCESS : PSA_ERROR_NOT_PERMITTED );
     }
 
+    mbedtls_test_set_step( 1 );
     PSA_ASSERT( psa_register_se_driver( MIN_DRIVER_LOCATION, &driver ) );
     PSA_ASSERT( psa_crypto_init( ) );
 
@@ -1525,7 +1497,8 @@
     PSA_ASSERT( psa_purge_key( id ) );
 
     /* Restart and try again. */
-    PSA_DONE( );
+    mbedtls_test_set_step( 2 );
+    PSA_SESSION_DONE( );
     PSA_ASSERT( psa_register_se_driver( location, &driver ) );
     PSA_ASSERT( psa_crypto_init( ) );
     if( ! check_key_attributes( id, &attributes ) )
diff --git a/tests/suites/test_suite_psa_crypto_slot_management.function b/tests/suites/test_suite_psa_crypto_slot_management.function
index dbf05d2..bafb7d8 100644
--- a/tests/suites/test_suite_psa_crypto_slot_management.function
+++ b/tests/suites/test_suite_psa_crypto_slot_management.function
@@ -43,51 +43,6 @@
     INVALID_HANDLE_HUGE,
 } invalid_handle_construction_t;
 
-/* All test functions that create persistent keys must call
- * `TEST_USES_KEY_ID( key_id )` before creating a persistent key with this
- * identifier, and must call psa_purge_key_storage() in their cleanup
- * code. */
-
-#if defined(MBEDTLS_PSA_CRYPTO_STORAGE_C)
-static mbedtls_svc_key_id_t key_ids_used_in_test[9];
-static size_t num_key_ids_used;
-
-/* Record a key id as potentially used in a test case. */
-static int test_uses_key_id( mbedtls_svc_key_id_t key_id )
-{
-    size_t i;
-    if( MBEDTLS_SVC_KEY_ID_GET_KEY_ID( key_id ) >
-        PSA_MAX_PERSISTENT_KEY_IDENTIFIER )
-    {
-        /* Don't touch key id values that designate non-key files. */
-        return( 1 );
-    }
-    for( i = 0; i < num_key_ids_used ; i++ )
-    {
-        if( mbedtls_svc_key_id_equal( key_id, key_ids_used_in_test[i] ) )
-            return( 1 );
-    }
-    if( num_key_ids_used == ARRAY_LENGTH( key_ids_used_in_test ) )
-        return( 0 );
-    key_ids_used_in_test[num_key_ids_used] = key_id;
-    ++num_key_ids_used;
-    return( 1 );
-}
-#define TEST_USES_KEY_ID( key_id )                       \
-    TEST_ASSERT( test_uses_key_id( key_id ) )
-
-/* Destroy all key ids that may have been created by the current test case. */
-static void psa_purge_key_storage( void )
-{
-    size_t i;
-    for( i = 0; i < num_key_ids_used; i++ )
-        psa_destroy_persistent_key( key_ids_used_in_test[i] );
-    num_key_ids_used = 0;
-}
-#else
-#define TEST_USES_KEY_ID( key_id ) ( (void) ( key_id ) )
-#endif /* MBEDTLS_PSA_CRYPTO_STORAGE_C */
-
 /** Apply \p invalidate_method to invalidate the specified key:
  * close it, destroy it, or do nothing;
  */
@@ -131,7 +86,7 @@
         case INVALIDATE_BY_DESTROYING_WITH_SHUTDOWN:
         case INVALIDATE_BY_PURGING_WITH_SHUTDOWN:
             /* All keys must have been closed. */
-            PSA_DONE( );
+            PSA_SESSION_DONE( );
             break;
         case INVALIDATE_BY_SHUTDOWN:
             /* Some keys may remain behind, and we're testing that this
@@ -168,6 +123,7 @@
     mbedtls_svc_key_id_t key = MBEDTLS_SVC_KEY_ID_INIT;
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
 
+    mbedtls_test_set_step( 1 );
     PSA_ASSERT( psa_crypto_init( ) );
 
     /* Import a key. */
@@ -214,6 +170,7 @@
     psa_reset_key_attributes( &attributes );
 
     /* Do something that invalidates the key. */
+    mbedtls_test_set_step( 2 );
     if( ! invalidate_key( invalidate_method, key ) )
         goto exit;
     if( ! invalidate_psa( invalidate_method ) )
@@ -263,6 +220,7 @@
 
     TEST_USES_KEY_ID( id );
 
+    mbedtls_test_set_step( 1 );
     PSA_ASSERT( psa_crypto_init( ) );
 
     psa_set_key_id( &attributes, id );
@@ -312,6 +270,7 @@
      * Do something that wipes key data in volatile memory or destroy the
      * key.
      */
+    mbedtls_test_set_step( 2 );
     if( ! invalidate_key( invalidate_method, id ) )
         goto exit;
     if( ! invalidate_psa( invalidate_method ) )
@@ -383,7 +342,6 @@
     psa_reset_key_attributes( &read_attributes );
 
     PSA_DONE( );
-    psa_purge_key_storage( );
     mbedtls_free( reexported );
 }
 /* END_CASE */
@@ -457,7 +415,6 @@
     psa_reset_key_attributes( &attributes );
 
     PSA_DONE( );
-    psa_purge_key_storage( );
 }
 /* END_CASE */
 
@@ -516,9 +473,6 @@
 
 exit:
     PSA_DONE( );
-#if defined(MBEDTLS_PSA_CRYPTO_STORAGE_C)
-    psa_purge_key_storage( );
-#endif
 }
 /* END_CASE */
 
@@ -654,9 +608,6 @@
 
     PSA_DONE( );
     mbedtls_free( export_buffer );
-#if defined(MBEDTLS_PSA_CRYPTO_STORAGE_C)
-    psa_purge_key_storage( );
-#endif
 }
 /* END_CASE */
 
@@ -772,9 +723,6 @@
 
     PSA_DONE( );
     mbedtls_free( export_buffer );
-#if defined(MBEDTLS_PSA_CRYPTO_STORAGE_C)
-    psa_purge_key_storage( );
-#endif
 }
 /* END_CASE */
 
diff --git a/visualc/VS2010/mbedTLS.vcxproj b/visualc/VS2010/mbedTLS.vcxproj
index da026d9..c2051e6 100644
--- a/visualc/VS2010/mbedTLS.vcxproj
+++ b/visualc/VS2010/mbedTLS.vcxproj
@@ -232,11 +232,13 @@
     <ClInclude Include="..\..\include\psa\crypto_struct.h" />

     <ClInclude Include="..\..\include\psa\crypto_types.h" />

     <ClInclude Include="..\..\include\psa\crypto_values.h" />

+    <ClInclude Include="..\..\tests\include\test\asn1_helpers.h" />

     <ClInclude Include="..\..\tests\include\test\constant_flow.h" />

     <ClInclude Include="..\..\tests\include\test\fake_external_rng_for_test.h" />

     <ClInclude Include="..\..\tests\include\test\helpers.h" />

     <ClInclude Include="..\..\tests\include\test\macros.h" />

     <ClInclude Include="..\..\tests\include\test\psa_crypto_helpers.h" />

+    <ClInclude Include="..\..\tests\include\test\psa_exercise_key.h" />

     <ClInclude Include="..\..\tests\include\test\psa_helpers.h" />

     <ClInclude Include="..\..\tests\include\test\random.h" />

     <ClInclude Include="..\..\tests\include\test\drivers\cipher.h" />

@@ -354,9 +356,11 @@
     <ClCompile Include="..\..\library\x509write_crt.c" />

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

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

+    <ClCompile Include="..\..\tests\src\asn1_helpers.c" />

     <ClCompile Include="..\..\tests\src\fake_external_rng_for_test.c" />

     <ClCompile Include="..\..\tests\src\helpers.c" />

     <ClCompile Include="..\..\tests\src\psa_crypto_helpers.c" />

+    <ClCompile Include="..\..\tests\src\psa_exercise_key.c" />

     <ClCompile Include="..\..\tests\src\random.c" />

     <ClCompile Include="..\..\tests\src\threading_helpers.c" />

     <ClCompile Include="..\..\3rdparty\everest\library\everest.c" />