/* BEGIN_HEADER */
#include "test/drivers/test_driver.h"
/* END_HEADER */

/* BEGIN_DEPENDENCIES
 * depends_on:MBEDTLS_PSA_CRYPTO_C:MBEDTLS_PSA_CRYPTO_DRIVERS:PSA_CRYPTO_DRIVER_TEST
 * END_DEPENDENCIES
 */

/* BEGIN_CASE depends_on:PSA_WANT_ALG_DETERMINISTIC_ECDSA:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256 */
void ecdsa_sign_hash( int force_status_arg,
                      data_t *key_input,
                      data_t *data_input,
                      data_t *expected_output,
                      int fake_output,
                      int expected_status_arg )
{
    psa_status_t force_status = force_status_arg;
    psa_status_t expected_status = expected_status_arg;
    mbedtls_svc_key_id_t key = MBEDTLS_SVC_KEY_ID_INIT;
    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
    psa_algorithm_t alg = PSA_ALG_DETERMINISTIC_ECDSA( PSA_ALG_SHA_256 );
    uint8_t signature[64];
    size_t signature_length = 0xdeadbeef;
    psa_status_t actual_status;
    mbedtls_test_driver_signature_sign_hooks =
        mbedtls_test_driver_signature_hooks_init();

    PSA_ASSERT( psa_crypto_init( ) );
    psa_set_key_type( &attributes,
                      PSA_KEY_TYPE_ECC_KEY_PAIR( PSA_ECC_FAMILY_SECP_R1 ) );
    psa_set_key_usage_flags( &attributes, PSA_KEY_USAGE_SIGN_HASH );
    psa_set_key_algorithm( &attributes, alg );
    psa_import_key( &attributes,
                    key_input->x, key_input->len,
                    &key );

    mbedtls_test_driver_signature_sign_hooks.forced_status = force_status;
    if( fake_output == 1 )
    {
        mbedtls_test_driver_signature_sign_hooks.forced_output =
            expected_output->x;
        mbedtls_test_driver_signature_sign_hooks.forced_output_length =
            expected_output->len;
    }

    actual_status = psa_sign_hash( key, alg,
                                   data_input->x, data_input->len,
                                   signature, sizeof( signature ),
                                   &signature_length );
    TEST_EQUAL( actual_status, expected_status );
    if( expected_status == PSA_SUCCESS )
    {
        ASSERT_COMPARE( signature, signature_length,
                        expected_output->x, expected_output->len );
    }
    TEST_EQUAL( mbedtls_test_driver_signature_sign_hooks.hits, 1 );

exit:
    psa_reset_key_attributes( &attributes );
    psa_destroy_key( key );
    PSA_DONE( );
    mbedtls_test_driver_signature_sign_hooks =
        mbedtls_test_driver_signature_hooks_init();
}
/* END_CASE */

/* BEGIN_CASE depends_on:PSA_WANT_ALG_ECDSA:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256 */
void ecdsa_verify_hash( int force_status_arg,
                        int register_public_key,
                        data_t *key_input,
                        data_t *data_input,
                        data_t *signature_input,
                        int expected_status_arg )
{
    psa_status_t force_status = force_status_arg;
    psa_status_t expected_status = expected_status_arg;
    mbedtls_svc_key_id_t key = MBEDTLS_SVC_KEY_ID_INIT;
    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
    psa_algorithm_t alg = PSA_ALG_ECDSA( PSA_ALG_SHA_256 );
    psa_status_t actual_status;
    mbedtls_test_driver_signature_verify_hooks =
        mbedtls_test_driver_signature_hooks_init();

    PSA_ASSERT( psa_crypto_init( ) );
    if( register_public_key )
    {
        psa_set_key_type( &attributes,
                      PSA_KEY_TYPE_ECC_PUBLIC_KEY( PSA_ECC_FAMILY_SECP_R1 ) );
        psa_set_key_usage_flags( &attributes, PSA_KEY_USAGE_VERIFY_HASH );
        psa_set_key_algorithm( &attributes, alg );
        psa_import_key( &attributes,
                        key_input->x, key_input->len,
                        &key );
    }
    else
    {
        psa_set_key_type( &attributes,
                      PSA_KEY_TYPE_ECC_KEY_PAIR( PSA_ECC_FAMILY_SECP_R1 ) );
        psa_set_key_usage_flags( &attributes, PSA_KEY_USAGE_VERIFY_HASH );
        psa_set_key_algorithm( &attributes, alg );
        psa_import_key( &attributes,
                        key_input->x, key_input->len,
                        &key );
    }

    mbedtls_test_driver_signature_verify_hooks.forced_status = force_status;

    actual_status = psa_verify_hash( key, alg,
                                     data_input->x, data_input->len,
                                     signature_input->x, signature_input->len );
    TEST_EQUAL( actual_status, expected_status );
    TEST_EQUAL( mbedtls_test_driver_signature_verify_hooks.hits, 1 );

exit:
    psa_reset_key_attributes( &attributes );
    psa_destroy_key( key );
    PSA_DONE( );
    mbedtls_test_driver_signature_verify_hooks =
        mbedtls_test_driver_signature_hooks_init();
}
/* END_CASE */

/* BEGIN_CASE depends_on:PSA_WANT_ALG_DETERMINISTIC_ECDSA:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256 */
void ecdsa_sign_message( int force_status_arg,
                         data_t *key_input,
                         data_t *data_input,
                         data_t *expected_output,
                         int fake_output,
                         int expected_status_arg )
{
    psa_status_t force_status = force_status_arg;
    psa_status_t expected_status = expected_status_arg;
    mbedtls_svc_key_id_t key = MBEDTLS_SVC_KEY_ID_INIT;
    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
    psa_algorithm_t alg = PSA_ALG_DETERMINISTIC_ECDSA( PSA_ALG_SHA_256 );
    uint8_t signature[64];
    size_t signature_length = 0xdeadbeef;
    psa_status_t actual_status;
    mbedtls_test_driver_signature_sign_hooks =
        mbedtls_test_driver_signature_hooks_init();

    PSA_ASSERT( psa_crypto_init( ) );
    psa_set_key_type( &attributes,
                      PSA_KEY_TYPE_ECC_KEY_PAIR( PSA_ECC_FAMILY_SECP_R1 ) );
    psa_set_key_usage_flags( &attributes, PSA_KEY_USAGE_SIGN_MESSAGE );
    psa_set_key_algorithm( &attributes, alg );
    psa_import_key( &attributes,
                    key_input->x, key_input->len,
                    &key );

    mbedtls_test_driver_signature_sign_hooks.forced_status = force_status;
    if( fake_output == 1 )
    {
        mbedtls_test_driver_signature_sign_hooks.forced_output =
            expected_output->x;
        mbedtls_test_driver_signature_sign_hooks.forced_output_length =
            expected_output->len;
    }

    actual_status = psa_sign_message( key, alg,
                                      data_input->x, data_input->len,
                                      signature, sizeof( signature ),
                                      &signature_length );
    TEST_EQUAL( actual_status, expected_status );
    if( expected_status == PSA_SUCCESS )
    {
        ASSERT_COMPARE( signature, signature_length,
                        expected_output->x, expected_output->len );
    }
    /* In the builtin algorithm the driver is called twice. */
    TEST_EQUAL( mbedtls_test_driver_signature_sign_hooks.hits,
                force_status == PSA_ERROR_NOT_SUPPORTED ? 2 : 1 );

exit:
    psa_reset_key_attributes( &attributes );
    psa_destroy_key( key );
    PSA_DONE( );
    mbedtls_test_driver_signature_sign_hooks =
        mbedtls_test_driver_signature_hooks_init();
}
/* END_CASE */

/* BEGIN_CASE depends_on:PSA_WANT_ALG_ECDSA:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256 */
void ecdsa_verify_message( int force_status_arg,
                           int register_public_key,
                           data_t *key_input,
                           data_t *data_input,
                           data_t *signature_input,
                           int expected_status_arg )
{
    psa_status_t force_status = force_status_arg;
    psa_status_t expected_status = expected_status_arg;
    mbedtls_svc_key_id_t key = MBEDTLS_SVC_KEY_ID_INIT;
    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
    psa_algorithm_t alg = PSA_ALG_ECDSA( PSA_ALG_SHA_256 );
    psa_status_t actual_status;
    mbedtls_test_driver_signature_verify_hooks =
        mbedtls_test_driver_signature_hooks_init();

    PSA_ASSERT( psa_crypto_init( ) );
    if( register_public_key )
    {
        psa_set_key_type( &attributes,
                      PSA_KEY_TYPE_ECC_PUBLIC_KEY( PSA_ECC_FAMILY_SECP_R1 ) );
        psa_set_key_usage_flags( &attributes, PSA_KEY_USAGE_VERIFY_MESSAGE );
        psa_set_key_algorithm( &attributes, alg );
        psa_import_key( &attributes,
                        key_input->x, key_input->len,
                        &key );
    }
    else
    {
        psa_set_key_type( &attributes,
                      PSA_KEY_TYPE_ECC_KEY_PAIR( PSA_ECC_FAMILY_SECP_R1 ) );
        psa_set_key_usage_flags( &attributes, PSA_KEY_USAGE_VERIFY_MESSAGE );
        psa_set_key_algorithm( &attributes, alg );
        psa_import_key( &attributes,
                        key_input->x, key_input->len,
                        &key );
    }

    mbedtls_test_driver_signature_verify_hooks.forced_status = force_status;

    actual_status = psa_verify_message( key, alg,
                                        data_input->x, data_input->len,
                                        signature_input->x, signature_input->len );
    TEST_EQUAL( actual_status, expected_status );
    /* In the builtin algorithm the driver is called twice. */
    TEST_EQUAL( mbedtls_test_driver_signature_verify_hooks.hits,
                force_status == PSA_ERROR_NOT_SUPPORTED ? 2 : 1 );

exit:
    psa_reset_key_attributes( &attributes );
    psa_destroy_key( key );
    PSA_DONE( );
    mbedtls_test_driver_signature_verify_hooks =
        mbedtls_test_driver_signature_hooks_init();
}
/* END_CASE */

