Merge pull request #202 from gilles-peskine-arm/psa-se_driver-choose_key_slot_number

Let applications create a key in a specific secure element slot
diff --git a/include/psa/crypto_se_driver.h b/include/psa/crypto_se_driver.h
index 69cdaba..9a5d97d 100644
--- a/include/psa/crypto_se_driver.h
+++ b/include/psa/crypto_se_driver.h
@@ -812,6 +812,42 @@
 
 /** \brief A function that allocates a slot for a key.
  *
+ * To create a key in a specific slot in a secure element, the core
+ * first calls this function to determine a valid slot number,
+ * then calls a function to create the key material in that slot.
+ * For example, in nominal conditions (that is, if no error occurs),
+ * the effect of a call to psa_import_key() with a lifetime that places
+ * the key in a secure element is the following:
+ * -# The core calls psa_drv_se_key_management_t::p_allocate
+ *    (or in some implementations
+ *    psa_drv_se_key_management_t::p_validate_slot_number). The driver
+ *    selects (or validates) a suitable slot number given the key attributes
+ *    and the state of the secure element.
+ * -# The core calls psa_drv_se_key_management_t::p_import to import
+ *    the key material in the selected slot.
+ *
+ * Other key creation methods lead to similar sequences. For example, the
+ * sequence for psa_generate_key() is the same except that the second step
+ * is a call to psa_drv_se_key_management_t::p_generate.
+ *
+ * In case of errors, other behaviors are possible.
+ * - If the PSA Cryptography subsystem dies after the first step,
+ *   for example because the device has lost power abruptly,
+ *   the second step may never happen, or may happen after a reset
+ *   and re-initialization. Alternatively, after a reset and
+ *   re-initialization, the core may call
+ *   psa_drv_se_key_management_t::p_destroy on the slot number that
+ *   was allocated (or validated) instead of calling a key creation function.
+ * - If an error occurs, the core may call
+ *   psa_drv_se_key_management_t::p_destroy on the slot number that
+ *   was allocated (or validated) instead of calling a key creation function.
+ *
+ * Errors and system resets also have an impact on the driver's persistent
+ * data. If a reset happens before the overall key creation process is
+ * completed (before or after the second step above), it is unspecified
+ * whether the persistent data after the reset is identical to what it
+ * was before or after the call to `p_allocate` (or `p_validate_slot_number`).
+ *
  * \param[in,out] drv_context       The driver context structure.
  * \param[in,out] persistent_data   A pointer to the persistent data
  *                                  that allows writing.
@@ -833,6 +869,42 @@
     const psa_key_attributes_t *attributes,
     psa_key_slot_number_t *key_slot);
 
+/** \brief A function that determines whether a slot number is valid
+ * for a key.
+ *
+ * To create a key in a specific slot in a secure element, the core
+ * first calls this function to validate the choice of slot number,
+ * then calls a function to create the key material in that slot.
+ * See the documentation of #psa_drv_se_allocate_key_t for more details.
+ *
+ * As of the PSA Cryptography API specification version 1.0, there is no way
+ * for applications to trigger a call to this function. However some
+ * implementations offer the capability to create or declare a key in
+ * a specific slot via implementation-specific means, generally for the
+ * sake of initial device provisioning or onboarding. Such a mechanism may
+ * be added to a future version of the PSA Cryptography API specification.
+ *
+ * \param[in,out] drv_context       The driver context structure.
+ * \param[in] attributes    Attributes of the key.
+ * \param[in] key_slot      Slot where the key is to be stored.
+ *
+ * \retval #PSA_SUCCESS
+ *         The given slot number is valid for a key with the given
+ *         attributes.
+ * \retval #PSA_ERROR_INVALID_ARGUMENT
+ *         The given slot number is not valid for a key with the
+ *         given attributes. This includes the case where the slot
+ *         number is not valid at all.
+ * \retval #PSA_ERROR_ALREADY_EXISTS
+ *         There is already a key with the specified slot number.
+ *         Drivers may choose to return this error from the key
+ *         creation function instead.
+ */
+typedef psa_status_t (*psa_drv_se_validate_slot_number_t)(
+    psa_drv_se_context_t *drv_context,
+    const psa_key_attributes_t *attributes,
+    psa_key_slot_number_t key_slot);
+
 /** \brief A function that imports a key into a secure element in binary format
  *
  * This function can support any output from psa_export_key(). Refer to the
@@ -977,8 +1049,10 @@
  * If one of the functions is not implemented, it should be set to NULL.
  */
 typedef struct {
-    /** Function that allocates a slot. */
+    /** Function that allocates a slot for a key. */
     psa_drv_se_allocate_key_t   p_allocate;
+    /** Function that checks the validity of a slot for a key. */
+    psa_drv_se_validate_slot_number_t p_validate_slot_number;
     /** Function that performs a key import operation */
     psa_drv_se_import_key_t     p_import;
     /** Function that performs a generation */
diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 5cb88de..856d862 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -1582,10 +1582,6 @@
      * we can roll back to a state where the key doesn't exist. */
     if( *p_drv != NULL )
     {
-        /* Choosing a slot number is not supported yet. */
-        if( attributes->core.flags & MBEDTLS_PSA_KA_FLAG_HAS_SLOT_NUMBER )
-            return( PSA_ERROR_NOT_SUPPORTED );
-
         status = psa_find_se_slot_for_key( attributes, *p_drv,
                                            &slot->data.se.slot_number );
         if( status != PSA_SUCCESS )
diff --git a/library/psa_crypto_se.c b/library/psa_crypto_se.c
index bc73251..ca38e20 100644
--- a/library/psa_crypto_se.c
+++ b/library/psa_crypto_se.c
@@ -201,7 +201,6 @@
     psa_key_slot_number_t *slot_number )
 {
     psa_status_t status;
-    psa_drv_se_allocate_key_t p_allocate = NULL;
 
     /* If the lifetime is wrong, it's a bug in the library. */
     if( driver->lifetime != psa_get_key_lifetime( attributes ) )
@@ -210,17 +209,33 @@
     /* If the driver doesn't support key creation in any way, give up now. */
     if( driver->methods->key_management == NULL )
         return( PSA_ERROR_NOT_SUPPORTED );
-    p_allocate = driver->methods->key_management->p_allocate;
 
-    /* If the driver doesn't tell us how to allocate a slot, that's
-     * not supported for the time being. */
-    if( p_allocate == NULL )
-        return( PSA_ERROR_NOT_SUPPORTED );
-
-    status = p_allocate( &driver->context,
-                         driver->internal.persistent_data,
-                         attributes,
-                         slot_number );
+    if( psa_get_key_slot_number( attributes, slot_number ) == PSA_SUCCESS )
+    {
+        /* The application wants to use a specific slot. Allow it if
+         * the driver supports it. On a system with isolation,
+         * the crypto service must check that the application is
+         * permitted to request this slot. */
+        psa_drv_se_validate_slot_number_t p_validate_slot_number =
+            driver->methods->key_management->p_validate_slot_number;
+        if( p_validate_slot_number == NULL )
+            return( PSA_ERROR_NOT_SUPPORTED );
+        status = p_validate_slot_number( &driver->context, attributes,
+                                         *slot_number );
+    }
+    else
+    {
+        /* The application didn't tell us which slot to use. Let the driver
+         * choose. This is the normal case. */
+        psa_drv_se_allocate_key_t p_allocate =
+            driver->methods->key_management->p_allocate;
+        if( p_allocate == NULL )
+            return( PSA_ERROR_NOT_SUPPORTED );
+        status = p_allocate( &driver->context,
+                             driver->internal.persistent_data,
+                             attributes,
+                             slot_number );
+    }
     return( status );
 }
 
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 6fb65f0..e6482dd 100644
--- a/tests/suites/test_suite_psa_crypto_se_driver_hal.data
+++ b/tests/suites/test_suite_psa_crypto_se_driver_hal.data
@@ -39,6 +39,21 @@
 SE key import-export, check after restart (slot 3)
 key_creation_import_export:3:1
 
+Key creation in a specific slot (0)
+key_creation_in_chosen_slot:0:0:PSA_SUCCESS
+
+Key creation in a specific slot (max)
+key_creation_in_chosen_slot:ARRAY_LENGTH( ram_slots ) - 1:0:PSA_SUCCESS
+
+Key creation in a specific slot (0, restart)
+key_creation_in_chosen_slot:0:1:PSA_SUCCESS
+
+Key creation in a specific slot (max, restart)
+key_creation_in_chosen_slot:ARRAY_LENGTH( ram_slots ) - 1:1:PSA_SUCCESS
+
+Key creation in a specific slot (too large)
+key_creation_in_chosen_slot:ARRAY_LENGTH( ram_slots ):0:PSA_ERROR_INVALID_ARGUMENT
+
 Key creation smoke test: AES-CTR
 key_creation_smoke:PSA_KEY_TYPE_AES:PSA_ALG_CTR:"aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"
 
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 9a57464..0fab043 100644
--- a/tests/suites/test_suite_psa_crypto_se_driver_hal.function
+++ b/tests/suites/test_suite_psa_crypto_se_driver_hal.function
@@ -177,6 +177,18 @@
     return( PSA_ERROR_INSUFFICIENT_STORAGE );
 }
 
