Move exercise_key and related functions to their own module

Move mbedtls_test_psa_exercise_key() (formerly exercise_key()) and
related functions to its own module. Export the few auxiliary
functions that are also called directly.

Signed-off-by: Gilles Peskine <Gilles.Peskine@arm.com>
diff --git a/tests/include/test/psa_exercise_key.h b/tests/include/test/psa_exercise_key.h
index d4a0ea6..8fd152c 100644
--- a/tests/include/test/psa_exercise_key.h
+++ b/tests/include/test/psa_exercise_key.h
@@ -26,5 +26,143 @@
 
 #include <psa/crypto.h>
 
+/* 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
+
+int mbedtls_test_psa_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_status_t mbedtls_test_psa_raw_key_agreement_with_self(
+    psa_algorithm_t alg,
+    mbedtls_svc_key_id_t key );
+
+psa_status_t mbedtls_test_psa_key_agreement_with_self(
+    psa_key_derivation_operation_t *operation,
+    mbedtls_svc_key_id_t key );
+
+int mbedtls_test_psa_exported_key_sanity_check(
+    psa_key_type_t type, size_t bits,
+    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/psa_exercise_key.c b/tests/src/psa_exercise_key.c
index 78e46d9..79ba325 100644
--- a/tests/src/psa_exercise_key.c
+++ b/tests/src/psa_exercise_key.c
@@ -21,11 +21,876 @@
 
 #include <test/helpers.h>
 #include <test/macros.h>
-#include <test/psa_crypto_helpers.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;
+    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 );
+}
+
+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 );
+}
+
+int mbedtls_test_psa_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( !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,
+    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( ! 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( 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( ! 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( 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 = 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;
+
+    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 );
+}
+
+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/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index f0b8647..7ae6725 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -14,103 +14,11 @@
 
 #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)
-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
-
 /** Test if a buffer contains a constant byte value.
  *
  * `mem_is_char(buffer, c, size)` is true after `memset(buffer, c, size)`.
@@ -217,80 +125,6 @@
     return( len );
 }
 
-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;
-    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,
@@ -358,813 +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 );
-}
-
-int mbedtls_test_psa_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( !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,
-    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( ! 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( 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( ! 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( 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 = 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 );
-}
-
-/** 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 )
-{
-    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 );
-}
-
-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 );
-    }
-
-}
-
 static int test_operations_on_invalid_key( mbedtls_svc_key_id_t key )
 {
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;