/* BEGIN_CASE depends_on:PSA_WANT_ALG_ECDSA:PSA_WANT_ECC_SECP_R1_256 */
void generate_key( int force_status_arg,
                   data_t *fake_output,
                   int expected_status_arg )
{
    psa_status_t force_status = force_status_arg;
    psa_status_t expected_status = expected_status_arg;
    mbedtls_svc_key_id_t key = MBEDTLS_SVC_KEY_ID_INIT;
    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
    psa_algorithm_t alg = PSA_ALG_ECDSA( PSA_ALG_SHA_256 );
    const uint8_t *expected_output = NULL;
    size_t expected_output_length = 0;
    psa_status_t actual_status;
    uint8_t actual_output[PSA_KEY_EXPORT_ECC_KEY_PAIR_MAX_SIZE(256)] = {0};
    size_t actual_output_length;
    mbedtls_test_driver_key_management_hooks =
        mbedtls_test_driver_key_management_hooks_init();

    psa_set_key_type( &attributes,
                      PSA_KEY_TYPE_ECC_KEY_PAIR( PSA_ECC_FAMILY_SECP_R1 ) );
    psa_set_key_bits( &attributes, 256 );
    psa_set_key_usage_flags( &attributes, PSA_KEY_USAGE_SIGN_HASH | PSA_KEY_USAGE_EXPORT );
    psa_set_key_algorithm( &attributes, alg );

    if( fake_output->len > 0 )
    {
        expected_output =
            mbedtls_test_driver_key_management_hooks.forced_output =
            fake_output->x;

        expected_output_length =
            mbedtls_test_driver_key_management_hooks.forced_output_length =
            fake_output->len;
    }

    mbedtls_test_driver_key_management_hooks.hits = 0;
    mbedtls_test_driver_key_management_hooks.forced_status = force_status;

    PSA_ASSERT( psa_crypto_init( ) );

    actual_status = psa_generate_key( &attributes, &key );
    TEST_EQUAL( mbedtls_test_driver_key_management_hooks.hits, 1 );
    TEST_EQUAL( actual_status, expected_status );

    if( actual_status == PSA_SUCCESS )
    {
        psa_export_key( key, actual_output, sizeof(actual_output), &actual_output_length );

        if( fake_output->len > 0 )
        {
            ASSERT_COMPARE( actual_output, actual_output_length,
                            expected_output, expected_output_length );
        }
        else
        {
            size_t zeroes = 0;
            for( size_t i = 0; i < sizeof(actual_output); i++ )
            {
                if( actual_output[i] == 0)
                    zeroes++;
            }
            TEST_ASSERT( zeroes != sizeof(actual_output) );
        }
    }
exit:
    psa_reset_key_attributes( &attributes );
    psa_destroy_key( key );
    PSA_DONE( );
    mbedtls_test_driver_key_management_hooks =
        mbedtls_test_driver_key_management_hooks_init();
}
/* END_CASE */

/* BEGIN_CASE */
void validate_key( int force_status_arg,
                   int location,
                   int owner_id_arg,
                   int id_arg,
                   int key_type_arg,
                   data_t *key_input,
                   int expected_status_arg )
{
    psa_key_lifetime_t lifetime =
                              PSA_KEY_LIFETIME_FROM_PERSISTENCE_AND_LOCATION( \
                                  PSA_KEY_PERSISTENCE_DEFAULT, location);
    mbedtls_svc_key_id_t id = mbedtls_svc_key_id_make( owner_id_arg, id_arg );
    psa_status_t force_status = force_status_arg;
    psa_status_t expected_status = expected_status_arg;
    psa_key_type_t key_type = key_type_arg;
    mbedtls_svc_key_id_t key = MBEDTLS_SVC_KEY_ID_INIT;
    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
    psa_status_t actual_status;
    mbedtls_test_driver_key_management_hooks =
        mbedtls_test_driver_key_management_hooks_init();

    psa_set_key_id( &attributes, id );
    psa_set_key_type( &attributes,
                      key_type );
    psa_set_key_lifetime( &attributes, lifetime );
    psa_set_key_bits( &attributes, 0 );
    psa_set_key_usage_flags( &attributes, PSA_KEY_USAGE_EXPORT );

    mbedtls_test_driver_key_management_hooks.forced_status = force_status;

    PSA_ASSERT( psa_crypto_init( ) );

    actual_status = psa_import_key( &attributes, key_input->x, key_input->len, &key );
    TEST_EQUAL( mbedtls_test_driver_key_management_hooks.hits, 1 );
    TEST_EQUAL( actual_status, expected_status );
    TEST_EQUAL( mbedtls_test_driver_key_management_hooks.location, location );
exit:
    psa_reset_key_attributes( &attributes );
    psa_destroy_key( key );
    PSA_DONE( );
    mbedtls_test_driver_key_management_hooks =
        mbedtls_test_driver_key_management_hooks_init();
}
/* END_CASE */

/* BEGIN_CASE */
void export_key( int force_status_arg,
                 data_t *fake_output,
                 int key_in_type_arg,
                 data_t *key_in,
                 int key_out_type_arg,
                 data_t *expected_output,
                 int expected_status_arg )
{
    psa_status_t force_status = force_status_arg;
    psa_status_t expected_status = expected_status_arg;
    psa_key_handle_t handle = 0;
    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
    psa_key_type_t input_key_type = key_in_type_arg;
    psa_key_type_t output_key_type = key_out_type_arg;
    const uint8_t *expected_output_ptr = NULL;
    size_t expected_output_length = 0;
    psa_status_t actual_status;
    uint8_t actual_output[PSA_KEY_EXPORT_ECC_PUBLIC_KEY_MAX_SIZE(256)] = {0};
    size_t actual_output_length;
    mbedtls_test_driver_key_management_hooks =
        mbedtls_test_driver_key_management_hooks_init();

    psa_set_key_type( &attributes, input_key_type );
    psa_set_key_bits( &attributes, 256 );
    psa_set_key_usage_flags( &attributes, PSA_KEY_USAGE_EXPORT );

    PSA_ASSERT( psa_crypto_init( ) );
    PSA_ASSERT( psa_import_key( &attributes, key_in->x, key_in->len, &handle ) );

    if( fake_output->len > 0 )
    {
        expected_output_ptr =
            mbedtls_test_driver_key_management_hooks.forced_output =
            fake_output->x;

        expected_output_length =
            mbedtls_test_driver_key_management_hooks.forced_output_length =
            fake_output->len;
    }
    else
    {
        expected_output_ptr = expected_output->x;
        expected_output_length = expected_output->len;
    }

    mbedtls_test_driver_key_management_hooks.hits = 0;
    mbedtls_test_driver_key_management_hooks.forced_status = force_status;

    if( PSA_KEY_TYPE_IS_ECC_PUBLIC_KEY( output_key_type ) )
        actual_status = psa_export_public_key( handle, actual_output, sizeof(actual_output), &actual_output_length );
    else
        actual_status = psa_export_key( handle, actual_output, sizeof(actual_output), &actual_output_length );
    TEST_EQUAL( actual_status, expected_status );

    if( PSA_KEY_TYPE_IS_ECC_PUBLIC_KEY( output_key_type ) &&
        !PSA_KEY_TYPE_IS_ECC_PUBLIC_KEY( input_key_type ) )
        TEST_EQUAL( mbedtls_test_driver_key_management_hooks.hits, 1 );

    if( actual_status == PSA_SUCCESS )
    {
        ASSERT_COMPARE( actual_output, actual_output_length,
                        expected_output_ptr, expected_output_length );
    }
exit:
    psa_reset_key_attributes( &attributes );
    psa_destroy_key( handle );
    PSA_DONE( );
    mbedtls_test_driver_key_management_hooks =
        mbedtls_test_driver_key_management_hooks_init();
}
/* END_CASE */

/* BEGIN_CASE */
void cipher_encrypt_validation( int alg_arg,
                                int key_type_arg,
                                data_t *key_data,
                                data_t *input )
{
    mbedtls_svc_key_id_t key = MBEDTLS_SVC_KEY_ID_INIT;
    psa_key_type_t key_type = key_type_arg;
    psa_algorithm_t alg = alg_arg;
    size_t iv_size = PSA_CIPHER_IV_LENGTH ( key_type, alg );
    unsigned char *output1 = NULL;
    size_t output1_buffer_size = 0;
    size_t output1_length = 0;
    unsigned char *output2 = NULL;
    size_t output2_buffer_size = 0;
    size_t output2_length = 0;
    size_t function_output_length = 0;
    psa_cipher_operation_t operation = PSA_CIPHER_OPERATION_INIT;
    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
    mbedtls_test_driver_cipher_hooks = mbedtls_test_driver_cipher_hooks_init();

    PSA_ASSERT( psa_crypto_init( ) );

    psa_set_key_usage_flags( &attributes, PSA_KEY_USAGE_ENCRYPT );
    psa_set_key_algorithm( &attributes, alg );
    psa_set_key_type( &attributes, key_type );

    output1_buffer_size = PSA_CIPHER_ENCRYPT_OUTPUT_SIZE( key_type, alg, input->len );
    output2_buffer_size = PSA_CIPHER_UPDATE_OUTPUT_SIZE( key_type, alg, input->len ) +
                          PSA_CIPHER_FINISH_OUTPUT_SIZE( key_type, alg );
    ASSERT_ALLOC( output1, output1_buffer_size );
    ASSERT_ALLOC( output2, output2_buffer_size );

    PSA_ASSERT( psa_import_key( &attributes, key_data->x, key_data->len,
                                &key ) );

    PSA_ASSERT( psa_cipher_encrypt( key, alg, input->x, input->len, output1,
                                    output1_buffer_size, &output1_length ) );
    TEST_EQUAL( mbedtls_test_driver_cipher_hooks.hits, 1 );
    mbedtls_test_driver_cipher_hooks.hits = 0;

    PSA_ASSERT( psa_cipher_encrypt_setup( &operation, key, alg ) );
    TEST_EQUAL( mbedtls_test_driver_cipher_hooks.hits, 1 );
    mbedtls_test_driver_cipher_hooks.hits = 0;

    PSA_ASSERT( psa_cipher_set_iv( &operation, output1, iv_size ) );
    TEST_EQUAL( mbedtls_test_driver_cipher_hooks.hits, 1 );
    mbedtls_test_driver_cipher_hooks.hits = 0;

    PSA_ASSERT( psa_cipher_update( &operation,
                                   input->x, input->len,
                                   output2, output2_buffer_size,
                                   &function_output_length ) );
    TEST_EQUAL( mbedtls_test_driver_cipher_hooks.hits, 1 );
    mbedtls_test_driver_cipher_hooks.hits = 0;

    output2_length += function_output_length;
    PSA_ASSERT( psa_cipher_finish( &operation,
                                   output2 + output2_length,
                                   output2_buffer_size - output2_length,
                                   &function_output_length ) );
    /* Finish will have called abort as well, so expecting two hits here */
    TEST_EQUAL( mbedtls_test_driver_cipher_hooks.hits, 2 );
    mbedtls_test_driver_cipher_hooks.hits = 0;

    output2_length += function_output_length;

    PSA_ASSERT( psa_cipher_abort( &operation ) );
    // driver function should've been called as part of the finish() core routine
    TEST_EQUAL( mbedtls_test_driver_cipher_hooks.hits, 0 );
    ASSERT_COMPARE( output1 + iv_size, output1_length - iv_size,
                    output2, output2_length );

exit:
    psa_cipher_abort( &operation );
    mbedtls_free( output1 );
    mbedtls_free( output2 );
    psa_destroy_key( key );
    PSA_DONE( );
    mbedtls_test_driver_cipher_hooks = mbedtls_test_driver_cipher_hooks_init();
}
/* END_CASE */

