Merge pull request #3438 from ronald-cron-arm/programs-use-common-test-code

Add support to build and link test common code in programs 
diff --git a/ChangeLog.d/do_not_persist_volatile_external_keys.txt b/ChangeLog.d/do_not_persist_volatile_external_keys.txt
new file mode 100644
index 0000000..b27292c
--- /dev/null
+++ b/ChangeLog.d/do_not_persist_volatile_external_keys.txt
@@ -0,0 +1,4 @@
+Default behavior changes
+   * Stop storing persistent information about externally stored keys created
+     through PSA Crypto with a volatile lifetime. Reported in #3288 and
+     contributed by Steven Cooreman in #3382.
diff --git a/include/psa/crypto_values.h b/include/psa/crypto_values.h
index 9fed276..f33946a 100644
--- a/include/psa/crypto_values.h
+++ b/include/psa/crypto_values.h
@@ -1611,7 +1611,7 @@
  */
 #define PSA_KEY_LIFETIME_IS_VOLATILE(lifetime)  \
     (PSA_KEY_LIFETIME_GET_PERSISTENCE(lifetime) == \
-     PSA_KEY_LIFETIME_PERSISTENCE_VOLATILE)
+     PSA_KEY_PERSISTENCE_VOLATILE)
 
 /** Construct a lifetime from a persistence level and a location.
  *
diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 3dc3b86..8cd8079 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -1499,16 +1499,17 @@
     const psa_key_attributes_t *attributes,
     psa_se_drv_table_entry_t **p_drv )
 {
-    psa_status_t status;
+    psa_status_t status = PSA_ERROR_INVALID_ARGUMENT;
 
-    if( attributes->core.lifetime != PSA_KEY_LIFETIME_VOLATILE )
-    {
-        status = psa_validate_persistent_key_parameters(
-            attributes->core.lifetime, attributes->core.id,
-            p_drv, 1 );
-        if( status != PSA_SUCCESS )
-            return( status );
-    }
+    status = psa_validate_key_location( psa_get_key_lifetime( attributes ),
+                                        p_drv );
+    if( status != PSA_SUCCESS )
+        return( status );
+
+    status = psa_validate_key_persistence( psa_get_key_lifetime( attributes ),
+                                           psa_get_key_id( attributes ) );
+    if( status != PSA_SUCCESS )
+        return( status );
 
     status = psa_validate_key_policy( &attributes->core.policy );
     if( status != PSA_SUCCESS )
@@ -1594,11 +1595,14 @@
 
 #if defined(MBEDTLS_PSA_CRYPTO_SE_C)
     /* For a key in a secure element, we need to do three things
-     * when creating or registering a key:
+     * when creating or registering a persistent key:
      * create the key file in internal storage, create the
      * key inside the secure element, and update the driver's
-     * persistent data. Start a transaction that will encompass these
-     * three actions. */
+     * persistent data. This is done by starting a transaction that will
+     * encompass these three actions.
+     * For registering a volatile key, we just need to find an appropriate
+     * slot number inside the SE. Since the key is designated volatile, creating
+     * a transaction is not required. */
     /* The first thing to do is to find a slot number for the new key.
      * We save the slot number in persistent storage as part of the
      * transaction data. It will be needed to recover if the power
@@ -1613,15 +1617,19 @@
                                            &slot->data.se.slot_number );
         if( status != PSA_SUCCESS )
             return( status );
-        psa_crypto_prepare_transaction( PSA_CRYPTO_TRANSACTION_CREATE_KEY );
-        psa_crypto_transaction.key.lifetime = slot->attr.lifetime;
-        psa_crypto_transaction.key.slot = slot->data.se.slot_number;
-        psa_crypto_transaction.key.id = slot->attr.id;
-        status = psa_crypto_save_transaction( );
-        if( status != PSA_SUCCESS )
+
+        if( ! PSA_KEY_LIFETIME_IS_VOLATILE( attributes->core.lifetime ) )
         {
-            (void) psa_crypto_stop_transaction( );
-            return( status );
+            psa_crypto_prepare_transaction( PSA_CRYPTO_TRANSACTION_CREATE_KEY );
+            psa_crypto_transaction.key.lifetime = slot->attr.lifetime;
+            psa_crypto_transaction.key.slot = slot->data.se.slot_number;
+            psa_crypto_transaction.key.id = slot->attr.id;
+            status = psa_crypto_save_transaction( );
+            if( status != PSA_SUCCESS )
+            {
+                (void) psa_crypto_stop_transaction( );
+                return( status );
+            }
         }
     }
 
@@ -1661,7 +1669,7 @@
     (void) driver;
 
 #if defined(MBEDTLS_PSA_CRYPTO_STORAGE_C)
-    if( slot->attr.lifetime != PSA_KEY_LIFETIME_VOLATILE )
+    if( ! PSA_KEY_LIFETIME_IS_VOLATILE( slot->attr.lifetime ) )
     {
 #if defined(MBEDTLS_PSA_CRYPTO_SE_C)
         if( driver != NULL )
@@ -1709,8 +1717,8 @@
     /* Finish the transaction for a key creation. This does not
      * happen when registering an existing key. Detect this case
      * by checking whether a transaction is in progress (actual
-     * creation of a key in a secure element requires a transaction,
-     * but registration doesn't use one). */
+     * creation of a persistent key in a secure element requires a transaction,
+     * but registration or volatile key creation doesn't use one). */
     if( driver != NULL &&
         psa_crypto_transaction.unknown.type == PSA_CRYPTO_TRANSACTION_CREATE_KEY )
     {
diff --git a/library/psa_crypto_slot_management.c b/library/psa_crypto_slot_management.c
index 8ffb5a0..801caf0 100644
--- a/library/psa_crypto_slot_management.c
+++ b/library/psa_crypto_slot_management.c
@@ -184,36 +184,53 @@
 }
 #endif /* defined(MBEDTLS_PSA_CRYPTO_STORAGE_C) */
 
