Merge pull request #8815 from gilles-peskine-arm/psa_generate_key_ext-prototype

Introduce psa_generate_key_ext
diff --git a/ChangeLog.d/psa_generate_key_ext.txt b/ChangeLog.d/psa_generate_key_ext.txt
new file mode 100644
index 0000000..8340f01
--- /dev/null
+++ b/ChangeLog.d/psa_generate_key_ext.txt
@@ -0,0 +1,3 @@
+Features
+   * The new function psa_generate_key_ext() allows generating an RSA
+     key pair with a custom public exponent.
diff --git a/include/psa/crypto.h b/include/psa/crypto.h
index fe10ee0..e54af34 100644
--- a/include/psa/crypto.h
+++ b/include/psa/crypto.h
@@ -119,8 +119,9 @@
  * value in the structure.
  * The persistent key will be written to storage when the attribute
  * structure is passed to a key creation function such as
- * psa_import_key(), psa_generate_key(),
- * psa_key_derivation_output_key() or psa_copy_key().
+ * psa_import_key(), psa_generate_key(), psa_generate_key_ext(),
+ * psa_key_derivation_output_key(), psa_key_derivation_output_key_ext()
+ * or psa_copy_key().
  *
  * This function may be declared as `static` (i.e. without external
  * linkage). This function may be provided as a function-like macro,
@@ -163,8 +164,9 @@
  * value in the structure.
  * The persistent key will be written to storage when the attribute
  * structure is passed to a key creation function such as
- * psa_import_key(), psa_generate_key(),
- * psa_key_derivation_output_key() or psa_copy_key().
+ * psa_import_key(), psa_generate_key(), psa_generate_key_ext(),
+ * psa_key_derivation_output_key(), psa_key_derivation_output_key_ext()
+ * or psa_copy_key().
  *
  * This function may be declared as `static` (i.e. without external
  * linkage). This function may be provided as a function-like macro,
@@ -3226,7 +3228,8 @@
  *    psa_key_derivation_set_capacity(). You may do this before, in the middle
  *    of or after providing inputs. For some algorithms, this step is mandatory
  *    because the output depends on the maximum capacity.
- * -# To derive a key, call psa_key_derivation_output_key().
+ * -# To derive a key, call psa_key_derivation_output_key() or
+ *    psa_key_derivation_output_key_ext().
  *    To derive a byte string for a different purpose, call
  *    psa_key_derivation_output_bytes().
  *    Successive calls to these functions use successive output bytes
@@ -3449,7 +3452,8 @@
  * \note Once all inputs steps are completed, the operations will allow:
  * - psa_key_derivation_output_bytes() if each input was either a direct input
  *   or  a key with #PSA_KEY_USAGE_DERIVE set;
- * - psa_key_derivation_output_key() if the input for step
+ * - psa_key_derivation_output_key() or psa_key_derivation_output_key_ext()
+ *   if the input for step
  *   #PSA_KEY_DERIVATION_INPUT_SECRET or #PSA_KEY_DERIVATION_INPUT_PASSWORD
  *   was from a key slot with #PSA_KEY_USAGE_DERIVE and each other input was
  *   either a direct input or a key with #PSA_KEY_USAGE_DERIVE set;
@@ -3697,6 +3701,11 @@
  * Future versions of this specification may include additional restrictions
  * on the derived key based on the attributes and strength of the secret key.
  *
+ * \note This function is equivalent to calling
+ *       psa_key_derivation_output_key_ext()
+ *       with the production parameters #PSA_KEY_PRODUCTION_PARAMETERS_INIT
+ *       and `params_data_length == 0` (i.e. `params->data` is empty).
+ *
  * \param[in] attributes    The attributes for the new key.
  *                          If the key type to be created is
  *                          #PSA_KEY_TYPE_PASSWORD_HASH then the algorithm in
@@ -3750,6 +3759,83 @@
     psa_key_derivation_operation_t *operation,
     mbedtls_svc_key_id_t *key);
 
