New function mbedtls_psa_register_se_key

Register an existing key in a secure element.

Minimal implementation that doesn't call any driver method and just
lets the application declare whatever it wants.
diff --git a/include/psa/crypto_extra.h b/include/psa/crypto_extra.h
index 130ce75..3550122 100644
--- a/include/psa/crypto_extra.h
+++ b/include/psa/crypto_extra.h
@@ -175,6 +175,44 @@
     attributes->core.flags &= ~MBEDTLS_PSA_KA_FLAG_HAS_SLOT_NUMBER;
 }
 
+/** Register a key that is already present in a secure element.
+ *
+ * The key must be located in a secure element designated by the
+ * lifetime field in \p attributes, in the slot set with
+ * psa_set_key_slot_number() in the attribute structure.
+ * This function makes the key available through the key identifier
+ * specified in \p attributes.
+ *
+ * \param[in] attributes        The attributes of the existing key.
+ *
+ * \retval #PSA_SUCCESS
+ *         The key was successfully registered.
+ *         Note that depending on the design of the driver, this may or may
+ *         not guarantee that a key actually exists in the designated slot
+ *         and is compatible with the specified attributes.
+ * \retval #PSA_ERROR_ALREADY_EXISTS
+ *         There is already a key with the identifier specified in
+ *         \p attributes.
+ * \retval #PSA_ERROR_INVALID_ARGUMENT
+ *         \p attributes specifies a lifetime which is not located
+ *         in a secure element.
+ * \retval #PSA_ERROR_INVALID_ARGUMENT
+ *         No slot number is specified in \p attributes,
+ *         or the specified slot number is not valid.
+ * \retval #PSA_ERROR_NOT_PERMITTED
+ *         The caller is not authorized to register the specified key slot.
+ * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
+ * \retval #PSA_ERROR_COMMUNICATION_FAILURE
+ * \retval #PSA_ERROR_HARDWARE_FAILURE
+ * \retval #PSA_ERROR_CORRUPTION_DETECTED
+ * \retval #PSA_ERROR_BAD_STATE
+ *         The library has not been previously initialized by psa_crypto_init().
+ *         It is implementation-dependent whether a failure to initialize
+ *         results in this error code.
+ */
+psa_status_t mbedtls_psa_register_se_key(
+    const psa_key_attributes_t *attributes);
+
 #endif /* MBEDTLS_PSA_CRYPTO_SE_C */
 
 /**@}*/
diff --git a/include/psa/crypto_se_driver.h b/include/psa/crypto_se_driver.h
index cdf0de1..1b0b3b2 100644
--- a/include/psa/crypto_se_driver.h
+++ b/include/psa/crypto_se_driver.h
@@ -818,6 +818,7 @@
     PSA_KEY_CREATION_GENERATE, /**< During psa_generate_key() */
     PSA_KEY_CREATION_DERIVE, /**< During psa_key_derivation_output_key() */
     PSA_KEY_CREATION_COPY, /**< During psa_copy_key() */
+    PSA_KEY_CREATION_REGISTER, /*TEMPORARY*/
 } psa_key_creation_method_t;
 
 /** \brief A function that allocates a slot for a key.
diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 08f9476..086ba82 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -1570,7 +1570,8 @@
     slot->attr.flags &= ~MBEDTLS_PSA_KA_MASK_EXTERNAL_ONLY;
 
 #if defined(MBEDTLS_PSA_CRYPTO_SE_C)
-    /* For a key in a secure element, we need to do three things:
+    /* For a key in a secure element, we need to do three things
+     * when creating a key (but not when registering an existing 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
@@ -1583,7 +1584,7 @@
      * secure element driver updates its persistent state, but we do not yet
      * save the driver's persistent state, so that if the power fails,
      * we can roll back to a state where the key doesn't exist. */
-    if( *p_drv != NULL )
+    if( *p_drv != NULL && method != PSA_KEY_CREATION_REGISTER )
     {
         status = psa_find_se_slot_for_key( attributes, method, *p_drv,
                                            &slot->data.se.slot_number );
@@ -1677,7 +1678,13 @@
 #endif /* defined(MBEDTLS_PSA_CRYPTO_STORAGE_C) */
 
 #if defined(MBEDTLS_PSA_CRYPTO_SE_C)
-    if( driver != NULL )
+    /* 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). */
+    if( driver != NULL &&
+        psa_crypto_transaction.unknown.type == PSA_CRYPTO_TRANSACTION_CREATE_KEY )
     {
         status = psa_save_se_persistent_data( driver );
         if( status != PSA_SUCCESS )
@@ -1720,9 +1727,12 @@
      * to internal storage), we need to destroy the key in the secure
      * element. */
 
-    /* Abort the ongoing transaction if any. We already did what it
-     * takes to undo any partial creation. All that's left is to update
-     * the transaction data itself. */
+    /* Abort the ongoing transaction if any (there may not be one if
+     * the creation process failed before starting one, or if the
+     * key creation is a registration of a key in a secure element).
+     * Earlier functions must already have done what it takes to undo any
+     * partial creation. All that's left is to update the transaction data
+     * itself. */
     (void) psa_crypto_stop_transaction( );
 #endif /* MBEDTLS_PSA_CRYPTO_SE_C */
 
@@ -1852,6 +1862,59 @@
     return( status );
 }
 