+static psa_status_t ram_validate_slot_number(
+    psa_drv_se_context_t *context,
+    const psa_key_attributes_t *attributes,
+    psa_key_slot_number_t slot_number )
+{
+    (void) context;
+    (void) attributes;
+    if( slot_number >= ARRAY_LENGTH( ram_slots ) )
+        return( PSA_ERROR_INVALID_ARGUMENT );
+    return( PSA_SUCCESS );
+}
+
 
 
 /****************************************************************/
@@ -537,6 +549,74 @@
 /* END_CASE */
 
 /* BEGIN_CASE */
+void key_creation_in_chosen_slot( int slot_arg,
+                                  int restart,
+                                  int expected_status_arg )
+{
+    psa_key_slot_number_t wanted_slot = slot_arg;
+    psa_status_t expected_status = expected_status_arg;
+    psa_status_t status;
+    psa_drv_se_t driver;
+    psa_drv_se_key_management_t key_management;
+    psa_key_lifetime_t lifetime = 2;
+    psa_key_id_t id = 1;
+    psa_key_handle_t handle = 0;
+    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
+    const uint8_t key_material[3] = {0xfa, 0xca, 0xde};
+
+    memset( &driver, 0, sizeof( driver ) );
+    memset( &key_management, 0, sizeof( key_management ) );
+    driver.hal_version = PSA_DRV_SE_HAL_VERSION;
+    driver.key_management = &key_management;
+    driver.persistent_data_size = sizeof( ram_slot_usage_t );
+    key_management.p_validate_slot_number = ram_validate_slot_number;
+    key_management.p_import = ram_import;
+    key_management.p_destroy = ram_destroy;
+    key_management.p_export = ram_export;
+
+    PSA_ASSERT( psa_register_se_driver( lifetime, &driver ) );
+    PSA_ASSERT( psa_crypto_init( ) );
+
+    /* Create a key. */
+    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_slot_number( &attributes, wanted_slot );
+    status = psa_import_key( &attributes,
+                             key_material, sizeof( key_material ),
+                             &handle );
+    TEST_EQUAL( status, expected_status );
+
+    if( status != PSA_SUCCESS )
+        goto exit;
+
+    /* Maybe restart, to check that the information is saved correctly. */
+    if( restart )
+    {
+        mbedtls_psa_crypto_free( );
+        PSA_ASSERT( psa_register_se_driver( lifetime, &driver ) );
+        PSA_ASSERT( psa_crypto_init( ) );
+        PSA_ASSERT( psa_open_key( id, &handle ) );
+    }
+
+    /* Test that the key was created in the expected slot. */
+    TEST_EQUAL( ram_slots[wanted_slot].type, PSA_KEY_TYPE_RAW_DATA );
+
+    /* Test that the key is reported with the correct attributes,
+     * including the expected slot. */
+    PSA_ASSERT( psa_get_key_attributes( handle, &attributes ) );
+
+    PSA_ASSERT( psa_destroy_key( handle ) );
+
+exit:
+    PSA_DONE( );
+    ram_slots_reset( );
+    psa_purge_storage( );
+}
+/* END_CASE */
+
+/* BEGIN_CASE */
 void key_creation_smoke( int type_arg, int alg_arg,
                          data_t *key_material )
 {