+/** Derive a key from an ongoing key derivation operation with custom
+ *  production parameters.
+ *
+ * See the description of psa_key_derivation_out_key() for the operation of
+ * this function with the default production parameters.
+ * Mbed TLS currently does not currently support any non-default production
+ * parameters.
+ *
+ * \note This function is experimental and may change in future minor
+ *       versions of Mbed TLS.
+ *
+ * \param[in] attributes    The attributes for the new key.
+ *                          If the key type to be created is
+ *                          #PSA_KEY_TYPE_PASSWORD_HASH then the algorithm in
+ *                          the policy must be the same as in the current
+ *                          operation.
+ * \param[in,out] operation The key derivation operation object to read from.
+ * \param[in] params        Customization parameters for the key derivation.
+ *                          When this is #PSA_KEY_PRODUCTION_PARAMETERS_INIT
+ *                          with \p params_data_length = 0,
+ *                          this function is equivalent to
+ *                          psa_key_derivation_output_key().
+ *                          Mbed TLS currently only supports the default
+ *                          production parameters, i.e.
+ *                          #PSA_KEY_PRODUCTION_PARAMETERS_INIT,
+ *                          for all key types.
+ * \param params_data_length
+ *                          Length of `params->data` in bytes.
+ * \param[out] key          On success, an identifier for the newly created
+ *                          key. For persistent keys, this is the key
+ *                          identifier defined in \p attributes.
+ *                          \c 0 on failure.
+ *
+ * \retval #PSA_SUCCESS
+ *         Success.
+ *         If the key is persistent, the key material and the key's metadata
+ *         have been saved to persistent storage.
+ * \retval #PSA_ERROR_ALREADY_EXISTS
+ *         This is an attempt to create a persistent key, and there is
+ *         already a persistent key with the given identifier.
+ * \retval #PSA_ERROR_INSUFFICIENT_DATA
+ *         There was not enough data to create the desired key.
+ *         Note that in this case, no output is written to the output buffer.
+ *         The operation's capacity is set to 0, thus subsequent calls to
+ *         this function will not succeed, even with a smaller output buffer.
+ * \retval #PSA_ERROR_NOT_SUPPORTED
+ *         The key type or key size is not supported, either by the
+ *         implementation in general or in this particular location.
+ * \retval #PSA_ERROR_INVALID_ARGUMENT
+ *         The provided key attributes are not valid for the operation.
+ * \retval #PSA_ERROR_NOT_PERMITTED
+ *         The #PSA_KEY_DERIVATION_INPUT_SECRET or
+ *         #PSA_KEY_DERIVATION_INPUT_PASSWORD input was not provided through a
+ *         key; or one of the inputs was a key whose policy didn't allow
+ *         #PSA_KEY_USAGE_DERIVE.
+ * \retval #PSA_ERROR_INSUFFICIENT_MEMORY \emptydescription
+ * \retval #PSA_ERROR_INSUFFICIENT_STORAGE \emptydescription
+ * \retval #PSA_ERROR_COMMUNICATION_FAILURE \emptydescription
+ * \retval #PSA_ERROR_HARDWARE_FAILURE \emptydescription
+ * \retval #PSA_ERROR_CORRUPTION_DETECTED \emptydescription
+ * \retval #PSA_ERROR_DATA_INVALID \emptydescription
+ * \retval #PSA_ERROR_DATA_CORRUPT \emptydescription
+ * \retval #PSA_ERROR_STORAGE_FAILURE \emptydescription
+ * \retval #PSA_ERROR_BAD_STATE
+ *         The operation state is not valid (it must be active and completed
+ *         all required input steps), or 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 psa_key_derivation_output_key_ext(
+    const psa_key_attributes_t *attributes,
+    psa_key_derivation_operation_t *operation,
+    const psa_key_production_parameters_t *params,
+    size_t params_data_length,
+    mbedtls_svc_key_id_t *key);
+
 /** Compare output data from a key derivation operation to an expected value.
  *
  * This function calculates output bytes from a key derivation algorithm and
@@ -3835,7 +3921,8 @@
  *                          and the permitted algorithm must match the
  *                          operation. The value of this key was likely
  *                          computed by a previous call to
- *                          psa_key_derivation_output_key().
+ *                          psa_key_derivation_output_key() or
+ *                          psa_key_derivation_output_key_ext().
  *
  * \retval #PSA_SUCCESS \emptydescription
  * \retval #PSA_ERROR_INVALID_SIGNATURE
@@ -4003,6 +4090,10 @@
  *   between 2^{n-1} and 2^n where n is the bit size specified in the
  *   attributes.
  *
+ * \note This function is equivalent to calling psa_generate_key_ext()
+ *       with the production parameters #PSA_KEY_PRODUCTION_PARAMETERS_INIT
+ *       and `params_data_length == 0` (i.e. `params->data` is empty).
+ *
  * \param[in] attributes    The attributes for the new key.
  * \param[out] key          On success, an identifier for the newly created
  *                          key. For persistent keys, this is the key
@@ -4035,6 +4126,60 @@
 psa_status_t psa_generate_key(const psa_key_attributes_t *attributes,
                               mbedtls_svc_key_id_t *key);
 
+/**
+ * \brief Generate a key or key pair using custom production parameters.
+ *
+ * See the description of psa_generate_key() for the operation of this
+ * function with the default production parameters. In addition, this function
+ * supports the following production customizations, described in more detail
+ * in the documentation of ::psa_key_production_parameters_t:
+ *
+ * - RSA keys: generation with a custom public exponent.
+ *
+ * \note This function is experimental and may change in future minor
+ *       versions of Mbed TLS.
+ *
+ * \param[in] attributes    The attributes for the new key.
+ * \param[in] params        Customization parameters for the key generation.
+ *                          When this is #PSA_KEY_PRODUCTION_PARAMETERS_INIT
+ *                          with \p params_data_length = 0,
+ *                          this function is equivalent to
+ *                          psa_key_generation_output_key().
+ * \param params_data_length
+ *                          Length of `params->data` in bytes.
+ * \param[out] key          On success, an identifier for the newly created
+ *                          key. For persistent keys, this is the key
+ *                          identifier defined in \p attributes.
+ *                          \c 0 on failure.
+ *
+ * \retval #PSA_SUCCESS
+ *         Success.
+ *         If the key is persistent, the key material and the key's metadata
+ *         have been saved to persistent storage.
+ * \retval #PSA_ERROR_ALREADY_EXISTS
+ *         This is an attempt to create a persistent key, and there is
+ *         already a persistent key with the given identifier.
+ * \retval #PSA_ERROR_NOT_SUPPORTED \emptydescription
+ * \retval #PSA_ERROR_INVALID_ARGUMENT \emptydescription
+ * \retval #PSA_ERROR_INSUFFICIENT_MEMORY \emptydescription
+ * \retval #PSA_ERROR_INSUFFICIENT_ENTROPY \emptydescription
+ * \retval #PSA_ERROR_COMMUNICATION_FAILURE \emptydescription
+ * \retval #PSA_ERROR_HARDWARE_FAILURE \emptydescription
+ * \retval #PSA_ERROR_CORRUPTION_DETECTED \emptydescription
+ * \retval #PSA_ERROR_INSUFFICIENT_STORAGE \emptydescription
+ * \retval #PSA_ERROR_DATA_INVALID \emptydescription
+ * \retval #PSA_ERROR_DATA_CORRUPT \emptydescription
+ * \retval #PSA_ERROR_STORAGE_FAILURE \emptydescription
+ * \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 psa_generate_key_ext(const psa_key_attributes_t *attributes,
+                                  const psa_key_production_parameters_t *params,
+                                  size_t params_data_length,
+                                  mbedtls_svc_key_id_t *key);
+
 /**@}*/
 
 /** \defgroup interruptible_hash Interruptible sign/verify hash
diff --git a/include/psa/crypto_struct.h b/include/psa/crypto_struct.h
index 34d072b..e2068e8 100644
--- a/include/psa/crypto_struct.h
+++ b/include/psa/crypto_struct.h
@@ -223,6 +223,22 @@
     return v;
 }
 
+struct psa_key_production_parameters_s {
+    /* Future versions may add other fields in this structure. */
+    uint32_t flags;
+    uint8_t data[];
+};
+
+/** The default production parameters for key generation or key derivation.
+ *
+ * Calling psa_generate_key_ext() or psa_key_derivation_output_key_ext()
+ * with `params=PSA_KEY_PRODUCTION_PARAMETERS_INIT` and
+ * `params_data_length == 0` is equivalent to
+ * calling psa_generate_key() or psa_key_derivation_output_key()
+ * respectively.
+ */
+#define PSA_KEY_PRODUCTION_PARAMETERS_INIT { 0 }
+
 struct psa_key_policy_s {
     psa_key_usage_t MBEDTLS_PRIVATE(usage);
     psa_algorithm_t MBEDTLS_PRIVATE(alg);
diff --git a/include/psa/crypto_types.h b/include/psa/crypto_types.h
index e2ebd8a..31ea686 100644
--- a/include/psa/crypto_types.h
+++ b/include/psa/crypto_types.h
@@ -455,6 +455,30 @@
  */
 typedef uint16_t psa_key_derivation_step_t;
 
+/** \brief Custom parameters for key generation or key derivation.
+ *
+ * This is a structure type with at least the following fields:
+ *
+ * - \c flags: an unsigned integer type. 0 for the default production parameters.
+ * - \c data: a flexible array of bytes.
+ *
+ * The interpretation of this structure depend on the type of the
+ * created key.
+ *
+ * - #PSA_KEY_TYPE_RSA_KEY_PAIR:
+ *     - \c flags: must be 0.
+ *     - \c data: the public exponent, in little-endian order.
+ *       This must be an odd integer and must not be 1.
+ *       Implementations must support 65535, should support 3 and may
+ *       support other values.
+ *       When not using a driver, Mbed TLS supports values up to \c INT_MAX.
+ *       If this is empty or if the custom production parameters are omitted
+ *       altogether, the default value 65537 is used.
+ * - Other key types: reserved for future use. \c flags must be 0.
+ *
+ */
+typedef struct psa_key_production_parameters_s psa_key_production_parameters_t;
+
 /**@}*/
 
 #endif /* PSA_CRYPTO_TYPES_H */
diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 87444e1..e9061f5 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -6066,9 +6066,28 @@
     return status;
 }
 