/* BEGIN_CASE */
void cipher_encrypt_multipart( int alg_arg,
                               int key_type_arg,
                               data_t *key_data,
                               data_t *iv,
                               data_t *input,
                               int first_part_size_arg,
                               int output1_length_arg,
                               int output2_length_arg,
                               data_t *expected_output,
                               int mock_output_arg,
                               int force_status_arg,
                               int expected_status_arg )
{
    mbedtls_svc_key_id_t key = MBEDTLS_SVC_KEY_ID_INIT;
    psa_key_type_t key_type = key_type_arg;
    psa_algorithm_t alg = alg_arg;
    psa_status_t status;
    psa_status_t expected_status = expected_status_arg;
    psa_status_t force_status = force_status_arg;
    size_t first_part_size = first_part_size_arg;
    size_t output1_length = output1_length_arg;
    size_t output2_length = output2_length_arg;
    unsigned char *output = NULL;
    size_t output_buffer_size = 0;
    size_t function_output_length = 0;
    size_t total_output_length = 0;
    psa_cipher_operation_t operation = PSA_CIPHER_OPERATION_INIT;
    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
    mbedtls_test_driver_cipher_hooks = mbedtls_test_driver_cipher_hooks_init();
    mbedtls_test_driver_cipher_hooks.forced_status = force_status;

    /* Test operation initialization */
    mbedtls_psa_cipher_operation_t mbedtls_operation =
            MBEDTLS_PSA_CIPHER_OPERATION_INIT;

    mbedtls_transparent_test_driver_cipher_operation_t tranparent_operation =
            MBEDTLS_TRANSPARENT_TEST_DRIVER_CIPHER_OPERATION_INIT;

    mbedtls_opaque_test_driver_cipher_operation_t opaque_operation =
            MBEDTLS_OPAQUE_TEST_DRIVER_CIPHER_OPERATION_INIT;

    operation.ctx.mbedtls_ctx = mbedtls_operation;
    operation.ctx.transparent_test_driver_ctx = tranparent_operation;
    operation.ctx.opaque_test_driver_ctx = opaque_operation;

    PSA_ASSERT( psa_crypto_init( ) );

    psa_set_key_usage_flags( &attributes, PSA_KEY_USAGE_ENCRYPT );
    psa_set_key_algorithm( &attributes, alg );
    psa_set_key_type( &attributes, key_type );

    PSA_ASSERT( psa_import_key( &attributes, key_data->x, key_data->len,
                                &key ) );

    PSA_ASSERT( psa_cipher_encrypt_setup( &operation, key, alg ) );
    TEST_EQUAL( mbedtls_test_driver_cipher_hooks.hits, 1 );
    mbedtls_test_driver_cipher_hooks.hits = 0;

    PSA_ASSERT( psa_cipher_set_iv( &operation, iv->x, iv->len ) );
    TEST_EQUAL( mbedtls_test_driver_cipher_hooks.hits, ( force_status == PSA_SUCCESS ? 1 : 0 ) );
    mbedtls_test_driver_cipher_hooks.hits = 0;

    output_buffer_size = ( (size_t) input->len +
                           PSA_BLOCK_CIPHER_BLOCK_LENGTH( key_type ) );
    ASSERT_ALLOC( output, output_buffer_size );

    if( mock_output_arg )
    {
        mbedtls_test_driver_cipher_hooks.forced_output = expected_output->x;
        mbedtls_test_driver_cipher_hooks.forced_output_length = expected_output->len;
    }

    TEST_ASSERT( first_part_size <= input->len );
    PSA_ASSERT( psa_cipher_update( &operation, input->x, first_part_size,
                                   output, output_buffer_size,
                                   &function_output_length ) );
    TEST_EQUAL( mbedtls_test_driver_cipher_hooks.hits, ( force_status == PSA_SUCCESS ? 1 : 0 ) );
    mbedtls_test_driver_cipher_hooks.hits = 0;

    TEST_ASSERT( function_output_length == output1_length );
    total_output_length += function_output_length;

    if( first_part_size < input->len )
    {
        PSA_ASSERT( psa_cipher_update( &operation,
                                       input->x + first_part_size,
                                       input->len - first_part_size,
                                       output + total_output_length,
                                       output_buffer_size - total_output_length,
                                       &function_output_length ) );
        TEST_EQUAL( mbedtls_test_driver_cipher_hooks.hits, 1 );
        mbedtls_test_driver_cipher_hooks.hits = 0;

        TEST_ASSERT( function_output_length == output2_length );
        total_output_length += function_output_length;
    }

    if( mock_output_arg )
    {
        mbedtls_test_driver_cipher_hooks.forced_output = NULL;
        mbedtls_test_driver_cipher_hooks.forced_output_length = 0;
    }

    status =  psa_cipher_finish( &operation,
                                 output + total_output_length,
                                 output_buffer_size - total_output_length,
                                 &function_output_length );
    /* Finish will have called abort as well, so expecting two hits here */
    TEST_EQUAL( mbedtls_test_driver_cipher_hooks.hits, ( force_status == PSA_SUCCESS ? 2 : 0 ) );
    mbedtls_test_driver_cipher_hooks.hits = 0 ;
    total_output_length += function_output_length;
    TEST_EQUAL( status, expected_status );

    if( expected_status == PSA_SUCCESS )
    {
        PSA_ASSERT( psa_cipher_abort( &operation ) );
        TEST_EQUAL( mbedtls_test_driver_cipher_hooks.hits, 0 );

        ASSERT_COMPARE( expected_output->x, expected_output->len,
                        output, total_output_length );
    }

exit:
    psa_cipher_abort( &operation );
    mbedtls_free( output );
    psa_destroy_key( key );
    PSA_DONE( );
    mbedtls_test_driver_cipher_hooks = mbedtls_test_driver_cipher_hooks_init();
}
/* END_CASE */

/* BEGIN_CASE */
void cipher_decrypt_multipart( int alg_arg,
                               int key_type_arg,
                               data_t *key_data,
                               data_t *iv,
                               data_t *input,
                               int first_part_size_arg,
                               int output1_length_arg,
                               int output2_length_arg,
                               data_t *expected_output,
                               int mock_output_arg,
                               int force_status_arg,
                               int expected_status_arg )
{
    mbedtls_svc_key_id_t key = MBEDTLS_SVC_KEY_ID_INIT;
    psa_key_type_t key_type = key_type_arg;
    psa_algorithm_t alg = alg_arg;
    psa_status_t status;
    psa_status_t expected_status = expected_status_arg;
    psa_status_t force_status = force_status_arg;
    size_t first_part_size = first_part_size_arg;
    size_t output1_length = output1_length_arg;
    size_t output2_length = output2_length_arg;
    unsigned char *output = NULL;
    size_t output_buffer_size = 0;
    size_t function_output_length = 0;
    size_t total_output_length = 0;
    psa_cipher_operation_t operation = PSA_CIPHER_OPERATION_INIT;
    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
    mbedtls_test_driver_cipher_hooks = mbedtls_test_driver_cipher_hooks_init();
    mbedtls_test_driver_cipher_hooks.forced_status = force_status;

    /* Test operation initialization */
    mbedtls_psa_cipher_operation_t mbedtls_operation =
            MBEDTLS_PSA_CIPHER_OPERATION_INIT;

    mbedtls_transparent_test_driver_cipher_operation_t tranparent_operation =
            MBEDTLS_TRANSPARENT_TEST_DRIVER_CIPHER_OPERATION_INIT;

    mbedtls_opaque_test_driver_cipher_operation_t opaque_operation =
            MBEDTLS_OPAQUE_TEST_DRIVER_CIPHER_OPERATION_INIT;

    operation.ctx.mbedtls_ctx = mbedtls_operation;
    operation.ctx.transparent_test_driver_ctx = tranparent_operation;
    operation.ctx.opaque_test_driver_ctx = opaque_operation;

    PSA_ASSERT( psa_crypto_init( ) );

    psa_set_key_usage_flags( &attributes, PSA_KEY_USAGE_DECRYPT );
    psa_set_key_algorithm( &attributes, alg );
    psa_set_key_type( &attributes, key_type );

    PSA_ASSERT( psa_import_key( &attributes, key_data->x, key_data->len,
                                &key ) );

    PSA_ASSERT( psa_cipher_decrypt_setup( &operation, key, alg ) );
    TEST_EQUAL( mbedtls_test_driver_cipher_hooks.hits, 1 );
    mbedtls_test_driver_cipher_hooks.hits = 0;

    PSA_ASSERT( psa_cipher_set_iv( &operation, iv->x, iv->len ) );
    TEST_EQUAL( mbedtls_test_driver_cipher_hooks.hits, ( force_status == PSA_SUCCESS ? 1 : 0 ) );
    mbedtls_test_driver_cipher_hooks.hits = 0;

    output_buffer_size = ( (size_t) input->len +
                           PSA_BLOCK_CIPHER_BLOCK_LENGTH( key_type ) );
    ASSERT_ALLOC( output, output_buffer_size );

    if( mock_output_arg )
    {
        mbedtls_test_driver_cipher_hooks.forced_output = expected_output->x;
        mbedtls_test_driver_cipher_hooks.forced_output_length = expected_output->len;
    }

    TEST_ASSERT( first_part_size <= input->len );
    PSA_ASSERT( psa_cipher_update( &operation,
                                   input->x, first_part_size,
                                   output, output_buffer_size,
                                   &function_output_length ) );
    TEST_EQUAL( mbedtls_test_driver_cipher_hooks.hits, ( force_status == PSA_SUCCESS ? 1 : 0 ) );
    mbedtls_test_driver_cipher_hooks.hits = 0;

    TEST_ASSERT( function_output_length == output1_length );
    total_output_length += function_output_length;

    if( first_part_size < input->len )
    {
        PSA_ASSERT( psa_cipher_update( &operation,
                                       input->x + first_part_size,
                                       input->len - first_part_size,
                                       output + total_output_length,
                                       output_buffer_size - total_output_length,
                                       &function_output_length ) );
        TEST_EQUAL( mbedtls_test_driver_cipher_hooks.hits, ( force_status == PSA_SUCCESS ? 1 : 0 ) );
        mbedtls_test_driver_cipher_hooks.hits = 0;

        TEST_ASSERT( function_output_length == output2_length );
        total_output_length += function_output_length;
    }

    if( mock_output_arg )
    {
        mbedtls_test_driver_cipher_hooks.forced_output = NULL;
        mbedtls_test_driver_cipher_hooks.forced_output_length = 0;
    }

    status = psa_cipher_finish( &operation,
                                output + total_output_length,
                                output_buffer_size - total_output_length,
                                &function_output_length );
    /* Finish will have called abort as well, so expecting two hits here */
    TEST_EQUAL( mbedtls_test_driver_cipher_hooks.hits, ( force_status == PSA_SUCCESS ? 2 : 0 ) );
    mbedtls_test_driver_cipher_hooks.hits = 0;
    total_output_length += function_output_length;
    TEST_EQUAL( status, expected_status );

    if( expected_status == PSA_SUCCESS )
    {
        PSA_ASSERT( psa_cipher_abort( &operation ) );
        TEST_EQUAL( mbedtls_test_driver_cipher_hooks.hits, 0 );

        ASSERT_COMPARE( expected_output->x, expected_output->len,
                        output, total_output_length );
    }

exit:
    psa_cipher_abort( &operation );
    mbedtls_free( output );
    psa_destroy_key( key );
    PSA_DONE( );
    mbedtls_test_driver_cipher_hooks = mbedtls_test_driver_cipher_hooks_init();
}
/* END_CASE */

