Fix PSA code and unit tests

Fix PSA code and unit tests for the unit tests
to pass with key identifiers encoding owner
identifiers.

The changes in PSA code just make the enablement
of key identifiers encoding owner identifiers
platform independent. Previous to this commit,
such key identifiers were used only in the case
of PSA SPM platforms.

Signed-off-by: Ronald Cron <ronald.cron@arm.com>
diff --git a/tests/suites/test_suite_psa_crypto.data b/tests/suites/test_suite_psa_crypto.data
index cd26017..723eebf 100644
--- a/tests/suites/test_suite_psa_crypto.data
+++ b/tests/suites/test_suite_psa_crypto.data
@@ -5,19 +5,19 @@
 attributes_set_get:0x6963:PSA_KEY_LIFETIME_PERSISTENT:PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT:PSA_ALG_CCM:PSA_KEY_TYPE_AES:128
 
 PSA key attributes: id only
-persistence_attributes:0x1234:-1:-1:0x1234:PSA_KEY_LIFETIME_PERSISTENT
+persistence_attributes:0x1234:0x5678:-1:-1:0:0x1234:0x5678:PSA_KEY_LIFETIME_PERSISTENT
 
 PSA key attributes: lifetime=3 only
-persistence_attributes:-1:3:-1:0:3
+persistence_attributes:-1:0:3:-1:0:0:0:3
 
 PSA key attributes: id then back to volatile
-persistence_attributes:0x1234:PSA_KEY_LIFETIME_VOLATILE:-1:0:PSA_KEY_LIFETIME_VOLATILE
+persistence_attributes:0x1234:0x5678:PSA_KEY_LIFETIME_VOLATILE:-1:0:0:0x5678:PSA_KEY_LIFETIME_VOLATILE
 
 PSA key attributes: id then lifetime
-persistence_attributes:0x1234:3:-1:0x1234:3
+persistence_attributes:0x1234:0x5678:3:-1:0:0x1234:0x5678:3
 
 PSA key attributes: lifetime then id
-persistence_attributes:0x1234:3:0x1235:0x1235:3
+persistence_attributes:0x1234:0x5678:3:0x1235:0x5679:0x1235:0x5679:3
 
 PSA key attributes: slot number
 slot_number_attribute:
diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index 3f34211..3e3a7a2 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -245,12 +245,12 @@
 
     /* Persistence */
     if( lifetime == PSA_KEY_LIFETIME_VOLATILE )
-        TEST_ASSERT( id == 0 );
+        TEST_ASSERT( MBEDTLS_SVC_KEY_ID_GET_KEY_ID( id ) == 0 );
     else
     {
         TEST_ASSERT(
-            ( PSA_KEY_ID_USER_MIN <= id && id <= PSA_KEY_ID_USER_MAX ) ||
-            ( PSA_KEY_ID_USER_MIN <= id && id <= PSA_KEY_ID_USER_MAX ) );
+            ( PSA_KEY_ID_USER_MIN <= MBEDTLS_SVC_KEY_ID_GET_KEY_ID( id ) ) &&
+            ( MBEDTLS_SVC_KEY_ID_GET_KEY_ID( id ) <= PSA_KEY_ID_USER_MAX ) );
     }
 #if defined(MBEDTLS_PSA_CRYPTO_SE_C)
     /* randomly-generated 64-bit constant, should never appear in test data */