-psa_status_t psa_key_derivation_output_key(const psa_key_attributes_t *attributes,
-                                           psa_key_derivation_operation_t *operation,
-                                           mbedtls_svc_key_id_t *key)
+static const psa_key_production_parameters_t default_production_parameters =
+    PSA_KEY_PRODUCTION_PARAMETERS_INIT;
+
+int psa_key_production_parameters_are_default(
+    const psa_key_production_parameters_t *params,
+    size_t params_data_length)
+{
+    if (params->flags != 0) {
+        return 0;
+    }
+    if (params_data_length != 0) {
+        return 0;
+    }
+    return 1;
+}
+
+psa_status_t psa_key_derivation_output_key_ext(
+    const psa_key_attributes_t *attributes,
+    psa_key_derivation_operation_t *operation,
+    const psa_key_production_parameters_t *params,
+    size_t params_data_length,
+    mbedtls_svc_key_id_t *key)
 {
     psa_status_t status;
     psa_key_slot_t *slot = NULL;
@@ -6082,6 +6101,10 @@
         return PSA_ERROR_INVALID_ARGUMENT;
     }
 
+    if (!psa_key_production_parameters_are_default(params, params_data_length)) {
+        return PSA_ERROR_INVALID_ARGUMENT;
+    }
+
     if (operation->alg == PSA_ALG_NONE) {
         return PSA_ERROR_BAD_STATE;
     }
@@ -6113,6 +6136,15 @@
     return status;
 }
 
+psa_status_t psa_key_derivation_output_key(
+    const psa_key_attributes_t *attributes,
+    psa_key_derivation_operation_t *operation,
+    mbedtls_svc_key_id_t *key)
+{
+    return psa_key_derivation_output_key_ext(attributes, operation,
+                                             &default_production_parameters, 0,
+                                             key);
+}
 
 
 /****************************************************************/
@@ -7509,11 +7541,16 @@
 
 psa_status_t psa_generate_key_internal(
     const psa_key_attributes_t *attributes,
+    const psa_key_production_parameters_t *params, size_t params_data_length,
     uint8_t *key_buffer, size_t key_buffer_size, size_t *key_buffer_length)
 {
     psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
     psa_key_type_t type = attributes->core.type;
 
+    /* Only used for RSA */
+    (void) params;
+    (void) params_data_length;
+
     if ((attributes->domain_parameters == NULL) &&
         (attributes->domain_parameters_size != 0)) {
         return PSA_ERROR_INVALID_ARGUMENT;
@@ -7534,7 +7571,16 @@
 
 #if defined(MBEDTLS_PSA_BUILTIN_KEY_TYPE_RSA_KEY_PAIR_GENERATE)
     if (type == PSA_KEY_TYPE_RSA_KEY_PAIR) {
-        return mbedtls_psa_rsa_generate_key(attributes,
+        /* Hack: if the method specifies a non-default e, pass it
+         * via the domain parameters. TODO: refactor this code so
+         * that mbedtls_psa_rsa_generate_key() gets e via a new
+         * parameter instead. */
+        psa_key_attributes_t override_attributes = *attributes;
+        if (params_data_length != 0) {
+            override_attributes.domain_parameters_size = params_data_length;
+            override_attributes.domain_parameters = (uint8_t *) &params->data;
+        }
+        return mbedtls_psa_rsa_generate_key(&override_attributes,
                                             key_buffer,
                                             key_buffer_size,
                                             key_buffer_length);
@@ -7566,8 +7612,10 @@
     return PSA_SUCCESS;
 }
 
-psa_status_t psa_generate_key(const psa_key_attributes_t *attributes,
-                              mbedtls_svc_key_id_t *key)
+psa_status_t psa_generate_key_ext(const psa_key_attributes_t *attributes,
+                                  const psa_key_production_parameters_t *params,
+                                  size_t params_data_length,
+                                  mbedtls_svc_key_id_t *key)
 {
     psa_status_t status;
     psa_key_slot_t *slot = NULL;
@@ -7587,6 +7635,17 @@
         return PSA_ERROR_INVALID_ARGUMENT;
     }
 
+#if defined(PSA_WANT_KEY_TYPE_RSA_KEY_PAIR_GENERATE)
+    if (attributes->core.type == PSA_KEY_TYPE_RSA_KEY_PAIR) {
+        if (params->flags != 0) {
+            return PSA_ERROR_INVALID_ARGUMENT;
+        }
+    } else
+#endif
+    if (!psa_key_production_parameters_are_default(params, params_data_length)) {
+        return PSA_ERROR_INVALID_ARGUMENT;
+    }
+
     status = psa_start_key_creation(PSA_KEY_CREATION_GENERATE, attributes,
                                     &slot, &driver);
     if (status != PSA_SUCCESS) {
@@ -7624,8 +7683,9 @@
     }
 
     status = psa_driver_wrapper_generate_key(attributes,
-                                             slot->key.data, slot->key.bytes, &slot->key.bytes);
-
+                                             params, params_data_length,
+                                             slot->key.data, slot->key.bytes,
+                                             &slot->key.bytes);
     if (status != PSA_SUCCESS) {
         psa_remove_key_data_from_memory(slot);
     }
@@ -7641,6 +7701,14 @@
     return status;
 }
 
+psa_status_t psa_generate_key(const psa_key_attributes_t *attributes,
+                              mbedtls_svc_key_id_t *key)
+{
+    return psa_generate_key_ext(attributes,
+                                &default_production_parameters, 0,
+                                key);
+}
+
 /****************************************************************/
 /* Module setup */
 /****************************************************************/
diff --git a/library/psa_crypto_core.h b/library/psa_crypto_core.h
index 0d7322c..afa8659 100644
--- a/library/psa_crypto_core.h
+++ b/library/psa_crypto_core.h
@@ -398,6 +398,18 @@
     const uint8_t *key_buffer, size_t key_buffer_size,
     uint8_t *data, size_t data_size, size_t *data_length);
 
