/* BEGIN_HEADER */
#include "psa_crypto_helpers.h"
#include "psa/crypto_se_driver.h"

#include "psa_crypto_se.h"
#include "psa_crypto_storage.h"

/** The minimum valid lifetime value for a secure element driver. */
#define MIN_DRIVER_LIFETIME 2

/** The driver detected a condition that shouldn't happen.
 * This is probably a bug in the library. */
#define PSA_ERROR_DETECTED_BY_DRIVER ((psa_status_t)( -500 ))

/** Like #TEST_ASSERT for use in a driver method.
 *
 * Use this macro to assert on guarantees provided by the core.
 */
#define DRIVER_ASSERT( TEST )                               \
    do {                                                    \
       if( ! (TEST) )                                       \
       {                                                    \
          test_fail( #TEST, __LINE__, __FILE__ );           \
          return( PSA_ERROR_DETECTED_BY_DRIVER );           \
       }                                                    \
    } while( 0 )

#define RAM_MAX_KEY_SIZE 64
typedef struct
{
    psa_key_lifetime_t lifetime;
    psa_key_type_t type;
    size_t bits;
    uint8_t content[RAM_MAX_KEY_SIZE];
} ram_slot_t;
static ram_slot_t ram_slots[16];

/* A type with at least ARRAY_LENGTH(ram_slots) bits, containing a
 * bit vector indicating which slots are in use. */
typedef uint16_t ram_slot_usage_t;

static uint8_t ram_min_slot = 0;

static void ram_slots_reset( void )
{
    memset( ram_slots, 0, sizeof( ram_slots ) );
    ram_min_slot = 0;
}

static psa_status_t ram_import( psa_drv_se_context_t *context,
                                psa_key_slot_number_t slot_number,
                                psa_key_lifetime_t lifetime,
                                psa_key_type_t type,
                                psa_algorithm_t algorithm,
                                psa_key_usage_t usage,
                                const uint8_t *p_data,
                                size_t data_length )
{
    (void) context;
    DRIVER_ASSERT( slot_number < ARRAY_LENGTH( ram_slots ) );
    if( data_length > sizeof( ram_slots[slot_number].content ) )
        return( PSA_ERROR_INSUFFICIENT_STORAGE );
    ram_slots[slot_number].lifetime = lifetime;
    ram_slots[slot_number].type = type;
    ram_slots[slot_number].bits = PSA_BYTES_TO_BITS( data_length );
    (void) algorithm;
    (void) usage;
    memcpy( ram_slots[slot_number].content, p_data, data_length );
    return( PSA_SUCCESS );
}

psa_status_t ram_export( psa_drv_se_context_t *context,
                         psa_key_slot_number_t slot_number,
                         uint8_t *p_data,
                         size_t data_size,
                         size_t *p_data_length )
{
    size_t actual_size;
    (void) context;
    DRIVER_ASSERT( slot_number < ARRAY_LENGTH( ram_slots ) );
    actual_size = PSA_BITS_TO_BYTES( ram_slots[slot_number].bits );
    if( actual_size > data_size )
        return( PSA_ERROR_BUFFER_TOO_SMALL );
    *p_data_length = actual_size;
    memcpy( p_data, ram_slots[slot_number].content, actual_size );
    return( PSA_SUCCESS );
}

psa_status_t ram_destroy( psa_drv_se_context_t *context,
                          void *persistent_data,
                          psa_key_slot_number_t slot_number )
{
    ram_slot_usage_t *slot_usage = persistent_data;
    DRIVER_ASSERT( context->persistent_data_size == sizeof( ram_slot_usage_t ) );
    DRIVER_ASSERT( slot_number < ARRAY_LENGTH( ram_slots ) );
    memset( &ram_slots[slot_number], 0, sizeof( ram_slots[slot_number] ) );
    *slot_usage &= ~(ram_slot_usage_t)( 1 << slot_number );
    return( PSA_SUCCESS );
}

psa_status_t ram_allocate( psa_drv_se_context_t *context,
                           void *persistent_data,
                           const psa_key_attributes_t *attributes,
                           psa_key_slot_number_t *slot_number )
{
    ram_slot_usage_t *slot_usage = persistent_data;
    (void) attributes;
    DRIVER_ASSERT( context->persistent_data_size == sizeof( ram_slot_usage_t ) );
    for( *slot_number = ram_min_slot;
         *slot_number < ARRAY_LENGTH( ram_slots );
         ++( *slot_number ) )
    {
        if( ! ( *slot_usage & 1 << *slot_number ) )
            return( PSA_SUCCESS );
    }
    return( PSA_ERROR_INSUFFICIENT_STORAGE );
}

#define MAX_KEY_ID_FOR_TEST 10
void psa_purge_storage( void )
{
    psa_key_id_t id;
    psa_key_lifetime_t lifetime;
    /* The tests may have potentially created key ids from 1 to
     * MAX_KEY_ID_FOR_TEST. In addition, run the destroy function on key id
     * 0, which file-based storage uses as a temporary file. */
    for( id = 0; id <= MAX_KEY_ID_FOR_TEST; id++ )
        psa_destroy_persistent_key( id );
    /* Purge the transaction file. */
    psa_crypto_stop_transaction( );
    /* Purge driver persistent data. */
    for( lifetime = 0; lifetime < PSA_MAX_SE_LIFETIME; lifetime++ )
        psa_destroy_se_persistent_data( lifetime );
}

/* END_HEADER */

/* BEGIN_DEPENDENCIES
 * depends_on:MBEDTLS_PSA_CRYPTO_SE_C
 * END_DEPENDENCIES
 */

/* BEGIN_CASE */
void register_one( int lifetime, int version, int expected_status_arg )
{
    psa_status_t expected_status = expected_status_arg;
    psa_drv_se_t driver;

    memset( &driver, 0, sizeof( driver ) );
    driver.hal_version = version;

    TEST_EQUAL( psa_register_se_driver( lifetime, &driver ),
                expected_status );

    PSA_ASSERT( psa_crypto_init( ) );

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

/* BEGIN_CASE */
void register_twice( int count )
{
    psa_drv_se_t driver;
    psa_key_lifetime_t lifetime;
    psa_key_lifetime_t max = MIN_DRIVER_LIFETIME + count;

    memset( &driver, 0, sizeof( driver ) );
    driver.hal_version = PSA_DRV_SE_HAL_VERSION;

    for( lifetime = MIN_DRIVER_LIFETIME; lifetime < max; lifetime++ )
        PSA_ASSERT( psa_register_se_driver( lifetime, &driver ) );
    for( lifetime = MIN_DRIVER_LIFETIME; lifetime < max; lifetime++ )
        TEST_EQUAL( psa_register_se_driver( lifetime, &driver ),
                    PSA_ERROR_ALREADY_EXISTS );

    PSA_ASSERT( psa_crypto_init( ) );

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

/* BEGIN_CASE */
void register_max( )
{
    psa_drv_se_t driver;
    psa_key_lifetime_t lifetime;
    psa_key_lifetime_t max = MIN_DRIVER_LIFETIME + PSA_MAX_SE_DRIVERS;

    memset( &driver, 0, sizeof( driver ) );
    driver.hal_version = PSA_DRV_SE_HAL_VERSION;

    for( lifetime = MIN_DRIVER_LIFETIME; lifetime < max; lifetime++ )
        PSA_ASSERT( psa_register_se_driver( lifetime, &driver ) );

    TEST_EQUAL( psa_register_se_driver( lifetime, &driver ),
                PSA_ERROR_INSUFFICIENT_MEMORY );

    PSA_ASSERT( psa_crypto_init( ) );

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

/* BEGIN_CASE */
void key_creation_import_export( int min_slot, int restart )
{
    psa_drv_se_t driver;
    psa_drv_se_key_management_t key_management;
    psa_key_lifetime_t lifetime = 2;
    psa_key_id_t id = 1;
    psa_key_handle_t handle = 0;
    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
    const uint8_t key_material[3] = {0xfa, 0xca, 0xde};
    uint8_t exported[sizeof( key_material )];
    size_t exported_length;

    memset( &driver, 0, sizeof( driver ) );
    memset( &key_management, 0, sizeof( key_management ) );
    driver.hal_version = PSA_DRV_SE_HAL_VERSION;
    driver.key_management = &key_management;
    driver.persistent_data_size = sizeof( ram_slot_usage_t );
    key_management.p_allocate = ram_allocate;
    key_management.p_import = ram_import;
    key_management.p_destroy = ram_destroy;
    key_management.p_export = ram_export;
    ram_min_slot = min_slot;

    PSA_ASSERT( psa_register_se_driver( lifetime, &driver ) );
    PSA_ASSERT( psa_crypto_init( ) );

    /* Create a key. */
    psa_set_key_id( &attributes, id );
    psa_set_key_lifetime( &attributes, lifetime );
    psa_set_key_usage_flags( &attributes, PSA_KEY_USAGE_EXPORT );
    psa_set_key_type( &attributes, PSA_KEY_TYPE_RAW_DATA );
    PSA_ASSERT( psa_import_key( &attributes,
                                key_material, sizeof( key_material ),
                                &handle ) );

    /* Maybe restart, to check that the information is saved correctly. */
    if( restart )
    {
        mbedtls_psa_crypto_free( );
        PSA_ASSERT( psa_register_se_driver( lifetime, &driver ) );
        PSA_ASSERT( psa_crypto_init( ) );
        PSA_ASSERT( psa_open_key( id, &handle ) );
    }

    /* Test that the key was created in the expected slot. */
    TEST_ASSERT( ram_slots[min_slot].type == PSA_KEY_TYPE_RAW_DATA );

    PSA_ASSERT( psa_export_key( handle,
                                exported, sizeof( exported ),
                                &exported_length ) );
    ASSERT_COMPARE( key_material, sizeof( key_material ),
                    exported, exported_length );

    PSA_ASSERT( psa_destroy_key( handle ) );

    /* Test that the key has been erased from the designated slot. */
    TEST_ASSERT( ram_slots[min_slot].type == 0 );

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