-psa_status_t psa_validate_persistent_key_parameters(
-    psa_key_lifetime_t lifetime,
-    psa_key_file_id_t id,
-    psa_se_drv_table_entry_t **p_drv,
-    int creating )
+psa_status_t psa_validate_key_location( psa_key_lifetime_t lifetime,
+                                        psa_se_drv_table_entry_t **p_drv )
 {
-    if( p_drv != NULL )
-        *p_drv = NULL;
-#if defined(MBEDTLS_PSA_CRYPTO_SE_C)
-    if( psa_key_lifetime_is_external( lifetime ) )
+    if ( psa_key_lifetime_is_external( lifetime ) )
     {
-        *p_drv = psa_get_se_driver_entry( lifetime );
-        if( *p_drv == NULL )
+#if defined(MBEDTLS_PSA_CRYPTO_SE_C)
+        psa_se_drv_table_entry_t *driver = psa_get_se_driver_entry( lifetime );
+        if( driver == NULL )
             return( PSA_ERROR_INVALID_ARGUMENT );
+        else
+        {
+            if (p_drv != NULL)
+                *p_drv = driver;
+            return( PSA_SUCCESS );
+        }
+#else
+        (void) p_drv;
+        return( PSA_ERROR_INVALID_ARGUMENT );
+#endif /* MBEDTLS_PSA_CRYPTO_SE_C */
     }
     else
-#endif /* MBEDTLS_PSA_CRYPTO_SE_C */
-    if( lifetime != PSA_KEY_LIFETIME_PERSISTENT )
-        return( PSA_ERROR_INVALID_ARGUMENT );
+        /* Local/internal keys are always valid */
+        return( PSA_SUCCESS );
+}
 
+psa_status_t psa_validate_key_persistence( psa_key_lifetime_t lifetime,
+                                           psa_key_id_t key_id )
+{
+    if ( PSA_KEY_LIFETIME_IS_VOLATILE( lifetime ) )
+    {
+        /* Volatile keys are always supported */
+        return( PSA_SUCCESS );
+    }
+    else
+    {
+        /* Persistent keys require storage support */
 #if defined(MBEDTLS_PSA_CRYPTO_STORAGE_C)
-    if( ! psa_is_key_id_valid( id, ! creating ) )
-        return( PSA_ERROR_INVALID_ARGUMENT );
-    return( PSA_SUCCESS );
-
+        if( psa_is_key_id_valid( key_id,
+                                 psa_key_lifetime_is_external( lifetime ) ) )
+            return( PSA_SUCCESS );
+        else
+            return( PSA_ERROR_INVALID_ARGUMENT );
 #else /* MBEDTLS_PSA_CRYPTO_STORAGE_C */
-    (void) id;
-    (void) creating;
-    return( PSA_ERROR_NOT_SUPPORTED );
+        (void) key_id;
+        return( PSA_ERROR_NOT_SUPPORTED );
 #endif /* !MBEDTLS_PSA_CRYPTO_STORAGE_C */
+    }
 }
 
 psa_status_t psa_open_key( psa_key_file_id_t id, psa_key_handle_t *handle )