+/** Whether a key production parameters structure is the default.
+ *
+ * Calls to a key generation driver with non-default production parameters
+ * require a driver supporting custom production parameters.
+ *
+ * \param[in] params            The key production parameters to check.
+ * \param params_data_length    Size of `params->data` in bytes.
+ */
+int psa_key_production_parameters_are_default(
+    const psa_key_production_parameters_t *params,
+    size_t params_data_length);
+
 /**
  * \brief Generate a key.
  *
@@ -405,6 +417,9 @@
  *       entry point.
  *
  * \param[in]  attributes         The attributes for the key to generate.
+ * \param[in]  params             The production parameters from
+ *                                psa_generate_key_ext().
+ * \param      params_data_length The size of `params->data` in bytes.
  * \param[out] key_buffer         Buffer where the key data is to be written.
  * \param[in]  key_buffer_size    Size of \p key_buffer in bytes.
  * \param[out] key_buffer_length  On success, the number of bytes written in
@@ -419,6 +434,8 @@
  *         The size of \p key_buffer is too small.
  */
 psa_status_t psa_generate_key_internal(const psa_key_attributes_t *attributes,
+                                       const psa_key_production_parameters_t *params,
+                                       size_t params_data_length,
                                        uint8_t *key_buffer,
                                        size_t key_buffer_size,
                                        size_t *key_buffer_length);
diff --git a/scripts/data_files/driver_templates/psa_crypto_driver_wrappers.h.jinja b/scripts/data_files/driver_templates/psa_crypto_driver_wrappers.h.jinja
index 924b08c..4f9764d 100644
--- a/scripts/data_files/driver_templates/psa_crypto_driver_wrappers.h.jinja
+++ b/scripts/data_files/driver_templates/psa_crypto_driver_wrappers.h.jinja
@@ -731,12 +731,27 @@
 
 static inline psa_status_t psa_driver_wrapper_generate_key(
     const psa_key_attributes_t *attributes,
+    const psa_key_production_parameters_t *params, size_t params_data_length,
     uint8_t *key_buffer, size_t key_buffer_size, size_t *key_buffer_length )
 {
     psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
     psa_key_location_t location =
         PSA_KEY_LIFETIME_GET_LOCATION(attributes->core.lifetime);
 
+#if defined(PSA_WANT_KEY_TYPE_RSA_KEY_PAIR_GENERATE)
+    int is_default_production =
+        psa_key_production_parameters_are_default(params, params_data_length);
+    if( location != PSA_KEY_LOCATION_LOCAL_STORAGE && !is_default_production )
+    {
+        /* We don't support passing custom production parameters
+         * to drivers yet. */
+        return PSA_ERROR_NOT_SUPPORTED;
+    }
+#else
+    int is_default_production = 1;
+    (void) is_default_production;
+#endif
+
     /* Try dynamically-registered SE interface first */
 #if defined(MBEDTLS_PSA_CRYPTO_SE_C)
     const psa_drv_se_t *drv;
@@ -762,8 +777,11 @@
     {
         case PSA_KEY_LOCATION_LOCAL_STORAGE:
 #if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
-            /* Transparent drivers are limited to generating asymmetric keys */
-            if( PSA_KEY_TYPE_IS_ASYMMETRIC( attributes->core.type ) )
+            /* Transparent drivers are limited to generating asymmetric keys. */
+            /* We don't support passing custom production parameters
+             * to drivers yet. */
+            if( PSA_KEY_TYPE_IS_ASYMMETRIC( attributes->core.type ) &&
+                is_default_production )
             {
             /* Cycle through all known transparent accelerators */
 #if defined(PSA_CRYPTO_DRIVER_TEST)
@@ -793,7 +811,8 @@
 
             /* Software fallback */
             status = psa_generate_key_internal(
-                attributes, key_buffer, key_buffer_size, key_buffer_length );
+                attributes, params, params_data_length,
+                key_buffer, key_buffer_size, key_buffer_length );
             break;
 
         /* Add cases for opaque driver here */
diff --git a/tests/scripts/analyze_outcomes.py b/tests/scripts/analyze_outcomes.py
index 2a29f71..5b4deb6 100755
--- a/tests/scripts/analyze_outcomes.py
+++ b/tests/scripts/analyze_outcomes.py
@@ -570,6 +570,11 @@
                     re.compile(r'mbedtls_ct_zeroize_if .*'),
                     re.compile(r'mbedtls_ct_memmove_left .*')
                 ],
+                'test_suite_psa_crypto': [
+                    # We don't support generate_key_ext entry points
+                    # in drivers yet.
+                    re.compile(r'PSA generate key ext: RSA, e=.*'),
+                ],
             }
         }
     },
diff --git a/tests/suites/test_suite_psa_crypto.data b/tests/suites/test_suite_psa_crypto.data
index 38e4046..e9b5717 100644
--- a/tests/suites/test_suite_psa_crypto.data
+++ b/tests/suites/test_suite_psa_crypto.data
@@ -6870,6 +6870,18 @@
 depends_on:PSA_WANT_ALG_PBKDF2_AES_CMAC_PRF_128:PSA_WANT_ALG_CMAC:PSA_WANT_KEY_TYPE_AES:!MBEDTLS_AES_ONLY_128_BIT_KEY_LENGTH
 derive_key_type:PSA_ALG_PBKDF2_AES_CMAC_PRF_128:"706173737764":"01":"73616c74":PSA_KEY_TYPE_AES:256:"28e288c6345bb5ecf7ca70274208a3ba0f1148b5868537d5e09d3ee6813b1f52"
 
