Allow drivers to choose key slots

Add a function that allows the driver to allocate a key slot, and a
callback function implemented in the core that allows the driver to
search the table of occupied slots.
diff --git a/include/psa/crypto_se_driver.h b/include/psa/crypto_se_driver.h
index 76b7c48..f1bce73 100644
--- a/include/psa/crypto_se_driver.h
+++ b/include/psa/crypto_se_driver.h
@@ -677,6 +677,49 @@
 } psa_drv_se_aead_t;
 /**@}*/
 
+/** \defgroup se_slot_usage Secure element slot usage mask
+ */
+/**@{*/
+
+/** The type of slot usage data for a driver.
+ *
+ * A driver can use this data to keep track of which slot numbers are in use
+ * and which are available for new keys. The core stores this data to
+ * internal persistent storage.
+ *
+ * The representation of this type is opaque. To access it,
+ * the driver can call psa_drv_cb_find_free_slot().
+ */
+typedef struct psa_drv_se_slot_usage_s psa_drv_se_slot_usage_t;
+
+/** Callback function to find a free slot within a range.
+ *
+ * The PSA Crypto core provides this function to access the opaque
+ * data type that represents the slot usage.
+ *
+ * \param[in] slot_usage    The opaque data structure containing the
+ *                          driver's slot usage table.
+ * \param from              The start of the range.
+ *                          It is included in the search.
+ * \param before            The end of the range.
+ *                          It is not included in the search.
+ * \param[out] found        On success, a free slot number such that
+ *                          `from <= found < before`.
+ *
+ * \retval #PSA_SUCCESS
+ *         Success. \c *found contains a slot number which is between
+ *         \p from and \p before-1 inclusive and which is currently free.
+ * \retval #PSA_ERROR_INSUFFICIENT_STORAGE
+ *         All slots in the indicated range are occupied.
+ */
+psa_status_t psa_drv_cb_find_free_slot(
+    const psa_drv_se_slot_usage_t *slot_usage,
+    psa_key_slot_number_t from,
+    psa_key_slot_number_t before,
+    psa_key_slot_number_t *found);
+
+/**@}*/
+
 /** \defgroup se_key_management Secure Element Key Management
  * Currently, key management is limited to importing keys in the clear,
  * destroying keys, and exporting keys in the clear.
@@ -685,6 +728,32 @@
  */
 /**@{*/
 