/* BEGIN_CASE */
void cipher_decrypt( int alg_arg,
                     int key_type_arg,
                     data_t *key_data,
                     data_t *iv,
                     data_t *input_arg,
                     data_t *expected_output,
                     int mock_output_arg,
                     int force_status_arg,
                     int expected_status_arg )
{
    mbedtls_svc_key_id_t key = MBEDTLS_SVC_KEY_ID_INIT;
    psa_status_t status;
    psa_key_type_t key_type = key_type_arg;
    psa_algorithm_t alg = alg_arg;
    psa_status_t expected_status = expected_status_arg;
    psa_status_t force_status = force_status_arg;
    unsigned char *input = NULL;
    size_t input_buffer_size = 0;
    unsigned char *output = NULL;
    size_t output_buffer_size = 0;
    size_t output_length = 0;
    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
    mbedtls_test_driver_cipher_hooks = mbedtls_test_driver_cipher_hooks_init();
    mbedtls_test_driver_cipher_hooks.forced_status = force_status;

    PSA_ASSERT( psa_crypto_init( ) );

    psa_set_key_usage_flags( &attributes, PSA_KEY_USAGE_DECRYPT );
    psa_set_key_algorithm( &attributes, alg );
    psa_set_key_type( &attributes, key_type );

    /* Allocate input buffer and copy the iv and the plaintext */
    input_buffer_size = ( (size_t) input_arg->len + (size_t) iv->len );
    if ( input_buffer_size > 0 )
    {
        ASSERT_ALLOC( input, input_buffer_size );
        memcpy( input, iv->x, iv->len );
        memcpy( input + iv->len, input_arg->x, input_arg->len );
    }

    output_buffer_size = PSA_CIPHER_DECRYPT_OUTPUT_SIZE( key_type, alg, input_buffer_size );
    ASSERT_ALLOC( output, output_buffer_size );

    PSA_ASSERT( psa_import_key( &attributes, key_data->x, key_data->len,
                                &key ) );

    if( mock_output_arg )
    {
        mbedtls_test_driver_cipher_hooks.forced_output = expected_output->x;
        mbedtls_test_driver_cipher_hooks.forced_output_length = expected_output->len;
    }

    status = psa_cipher_decrypt( key, alg, input, input_buffer_size, output,
                                 output_buffer_size, &output_length );
    TEST_EQUAL( mbedtls_test_driver_cipher_hooks.hits, 1 );
    mbedtls_test_driver_cipher_hooks.hits = 0;

    TEST_EQUAL( status, expected_status );

    if( expected_status == PSA_SUCCESS )
    {
        ASSERT_COMPARE( expected_output->x, expected_output->len,
                        output, output_length );
    }

exit:
    mbedtls_free( input );
    mbedtls_free( output );
    psa_destroy_key( key );
    PSA_DONE( );
    mbedtls_test_driver_cipher_hooks = mbedtls_test_driver_cipher_hooks_init();
}
/* END_CASE */