+PSA key derivation: default params -> AES-128
+depends_on:PSA_WANT_ALG_HKDF:PSA_WANT_ALG_SHA_256:PSA_WANT_KEY_TYPE_AES
+derive_key_ext:PSA_ALG_HKDF(PSA_ALG_SHA_256):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":PSA_KEY_TYPE_AES:128:0:"":PSA_SUCCESS:"3cb25f25faacd57a90434f64d0362f2a"
+
+PSA key derivation: params.flags=1 -> AES-128
+depends_on:PSA_WANT_ALG_HKDF:PSA_WANT_ALG_SHA_256:PSA_WANT_KEY_TYPE_AES
+derive_key_ext:PSA_ALG_HKDF(PSA_ALG_SHA_256):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":PSA_KEY_TYPE_AES:128:1:"":PSA_ERROR_INVALID_ARGUMENT:""
+
+PSA key derivation: params.data non-empty -> AES-128
+depends_on:PSA_WANT_ALG_HKDF:PSA_WANT_ALG_SHA_256:PSA_WANT_KEY_TYPE_AES
+derive_key_ext:PSA_ALG_HKDF(PSA_ALG_SHA_256):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":PSA_KEY_TYPE_AES:128:0:"2a":PSA_ERROR_INVALID_ARGUMENT:""
+
 PSA key derivation: invalid type (0)
 depends_on:PSA_WANT_ALG_HKDF:PSA_WANT_ALG_SHA_256
 derive_key:PSA_ALG_HKDF(PSA_ALG_SHA_256):"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":"000102030405060708090a0b0c":"f0f1f2f3f4f5f6f7f8f9":PSA_KEY_TYPE_NONE:128:PSA_ERROR_NOT_SUPPORTED:0
@@ -7419,22 +7431,22 @@
 depends_on:PSA_WANT_ALG_ECDH:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR_GENERATE:PSA_WANT_ECC_MONTGOMERY_448
 generate_key:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_MONTGOMERY):448:PSA_KEY_USAGE_EXPORT | PSA_KEY_USAGE_DERIVE:PSA_ALG_ECDH:PSA_SUCCESS:0
 
-PSA generate key: RSA, default e
+PSA generate key: RSA, domain parameters: default e
 generate_key_rsa:PSA_VENDOR_RSA_GENERATE_MIN_KEY_BITS:"":PSA_SUCCESS
 
-PSA generate key: RSA, e=3
+PSA generate key: RSA, domain parameters: e=3
 generate_key_rsa:PSA_VENDOR_RSA_GENERATE_MIN_KEY_BITS:"03":PSA_SUCCESS
 
-PSA generate key: RSA, e=65537
+PSA generate key: RSA, domain parameters: e=65537
 generate_key_rsa:PSA_VENDOR_RSA_GENERATE_MIN_KEY_BITS:"010001":PSA_SUCCESS
 
-PSA generate key: RSA, e=513
+PSA generate key: RSA, domain parameters: e=513
 generate_key_rsa:PSA_VENDOR_RSA_GENERATE_MIN_KEY_BITS:"0201":PSA_SUCCESS
 
-PSA generate key: RSA, e=1
+PSA generate key: RSA, domain parameters: e=1
 generate_key_rsa:PSA_VENDOR_RSA_GENERATE_MIN_KEY_BITS:"01":PSA_ERROR_INVALID_ARGUMENT
 
-PSA generate key: RSA, e=2
+PSA generate key: RSA, domain parameters: e=2
 generate_key_rsa:PSA_VENDOR_RSA_GENERATE_MIN_KEY_BITS:"02":PSA_ERROR_INVALID_ARGUMENT
 
 PSA generate key: FFDH, 2048 bits, good
@@ -7461,6 +7473,86 @@
 depends_on:PSA_WANT_ALG_FFDH:PSA_WANT_KEY_TYPE_DH_KEY_PAIR_GENERATE
 generate_key:PSA_KEY_TYPE_DH_KEY_PAIR(PSA_DH_FAMILY_RFC7919):1024:PSA_KEY_USAGE_EXPORT:PSA_ALG_FFDH:PSA_ERROR_NOT_SUPPORTED:0
 