@@ -224,10 +241,8 @@
 
     *handle = 0;
 
-    status = psa_validate_persistent_key_parameters(
-        PSA_KEY_LIFETIME_PERSISTENT, id, NULL, 0 );
-    if( status != PSA_SUCCESS )
-        return( status );
+    if( ! psa_is_key_id_valid( id, 1 ) )
+        return( PSA_ERROR_INVALID_ARGUMENT );
 
     status = psa_get_empty_key_slot( handle, &slot );
     if( status != PSA_SUCCESS )
diff --git a/library/psa_crypto_slot_management.h b/library/psa_crypto_slot_management.h
index 6cb02f5..5bf0c0e 100644
--- a/library/psa_crypto_slot_management.h
+++ b/library/psa_crypto_slot_management.h
@@ -89,42 +89,40 @@
  */
 static inline int psa_key_lifetime_is_external( psa_key_lifetime_t lifetime )
 {
-    return( lifetime != PSA_KEY_LIFETIME_VOLATILE &&
-            lifetime != PSA_KEY_LIFETIME_PERSISTENT );
+    return( PSA_KEY_LIFETIME_GET_LOCATION( lifetime )
+                != PSA_KEY_LOCATION_LOCAL_STORAGE );
 }
 
-/** Test whether the given parameters are acceptable for a persistent key.
+/** Validate a key's location.
  *
- * This function does not access the storage in any way. It only tests
- * whether the parameters are meaningful and permitted by general policy.
- * It does not test whether the a file by the given id exists or could be
- * created.
+ * This function checks whether the key's attributes point to a location that
+ * is known to the PSA Core, and returns the driver function table if the key
+ * is to be found in an external location.
  *
- * If the key is in external storage, this function returns the corresponding
- * driver.
+ * \param[in] lifetime      The key lifetime attribute.
+ * \param[out] p_drv        On success, when a key is located in external
+ *                          storage, returns a pointer to the driver table
+ *                          associated with the key's storage location.
  *
- * \param lifetime      The lifetime to test.
- * \param id            The key id to test.
- * \param[out] p_drv    On output, if \p lifetime designates a key
- *                      in an external processor, \c *p_drv is a pointer
- *                      to the driver table entry fot this lifetime.
- *                      If \p lifetime designates a transparent key,
- *                      \c *p_drv is \c NULL.
- * \param creating      0 if attempting to open an existing key.
- *                      Nonzero if attempting to create a key.
- *
- * \retval PSA_SUCCESS
- *         The given parameters are valid.
- * \retval PSA_ERROR_INVALID_ARGUMENT
- *         \p lifetime is volatile or is invalid.
- * \retval PSA_ERROR_INVALID_ARGUMENT
- *         \p id is invalid.
+ * \retval #PSA_SUCCESS
+ * \retval #PSA_ERROR_INVALID_ARGUMENT
  */
-psa_status_t psa_validate_persistent_key_parameters(
-    psa_key_lifetime_t lifetime,
-    psa_key_file_id_t id,
-    psa_se_drv_table_entry_t **p_drv,
-    int creating );
+psa_status_t psa_validate_key_location( psa_key_lifetime_t lifetime,
+                                        psa_se_drv_table_entry_t **p_drv );
+
+/** Validate that a key's persistence attributes are valid.
+ *
+ * This function checks whether a key's declared persistence level and key ID
+ * attributes are valid and known to the PSA Core in its actual configuration.
+ *
+ * \param[in] lifetime    The key lifetime attribute.
+ * \param[in] key_id      The key ID attribute
+ *
+ * \retval #PSA_SUCCESS
+ * \retval #PSA_ERROR_INVALID_ARGUMENT
+ */
+psa_status_t psa_validate_key_persistence( psa_key_lifetime_t lifetime,
+                                           psa_key_id_t key_id );
 
 
 #endif /* PSA_CRYPTO_SLOT_MANAGEMENT_H */