/* BEGIN_CASE */
void cipher_entry_points( int alg_arg, int key_type_arg,
                          data_t *key_data, data_t *iv,
                          data_t *input )
{
    mbedtls_svc_key_id_t key = MBEDTLS_SVC_KEY_ID_INIT;
    psa_status_t status;
    psa_key_type_t key_type = key_type_arg;
    psa_algorithm_t alg = alg_arg;
    unsigned char *output = NULL;
    size_t output_buffer_size = 0;
    size_t function_output_length = 0;
    psa_cipher_operation_t operation = PSA_CIPHER_OPERATION_INIT;
    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
    mbedtls_test_driver_cipher_hooks = mbedtls_test_driver_cipher_hooks_init();

    ASSERT_ALLOC( output, input->len + 16 );
    output_buffer_size = input->len + 16;

    PSA_ASSERT( psa_crypto_init( ) );

    psa_set_key_usage_flags( &attributes, PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT );
    psa_set_key_algorithm( &attributes, alg );
    psa_set_key_type( &attributes, key_type );

    PSA_ASSERT( psa_import_key( &attributes, key_data->x, key_data->len,
                                &key ) );

    /*
     * Test encrypt failure
     * First test that if we don't force a driver error, encryption is
     * successfull, then force driver error.
     */
    status = psa_cipher_encrypt(
        key, alg, input->x, input->len,
        output, output_buffer_size, &function_output_length );
    TEST_EQUAL( mbedtls_test_driver_cipher_hooks.hits, 1 );
    TEST_EQUAL( status, PSA_SUCCESS );
    mbedtls_test_driver_cipher_hooks.hits = 0;

    mbedtls_test_driver_cipher_hooks.forced_status = PSA_ERROR_GENERIC_ERROR;
    /* Set the output buffer in a given state. */
    for( size_t i = 0; i < output_buffer_size; i++ )
        output[i] = 0xa5;

    status = psa_cipher_encrypt(
        key, alg, input->x, input->len,
        output, output_buffer_size, &function_output_length );
    TEST_EQUAL( mbedtls_test_driver_cipher_hooks.hits, 1 );
    TEST_EQUAL( status, PSA_ERROR_GENERIC_ERROR );
    /*
     * Check that the output buffer is still in the same state.
     * This will fail if the output buffer is used by the core to pass the IV
     * it generated to the driver (and is not restored).
     */
    for( size_t i = 0; i < output_buffer_size; i++ )
    {
        TEST_EQUAL( output[i], 0xa5 );
    }
    mbedtls_test_driver_cipher_hooks.hits = 0;

    /* Test setup call, encrypt */
    mbedtls_test_driver_cipher_hooks.forced_status = PSA_ERROR_GENERIC_ERROR;
    status = psa_cipher_encrypt_setup( &operation, key, alg );
    /* When setup fails, it shouldn't call any further entry points */
    TEST_EQUAL( mbedtls_test_driver_cipher_hooks.hits, 1 );
    TEST_EQUAL( status, mbedtls_test_driver_cipher_hooks.forced_status );
    mbedtls_test_driver_cipher_hooks.hits = 0;
    status = psa_cipher_set_iv( &operation, iv->x, iv->len );
    TEST_EQUAL( status, PSA_ERROR_BAD_STATE );
    TEST_EQUAL( mbedtls_test_driver_cipher_hooks.hits, 0 );

    /* Test setup call failure, decrypt */
    status = psa_cipher_decrypt_setup( &operation, key, alg );
    /* When setup fails, it shouldn't call any further entry points */
    TEST_EQUAL( mbedtls_test_driver_cipher_hooks.hits, 1 );
    TEST_EQUAL( status, mbedtls_test_driver_cipher_hooks.forced_status );
    mbedtls_test_driver_cipher_hooks.hits = 0;
    status = psa_cipher_set_iv( &operation, iv->x, iv->len );
    TEST_EQUAL( status, PSA_ERROR_BAD_STATE );
    TEST_EQUAL( mbedtls_test_driver_cipher_hooks.hits, 0 );

    /* Test IV setting failure */
    mbedtls_test_driver_cipher_hooks.forced_status = PSA_SUCCESS;
    status = psa_cipher_encrypt_setup( &operation, key, alg );
    TEST_EQUAL( mbedtls_test_driver_cipher_hooks.hits, 1 );
    TEST_EQUAL( status, mbedtls_test_driver_cipher_hooks.forced_status );
    mbedtls_test_driver_cipher_hooks.hits = 0;

    mbedtls_test_driver_cipher_hooks.forced_status = PSA_ERROR_GENERIC_ERROR;
    status = psa_cipher_set_iv( &operation, iv->x, iv->len );
    /* When setting the IV fails, it should call abort too */
    TEST_EQUAL( mbedtls_test_driver_cipher_hooks.hits, 2 );
    TEST_EQUAL( status, mbedtls_test_driver_cipher_hooks.forced_status );
    /* Failure should prevent further operations from executing on the driver */
    mbedtls_test_driver_cipher_hooks.hits = 0;
    status = psa_cipher_update( &operation,
                                input->x, input->len,
                                output, output_buffer_size,
                                &function_output_length );
    TEST_EQUAL( status, PSA_ERROR_BAD_STATE );
    TEST_EQUAL( mbedtls_test_driver_cipher_hooks.hits, 0 );
    psa_cipher_abort( &operation );

    /* Test IV generation failure */
    mbedtls_test_driver_cipher_hooks.forced_status = PSA_SUCCESS;
    status = psa_cipher_encrypt_setup( &operation, key, alg );
    TEST_EQUAL( mbedtls_test_driver_cipher_hooks.hits, 1 );
    TEST_EQUAL( status, mbedtls_test_driver_cipher_hooks.forced_status );
    mbedtls_test_driver_cipher_hooks.hits = 0;

    mbedtls_test_driver_cipher_hooks.forced_status = PSA_ERROR_GENERIC_ERROR;
    /* Set the output buffer in a given state. */
    for( size_t i = 0; i < 16; i++ )
        output[i] = 0xa5;

    status = psa_cipher_generate_iv( &operation, output, 16, &function_output_length );
    /* When generating the IV fails, it should call abort too */
    TEST_EQUAL( mbedtls_test_driver_cipher_hooks.hits, 2 );
    TEST_EQUAL( status, mbedtls_test_driver_cipher_hooks.forced_status );
    /*
     * Check that the output buffer is still in the same state.
     * This will fail if the output buffer is used by the core to pass the IV
     * it generated to the driver (and is not restored).
     */
    for( size_t i = 0; i < 16; i++ )
    {
        TEST_EQUAL( output[i], 0xa5 );
    }
    /* Failure should prevent further operations from executing on the driver */
    mbedtls_test_driver_cipher_hooks.hits = 0;
    status = psa_cipher_update( &operation,
                                input->x, input->len,
                                output, output_buffer_size,
                                &function_output_length );
    TEST_EQUAL( status, PSA_ERROR_BAD_STATE );
    TEST_EQUAL( mbedtls_test_driver_cipher_hooks.hits, 0 );
    psa_cipher_abort( &operation );

    /* Test update failure */
    mbedtls_test_driver_cipher_hooks.forced_status = PSA_SUCCESS;
    status = psa_cipher_encrypt_setup( &operation, key, alg );
    TEST_EQUAL( mbedtls_test_driver_cipher_hooks.hits, 1 );
    TEST_EQUAL( status, mbedtls_test_driver_cipher_hooks.forced_status );
    mbedtls_test_driver_cipher_hooks.hits = 0;

    status = psa_cipher_set_iv( &operation, iv->x, iv->len );
    TEST_EQUAL( mbedtls_test_driver_cipher_hooks.hits, 1 );
    TEST_EQUAL( status, mbedtls_test_driver_cipher_hooks.forced_status );
    mbedtls_test_driver_cipher_hooks.hits = 0;

    mbedtls_test_driver_cipher_hooks.forced_status = PSA_ERROR_GENERIC_ERROR;
    status = psa_cipher_update( &operation,
                                input->x, input->len,
                                output, output_buffer_size,
                                &function_output_length );
    /* When the update call fails, it should call abort too */
    TEST_EQUAL( mbedtls_test_driver_cipher_hooks.hits, 2 );
    TEST_EQUAL( status, mbedtls_test_driver_cipher_hooks.forced_status );
    /* Failure should prevent further operations from executing on the driver */
    mbedtls_test_driver_cipher_hooks.hits = 0;
    status = psa_cipher_update( &operation,
                                input->x, input->len,
                                output, output_buffer_size,
                                &function_output_length );
    TEST_EQUAL( status, PSA_ERROR_BAD_STATE );
    TEST_EQUAL( mbedtls_test_driver_cipher_hooks.hits, 0 );
    psa_cipher_abort( &operation );

    /* Test finish failure */
    mbedtls_test_driver_cipher_hooks.forced_status = PSA_SUCCESS;
    status = psa_cipher_encrypt_setup( &operation, key, alg );
    TEST_EQUAL( mbedtls_test_driver_cipher_hooks.hits, 1 );
    TEST_EQUAL( status, mbedtls_test_driver_cipher_hooks.forced_status );
    mbedtls_test_driver_cipher_hooks.hits = 0;

    status = psa_cipher_set_iv( &operation, iv->x, iv->len );
    TEST_EQUAL( mbedtls_test_driver_cipher_hooks.hits, 1 );
    TEST_EQUAL( status, mbedtls_test_driver_cipher_hooks.forced_status );
    mbedtls_test_driver_cipher_hooks.hits = 0;

    status = psa_cipher_update( &operation,
                                input->x, input->len,
                                output, output_buffer_size,
                                &function_output_length );
    TEST_EQUAL( mbedtls_test_driver_cipher_hooks.hits, 1 );
    TEST_EQUAL( status, mbedtls_test_driver_cipher_hooks.forced_status );
    mbedtls_test_driver_cipher_hooks.hits = 0;

    mbedtls_test_driver_cipher_hooks.forced_status = PSA_ERROR_GENERIC_ERROR;
    status = psa_cipher_finish( &operation,
                                output + function_output_length,
                                output_buffer_size - function_output_length,
                                &function_output_length );
    /* When the finish call fails, it should call abort too */
    TEST_EQUAL( mbedtls_test_driver_cipher_hooks.hits, 2 );
    TEST_EQUAL( status, mbedtls_test_driver_cipher_hooks.forced_status );
    /* Failure should prevent further operations from executing on the driver */
    mbedtls_test_driver_cipher_hooks.hits = 0;
    status = psa_cipher_update( &operation,
                                input->x, input->len,
                                output, output_buffer_size,
                                &function_output_length );
    TEST_EQUAL( status, PSA_ERROR_BAD_STATE );
    TEST_EQUAL( mbedtls_test_driver_cipher_hooks.hits, 0 );
    psa_cipher_abort( &operation );

exit:
    psa_cipher_abort( &operation );
    mbedtls_free( output );
    psa_destroy_key( key );
    PSA_DONE( );
    mbedtls_test_driver_cipher_hooks = mbedtls_test_driver_cipher_hooks_init();
}
/* END_CASE */

/* BEGIN_CASE */
void aead_encrypt( int key_type_arg, data_t *key_data,
                   int alg_arg,
                   data_t *nonce,
                   data_t *additional_data,
                   data_t *input_data,
                   data_t *expected_result,
                   int forced_status_arg )
{
    mbedtls_svc_key_id_t key = MBEDTLS_SVC_KEY_ID_INIT;
    psa_key_type_t key_type = key_type_arg;
    psa_algorithm_t alg = alg_arg;
    size_t key_bits;
    psa_status_t forced_status = forced_status_arg;
    unsigned char *output_data = NULL;
    size_t output_size = 0;
    size_t output_length = 0;
    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
    psa_status_t status = PSA_ERROR_GENERIC_ERROR;
    mbedtls_test_driver_aead_hooks = mbedtls_test_driver_aead_hooks_init();

    PSA_ASSERT( psa_crypto_init( ) );

    psa_set_key_usage_flags( &attributes, PSA_KEY_USAGE_ENCRYPT  );
    psa_set_key_algorithm( &attributes, alg );
    psa_set_key_type( &attributes, key_type );

    PSA_ASSERT( psa_import_key( &attributes, key_data->x, key_data->len,
                                &key ) );
    PSA_ASSERT( psa_get_key_attributes( key, &attributes ) );
    key_bits = psa_get_key_bits( &attributes );

    output_size = input_data->len + PSA_AEAD_TAG_LENGTH( key_type, key_bits,
                                                         alg );
    /* For all currently defined algorithms, PSA_AEAD_ENCRYPT_OUTPUT_SIZE
     * should be exact. */
    TEST_EQUAL( output_size,
                PSA_AEAD_ENCRYPT_OUTPUT_SIZE( key_type, alg, input_data->len ) );
    TEST_ASSERT( output_size <=
                 PSA_AEAD_ENCRYPT_OUTPUT_MAX_SIZE( input_data->len ) );
    ASSERT_ALLOC( output_data, output_size );

    mbedtls_test_driver_aead_hooks.forced_status = forced_status;
    status = psa_aead_encrypt( key, alg,
                               nonce->x, nonce->len,
                               additional_data->x, additional_data->len,
                               input_data->x, input_data->len,
                               output_data, output_size,
                               &output_length );
    TEST_EQUAL( mbedtls_test_driver_aead_hooks.hits_encrypt, 1 );
    TEST_EQUAL( mbedtls_test_driver_aead_hooks.driver_status, forced_status );

    TEST_EQUAL( status, ( forced_status == PSA_ERROR_NOT_SUPPORTED ) ?
                        PSA_SUCCESS : forced_status );

    if( status == PSA_SUCCESS )
    {
        ASSERT_COMPARE( expected_result->x, expected_result->len,
                        output_data, output_length );
    }

exit:
    psa_destroy_key( key );
    mbedtls_free( output_data );
    PSA_DONE( );
    mbedtls_test_driver_aead_hooks = mbedtls_test_driver_aead_hooks_init();
}
/* END_CASE */