+PSA generate key ext: RSA, params.flags=1
+depends_on:PSA_WANT_KEY_TYPE_RSA_KEY_PAIR_GENERATE
+generate_key_ext:PSA_KEY_TYPE_RSA_KEY_PAIR:PSA_VENDOR_RSA_GENERATE_MIN_KEY_BITS:PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT:0:1:"":PSA_ERROR_INVALID_ARGUMENT
+
+PSA generate key ext: RSA, empty e
+depends_on:PSA_WANT_KEY_TYPE_RSA_KEY_PAIR_GENERATE:PSA_WANT_ALG_RSA_PKCS1V15_CRYPT
+generate_key_ext:PSA_KEY_TYPE_RSA_KEY_PAIR:PSA_VENDOR_RSA_GENERATE_MIN_KEY_BITS:PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT:PSA_ALG_RSA_PKCS1V15_CRYPT:0:"":PSA_SUCCESS
+
+PSA generate key ext: RSA, e=3
+depends_on:PSA_WANT_KEY_TYPE_RSA_KEY_PAIR_GENERATE:MBEDTLS_PSA_BUILTIN_KEY_TYPE_RSA_KEY_PAIR_GENERATE:PSA_WANT_ALG_RSA_PKCS1V15_CRYPT
+generate_key_ext:PSA_KEY_TYPE_RSA_KEY_PAIR:PSA_VENDOR_RSA_GENERATE_MIN_KEY_BITS:PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT:PSA_ALG_RSA_PKCS1V15_CRYPT:0:"03":PSA_SUCCESS
+
+PSA generate key ext: RSA, e=3 with leading zeros
+depends_on:PSA_WANT_KEY_TYPE_RSA_KEY_PAIR_GENERATE:MBEDTLS_PSA_BUILTIN_KEY_TYPE_RSA_KEY_PAIR_GENERATE:PSA_WANT_ALG_RSA_PKCS1V15_CRYPT
+generate_key_ext:PSA_KEY_TYPE_RSA_KEY_PAIR:PSA_VENDOR_RSA_GENERATE_MIN_KEY_BITS:PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT:PSA_ALG_RSA_PKCS1V15_CRYPT:0:"000003":PSA_SUCCESS
+
+# TODO: currently errors with NOT_SUPPORTED because e is converted to an int
+# and the conversion errors out if there are too many digits without checking
+# for leading zeros. This is a very minor bug. Re-enable this test when this
+# bug is fixed.
+#PSA generate key ext: RSA, e=3 with many leading zeros
+#depends_on:PSA_WANT_KEY_TYPE_RSA_KEY_PAIR_GENERATE:MBEDTLS_PSA_BUILTIN_KEY_TYPE_RSA_KEY_PAIR_GENERATE:PSA_WANT_ALG_RSA_PKCS1V15_CRYPT
+#generate_key_ext:PSA_KEY_TYPE_RSA_KEY_PAIR:PSA_VENDOR_RSA_GENERATE_MIN_KEY_BITS:PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT:PSA_ALG_RSA_PKCS1V15_CRYPT:0:"0000000000000000000000000000000003":PSA_SUCCESS
+
+PSA generate key ext: RSA, e=513
+depends_on:PSA_WANT_KEY_TYPE_RSA_KEY_PAIR_GENERATE:MBEDTLS_PSA_BUILTIN_KEY_TYPE_RSA_KEY_PAIR_GENERATE:PSA_WANT_ALG_RSA_PKCS1V15_CRYPT
+generate_key_ext:PSA_KEY_TYPE_RSA_KEY_PAIR:PSA_VENDOR_RSA_GENERATE_MIN_KEY_BITS:PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT:PSA_ALG_RSA_PKCS1V15_CRYPT:0:"0201":PSA_SUCCESS
+
+PSA generate key ext: RSA, e=65537
+depends_on:PSA_WANT_KEY_TYPE_RSA_KEY_PAIR_GENERATE:MBEDTLS_PSA_BUILTIN_KEY_TYPE_RSA_KEY_PAIR_GENERATE:PSA_WANT_ALG_RSA_PKCS1V15_CRYPT
+generate_key_ext:PSA_KEY_TYPE_RSA_KEY_PAIR:PSA_VENDOR_RSA_GENERATE_MIN_KEY_BITS:PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT:PSA_ALG_RSA_PKCS1V15_CRYPT:0:"010001":PSA_SUCCESS
+
+PSA generate key ext: RSA, e=2^31-1
+depends_on:PSA_WANT_KEY_TYPE_RSA_KEY_PAIR_GENERATE:MBEDTLS_PSA_BUILTIN_KEY_TYPE_RSA_KEY_PAIR_GENERATE:PSA_WANT_ALG_RSA_PKCS1V15_CRYPT:INT_MAX>=0x7fffffff
+generate_key_ext:PSA_KEY_TYPE_RSA_KEY_PAIR:PSA_VENDOR_RSA_GENERATE_MIN_KEY_BITS:PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT:PSA_ALG_RSA_PKCS1V15_CRYPT:0:"7fffffff":PSA_SUCCESS
+
+PSA generate key ext: RSA, e=2^31+3 (too large for built-in RSA)
+depends_on:PSA_WANT_KEY_TYPE_RSA_KEY_PAIR_GENERATE:MBEDTLS_PSA_BUILTIN_KEY_TYPE_RSA_KEY_PAIR_GENERATE:!MBEDTLS_PSA_ACCEL_KEY_TYPE_RSA_KEY_PAIR_GENERATE:INT_MAX<=0x7fffffff
+generate_key_ext:PSA_KEY_TYPE_RSA_KEY_PAIR:PSA_VENDOR_RSA_GENERATE_MIN_KEY_BITS:PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT:0:0:"80000003":PSA_ERROR_NOT_SUPPORTED
+
+PSA generate key ext: RSA, e=2^64+3 (too large for built-in RSA)
+depends_on:PSA_WANT_KEY_TYPE_RSA_KEY_PAIR_GENERATE:MBEDTLS_PSA_BUILTIN_KEY_TYPE_RSA_KEY_PAIR_GENERATE:!MBEDTLS_PSA_ACCEL_KEY_TYPE_RSA_KEY_PAIR_GENERATE:INT_MAX<=0xffffffffffffffff
+generate_key_ext:PSA_KEY_TYPE_RSA_KEY_PAIR:PSA_VENDOR_RSA_GENERATE_MIN_KEY_BITS:PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT:0:0:"010000000000000003":PSA_ERROR_NOT_SUPPORTED
+
+PSA generate key ext: RSA, e=1
+depends_on:PSA_WANT_KEY_TYPE_RSA_KEY_PAIR_GENERATE:MBEDTLS_PSA_BUILTIN_KEY_TYPE_RSA_KEY_PAIR_GENERATE
+generate_key_ext:PSA_KEY_TYPE_RSA_KEY_PAIR:PSA_VENDOR_RSA_GENERATE_MIN_KEY_BITS:PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT:0:0:"01":PSA_ERROR_INVALID_ARGUMENT
+
+PSA generate key ext: RSA, e=0
+depends_on:PSA_WANT_KEY_TYPE_RSA_KEY_PAIR_GENERATE:MBEDTLS_PSA_BUILTIN_KEY_TYPE_RSA_KEY_PAIR_GENERATE
+generate_key_ext:PSA_KEY_TYPE_RSA_KEY_PAIR:PSA_VENDOR_RSA_GENERATE_MIN_KEY_BITS:PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT:0:0:"00":PSA_ERROR_INVALID_ARGUMENT
+
+PSA generate key ext: RSA, e=2
+depends_on:PSA_WANT_KEY_TYPE_RSA_KEY_PAIR_GENERATE:MBEDTLS_PSA_BUILTIN_KEY_TYPE_RSA_KEY_PAIR_GENERATE
+generate_key_ext:PSA_KEY_TYPE_RSA_KEY_PAIR:PSA_VENDOR_RSA_GENERATE_MIN_KEY_BITS:PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT:0:0:"02":PSA_ERROR_INVALID_ARGUMENT
+
+# Check that with a driver, we reject a custom e as unsupported,
+# as opposed to silently using the default e.
+# When we add proper driver support, remove this test case and remove
+# the dependency on MBEDTLS_PSA_BUILTIN_KEY_TYPE_RSA_KEY_PAIR_GENERATE from
+# the positive/invalid_argument test cases.
+PSA generate key ext: RSA, e=3 with driver and no fallback (not yet supported)
+depends_on:PSA_WANT_KEY_TYPE_RSA_KEY_PAIR_GENERATE:!MBEDTLS_PSA_BUILTIN_KEY_TYPE_RSA_KEY_PAIR_GENERATE
+generate_key_ext:PSA_KEY_TYPE_RSA_KEY_PAIR:PSA_VENDOR_RSA_GENERATE_MIN_KEY_BITS:PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT:0:0:"03":PSA_ERROR_NOT_SUPPORTED
+
+PSA generate key ext: ECC, flags=0
+depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR_GENERATE:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_ECDH
+generate_key_ext:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1):256:PSA_KEY_USAGE_DERIVE:PSA_ALG_ECDH:0:"":PSA_SUCCESS
+
+PSA generate key ext: ECC, flags=1
+depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR_GENERATE:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_ECDH
+generate_key_ext:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1):256:PSA_KEY_USAGE_DERIVE:PSA_ALG_ECDH:1:"":PSA_ERROR_INVALID_ARGUMENT
+
+PSA generate key ext: ECC, params.data non-empty
+depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR_GENERATE:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_ECDH
+generate_key_ext:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1):256:PSA_KEY_USAGE_DERIVE:PSA_ALG_ECDH:0:"2a":PSA_ERROR_INVALID_ARGUMENT
+
+Key production parameters initializers
+key_production_parameters_init:
+
 PSA import persistent key: raw data, 8 bits
 depends_on:MBEDTLS_PSA_CRYPTO_STORAGE_C
 persistent_key_load_key_from_storage:"2a":PSA_KEY_TYPE_RAW_DATA:8:PSA_KEY_USAGE_EXPORT:0:IMPORT_KEY
diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index 95e7a2d..b40b5f8 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -1255,6 +1255,84 @@
 }
 #endif /* MBEDTLS_ECP_RESTARTABLE */
 
+#if defined(PSA_WANT_KEY_TYPE_RSA_KEY_PAIR_GENERATE)
+static int rsa_test_e(mbedtls_svc_key_id_t key,
+                      size_t bits,
+                      const data_t *e_arg)
+{
+    uint8_t *exported = NULL;
+    size_t exported_size =
+        PSA_EXPORT_KEY_OUTPUT_SIZE(PSA_KEY_TYPE_RSA_PUBLIC_KEY, bits);
+    size_t exported_length = SIZE_MAX;
+    int ok = 0;
+
+    TEST_CALLOC(exported, exported_size);
+    PSA_ASSERT(psa_export_public_key(key,
+                                     exported, exported_size,
+                                     &exported_length));
+    uint8_t *p = exported;
+    uint8_t *end = exported + exported_length;
+    size_t len;
+    /*   RSAPublicKey ::= SEQUENCE {
+     *      modulus            INTEGER,    -- n
+     *      publicExponent     INTEGER  }  -- e
+     */
+    TEST_EQUAL(0, mbedtls_asn1_get_tag(&p, end, &len,
+                                       MBEDTLS_ASN1_SEQUENCE |
+                                       MBEDTLS_ASN1_CONSTRUCTED));
+    TEST_ASSERT(mbedtls_test_asn1_skip_integer(&p, end, bits, bits, 1));
+    TEST_EQUAL(0, mbedtls_asn1_get_tag(&p, end, &len,
+                                       MBEDTLS_ASN1_INTEGER));
+    if (len >= 1 && p[0] == 0) {
+        ++p;
+        --len;
+    }
+    if (e_arg->len == 0) {
+        TEST_EQUAL(len, 3);
+        TEST_EQUAL(p[0], 1);
+        TEST_EQUAL(p[1], 0);
+        TEST_EQUAL(p[2], 1);
+    } else {
+        const uint8_t *expected = e_arg->x;
+        size_t expected_len = e_arg->len;
+        while (expected_len > 0 && *expected == 0) {
+            ++expected;
+            --expected_len;
+        }
+        TEST_MEMORY_COMPARE(p, len, expected, expected_len);
+    }
+    ok = 1;
+
+exit:
+    mbedtls_free(exported);
+    return ok;
+}
+#endif /* PSA_WANT_KEY_TYPE_RSA_KEY_PAIR_GENERATE */
+
+static int setup_key_production_parameters(
+    psa_key_production_parameters_t **params, size_t *params_data_length,
+    int flags_arg, const data_t *params_data)
+{
+    *params_data_length = params_data->len;
+    /* If there are N bytes of padding at the end of
+     * psa_key_production_parameters_t, then it's enough to allocate
+     * MIN(sizeof(psa_key_production_parameters_t),
+     *     offsetof(psa_key_production_parameters_t, data) + params_data_length).
+     *
+     * For simplicity, here, we allocate up to N more bytes than necessary.
+     * In practice, the current layout of psa_key_production_parameters_t
+     * makes padding extremely unlikely, so we don't worry about testing
+     * that the library code doesn't try to access these extra N bytes.
+     */
+    *params = mbedtls_calloc(1, sizeof(**params) + *params_data_length);
+    TEST_ASSERT(*params != NULL);
+    (*params)->flags = (uint32_t) flags_arg;
+    memcpy((*params)->data, params_data->x, params_data->len);
+    return 1;
+exit:
+    return 0;
+}
+
 /* END_HEADER */
 
 /* BEGIN_DEPENDENCIES
@@ -9255,6 +9333,81 @@
 /* END_CASE */
 
 /* BEGIN_CASE */
+void derive_key_ext(int alg_arg,
+                    data_t *key_data,
+                    data_t *input1,
+                    data_t *input2,
+                    int key_type_arg, int bits_arg,
+                    int flags_arg,
+                    data_t *params_data,
+                    psa_status_t expected_status,
+                    data_t *expected_export)
+{
+    mbedtls_svc_key_id_t base_key = MBEDTLS_SVC_KEY_ID_INIT;
+    mbedtls_svc_key_id_t derived_key = MBEDTLS_SVC_KEY_ID_INIT;
+    const psa_algorithm_t alg = alg_arg;
+    const psa_key_type_t key_type = key_type_arg;
+    const size_t bits = bits_arg;
+    psa_key_production_parameters_t *params = NULL;
+    size_t params_data_length = 0;
+    psa_key_derivation_operation_t operation = PSA_KEY_DERIVATION_OPERATION_INIT;
+    const size_t export_buffer_size =
+        PSA_EXPORT_KEY_OUTPUT_SIZE(key_type, bits);
+    uint8_t *export_buffer = NULL;
+    psa_key_attributes_t base_attributes = PSA_KEY_ATTRIBUTES_INIT;
+    psa_key_attributes_t derived_attributes = PSA_KEY_ATTRIBUTES_INIT;
+    size_t export_length;
+
+    TEST_CALLOC(export_buffer, export_buffer_size);
+    PSA_ASSERT(psa_crypto_init());
+
+    psa_set_key_usage_flags(&base_attributes, PSA_KEY_USAGE_DERIVE);
+    psa_set_key_algorithm(&base_attributes, alg);
+    psa_set_key_type(&base_attributes, PSA_KEY_TYPE_DERIVE);
+    PSA_ASSERT(psa_import_key(&base_attributes, key_data->x, key_data->len,
+                              &base_key));
+
+    if (mbedtls_test_psa_setup_key_derivation_wrap(
+            &operation, base_key, alg,
+            input1->x, input1->len,
+            input2->x, input2->len,
+            PSA_KEY_DERIVATION_UNLIMITED_CAPACITY) == 0) {
+        goto exit;
+    }
+
+    psa_set_key_usage_flags(&derived_attributes, PSA_KEY_USAGE_EXPORT);
+    psa_set_key_algorithm(&derived_attributes, 0);
+    psa_set_key_type(&derived_attributes, key_type);
+    psa_set_key_bits(&derived_attributes, bits);
+    if (!setup_key_production_parameters(&params, &params_data_length,
+                                         flags_arg, params_data)) {
+        goto exit;
+    }
+
+    TEST_EQUAL(psa_key_derivation_output_key_ext(&derived_attributes, &operation,
+                                                 params, params_data_length,
+                                                 &derived_key),
+               expected_status);
+
+    if (expected_status == PSA_SUCCESS) {
+        PSA_ASSERT(psa_export_key(derived_key,
+                                  export_buffer, export_buffer_size,
+                                  &export_length));
+        TEST_MEMORY_COMPARE(export_buffer, export_length,
+                            expected_export->x, expected_export->len);
+    }
+
+exit:
+    mbedtls_free(export_buffer);
+    mbedtls_free(params);
+    psa_key_derivation_abort(&operation);
+    psa_destroy_key(base_key);
+    psa_destroy_key(derived_key);
+    PSA_DONE();
+}
+/* END_CASE */
+
+/* BEGIN_CASE */
 void derive_key(int alg_arg,
                 data_t *key_data, data_t *input1, data_t *input2,
                 int type_arg, int bits_arg,
@@ -9699,10 +9852,6 @@
     psa_algorithm_t alg = PSA_ALG_RSA_PKCS1V15_SIGN_RAW;
     psa_status_t expected_status = expected_status_arg;
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
-    uint8_t *exported = NULL;
-    size_t exported_size =
-        PSA_EXPORT_KEY_OUTPUT_SIZE(PSA_KEY_TYPE_RSA_PUBLIC_KEY, bits);
-    size_t exported_length = SIZE_MAX;
     uint8_t *e_read_buffer = NULL;
     int is_default_public_exponent = 0;
     size_t e_read_size = PSA_KEY_DOMAIN_PARAMETERS_SIZE(type, bits);
@@ -9715,7 +9864,6 @@
         e_read_size = 0;
     }
     TEST_CALLOC(e_read_buffer, e_read_size);