@@ -1178,17 +1178,21 @@
 static int test_operations_on_invalid_handle( psa_key_handle_t handle )
 {
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
+    mbedtls_svc_key_id_t key_id = mbedtls_svc_key_id_make( 1, 0x6964 );
     uint8_t buffer[1];
     size_t length;
     int ok = 0;
 
-    psa_set_key_id( &attributes, 0x6964 );
+    psa_set_key_id( &attributes, key_id );
     psa_set_key_usage_flags( &attributes, PSA_KEY_USAGE_ENCRYPT );
     psa_set_key_algorithm( &attributes, PSA_ALG_CTR );
     psa_set_key_type( &attributes, PSA_KEY_TYPE_AES );
     TEST_EQUAL( psa_get_key_attributes( handle, &attributes ),
                 PSA_ERROR_INVALID_HANDLE );
-    TEST_EQUAL( psa_get_key_id( &attributes ), 0 );
+    TEST_EQUAL(
+        MBEDTLS_SVC_KEY_ID_GET_KEY_ID( psa_get_key_id( &attributes ) ), 0 );
+    TEST_EQUAL(
+        MBEDTLS_SVC_KEY_ID_GET_OWNER_ID( psa_get_key_id( &attributes ) ), 0 );
     TEST_EQUAL( psa_get_key_lifetime( &attributes ), 0 );
     TEST_EQUAL( psa_get_key_usage_flags( &attributes ), 0 );
     TEST_EQUAL( psa_get_key_algorithm( &attributes ), 0 );
@@ -1333,7 +1337,10 @@
     psa_key_type_t type = type_arg;
     size_t bits = bits_arg;
 
-    TEST_EQUAL( psa_get_key_id( &attributes ), 0 );
+    TEST_EQUAL(
+        MBEDTLS_SVC_KEY_ID_GET_KEY_ID( psa_get_key_id( &attributes ) ), 0 );
+    TEST_EQUAL(
+        MBEDTLS_SVC_KEY_ID_GET_OWNER_ID( psa_get_key_id( &attributes ) ), 0 );
     TEST_EQUAL( psa_get_key_lifetime( &attributes ), 0 );
     TEST_EQUAL( psa_get_key_usage_flags( &attributes ), 0 );
     TEST_EQUAL( psa_get_key_algorithm( &attributes ), 0 );
@@ -1347,7 +1354,8 @@
     psa_set_key_type( &attributes, type );
     psa_set_key_bits( &attributes, bits );
 
-    TEST_EQUAL( psa_get_key_id( &attributes ), id );
+    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_usage_flags( &attributes ), usage_flags );
     TEST_EQUAL( psa_get_key_algorithm( &attributes ), alg );
@@ -1356,7 +1364,10 @@
 
     psa_reset_key_attributes( &attributes );
 
-    TEST_EQUAL( psa_get_key_id( &attributes ), 0 );
+    TEST_EQUAL(
+        MBEDTLS_SVC_KEY_ID_GET_KEY_ID( psa_get_key_id( &attributes ) ), 0 );
+    TEST_EQUAL(
+        MBEDTLS_SVC_KEY_ID_GET_OWNER_ID( psa_get_key_id( &attributes ) ), 0 );
     TEST_EQUAL( psa_get_key_lifetime( &attributes ), 0 );
     TEST_EQUAL( psa_get_key_usage_flags( &attributes ), 0 );
     TEST_EQUAL( psa_get_key_algorithm( &attributes ), 0 );
@@ -1366,15 +1377,19 @@
 /* END_CASE */
 
 /* BEGIN_CASE */
-void persistence_attributes( int id1_arg, int lifetime_arg, int id2_arg,
-                             int expected_id_arg, int expected_lifetime_arg )
+void persistence_attributes( int id1_arg, int owner_id1_arg, int lifetime_arg,
+                             int id2_arg, int owner_id2_arg,
+                             int expected_id_arg, int expected_owner_id_arg,
+                             int expected_lifetime_arg )
 {
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
-    mbedtls_svc_key_id_t id1 = mbedtls_svc_key_id_make( 1, id1_arg );
+    mbedtls_svc_key_id_t id1 =
+        mbedtls_svc_key_id_make( owner_id1_arg, id1_arg );
     psa_key_lifetime_t lifetime = lifetime_arg;
-    mbedtls_svc_key_id_t id2 = mbedtls_svc_key_id_make( 1, id2_arg );
+    mbedtls_svc_key_id_t id2 =
+        mbedtls_svc_key_id_make( owner_id2_arg, id2_arg );
     mbedtls_svc_key_id_t expected_id =
-        mbedtls_svc_key_id_make( 1, expected_id_arg );
+        mbedtls_svc_key_id_make( expected_owner_id_arg, expected_id_arg );
     psa_key_lifetime_t expected_lifetime = expected_lifetime_arg;
 
     if( id1_arg != -1 )
@@ -1384,7 +1399,8 @@
     if( id2_arg != -1 )
         psa_set_key_id( &attributes, id2 );
 
-    TEST_EQUAL( psa_get_key_id( &attributes ), expected_id );
+    TEST_ASSERT( mbedtls_svc_key_id_equal(
+                     psa_get_key_id( &attributes ), expected_id ) );
     TEST_EQUAL( psa_get_key_lifetime( &attributes ), expected_lifetime );
 }
 /* END_CASE */