+#if defined(MBEDTLS_PSA_CRYPTO_SE_C)
+psa_status_t mbedtls_psa_register_se_key(
+    const psa_key_attributes_t *attributes )
+{
+    psa_status_t status;
+    psa_key_slot_t *slot = NULL;
+    psa_se_drv_table_entry_t *driver = NULL;
+    const psa_drv_se_t *drv;
+    psa_key_handle_t handle = 0;
+
+    /* Leaving attributes unspecified is not currently supported.
+     * It could make sense to query the key type and size from the
+     * secure element, but not all secure elements support this
+     * and the driver HAL doesn't currently support it. */
+    if( psa_get_key_type( attributes ) == PSA_KEY_TYPE_NONE )
+        return( PSA_ERROR_NOT_SUPPORTED );
+    if( psa_get_key_bits( attributes ) == 0 )
+        return( PSA_ERROR_NOT_SUPPORTED );
+
+    status = psa_start_key_creation( PSA_KEY_CREATION_REGISTER, attributes,
+                                     &handle, &slot, &driver );
+    if( status != PSA_SUCCESS )
+        goto exit;
+
+    if( driver == NULL )
+    {
+        status = PSA_ERROR_INVALID_ARGUMENT;
+        goto exit;
+    }
+    drv = psa_get_se_driver_methods( driver );
+
+    if ( psa_get_key_slot_number( attributes,
+                                  &slot->data.se.slot_number ) != PSA_SUCCESS )
+    {
+        /* The application didn't specify a slot number. This doesn't
+         * make sense when registering a slot. */
+        status = PSA_ERROR_INVALID_ARGUMENT;
+        goto exit;
+    }
+
+    status = psa_finish_key_creation( slot, driver );
+
+exit:
+    if( status != PSA_SUCCESS )
+    {
+        psa_fail_key_creation( slot, driver );
+    }
+    /* Registration doesn't keep the key in RAM. */
+    psa_close_key( handle );
+    return( status );
+}
+#endif /* MBEDTLS_PSA_CRYPTO_SE_C */
+
 static psa_status_t psa_copy_key_material( const psa_key_slot_t *source,
                                            psa_key_slot_t *target )
 {
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 e6482dd..a8dd0c7 100644
--- a/tests/suites/test_suite_psa_crypto_se_driver_hal.data
+++ b/tests/suites/test_suite_psa_crypto_se_driver_hal.data
@@ -110,3 +110,15 @@
 
 Generate key: not supported
 generate_key_not_supported:PSA_KEY_TYPE_AES:128
+
+Key registration: smoke test
+register_key_smoke_test:MIN_DRIVER_LIFETIME:PSA_SUCCESS
+
+Key registration: invalid lifetime (volatile)
+register_key_smoke_test:PSA_KEY_LIFETIME_VOLATILE:PSA_ERROR_INVALID_ARGUMENT
+
+Key registration: invalid lifetime (internal storage)
+register_key_smoke_test:PSA_KEY_LIFETIME_PERSISTENT:PSA_ERROR_INVALID_ARGUMENT
+
+Key registration: invalid lifetime (no registered driver)
+register_key_smoke_test:MIN_DRIVER_LIFETIME + 1:PSA_ERROR_INVALID_ARGUMENT
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 19b421d..2edf94f 100644
--- a/tests/suites/test_suite_psa_crypto_se_driver_hal.function
+++ b/tests/suites/test_suite_psa_crypto_se_driver_hal.function
@@ -716,3 +716,59 @@
     psa_purge_storage( );
 }
 /* END_CASE */
+
+/* BEGIN_CASE */
+void register_key_smoke_test( int lifetime_arg, int expected_status_arg )
+{
+    psa_key_lifetime_t lifetime = lifetime_arg;
+    psa_status_t expected_status = expected_status_arg;
+    psa_drv_se_t driver;
+    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
+    psa_key_id_t id = 1;
+    size_t bit_size = 48;
+    psa_key_slot_number_t wanted_slot = 0x123456789;
+    psa_key_handle_t handle = 0;
+    psa_status_t status;
+
+    memset( &driver, 0, sizeof( driver ) );
+    driver.hal_version = PSA_DRV_SE_HAL_VERSION;
+
+    PSA_ASSERT( psa_register_se_driver( MIN_DRIVER_LIFETIME, &driver ) );
+    PSA_ASSERT( psa_crypto_init( ) );
+
+    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_set_key_bits( &attributes, bit_size );
+    psa_set_key_slot_number( &attributes, wanted_slot );
+
+    status = mbedtls_psa_register_se_key( &attributes );
+    TEST_EQUAL( status, expected_status );
+
+    if( status != PSA_SUCCESS )
+        goto exit;
+
+    /* Test that the key exists and has the expected attributes. */
+    PSA_ASSERT( psa_open_key( id, &handle ) );
+    if( ! check_key_attributes( handle, &attributes ) )
+        goto exit;
+    PSA_ASSERT( psa_close_key( handle ) );
+
+    /* Restart and try again. */
+    PSA_DONE( );
+    PSA_ASSERT( psa_register_se_driver( lifetime, &driver ) );
+    PSA_ASSERT( psa_crypto_init( ) );
+    PSA_ASSERT( psa_open_key( id, &handle ) );
+    if( ! check_key_attributes( handle, &attributes ) )
+        goto exit;
+    /* This time, destroy the key. */
+    PSA_ASSERT( psa_destroy_key( handle ) );
+
+exit:
+    psa_reset_key_attributes( &attributes );
+    psa_destroy_key( handle );
+    PSA_DONE( );
+    psa_purge_storage( );
+}
+/* END_CASE */