/* BEGIN_CASE */
void aead_decrypt( int key_type_arg, data_t *key_data,
                   int alg_arg,
                   data_t *nonce,
                   data_t *additional_data,
                   data_t *input_data,
                   data_t *expected_data,
                   int forced_status_arg )
{
    mbedtls_svc_key_id_t key = MBEDTLS_SVC_KEY_ID_INIT;
    psa_key_type_t key_type = key_type_arg;
    psa_algorithm_t alg = alg_arg;
    size_t key_bits;
    psa_status_t forced_status = forced_status_arg;
    unsigned char *output_data = NULL;
    size_t output_size = 0;
    size_t output_length = 0;
    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
    psa_status_t status = PSA_ERROR_GENERIC_ERROR;
    mbedtls_test_driver_aead_hooks = mbedtls_test_driver_aead_hooks_init();

    PSA_ASSERT( psa_crypto_init( ) );

    psa_set_key_usage_flags( &attributes, PSA_KEY_USAGE_DECRYPT  );
    psa_set_key_algorithm( &attributes, alg );
    psa_set_key_type( &attributes, key_type );

    PSA_ASSERT( psa_import_key( &attributes, key_data->x, key_data->len,
                                &key ) );
    PSA_ASSERT( psa_get_key_attributes( key, &attributes ) );
    key_bits = psa_get_key_bits( &attributes );

    output_size = input_data->len - PSA_AEAD_TAG_LENGTH( key_type, key_bits,
                                                         alg );
    ASSERT_ALLOC( output_data, output_size );

    mbedtls_test_driver_aead_hooks.forced_status = forced_status;
    status = psa_aead_decrypt( key, alg,
                               nonce->x, nonce->len,
                               additional_data->x,
                               additional_data->len,
                               input_data->x, input_data->len,
                               output_data, output_size,
                               &output_length );
    TEST_EQUAL( mbedtls_test_driver_aead_hooks.hits_decrypt, 1 );
    TEST_EQUAL( mbedtls_test_driver_aead_hooks.driver_status, forced_status );

    TEST_EQUAL( status, ( forced_status == PSA_ERROR_NOT_SUPPORTED ) ?
                        PSA_SUCCESS : forced_status );

    if( status == PSA_SUCCESS )
    {
        ASSERT_COMPARE( expected_data->x, expected_data->len,
                        output_data, output_length );
    }

exit:
    psa_destroy_key( key );
    mbedtls_free( output_data );
    PSA_DONE( );
    mbedtls_test_driver_aead_hooks = mbedtls_test_driver_aead_hooks_init();
}
/* END_CASE */

/* BEGIN_CASE */
void mac_sign( int key_type_arg,
               data_t *key_data,
               int alg_arg,
               data_t *input,
               data_t *expected_mac,
               int forced_status_arg )
{
    mbedtls_svc_key_id_t key = MBEDTLS_SVC_KEY_ID_INIT;
    psa_key_type_t key_type = key_type_arg;
    psa_algorithm_t alg = alg_arg;
    psa_mac_operation_t operation = PSA_MAC_OPERATION_INIT;
    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
    uint8_t *actual_mac = NULL;
    size_t mac_buffer_size =
        PSA_MAC_LENGTH( key_type, PSA_BYTES_TO_BITS( key_data->len ), alg );
    size_t mac_length = 0;
    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
    psa_status_t forced_status = forced_status_arg;
    mbedtls_test_driver_mac_hooks = mbedtls_test_driver_mac_hooks_init();

    TEST_ASSERT( mac_buffer_size <= PSA_MAC_MAX_SIZE );
    /* We expect PSA_MAC_LENGTH to be exact. */
    TEST_ASSERT( expected_mac->len == mac_buffer_size );

    PSA_ASSERT( psa_crypto_init( ) );

    psa_set_key_usage_flags( &attributes, PSA_KEY_USAGE_SIGN_HASH );
    psa_set_key_algorithm( &attributes, alg );
    psa_set_key_type( &attributes, key_type );

    PSA_ASSERT( psa_import_key( &attributes, key_data->x, key_data->len,
                                &key ) );

    ASSERT_ALLOC( actual_mac, mac_buffer_size );
    mbedtls_test_driver_mac_hooks.forced_status = forced_status;

    /*
     * Calculate the MAC, one-shot case.
     */
    status = psa_mac_compute( key, alg,
                              input->x, input->len,
                              actual_mac, mac_buffer_size,
                              &mac_length );

    TEST_EQUAL( mbedtls_test_driver_mac_hooks.hits, 1 );
    if( forced_status == PSA_SUCCESS ||
        forced_status == PSA_ERROR_NOT_SUPPORTED )
    {
        PSA_ASSERT( status );
    }
    else
        TEST_EQUAL( forced_status, status );

    PSA_ASSERT( psa_mac_abort( &operation ) );
    TEST_EQUAL( mbedtls_test_driver_mac_hooks.hits, 1 );

    if( forced_status == PSA_SUCCESS )
    {
        ASSERT_COMPARE( expected_mac->x, expected_mac->len,
                        actual_mac, mac_length );
    }

    mbedtls_free( actual_mac );
    actual_mac = NULL;

exit:
    psa_mac_abort( &operation );
    psa_destroy_key( key );
    PSA_DONE( );
    mbedtls_free( actual_mac );
    mbedtls_test_driver_mac_hooks = mbedtls_test_driver_mac_hooks_init();
}
/* END_CASE */

/* BEGIN_CASE */
void mac_sign_multipart( int key_type_arg,
                         data_t *key_data,
                         int alg_arg,
                         data_t *input,
                         data_t *expected_mac,
                         int fragments_count,
                         int forced_status_arg )
{
    mbedtls_svc_key_id_t key = MBEDTLS_SVC_KEY_ID_INIT;
    psa_key_type_t key_type = key_type_arg;
    psa_algorithm_t alg = alg_arg;
    psa_mac_operation_t operation = PSA_MAC_OPERATION_INIT;
    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
    uint8_t *actual_mac = NULL;
    size_t mac_buffer_size =
        PSA_MAC_LENGTH( key_type, PSA_BYTES_TO_BITS( key_data->len ), alg );
    size_t mac_length = 0;
    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
    psa_status_t forced_status = forced_status_arg;
    uint8_t *input_x = input->x;
    mbedtls_test_driver_mac_hooks = mbedtls_test_driver_mac_hooks_init();

    TEST_ASSERT( mac_buffer_size <= PSA_MAC_MAX_SIZE );
    /* We expect PSA_MAC_LENGTH to be exact. */
    TEST_ASSERT( expected_mac->len == mac_buffer_size );

    PSA_ASSERT( psa_crypto_init( ) );

    psa_set_key_usage_flags( &attributes, PSA_KEY_USAGE_SIGN_HASH );
    psa_set_key_algorithm( &attributes, alg );
    psa_set_key_type( &attributes, key_type );

    PSA_ASSERT( psa_import_key( &attributes, key_data->x, key_data->len,
                                &key ) );

    ASSERT_ALLOC( actual_mac, mac_buffer_size );
    mbedtls_test_driver_mac_hooks.forced_status = forced_status;

    /*
     * Calculate the MAC, multipart case.
     */
    status = psa_mac_sign_setup( &operation, key, alg );
    TEST_EQUAL( mbedtls_test_driver_mac_hooks.hits, 1 );

    if( forced_status == PSA_SUCCESS ||
        forced_status == PSA_ERROR_NOT_SUPPORTED )
    {
        PSA_ASSERT( status );
    }
    else
        TEST_EQUAL( forced_status, status );

    if ( fragments_count )
    {
        TEST_ASSERT( ( input->len / fragments_count ) > 0 );
    }

    for ( int i = 0; i < fragments_count; i++)
    {
        int fragment_size = input->len / fragments_count;
        if ( i == fragments_count - 1 )
            fragment_size += ( input->len % fragments_count );

        status = psa_mac_update( &operation,
                                input_x, fragment_size );
        if( forced_status == PSA_SUCCESS )
            TEST_EQUAL( mbedtls_test_driver_mac_hooks.hits, 2 + i );
        else
            TEST_EQUAL( mbedtls_test_driver_mac_hooks.hits, 1 );
        if( forced_status == PSA_SUCCESS ||
            forced_status == PSA_ERROR_NOT_SUPPORTED )
        {
            PSA_ASSERT( status );
        }
        else
            TEST_EQUAL( PSA_ERROR_BAD_STATE, status );
        input_x += fragment_size;
    }

    status = psa_mac_sign_finish( &operation,
                                  actual_mac, mac_buffer_size,
                                  &mac_length );
    if( forced_status == PSA_SUCCESS )
        TEST_EQUAL( mbedtls_test_driver_mac_hooks.hits, 3 + fragments_count );
    else
        TEST_EQUAL( mbedtls_test_driver_mac_hooks.hits, 1 );

    if( forced_status == PSA_SUCCESS ||
        forced_status == PSA_ERROR_NOT_SUPPORTED )
    {
        PSA_ASSERT( status );
    }
    else
        TEST_EQUAL( PSA_ERROR_BAD_STATE, status );

    PSA_ASSERT( psa_mac_abort( &operation ) );
    if( forced_status == PSA_SUCCESS )
        TEST_EQUAL( mbedtls_test_driver_mac_hooks.hits, 3 + fragments_count );
    else
        TEST_EQUAL( mbedtls_test_driver_mac_hooks.hits, 1 );

    if( forced_status == PSA_SUCCESS )
    {
        ASSERT_COMPARE( expected_mac->x, expected_mac->len,
                        actual_mac, mac_length );
    }

    mbedtls_free( actual_mac );
    actual_mac = NULL;

exit:
    psa_mac_abort( &operation );
    psa_destroy_key( key );
    PSA_DONE( );
    mbedtls_free( actual_mac );
    mbedtls_test_driver_mac_hooks = mbedtls_test_driver_mac_hooks_init();
}
/* END_CASE */

/* BEGIN_CASE */
void mac_verify( int key_type_arg,
                 data_t *key_data,
                 int alg_arg,
                 data_t *input,
                 data_t *expected_mac,
                 int forced_status_arg )
{
    mbedtls_svc_key_id_t key = MBEDTLS_SVC_KEY_ID_INIT;
    psa_key_type_t key_type = key_type_arg;
    psa_algorithm_t alg = alg_arg;
    psa_mac_operation_t operation = PSA_MAC_OPERATION_INIT;
    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
    psa_status_t status = PSA_ERROR_GENERIC_ERROR;
    psa_status_t forced_status = forced_status_arg;
    mbedtls_test_driver_mac_hooks = mbedtls_test_driver_mac_hooks_init();

    TEST_ASSERT( expected_mac->len <= PSA_MAC_MAX_SIZE );

    PSA_ASSERT( psa_crypto_init( ) );

    psa_set_key_usage_flags( &attributes, PSA_KEY_USAGE_VERIFY_HASH );
    psa_set_key_algorithm( &attributes, alg );
    psa_set_key_type( &attributes, key_type );

    PSA_ASSERT( psa_import_key( &attributes, key_data->x, key_data->len,
                                &key ) );

    mbedtls_test_driver_mac_hooks.forced_status = forced_status;

    /*
     * Verify the MAC, one-shot case.
     */
    status = psa_mac_verify( key, alg,
                             input->x, input->len,
                             expected_mac->x, expected_mac->len );
    TEST_EQUAL( mbedtls_test_driver_mac_hooks.hits, 1 );
    if( forced_status == PSA_SUCCESS ||
        forced_status == PSA_ERROR_NOT_SUPPORTED )
    {
        PSA_ASSERT( status );
    }
    else
        TEST_EQUAL( forced_status, status );

    PSA_ASSERT( psa_mac_abort( &operation ) );
    TEST_EQUAL( mbedtls_test_driver_mac_hooks.hits, 1 );
exit:
    psa_mac_abort( &operation );
    psa_destroy_key( key );
    PSA_DONE( );
    mbedtls_test_driver_mac_hooks = mbedtls_test_driver_mac_hooks_init();
}
/* END_CASE */