-    TEST_CALLOC(exported, exported_size);
 
     PSA_ASSERT(psa_crypto_init());
 
@@ -9759,37 +9907,7 @@
         goto exit;
     }
 
-    /* Export the key and check the public exponent. */
-    PSA_ASSERT(psa_export_public_key(key,
-                                     exported, exported_size,
-                                     &exported_length));
-    {
-        uint8_t *p = exported;
-        uint8_t *end = exported + exported_length;
-        size_t len;
-        /*   RSAPublicKey ::= SEQUENCE {
-         *      modulus            INTEGER,    -- n
-         *      publicExponent     INTEGER  }  -- e
-         */
-        TEST_EQUAL(0, mbedtls_asn1_get_tag(&p, end, &len,
-                                           MBEDTLS_ASN1_SEQUENCE |
-                                           MBEDTLS_ASN1_CONSTRUCTED));
-        TEST_ASSERT(mbedtls_test_asn1_skip_integer(&p, end, bits, bits, 1));
-        TEST_EQUAL(0, mbedtls_asn1_get_tag(&p, end, &len,
-                                           MBEDTLS_ASN1_INTEGER));
-        if (len >= 1 && p[0] == 0) {
-            ++p;
-            --len;
-        }
-        if (e_arg->len == 0) {
-            TEST_EQUAL(len, 3);
-            TEST_EQUAL(p[0], 1);
-            TEST_EQUAL(p[1], 0);
-            TEST_EQUAL(p[2], 1);
-        } else {
-            TEST_MEMORY_COMPARE(p, len, e_arg->x, e_arg->len);
-        }
-    }
+    TEST_ASSERT(rsa_test_e(key, bits, e_arg));
 
 exit:
     /*
@@ -9801,7 +9919,88 @@
     psa_destroy_key(key);
     PSA_DONE();
     mbedtls_free(e_read_buffer);
-    mbedtls_free(exported);
+}
+/* END_CASE */
+
+/* BEGIN_CASE */
+void generate_key_ext(int type_arg,
+                      int bits_arg,
+                      int usage_arg,
+                      int alg_arg,
+                      int flags_arg,
+                      data_t *params_data,
+                      int expected_status_arg)
+{
+    mbedtls_svc_key_id_t key = MBEDTLS_SVC_KEY_ID_INIT;
+    psa_key_type_t type = type_arg;
+    psa_key_usage_t usage = usage_arg;
+    size_t bits = bits_arg;
+    psa_algorithm_t alg = alg_arg;
+    psa_status_t expected_status = expected_status_arg;
+    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
+    psa_key_production_parameters_t *params = NULL;
+    size_t params_data_length = 0;
+    psa_key_attributes_t got_attributes = PSA_KEY_ATTRIBUTES_INIT;
+
+    PSA_ASSERT(psa_crypto_init());
+
+    psa_set_key_usage_flags(&attributes, usage);
+    psa_set_key_algorithm(&attributes, alg);
+    psa_set_key_type(&attributes, type);
+    psa_set_key_bits(&attributes, bits);
+
+    if (!setup_key_production_parameters(&params, &params_data_length,
+                                         flags_arg, params_data)) {
+        goto exit;
+    }
+
+    /* Generate a key */
+    psa_status_t status = psa_generate_key_ext(&attributes,
+                                               params, params_data_length,
+                                               &key);
+
+    TEST_EQUAL(status, expected_status);
+    if (expected_status != PSA_SUCCESS) {
+        goto exit;
+    }
+
+    /* Test the key information */
+    PSA_ASSERT(psa_get_key_attributes(key, &got_attributes));
+    TEST_EQUAL(psa_get_key_type(&got_attributes), type);
+    TEST_EQUAL(psa_get_key_bits(&got_attributes), bits);
+
+#if defined(PSA_WANT_KEY_TYPE_RSA_KEY_PAIR_GENERATE)
+    if (type == PSA_KEY_TYPE_RSA_KEY_PAIR) {
+        TEST_ASSERT(rsa_test_e(key, bits, params_data));
+    }
+#endif
+
+    /* Do something with the key according to its type and permitted usage. */
+    if (!mbedtls_test_psa_exercise_key(key, usage, alg)) {
+        goto exit;
+    }
+
+exit:
+    /*
+     * Key attributes may have been returned by psa_get_key_attributes()
+     * thus reset them as required.
+     */
+    psa_reset_key_attributes(&got_attributes);
+    mbedtls_free(params);
+    psa_destroy_key(key);
+    PSA_DONE();
+}
+/* END_CASE */
+
+/* BEGIN_CASE */
+void key_production_parameters_init()
+{
+    psa_key_production_parameters_t init = PSA_KEY_PRODUCTION_PARAMETERS_INIT;
+    psa_key_production_parameters_t zero;
+    memset(&zero, 0, sizeof(zero));
+
+    TEST_EQUAL(init.flags, 0);
+    TEST_EQUAL(zero.flags, 0);
 }
 /* END_CASE */