diff --git a/tests/scripts/docker_env.sh b/tests/scripts/docker_env.sh
index 8bdc425..582a17d 100755
--- a/tests/scripts/docker_env.sh
+++ b/tests/scripts/docker_env.sh
@@ -60,12 +60,19 @@
     DOCKER="sudo docker"
 fi
 
+# Figure out the number of processors available
+if [ "$(uname)" == "Darwin" ]; then
+    NUM_PROC="$(sysctl -n hw.logicalcpu)"
+else
+    NUM_PROC="$(nproc)"
+fi
+
 # Build the Docker image
 echo "Getting docker image up to date (this may take a few minutes)..."
 ${DOCKER} image build \
     -t ${DOCKER_IMAGE_TAG} \
     --cache-from=${DOCKER_IMAGE_TAG} \
-    --build-arg MAKEFLAGS_PARALLEL="-j $(nproc)" \
+    --build-arg MAKEFLAGS_PARALLEL="-j ${NUM_PROC}" \
     --network host \
     ${http_proxy+--build-arg http_proxy=${http_proxy}} \
     ${https_proxy+--build-arg https_proxy=${https_proxy}} \
diff --git a/tests/suites/test_suite_psa_crypto_se_driver_hal.data b/tests/suites/test_suite_psa_crypto_se_driver_hal.data
index 55c3426..023024d 100644
--- a/tests/suites/test_suite_psa_crypto_se_driver_hal.data
+++ b/tests/suites/test_suite_psa_crypto_se_driver_hal.data
@@ -24,17 +24,29 @@
 Register SE driver: maximum number of drivers
 register_max:
 
-SE key import-export (p_allocate allows all slots)
-key_creation_import_export:0:0
+SE key import-export persistent (p_allocate allows all slots)
+key_creation_import_export:TEST_SE_PERSISTENT_LIFETIME:0:0
 
-SE key import-export (p_allocate allows 1 slot)
-key_creation_import_export:ARRAY_LENGTH( ram_slots ) - 1:0
+SE key import-export persistent (p_allocate allows 1 slot)
+key_creation_import_export:TEST_SE_PERSISTENT_LIFETIME:ARRAY_LENGTH( ram_slots ) - 1:0
 
-SE key import-export, check after restart (slot 0)
-key_creation_import_export:0:1
+SE key import-export persistent, check after restart (slot 0)
+key_creation_import_export:TEST_SE_PERSISTENT_LIFETIME:0:1
 
-SE key import-export, check after restart (slot 3)
-key_creation_import_export:3:1
+SE key import-export persistent, check after restart (slot 3)
+key_creation_import_export:TEST_SE_PERSISTENT_LIFETIME:3:1
+
+SE key import-export volatile (p_allocate allows all slots)
+key_creation_import_export:TEST_SE_VOLATILE_LIFETIME:0:0
+
+SE key import-export volatile (p_allocate allows 1 slot)
+key_creation_import_export:TEST_SE_VOLATILE_LIFETIME:ARRAY_LENGTH( ram_slots ) - 1:0
+
+SE key import-export volatile, check after restart (slot 0)
+key_creation_import_export:TEST_SE_VOLATILE_LIFETIME:0:1
+
+SE key import-export volatile, check after restart (slot 3)
+key_creation_import_export:TEST_SE_VOLATILE_LIFETIME:3:1
 
 Key creation in a specific slot (0)
 key_creation_in_chosen_slot:0:0:PSA_SUCCESS
@@ -118,22 +130,28 @@
 generate_key_smoke:PSA_KEY_TYPE_HMAC:256:PSA_ALG_HMAC( PSA_ALG_SHA_256 )
 
 Key registration: smoke test
-register_key_smoke_test:TEST_SE_PERSISTENT_LIFETIME:1:PSA_SUCCESS
+register_key_smoke_test:TEST_SE_PERSISTENT_LIFETIME:1:1:PSA_SUCCESS
 
-Key registration: invalid lifetime (volatile)
-register_key_smoke_test:PSA_KEY_LIFETIME_VOLATILE:1:PSA_ERROR_INVALID_ARGUMENT
+Key registration: invalid lifetime (volatile internal storage)
+register_key_smoke_test:PSA_KEY_LIFETIME_VOLATILE:1:1:PSA_ERROR_INVALID_ARGUMENT
 
 Key registration: invalid lifetime (internal storage)