+/* This type is documented in crypto.h. As far as drivers are concerned,
+ * this is an opaque type. */
+typedef struct psa_key_attributes_s psa_key_attributes_t;
+
+/** \brief A function that allocates a slot number for a key.
+ *
+ * This function is typically implemented as one or more calls to
+ * psa_drv_cb_find_free_slot(), with bounds determined by the key
+ * attributes and the secure element configuration.
+ *
+ * \param[in] attributes    Attributes of the key.
+ * \param[in] slot_usage    Slot usage data of the driver.
+ * \param[out] key_slot     Slot where the key will be stored.
+ *                          This must be a valid slot for a key of the
+ *                          chosen type. It must be unoccupied.
+ *
+ * \retval #PSA_SUCCESS
+ *         Success.
+ * \retval #PSA_ERROR_NOT_SUPPORTED
+ * \retval #PSA_ERROR_INSUFFICIENT_STORAGE
+ */
+typedef psa_status_t (*psa_drv_se_allocate_key_t)(
+    const psa_key_attributes_t *attributes,
+    const psa_drv_se_slot_usage_t *slot_usage,
+    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
@@ -815,6 +884,13 @@
      * slots numbered 0 through `slot_count - 1`.
      */
     psa_key_slot_number_t slot_count;
+    /** Function that allocates a slot number.
+     *
+     * If the secure element has no constraints regarding which keys
+     * can go into which slots, this should be \c NULL. In this case
+     * the core will pick any free slot when creating a key.
+     */
+    psa_drv_se_allocate_key_t   p_allocate;
     /** 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_se.c b/library/psa_crypto_se.c
index b89b0b2..b09888c 100644
--- a/library/psa_crypto_se.c
+++ b/library/psa_crypto_se.c
@@ -48,8 +48,6 @@
 /* Driver lookup */
 /****************************************************************/
 
-typedef struct psa_drv_se_slot_usage_s psa_drv_se_slot_usage_t;
-
 typedef struct psa_se_drv_table_entry_s
 {
     psa_key_lifetime_t lifetime;
@@ -177,6 +175,8 @@
     const psa_se_drv_table_entry_t *drv,
     psa_key_slot_number_t *slot_number )
 {
+    psa_status_t status;
+
     /* The maximum possible value of the type is never a valid slot number
      * because it's too large. (0 is valid.) */
     *slot_number = -1;
@@ -188,10 +188,20 @@
     if( drv->methods->key_management == NULL )
         return( PSA_ERROR_NOT_SUPPORTED );
 
-    return( psa_drv_cb_find_free_slot(
-                drv->slot_usage,
-                0, drv->methods->key_management->slot_count,
-                slot_number ) );
+    if( drv->methods->key_management->p_allocate == NULL )
+    {
+        status = psa_drv_cb_find_free_slot(
+            drv->slot_usage,
+            0, drv->methods->key_management->slot_count,
+            slot_number );
+    }
+    else
+    {
+        status = drv->methods->key_management->p_allocate( attributes,
+                                                           drv->slot_usage,
+                                                           slot_number );
+    }
+    return( status );
 }
 
 psa_status_t psa_update_se_slot_usage(
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 f6af2c0..035d79d 100644
--- a/tests/suites/test_suite_psa_crypto_se_driver_hal.data
+++ b/tests/suites/test_suite_psa_crypto_se_driver_hal.data
@@ -27,5 +27,11 @@
 Register SE driver: maximum number of drivers
 register_max:
 
-Key creation smoke test
-key_creation_import_export:
+Key creation smoke test (no p_allocate)
+key_creation_import_export:-1
+
+Key creation smoke test (p_allocate allows all slots)
+key_creation_import_export:0
+
+Key creation smoke test (p_allocate allows 1 slot)
+key_creation_import_export:ARRAY_LENGTH( ram_slots ) - 1
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 bbafdca..33a8c4e 100644
--- a/tests/suites/test_suite_psa_crypto_se_driver_hal.function
+++ b/tests/suites/test_suite_psa_crypto_se_driver_hal.function
@@ -21,9 +21,12 @@
 } ram_slot_t;
 static ram_slot_t ram_slots[16];
 
+static uint8_t ram_min_slot = 0;
+
 static void ram_slots_reset( void )
 {
     memset( ram_slots, 0, sizeof( ram_slots ) );
+    ram_min_slot = 0;
 }
 
 static psa_status_t ram_import( psa_key_slot_number_t slot_number,
@@ -71,6 +74,16 @@
     return( PSA_SUCCESS );
 }
 
+psa_status_t ram_allocate( const psa_key_attributes_t *attributes,
+                           const psa_drv_se_slot_usage_t *slot_usage,
+                           psa_key_slot_number_t *slot_number )
+{
+    (void) attributes;
+    return( psa_drv_cb_find_free_slot( slot_usage,
+                                       ram_min_slot, -1,
+                                       slot_number ) );
+}
+
 /* END_HEADER */
 
 /* BEGIN_DEPENDENCIES
@@ -144,7 +157,7 @@
 /* END_CASE */
 
 /* BEGIN_CASE */
-void key_creation_import_export( )
+void key_creation_import_export( int min_slot )
 {
     psa_drv_se_t driver;
     psa_drv_se_key_management_t key_management;
@@ -155,7 +168,7 @@
     const uint8_t key_material[3] = {0xfa, 0xca, 0xde};
     uint8_t exported[sizeof( key_material )];
     size_t exported_length;
-    psa_key_slot_number_t expected_slot = 0;
+    psa_key_slot_number_t expected_slot = ( min_slot >= 0 ? min_slot : 0 );
 
     memset( &driver, 0, sizeof( driver ) );
     memset( &key_management, 0, sizeof( key_management ) );
@@ -165,6 +178,11 @@
     key_management.p_import = ram_import;
     key_management.p_export = ram_export;
     key_management.p_destroy = ram_destroy;
+    if( min_slot >= 0 )
+    {
+        key_management.p_allocate = ram_allocate;
+        ram_min_slot = min_slot;
+    }
 
     PSA_ASSERT( psa_register_se_driver( lifetime, &driver ) );
     PSA_ASSERT( psa_crypto_init( ) );