/* BEGIN_HEADER */
#include <stdint.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;

/** 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_SESSION_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;

    mbedtls_test_set_step(1);
    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)
    {
        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_get_key_attributes(key_with_invalid_owner, &attributes),
                   PSA_ERROR_INVALID_HANDLE);
    }
#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. */
    mbedtls_test_set_step(2);
    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_INVALID_HANDLE);
    TEST_EQUAL(psa_close_key(key), PSA_ERROR_INVALID_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 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);

    mbedtls_test_set_step(1);
    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),
               mbedtls_test_update_key_usage_flags(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),
               mbedtls_test_update_key_usage_flags(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.
     */
    mbedtls_test_set_step(2);
    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),
                       mbedtls_test_update_key_usage_flags(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));
            TEST_CALLOC(reexported, key_data->len);
            if (usage_flags & PSA_KEY_USAGE_EXPORT) {
                PSA_ASSERT(psa_export_key(id, reexported, key_data->len,
                                          &reexported_length));
                TEST_MEMORY_COMPARE(key_data->x, key_data->len,
                                    reexported, reexported_length);
            } else {
                TEST_EQUAL(psa_export_key(id, reexported,
                                          key_data->len, &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 referring to an
             * existing key.
             */
            TEST_EQUAL(psa_get_key_attributes(handle, &read_attributes),
                       PSA_ERROR_INVALID_HANDLE);
            TEST_EQUAL(psa_close_key(handle), PSA_ERROR_INVALID_HANDLE);
            TEST_EQUAL(psa_get_key_attributes(id, &read_attributes),
                       PSA_ERROR_INVALID_HANDLE);
            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();
    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));
    TEST_MEMORY_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();
}
/* 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.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();
}
/* 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;
        TEST_CALLOC(export_buffer, material->len);
        PSA_ASSERT(psa_export_key(returned_target_id, export_buffer,
                                  material->len, &length));
        TEST_MEMORY_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);
}
/* 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;
        TEST_CALLOC(export_buffer, target_material->len);
        PSA_ASSERT(psa_export_key(returned_target_id, export_buffer,
                                  target_material->len, &length));
        TEST_MEMORY_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);
}
/* END_CASE */

/* BEGIN_CASE */
void invalid_handle(int handle_construction,
                    int close_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_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. It is very unlikely that
             * all IDs are used up to the last one, so pick
             * PSA_KEY_ID_VOLATILE_MAX to build an unopened and thus invalid
             * identifier.
             */
            key_id = PSA_KEY_ID_VOLATILE_MAX;

            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_FAIL("unknown handle construction");
    }

    /* Attempt to use the invalid handle. */
    TEST_EQUAL(psa_get_key_attributes(invalid_handle, &attributes),
               PSA_ERROR_INVALID_HANDLE);
    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;

    TEST_CALLOC(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));
        TEST_MEMORY_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 MBEDTLS_PSA_KEY_SLOT_COUNT persistent keys.
     */
    for (i = 0; i < MBEDTLS_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 = MBEDTLS_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 ( MBEDTLS_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 <= MBEDTLS_PSA_KEY_SLOT_COUNT; i++) {
        if (i < MBEDTLS_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));
        TEST_MEMORY_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;
    mbedtls_psa_stats_t psa_key_slots_stats;
    size_t available_key_slots = 0;

    TEST_ASSERT(MBEDTLS_PSA_KEY_SLOT_COUNT >= 1);

    PSA_ASSERT(psa_crypto_init());
    mbedtls_psa_get_stats(&psa_key_slots_stats);
    available_key_slots = psa_key_slots_stats.empty_slots;

    TEST_CALLOC(keys, available_key_slots);

    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 the maximum available number of volatile keys
     */
    psa_set_key_lifetime(&attributes, PSA_KEY_LIFETIME_VOLATILE);
    for (i = 0; i < available_key_slots; 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[available_key_slots - 1],
                              exported, sizeof(exported),
                              &exported_length));
    i = available_key_slots - 1;
    TEST_MEMORY_COMPARE(exported, exported_length, (uint8_t *) &i, sizeof(i));
    PSA_ASSERT(psa_destroy_key(keys[available_key_slots - 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.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 < (available_key_slots - 1); i++) {
        PSA_ASSERT(psa_export_key(keys[i],
                                  exported, sizeof(exported),
                                  &exported_length));
        TEST_MEMORY_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));
    TEST_MEMORY_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 */