/* BEGIN_HEADER */
#include <stdint.h>

#include "test/psa_crypto_helpers.h"
#include "psa_crypto_slot_management.h"
#include "psa_crypto_storage.h"

typedef enum
{
    /**< Close key(s) */
    INVALIDATE_BY_CLOSING,

    /**< Destroy key(s) */
    INVALIDATE_BY_DESTROYING,

    /**< Purge key(s) */
    INVALIDATE_BY_PURGING,

    /**< Terminate and reinitialize without closing/destroying keys */
    INVALIDATE_BY_SHUTDOWN,

    /**< Close key(s) then terminate and re-initialize */
    INVALIDATE_BY_CLOSING_WITH_SHUTDOWN,

    /**< Destroy key(s) then terminate and re-initialize */
    INVALIDATE_BY_DESTROYING_WITH_SHUTDOWN,

    /**< Purge key(s) then terminate and re-initialize */
    INVALIDATE_BY_PURGING_WITH_SHUTDOWN,
} invalidate_method_t;

typedef enum
{
    KEEP_OPEN,
    CLOSE_BEFORE,
    CLOSE_AFTER,
} reopen_policy_t;

typedef enum
{
    INVALID_HANDLE_0,
    INVALID_HANDLE_UNOPENED,
    INVALID_HANDLE_CLOSED,
    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;
 */
static int invalidate_key( invalidate_method_t invalidate_method,
                           mbedtls_svc_key_id_t key )
{
    switch( invalidate_method )
    {
        /* Closing the key invalidate only volatile keys, not persistent ones. */
        case INVALIDATE_BY_CLOSING:
        case INVALIDATE_BY_CLOSING_WITH_SHUTDOWN:
            PSA_ASSERT( psa_close_key( key ) );
            break;
        case INVALIDATE_BY_DESTROYING:
        case INVALIDATE_BY_DESTROYING_WITH_SHUTDOWN:
            PSA_ASSERT( psa_destroy_key( key ) );
            break;
        /* Purging the key just purges RAM data of persistent keys. */
        case INVALIDATE_BY_PURGING:
        case INVALIDATE_BY_PURGING_WITH_SHUTDOWN:
            PSA_ASSERT( psa_purge_key( key ) );
            break;
        case INVALIDATE_BY_SHUTDOWN:
            break;
    }
    return( 1 );
exit:
    return( 0 );
}

/** Restart the PSA subsystem if \p invalidate_method says so. */
static int invalidate_psa( invalidate_method_t invalidate_method )
{
    switch( invalidate_method )
    {
        case INVALIDATE_BY_CLOSING:
        case INVALIDATE_BY_DESTROYING:
        case INVALIDATE_BY_PURGING:
            return( 1 );
        case INVALIDATE_BY_CLOSING_WITH_SHUTDOWN:
        case INVALIDATE_BY_DESTROYING_WITH_SHUTDOWN:
        case INVALIDATE_BY_PURGING_WITH_SHUTDOWN:
            /* All keys must have been closed. */
            PSA_DONE( );
            break;
        case INVALIDATE_BY_SHUTDOWN:
            /* Some keys may remain behind, and we're testing that this
             * properly closes them. */
            mbedtls_psa_crypto_free( );
            break;
    }

    PSA_ASSERT( psa_crypto_init( ) );
    ASSERT_PSA_PRISTINE( );
    return( 1 );

exit:
    return( 0 );
}

/* END_HEADER */

/* BEGIN_DEPENDENCIES
 * depends_on:MBEDTLS_PSA_CRYPTO_C
 * END_DEPENDENCIES
 */

/* BEGIN_CASE */
void transient_slot_lifecycle( int owner_id_arg,
                               int usage_arg, int alg_arg,
                               int type_arg, data_t *key_data,
                               int invalidate_method_arg )
{
    psa_algorithm_t alg = alg_arg;
    psa_key_usage_t usage_flags = usage_arg;
    psa_key_type_t type = type_arg;
    invalidate_method_t invalidate_method = invalidate_method_arg;
    mbedtls_svc_key_id_t key = MBEDTLS_SVC_KEY_ID_INIT;
    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;

    PSA_ASSERT( psa_crypto_init( ) );

    /* Import a key. */
#if defined(MBEDTLS_PSA_CRYPTO_KEY_ID_ENCODES_OWNER)
    mbedtls_key_owner_id_t owner_id = owner_id_arg;

    mbedtls_set_key_owner_id( &attributes, owner_id );
#else
    (void)owner_id_arg;
#endif

    psa_set_key_usage_flags( &attributes, usage_flags );
    psa_set_key_algorithm( &attributes, alg );
    psa_set_key_type( &attributes, type );
    PSA_ASSERT( psa_import_key( &attributes, key_data->x, key_data->len,
                                &key ) );
    TEST_ASSERT( ! mbedtls_svc_key_id_is_null( key ) );
    PSA_ASSERT( psa_get_key_attributes( key, &attributes ) );
    TEST_EQUAL( psa_get_key_type( &attributes ), type );
    psa_reset_key_attributes( &attributes );

#if defined(MBEDTLS_PSA_CRYPTO_KEY_ID_ENCODES_OWNER)
    {
        psa_key_handle_t handle;
        mbedtls_svc_key_id_t key_with_invalid_owner =
            mbedtls_svc_key_id_make( owner_id + 1,
                                     MBEDTLS_SVC_KEY_ID_GET_KEY_ID( key ) );

        TEST_ASSERT( mbedtls_key_owner_id_equal(
                         owner_id,
                         MBEDTLS_SVC_KEY_ID_GET_OWNER_ID( key ) ) );
        TEST_EQUAL( psa_open_key( key_with_invalid_owner, &handle ),
                    PSA_ERROR_DOES_NOT_EXIST );
    }
#endif

    /*
     * Purge the key and make sure that it is still valid, as purging a
     * volatile key shouldn't invalidate/destroy it.
     */
    PSA_ASSERT( psa_purge_key( key ) );
    PSA_ASSERT( psa_get_key_attributes( key, &attributes ) );
    TEST_EQUAL( psa_get_key_type( &attributes ), type );
    psa_reset_key_attributes( &attributes );

    /* Do something that invalidates the key. */
    if( ! invalidate_key( invalidate_method, key ) )
        goto exit;
    if( ! invalidate_psa( invalidate_method ) )
        goto exit;

    /* Test that the key is now invalid. */
    TEST_EQUAL( psa_get_key_attributes( key, &attributes ),
                PSA_ERROR_DOES_NOT_EXIST );
    TEST_EQUAL( psa_close_key( key ), PSA_ERROR_DOES_NOT_EXIST );

exit:
    /*
     * Key attributes may have been returned by psa_get_key_attributes()
     * thus reset them as required.
     */
    psa_reset_key_attributes( &attributes );

    PSA_DONE( );
}
/* END_CASE */

/* BEGIN_CASE depends_on:MBEDTLS_PSA_CRYPTO_STORAGE_C */
void persistent_slot_lifecycle( int lifetime_arg, int owner_id_arg, int id_arg,
                                int usage_arg, int alg_arg, int alg2_arg,
                                int type_arg, data_t *key_data,
                                int invalidate_method_arg )
{
    psa_key_lifetime_t lifetime = lifetime_arg;
    mbedtls_svc_key_id_t id = mbedtls_svc_key_id_make( owner_id_arg, id_arg );
    psa_algorithm_t alg = alg_arg;
    psa_algorithm_t alg2 = alg2_arg;
    psa_key_usage_t usage_flags = usage_arg;
    psa_key_type_t type = type_arg;
    invalidate_method_t invalidate_method = invalidate_method_arg;
    mbedtls_svc_key_id_t returned_id = MBEDTLS_SVC_KEY_ID_INIT;
    psa_key_handle_t handle = PSA_KEY_HANDLE_INIT;
    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
    psa_key_attributes_t read_attributes = PSA_KEY_ATTRIBUTES_INIT;
    uint8_t *reexported = NULL;
    size_t reexported_length = -1;

#if defined(MBEDTLS_PSA_CRYPTO_KEY_ID_ENCODES_OWNER)
    mbedtls_svc_key_id_t wrong_owner_id =
        mbedtls_svc_key_id_make( owner_id_arg + 1, id_arg );
    mbedtls_svc_key_id_t invalid_svc_key_id = MBEDTLS_SVC_KEY_ID_INIT;
#endif

    TEST_USES_KEY_ID( id );

    PSA_ASSERT( psa_crypto_init( ) );

    psa_set_key_id( &attributes, id );
    psa_set_key_lifetime( &attributes, lifetime );
    psa_set_key_type( &attributes, type );
    psa_set_key_usage_flags( &attributes, usage_flags );
    psa_set_key_algorithm( &attributes, alg );
    psa_set_key_enrollment_algorithm( &attributes, alg2 );
    PSA_ASSERT( psa_import_key( &attributes, key_data->x, key_data->len,
                                &returned_id ) );
    TEST_ASSERT( mbedtls_svc_key_id_equal( id, returned_id ) );

#if defined(MBEDTLS_PSA_CRYPTO_KEY_ID_ENCODES_OWNER)
    TEST_EQUAL( psa_open_key( wrong_owner_id, &invalid_svc_key_id ),
                PSA_ERROR_DOES_NOT_EXIST );
#endif

    PSA_ASSERT( psa_get_key_attributes( id, &attributes ) );
    TEST_EQUAL( psa_get_key_lifetime( &attributes ), lifetime );
    TEST_ASSERT( mbedtls_svc_key_id_equal(
                     psa_get_key_id( &attributes ), id ) );
    TEST_EQUAL( psa_get_key_usage_flags( &attributes ), usage_flags );
    TEST_EQUAL( psa_get_key_algorithm( &attributes ), alg );
    TEST_EQUAL( psa_get_key_enrollment_algorithm( &attributes ), alg2 );
    TEST_EQUAL( psa_get_key_type( &attributes ), type );

    /* Close the key and then open it. */
    PSA_ASSERT( psa_close_key( id ) );

#if defined(MBEDTLS_PSA_CRYPTO_KEY_ID_ENCODES_OWNER)
    TEST_EQUAL( psa_open_key( wrong_owner_id, &invalid_svc_key_id ),
                PSA_ERROR_DOES_NOT_EXIST );
#endif

    PSA_ASSERT( psa_open_key( id, &handle ) );
    TEST_ASSERT( ! psa_key_handle_is_null( handle ) );
    PSA_ASSERT( psa_get_key_attributes( handle, &attributes ) );
    TEST_EQUAL( psa_get_key_lifetime( &attributes ), lifetime );
    TEST_ASSERT( mbedtls_svc_key_id_equal(
                     psa_get_key_id( &attributes ), id ) );
    TEST_EQUAL( psa_get_key_usage_flags( &attributes ), usage_flags );
    TEST_EQUAL( psa_get_key_algorithm( &attributes ), alg );
    TEST_EQUAL( psa_get_key_enrollment_algorithm( &attributes ), alg2 );
    TEST_EQUAL( psa_get_key_type( &attributes ), type );

    /*
     * Do something that wipes key data in volatile memory or destroy the
     * key.
     */
    if( ! invalidate_key( invalidate_method, id ) )
        goto exit;
    if( ! invalidate_psa( invalidate_method ) )
        goto exit;

    /* Try to reaccess the key. If we destroyed it, check that it doesn't
     * exist. Otherwise check that it still exists and has the expected
     * content. */
    switch( invalidate_method )
    {
        case INVALIDATE_BY_CLOSING:
        case INVALIDATE_BY_CLOSING_WITH_SHUTDOWN:
        case INVALIDATE_BY_PURGING:
        case INVALIDATE_BY_PURGING_WITH_SHUTDOWN:
        case INVALIDATE_BY_SHUTDOWN:
            PSA_ASSERT( psa_open_key( id, &handle ) );
            PSA_ASSERT( psa_get_key_attributes( id, &read_attributes ) );
            TEST_EQUAL( psa_get_key_lifetime( &attributes ),
                        psa_get_key_lifetime( &read_attributes ) );
            TEST_ASSERT( mbedtls_svc_key_id_equal(
                             psa_get_key_id( &attributes ),
                             psa_get_key_id( &read_attributes ) ) );
            TEST_EQUAL( psa_get_key_usage_flags( &attributes ), usage_flags );
            TEST_EQUAL( psa_get_key_algorithm( &attributes ),
                        psa_get_key_algorithm( &read_attributes ) );
            TEST_EQUAL( psa_get_key_enrollment_algorithm( &attributes ),
                        psa_get_key_enrollment_algorithm( &read_attributes ) );
            TEST_EQUAL( psa_get_key_type( &attributes ),
                        psa_get_key_type( &read_attributes ) );
            TEST_EQUAL( psa_get_key_bits( &attributes ),
                        psa_get_key_bits( &read_attributes ) );
            if( usage_flags & PSA_KEY_USAGE_EXPORT )
            {
                ASSERT_ALLOC( reexported, key_data->len );
                PSA_ASSERT( psa_export_key( id, reexported, key_data->len,
                                            &reexported_length ) );
                ASSERT_COMPARE( key_data->x, key_data->len,
                                reexported, reexported_length );
            }
            else
            {
                TEST_EQUAL( psa_export_key( id, NULL, 0, &reexported_length ),
                            PSA_ERROR_NOT_PERMITTED );
            }
            PSA_ASSERT( psa_close_key( handle ) );
            break;

        case INVALIDATE_BY_DESTROYING:
        case INVALIDATE_BY_DESTROYING_WITH_SHUTDOWN:
            /*
             * Test that the key handle and identifier are now not refering to an
             * existing key.
             */
            TEST_EQUAL( psa_get_key_attributes( handle, &read_attributes ),
                        PSA_ERROR_DOES_NOT_EXIST );
            TEST_EQUAL( psa_close_key( handle ), PSA_ERROR_DOES_NOT_EXIST );
            TEST_EQUAL( psa_get_key_attributes( id, &read_attributes ),
                        PSA_ERROR_DOES_NOT_EXIST );
            break;
    }

exit:
    /*
     * Key attributes may have been returned by psa_get_key_attributes()
     * thus reset them as required.
     */
    psa_reset_key_attributes( &attributes );
    psa_reset_key_attributes( &read_attributes );

    PSA_DONE( );
    psa_purge_key_storage( );
    mbedtls_free( reexported );
}
/* END_CASE */

/* BEGIN_CASE depends_on:MBEDTLS_PSA_CRYPTO_STORAGE_C */
void create_existent( int lifetime_arg, int owner_id_arg, int id_arg,
                      int reopen_policy_arg )
{
    psa_key_lifetime_t lifetime = lifetime_arg;
    mbedtls_svc_key_id_t id = mbedtls_svc_key_id_make( owner_id_arg, id_arg );
    mbedtls_svc_key_id_t returned_id = MBEDTLS_SVC_KEY_ID_INIT;
    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
    psa_key_type_t type1 = PSA_KEY_TYPE_RAW_DATA;
    const uint8_t material1[5] = "a key";
    const uint8_t material2[5] = "b key";
    size_t bits1 = PSA_BYTES_TO_BITS( sizeof( material1 ) );
    uint8_t reexported[sizeof( material1 )];
    size_t reexported_length;
    reopen_policy_t reopen_policy = reopen_policy_arg;

    TEST_USES_KEY_ID( id );

    PSA_ASSERT( psa_crypto_init( ) );

    /* Create a key. */
    psa_set_key_id( &attributes, id );
    psa_set_key_lifetime( &attributes, lifetime );
    psa_set_key_type( &attributes, type1 );
    psa_set_key_usage_flags( &attributes, PSA_KEY_USAGE_EXPORT );
    psa_set_key_algorithm( &attributes, 0 );
    PSA_ASSERT( psa_import_key( &attributes, material1, sizeof( material1 ),
                                &returned_id ) );
    TEST_ASSERT( mbedtls_svc_key_id_equal( id, returned_id ) );

    if( reopen_policy == CLOSE_BEFORE )
        PSA_ASSERT( psa_close_key( id ) );

    /* Attempt to create a new key in the same slot. */
    TEST_EQUAL( psa_import_key( &attributes, material2, sizeof( material2 ),
                                &returned_id ),
                PSA_ERROR_ALREADY_EXISTS );
    TEST_ASSERT( mbedtls_svc_key_id_is_null( returned_id ) );

    if( reopen_policy == CLOSE_AFTER )
        PSA_ASSERT( psa_close_key( id ) );

    /* Check that the original key hasn't changed. */
    psa_reset_key_attributes( &attributes );
    PSA_ASSERT( psa_get_key_attributes( id, &attributes ) );
    TEST_ASSERT( mbedtls_svc_key_id_equal(
                     psa_get_key_id( &attributes ), id ) );
    TEST_EQUAL( psa_get_key_lifetime( &attributes ), lifetime );
    TEST_EQUAL( psa_get_key_type( &attributes ), type1 );
    TEST_EQUAL( psa_get_key_bits( &attributes ), bits1 );
    TEST_EQUAL( psa_get_key_usage_flags( &attributes ), PSA_KEY_USAGE_EXPORT );
    TEST_EQUAL( psa_get_key_algorithm( &attributes ), 0 );

    PSA_ASSERT( psa_export_key( id,
                                reexported, sizeof( reexported ),
                                &reexported_length ) );
    ASSERT_COMPARE( material1, sizeof( material1 ),
                    reexported, reexported_length );

    PSA_ASSERT( psa_close_key( id ) );

exit:
    /*
     * Key attributes may have been returned by psa_get_key_attributes()
     * thus reset them as required.
     */
    psa_reset_key_attributes( &attributes );

    PSA_DONE( );
    psa_purge_key_storage( );
}
/* END_CASE */

/* BEGIN_CASE */
void open_fail( int id_arg,
                int expected_status_arg )
{
    mbedtls_svc_key_id_t id = mbedtls_svc_key_id_make( 1, id_arg );
    psa_status_t expected_status = expected_status_arg;
    psa_key_handle_t handle = mbedtls_svc_key_id_make( 0xdead, 0xdead );

    PSA_ASSERT( psa_crypto_init( ) );

    TEST_EQUAL( psa_open_key( id, &handle ), expected_status );
    TEST_ASSERT( psa_key_handle_is_null( handle ) );

exit:
    PSA_DONE( );
}
/* END_CASE */

/* BEGIN_CASE */
void create_fail( int lifetime_arg, int id_arg,
                  int expected_status_arg )
{
    psa_key_lifetime_t lifetime = lifetime_arg;
    mbedtls_svc_key_id_t id = mbedtls_svc_key_id_make( 1, id_arg );
    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
    psa_status_t expected_status = expected_status_arg;
    mbedtls_svc_key_id_t returned_id =
        mbedtls_svc_key_id_make( 0xdead, 0xdead );
    uint8_t material[1] = {'k'};

    TEST_USES_KEY_ID( id );

    PSA_ASSERT( psa_crypto_init( ) );

    psa_set_key_lifetime( &attributes, lifetime );
    if( PSA_KEY_LIFETIME_IS_VOLATILE( lifetime ) )
    {
        /*
         * Not possible to set a key identifier different from 0 through
         * PSA key attributes APIs thus accessing to the attributes
         * directly.
         */
        attributes.core.id = id;
    }
    else
        psa_set_key_id( &attributes, id );

    psa_set_key_type( &attributes, PSA_KEY_TYPE_RAW_DATA );
    TEST_EQUAL( psa_import_key( &attributes, material, sizeof( material ),
                                &returned_id ),
                expected_status );
    TEST_ASSERT( mbedtls_svc_key_id_is_null( returned_id ) );

exit:
    PSA_DONE( );
#if defined(MBEDTLS_PSA_CRYPTO_STORAGE_C)
    psa_purge_key_storage( );
#endif
}
/* END_CASE */

/* BEGIN_CASE */
void copy_across_lifetimes( int source_lifetime_arg, int source_owner_id_arg,
                            int source_id_arg, int source_usage_arg,
                            int source_alg_arg, int source_alg2_arg,
                            int type_arg, data_t *material,
                            int target_lifetime_arg, int target_owner_id_arg,
                            int target_id_arg, int target_usage_arg,
                            int target_alg_arg, int target_alg2_arg,
                            int expected_usage_arg,
                            int expected_alg_arg, int expected_alg2_arg )
{
    psa_key_lifetime_t source_lifetime = source_lifetime_arg;
    mbedtls_svc_key_id_t source_id =
        mbedtls_svc_key_id_make( source_owner_id_arg, source_id_arg );
    psa_key_usage_t source_usage = source_usage_arg;
    psa_algorithm_t source_alg = source_alg_arg;
    psa_key_attributes_t source_attributes = PSA_KEY_ATTRIBUTES_INIT;
    psa_key_type_t source_type = type_arg;
    mbedtls_svc_key_id_t returned_source_id = MBEDTLS_SVC_KEY_ID_INIT;
    psa_key_lifetime_t target_lifetime = target_lifetime_arg;
    mbedtls_svc_key_id_t target_id =
        mbedtls_svc_key_id_make( target_owner_id_arg, target_id_arg );
    psa_key_usage_t target_usage = target_usage_arg;
    psa_algorithm_t target_alg = target_alg_arg;
    psa_key_attributes_t target_attributes = PSA_KEY_ATTRIBUTES_INIT;
    mbedtls_svc_key_id_t returned_target_id = MBEDTLS_SVC_KEY_ID_INIT;
    psa_key_handle_t target_handle = PSA_KEY_HANDLE_INIT;
    psa_key_usage_t expected_usage = expected_usage_arg;
    psa_algorithm_t expected_alg = expected_alg_arg;
    psa_algorithm_t expected_alg2 = expected_alg2_arg;
    uint8_t *export_buffer = NULL;

    TEST_USES_KEY_ID( source_id );
    TEST_USES_KEY_ID( target_id );

    PSA_ASSERT( psa_crypto_init( ) );

    /* Populate the source slot. */
    psa_set_key_id( &source_attributes, source_id );
    psa_set_key_lifetime( &source_attributes, source_lifetime );

    psa_set_key_type( &source_attributes, source_type );
    psa_set_key_usage_flags( &source_attributes, source_usage );
    psa_set_key_algorithm( &source_attributes, source_alg );
    psa_set_key_enrollment_algorithm( &source_attributes, source_alg2_arg );
    PSA_ASSERT( psa_import_key( &source_attributes,
                                material->x, material->len,
                                &returned_source_id ) );
    /* Update the attributes with the bit size. */
    PSA_ASSERT( psa_get_key_attributes( returned_source_id,
                                        &source_attributes ) );

    /* Prepare the target slot. */
    psa_set_key_id( &target_attributes, target_id );
    psa_set_key_lifetime( &target_attributes, target_lifetime );

    psa_set_key_usage_flags( &target_attributes, target_usage );
    psa_set_key_algorithm( &target_attributes, target_alg );
    psa_set_key_enrollment_algorithm( &target_attributes, target_alg2_arg );

    /* Copy the key. */
    PSA_ASSERT( psa_copy_key( returned_source_id,
                              &target_attributes, &returned_target_id ) );

    /* Destroy the source to ensure that this doesn't affect the target. */
    PSA_ASSERT( psa_destroy_key( returned_source_id ) );

    /* If the target key is persistent, restart the system to make
     * sure that the material is still alive. */
    if( ! PSA_KEY_LIFETIME_IS_VOLATILE( target_lifetime ) )
    {
        mbedtls_psa_crypto_free( );
        PSA_ASSERT( psa_crypto_init( ) );
        PSA_ASSERT( psa_open_key( target_id, &target_handle ) );
    }

    /* Test that the target slot has the expected content. */
    psa_reset_key_attributes( &target_attributes );
    PSA_ASSERT( psa_get_key_attributes( returned_target_id,
                                        &target_attributes ) );

    if( ! PSA_KEY_LIFETIME_IS_VOLATILE( target_lifetime ) )
    {
        TEST_ASSERT( mbedtls_svc_key_id_equal(
                         target_id, psa_get_key_id( &target_attributes ) ) );
    }
    else
    {
#if defined(MBEDTLS_PSA_CRYPTO_KEY_ID_ENCODES_OWNER)
        TEST_EQUAL( MBEDTLS_SVC_KEY_ID_GET_OWNER_ID( returned_target_id ),
                    target_owner_id_arg );
#endif
    }

    TEST_EQUAL( target_lifetime, psa_get_key_lifetime( &target_attributes ) );
    TEST_EQUAL( source_type, psa_get_key_type( &target_attributes ) );
    TEST_EQUAL( psa_get_key_bits( &source_attributes ),
                psa_get_key_bits( &target_attributes ) );
    TEST_EQUAL( expected_usage, psa_get_key_usage_flags( &target_attributes ) );
    TEST_EQUAL( expected_alg, psa_get_key_algorithm( &target_attributes ) );
    TEST_EQUAL( expected_alg2,
                psa_get_key_enrollment_algorithm( &target_attributes ) );
    if( expected_usage & PSA_KEY_USAGE_EXPORT )
    {
        size_t length;
        ASSERT_ALLOC( export_buffer, material->len );
        PSA_ASSERT( psa_export_key( returned_target_id, export_buffer,
                                    material->len, &length ) );
        ASSERT_COMPARE( material->x, material->len,
                        export_buffer, length );
    }
    else
    {
        size_t length;
        /* Check that the key is actually non-exportable. */
        TEST_EQUAL( psa_export_key( returned_target_id, export_buffer,
                                    material->len, &length ),
                    PSA_ERROR_NOT_PERMITTED );
    }

    PSA_ASSERT( psa_destroy_key( returned_target_id ) );

exit:
    /*
     * Source and target key attributes may have been returned by
     * psa_get_key_attributes() thus reset them as required.
     */
    psa_reset_key_attributes( &source_attributes );
    psa_reset_key_attributes( &target_attributes );

    PSA_DONE( );
    mbedtls_free( export_buffer );
#if defined(MBEDTLS_PSA_CRYPTO_STORAGE_C)
    psa_purge_key_storage( );
#endif
}
/* END_CASE */

/* BEGIN_CASE */
void copy_to_occupied( int source_lifetime_arg, int source_id_arg,
                       int source_usage_arg, int source_alg_arg,
                       int source_type_arg, data_t *source_material,
                       int target_lifetime_arg, int target_id_arg,
                       int target_usage_arg, int target_alg_arg,
                       int target_type_arg, data_t *target_material )
{
    psa_key_lifetime_t source_lifetime = source_lifetime_arg;
    mbedtls_svc_key_id_t source_id =
        mbedtls_svc_key_id_make( 1, source_id_arg );
    psa_key_usage_t source_usage = source_usage_arg;
    psa_algorithm_t source_alg = source_alg_arg;
    psa_key_type_t source_type = source_type_arg;
    mbedtls_svc_key_id_t returned_source_id = MBEDTLS_SVC_KEY_ID_INIT;
    psa_key_lifetime_t target_lifetime = target_lifetime_arg;
    mbedtls_svc_key_id_t target_id =
        mbedtls_svc_key_id_make( 1, target_id_arg );
    psa_key_usage_t target_usage = target_usage_arg;
    psa_algorithm_t target_alg = target_alg_arg;
    psa_key_type_t target_type = target_type_arg;
    mbedtls_svc_key_id_t returned_target_id = MBEDTLS_SVC_KEY_ID_INIT;
    mbedtls_svc_key_id_t new_key = MBEDTLS_SVC_KEY_ID_INIT;
    uint8_t *export_buffer = NULL;
    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
    psa_key_attributes_t attributes1 = PSA_KEY_ATTRIBUTES_INIT;
    psa_key_attributes_t attributes2 = PSA_KEY_ATTRIBUTES_INIT;

    TEST_USES_KEY_ID( source_id );
    TEST_USES_KEY_ID( target_id );

    PSA_ASSERT( psa_crypto_init( ) );

    /* Populate the source slot. */
    if( ! PSA_KEY_LIFETIME_IS_VOLATILE( source_lifetime ) )
    {
        psa_set_key_id( &attributes, source_id );
        psa_set_key_lifetime( &attributes, source_lifetime );
    }
    psa_set_key_type( &attributes, source_type );
    psa_set_key_usage_flags( &attributes, source_usage );
    psa_set_key_algorithm( &attributes, source_alg );
    PSA_ASSERT( psa_import_key( &attributes,
                                source_material->x, source_material->len,
                                &returned_source_id ) );

    /* Populate the target slot. */
    if( mbedtls_svc_key_id_equal( target_id, source_id ) )
    {
        returned_target_id = returned_source_id;
    }
    else
    {
        psa_set_key_id( &attributes1, target_id );
        psa_set_key_lifetime( &attributes1, target_lifetime );
        psa_set_key_type( &attributes1, target_type );
        psa_set_key_usage_flags( &attributes1, target_usage );
        psa_set_key_algorithm( &attributes1, target_alg );
        PSA_ASSERT( psa_import_key( &attributes1,
                                    target_material->x, target_material->len,
                                    &returned_target_id ) );
    }

    PSA_ASSERT( psa_get_key_attributes( returned_target_id, &attributes1 ) );

    /* Make a copy attempt. */
    psa_set_key_id( &attributes, target_id );
    psa_set_key_lifetime( &attributes, target_lifetime );
    TEST_EQUAL( psa_copy_key( returned_source_id,
                              &attributes, &new_key ),
                PSA_ERROR_ALREADY_EXISTS );
    TEST_ASSERT( mbedtls_svc_key_id_is_null( new_key ) );

    /* Test that the target slot is unaffected. */
    PSA_ASSERT( psa_get_key_attributes( returned_target_id, &attributes2 ) );
    TEST_ASSERT( mbedtls_svc_key_id_equal(
                     psa_get_key_id( &attributes1 ),
                     psa_get_key_id( &attributes2 ) ) );
    TEST_EQUAL( psa_get_key_lifetime( &attributes1 ),
                psa_get_key_lifetime( &attributes2 ) );
    TEST_EQUAL( psa_get_key_type( &attributes1 ),
                psa_get_key_type( &attributes2 ) );
    TEST_EQUAL( psa_get_key_bits( &attributes1 ),
                psa_get_key_bits( &attributes2 ) );
    TEST_EQUAL( psa_get_key_usage_flags( &attributes1 ),
                psa_get_key_usage_flags( &attributes2 ) );
    TEST_EQUAL( psa_get_key_algorithm( &attributes1 ),
                psa_get_key_algorithm( &attributes2 ) );
    if( target_usage & PSA_KEY_USAGE_EXPORT )
    {
        size_t length;
        ASSERT_ALLOC( export_buffer, target_material->len );
        PSA_ASSERT( psa_export_key( returned_target_id, export_buffer,
                                    target_material->len, &length ) );
        ASSERT_COMPARE( target_material->x, target_material->len,
                        export_buffer, length );
    }

    PSA_ASSERT( psa_destroy_key( returned_source_id ) );
    if( ! mbedtls_svc_key_id_equal( target_id, source_id ) )
        PSA_ASSERT( psa_destroy_key( returned_target_id ) );

exit:
    /*
     * Key attributes may have been returned by psa_get_key_attributes()
     * thus reset them as required.
     */
    psa_reset_key_attributes( &attributes1 );
    psa_reset_key_attributes( &attributes2 );

    PSA_DONE( );
    mbedtls_free( export_buffer );
#if defined(MBEDTLS_PSA_CRYPTO_STORAGE_C)
    psa_purge_key_storage( );
#endif
}
/* END_CASE */

/* BEGIN_CASE */
void invalid_handle( int handle_construction,
                     int close_status_arg, int usage_status_arg )
{
    psa_key_handle_t valid_handle = PSA_KEY_HANDLE_INIT;
    psa_key_handle_t invalid_handle = PSA_KEY_HANDLE_INIT;
    psa_key_id_t key_id;
    psa_status_t close_status = close_status_arg;
    psa_status_t usage_status = usage_status_arg;
    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
    uint8_t material[1] = "a";

    PSA_ASSERT( psa_crypto_init( ) );

    /* Allocate a handle and store a key in it. */
    psa_set_key_type( &attributes, PSA_KEY_TYPE_RAW_DATA );
    psa_set_key_usage_flags( &attributes, 0 );
    psa_set_key_algorithm( &attributes, 0 );
    PSA_ASSERT( psa_import_key( &attributes,
                                material, sizeof( material ),
                                &valid_handle ) );
    TEST_ASSERT( ! psa_key_handle_is_null( valid_handle ) );

    /* Construct an invalid handle as specified in the test case data. */
    switch( handle_construction )
    {
        case INVALID_HANDLE_0:
            invalid_handle = PSA_KEY_HANDLE_INIT;
            break;
        case INVALID_HANDLE_UNOPENED:

            /*
             * MBEDTLS_SVC_KEY_ID_GET_KEY_ID( valid_handle ) is a volatile
             * key identifier as the imported key is a volatile key. Volatile
             * key identifiers are in the range from PSA_KEY_ID_VOLATILE_MIN
             * to PSA_KEY_ID_VOLATILE_MAX included. Thus pick a key identifier
             * in the range from PSA_KEY_ID_VOLATILE_MIN to
             * PSA_KEY_ID_VOLATILE_MAX different from
             * MBEDTLS_SVC_KEY_ID_GET_KEY_ID( valid_handle ) to build an
             * unopened and thus invalid identifier.
             */

            if( MBEDTLS_SVC_KEY_ID_GET_KEY_ID( valid_handle ) ==
                PSA_KEY_ID_VOLATILE_MIN )
                key_id = PSA_KEY_ID_VOLATILE_MIN + 1;
            else
                key_id = MBEDTLS_SVC_KEY_ID_GET_KEY_ID( valid_handle ) - 1;

            invalid_handle =
                mbedtls_svc_key_id_make( 0, key_id );
            break;
        case INVALID_HANDLE_CLOSED:
            PSA_ASSERT( psa_import_key( &attributes,
                                        material, sizeof( material ),
                                        &invalid_handle ) );
            PSA_ASSERT( psa_destroy_key( invalid_handle ) );
            break;
        case INVALID_HANDLE_HUGE:
            invalid_handle =
                mbedtls_svc_key_id_make( 0, PSA_KEY_ID_VENDOR_MAX + 1 );
            break;
        default:
            TEST_ASSERT( ! "unknown handle construction" );
    }

    /* Attempt to use the invalid handle. */
    TEST_EQUAL( psa_get_key_attributes( invalid_handle, &attributes ),
                usage_status );
    TEST_EQUAL( psa_close_key( invalid_handle ), close_status );
    TEST_EQUAL( psa_destroy_key( invalid_handle ), close_status );

    /* After all this, check that the original handle is intact. */
    PSA_ASSERT( psa_get_key_attributes( valid_handle, &attributes ) );
    TEST_EQUAL( psa_get_key_type( &attributes ), PSA_KEY_TYPE_RAW_DATA );
    TEST_EQUAL( psa_get_key_bits( &attributes ),
                PSA_BYTES_TO_BITS( sizeof( material ) ) );
    PSA_ASSERT( psa_close_key( valid_handle ) );

exit:
    /*
     * Key attributes may have been returned by psa_get_key_attributes()
     * thus reset them as required.
     */
    psa_reset_key_attributes( &attributes );

    PSA_DONE( );
}
/* END_CASE */

/* BEGIN_CASE */
void many_transient_keys( int max_keys_arg )
{
    mbedtls_svc_key_id_t *keys = NULL;
    size_t max_keys = max_keys_arg;
    size_t i, j;
    psa_status_t status;
    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
    uint8_t exported[sizeof( size_t )];
    size_t exported_length;

    ASSERT_ALLOC( keys, max_keys );
    PSA_ASSERT( psa_crypto_init( ) );

    psa_set_key_usage_flags( &attributes, PSA_KEY_USAGE_EXPORT );
    psa_set_key_algorithm( &attributes, 0 );
    psa_set_key_type( &attributes, PSA_KEY_TYPE_RAW_DATA );

    for( i = 0; i < max_keys; i++ )
    {
        status = psa_import_key( &attributes,
                                 (uint8_t *) &i, sizeof( i ),
                                 &keys[i] );
        if( status == PSA_ERROR_INSUFFICIENT_MEMORY )
            break;
        PSA_ASSERT( status );
        TEST_ASSERT( ! mbedtls_svc_key_id_is_null( keys[i] ) );
        for( j = 0; j < i; j++ )
            TEST_ASSERT( ! mbedtls_svc_key_id_equal( keys[i], keys[j] ) );
    }
    max_keys = i;

    for( i = 1; i < max_keys; i++ )
    {
        PSA_ASSERT( psa_close_key( keys[i - 1] ) );
        PSA_ASSERT( psa_export_key( keys[i],
                                    exported, sizeof( exported ),
                                    &exported_length ) );
        ASSERT_COMPARE( exported, exported_length,
                        (uint8_t *) &i, sizeof( i ) );
    }
    PSA_ASSERT( psa_close_key( keys[i - 1] ) );

exit:
    PSA_DONE( );
    mbedtls_free( keys );
}
/* END_CASE */

/* BEGIN_CASE depends_on:MBEDTLS_PSA_CRYPTO_STORAGE_C */
void key_slot_eviction_to_import_new_key( int lifetime_arg )
{
    psa_key_lifetime_t lifetime = (psa_key_lifetime_t)lifetime_arg;
    size_t i;
    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
    uint8_t exported[sizeof( size_t )];
    size_t exported_length;
    mbedtls_svc_key_id_t key, returned_key_id;

    PSA_ASSERT( psa_crypto_init( ) );

    psa_set_key_usage_flags( &attributes, PSA_KEY_USAGE_EXPORT );
    psa_set_key_algorithm( &attributes, 0 );
    psa_set_key_type( &attributes, PSA_KEY_TYPE_RAW_DATA );

    /*
     * Create PSA_KEY_SLOT_COUNT persistent keys.
     */
    for( i = 0; i < PSA_KEY_SLOT_COUNT; i++ )
    {
        key = mbedtls_svc_key_id_make( i, i + 1 );
        psa_set_key_id( &attributes, key );
        PSA_ASSERT( psa_import_key( &attributes,
                                    (uint8_t *) &i, sizeof( i ),
                                    &returned_key_id ) );
       TEST_ASSERT( mbedtls_svc_key_id_equal( returned_key_id, key ) );
    }

    /*
     * Create a new persistent or volatile key. When creating the key,
     * one of the descriptions of the previously created persistent keys
     * is removed from the RAM key slots. This makes room to store its
     * description in RAM.
     */
    i = PSA_KEY_SLOT_COUNT;
    key = mbedtls_svc_key_id_make( i, i + 1 );
    psa_set_key_id( &attributes, key );
    psa_set_key_lifetime( &attributes, lifetime );

    PSA_ASSERT( psa_import_key( &attributes,
                                (uint8_t *) &i, sizeof( i ),
                                &returned_key_id ) );
    if( lifetime != PSA_KEY_LIFETIME_VOLATILE )
        TEST_ASSERT( mbedtls_svc_key_id_equal( returned_key_id, key ) );
    else
        TEST_ASSERT( psa_key_id_is_volatile(
                     MBEDTLS_SVC_KEY_ID_GET_KEY_ID( returned_key_id ) ) );

    /*
     * Check that we can export all ( PSA_KEY_SLOT_COUNT + 1 ) keys,
     * that they have the expected value and destroy them. In that process,
     * the description of the persistent key that was evicted from the RAM
     * slots when creating the last key is restored in a RAM slot to export
     * its value.
     */
    for( i = 0; i <= PSA_KEY_SLOT_COUNT; i++ )
    {
        if( i < PSA_KEY_SLOT_COUNT )
            key = mbedtls_svc_key_id_make( i, i + 1 );
        else
            key = returned_key_id;

        PSA_ASSERT( psa_export_key( key,
                                    exported, sizeof( exported ),
                                    &exported_length ) );
        ASSERT_COMPARE( exported, exported_length,
                        (uint8_t *) &i, sizeof( i ) );
        PSA_ASSERT( psa_destroy_key( key ) );
    }

exit:
    PSA_DONE( );
}
/* END_CASE */

/* BEGIN_CASE depends_on:MBEDTLS_PSA_CRYPTO_STORAGE_C */
void non_reusable_key_slots_integrity_in_case_of_key_slot_starvation( )
{
    psa_status_t status;
    size_t i;
    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
    uint8_t exported[sizeof( size_t )];
    size_t exported_length;
    mbedtls_svc_key_id_t persistent_key = MBEDTLS_SVC_KEY_ID_INIT;
    mbedtls_svc_key_id_t persistent_key2 = MBEDTLS_SVC_KEY_ID_INIT;
    mbedtls_svc_key_id_t returned_key_id = MBEDTLS_SVC_KEY_ID_INIT;
    mbedtls_svc_key_id_t *keys = NULL;

    TEST_ASSERT( PSA_KEY_SLOT_COUNT >= 1 );

    ASSERT_ALLOC( keys, PSA_KEY_SLOT_COUNT );
    PSA_ASSERT( psa_crypto_init( ) );

    psa_set_key_usage_flags( &attributes,
                             PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_COPY );
    psa_set_key_algorithm( &attributes, 0 );
    psa_set_key_type( &attributes, PSA_KEY_TYPE_RAW_DATA );

    /*
     * Create a persistent key
     */
    persistent_key = mbedtls_svc_key_id_make( 0x100, 0x205 );
    psa_set_key_id( &attributes, persistent_key );
    PSA_ASSERT( psa_import_key( &attributes,
                                (uint8_t *) &persistent_key,
                                sizeof( persistent_key ),
                                &returned_key_id ) );
    TEST_ASSERT( mbedtls_svc_key_id_equal( returned_key_id, persistent_key ) );

    /*
     * Create PSA_KEY_SLOT_COUNT volatile keys
     */
    psa_set_key_lifetime( &attributes, PSA_KEY_LIFETIME_VOLATILE );
    for( i = 0; i < PSA_KEY_SLOT_COUNT; i++ )
    {
        PSA_ASSERT( psa_import_key( &attributes,
                                    (uint8_t *) &i, sizeof( i ),
                                    &keys[i]) );
    }
    psa_reset_key_attributes( &attributes );

    /*
     * Check that we cannot access the persistent key as all slots are
     * occupied by volatile keys and the implementation needs to load the
     * persistent key description in a slot to be able to access it.
     */
    status = psa_get_key_attributes( persistent_key, &attributes );
    TEST_EQUAL( status, PSA_ERROR_INSUFFICIENT_MEMORY );

    /*
     * Check we can export the volatile key created last and that it has the
     * expected value. Then, destroy it.
     */
    PSA_ASSERT( psa_export_key( keys[PSA_KEY_SLOT_COUNT - 1],
                                exported, sizeof( exported ),
                                &exported_length ) );
    i = PSA_KEY_SLOT_COUNT - 1;
    ASSERT_COMPARE( exported, exported_length, (uint8_t *) &i, sizeof( i ) );
    PSA_ASSERT( psa_destroy_key( keys[PSA_KEY_SLOT_COUNT - 1] ) );

    /*
     * Check that we can now access the persistent key again.
     */
    PSA_ASSERT( psa_get_key_attributes( persistent_key, &attributes ) );
    TEST_ASSERT( mbedtls_svc_key_id_equal( attributes.core.id,
                                           persistent_key ) );

    /*
     * Check that we cannot copy the persistent key as all slots are occupied
     * by the persistent key and the volatile keys and the slot containing the
     * persistent key cannot be reclaimed as it contains the key to copy.
     */
    persistent_key2 = mbedtls_svc_key_id_make( 0x100, 0x204 );
    psa_set_key_id( &attributes, persistent_key2 );
    status = psa_copy_key( persistent_key, &attributes, &returned_key_id );
    TEST_EQUAL( status, PSA_ERROR_INSUFFICIENT_MEMORY );

    /*
     * Check we can export the remaining volatile keys and that they have the
     * expected values.
     */
    for( i = 0; i < ( PSA_KEY_SLOT_COUNT - 1 ); i++ )
    {
        PSA_ASSERT( psa_export_key( keys[i],
                                    exported, sizeof( exported ),
                                    &exported_length ) );
        ASSERT_COMPARE( exported, exported_length,
                        (uint8_t *) &i, sizeof( i ) );
        PSA_ASSERT( psa_destroy_key( keys[i] ) );
    }

    /*
     * Check we can export the persistent key and that it have the expected
     * value.
     */

    PSA_ASSERT( psa_export_key( persistent_key, exported, sizeof( exported ),
                                &exported_length ) );
    ASSERT_COMPARE( exported, exported_length,
                    (uint8_t *) &persistent_key, sizeof( persistent_key ) );
exit:
    /*
     * Key attributes may have been returned by psa_get_key_attributes()
     * thus reset them as required.
     */
    psa_reset_key_attributes( &attributes );

    psa_destroy_key( persistent_key );
    PSA_DONE( );
    mbedtls_free( keys );
}
/* END_CASE */