/* BEGIN_CASE */
void mac_verify_multipart( int key_type_arg,
                           data_t *key_data,
                           int alg_arg,
                           data_t *input,
                           data_t *expected_mac,
                           int fragments_count,
                           int forced_status_arg )
{
    mbedtls_svc_key_id_t key = MBEDTLS_SVC_KEY_ID_INIT;
    psa_key_type_t key_type = key_type_arg;
    psa_algorithm_t alg = alg_arg;
    psa_mac_operation_t operation = PSA_MAC_OPERATION_INIT;
    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
    psa_status_t status = PSA_ERROR_GENERIC_ERROR;
    psa_status_t forced_status = forced_status_arg;
    uint8_t *input_x = input->x;
    mbedtls_test_driver_mac_hooks = mbedtls_test_driver_mac_hooks_init();

    TEST_ASSERT( expected_mac->len <= PSA_MAC_MAX_SIZE );

    PSA_ASSERT( psa_crypto_init( ) );

    psa_set_key_usage_flags( &attributes, PSA_KEY_USAGE_VERIFY_HASH );
    psa_set_key_algorithm( &attributes, alg );
    psa_set_key_type( &attributes, key_type );

    PSA_ASSERT( psa_import_key( &attributes, key_data->x, key_data->len,
                                &key ) );

    mbedtls_test_driver_mac_hooks.forced_status = forced_status;

    /*
     * Verify the MAC, multi-part case.
     */
    status = psa_mac_verify_setup( &operation, key, alg );
    TEST_EQUAL( mbedtls_test_driver_mac_hooks.hits, 1 );

    if( forced_status == PSA_SUCCESS ||
        forced_status == PSA_ERROR_NOT_SUPPORTED )
    {
        PSA_ASSERT( status );
    }
    else
        TEST_EQUAL( forced_status, status );

    if ( fragments_count )
    {
        TEST_ASSERT( ( input->len / fragments_count ) > 0 );
    }

    for ( int i = 0; i < fragments_count; i++)
    {
        int fragment_size = input->len / fragments_count;
        if ( i == fragments_count - 1 )
            fragment_size += ( input->len % fragments_count );

        status = psa_mac_update( &operation,
                                input_x, fragment_size );
        if( forced_status == PSA_SUCCESS )
            TEST_EQUAL( mbedtls_test_driver_mac_hooks.hits, 2 + i );
        else
            TEST_EQUAL( mbedtls_test_driver_mac_hooks.hits, 1 );

        if( forced_status == PSA_SUCCESS ||
            forced_status == PSA_ERROR_NOT_SUPPORTED )
        {
            PSA_ASSERT( status );
        }
        else
            TEST_EQUAL( PSA_ERROR_BAD_STATE, status );
        input_x += fragment_size;
    }

    status = psa_mac_verify_finish( &operation,
                                    expected_mac->x,
                                    expected_mac->len );
    if( forced_status == PSA_SUCCESS )
        TEST_EQUAL( mbedtls_test_driver_mac_hooks.hits, 3 + fragments_count );
    else
        TEST_EQUAL( mbedtls_test_driver_mac_hooks.hits, 1 );

    if( forced_status == PSA_SUCCESS ||
        forced_status == PSA_ERROR_NOT_SUPPORTED )
    {
        PSA_ASSERT( status );
    }
    else
        TEST_EQUAL( PSA_ERROR_BAD_STATE, status );


    PSA_ASSERT( psa_mac_abort( &operation ) );
    if( forced_status == PSA_SUCCESS )
        TEST_EQUAL( mbedtls_test_driver_mac_hooks.hits, 3 + fragments_count );
    else
        TEST_EQUAL( mbedtls_test_driver_mac_hooks.hits, 1 );

exit:
    psa_mac_abort( &operation );
    psa_destroy_key( key );
    PSA_DONE( );
    mbedtls_test_driver_mac_hooks = mbedtls_test_driver_mac_hooks_init();
}
/* END_CASE */

/* BEGIN_CASE depends_on:PSA_CRYPTO_DRIVER_TEST:MBEDTLS_PSA_CRYPTO_DRIVERS:MBEDTLS_PSA_CRYPTO_BUILTIN_KEYS */
void builtin_key_export( int builtin_key_id_arg,
                         int builtin_key_type_arg,
                         int builtin_key_bits_arg,
                         int builtin_key_algorithm_arg,
                         data_t *expected_output,
                         int expected_status_arg )
{
    psa_key_id_t builtin_key_id = (psa_key_id_t) builtin_key_id_arg;
    psa_key_type_t builtin_key_type = (psa_key_type_t) builtin_key_type_arg;
    psa_algorithm_t builtin_key_alg = (psa_algorithm_t) builtin_key_algorithm_arg;
    size_t builtin_key_bits = (size_t) builtin_key_bits_arg;
    psa_status_t expected_status = expected_status_arg;
    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;

    mbedtls_svc_key_id_t key = mbedtls_svc_key_id_make( 0, builtin_key_id );
    uint8_t* output_buffer = NULL;
    size_t output_size = 0;
    psa_status_t actual_status;

    PSA_ASSERT( psa_crypto_init( ) );
    ASSERT_ALLOC( output_buffer, expected_output->len );

    actual_status = psa_export_key( key, output_buffer, expected_output->len, &output_size );

    if( expected_status == PSA_SUCCESS )
    {
        PSA_ASSERT( actual_status );
        TEST_EQUAL( output_size, expected_output->len );
        ASSERT_COMPARE( output_buffer, output_size,
                        expected_output->x, expected_output->len );

        PSA_ASSERT( psa_get_key_attributes( key, &attributes ) );
        TEST_EQUAL( psa_get_key_bits( &attributes ), builtin_key_bits );
        TEST_EQUAL( psa_get_key_type( &attributes ), builtin_key_type );
        TEST_EQUAL( psa_get_key_algorithm( &attributes ), builtin_key_alg );
    }
    else
    {
        if( actual_status != expected_status )
            fprintf( stderr, "Expected %d but got %d\n", expected_status, actual_status );
        TEST_EQUAL( actual_status, expected_status );
        TEST_EQUAL( output_size, 0 );
    }

exit:
    mbedtls_free( output_buffer );
    psa_reset_key_attributes( &attributes );
    psa_destroy_key( key );
    PSA_DONE( );
}
/* END_CASE */

/* BEGIN_CASE depends_on:PSA_CRYPTO_DRIVER_TEST:MBEDTLS_PSA_CRYPTO_DRIVERS:MBEDTLS_PSA_CRYPTO_BUILTIN_KEYS */
void builtin_pubkey_export( int builtin_key_id_arg,
                            int builtin_key_type_arg,
                            int builtin_key_bits_arg,
                            int builtin_key_algorithm_arg,
                            data_t *expected_output,
                            int expected_status_arg )
{
    psa_key_id_t builtin_key_id = (psa_key_id_t) builtin_key_id_arg;
    psa_key_type_t builtin_key_type = (psa_key_type_t) builtin_key_type_arg;
    psa_algorithm_t builtin_key_alg = (psa_algorithm_t) builtin_key_algorithm_arg;
    size_t builtin_key_bits = (size_t) builtin_key_bits_arg;
    psa_status_t expected_status = expected_status_arg;
    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;

    mbedtls_svc_key_id_t key = mbedtls_svc_key_id_make( 0, builtin_key_id );
    uint8_t* output_buffer = NULL;
    size_t output_size = 0;
    psa_status_t actual_status;

    PSA_ASSERT( psa_crypto_init( ) );
    ASSERT_ALLOC( output_buffer, expected_output->len );

    actual_status = psa_export_public_key( key, output_buffer, expected_output->len, &output_size );

    if( expected_status == PSA_SUCCESS )
    {
        PSA_ASSERT( actual_status );
        TEST_EQUAL( output_size, expected_output->len );
        ASSERT_COMPARE( output_buffer, output_size,
                        expected_output->x, expected_output->len );

        PSA_ASSERT( psa_get_key_attributes( key, &attributes ) );
        TEST_EQUAL( psa_get_key_bits( &attributes ), builtin_key_bits );
        TEST_EQUAL( psa_get_key_type( &attributes ), builtin_key_type );
        TEST_EQUAL( psa_get_key_algorithm( &attributes ), builtin_key_alg );
    }
    else
    {
        TEST_EQUAL( actual_status, expected_status );
        TEST_EQUAL( output_size, 0 );
    }

exit:
    mbedtls_free( output_buffer );
    psa_reset_key_attributes( &attributes );
    psa_destroy_key( key );
    PSA_DONE( );
}
/* END_CASE */

/* BEGIN_CASE */
void hash_compute( int alg_arg,
                   data_t *input, data_t *hash,
                   int forced_status_arg,
                   int expected_status_arg )
{
    psa_algorithm_t alg = alg_arg;
    psa_status_t forced_status = forced_status_arg;
    psa_status_t expected_status = expected_status_arg;
    unsigned char *output = NULL;
    size_t output_length;

    mbedtls_test_driver_hash_hooks = mbedtls_test_driver_hash_hooks_init();
    mbedtls_test_driver_hash_hooks.forced_status = forced_status;

    PSA_ASSERT( psa_crypto_init( ) );
    ASSERT_ALLOC( output, PSA_HASH_LENGTH( alg ) );

    TEST_EQUAL( psa_hash_compute( alg, input->x, input->len,
                                  output, PSA_HASH_LENGTH( alg ),
                                  &output_length ), expected_status );
    TEST_EQUAL( mbedtls_test_driver_hash_hooks.hits, 1 );
    TEST_EQUAL( mbedtls_test_driver_hash_hooks.driver_status, forced_status );

    if( expected_status == PSA_SUCCESS )
    {
        ASSERT_COMPARE( output, output_length, hash->x, hash->len );
    }

exit:
    mbedtls_free( output );
    PSA_DONE( );
    mbedtls_test_driver_hash_hooks = mbedtls_test_driver_hash_hooks_init();
}
/* END_CASE */