-register_key_smoke_test:PSA_KEY_LIFETIME_PERSISTENT:1:PSA_ERROR_INVALID_ARGUMENT
+register_key_smoke_test:PSA_KEY_LIFETIME_PERSISTENT:1:1:PSA_ERROR_INVALID_ARGUMENT
 
 Key registration: invalid lifetime (no registered driver)
-register_key_smoke_test:PSA_KEY_LIFETIME_FROM_PERSISTENCE_AND_LOCATION( PSA_KEY_PERSISTENCE_DEFAULT, TEST_DRIVER_LOCATION + 1 ):1:PSA_ERROR_INVALID_ARGUMENT
+register_key_smoke_test:PSA_KEY_LIFETIME_FROM_PERSISTENCE_AND_LOCATION( PSA_KEY_PERSISTENCE_DEFAULT, TEST_DRIVER_LOCATION + 1 ):1:1:PSA_ERROR_INVALID_ARGUMENT
 
 Key registration: rejected
-register_key_smoke_test:TEST_SE_PERSISTENT_LIFETIME:0:PSA_ERROR_NOT_PERMITTED
+register_key_smoke_test:TEST_SE_PERSISTENT_LIFETIME:1:0:PSA_ERROR_NOT_PERMITTED
 
 Key registration: not supported
-register_key_smoke_test:TEST_SE_PERSISTENT_LIFETIME:-1:PSA_ERROR_NOT_SUPPORTED
+register_key_smoke_test:TEST_SE_PERSISTENT_LIFETIME:1:-1:PSA_ERROR_NOT_SUPPORTED
+
+Key registration: key id out of range
+register_key_smoke_test:TEST_SE_PERSISTENT_LIFETIME:PSA_KEY_ID_VENDOR_MAX+1:-1:PSA_ERROR_INVALID_ARGUMENT
+
+Key registration: key id in vendor range
+register_key_smoke_test:TEST_SE_PERSISTENT_LIFETIME:PSA_KEY_ID_VENDOR_MAX:1:PSA_SUCCESS
 
 Import-sign-verify: sign in driver, ECDSA
 depends_on:MBEDTLS_ECDSA_C:MBEDTLS_ECP_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED
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 9f44b88..7f9b4c2 100644
--- a/tests/suites/test_suite_psa_crypto_se_driver_hal.function
+++ b/tests/suites/test_suite_psa_crypto_se_driver_hal.function
@@ -27,6 +27,10 @@
     ( PSA_KEY_LIFETIME_FROM_PERSISTENCE_AND_LOCATION(           \
         PSA_KEY_PERSISTENCE_DEFAULT, TEST_DRIVER_LOCATION ) )
 
+#define TEST_SE_VOLATILE_LIFETIME                               \
+    ( PSA_KEY_LIFETIME_FROM_PERSISTENCE_AND_LOCATION(           \
+        PSA_KEY_PERSISTENCE_VOLATILE, TEST_DRIVER_LOCATION ) )
+
 /** 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 ))
@@ -609,6 +613,20 @@
     return( ok );
 }
 
+/* Check that no persistent data exists for the given location. */
+static int check_no_persistent_data( psa_key_location_t location )
+{
+    psa_storage_uid_t uid = file_uid_for_location( location );
+    struct psa_storage_info_t info;
+    int ok = 0;
+
+    TEST_EQUAL( psa_its_get_info( uid, &info ), PSA_ERROR_DOES_NOT_EXIST );
+    ok = 1;
+
+exit:
+    return( ok );
+}
+
 /* Check that a function's return status is "smoke-free", i.e. that
  * it's an acceptable error code when calling an API function that operates
  * on a key with potentially bogus parameters. */
@@ -829,11 +847,11 @@
 /* END_CASE */
 
 /* BEGIN_CASE */