@@ -5677,7 +5693,8 @@
     /* Check key slot still contains key data */
     PSA_ASSERT( psa_open_key( key_id, &handle ) );
     PSA_ASSERT( psa_get_key_attributes( handle, &attributes ) );
-    TEST_EQUAL( psa_get_key_id( &attributes ), key_id );
+    TEST_ASSERT( mbedtls_svc_key_id_equal(
+                     psa_get_key_id( &attributes ), key_id ) );
     TEST_EQUAL( psa_get_key_lifetime( &attributes ),
                 PSA_KEY_LIFETIME_PERSISTENT );
     TEST_EQUAL( psa_get_key_type( &attributes ), type );
diff --git a/tests/suites/test_suite_psa_crypto_persistent_key.function b/tests/suites/test_suite_psa_crypto_persistent_key.function
index ffb4f46..5087034 100644
--- a/tests/suites/test_suite_psa_crypto_persistent_key.function
+++ b/tests/suites/test_suite_psa_crypto_persistent_key.function
@@ -224,7 +224,8 @@
 
     psa_reset_key_attributes( &attributes );
     PSA_ASSERT( psa_get_key_attributes( handle, &attributes ) );
-    TEST_EQUAL( psa_get_key_id( &attributes ), key_id );
+    TEST_ASSERT( mbedtls_svc_key_id_equal(
+                     psa_get_key_id( &attributes ), key_id ) );
     TEST_EQUAL( psa_get_key_lifetime( &attributes ),
                 PSA_KEY_LIFETIME_PERSISTENT );
     TEST_EQUAL( psa_get_key_type( &attributes ), type );
@@ -276,7 +277,8 @@
     /* Test the key information */
     psa_reset_key_attributes( &attributes );
     PSA_ASSERT( psa_get_key_attributes( handle, &attributes ) );
-    TEST_EQUAL( psa_get_key_id( &attributes ), key_id );
+    TEST_ASSERT( mbedtls_svc_key_id_equal(
+                     psa_get_key_id( &attributes ), key_id ) );
     TEST_EQUAL( psa_get_key_lifetime( &attributes ),
                 PSA_KEY_LIFETIME_PERSISTENT );
     TEST_EQUAL( psa_get_key_type( &attributes ), type );
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 d117738..6f7cfa9 100644
--- a/tests/suites/test_suite_psa_crypto_se_driver_hal.function
+++ b/tests/suites/test_suite_psa_crypto_se_driver_hal.function
@@ -532,8 +532,9 @@
 
     PSA_ASSERT( psa_get_key_attributes( handle, &actual_attributes ) );
 
