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/include/psa/crypto_platform.h b/include/psa/crypto_platform.h
index aa55aea..4e2f162 100644
--- a/include/psa/crypto_platform.h
+++ b/include/psa/crypto_platform.h
@@ -44,16 +44,32 @@
 /* PSA requires several types which C99 provides in stdint.h. */
 #include <stdint.h>
 
+#if ( defined(__ARMCC_VERSION) || defined(_MSC_VER) ) && \
+    !defined(inline) && !defined(__cplusplus)
+#define inline __inline
+#endif
+
 /* Integral type representing a key handle. */
 typedef uint16_t psa_key_handle_t;
 
 #if defined(MBEDTLS_PSA_CRYPTO_KEY_ID_ENCODES_OWNER)
 
-#if defined(PSA_CRYPTO_SECURE)
 /* Building for the PSA Crypto service on a PSA platform. */
 /* A key owner is a PSA partition identifier. */
 typedef int32_t mbedtls_key_owner_id_t;
-#endif
+
+/** Compare two key owner identifiers.
+ *
+ * \param id1 First key owner identifier.
+ * \param id2 Second key owner identifier.
+ *
+ * \return Non-zero if the two key owner identifiers are equal, zero otherwise.
+ */
+static inline int mbedtls_key_owner_id_equal( mbedtls_key_owner_id_t id1,
+                                              mbedtls_key_owner_id_t id2 )
+{
+    return( id1 == id2 );
+}
 
 #endif /* MBEDTLS_PSA_CRYPTO_KEY_ID_ENCODES_OWNER */
 
diff --git a/include/psa/crypto_types.h b/include/psa/crypto_types.h
index fb61e9a..864e937 100644
--- a/include/psa/crypto_types.h
+++ b/include/psa/crypto_types.h
@@ -240,6 +240,7 @@
 
 #define MBEDTLS_SVC_KEY_ID_INIT ( (psa_key_id_t)0 )
 #define MBEDTLS_SVC_KEY_ID_GET_KEY_ID( id ) ( id )
+#define MBEDTLS_SVC_KEY_ID_GET_OWNER_ID( id ) ( 0 )
 
 /** Utility to initialize a key identifier at runtime.
  *
@@ -254,6 +255,19 @@
     return( key_id );
 }
 
+/** Compare two key identifiers.
+ *
+ * \param id1 First key identifier.
+ * \param id2 Second key identifier.
+ *
+ * \return Non-zero if the two key identifier are equal, zero otherwise.
+ */
+static inline int mbedtls_svc_key_id_equal( mbedtls_svc_key_id_t id1,
+                                            mbedtls_svc_key_id_t id2 )
+{
+    return( id1 == id2 );
+}
+
 #else /* MBEDTLS_PSA_CRYPTO_KEY_ID_ENCODES_OWNER */
 /* Implementation-specific: The Mbed Cryptography library can be built as
  * part of a multi-client service that exposes the PSA Cryptograpy API in each
@@ -268,6 +282,7 @@
 
 #define MBEDTLS_SVC_KEY_ID_INIT ( (mbedtls_svc_key_id_t){ 0, 0 } )
 #define MBEDTLS_SVC_KEY_ID_GET_KEY_ID( id ) ( ( id ).key_id )
+#define MBEDTLS_SVC_KEY_ID_GET_OWNER_ID( id ) ( ( id ).owner )
 
 /** Utility to initialize a key identifier at runtime.
  *
@@ -281,6 +296,20 @@
                                     .owner = owner_id } );
 }
 
+/** Compare two key identifiers.
+ *
+ * \param id1 First key identifier.
+ * \param id2 Second key identifier.
+ *
+ * \return Non-zero if the two key identifier are equal, zero otherwise.
+ */
+static inline int mbedtls_svc_key_id_equal( mbedtls_svc_key_id_t id1,
+                                            mbedtls_svc_key_id_t id2 )
+{
+    return( ( id1.key_id == id2.key_id ) &&
+            mbedtls_key_owner_id_equal( id1.owner, id2.owner ) );
+}
+
 #endif /* !MBEDTLS_PSA_CRYPTO_KEY_ID_ENCODES_OWNER */
 
 /**@}*/
diff --git a/library/psa_crypto_storage.c b/library/psa_crypto_storage.c
index b94cfd0..e48bc28 100644
--- a/library/psa_crypto_storage.c
+++ b/library/psa_crypto_storage.c
@@ -62,8 +62,7 @@
  * 0xFFFFFF52. */
 static psa_storage_uid_t psa_its_identifier_of_slot( mbedtls_svc_key_id_t key )
 {
-#if defined(MBEDTLS_PSA_CRYPTO_KEY_ID_ENCODES_OWNER) && \
-    defined(PSA_CRYPTO_SECURE)
+#if defined(MBEDTLS_PSA_CRYPTO_KEY_ID_ENCODES_OWNER)
     /* Encode the owner in the upper 32 bits. This means that if
      * owner values are nonzero (as they are on a PSA platform),
      * no key file will ever have a value less than 0x100000000, so
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 ),