/* BEGIN_CASE */
void hash_multipart_setup( int alg_arg,
                           data_t *input, data_t *hash,
                           int forced_status_arg,
                           int expected_status_arg )
{
    psa_algorithm_t alg = alg_arg;
    psa_status_t forced_status = forced_status_arg;
    psa_status_t expected_status = expected_status_arg;
    unsigned char *output = NULL;
    psa_hash_operation_t operation = PSA_HASH_OPERATION_INIT;
    size_t output_length;

    mbedtls_test_driver_hash_hooks = mbedtls_test_driver_hash_hooks_init();
    ASSERT_ALLOC( output, PSA_HASH_LENGTH( alg ) );

    PSA_ASSERT( psa_crypto_init( ) );

    mbedtls_test_driver_hash_hooks.forced_status = forced_status;
    TEST_EQUAL( psa_hash_setup( &operation, alg ), expected_status );
    TEST_EQUAL( mbedtls_test_driver_hash_hooks.hits, 1 );
    TEST_EQUAL( mbedtls_test_driver_hash_hooks.driver_status, forced_status );

    if( expected_status == PSA_SUCCESS )
    {
        PSA_ASSERT( psa_hash_update( &operation, input->x, input->len ) );
        TEST_EQUAL( mbedtls_test_driver_hash_hooks.hits,
                    forced_status == PSA_ERROR_NOT_SUPPORTED ? 1 : 2 );
        TEST_EQUAL( mbedtls_test_driver_hash_hooks.driver_status, forced_status );

        PSA_ASSERT( psa_hash_finish( &operation,
                                     output, PSA_HASH_LENGTH( alg ),
                                     &output_length ) );
        TEST_EQUAL( mbedtls_test_driver_hash_hooks.hits,
                    forced_status == PSA_ERROR_NOT_SUPPORTED ? 1 : 4 );
        TEST_EQUAL( mbedtls_test_driver_hash_hooks.driver_status, forced_status );

        ASSERT_COMPARE( output, output_length, hash->x, hash->len );
    }

exit:
    psa_hash_abort( &operation );
    mbedtls_free( output );
    PSA_DONE( );
    mbedtls_test_driver_hash_hooks = mbedtls_test_driver_hash_hooks_init();
}
/* END_CASE */

/* BEGIN_CASE */
void hash_multipart_update( int alg_arg,
                            data_t *input, data_t *hash,
                            int forced_status_arg )
{
    psa_algorithm_t alg = alg_arg;
    psa_status_t forced_status = forced_status_arg;
    unsigned char *output = NULL;
    psa_hash_operation_t operation = PSA_HASH_OPERATION_INIT;
    size_t output_length;

    mbedtls_test_driver_hash_hooks = mbedtls_test_driver_hash_hooks_init();
    ASSERT_ALLOC( output, PSA_HASH_LENGTH( alg ) );

    PSA_ASSERT( psa_crypto_init( ) );

    /*
     * Update inactive operation, the driver shouldn't be called.
     */
    TEST_EQUAL( psa_hash_update( &operation, input->x, input->len ),
                PSA_ERROR_BAD_STATE );
    TEST_EQUAL( mbedtls_test_driver_hash_hooks.hits, 0 );

    PSA_ASSERT( psa_hash_setup( &operation, alg ) );
    TEST_EQUAL( mbedtls_test_driver_hash_hooks.hits, 1 );
    TEST_EQUAL( mbedtls_test_driver_hash_hooks.driver_status, PSA_SUCCESS );

    mbedtls_test_driver_hash_hooks.forced_status = forced_status;
    TEST_EQUAL( psa_hash_update( &operation, input->x, input->len ),
                forced_status );
    /* One or two more calls to the driver interface: update or update + abort */
    TEST_EQUAL( mbedtls_test_driver_hash_hooks.hits,
                forced_status == PSA_SUCCESS ? 2 : 3 );
    TEST_EQUAL( mbedtls_test_driver_hash_hooks.driver_status, forced_status );

    if( forced_status == PSA_SUCCESS )
    {
        mbedtls_test_driver_hash_hooks = mbedtls_test_driver_hash_hooks_init();
        PSA_ASSERT( psa_hash_finish( &operation,
                                     output, PSA_HASH_LENGTH( alg ),
                                     &output_length ) );
        /* Two calls to the driver interface: update + abort */
        TEST_EQUAL( mbedtls_test_driver_hash_hooks.hits, 2 );
        TEST_EQUAL( mbedtls_test_driver_hash_hooks.driver_status, PSA_SUCCESS );

        ASSERT_COMPARE( output, output_length, hash->x, hash->len );
    }

exit:
    psa_hash_abort( &operation );
    mbedtls_free( output );
    PSA_DONE( );
    mbedtls_test_driver_hash_hooks = mbedtls_test_driver_hash_hooks_init();
}
/* END_CASE */

/* BEGIN_CASE */
void hash_multipart_finish( int alg_arg,
                            data_t *input, data_t *hash,
                            int forced_status_arg )
{
    psa_algorithm_t alg = alg_arg;
    psa_status_t forced_status = forced_status_arg;
    unsigned char *output = NULL;
    psa_hash_operation_t operation = PSA_HASH_OPERATION_INIT;
    size_t output_length;

    mbedtls_test_driver_hash_hooks = mbedtls_test_driver_hash_hooks_init();
    ASSERT_ALLOC( output, PSA_HASH_LENGTH( alg ) );

    PSA_ASSERT( psa_crypto_init( ) );

    /*
     * Finish inactive operation, the driver shouldn't be called.
     */
    TEST_EQUAL( psa_hash_finish( &operation, output, PSA_HASH_LENGTH( alg ),
                                 &output_length ),
                PSA_ERROR_BAD_STATE );
    TEST_EQUAL( mbedtls_test_driver_hash_hooks.hits, 0 );

    PSA_ASSERT( psa_hash_setup( &operation, alg ) );
    TEST_EQUAL( mbedtls_test_driver_hash_hooks.hits, 1 );
    TEST_EQUAL( mbedtls_test_driver_hash_hooks.driver_status, PSA_SUCCESS );

    PSA_ASSERT( psa_hash_update( &operation, input->x, input->len ) );
    TEST_EQUAL( mbedtls_test_driver_hash_hooks.hits, 2 );
    TEST_EQUAL( mbedtls_test_driver_hash_hooks.driver_status, PSA_SUCCESS );

    mbedtls_test_driver_hash_hooks.forced_status = forced_status;
    TEST_EQUAL( psa_hash_finish( &operation,
                                 output, PSA_HASH_LENGTH( alg ),
                                 &output_length ),
                forced_status );
    /* Two more calls to the driver interface: finish + abort */
    TEST_EQUAL( mbedtls_test_driver_hash_hooks.hits, 4 );
    TEST_EQUAL( mbedtls_test_driver_hash_hooks.driver_status, forced_status );

    if( forced_status == PSA_SUCCESS )
        ASSERT_COMPARE( output, output_length, hash->x, hash->len );

exit:
    psa_hash_abort( &operation );
    mbedtls_free( output );
    PSA_DONE( );
    mbedtls_test_driver_hash_hooks = mbedtls_test_driver_hash_hooks_init();
}
/* END_CASE */

/* BEGIN_CASE */
void hash_clone( int alg_arg,
                 data_t *input, data_t *hash,
                 int forced_status_arg )
{
    psa_algorithm_t alg = alg_arg;
    psa_status_t forced_status = forced_status_arg;
    unsigned char *output = NULL;
    psa_hash_operation_t source_operation = PSA_HASH_OPERATION_INIT;
    psa_hash_operation_t target_operation = PSA_HASH_OPERATION_INIT;
    size_t output_length;

    mbedtls_test_driver_hash_hooks = mbedtls_test_driver_hash_hooks_init();
    ASSERT_ALLOC( output, PSA_HASH_LENGTH( alg ) );

    PSA_ASSERT( psa_crypto_init( ) );

    /*
     * Clone inactive operation, the driver shouldn't be called.
     */
    TEST_EQUAL( psa_hash_clone( &source_operation, &target_operation ),
                PSA_ERROR_BAD_STATE );
    TEST_EQUAL( mbedtls_test_driver_hash_hooks.hits, 0 );

    PSA_ASSERT( psa_hash_setup( &source_operation, alg ) );
    TEST_EQUAL( mbedtls_test_driver_hash_hooks.hits, 1 );
    TEST_EQUAL( mbedtls_test_driver_hash_hooks.driver_status, PSA_SUCCESS );

    mbedtls_test_driver_hash_hooks.forced_status = forced_status;
    TEST_EQUAL( psa_hash_clone( &source_operation, &target_operation ),
                forced_status );
    TEST_EQUAL( mbedtls_test_driver_hash_hooks.hits,
                forced_status == PSA_SUCCESS ? 2 : 3 );
    TEST_EQUAL( mbedtls_test_driver_hash_hooks.driver_status, forced_status );

    if( forced_status == PSA_SUCCESS )
    {
        mbedtls_test_driver_hash_hooks = mbedtls_test_driver_hash_hooks_init();
        PSA_ASSERT( psa_hash_update( &target_operation,
                                     input->x, input->len ) );
        TEST_EQUAL( mbedtls_test_driver_hash_hooks.hits, 1 );
        TEST_EQUAL( mbedtls_test_driver_hash_hooks.driver_status, PSA_SUCCESS );

        PSA_ASSERT( psa_hash_finish( &target_operation,
                                     output, PSA_HASH_LENGTH( alg ),
                                     &output_length ) );
        TEST_EQUAL( mbedtls_test_driver_hash_hooks.hits, 3 );
        TEST_EQUAL( mbedtls_test_driver_hash_hooks.driver_status, PSA_SUCCESS );

        ASSERT_COMPARE( output, output_length, hash->x, hash->len );
    }

exit:
    psa_hash_abort( &source_operation );
    psa_hash_abort( &target_operation );
    mbedtls_free( output );
    PSA_DONE( );
    mbedtls_test_driver_hash_hooks = mbedtls_test_driver_hash_hooks_init();
}
/* END_CASE */