-    TEST_EQUAL( psa_get_key_id( &actual_attributes ),
-                psa_get_key_id( reference_attributes ) );
+    TEST_ASSERT( mbedtls_svc_key_id_equal(
+                     psa_get_key_id( &actual_attributes ),
+                     psa_get_key_id( reference_attributes ) ) );
     TEST_EQUAL( psa_get_key_lifetime( &actual_attributes ),
                 psa_get_key_lifetime( reference_attributes ) );
     TEST_EQUAL( psa_get_key_type( &actual_attributes ),
@@ -762,6 +763,7 @@
 {
     psa_key_id_t id;
     psa_key_location_t location;
+
     /* 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. */
diff --git a/tests/suites/test_suite_psa_crypto_se_driver_hal_mocks.function b/tests/suites/test_suite_psa_crypto_se_driver_hal_mocks.function
index c4c89cd..7d4a591 100644
--- a/tests/suites/test_suite_psa_crypto_se_driver_hal_mocks.function
+++ b/tests/suites/test_suite_psa_crypto_se_driver_hal_mocks.function
@@ -91,11 +91,13 @@
 {
     psa_key_id_t id;
     psa_key_location_t location;
+
     /* 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( mbedtls_svc_key_id_make( 1, id ) );
+
     /* Purge the transaction file. */
     psa_crypto_stop_transaction( );
     /* Purge driver persistent data. */
@@ -360,8 +362,20 @@
     TEST_ASSERT( mock_allocate_data.called == 1 );
     TEST_ASSERT( mock_import_data.called ==
         ( mock_alloc_return_value == PSA_SUCCESS? 1 : 0 ) );
-    TEST_ASSERT( mock_import_data.attributes.core.id ==
-        ( mock_alloc_return_value == PSA_SUCCESS? id : 0 ) );
+
+    if( mock_alloc_return_value == PSA_SUCCESS )
+    {
+        TEST_ASSERT( mbedtls_svc_key_id_equal(
+                         mock_import_data.attributes.core.id, id ) );
+    }
+    else
+    {
+        TEST_ASSERT( MBEDTLS_SVC_KEY_ID_GET_KEY_ID(
+                         mock_import_data.attributes.core.id ) == 0 );
+        TEST_ASSERT( MBEDTLS_SVC_KEY_ID_GET_OWNER_ID(
+                         mock_import_data.attributes.core.id ) == 0 );
+    }
+
     TEST_ASSERT( mock_import_data.attributes.core.lifetime ==
         ( mock_alloc_return_value == PSA_SUCCESS? lifetime : 0 ) );
     TEST_ASSERT( mock_import_data.attributes.core.policy.usage ==
@@ -467,8 +481,20 @@
     TEST_ASSERT( mock_allocate_data.called == 1 );
     TEST_ASSERT( mock_generate_data.called ==
         ( mock_alloc_return_value == PSA_SUCCESS? 1 : 0 ) );
-    TEST_ASSERT( mock_generate_data.attributes.core.id ==
-        ( mock_alloc_return_value == PSA_SUCCESS? id : 0 ) );
+
+    if( mock_alloc_return_value == PSA_SUCCESS )
+    {
+        TEST_ASSERT( mbedtls_svc_key_id_equal(
+                         mock_generate_data.attributes.core.id, id ) );
+    }
+    else
+    {
+        TEST_ASSERT( MBEDTLS_SVC_KEY_ID_GET_KEY_ID(
+                         mock_generate_data.attributes.core.id ) == 0 );
+        TEST_ASSERT( MBEDTLS_SVC_KEY_ID_GET_OWNER_ID(
+                         mock_generate_data.attributes.core.id ) == 0 );
+    }
+
     TEST_ASSERT( mock_generate_data.attributes.core.lifetime ==
         ( mock_alloc_return_value == PSA_SUCCESS? lifetime : 0 ) );
     TEST_ASSERT( mock_generate_data.attributes.core.policy.usage ==
diff --git a/tests/suites/test_suite_psa_crypto_slot_management.function b/tests/suites/test_suite_psa_crypto_slot_management.function
index bc6ecda..c9d4404 100644
--- a/tests/suites/test_suite_psa_crypto_slot_management.function
+++ b/tests/suites/test_suite_psa_crypto_slot_management.function
@@ -41,14 +41,15 @@
 static int test_uses_key_id( mbedtls_svc_key_id_t key_id )
 {
     size_t i;
-    if( key_id > PSA_MAX_PERSISTENT_KEY_IDENTIFIER )
+    if( MBEDTLS_SVC_KEY_ID_GET_KEY_ID( key_id ) >
+        PSA_MAX_PERSISTENT_KEY_IDENTIFIER )
     {
         /* Don't touch key id values that designate non-key files. */
         return( 1 );
     }
     for( i = 0; i < num_key_ids_used ; i++ )
     {
-        if( key_id == key_ids_used_in_test[i] )
+        if( mbedtls_svc_key_id_equal( key_id, key_ids_used_in_test[i] ) )
             return( 1 );
     }
     if( num_key_ids_used == ARRAY_LENGTH( key_ids_used_in_test ) )
@@ -206,7 +207,8 @@
     TEST_ASSERT( handle != 0 );
     PSA_ASSERT( psa_get_key_attributes( handle, &attributes ) );
     TEST_EQUAL( psa_get_key_lifetime( &attributes ), lifetime );
-    TEST_EQUAL( psa_get_key_id( &attributes ), id );
+    TEST_ASSERT( mbedtls_svc_key_id_equal(
+                     psa_get_key_id( &attributes ), id ) );
     TEST_EQUAL( psa_get_key_usage_flags( &attributes ), usage_flags );
     TEST_EQUAL( psa_get_key_algorithm( &attributes ), alg );
     TEST_EQUAL( psa_get_key_enrollment_algorithm( &attributes ), alg2 );
@@ -217,7 +219,8 @@
     PSA_ASSERT( psa_open_key( id, &handle ) );
     PSA_ASSERT( psa_get_key_attributes( handle, &attributes ) );
     TEST_EQUAL( psa_get_key_lifetime( &attributes ), lifetime );
-    TEST_EQUAL( psa_get_key_id( &attributes ), id );
+    TEST_ASSERT( mbedtls_svc_key_id_equal(
+                     psa_get_key_id( &attributes ), id ) );
     TEST_EQUAL( psa_get_key_usage_flags( &attributes ), usage_flags );
     TEST_EQUAL( psa_get_key_algorithm( &attributes ), alg );
     TEST_EQUAL( psa_get_key_enrollment_algorithm( &attributes ), alg2 );
@@ -247,8 +250,9 @@
             PSA_ASSERT( psa_get_key_attributes( handle, &read_attributes ) );
             TEST_EQUAL( psa_get_key_lifetime( &attributes ),
                         psa_get_key_lifetime( &read_attributes ) );
-            TEST_EQUAL( psa_get_key_id( &attributes ),
-                        psa_get_key_id( &read_attributes ) );
+            TEST_ASSERT( mbedtls_svc_key_id_equal(
+                             psa_get_key_id( &attributes ),
+                             psa_get_key_id( &read_attributes ) ) );
             TEST_EQUAL( psa_get_key_usage_flags( &attributes ), usage_flags );
             TEST_EQUAL( psa_get_key_algorithm( &attributes ),
                         psa_get_key_algorithm( &read_attributes ) );
@@ -338,7 +342,8 @@
     /* Check that the original key hasn't changed. */
     psa_reset_key_attributes( &attributes );
     PSA_ASSERT( psa_get_key_attributes( handle1, &attributes ) );
-    TEST_EQUAL( psa_get_key_id( &attributes ), id );
+    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 );
@@ -445,11 +450,9 @@
     PSA_ASSERT( psa_crypto_init( ) );
 
     /* Populate the source slot. */
-    if( source_lifetime != PSA_KEY_LIFETIME_VOLATILE )
-    {
-        psa_set_key_id( &source_attributes, source_id );
-        psa_set_key_lifetime( &source_attributes, source_lifetime );
-    }
+    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 );
@@ -461,11 +464,9 @@
     PSA_ASSERT( psa_get_key_attributes( source_handle, &source_attributes ) );
 
     /* Prepare the target slot. */
-    if( target_lifetime != PSA_KEY_LIFETIME_VOLATILE )
-    {
-        psa_set_key_id( &target_attributes, target_id );
-        psa_set_key_lifetime( &target_attributes, target_lifetime );
-    }
+    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 );
@@ -489,7 +490,20 @@
     /* Test that the target slot has the expected content. */
     psa_reset_key_attributes( &target_attributes );
     PSA_ASSERT( psa_get_key_attributes( target_handle, &target_attributes ) );
-    TEST_EQUAL( target_id, psa_get_key_id( &target_attributes ) );
+
+    if( target_lifetime != PSA_KEY_LIFETIME_VOLATILE )
+    {
+        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( target_id ), 1 );
+#endif
+        TEST_EQUAL( MBEDTLS_SVC_KEY_ID_GET_KEY_ID( target_id ), 0 );
+    }
+
     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 ),
@@ -574,7 +588,7 @@
                                 &source_handle ) );
 
     /* Populate the target slot. */
-    if( target_id == source_id )
+    if( mbedtls_svc_key_id_equal( target_id, source_id ) )
     {
         target_handle = source_handle;
     }
@@ -601,8 +615,9 @@
 
     /* Test that the target slot is unaffected. */
     PSA_ASSERT( psa_get_key_attributes( target_handle, &attributes2 ) );
-    TEST_EQUAL( psa_get_key_id( &attributes1 ),
-                psa_get_key_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 ),