-void key_creation_import_export( int min_slot, int restart )
+void key_creation_import_export( int lifetime_arg, int min_slot, int restart )
 {
     psa_drv_se_t driver;
     psa_drv_se_key_management_t key_management;
-    psa_key_lifetime_t lifetime = TEST_SE_PERSISTENT_LIFETIME;
+    psa_key_lifetime_t lifetime = (psa_key_lifetime_t) lifetime_arg;
     psa_key_location_t location = PSA_KEY_LIFETIME_GET_LOCATION( lifetime );
     psa_key_id_t id = 1;
     psa_key_handle_t handle = 0;
@@ -864,10 +882,25 @@
     PSA_ASSERT( psa_import_key( &attributes,
                                 key_material, sizeof( key_material ),
                                 &handle ) );
-    if( ! check_persistent_data( location,
-                                 &ram_shadow_slot_usage,
-                                 sizeof( ram_shadow_slot_usage ) ) )
-        goto exit;
+
+
+    if( PSA_KEY_LIFETIME_IS_VOLATILE( lifetime ) )
+    {
+        /* For volatile keys, check no persistent data was created */
+        if( ! check_no_persistent_data( location ) )
+            goto exit;
+    }
+    else
+    {
+        /* For persistent keys, check persistent data */
+        if( ! check_persistent_data( location,
+                             &ram_shadow_slot_usage,
+                             sizeof( ram_shadow_slot_usage ) ) )
+            goto exit;
+    }
+
+    /* Test that the key was created in the expected slot. */
+    TEST_EQUAL( ram_slots[min_slot].type, PSA_KEY_TYPE_RAW_DATA );
 
     /* Maybe restart, to check that the information is saved correctly. */
     if( restart )
@@ -875,15 +908,37 @@
         mbedtls_psa_crypto_free( );
         PSA_ASSERT( psa_register_se_driver( location, &driver ) );
         PSA_ASSERT( psa_crypto_init( ) );
-        if( ! check_persistent_data( location,
-                                     &ram_shadow_slot_usage,
-                                     sizeof( ram_shadow_slot_usage ) ) )
-            goto exit;
-        PSA_ASSERT( psa_open_key( id, &handle ) );
+
+        if( PSA_KEY_LIFETIME_IS_VOLATILE( lifetime ) )
+        {
+            /* Check that the PSA core has no knowledge of the volatile key */
+            TEST_ASSERT( psa_open_key( id, &handle ) == PSA_ERROR_DOES_NOT_EXIST );
+
+            /* Drop data from our mockup driver */
+            ram_slots_reset();
+            ram_min_slot = min_slot;
+
+            /* Re-import key */
+            PSA_ASSERT( psa_import_key( &attributes,
+                                key_material, sizeof( key_material ),
+                                &handle ) );
+        }
+        else
+        {
+
+            /* Check we can re-open the persistent key */
+            if( ! check_persistent_data( location,
+                                         &ram_shadow_slot_usage,
+                                         sizeof( ram_shadow_slot_usage ) ) )
+                goto exit;
+
+            /* Check that the PSA core still knows about the key */
+            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 );
+    TEST_EQUAL( ram_slots[min_slot].type, PSA_KEY_TYPE_RAW_DATA );
 
     /* Test the key attributes, including the reported slot number. */
     psa_set_key_bits( &attributes,
@@ -909,7 +964,7 @@
                 PSA_ERROR_DOES_NOT_EXIST );
 
     /* Test that the key has been erased from the designated slot. */
-    TEST_ASSERT( ram_slots[min_slot].type == 0 );
+    TEST_EQUAL( ram_slots[min_slot].type, 0 );
 
 exit:
     PSA_DONE( );
@@ -1263,7 +1318,7 @@
          * generate material, store the desired result of generation in
          * the mock secure element storage. */
         PSA_ASSERT( psa_get_key_attributes( drv_handle, &drv_attributes ) );
-        TEST_ASSERT( key_material->len == PSA_BITS_TO_BYTES( bits ) );
+        TEST_EQUAL( key_material->len, PSA_BITS_TO_BYTES( bits ) );
         memcpy( ram_slots[ram_min_slot].content, key_material->x,
                 key_material->len );
     }
@@ -1355,6 +1410,7 @@
 
 /* BEGIN_CASE */
 void register_key_smoke_test( int lifetime_arg,
+                              int id_arg,
                               int validate,
                               int expected_status_arg )
 {
@@ -1364,7 +1420,7 @@
     psa_drv_se_t driver;
     psa_drv_se_key_management_t key_management;
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
-    psa_key_id_t id = 1;
+    psa_key_id_t id = id_arg;
     size_t bit_size = 48;
     psa_key_slot_number_t wanted_slot = 0x123456789;
     psa_key_handle_t handle = 0;