Pass the key creation method to drivers

Pass the key creation method (import/generate/derive/copy) to the
driver methods to allocate or validate a slot number. This allows
drivers to enforce policies such as "this key slot can only be used
for keys generated inside the secure element".
diff --git a/include/psa/crypto_se_driver.h b/include/psa/crypto_se_driver.h
index 9a5d97d..cdf0de1 100644
--- a/include/psa/crypto_se_driver.h
+++ b/include/psa/crypto_se_driver.h
@@ -810,25 +810,45 @@
  */
 /**@{*/
 
+/** An enumeration indicating how a key is created.
+ */
+typedef enum
+{
+    PSA_KEY_CREATION_IMPORT, /**< During psa_import_key() */
+    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_method_t;
+
 /** \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:
+ * In nominal conditions (that is, if no error occurs),
+ * the effect of a call to a key creation function in the PSA Cryptography
+ * API 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.
+ * -# The core calls a key creation function in the driver.
  *
- * 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.
+ * The key creation functions in the PSA Cryptography API are:
+ * - psa_import_key(), which causes
+ *   a call to `p_allocate` with \p method = #PSA_KEY_CREATION_IMPORT
+ *   then a call to psa_drv_se_key_management_t::p_import.
+ * - psa_generate_key(), which causes
+ *   a call to `p_allocate` with \p method = #PSA_KEY_CREATION_GENERATE
+ *   then a call to psa_drv_se_key_management_t::p_import.
+ * - psa_key_derivation_output_key(), which causes
+ *   a call to `p_allocate` with \p method = #PSA_KEY_CREATION_DERIVE
+ *   then a call to psa_drv_se_key_derivation_t::p_derive.
+ * - psa_copy_key(), which causes
+ *   a call to `p_allocate` with \p method = #PSA_KEY_CREATION_COPY
+ *   then a call to psa_drv_se_key_management_t::p_export.
  *
  * In case of errors, other behaviors are possible.
  * - If the PSA Cryptography subsystem dies after the first step,
@@ -852,6 +872,7 @@
  * \param[in,out] persistent_data   A pointer to the persistent data
  *                                  that allows writing.
  * \param[in] attributes            Attributes of the key.
+ * \param method                    The way in which the key is being created.
  * \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.
@@ -867,6 +888,7 @@
     psa_drv_se_context_t *drv_context,
     void *persistent_data,
     const psa_key_attributes_t *attributes,
+    psa_key_creation_method_t method,
     psa_key_slot_number_t *key_slot);
 
 /** \brief A function that determines whether a slot number is valid
@@ -884,9 +906,10 @@
  * 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.
+ * \param[in,out] drv_context   The driver context structure.
+ * \param[in] attributes        Attributes of the key.
+ * \param method                The way in which the key is being created.
+ * \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
@@ -903,6 +926,7 @@
 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_creation_method_t method,
     psa_key_slot_number_t key_slot);
 
 /** \brief A function that imports a key into a secure element in binary format
diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 0c8b99b..08f9476 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -1506,16 +1506,6 @@
     return( PSA_SUCCESS );
 }
 
-/** An enumeration indicating how a key is created.
- */
-typedef enum
-{
-    PSA_KEY_CREATION_IMPORT,
-    PSA_KEY_CREATION_GENERATE,
-    PSA_KEY_CREATION_DERIVE,
-    PSA_KEY_CREATION_COPY,
-} psa_key_creation_method_t;
-
 /** Prepare a key slot to receive key material.
  *
  * This function allocates a key slot and sets its metadata.
@@ -1595,7 +1585,7 @@
      * we can roll back to a state where the key doesn't exist. */
     if( *p_drv != NULL )
     {
-        status = psa_find_se_slot_for_key( attributes, *p_drv,
+        status = psa_find_se_slot_for_key( attributes, method, *p_drv,
                                            &slot->data.se.slot_number );
         if( status != PSA_SUCCESS )
             return( status );
diff --git a/library/psa_crypto_se.c b/library/psa_crypto_se.c
index ca38e20..523c621 100644
--- a/library/psa_crypto_se.c
+++ b/library/psa_crypto_se.c
@@ -197,6 +197,7 @@
 
 psa_status_t psa_find_se_slot_for_key(
     const psa_key_attributes_t *attributes,
+    psa_key_creation_method_t method,
     psa_se_drv_table_entry_t *driver,
     psa_key_slot_number_t *slot_number )
 {
@@ -220,7 +221,8 @@
             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,
+        status = p_validate_slot_number( &driver->context,
+                                         attributes, method,
                                          *slot_number );
     }
     else
@@ -233,7 +235,7 @@
             return( PSA_ERROR_NOT_SUPPORTED );
         status = p_allocate( &driver->context,
                              driver->internal.persistent_data,
-                             attributes,
+                             attributes, method,
                              slot_number );
     }
     return( status );
diff --git a/library/psa_crypto_se.h b/library/psa_crypto_se.h
index 378c78f..900a72b 100644
--- a/library/psa_crypto_se.h
+++ b/library/psa_crypto_se.h
@@ -135,6 +135,7 @@
  */
 psa_status_t psa_find_se_slot_for_key(
     const psa_key_attributes_t *attributes,
+    psa_key_creation_method_t method,
     psa_se_drv_table_entry_t *driver,
     psa_key_slot_number_t *slot_number );
 
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 0fab043..19b421d 100644
--- a/tests/suites/test_suite_psa_crypto_se_driver_hal.function
+++ b/tests/suites/test_suite_psa_crypto_se_driver_hal.function
@@ -41,10 +41,12 @@
 static psa_status_t counter_allocate( psa_drv_se_context_t *context,
                                       void *persistent_data,
                                       const psa_key_attributes_t *attributes,
+                                      psa_key_creation_method_t method,
                                       psa_key_slot_number_t *slot_number )
 {
     psa_key_slot_number_t *p_counter = persistent_data;
     (void) attributes;
+    (void) method;
     if( context->persistent_data_size != sizeof( psa_key_slot_number_t ) )
         return( PSA_ERROR_DETECTED_BY_DRIVER );
     ++*p_counter;
@@ -162,10 +164,12 @@
 static psa_status_t ram_allocate( psa_drv_se_context_t *context,
                                   void *persistent_data,
                                   const psa_key_attributes_t *attributes,
+                                  psa_key_creation_method_t method,
                                   psa_key_slot_number_t *slot_number )
 {
     ram_slot_usage_t *slot_usage = persistent_data;
     (void) attributes;
+    (void) method;
     DRIVER_ASSERT( context->persistent_data_size == sizeof( ram_slot_usage_t ) );
     for( *slot_number = ram_min_slot;
          *slot_number < ARRAY_LENGTH( ram_slots );
@@ -180,10 +184,12 @@
 static psa_status_t ram_validate_slot_number(
     psa_drv_se_context_t *context,
     const psa_key_attributes_t *attributes,
+    psa_key_creation_method_t method,
     psa_key_slot_number_t slot_number )
 {
     (void) context;
     (void) attributes;
+    (void) method;
     if( slot_number >= ARRAY_LENGTH( ram_slots ) )
         return( PSA_ERROR_INVALID_ARGUMENT );
     return( PSA_SUCCESS );