Add key id check when creating a volatile key

Signed-off-by: Ronald Cron <ronald.cron@arm.com>
diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 82e2549..e45c52e 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -1777,6 +1777,7 @@
 {
     psa_status_t status = PSA_ERROR_INVALID_ARGUMENT;
     psa_key_lifetime_t lifetime = psa_get_key_lifetime( attributes );
+    mbedtls_svc_key_id_t key = psa_get_key_id( attributes );
 
     status = psa_validate_key_location( psa_get_key_lifetime( attributes ),
                                         p_drv );
@@ -1787,8 +1788,12 @@
     if( status != PSA_SUCCESS )
         return( status );
 
-    /* Validate the key identifier only in the case of a persistent key. */
-    if ( ! PSA_KEY_LIFETIME_IS_VOLATILE( lifetime ) )
+    if ( PSA_KEY_LIFETIME_IS_VOLATILE( lifetime ) )
+    {
+        if( MBEDTLS_SVC_KEY_ID_GET_KEY_ID( key ) != 0 )
+            return( PSA_ERROR_INVALID_ARGUMENT );
+    }
+    else
     {
         status = psa_validate_key_id(
             psa_get_key_id( attributes ),
diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index 204e36e..8279768 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -2340,7 +2340,12 @@
 
     /* Prepare the target attributes. */
     if( copy_attributes )
+    {
         target_attributes = source_attributes;
+        /* Set volatile lifetime to reset the key identifier to 0. */
+        psa_set_key_lifetime( &target_attributes, PSA_KEY_LIFETIME_VOLATILE );
+    }
+
     if( target_usage_arg != -1 )
         psa_set_key_usage_flags( &target_attributes, target_usage_arg );
     if( target_alg_arg != -1 )
diff --git a/tests/suites/test_suite_psa_crypto_se_driver_hal.function b/tests/suites/test_suite_psa_crypto_se_driver_hal.function
index c9f9dbe..04aecb6 100644
--- a/tests/suites/test_suite_psa_crypto_se_driver_hal.function
+++ b/tests/suites/test_suite_psa_crypto_se_driver_hal.function
@@ -911,7 +911,6 @@
                                 key_material, sizeof( key_material ),
                                 &returned_id ) );
 
-
     if( PSA_KEY_LIFETIME_IS_VOLATILE( lifetime ) )
     {
         /* For volatile keys, check no persistent data was created */
diff --git a/tests/suites/test_suite_psa_crypto_slot_management.data b/tests/suites/test_suite_psa_crypto_slot_management.data
index 4f31a23..2533425 100644
--- a/tests/suites/test_suite_psa_crypto_slot_management.data
+++ b/tests/suites/test_suite_psa_crypto_slot_management.data
@@ -114,6 +114,9 @@
 depends_on:MBEDTLS_PSA_CRYPTO_STORAGE_C
 create_fail:PSA_KEY_LIFETIME_PERSISTENT:0:PSA_ERROR_INVALID_HANDLE
 
+Create failure: invalid key id (1) for a volatile key
+create_fail:PSA_KEY_LIFETIME_VOLATILE:1:PSA_ERROR_INVALID_ARGUMENT
+
 Create failure: invalid key id (random seed UID)
 depends_on:MBEDTLS_PSA_CRYPTO_STORAGE_C
 create_fail:PSA_KEY_LIFETIME_PERSISTENT:PSA_CRYPTO_ITS_RANDOM_SEED_UID:PSA_ERROR_INVALID_HANDLE
diff --git a/tests/suites/test_suite_psa_crypto_slot_management.function b/tests/suites/test_suite_psa_crypto_slot_management.function
index 817094b..66bf0a4 100644
--- a/tests/suites/test_suite_psa_crypto_slot_management.function
+++ b/tests/suites/test_suite_psa_crypto_slot_management.function
@@ -476,8 +476,19 @@
 
     PSA_ASSERT( psa_crypto_init( ) );
 
-    psa_set_key_id( &attributes, id );
     psa_set_key_lifetime( &attributes, lifetime );
+    if( PSA_KEY_LIFETIME_IS_VOLATILE( lifetime ) )
+    {
+        /*
+         * Not possible to set a key identifier different from 0 through
+         * PSA key attributes APIs thus accessing to the attributes
+         * directly.
+         */
+        attributes.core.id = id;
+    }
+    else
+        psa_set_key_id( &attributes, id );
+
     psa_set_key_type( &attributes, PSA_KEY_TYPE_RAW_DATA );
     TEST_EQUAL( psa_import_key( &attributes, material, sizeof( material ),
                                 &returned_id ),