diff --git a/include/psa/crypto.h b/include/psa/crypto.h
index b62788b..6a7bce8 100644
--- a/include/psa/crypto.h
+++ b/include/psa/crypto.h
@@ -348,6 +348,9 @@
  * with the key in volatile memory. The key slot in persistent storage is
  * not affected and can be opened again later with psa_open_key().
  *
+ * If the key is currently in use in a multipart operation,
+ * the multipart operation is aborted.
+ *
  * \param handle        The key handle to close.
  *
  * \retval #PSA_SUCCESS
@@ -430,6 +433,9 @@
  * This function also erases any metadata such as policies and frees all
  * resources associated with the key.
  *
+ * If the key is currently in use in a multipart operation,
+ * the multipart operation is aborted.
+ *
  * \param handle        Handle to the key slot to erase.
  *
  * \retval #PSA_SUCCESS
@@ -485,6 +491,106 @@
                                      size_t *bits);
 
 /**
+ * \brief Set domain parameters for a key.
+ *
+ * Some key types require additional domain parameters to be set before import
+ * or generation of the key. The domain parameters can be set with this
+ * function or, for key generation, through the \c extra parameter of
+ * psa_generate_key().
+ *
+ * The format for the required domain parameters varies by the key type.
+ * - For DSA public keys (#PSA_KEY_TYPE_DSA_PUBLIC_KEY),
+ *   the `Dss-Parms` format as defined by RFC 3279 &sect;2.3.2.
+ *   ```
+ *   Dss-Parms ::= SEQUENCE  {
+ *      p       INTEGER,
+ *      q       INTEGER,
+ *      g       INTEGER
+ *   }
+ *   ```
+ * - For Diffie-Hellman key exchange keys (#PSA_KEY_TYPE_DH_PUBLIC_KEY), the
+ *   `DomainParameters` format as defined by RFC 3279 &sect;2.3.3.
+ *   ```
+ *   DomainParameters ::= SEQUENCE {
+ *      p               INTEGER,                    -- odd prime, p=jq +1
+ *      g               INTEGER,                    -- generator, g
+ *      q               INTEGER,                    -- factor of p-1
+ *      j               INTEGER OPTIONAL,           -- subgroup factor
+ *      validationParms ValidationParms OPTIONAL
+ *   }
+ *   ValidationParms ::= SEQUENCE {
+ *      seed            BIT STRING,
+ *      pgenCounter     INTEGER
+ *   }
+ *   ```
+ *
+ * \param handle      Handle to the slot where the key will be stored.
+ *                    This must be a valid slot for a key of the chosen
+ *                    type: it must have been obtained by calling
+ *                    psa_allocate_key() or psa_create_key() with the
+ *                    correct \p type and with a maximum size that is
+ *                    compatible with \p data. It must not contain
+ *                    key material yet.
+ * \param type        Key type (a \c PSA_KEY_TYPE_XXX value). When
+ *                    subsequently creating key material into \p handle,
+ *                    the type must be compatible.
+ * \param[in] data    Buffer containing the key domain parameters. The content
+ *                    of this buffer is interpreted according to \p type. of
+ *                    psa_export_key() or psa_export_public_key() for the
+ *                    chosen type.
+ * \param data_length Size of the \p data buffer in bytes.
+ *
+ * \retval #PSA_SUCCESS
+ * \retval #PSA_ERROR_INVALID_HANDLE
+ * \retval #PSA_ERROR_OCCUPIED_SLOT
+ *         There is already a key in the specified slot.
+ * \retval #PSA_ERROR_INVALID_ARGUMENT
+ * \retval #PSA_ERROR_COMMUNICATION_FAILURE
+ * \retval #PSA_ERROR_HARDWARE_FAILURE
+ * \retval #PSA_ERROR_TAMPERING_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 psa_set_key_domain_parameters(psa_key_handle_t handle,
+                                           psa_key_type_t type,
+                                           const uint8_t *data,
+                                           size_t data_length);
+
+/**
+ * \brief Get domain parameters for a key.
+ *
+ * Get the domain parameters for a key with this function, if any. The format
+ * of the domain parameters written to \p data is specified in the
+ * documentation for psa_set_key_domain_parameters().
+ *
+ * \param handle            Handle to the key to get domain parameters from.
+ * \param[out] data         On success, the key domain parameters.
+ * \param data_size         Size of the \p data buffer in bytes.
+ * \param[out] data_length  On success, the number of bytes
+ *                          that make up the key domain parameters data.
+ *
+ * \retval #PSA_SUCCESS
+ * \retval #PSA_ERROR_INVALID_HANDLE
+ * \retval #PSA_ERROR_EMPTY_SLOT
+ *         There is no key in the specified slot.
+ * \retval #PSA_ERROR_INVALID_ARGUMENT
+ * \retval #PSA_ERROR_NOT_SUPPORTED
+ * \retval #PSA_ERROR_COMMUNICATION_FAILURE
+ * \retval #PSA_ERROR_HARDWARE_FAILURE
+ * \retval #PSA_ERROR_TAMPERING_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 psa_get_key_domain_parameters(psa_key_handle_t handle,
+                                           uint8_t *data,
+                                           size_t data_size,
+                                           size_t *data_length);
+
+/**
  * \brief Export a key in binary format.
  *
  * The output of this function can be passed to psa_import_key() to
@@ -519,19 +625,10 @@
  *       coefficient         INTEGER,  -- (inverse of q) mod p
  *   }
  *   ```
- * - For DSA private keys (#PSA_KEY_TYPE_DSA_KEYPAIR), the format
- *   is the non-encrypted DER encoding of the representation used by
- *   OpenSSL and OpenSSH, whose structure is described in ASN.1 as follows:
- *   ```
- *   DSAPrivateKey ::= SEQUENCE {
- *       version             INTEGER,  -- must be 0
- *       prime               INTEGER,  -- p
- *       subprime            INTEGER,  -- q
- *       generator           INTEGER,  -- g
- *       public              INTEGER,  -- y
- *       private             INTEGER,  -- x
- *   }
- *   ```
+ * - For DSA private keys (#PSA_KEY_TYPE_DSA_KEYPAIR), the format is the
+ *   representation of the private key `x` as a big-endian byte string. The
+ *   length of the byte string is the private key size in bytes (leading zeroes
+ *   are not stripped).
  * - For elliptic curve key pairs (key types for which
  *   #PSA_KEY_TYPE_IS_ECC_KEYPAIR is true), the format is
  *   a representation of the private value as a `ceiling(m/8)`-byte string
@@ -543,6 +640,10 @@
  *   and `PSA_ECC_CURVE_BRAINPOOL_PXXX`).
  *   This is the content of the `privateKey` field of the `ECPrivateKey`
  *   format defined by RFC 5915.
+ * - For Diffie-Hellman key exchange key pairs (#PSA_KEY_TYPE_DH_KEYPAIR), the
+ *   format is the representation of the private key `x` as a big-endian byte
+ *   string. The length of the byte string is the private key size in bytes
+ *   (leading zeroes are not stripped).
  * - For public keys (key types for which #PSA_KEY_TYPE_IS_PUBLIC_KEY is
  *   true), the format is the same as for psa_export_public_key().
  *
@@ -599,39 +700,20 @@
  *   ```
  * - For elliptic curve public keys (key types for which
  *   #PSA_KEY_TYPE_IS_ECC_PUBLIC_KEY is true), the format is the uncompressed
- *   representation defined by SEC1 &sect;2.3.3 as the content of an ECPoint:
+ *   representation defined by SEC1 &sect;2.3.3 as the content of an ECPoint.
  *   Let `m` be the bit size associated with the curve, i.e. the bit size of
  *   `q` for a curve over `F_q`. The representation consists of:
  *      - The byte 0x04;
  *      - `x_P` as a `ceiling(m/8)`-byte string, big-endian;
  *      - `y_P` as a `ceiling(m/8)`-byte string, big-endian.
- *
- * For other public key types, the format is the DER representation defined by
- * RFC 5280 as `SubjectPublicKeyInfo`, with the `subjectPublicKey` format
- * specified below.
- * ```
- * SubjectPublicKeyInfo  ::=  SEQUENCE  {
- *      algorithm          AlgorithmIdentifier,
- *      subjectPublicKey   BIT STRING  }
- * AlgorithmIdentifier  ::=  SEQUENCE  {
- *      algorithm          OBJECT IDENTIFIER,
- *      parameters         ANY DEFINED BY algorithm OPTIONAL  }
- * ```
- * - For DSA public keys (#PSA_KEY_TYPE_DSA_PUBLIC_KEY),
- *   the `subjectPublicKey` format is defined by RFC 3279 &sect;2.3.2 as
- *   `DSAPublicKey`,
- *   with the OID `id-dsa`,
- *   and with the parameters `DSS-Parms`.
- *   ```
- *   id-dsa OBJECT IDENTIFIER ::= {
- *      iso(1) member-body(2) us(840) x9-57(10040) x9cm(4) 1 }
- *
- *   Dss-Parms  ::=  SEQUENCE  {
- *      p                  INTEGER,
- *      q                  INTEGER,
- *      g                  INTEGER  }
- *   DSAPublicKey ::= INTEGER -- public key, Y
- *   ```
+ * - For DSA public keys (#PSA_KEY_TYPE_DSA_PUBLIC_KEY), the format is the
+ *   representation of the public key `y = g^x mod p` as a big-endian byte
+ *   string. The length of the byte string is the length of the base prime `p`
+ *   in bytes.
+ * - For Diffie-Hellman key exchange public keys (#PSA_KEY_TYPE_DH_PUBLIC_KEY),
+ *   the format is the representation of the public key `y = g^x mod p` as a
+ *   big-endian byte string. The length of the byte string is the length of the
+ *   base prime `p` in bytes.
  *
  * \param handle            Handle to the key to export.
  * \param[out] data         Buffer where the key data is to be written.
@@ -711,12 +793,12 @@
  * \retval #PSA_SUCCESS
  * \retval #PSA_ERROR_INVALID_HANDLE
  * \retval #PSA_ERROR_ALREADY_EXISTS
- *         \p target already contains key material.
+ *         \p target_handle already contains key material.
  * \retval #PSA_ERROR_DOES_NOT_EXIST
- *         \p source does not contain key material.
+ *         \p source_handle does not contain key material.
  * \retval #PSA_ERROR_INVALID_ARGUMENT
  *         The policy constraints on the source, on the target and
- *         \p constraints are incompatible.
+ *         \p constraint are incompatible.
  * \retval #PSA_ERROR_NOT_PERMITTED
  *         The source key is not exportable and its lifetime does not
  *         allow copying it to the target's lifetime.
@@ -736,6 +818,65 @@
  * @{
  */
 
+/** Calculate the hash (digest) of a message.
+ *
+ * \note To verify the hash of a message against an
+ *       expected value, use psa_hash_compare() instead.
+ *
+ * \param alg               The hash algorithm to compute (\c PSA_ALG_XXX value
+ *                          such that #PSA_ALG_IS_HASH(\p alg) is true).
+ * \param[in] input         Buffer containing the message to hash.
+ * \param input_length      Size of the \p input buffer in bytes.
+ * \param[out] hash         Buffer where the hash is to be written.
+ * \param hash_size         Size of the \p hash buffer in bytes.
+ * \param[out] hash_length  On success, the number of bytes
+ *                          that make up the hash value. This is always
+ *                          #PSA_HASH_SIZE(\p alg).
+ *
+ * \retval #PSA_SUCCESS
+ *         Success.
+ * \retval #PSA_ERROR_NOT_SUPPORTED
+ *         \p alg is not supported or is not a hash algorithm.
+ * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
+ * \retval #PSA_ERROR_COMMUNICATION_FAILURE
+ * \retval #PSA_ERROR_HARDWARE_FAILURE
+ * \retval #PSA_ERROR_TAMPERING_DETECTED
+ */
+psa_status_t psa_hash_compute(psa_algorithm_t alg,
+                              const uint8_t *input,
+                              size_t input_length,
+                              uint8_t *hash,
+                              size_t hash_size,
+                              size_t *hash_length);
+
+/** Calculate the hash (digest) of a message and compare it with a
+ * reference value.
+ *
+ * \param alg               The hash algorithm to compute (\c PSA_ALG_XXX value
+ *                          such that #PSA_ALG_IS_HASH(\p alg) is true).
+ * \param[in] input         Buffer containing the message to hash.
+ * \param input_length      Size of the \p input buffer in bytes.
+ * \param[out] hash         Buffer containing the expected hash value.
+ * \param hash_length       Size of the \p hash buffer in bytes.
+ *
+ * \retval #PSA_SUCCESS
+ *         The expected hash is identical to the actual hash of the input.
+ * \retval #PSA_ERROR_INVALID_SIGNATURE
+ *         The hash of the message was calculated successfully, but it
+ *         differs from the expected hash.
+ * \retval #PSA_ERROR_NOT_SUPPORTED
+ *         \p alg is not supported or is not a hash algorithm.
+ * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
+ * \retval #PSA_ERROR_COMMUNICATION_FAILURE
+ * \retval #PSA_ERROR_HARDWARE_FAILURE
+ * \retval #PSA_ERROR_TAMPERING_DETECTED
+ */
+psa_status_t psa_hash_compare(psa_algorithm_t alg,
+                              const uint8_t *input,
+                              size_t input_length,
+                              const uint8_t *hash,
+                              const size_t hash_length);
+
 /** The type of the state data structure for multipart hash operations.
  *
  * Before calling any function on a hash operation object, the application must
@@ -989,6 +1130,86 @@
  * @{
  */
 
+/** Calculate the MAC (message authentication code) of a message.
+ *
+ * \note To verify the MAC of a message against an
+ *       expected value, use psa_mac_verify() instead.
+ *       Beware that comparing integrity or authenticity data such as
+ *       MAC values with a function such as \c memcmp is risky
+ *       because the time taken by the comparison may leak information
+ *       about the MAC value which could allow an attacker to guess
+ *       a valid MAC and thereby bypass security controls.
+ *
+ * \param handle            Handle to the key to use for the operation.
+ * \param alg               The MAC algorithm to compute (\c PSA_ALG_XXX value
+ *                          such that #PSA_ALG_IS_MAC(\p alg) is true).
+ * \param[in] input         Buffer containing the input message.
+ * \param input_length      Size of the \p input buffer in bytes.
+ * \param[out] mac          Buffer where the MAC value is to be written.
+ * \param mac_size          Size of the \p mac buffer in bytes.
+ * \param[out] mac_length   On success, the number of bytes
+ *                          that make up the MAC value.
+ *
+ * \retval #PSA_SUCCESS
+ *         Success.
+ * \retval #PSA_ERROR_INVALID_HANDLE
+ * \retval #PSA_ERROR_EMPTY_SLOT
+ * \retval #PSA_ERROR_NOT_PERMITTED
+ * \retval #PSA_ERROR_INVALID_ARGUMENT
+ *         \p handle is not compatible with \p alg.
+ * \retval #PSA_ERROR_NOT_SUPPORTED
+ *         \p alg is not supported or is not a MAC algorithm.
+ * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
+ * \retval #PSA_ERROR_COMMUNICATION_FAILURE
+ * \retval #PSA_ERROR_HARDWARE_FAILURE
+ * \retval #PSA_ERROR_TAMPERING_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 psa_mac_compute(psa_key_handle_t handle,
+                             psa_algorithm_t alg,
+                             const uint8_t *input,
+                             size_t input_length,
+                             uint8_t *mac,
+                             size_t mac_size,
+                             size_t *mac_length);
+
+/** Calculate the MAC of a message and compare it with a reference value.
+ *
+ * \param handle            Handle to the key to use for the operation.
+ * \param alg               The MAC algorithm to compute (\c PSA_ALG_XXX value
+ *                          such that #PSA_ALG_IS_MAC(\p alg) is true).
+ * \param[in] input         Buffer containing the input message.
+ * \param input_length      Size of the \p input buffer in bytes.
+ * \param[out] mac          Buffer containing the expected MAC value.
+ * \param mac_length        Size of the \p mac buffer in bytes.
+ *
+ * \retval #PSA_SUCCESS
+ *         The expected MAC is identical to the actual MAC of the input.
+ * \retval #PSA_ERROR_INVALID_SIGNATURE
+ *         The MAC of the message was calculated successfully, but it
+ *         differs from the expected value.
+ * \retval #PSA_ERROR_INVALID_HANDLE
+ * \retval #PSA_ERROR_EMPTY_SLOT
+ * \retval #PSA_ERROR_NOT_PERMITTED
+ * \retval #PSA_ERROR_INVALID_ARGUMENT
+ *         \p handle is not compatible with \p alg.
+ * \retval #PSA_ERROR_NOT_SUPPORTED
+ *         \p alg is not supported or is not a MAC algorithm.
+ * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
+ * \retval #PSA_ERROR_COMMUNICATION_FAILURE
+ * \retval #PSA_ERROR_HARDWARE_FAILURE
+ * \retval #PSA_ERROR_TAMPERING_DETECTED
+ */
+psa_status_t psa_mac_verify(psa_key_handle_t handle,
+                            psa_algorithm_t alg,
+                            const uint8_t *input,
+                            size_t input_length,
+                            const uint8_t *mac,
+                            const size_t mac_length);
+
 /** The type of the state data structure for multipart MAC operations.
  *
  * Before calling any function on a MAC operation object, the application must
@@ -1048,8 +1269,6 @@
  * -# Initialize the operation object with one of the methods described in the
  *    documentation for #psa_mac_operation_t, e.g. PSA_MAC_OPERATION_INIT.
  * -# Call psa_mac_sign_setup() to specify the algorithm and key.
- *    The key remains associated with the operation even if the content
- *    of the key slot changes.
  * -# Call psa_mac_update() zero, one or more times, passing a fragment
  *    of the message each time. The MAC that is calculated is the MAC
  *    of the concatenation of these messages in order.
@@ -1068,8 +1287,10 @@
  *                          been initialized as per the documentation for
  *                          #psa_mac_operation_t and not yet in use.
  * \param handle            Handle to the key to use for the operation.
+ *                          It must remain valid until the operation
+ *                          terminates.
  * \param alg               The MAC algorithm to compute (\c PSA_ALG_XXX value
- *                          such that #PSA_ALG_IS_MAC(alg) is true).
+ *                          such that #PSA_ALG_IS_MAC(\p alg) is true).
  *
  * \retval #PSA_SUCCESS
  *         Success.
@@ -1077,7 +1298,7 @@
  * \retval #PSA_ERROR_DOES_NOT_EXIST
  * \retval #PSA_ERROR_NOT_PERMITTED
  * \retval #PSA_ERROR_INVALID_ARGUMENT
- *         \p key is not compatible with \p alg.
+ *         \p handle is not compatible with \p alg.
  * \retval #PSA_ERROR_NOT_SUPPORTED
  *         \p alg is not supported or is not a MAC algorithm.
  * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
@@ -1107,8 +1328,6 @@
  * -# Initialize the operation object with one of the methods described in the
  *    documentation for #psa_mac_operation_t, e.g. PSA_MAC_OPERATION_INIT.
  * -# Call psa_mac_verify_setup() to specify the algorithm and key.
- *    The key remains associated with the operation even if the content
- *    of the key slot changes.
  * -# Call psa_mac_update() zero, one or more times, passing a fragment
  *    of the message each time. The MAC that is calculated is the MAC
  *    of the concatenation of these messages in order.
@@ -1128,6 +1347,8 @@
  *                          been initialized as per the documentation for
  *                          #psa_mac_operation_t and not yet in use.
  * \param handle            Handle to the key to use for the operation.
+ *                          It must remain valid until the operation
+ *                          terminates.
  * \param alg               The MAC algorithm to compute (\c PSA_ALG_XXX value
  *                          such that #PSA_ALG_IS_MAC(\p alg) is true).
  *
@@ -1295,6 +1516,91 @@
  * @{
  */
 
+/** Encrypt a message using a symmetric cipher.
+ *
+ * This function encrypts a message with a random IV (initialization
+ * vector).
+ *
+ * \param handle                Handle to the key to use for the operation.
+ *                              It must remain valid until the operation
+ *                              terminates.
+ * \param alg                   The cipher algorithm to compute
+ *                              (\c PSA_ALG_XXX value such that
+ *                              #PSA_ALG_IS_CIPHER(\p alg) is true).
+ * \param[in] input             Buffer containing the message to encrypt.
+ * \param input_length          Size of the \p input buffer in bytes.
+ * \param[out] output           Buffer where the output is to be written.
+ *                              The output contains the IV followed by
+ *                              the ciphertext proper.
+ * \param output_size           Size of the \p output buffer in bytes.
+ * \param[out] output_length    On success, the number of bytes
+ *                              that make up the output.
+ *
+ * \retval #PSA_SUCCESS
+ *         Success.
+ * \retval #PSA_ERROR_INVALID_HANDLE
+ * \retval #PSA_ERROR_EMPTY_SLOT
+ * \retval #PSA_ERROR_NOT_PERMITTED
+ * \retval #PSA_ERROR_INVALID_ARGUMENT
+ *         \p handle is not compatible with \p alg.
+ * \retval #PSA_ERROR_NOT_SUPPORTED
+ *         \p alg is not supported or is not a cipher algorithm.
+ * \retval #PSA_ERROR_BUFFER_TOO_SMALL
+ * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
+ * \retval #PSA_ERROR_COMMUNICATION_FAILURE
+ * \retval #PSA_ERROR_HARDWARE_FAILURE
+ * \retval #PSA_ERROR_TAMPERING_DETECTED
+ */
+psa_status_t psa_cipher_encrypt(psa_key_handle_t handle,
+                                psa_algorithm_t alg,
+                                const uint8_t *input,
+                                size_t input_length,
+                                uint8_t *output,
+                                size_t output_size,
+                                size_t *output_length);
+
+/** Decrypt a message using a symmetric cipher.
+ *
+ * This function decrypts a message encrypted with a symmetric cipher.
+ *
+ * \param handle                Handle to the key to use for the operation.
+ *                              It must remain valid until the operation
+ *                              terminates.
+ * \param alg                   The cipher algorithm to compute
+ *                              (\c PSA_ALG_XXX value such that
+ *                              #PSA_ALG_IS_CIPHER(\p alg) is true).
+ * \param[in] input             Buffer containing the message to decrypt.
+ *                              This consists of the IV followed by the
+ *                              ciphertext proper.
+ * \param input_length          Size of the \p input buffer in bytes.
+ * \param[out] output           Buffer where the plaintext is to be written.
+ * \param output_size           Size of the \p output buffer in bytes.
+ * \param[out] output_length    On success, the number of bytes
+ *                              that make up the output.
+ *
+ * \retval #PSA_SUCCESS
+ *         Success.
+ * \retval #PSA_ERROR_INVALID_HANDLE
+ * \retval #PSA_ERROR_EMPTY_SLOT
+ * \retval #PSA_ERROR_NOT_PERMITTED
+ * \retval #PSA_ERROR_INVALID_ARGUMENT
+ *         \p handle is not compatible with \p alg.
+ * \retval #PSA_ERROR_NOT_SUPPORTED
+ *         \p alg is not supported or is not a cipher algorithm.
+ * \retval #PSA_ERROR_BUFFER_TOO_SMALL
+ * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
+ * \retval #PSA_ERROR_COMMUNICATION_FAILURE
+ * \retval #PSA_ERROR_HARDWARE_FAILURE
+ * \retval #PSA_ERROR_TAMPERING_DETECTED
+ */
+psa_status_t psa_cipher_decrypt(psa_key_handle_t handle,
+                                psa_algorithm_t alg,
+                                const uint8_t *input,
+                                size_t input_length,
+                                uint8_t *output,
+                                size_t output_size,
+                                size_t *output_length);
+
 /** The type of the state data structure for multipart cipher operations.
  *
  * Before calling any function on a cipher operation object, the application
@@ -1351,8 +1657,6 @@
  *    documentation for #psa_cipher_operation_t, e.g.
  *    PSA_CIPHER_OPERATION_INIT.
  * -# Call psa_cipher_encrypt_setup() to specify the algorithm and key.
- *    The key remains associated with the operation even if the content
- *    of the key slot changes.
  * -# Call either psa_cipher_generate_iv() or psa_cipher_set_iv() to
  *    generate or set the IV (initialization vector). You should use
  *    psa_cipher_generate_iv() unless the protocol you are implementing
@@ -1367,14 +1671,15 @@
  * After a successful call to psa_cipher_encrypt_setup(), the application must
  * eventually terminate the operation. The following events terminate an
  * operation:
- * - A failed call to psa_cipher_generate_iv(), psa_cipher_set_iv()
- *   or psa_cipher_update().
+ * - A failed call to any of the \c psa_cipher_xxx functions.
  * - A call to psa_cipher_finish() or psa_cipher_abort().
  *
  * \param[in,out] operation     The operation object to set up. It must have
  *                              been initialized as per the documentation for
  *                              #psa_cipher_operation_t and not yet in use.
  * \param handle                Handle to the key to use for the operation.
+ *                              It must remain valid until the operation
+ *                              terminates.
  * \param alg                   The cipher algorithm to compute
  *                              (\c PSA_ALG_XXX value such that
  *                              #PSA_ALG_IS_CIPHER(\p alg) is true).
@@ -1385,7 +1690,7 @@
  * \retval #PSA_ERROR_DOES_NOT_EXIST
  * \retval #PSA_ERROR_NOT_PERMITTED
  * \retval #PSA_ERROR_INVALID_ARGUMENT
- *         \p key is not compatible with \p alg.
+ *         \p handle is not compatible with \p alg.
  * \retval #PSA_ERROR_NOT_SUPPORTED
  *         \p alg is not supported or is not a cipher algorithm.
  * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
@@ -1414,9 +1719,7 @@
  *    documentation for #psa_cipher_operation_t, e.g.
  *    PSA_CIPHER_OPERATION_INIT.
  * -# Call psa_cipher_decrypt_setup() to specify the algorithm and key.
- *    The key remains associated with the operation even if the content
- *    of the key slot changes.
- * -# Call psa_cipher_update() with the IV (initialization vector) for the
+ * -# Call psa_cipher_set_iv() with the IV (initialization vector) for the
  *    decryption. If the IV is prepended to the ciphertext, you can call
  *    psa_cipher_update() on a buffer containing the IV followed by the
  *    beginning of the message.
@@ -1430,13 +1733,15 @@
  * After a successful call to psa_cipher_decrypt_setup(), the application must
  * eventually terminate the operation. The following events terminate an
  * operation:
- * - A failed call to psa_cipher_update().
+ * - A failed call to any of the \c psa_cipher_xxx functions.
  * - A call to psa_cipher_finish() or psa_cipher_abort().
  *
  * \param[in,out] operation     The operation object to set up. It must have
  *                              been initialized as per the documentation for
  *                              #psa_cipher_operation_t and not yet in use.
  * \param handle                Handle to the key to use for the operation.
+ *                              It must remain valid until the operation
+ *                              terminates.
  * \param alg                   The cipher algorithm to compute
  *                              (\c PSA_ALG_XXX value such that
  *                              #PSA_ALG_IS_CIPHER(\p alg) is true).
@@ -1447,7 +1752,7 @@
  * \retval #PSA_ERROR_DOES_NOT_EXIST
  * \retval #PSA_ERROR_NOT_PERMITTED
  * \retval #PSA_ERROR_INVALID_ARGUMENT
- *         \p key is not compatible with \p alg.
+ *         \p handle is not compatible with \p alg.
  * \retval #PSA_ERROR_NOT_SUPPORTED
  *         \p alg is not supported or is not a cipher algorithm.
  * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
@@ -1501,7 +1806,7 @@
 
 /** Set the IV for a symmetric encryption or decryption operation.
  *
- * This function sets the random IV (initialization vector), nonce
+ * This function sets the IV (initialization vector), nonce
  * or initial counter value for the encryption or decryption operation.
  *
  * The application must call psa_cipher_encrypt_setup() before
@@ -1670,7 +1975,7 @@
  *                                #PSA_AEAD_ENCRYPT_OUTPUT_SIZE(\p alg,
  *                                \p plaintext_length).
  * \param[out] ciphertext_length  On success, the size of the output
- *                                in the \b ciphertext buffer.
+ *                                in the \p ciphertext buffer.
  *
  * \retval #PSA_SUCCESS
  *         Success.
@@ -1678,7 +1983,7 @@
  * \retval #PSA_ERROR_DOES_NOT_EXIST
  * \retval #PSA_ERROR_NOT_PERMITTED
  * \retval #PSA_ERROR_INVALID_ARGUMENT
- *         \p key is not compatible with \p alg.
+ *         \p handle is not compatible with \p alg.
  * \retval #PSA_ERROR_NOT_SUPPORTED
  *         \p alg is not supported or is not an AEAD algorithm.
  * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
@@ -1726,7 +2031,7 @@
  *                                #PSA_AEAD_DECRYPT_OUTPUT_SIZE(\p alg,
  *                                \p ciphertext_length).
  * \param[out] plaintext_length   On success, the size of the output
- *                                in the \b plaintext buffer.
+ *                                in the \p plaintext buffer.
  *
  * \retval #PSA_SUCCESS
  *         Success.
@@ -1736,7 +2041,7 @@
  *         The ciphertext is not authentic.
  * \retval #PSA_ERROR_NOT_PERMITTED
  * \retval #PSA_ERROR_INVALID_ARGUMENT
- *         \p key is not compatible with \p alg.
+ *         \p handle is not compatible with \p alg.
  * \retval #PSA_ERROR_NOT_SUPPORTED
  *         \p alg is not supported or is not an AEAD algorithm.
  * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
@@ -1760,6 +2065,513 @@
                               size_t plaintext_size,
                               size_t *plaintext_length);
 
+/** The type of the state data structure for multipart AEAD operations.
+ *
+ * Before calling any function on an AEAD operation object, the application
+ * must initialize it by any of the following means:
+ * - Set the structure to all-bits-zero, for example:
+ *   \code
+ *   psa_aead_operation_t operation;
+ *   memset(&operation, 0, sizeof(operation));
+ *   \endcode
+ * - Initialize the structure to logical zero values, for example:
+ *   \code
+ *   psa_aead_operation_t operation = {0};
+ *   \endcode
+ * - Initialize the structure to the initializer #PSA_AEAD_OPERATION_INIT,
+ *   for example:
+ *   \code
+ *   psa_aead_operation_t operation = PSA_AEAD_OPERATION_INIT;
+ *   \endcode
+ * - Assign the result of the function psa_aead_operation_init()
+ *   to the structure, for example:
+ *   \code
+ *   psa_aead_operation_t operation;
+ *   operation = psa_aead_operation_init();
+ *   \endcode
+ *
+ * This is an implementation-defined \c struct. Applications should not
+ * make any assumptions about the content of this structure except
+ * as directed by the documentation of a specific implementation. */
+typedef struct psa_aead_operation_s psa_aead_operation_t;
+
+/** \def PSA_AEAD_OPERATION_INIT
+ *
+ * This macro returns a suitable initializer for an AEAD operation object of
+ * type #psa_aead_operation_t.
+ */
+#ifdef __DOXYGEN_ONLY__
+/* This is an example definition for documentation purposes.
+ * Implementations should define a suitable value in `crypto_struct.h`.
+ */
+#define PSA_AEAD_OPERATION_INIT {0}
+#endif
+
+/** Return an initial value for an AEAD operation object.
+ */
+static psa_aead_operation_t psa_aead_operation_init(void);
+
+/** Set the key for a multipart authenticated encryption operation.
+ *
+ * The sequence of operations to encrypt a message with authentication
+ * is as follows:
+ * -# Allocate an operation object which will be passed to all the functions
+ *    listed here.
+ * -# Initialize the operation object with one of the methods described in the
+ *    documentation for #psa_aead_operation_t, e.g.
+ *    PSA_AEAD_OPERATION_INIT.
+ * -# Call psa_aead_encrypt_setup() to specify the algorithm and key.
+ * -# If needed, call psa_aead_set_lengths() to specify the length of the
+ *    inputs to the subsequent calls to psa_aead_update_ad() and
+ *    psa_aead_update(). See the documentation of psa_aead_set_lengths()
+ *    for details.
+ * -# Call either psa_aead_generate_nonce() or psa_aead_set_nonce() to
+ *    generate or set the nonce. You should use
+ *    psa_aead_generate_nonce() unless the protocol you are implementing
+ *    requires a specific nonce value.
+ * -# Call psa_aead_update_ad() zero, one or more times, passing a fragment
+ *    of the non-encrypted additional authenticated data each time.
+ * -# Call psa_aead_update() zero, one or more times, passing a fragment
+ *    of the message to encrypt each time.
+ * -# Call psa_aead_finish().
+ *
+ * The application may call psa_aead_abort() at any time after the operation
+ * has been initialized.
+ *
+ * After a successful call to psa_aead_encrypt_setup(), the application must
+ * eventually terminate the operation. The following events terminate an
+ * operation:
+ * - A failed call to any of the \c psa_aead_xxx functions.
+ * - A call to psa_aead_finish(), psa_aead_verify() or psa_aead_abort().
+ *
+ * \param[in,out] operation     The operation object to set up. It must have
+ *                              been initialized as per the documentation for
+ *                              #psa_aead_operation_t and not yet in use.
+ * \param handle                Handle to the key to use for the operation.
+ *                              It must remain valid until the operation
+ *                              terminates.
+ * \param alg                   The AEAD algorithm to compute
+ *                              (\c PSA_ALG_XXX value such that
+ *                              #PSA_ALG_IS_AEAD(\p alg) is true).
+ *
+ * \retval #PSA_SUCCESS
+ *         Success.
+ * \retval #PSA_ERROR_INVALID_HANDLE
+ * \retval #PSA_ERROR_EMPTY_SLOT
+ * \retval #PSA_ERROR_NOT_PERMITTED
+ * \retval #PSA_ERROR_INVALID_ARGUMENT
+ *         \p handle is not compatible with \p alg.
+ * \retval #PSA_ERROR_NOT_SUPPORTED
+ *         \p alg is not supported or is not an AEAD algorithm.
+ * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
+ * \retval #PSA_ERROR_COMMUNICATION_FAILURE
+ * \retval #PSA_ERROR_HARDWARE_FAILURE
+ * \retval #PSA_ERROR_TAMPERING_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 psa_aead_encrypt_setup(psa_aead_operation_t *operation,
+                                    psa_key_handle_t handle,
+                                    psa_algorithm_t alg);
+
+/** Set the key for a multipart authenticated decryption operation.
+ *
+ * The sequence of operations to decrypt a message with authentication
+ * is as follows:
+ * -# Allocate an operation object which will be passed to all the functions
+ *    listed here.
+ * -# Initialize the operation object with one of the methods described in the
+ *    documentation for #psa_aead_operation_t, e.g.
+ *    PSA_AEAD_OPERATION_INIT.
+ * -# Call psa_aead_decrypt_setup() to specify the algorithm and key.
+ * -# If needed, call psa_aead_set_lengths() to specify the length of the
+ *    inputs to the subsequent calls to psa_aead_update_ad() and
+ *    psa_aead_update(). See the documentation of psa_aead_set_lengths()
+ *    for details.
+ * -# Call psa_aead_set_nonce() with the nonce for the decryption.
+ * -# Call psa_aead_update_ad() zero, one or more times, passing a fragment
+ *    of the non-encrypted additional authenticated data each time.
+ * -# Call psa_aead_update() zero, one or more times, passing a fragment
+ *    of the ciphertext to decrypt each time.
+ * -# Call psa_aead_verify().
+ *
+ * The application may call psa_aead_abort() at any time after the operation
+ * has been initialized.
+ *
+ * After a successful call to psa_aead_decrypt_setup(), the application must
+ * eventually terminate the operation. The following events terminate an
+ * operation:
+ * - A failed call to any of the \c psa_aead_xxx functions.
+ * - A call to psa_aead_finish(), psa_aead_verify() or psa_aead_abort().
+ *
+ * \param[in,out] operation     The operation object to set up. It must have
+ *                              been initialized as per the documentation for
+ *                              #psa_aead_operation_t and not yet in use.
+ * \param handle                Handle to the key to use for the operation.
+ *                              It must remain valid until the operation
+ *                              terminates.
+ * \param alg                   The AEAD algorithm to compute
+ *                              (\c PSA_ALG_XXX value such that
+ *                              #PSA_ALG_IS_AEAD(\p alg) is true).
+ *
+ * \retval #PSA_SUCCESS
+ *         Success.
+ * \retval #PSA_ERROR_INVALID_HANDLE
+ * \retval #PSA_ERROR_EMPTY_SLOT
+ * \retval #PSA_ERROR_NOT_PERMITTED
+ * \retval #PSA_ERROR_INVALID_ARGUMENT
+ *         \p handle is not compatible with \p alg.
+ * \retval #PSA_ERROR_NOT_SUPPORTED
+ *         \p alg is not supported or is not an AEAD algorithm.
+ * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
+ * \retval #PSA_ERROR_COMMUNICATION_FAILURE
+ * \retval #PSA_ERROR_HARDWARE_FAILURE
+ * \retval #PSA_ERROR_TAMPERING_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 psa_aead_decrypt_setup(psa_aead_operation_t *operation,
+                                    psa_key_handle_t handle,
+                                    psa_algorithm_t alg);
+
+/** Generate a random nonce for an authenticated encryption operation.
+ *
+ * This function generates a random nonce for the authenticated encryption
+ * operation with an appropriate size for the chosen algorithm, key type
+ * and key size.
+ *
+ * The application must call psa_aead_encrypt_setup() before
+ * calling this function.
+ *
+ * If this function returns an error status, the operation becomes inactive.
+ *
+ * \param[in,out] operation     Active AEAD operation.
+ * \param[out] nonce            Buffer where the generated nonce is to be
+ *                              written.
+ * \param nonce_size            Size of the \p nonce buffer in bytes.
+ * \param[out] nonce_length     On success, the number of bytes of the
+ *                              generated nonce.
+ *
+ * \retval #PSA_SUCCESS
+ *         Success.
+ * \retval #PSA_ERROR_BAD_STATE
+ *         The operation state is not valid (not set up, or nonce already set).
+ * \retval #PSA_ERROR_BUFFER_TOO_SMALL
+ *         The size of the \p nonce buffer is too small.
+ * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
+ * \retval #PSA_ERROR_COMMUNICATION_FAILURE
+ * \retval #PSA_ERROR_HARDWARE_FAILURE
+ * \retval #PSA_ERROR_TAMPERING_DETECTED
+ */
+psa_status_t psa_aead_generate_nonce(psa_aead_operation_t *operation,
+                                     unsigned char *nonce,
+                                     size_t nonce_size,
+                                     size_t *nonce_length);
+
+/** Set the nonce for an authenticated encryption or decryption operation.
+ *
+ * This function sets the nonce for the authenticated
+ * encryption or decryption operation.
+ *
+ * The application must call psa_aead_encrypt_setup() before
+ * calling this function.
+ *
+ * If this function returns an error status, the operation becomes inactive.
+ *
+ * \note When encrypting, applications should use psa_aead_generate_nonce()
+ * instead of this function, unless implementing a protocol that requires
+ * a non-random IV.
+ *
+ * \param[in,out] operation     Active AEAD operation.
+ * \param[in] nonce             Buffer containing the nonce to use.
+ * \param nonce_length          Size of the nonce in bytes.
+ *
+ * \retval #PSA_SUCCESS
+ *         Success.
+ * \retval #PSA_ERROR_BAD_STATE
+ *         The operation state is not valid (not set up, or nonce already set).
+ * \retval #PSA_ERROR_INVALID_ARGUMENT
+ *         The size of \p nonce is not acceptable for the chosen algorithm.
+ * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
+ * \retval #PSA_ERROR_COMMUNICATION_FAILURE
+ * \retval #PSA_ERROR_HARDWARE_FAILURE
+ * \retval #PSA_ERROR_TAMPERING_DETECTED
+ */
+psa_status_t psa_aead_set_nonce(psa_aead_operation_t *operation,
+                                const unsigned char *nonce,
+                                size_t nonce_length);
+
+/** Declare the lengths of the message and additional data for AEAD.
+ *
+ * The application must call this function before calling
+ * psa_aead_update_ad() or psa_aead_update() if the algorithm for
+ * the operation requires it. If the algorithm does not require it,
+ * calling this function is optional, but if this function is called
+ * then the implementation must enforce the lengths.
+ *
+ * You may call this function before or after setting the nonce with
+ * psa_aead_set_nonce() or psa_aead_generate_nonce().
+ *
+ * - For #PSA_ALG_CCM, calling this function is required.
+ * - For the other AEAD algorithms defined in this specification, calling
+ *   this function is not required.
+ * - For vendor-defined algorithm, refer to the vendor documentation.
+ *
+ * \param[in,out] operation     Active AEAD operation.
+ * \param ad_length             Size of the non-encrypted additional
+ *                              authenticated data in bytes.
+ * \param plaintext_length      Size of the plaintext to encrypt in bytes.
+ *
+ * \retval #PSA_SUCCESS
+ *         Success.
+ * \retval #PSA_ERROR_BAD_STATE
+ *         The operation state is not valid (not set up, already completed,
+ *         or psa_aead_update_ad() or psa_aead_update() already called).
+ * \retval #PSA_ERROR_INVALID_ARGUMENT
+ *         At least one of the lengths is not acceptable for the chosen
+ *         algorithm.
+ * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
+ * \retval #PSA_ERROR_COMMUNICATION_FAILURE
+ * \retval #PSA_ERROR_HARDWARE_FAILURE
+ * \retval #PSA_ERROR_TAMPERING_DETECTED
+ */
+psa_status_t psa_aead_set_lengths(psa_aead_operation_t *operation,
+                                  size_t ad_length,
+                                  size_t plaintext_length);
+
+/** Pass additional data to an active AEAD operation.
+ *
+ * Additional data is authenticated, but not encrypted.
+ *
+ * You may call this function multiple times to pass successive fragments
+ * of the additional data. You may not call this function after passing
+ * data to encrypt or decrypt with psa_aead_update().
+ *
+ * Before calling this function, you must:
+ * 1. Call either psa_aead_encrypt_setup() or psa_aead_decrypt_setup().
+ * 2. Set the nonce with psa_aead_generate_nonce() or psa_aead_set_nonce().
+ *
+ * If this function returns an error status, the operation becomes inactive.
+ *
+ * \warning When decrypting, until psa_aead_verify() has returned #PSA_SUCCESS,
+ *          there is no guarantee that the input is valid. Therefore, until
+ *          you have called psa_aead_verify() and it has returned #PSA_SUCCESS,
+ *          treat the input as untrusted and prepare to undo any action that
+ *          depends on the input if psa_aead_verify() returns an error status.
+ *
+ * \param[in,out] operation     Active AEAD operation.
+ * \param[in] input             Buffer containing the fragment of
+ *                              additional data.
+ * \param input_length          Size of the \p input buffer in bytes.
+ *
+ * \retval #PSA_SUCCESS
+ *         Success.
+ * \retval #PSA_ERROR_BAD_STATE
+ *         The operation state is not valid (not set up, nonce not set,
+ *         psa_aead_update() already called, or operation already completed).
+ * \retval #PSA_ERROR_INVALID_ARGUMENT
+ *         The total input length overflows the additional data length that
+ *         was previously specified with psa_aead_set_lengths().
+ * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
+ * \retval #PSA_ERROR_COMMUNICATION_FAILURE
+ * \retval #PSA_ERROR_HARDWARE_FAILURE
+ * \retval #PSA_ERROR_TAMPERING_DETECTED
+ */
+psa_status_t psa_aead_update_ad(psa_aead_operation_t *operation,
+                                const uint8_t *input,
+                                size_t input_length);
+
+/** Encrypt or decrypt a message fragment in an active AEAD operation.
+ *
+ * Before calling this function, you must:
+ * 1. Call either psa_aead_encrypt_setup() or psa_aead_decrypt_setup().
+ *    The choice of setup function determines whether this function
+ *    encrypts or decrypts its input.
+ * 2. Set the nonce with psa_aead_generate_nonce() or psa_aead_set_nonce().
+ * 3. Call psa_aead_update_ad() to pass all the additional data.
+ *
+ * If this function returns an error status, the operation becomes inactive.
+ *
+ * \warning When decrypting, until psa_aead_verify() has returned #PSA_SUCCESS,
+ *          there is no guarantee that the input is valid. Therefore, until
+ *          you have called psa_aead_verify() and it has returned #PSA_SUCCESS:
+ *          - Do not use the output in any way other than storing it in a
+ *            confidential location. If you take any action that depends
+ *            on the tentative decrypted data, this action will need to be
+ *            undone if the input turns out not to be valid. Furthermore,
+ *            if an adversary can observe that this action took place
+ *            (for example through timing), they may be able to use this
+ *            fact as an oracle to decrypt any message encrypted with the
+ *            same key.
+ *          - In particular, do not copy the output anywhere but to a
+ *            memory or storage space that you have exclusive access to.
+ *
+ * \param[in,out] operation     Active AEAD operation.
+ * \param[in] input             Buffer containing the message fragment to
+ *                              encrypt or decrypt.
+ * \param input_length          Size of the \p input buffer in bytes.
+ * \param[out] output           Buffer where the output is to be written.
+ * \param output_size           Size of the \p output buffer in bytes.
+ * \param[out] output_length    On success, the number of bytes
+ *                              that make up the returned output.
+ *
+ * \retval #PSA_SUCCESS
+ *         Success.
+ * \retval #PSA_ERROR_BAD_STATE
+ *         The operation state is not valid (not set up, nonce not set
+ *         or already completed).
+ * \retval #PSA_ERROR_BUFFER_TOO_SMALL
+ *         The size of the \p output buffer is too small.
+ * \retval #PSA_ERROR_INVALID_ARGUMENT
+ *         The total length of input to psa_aead_update_ad() so far is
+ *         less than the additional data length that was previously
+ *         specified with psa_aead_set_lengths().
+ * \retval #PSA_ERROR_INVALID_ARGUMENT
+ *         The total input length overflows the plaintext length that
+ *         was previously specified with psa_aead_set_lengths().
+ * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
+ * \retval #PSA_ERROR_COMMUNICATION_FAILURE
+ * \retval #PSA_ERROR_HARDWARE_FAILURE
+ * \retval #PSA_ERROR_TAMPERING_DETECTED
+ */
+psa_status_t psa_aead_update(psa_aead_operation_t *operation,
+                             const uint8_t *input,
+                             size_t input_length,
+                             unsigned char *output,
+                             size_t output_size,
+                             size_t *output_length);
+
+/** Finish encrypting a message in an AEAD operation.
+ *
+ * The operation must have been set up with psa_aead_encrypt_setup().
+ *
+ * This function finishes the authentication of the additional data
+ * formed by concatenating the inputs passed to preceding calls to
+ * psa_aead_update_ad() with the plaintext formed by concatenating the
+ * inputs passed to preceding calls to psa_aead_update().
+ *
+ * This function has two output buffers:
+ * - \p ciphertext contains trailing ciphertext that was buffered from
+ *   preceding calls to psa_aead_update(). For all standard AEAD algorithms,
+ *   psa_aead_update() does not buffer any output and therefore \p ciphertext
+ *   will not contain any output and can be a 0-sized buffer.
+ * - \p tag contains the authentication tag. Its length is always
+ *   #PSA_AEAD_TAG_LENGTH(\c alg) where \c alg is the AEAD algorithm
+ *   that the operation performs.
+ *
+ * When this function returns, the operation becomes inactive.
+ *
+ * \param[in,out] operation     Active AEAD operation.
+ * \param[out] ciphertext       Buffer where the last part of the ciphertext
+ *                              is to be written.
+ * \param ciphertext_size       Size of the \p ciphertext buffer in bytes.
+ * \param[out] ciphertext_length On success, the number of bytes of
+ *                              returned ciphertext.
+ * \param[out] tag              Buffer where the authentication tag is
+ *                              to be written.
+ * \param tag_size              Size of the \p tag buffer in bytes.
+ * \param[out] tag_length       On success, the number of bytes
+ *                              that make up the returned tag.
+ *
+ * \retval #PSA_SUCCESS
+ *         Success.
+ * \retval #PSA_ERROR_BAD_STATE
+ *         The operation state is not valid (not set up, nonce not set,
+ *         decryption, or already completed).
+ * \retval #PSA_ERROR_BUFFER_TOO_SMALL
+ *         The size of the \p ciphertext or \p tag buffer is too small.
+ * \retval #PSA_ERROR_INVALID_ARGUMENT
+ *         The total length of input to psa_aead_update_ad() so far is
+ *         less than the additional data length that was previously
+ *         specified with psa_aead_set_lengths().
+ * \retval #PSA_ERROR_INVALID_ARGUMENT
+ *         The total length of input to psa_aead_update() so far is
+ *         less than the plaintext length that was previously
+ *         specified with psa_aead_set_lengths().
+ * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
+ * \retval #PSA_ERROR_COMMUNICATION_FAILURE
+ * \retval #PSA_ERROR_HARDWARE_FAILURE
+ * \retval #PSA_ERROR_TAMPERING_DETECTED
+ */
+psa_status_t psa_aead_finish(psa_aead_operation_t *operation,
+                             uint8_t *ciphertext,
+                             size_t ciphertext_size,
+                             size_t *ciphertext_length,
+                             uint8_t *tag,
+                             size_t tag_size,
+                             size_t *tag_length);
+
+/** Finish authenticating and decrypting a message in an AEAD operation.
+ *
+ * The operation must have been set up with psa_aead_decrypt_setup().
+ *
+ * This function finishes the authentication of the additional data
+ * formed by concatenating the inputs passed to preceding calls to
+ * psa_aead_update_ad() with the ciphertext formed by concatenating the
+ * inputs passed to preceding calls to psa_aead_update().
+ *
+ * When this function returns, the operation becomes inactive.
+ *
+ * \param[in,out] operation     Active AEAD operation.
+ * \param[in] tag               Buffer containing the authentication tag.
+ * \param tag_length            Size of the \p tag buffer in bytes.
+ *
+ * \retval #PSA_SUCCESS
+ *         Success.
+ * \retval #PSA_ERROR_BAD_STATE
+ *         The operation state is not valid (not set up, nonce not set,
+ *         encryption, or already completed).
+ * \retval #PSA_ERROR_INVALID_ARGUMENT
+ *         The total length of input to psa_aead_update_ad() so far is
+ *         less than the additional data length that was previously
+ *         specified with psa_aead_set_lengths().
+ * \retval #PSA_ERROR_INVALID_ARGUMENT
+ *         The total length of input to psa_aead_update() so far is
+ *         less than the plaintext length that was previously
+ *         specified with psa_aead_set_lengths().
+ * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
+ * \retval #PSA_ERROR_COMMUNICATION_FAILURE
+ * \retval #PSA_ERROR_HARDWARE_FAILURE
+ * \retval #PSA_ERROR_TAMPERING_DETECTED
+ */
+psa_status_t psa_aead_verify(psa_aead_operation_t *operation,
+                             const uint8_t *tag,
+                             size_t tag_length);
+
+/** Abort an AEAD operation.
+ *
+ * Aborting an operation frees all associated resources except for the
+ * \p operation structure itself. Once aborted, the operation object
+ * can be reused for another operation by calling
+ * psa_aead_encrypt_setup() or psa_aead_decrypt_setup() again.
+ *
+ * You may call this function any time after the operation object has
+ * been initialized by any of the following methods:
+ * - A call to psa_aead_encrypt_setup() or psa_aead_decrypt_setup(),
+ *   whether it succeeds or not.
+ * - Initializing the \c struct to all-bits-zero.
+ * - Initializing the \c struct to logical zeros, e.g.
+ *   `psa_aead_operation_t operation = {0}`.
+ *
+ * In particular, calling psa_aead_abort() after the operation has been
+ * terminated by a call to psa_aead_abort() or psa_aead_finish()
+ * is safe and has no effect.
+ *
+ * \param[in,out] operation     Initialized AEAD operation.
+ *
+ * \retval #PSA_SUCCESS
+ * \retval #PSA_ERROR_BAD_STATE
+ *         \p operation is not an active AEAD operation.
+ * \retval #PSA_ERROR_COMMUNICATION_FAILURE
+ * \retval #PSA_ERROR_HARDWARE_FAILURE
+ * \retval #PSA_ERROR_TAMPERING_DETECTED
+ */
+psa_status_t psa_aead_abort(psa_aead_operation_t *operation);
+
 /**@}*/
 
 /** \defgroup asymmetric Asymmetric cryptography
@@ -1778,7 +2590,7 @@
  * \param handle                Handle to the key to use for the operation.
  *                              It must be an asymmetric key pair.
  * \param alg                   A signature algorithm that is compatible with
- *                              the type of \p key.
+ *                              the type of \p handle.
  * \param[in] hash              The hash or message to sign.
  * \param hash_length           Size of the \p hash buffer in bytes.
  * \param[out] signature        Buffer where the signature is to be written.
@@ -1792,7 +2604,7 @@
  *         determine a sufficient buffer size by calling
  *         #PSA_ASYMMETRIC_SIGN_OUTPUT_SIZE(\c key_type, \c key_bits, \p alg)
  *         where \c key_type and \c key_bits are the type and bit-size
- *         respectively of \p key.
+ *         respectively of \p handle.
  * \retval #PSA_ERROR_NOT_SUPPORTED
  * \retval #PSA_ERROR_INVALID_ARGUMENT
  * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
@@ -1825,7 +2637,7 @@
  * \param handle            Handle to the key to use for the operation.
  *                          It must be a public key or an asymmetric key pair.
  * \param alg               A signature algorithm that is compatible with
- *                          the type of \p key.
+ *                          the type of \p handle.
  * \param[in] hash          The hash or message whose signature is to be
  *                          verified.
  * \param hash_length       Size of the \p hash buffer in bytes.
@@ -1862,7 +2674,7 @@
  *                              It must be a public key or an asymmetric
  *                              key pair.
  * \param alg                   An asymmetric encryption algorithm that is
- *                              compatible with the type of \p key.
+ *                              compatible with the type of \p handle.
  * \param[in] input             The message to encrypt.
  * \param input_length          Size of the \p input buffer in bytes.
  * \param[in] salt              A salt or label, if supported by the
@@ -1889,7 +2701,7 @@
  *         determine a sufficient buffer size by calling
  *         #PSA_ASYMMETRIC_ENCRYPT_OUTPUT_SIZE(\c key_type, \c key_bits, \p alg)
  *         where \c key_type and \c key_bits are the type and bit-size
- *         respectively of \p key.
+ *         respectively of \p handle.
  * \retval #PSA_ERROR_NOT_SUPPORTED
  * \retval #PSA_ERROR_INVALID_ARGUMENT
  * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
@@ -1918,7 +2730,7 @@
  * \param handle                Handle to the key to use for the operation.
  *                              It must be an asymmetric key pair.
  * \param alg                   An asymmetric encryption algorithm that is
- *                              compatible with the type of \p key.
+ *                              compatible with the type of \p handle.
  * \param[in] input             The message to decrypt.
  * \param input_length          Size of the \p input buffer in bytes.
  * \param[in] salt              A salt or label, if supported by the
@@ -1945,7 +2757,7 @@
  *         determine a sufficient buffer size by calling
  *         #PSA_ASYMMETRIC_DECRYPT_OUTPUT_SIZE(\c key_type, \c key_bits, \p alg)
  *         where \c key_type and \c key_bits are the type and bit-size
- *         respectively of \p key.
+ *         respectively of \p handle.
  * \retval #PSA_ERROR_NOT_SUPPORTED
  * \retval #PSA_ERROR_INVALID_ARGUMENT
  * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
@@ -2037,6 +2849,22 @@
 psa_status_t psa_get_generator_capacity(const psa_crypto_generator_t *generator,
                                         size_t *capacity);
 
+/** Set the maximum capacity of a generator.
+ *
+ * \param[in,out] generator The generator object to modify.
+ * \param capacity          The new capacity of the generator.
+ *                          It must be less or equal to the generator's
+ *                          current capacity.
+ *
+ * \retval #PSA_SUCCESS
+ * \retval #PSA_ERROR_INVALID_ARGUMENT
+ *         \p capacity is larger than the generator's current capacity.
+ * \retval #PSA_ERROR_BAD_STATE
+ * \retval #PSA_ERROR_COMMUNICATION_FAILURE
+ */
+psa_status_t psa_set_generator_capacity(psa_crypto_generator_t *generator,
+                                        size_t capacity);
+
 /** Read some data from a generator.
  *
  * This function reads and returns a sequence of bytes from a generator.
@@ -2066,24 +2894,89 @@
                                 uint8_t *output,
                                 size_t output_length);
 
-/** Create a symmetric key from data read from a generator.
+/** Generate a key deterministically from data read from a generator.
  *
- * This function reads a sequence of bytes from a generator and imports
- * these bytes as a key.
- * The data that is read is discarded from the generator. The generator's
- * capacity is decreased by the number of bytes read.
+ * This function uses the output of a generator to derive a key.
+ * How much output it consumes and how the key is derived depends on the
+ * key type.
  *
- * This function is equivalent to calling #psa_generator_read and
- * passing the resulting output to #psa_import_key, but
- * if the implementation provides an isolation boundary then
- * the key material is not exposed outside the isolation boundary.
+ * - For key types for which the key is an arbitrary sequence of bytes
+ *   of a given size,
+ *   this function is functionally equivalent to calling #psa_generator_read
+ *   and passing the resulting output to #psa_import_key.
+ *   However, this function has a security benefit:
+ *   if the implementation provides an isolation boundary then
+ *   the key material is not exposed outside the isolation boundary.
+ *   As a consequence, for these key types, this function always consumes
+ *   exactly (\p bits / 8) bytes from the generator.
+ *   The following key types defined in this specification follow this scheme:
+ *
+ *     - #PSA_KEY_TYPE_AES;
+ *     - #PSA_KEY_TYPE_ARC4;
+ *     - #PSA_KEY_TYPE_CAMELLIA;
+ *     - #PSA_KEY_TYPE_DERIVE;
+ *     - #PSA_KEY_TYPE_HMAC.
+ *
+ * - For ECC keys on a Montgomery elliptic curve
+ *   (#PSA_KEY_TYPE_ECC_KEYPAIR(\c curve) where \c curve designates a
+ *   Montgomery curve), this function always draws a byte string whose
+ *   length is determined by the curve, and sets the mandatory bits
+ *   accordingly. That is:
+ *
+ *     - #PSA_ECC_CURVE_CURVE25519: draw a 32-byte string
+ *       and process it as specified in RFC 7748 &sect;5.
+ *     - #PSA_ECC_CURVE_CURVE448: draw a 56-byte string
+ *       and process it as specified in RFC 7748 &sect;5.
+ *
+ * - For key types for which the key is represented by a single sequence of
+ *   \p bits bits with constraints as to which bit sequences are acceptable,
+ *   this function draws a byte string of length (\p bits / 8) bytes rounded
+ *   up to the nearest whole number of bytes. If the resulting byte string
+ *   is acceptable, it becomes the key, otherwise the drawn bytes are discarded.
+ *   This process is repeated until an acceptable byte string is drawn.
+ *   The byte string drawn from the generator is interpreted as specified
+ *   for the output produced by psa_export_key().
+ *   The following key types defined in this specification follow this scheme:
+ *
+ *     - #PSA_KEY_TYPE_DES.
+ *       Force-set the parity bits, but discard forbidden weak keys.
+ *       For 2-key and 3-key triple-DES, the three keys are generated
+ *       successively (for example, for 3-key triple-DES,
+ *       if the first 8 bytes specify a weak key and the next 8 bytes do not,
+ *       discard the first 8 bytes, use the next 8 bytes as the first key,
+ *       and continue reading output from the generator to derive the other
+ *       two keys).
+ *     - Finite-field Diffie-Hellman keys (#PSA_KEY_TYPE_DH_KEYPAIR),
+ *       DSA keys (#PSA_KEY_TYPE_DSA_KEYPAIR), and
+ *       ECC keys on a Weierstrass elliptic curve
+ *       (#PSA_KEY_TYPE_ECC_KEYPAIR(\c curve) where \c curve designates a
+ *       Weierstrass curve).
+ *       For these key types, interpret the byte string as integer
+ *       in big-endian order. Discard it if it is not in the range
+ *       [0, *N* - 2] where *N* is the boundary of the private key domain
+ *       (the prime *p* for Diffie-Hellman, the subprime *q* for DSA,
+ *       or the order of the curve's base point for ECC).
+ *       Add 1 to the resulting integer and use this as the private key *x*.
+ *       This method allows compliance to NIST standards, specifically
+ *       the methods titled "key-pair generation by testing candidates"
+ *       in NIST SP 800-56A &sect;5.6.1.1.4 for Diffie-Hellman,
+ *       in FIPS 186-4 &sect;B.1.2 for DSA, and
+ *       in NIST SP 800-56A &sect;5.6.1.2.2 or
+ *       FIPS 186-4 &sect;B.4.2 for elliptic curve keys.
+ *
+ * - For other key types, including #PSA_KEY_TYPE_RSA_KEYPAIR,
+ *   the way in which the generator output is consumed is
+ *   implementation-defined.
+ *
+ * In all cases, the data that is read is discarded from the generator.
+ * The generator's capacity is decreased by the number of bytes read.
  *
  * \param handle            Handle to the slot where the key will be stored.
  *                          It must have been obtained by calling
  *                          psa_allocate_key() or psa_create_key() and must
  *                          not contain key material yet.
  * \param type              Key type (a \c PSA_KEY_TYPE_XXX value).
- *                          This must be a symmetric key type.
+ *                          This must be a secret key type or a key pair type.
  * \param bits              Key size in bits.
  * \param[in,out] generator The generator object to read from.
  *
@@ -2092,12 +2985,10 @@
  *         If the key is persistent, the key material and the key's metadata
  *         have been saved to persistent storage.
  * \retval #PSA_ERROR_INSUFFICIENT_DATA
- *                          There were fewer than \p output_length bytes
- *                          in the generator. Note that in this case, no
- *                          output is written to the output buffer.
- *                          The generator's capacity is set to 0, thus
- *                          subsequent calls to this function will not
- *                          succeed, even with a smaller output buffer.
+ *         There was not enough data to create the desired key.
+ *         Note that in this case, no output is written to the output buffer.
+ *         The generator'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 slot.
@@ -2161,37 +3052,36 @@
 
 /** Set up a key derivation operation.
  *
- * A key derivation algorithm takes three inputs: a secret input \p key and
- * two non-secret inputs \p label and p salt.
- * The result of this function is a byte generator which can
- * be used to produce keys and other cryptographic material.
+ * A key derivation algorithm takes some inputs and uses them to create
+ * a byte generator which can be used to produce keys and other
+ * cryptographic material.
  *
- * The role of \p label and \p salt is as follows:
- * - For HKDF (#PSA_ALG_HKDF), \p salt is the salt used in the "extract" step
- *   and \p label is the info string used in the "expand" step.
+ * To use a generator for key derivation:
+ * - Start with an initialized object of type #psa_crypto_generator_t.
+ * - Call psa_key_derivation_setup() to select the algorithm.
+ * - Provide the inputs for the key derivation by calling
+ *   psa_key_derivation_input_bytes() or psa_key_derivation_input_key()
+ *   as appropriate. Which inputs are needed, in what order, and whether
+ *   they may be keys and if so of what type depends on the algorithm.
+ * - Optionally set the generator's maximum capacity with
+ *   psa_set_generator_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.
+ * - Generate output with psa_generator_read() or
+ *   psa_generator_import_key(). Successive calls to these functions
+ *   use successive output bytes from the generator.
+ * - Clean up the generator object with psa_generator_abort().
  *
- * \param[in,out] generator       The generator object to set up. It must have
- *                                been initialized as per the documentation for
- *                                #psa_crypto_generator_t and not yet in use.
- * \param handle                  Handle to the secret key.
+ * \param[in,out] generator       The generator object to set up. It must
+ *                                have been initialized but not set up yet.
  * \param alg                     The key derivation algorithm to compute
  *                                (\c PSA_ALG_XXX value such that
  *                                #PSA_ALG_IS_KEY_DERIVATION(\p alg) is true).
- * \param[in] salt                Salt to use.
- * \param salt_length             Size of the \p salt buffer in bytes.
- * \param[in] label               Label to use.
- * \param label_length            Size of the \p label buffer in bytes.
- * \param capacity                The maximum number of bytes that the
- *                                generator will be able to provide.
  *
  * \retval #PSA_SUCCESS
  *         Success.
- * \retval #PSA_ERROR_INVALID_HANDLE
- * \retval #PSA_ERROR_DOES_NOT_EXIST
- * \retval #PSA_ERROR_NOT_PERMITTED
  * \retval #PSA_ERROR_INVALID_ARGUMENT
- *         \c key is not compatible with \c alg,
- *         or \p capacity is too large for the specified algorithm and key.
+ *         \c alg is not a key derivation algorithm.
  * \retval #PSA_ERROR_NOT_SUPPORTED
  *         \c alg is not supported or is not a key derivation algorithm.
  * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
@@ -2199,51 +3089,130 @@
  * \retval #PSA_ERROR_HARDWARE_FAILURE
  * \retval #PSA_ERROR_TAMPERING_DETECTED
  * \retval #PSA_ERROR_BAD_STATE
+ */
+psa_status_t psa_key_derivation_setup(psa_crypto_generator_t *generator,
+                                      psa_algorithm_t alg);
+
+/** Provide an input for key derivation or key agreement.
+ *
+ * Which inputs are required and in what order depends on the algorithm.
+ * Refer to the documentation of each key derivation or key agreement
+ * algorithm for information.
+ *
+ * This function passes direct inputs. Some inputs must be passed as keys
+ * using psa_key_derivation_input_key() instead of this function. Refer to
+ * the documentation of individual step types for information.
+ *
+ * \param[in,out] generator       The generator object to use. It must
+ *                                have been set up with
+ *                                psa_key_derivation_setup() and must not
+ *                                have produced any output yet.
+ * \param step                    Which step the input data is for.
+ * \param[in] data                Input data to use.
+ * \param data_length             Size of the \p data buffer in bytes.
+ *
+ * \retval #PSA_SUCCESS
+ *         Success.
+ * \retval #PSA_ERROR_INVALID_ARGUMENT
+ *         \c step is not compatible with the generator's algorithm.
+ * \retval #PSA_ERROR_INVALID_ARGUMENT
+ *         \c step does not allow direct inputs.
+ * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
+ * \retval #PSA_ERROR_COMMUNICATION_FAILURE
+ * \retval #PSA_ERROR_HARDWARE_FAILURE
+ * \retval #PSA_ERROR_TAMPERING_DETECTED
+ * \retval #PSA_ERROR_BAD_STATE
+ *         The value of \p step is not valid given the state of \p generator.
+ * \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_key_derivation(psa_crypto_generator_t *generator,
-                                psa_key_handle_t handle,
-                                psa_algorithm_t alg,
-                                const uint8_t *salt,
-                                size_t salt_length,
-                                const uint8_t *label,
-                                size_t label_length,
-                                size_t capacity);
+psa_status_t psa_key_derivation_input_bytes(psa_crypto_generator_t *generator,
+                                            psa_key_derivation_step_t step,
+                                            const uint8_t *data,
+                                            size_t data_length);
 
-/** Set up a key agreement operation.
+/** Provide an input for key derivation in the form of a key.
+ *
+ * Which inputs are required and in what order depends on the algorithm.
+ * Refer to the documentation of each key derivation or key agreement
+ * algorithm for information.
+ *
+ * This function passes key inputs. Some inputs must be passed as keys
+ * of the appropriate type using this function, while others must be
+ * passed as direct inputs using psa_key_derivation_input_bytes(). Refer to
+ * the documentation of individual step types for information.
+ *
+ * \param[in,out] generator       The generator object to use. It must
+ *                                have been set up with
+ *                                psa_key_derivation_setup() and must not
+ *                                have produced any output yet.
+ * \param step                    Which step the input data is for.
+ * \param handle                  Handle to the key. It must have an
+ *                                appropriate type for \p step and must
+ *                                allow the usage #PSA_KEY_USAGE_DERIVE.
+ *
+ * \retval #PSA_SUCCESS
+ *         Success.
+ * \retval #PSA_ERROR_INVALID_HANDLE
+ * \retval #PSA_ERROR_DOES_NOT_EXIST
+ * \retval #PSA_ERROR_NOT_PERMITTED
+ * \retval #PSA_ERROR_INVALID_ARGUMENT
+ *         \c step is not compatible with the generator's algorithm.
+ * \retval #PSA_ERROR_INVALID_ARGUMENT
+ *         \c step does not allow key inputs.
+ * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
+ * \retval #PSA_ERROR_COMMUNICATION_FAILURE
+ * \retval #PSA_ERROR_HARDWARE_FAILURE
+ * \retval #PSA_ERROR_TAMPERING_DETECTED
+ * \retval #PSA_ERROR_BAD_STATE
+ *         The value of \p step is not valid given the state of \p generator.
+ * \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_key_derivation_input_key(psa_crypto_generator_t *generator,
+                                          psa_key_derivation_step_t step,
+                                          psa_key_handle_t handle);
+
+/** Perform a key agreement and use the shared secret as input to a key
+ * derivation.
  *
  * A key agreement algorithm takes two inputs: a private key \p private_key
  * a public key \p peer_key.
- * The result of this function is a byte generator which can
- * be used to produce keys and other cryptographic material.
+ * The result of this function is passed as input to a key derivation.
+ * The output of this key derivation can be extracted by reading from the
+ * resulting generator to produce keys and other cryptographic material.
  *
- * The resulting generator always has the maximum capacity permitted by
- * the algorithm.
- *
- * \param[in,out] generator The generator object to set up. It must have been
- *                          initialized as per the documentation for
- *                          #psa_crypto_generator_t and not yet in use.
- * \param private_key       Handle to the private key to use.
+ * \param[in,out] generator       The generator object to use. It must
+ *                                have been set up with
+ *                                psa_key_derivation_setup() with a
+ *                                key agreement and derivation algorithm
+ *                                \c alg (\c PSA_ALG_XXX value such that
+ *                                #PSA_ALG_IS_KEY_AGREEMENT(\c alg) is true
+ *                                and #PSA_ALG_IS_RAW_KEY_AGREEMENT(\c alg)
+ *                                is false).
+ *                                The generator must be ready for an
+ *                                input of the type given by \p step.
+ * \param step                    Which step the input data is for.
+ * \param private_key             Handle to the private key to use.
  * \param[in] peer_key      Public key of the peer. The peer key must be in the
  *                          same format that psa_import_key() accepts for the
  *                          public key type corresponding to the type of
- *                          \p private_key. That is, this function performs the
+ *                          private_key. That is, this function performs the
  *                          equivalent of
- *                          `psa_import_key(internal_public_key_handle,
- *                          PSA_KEY_TYPE_PUBLIC_KEY_OF_KEYPAIR(private_key_type),
- *                          peer_key, peer_key_length)` where
- *                          `private_key_type` is the type of \p private_key.
- *                          For example, for EC keys, this means that \p
- *                          peer_key is interpreted as a point on the curve
- *                          that the private key is associated with. The
- *                          standard formats for public keys are documented in
- *                          the documentation of psa_export_public_key().
- * \param peer_key_length   Size of \p peer_key in bytes.
- * \param alg               The key agreement algorithm to compute
- *                          (\c PSA_ALG_XXX value such that
- *                          #PSA_ALG_IS_KEY_AGREEMENT(\p alg) is true).
+ *                          #psa_import_key(`internal_public_key_handle`,
+ *                          #PSA_KEY_TYPE_PUBLIC_KEY_OF_KEYPAIR(`private_key_type`),
+ *                          `peer_key`, `peer_key_length`) where
+ *                          `private_key_type` is the type of `private_key`.
+ *                          For example, for EC keys, this means that peer_key
+ *                          is interpreted as a point on the curve that the
+ *                          private key is on. The standard formats for public
+ *                          keys are documented in the documentation of
+ *                          psa_export_public_key().
+ * \param peer_key_length         Size of \p peer_key in bytes.
  *
  * \retval #PSA_SUCCESS
  *         Success.
@@ -2262,10 +3231,66 @@
  * \retval #PSA_ERROR_TAMPERING_DETECTED
  */
 psa_status_t psa_key_agreement(psa_crypto_generator_t *generator,
+                               psa_key_derivation_step_t step,
                                psa_key_handle_t private_key,
                                const uint8_t *peer_key,
-                               size_t peer_key_length,
-                               psa_algorithm_t alg);
+                               size_t peer_key_length);
+
+/** Perform a key agreement and use the shared secret as input to a key
+ * derivation.
+ *
+ * A key agreement algorithm takes two inputs: a private key \p private_key
+ * a public key \p peer_key.
+ *
+ * \warning The raw result of a key agreement algorithm such as finite-field
+ * Diffie-Hellman or elliptic curve Diffie-Hellman has biases and should
+ * not be used directly as key material. It should instead be passed as
+ * input to a key derivation algorithm. To chain a key agreement with
+ * a key derivation, use psa_key_agreement() and other functions from
+ * the key derivation and generator interface.
+ *
+ * \param alg                     The key agreement algorithm to compute
+ *                                (\c PSA_ALG_XXX value such that
+ *                                #PSA_ALG_IS_RAW_KEY_AGREEMENT(\p alg)
+ *                                is true).
+ * \param private_key             Handle to the private key to use.
+ * \param[in] peer_key            Public key of the peer. It must be
+ *                                in the same format that psa_import_key()
+ *                                accepts. The standard formats for public
+ *                                keys are documented in the documentation
+ *                                of psa_export_public_key().
+ * \param peer_key_length         Size of \p peer_key in bytes.
+ * \param[out] output             Buffer where the decrypted message is to
+ *                                be written.
+ * \param output_size             Size of the \c output buffer in bytes.
+ * \param[out] output_length      On success, the number of bytes
+ *                                that make up the returned output.
+ *
+ * \retval #PSA_SUCCESS
+ *         Success.
+ * \retval #PSA_ERROR_INVALID_HANDLE
+ * \retval #PSA_ERROR_EMPTY_SLOT
+ * \retval #PSA_ERROR_NOT_PERMITTED
+ * \retval #PSA_ERROR_INVALID_ARGUMENT
+ *         \p alg is not a key agreement algorithm
+ * \retval #PSA_ERROR_INVALID_ARGUMENT
+ *         \p private_key is not compatible with \p alg,
+ *         or \p peer_key is not valid for \p alg or not compatible with
+ *         \p private_key.
+ * \retval #PSA_ERROR_NOT_SUPPORTED
+ *         \p alg is not a supported key agreement algorithm.
+ * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
+ * \retval #PSA_ERROR_COMMUNICATION_FAILURE
+ * \retval #PSA_ERROR_HARDWARE_FAILURE
+ * \retval #PSA_ERROR_TAMPERING_DETECTED
+ */
+psa_status_t psa_key_agreement_raw_shared_secret(psa_algorithm_t alg,
+                                                 psa_key_handle_t private_key,
+                                                 const uint8_t *peer_key,
+                                                 size_t peer_key_length,
+                                                 uint8_t *output,
+                                                 size_t output_size,
+                                                 size_t *output_length);
 
 /**@}*/
 
@@ -2338,6 +3363,18 @@
  *                            specifying the public exponent. The
  *                            default public exponent used when \p extra
  *                            is \c NULL is 65537.
+ *                          - For an DSA key (\p type is
+ *                            #PSA_KEY_TYPE_DSA_KEYPAIR), \p extra is an
+ *                            optional structure specifying the key domain
+ *                            parameters. The key domain parameters can also be
+ *                            provided by psa_set_key_domain_parameters(),
+ *                            which documents the format of the structure.
+ *                          - For a DH key (\p type is
+ *                            #PSA_KEY_TYPE_DH_KEYPAIR), the \p extra is an
+ *                            optional structure specifying the key domain
+ *                            parameters. The key domain parameters can also be
+ *                            provided by psa_set_key_domain_parameters(),
+ *                            which documents the format of the structure.
  * \param extra_size        Size of the buffer that \p extra
  *                          points to, in bytes. Note that if \p extra is
  *                          \c NULL then \p extra_size must be zero.
diff --git a/include/psa/crypto_extra.h b/include/psa/crypto_extra.h
index c89c55d..093355d 100644
--- a/include/psa/crypto_extra.h
+++ b/include/psa/crypto_extra.h
@@ -143,6 +143,64 @@
 psa_status_t mbedtls_psa_inject_entropy(const unsigned char *seed,
                                         size_t seed_size);
 
+/** Set up a key derivation operation.
+ *
+ * FIMXE This function is no longer part of the official API. Its prototype
+ * is only kept around for the sake of tests that haven't been updated yet.
+ *
+ * A key derivation algorithm takes three inputs: a secret input \p handle and
+ * two non-secret inputs \p label and p salt.
+ * The result of this function is a byte generator which can
+ * be used to produce keys and other cryptographic material.
+ *
+ * The role of \p label and \p salt is as follows:
+ * - For HKDF (#PSA_ALG_HKDF), \p salt is the salt used in the "extract" step
+ *   and \p label is the info string used in the "expand" step.
+ *
+ * \param[in,out] generator       The generator object to set up. It must have
+ *                                been initialized as per the documentation for
+ *                                #psa_crypto_generator_t and not yet in use.
+ * \param handle                  Handle to the secret key.
+ * \param alg                     The key derivation algorithm to compute
+ *                                (\c PSA_ALG_XXX value such that
+ *                                #PSA_ALG_IS_KEY_DERIVATION(\p alg) is true).
+ * \param[in] salt                Salt to use.
+ * \param salt_length             Size of the \p salt buffer in bytes.
+ * \param[in] label               Label to use.
+ * \param label_length            Size of the \p label buffer in bytes.
+ * \param capacity                The maximum number of bytes that the
+ *                                generator will be able to provide.
+ *
+ * \retval #PSA_SUCCESS
+ *         Success.
+ * \retval #PSA_ERROR_INVALID_HANDLE
+ * \retval #PSA_ERROR_EMPTY_SLOT
+ * \retval #PSA_ERROR_NOT_PERMITTED
+ * \retval #PSA_ERROR_INVALID_ARGUMENT
+ *         \c key is not compatible with \c alg,
+ *         or \p capacity is too large for the specified algorithm and key.
+ * \retval #PSA_ERROR_NOT_SUPPORTED
+ *         \c alg is not supported or is not a key derivation algorithm.
+ * \retval #PSA_ERROR_INSUFFICIENT_MEMORY
+ * \retval #PSA_ERROR_COMMUNICATION_FAILURE
+ * \retval #PSA_ERROR_HARDWARE_FAILURE
+ * \retval #PSA_ERROR_TAMPERING_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 psa_key_derivation(psa_crypto_generator_t *generator,
+                                psa_key_handle_t handle,
+                                psa_algorithm_t alg,
+                                const uint8_t *salt,
+                                size_t salt_length,
+                                const uint8_t *label,
+                                size_t label_length,
+                                size_t capacity);
+
+/* FIXME Deprecated. Remove this as soon as all the tests are updated. */
+#define PSA_ALG_SELECT_RAW                      ((psa_algorithm_t)0x31000001)
 
 #ifdef __cplusplus
 }
diff --git a/include/psa/crypto_sizes.h b/include/psa/crypto_sizes.h
index 34664fc..3c879e8 100644
--- a/include/psa/crypto_sizes.h
+++ b/include/psa/crypto_sizes.h
@@ -228,7 +228,7 @@
  * \param key_type      The type of the MAC key.
  * \param key_bits      The size of the MAC key in bits.
  * \param alg           A MAC algorithm (\c PSA_ALG_XXX value such that
- *                      #PSA_ALG_IS_MAC(alg) is true).
+ *                      #PSA_ALG_IS_MAC(\p alg) is true).
  *
  * \return              The MAC size for the specified algorithm with
  *                      the specified key parameters.
@@ -253,7 +253,7 @@
  *
  * \param alg                 An AEAD algorithm
  *                            (\c PSA_ALG_XXX value such that
- *                            #PSA_ALG_IS_AEAD(alg) is true).
+ *                            #PSA_ALG_IS_AEAD(\p alg) is true).
  * \param plaintext_length    Size of the plaintext in bytes.
  *
  * \return                    The AEAD ciphertext size for the specified
@@ -268,6 +268,27 @@
      (plaintext_length) + PSA_AEAD_TAG_LENGTH(alg) :              \
      0)
 
+/** The maximum size of the output of psa_aead_finish(), in bytes.
+ *
+ * If the size of the ciphertext buffer is at least this large, it is
+ * guaranteed that psa_aead_finish() will not fail due to an
+ * insufficient buffer size. Depending on the algorithm, the actual size of
+ * the ciphertext may be smaller.
+ *
+ * \param alg                 An AEAD algorithm
+ *                            (\c PSA_ALG_XXX value such that
+ *                            #PSA_ALG_IS_AEAD(\p alg) is true).
+ *
+ * \return                    The maximum trailing ciphertext size for the
+ *                            specified algorithm.
+ *                            If the AEAD algorithm is not recognized, return 0.
+ *                            An implementation may return either 0 or a
+ *                            correct size for an AEAD algorithm that it
+ *                            recognizes, but does not support.
+ */
+#define PSA_AEAD_FINISH_OUTPUT_SIZE(alg)      \
+    ((size_t)0)
+
 /** The maximum size of the output of psa_aead_decrypt(), in bytes.
  *
  * If the size of the plaintext buffer is at least this large, it is
@@ -277,7 +298,7 @@
  *
  * \param alg                 An AEAD algorithm
  *                            (\c PSA_ALG_XXX value such that
- *                            #PSA_ALG_IS_AEAD(alg) is true).
+ *                            #PSA_ALG_IS_AEAD(\p alg) is true).
  * \param ciphertext_length   Size of the plaintext in bytes.
  *
  * \return                    The AEAD ciphertext size for the specified
diff --git a/include/psa/crypto_struct.h b/include/psa/crypto_struct.h
index ee3ecd7..2414ad5 100644
--- a/include/psa/crypto_struct.h
+++ b/include/psa/crypto_struct.h
@@ -152,6 +152,27 @@
     return( v );
 }
 
+struct psa_aead_operation_s
+{
+    psa_algorithm_t alg;
+    unsigned int key_set : 1;
+    unsigned int iv_set : 1;
+    uint8_t iv_size;
+    uint8_t block_size;
+    union
+    {
+        unsigned dummy; /* Enable easier initializing of the union. */
+        mbedtls_cipher_context_t cipher;
+    } ctx;
+};
+
+#define PSA_AEAD_OPERATION_INIT {0, 0, 0, 0, 0, {0}}
+static inline struct psa_aead_operation_s psa_aead_operation_init( void )
+{
+    const struct psa_aead_operation_s v = PSA_AEAD_OPERATION_INIT;
+    return( v );
+}
+
 #if defined(MBEDTLS_MD_C)
 typedef struct
 {
@@ -165,6 +186,8 @@
 #endif
     uint8_t offset_in_block;
     uint8_t block_number;
+    uint8_t state : 2;
+    uint8_t info_set : 1;
 } psa_hkdf_generator_t;
 #endif /* MBEDTLS_MD_C */
 
diff --git a/include/psa/crypto_types.h b/include/psa/crypto_types.h
index 923b94a..7054de7 100644
--- a/include/psa/crypto_types.h
+++ b/include/psa/crypto_types.h
@@ -110,4 +110,13 @@
 
 /**@}*/
 
+/** \defgroup derivation Key derivation
+ * @{
+ */
+
+/** \brief Encoding of the step of a key derivation. */
+typedef uint16_t psa_key_derivation_step_t;
+
+/**@}*/
+
 #endif /* PSA_CRYPTO_TYPES_H */
diff --git a/include/psa/crypto_values.h b/include/psa/crypto_values.h
index d42d8c2..f9ac00a 100644
--- a/include/psa/crypto_values.h
+++ b/include/psa/crypto_values.h
@@ -175,7 +175,7 @@
  * the global integrity of the keystore. Depending on the global
  * integrity guarantees offered by the implementation, access to other
  * data may or may not fail even if the data is still readable but
- * its integrity canont be guaranteed.
+ * its integrity cannot be guaranteed.
  *
  * Implementations should only use this error code to report a
  * permanent storage corruption. However application writers should
@@ -483,6 +483,15 @@
 #define PSA_ECC_CURVE_CURVE25519        ((psa_ecc_curve_t) 0x001d)
 #define PSA_ECC_CURVE_CURVE448          ((psa_ecc_curve_t) 0x001e)
 
+/** Diffie-Hellman key exchange public key. */
+#define PSA_KEY_TYPE_DH_PUBLIC_KEY             ((psa_key_type_t)0x60040000)
+/** Diffie-Hellman key exchange key pair (private and public key). */
+#define PSA_KEY_TYPE_DH_KEYPAIR                ((psa_key_type_t)0x70040000)
+/** Whether a key type is a Diffie-Hellman key exchange key (pair or
+ * public-only). */
+#define PSA_KEY_TYPE_IS_DH(type)                                       \
+    (PSA_KEY_TYPE_PUBLIC_KEY_OF_KEYPAIR(type) == PSA_KEY_TYPE_DH_PUBLIC_KEY)
+
 /** The block size of a block cipher.
  *
  * \param type  A cipher key type (value of type #psa_key_type_t).
@@ -517,9 +526,8 @@
 #define PSA_ALG_CATEGORY_AEAD                   ((psa_algorithm_t)0x06000000)
 #define PSA_ALG_CATEGORY_SIGN                   ((psa_algorithm_t)0x10000000)
 #define PSA_ALG_CATEGORY_ASYMMETRIC_ENCRYPTION  ((psa_algorithm_t)0x12000000)
-#define PSA_ALG_CATEGORY_KEY_AGREEMENT          ((psa_algorithm_t)0x22000000)
-#define PSA_ALG_CATEGORY_KEY_DERIVATION         ((psa_algorithm_t)0x30000000)
-#define PSA_ALG_CATEGORY_KEY_SELECTION          ((psa_algorithm_t)0x31000000)
+#define PSA_ALG_CATEGORY_KEY_DERIVATION         ((psa_algorithm_t)0x20000000)
+#define PSA_ALG_CATEGORY_KEY_AGREEMENT          ((psa_algorithm_t)0x30000000)
 
 #define PSA_ALG_IS_VENDOR_DEFINED(alg)                                  \
     (((alg) & PSA_ALG_VENDOR_FLAG) != 0)
@@ -591,7 +599,6 @@
 #define PSA_ALG_IS_ASYMMETRIC_ENCRYPTION(alg)                           \
     (((alg) & PSA_ALG_CATEGORY_MASK) == PSA_ALG_CATEGORY_ASYMMETRIC_ENCRYPTION)
 
-#define PSA_ALG_KEY_SELECTION_FLAG              ((psa_algorithm_t)0x01000000)
 /** Whether the specified algorithm is a key agreement algorithm.
  *
  * \param alg An algorithm identifier (value of type #psa_algorithm_t).
@@ -601,8 +608,7 @@
  *         algorithm identifier.
  */
 #define PSA_ALG_IS_KEY_AGREEMENT(alg)                                   \
-    (((alg) & PSA_ALG_CATEGORY_MASK & ~PSA_ALG_KEY_SELECTION_FLAG) ==   \
-     PSA_ALG_CATEGORY_KEY_AGREEMENT)
+    (((alg) & PSA_ALG_CATEGORY_MASK) == PSA_ALG_CATEGORY_KEY_AGREEMENT)
 
 /** Whether the specified algorithm is a key derivation algorithm.
  *
@@ -615,17 +621,6 @@
 #define PSA_ALG_IS_KEY_DERIVATION(alg)                                  \
     (((alg) & PSA_ALG_CATEGORY_MASK) == PSA_ALG_CATEGORY_KEY_DERIVATION)
 
-/** Whether the specified algorithm is a key selection algorithm.
- *
- * \param alg An algorithm identifier (value of type #psa_algorithm_t).
- *
- * \return 1 if \p alg is a key selection algorithm, 0 otherwise.
- *         This macro may return either 0 or 1 if \p alg is not a supported
- *         algorithm identifier.
- */
-#define PSA_ALG_IS_KEY_SELECTION(alg)                                   \
-    (((alg) & PSA_ALG_CATEGORY_MASK) == PSA_ALG_CATEGORY_KEY_SELECTION)
-
 #define PSA_ALG_HASH_MASK                       ((psa_algorithm_t)0x000000ff)
 
 #define PSA_ALG_MD2                             ((psa_algorithm_t)0x01000001)
@@ -685,7 +680,7 @@
  *
  * This value may not be used to build other algorithms that are
  * parametrized over a hash. For any valid use of this macro to build
- * an algorithm `\p alg`, #PSA_ALG_IS_HASH_AND_SIGN(\p alg) is true.
+ * an algorithm \c alg, #PSA_ALG_IS_HASH_AND_SIGN(\c alg) is true.
  *
  * This value may not be used to build an algorithm specification to
  * perform an operation. It is only valid to build policies.
@@ -702,7 +697,7 @@
  *                      #PSA_ALG_IS_HASH(\p hash_alg) is true).
  *
  * \return              The corresponding HMAC algorithm.
- * \return              Unspecified if \p alg is not a supported
+ * \return              Unspecified if \p hash_alg is not a supported
  *                      hash algorithm.
  */
 #define PSA_ALG_HMAC(hash_alg)                                  \
@@ -879,7 +874,12 @@
  */
 #define PSA_ALG_CBC_PKCS7                       ((psa_algorithm_t)0x04600101)
 
+/** The CCM authenticated encryption algorithm.
+ */
 #define PSA_ALG_CCM                             ((psa_algorithm_t)0x06001001)
+
+/** The GCM authenticated encryption algorithm.
+ */
 #define PSA_ALG_GCM                             ((psa_algorithm_t)0x06001002)
 
 /* In the encoding of a AEAD algorithm, the bits corresponding to
@@ -943,7 +943,7 @@
  *                      when specifying the algorithm in a usage policy.
  *
  * \return              The corresponding RSA PKCS#1 v1.5 signature algorithm.
- * \return              Unspecified if \p alg is not a supported
+ * \return              Unspecified if \p hash_alg is not a supported
  *                      hash algorithm.
  */
 #define PSA_ALG_RSA_PKCS1V15_SIGN(hash_alg)                             \
@@ -974,7 +974,7 @@
  *                      when specifying the algorithm in a usage policy.
  *
  * \return              The corresponding RSA PSS signature algorithm.
- * \return              Unspecified if \p alg is not a supported
+ * \return              Unspecified if \p hash_alg is not a supported
  *                      hash algorithm.
  */
 #define PSA_ALG_RSA_PSS(hash_alg)                               \
@@ -994,13 +994,27 @@
  *                      when specifying the algorithm in a usage policy.
  *
  * \return              The corresponding DSA signature algorithm.
- * \return              Unspecified if \p alg is not a supported
+ * \return              Unspecified if \p hash_alg is not a supported
  *                      hash algorithm.
  */
 #define PSA_ALG_DSA(hash_alg)                             \
     (PSA_ALG_DSA_BASE | ((hash_alg) & PSA_ALG_HASH_MASK))
 #define PSA_ALG_DETERMINISTIC_DSA_BASE          ((psa_algorithm_t)0x10050000)
 #define PSA_ALG_DSA_DETERMINISTIC_FLAG          ((psa_algorithm_t)0x00010000)
+/** Deterministic DSA signature with hashing.
+ *
+ * This is the deterministic variant defined by RFC 6979 of
+ * the signature scheme defined by FIPS 186-4.
+ *
+ * \param hash_alg      A hash algorithm (\c PSA_ALG_XXX value such that
+ *                      #PSA_ALG_IS_HASH(\p hash_alg) is true).
+ *                      This includes #PSA_ALG_ANY_HASH
+ *                      when specifying the algorithm in a usage policy.
+ *
+ * \return              The corresponding DSA signature algorithm.
+ * \return              Unspecified if \p hash_alg is not a supported
+ *                      hash algorithm.
+ */
 #define PSA_ALG_DETERMINISTIC_DSA(hash_alg)                             \
     (PSA_ALG_DETERMINISTIC_DSA_BASE | ((hash_alg) & PSA_ALG_HASH_MASK))
 #define PSA_ALG_IS_DSA(alg)                                             \
@@ -1031,7 +1045,7 @@
  *                      when specifying the algorithm in a usage policy.
  *
  * \return              The corresponding ECDSA signature algorithm.
- * \return              Unspecified if \p alg is not a supported
+ * \return              Unspecified if \p hash_alg is not a supported
  *                      hash algorithm.
  */
 #define PSA_ALG_ECDSA(hash_alg)                                 \
@@ -1066,7 +1080,7 @@
  *
  * \return              The corresponding deterministic ECDSA signature
  *                      algorithm.
- * \return              Unspecified if \p alg is not a supported
+ * \return              Unspecified if \p hash_alg is not a supported
  *                      hash algorithm.
  */
 #define PSA_ALG_DETERMINISTIC_ECDSA(hash_alg)                           \
@@ -1138,7 +1152,7 @@
  *                      for MGF1.
  *
  * \return              The corresponding RSA OAEP signature algorithm.
- * \return              Unspecified if \p alg is not a supported
+ * \return              Unspecified if \p hash_alg is not a supported
  *                      hash algorithm.
  */
 #define PSA_ALG_RSA_OAEP(hash_alg)                              \
@@ -1150,16 +1164,25 @@
      ((alg) & PSA_ALG_HASH_MASK) | PSA_ALG_CATEGORY_HASH :      \
      0)
 
-#define PSA_ALG_HKDF_BASE                       ((psa_algorithm_t)0x30000100)
+#define PSA_ALG_HKDF_BASE                       ((psa_algorithm_t)0x20000100)
 /** Macro to build an HKDF algorithm.
  *
  * For example, `PSA_ALG_HKDF(PSA_ALG_SHA256)` is HKDF using HMAC-SHA-256.
  *
+ * This key derivation algorithm uses the following inputs:
+ * - #PSA_KDF_STEP_SALT is the salt used in the "extract" step.
+ *   It is optional; if omitted, the derivation uses an empty salt.
+ * - #PSA_KDF_STEP_SECRET is the secret key used in the "extract" step.
+ * - #PSA_KDF_STEP_INFO is the info string used in the "expand" step.
+ * You must pass #PSA_KDF_STEP_SALT before #PSA_KDF_STEP_SECRET.
+ * You may pass #PSA_KDF_STEP_INFO at any time after steup and before
+ * starting to generate output.
+ *
  * \param hash_alg      A hash algorithm (\c PSA_ALG_XXX value such that
  *                      #PSA_ALG_IS_HASH(\p hash_alg) is true).
  *
  * \return              The corresponding HKDF algorithm.
- * \return              Unspecified if \p alg is not a supported
+ * \return              Unspecified if \p hash_alg is not a supported
  *                      hash algorithm.
  */
 #define PSA_ALG_HKDF(hash_alg)                                  \
@@ -1180,7 +1203,7 @@
 #define PSA_ALG_HKDF_GET_HASH(hkdf_alg)                         \
     (PSA_ALG_CATEGORY_HASH | ((hkdf_alg) & PSA_ALG_HASH_MASK))
 
-#define PSA_ALG_TLS12_PRF_BASE                     ((psa_algorithm_t)0x30000200)
+#define PSA_ALG_TLS12_PRF_BASE                  ((psa_algorithm_t)0x20000200)
 /** Macro to build a TLS-1.2 PRF algorithm.
  *
  * TLS 1.2 uses a custom pseudorandom function (PRF) for key schedule,
@@ -1200,7 +1223,7 @@
  *                      #PSA_ALG_IS_HASH(\p hash_alg) is true).
  *
  * \return              The corresponding TLS-1.2 PRF algorithm.
- * \return              Unspecified if \p alg is not a supported
+ * \return              Unspecified if \p hash_alg is not a supported
  *                      hash algorithm.
  */
 #define PSA_ALG_TLS12_PRF(hash_alg)                                  \
@@ -1219,7 +1242,7 @@
 #define PSA_ALG_TLS12_PRF_GET_HASH(hkdf_alg)                         \
     (PSA_ALG_CATEGORY_HASH | ((hkdf_alg) & PSA_ALG_HASH_MASK))
 
-#define PSA_ALG_TLS12_PSK_TO_MS_BASE ((psa_algorithm_t)0x30000300)
+#define PSA_ALG_TLS12_PSK_TO_MS_BASE            ((psa_algorithm_t)0x20000300)
 /** Macro to build a TLS-1.2 PSK-to-MasterSecret algorithm.
  *
  * In a pure-PSK handshake in TLS 1.2, the master secret is derived
@@ -1240,7 +1263,7 @@
  *                      #PSA_ALG_IS_HASH(\p hash_alg) is true).
  *
  * \return              The corresponding TLS-1.2 PSK to MS algorithm.
- * \return              Unspecified if \p alg is not a supported
+ * \return              Unspecified if \p hash_alg is not a supported
  *                      hash algorithm.
  */
 #define PSA_ALG_TLS12_PSK_TO_MS(hash_alg)                                  \
@@ -1259,55 +1282,67 @@
 #define PSA_ALG_TLS12_PSK_TO_MS_GET_HASH(hkdf_alg)                         \
     (PSA_ALG_CATEGORY_HASH | ((hkdf_alg) & PSA_ALG_HASH_MASK))
 
-#define PSA_ALG_KEY_DERIVATION_MASK             ((psa_algorithm_t)0x010fffff)
+#define PSA_ALG_KEY_DERIVATION_MASK             ((psa_algorithm_t)0x080fffff)
+#define PSA_ALG_KEY_AGREEMENT_MASK              ((psa_algorithm_t)0x10f00000)
 
-/** Use a shared secret as is.
+/** Macro to build a combined algorithm that chains a key agreement with
+ * a key derivation.
  *
- * Specify this algorithm as the selection component of a key agreement
- * to use the raw result of the key agreement as key material.
+ * \param ka_alg        A key agreement algorithm (\c PSA_ALG_XXX value such
+ *                      that #PSA_ALG_IS_KEY_AGREEMENT(\p ka_alg) is true).
+ * \param kdf_alg       A key derivation algorithm (\c PSA_ALG_XXX value such
+ *                      that #PSA_ALG_IS_KEY_DERIVATION(\p kdf_alg) is true).
  *
- * \warning The raw result of a key agreement algorithm such as finite-field
- * Diffie-Hellman or elliptic curve Diffie-Hellman has biases and should
- * not be used directly as key material. It can however be used as the secret
- * input in a key derivation algorithm.
+ * \return              The corresponding key agreement and derivation
+ *                      algorithm.
+ * \return              Unspecified if \p ka_alg is not a supported
+ *                      key agreement algorithm or \p kdf_alg is not a
+ *                      supported key derivation algorithm.
  */
-#define PSA_ALG_SELECT_RAW                      ((psa_algorithm_t)0x31000001)
+#define PSA_ALG_KEY_AGREEMENT(ka_alg, kdf_alg)  \
+    ((ka_alg) | (kdf_alg))
 
 #define PSA_ALG_KEY_AGREEMENT_GET_KDF(alg)                              \
     (((alg) & PSA_ALG_KEY_DERIVATION_MASK) | PSA_ALG_CATEGORY_KEY_DERIVATION)
 
-#define PSA_ALG_KEY_AGREEMENT_GET_BASE(alg)                              \
-    ((alg) & ~PSA_ALG_KEY_DERIVATION_MASK)
+#define PSA_ALG_KEY_AGREEMENT_GET_BASE(alg)                             \
+    (((alg) & PSA_ALG_KEY_AGREEMENT_MASK) | PSA_ALG_CATEGORY_KEY_AGREEMENT)
 
-#define PSA_ALG_FFDH_BASE                       ((psa_algorithm_t)0x22100000)
-/** The Diffie-Hellman key agreement algorithm.
+/** Whether the specified algorithm is a raw key agreement algorithm.
  *
- * This algorithm combines the finite-field Diffie-Hellman (DH) key
- * agreement, also known as Diffie-Hellman-Merkle (DHM) key agreement,
- * to produce a shared secret from a private key and the peer's
- * public key, with a key selection or key derivation algorithm to produce
- * one or more shared keys and other shared cryptographic material.
+ * A raw key agreement algorithm is one that does not specify
+ * a key derivation function.
+ * Usually, raw key agreement algorithms are constructed directly with
+ * a \c PSA_ALG_xxx macro while non-raw key agreement algorithms are
+ * constructed with PSA_ALG_KEY_AGREEMENT().
  *
- * The shared secret produced by key agreement and passed as input to the
- * derivation or selection algorithm \p kdf_alg is the shared secret
+ * \param alg An algorithm identifier (value of type #psa_algorithm_t).
+ *
+ * \return 1 if \p alg is a raw key agreement algorithm, 0 otherwise.
+ *         This macro may return either 0 or 1 if \p alg is not a supported
+ *         algorithm identifier.
+ */
+#define PSA_ALG_IS_RAW_KEY_AGREEMENT(alg)                               \
+    (PSA_ALG_IS_KEY_AGREEMENT(alg) &&                                   \
+     PSA_ALG_KEY_AGREEMENT_GET_KDF(alg) == PSA_ALG_CATEGORY_KEY_DERIVATION)
+
+#define PSA_ALG_IS_KEY_DERIVATION_OR_AGREEMENT(alg)     \
+    ((PSA_ALG_IS_KEY_DERIVATION(alg) || PSA_ALG_IS_KEY_AGREEMENT(alg)))
+
+/** The finite-field Diffie-Hellman (DH) key agreement algorithm.
+ *
+ * The shared secret produced by key agreement is
  * `g^{ab}` in big-endian format.
  * It is `ceiling(m / 8)` bytes long where `m` is the size of the prime `p`
  * in bits.
- *
- * \param kdf_alg       A key derivation algorithm (\c PSA_ALG_XXX value such
- *                      that #PSA_ALG_IS_KEY_DERIVATION(\p hash_alg) is true)
- *                      or a key selection algorithm (\c PSA_ALG_XXX value such
- *                      that #PSA_ALG_IS_KEY_SELECTION(\p hash_alg) is true).
- *
- * \return              The Diffie-Hellman algorithm with the specified
- *                      selection or derivation algorithm.
  */
-#define PSA_ALG_FFDH(kdf_alg) \
-    (PSA_ALG_FFDH_BASE | ((kdf_alg) & PSA_ALG_KEY_DERIVATION_MASK))
+#define PSA_ALG_FFDH                            ((psa_algorithm_t)0x30100000)
+
 /** Whether the specified algorithm is a finite field Diffie-Hellman algorithm.
  *
- * This includes every supported key selection or key agreement algorithm
- * for the output of the Diffie-Hellman calculation.
+ * This includes the raw finite field Diffie-Hellman algorithm as well as
+ * finite-field Diffie-Hellman followed by any supporter key derivation
+ * algorithm.
  *
  * \param alg An algorithm identifier (value of type #psa_algorithm_t).
  *
@@ -1316,18 +1351,11 @@
  *         key agreement algorithm identifier.
  */
 #define PSA_ALG_IS_FFDH(alg) \
-    (PSA_ALG_KEY_AGREEMENT_GET_BASE(alg) == PSA_ALG_FFDH_BASE)
+    (PSA_ALG_KEY_AGREEMENT_GET_BASE(alg) == PSA_ALG_FFDH)
 
-#define PSA_ALG_ECDH_BASE                       ((psa_algorithm_t)0x22200000)
 /** The elliptic curve Diffie-Hellman (ECDH) key agreement algorithm.
  *
- * This algorithm combines the elliptic curve Diffie-Hellman key
- * agreement to produce a shared secret from a private key and the peer's
- * public key, with a key selection or key derivation algorithm to produce
- * one or more shared keys and other shared cryptographic material.
- *
- * The shared secret produced by key agreement and passed as input to the
- * derivation or selection algorithm \p kdf_alg is the x-coordinate of
+ * The shared secret produced by key agreement is the x-coordinate of
  * the shared secret point. It is always `ceiling(m / 8)` bytes long where
  * `m` is the bit size associated with the curve, i.e. the bit size of the
  * order of the curve's coordinate field. When `m` is not a multiple of 8,
@@ -1349,22 +1377,15 @@
  *   the shared secret is the x-coordinate of `d_A Q_B = d_B Q_A`
  *   in big-endian byte order.
  *   The bit size is `m` for the field `F_{2^m}`.
- *
- * \param kdf_alg       A key derivation algorithm (\c PSA_ALG_XXX value such
- *                      that #PSA_ALG_IS_KEY_DERIVATION(\p hash_alg) is true)
- *                      or a selection algorithm (\c PSA_ALG_XXX value such
- *                      that #PSA_ALG_IS_KEY_SELECTION(\p hash_alg) is true).
- *
- * \return              The Diffie-Hellman algorithm with the specified
- *                      selection or derivation algorithm.
  */
-#define PSA_ALG_ECDH(kdf_alg) \
-    (PSA_ALG_ECDH_BASE | ((kdf_alg) & PSA_ALG_KEY_DERIVATION_MASK))
+#define PSA_ALG_ECDH                            ((psa_algorithm_t)0x30200000)
+
 /** Whether the specified algorithm is an elliptic curve Diffie-Hellman
  * algorithm.
  *
- * This includes every supported key selection or key agreement algorithm
- * for the output of the Diffie-Hellman calculation.
+ * This includes the raw elliptic curve Diffie-Hellman algorithm as well as
+ * elliptic curve Diffie-Hellman followed by any supporter key derivation
+ * algorithm.
  *
  * \param alg An algorithm identifier (value of type #psa_algorithm_t).
  *
@@ -1374,7 +1395,7 @@
  *         key agreement algorithm identifier.
  */
 #define PSA_ALG_IS_ECDH(alg) \
-    (PSA_ALG_KEY_AGREEMENT_GET_BASE(alg) == PSA_ALG_ECDH_BASE)
+    (PSA_ALG_KEY_AGREEMENT_GET_BASE(alg) == PSA_ALG_ECDH)
 
 /** Whether the specified algorithm encoding is a wildcard.
  *
@@ -1486,4 +1507,34 @@
 
 /**@}*/
 
+/** \defgroup derivation Key derivation
+ * @{
+ */
+
+/** A secret input for key derivation.
+ *
+ * This must be a key of type #PSA_KEY_TYPE_DERIVE.
+ */
+#define PSA_KDF_STEP_SECRET              ((psa_key_derivation_step_t)0x0101)
+
+/** A label for key derivation.
+ *
+ * This must be a direct input.
+ */
+#define PSA_KDF_STEP_LABEL               ((psa_key_derivation_step_t)0x0201)
+
+/** A salt for key derivation.
+ *
+ * This must be a direct input.
+ */
+#define PSA_KDF_STEP_SALT                ((psa_key_derivation_step_t)0x0202)
+
+/** An information string for key derivation.
+ *
+ * This must be a direct input.
+ */
+#define PSA_KDF_STEP_INFO                ((psa_key_derivation_step_t)0x0203)
+
+/**@}*/
+
 #endif /* PSA_CRYPTO_VALUES_H */
diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 3b9c78f..f6b034e 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -717,7 +717,7 @@
 }
 
 /* Retrieve an empty key slot (slot with no key data, but possibly
- * with some metadata such as a policy). */
+ * with some metadata such as a policy or domain parameters). */
 static psa_status_t psa_get_empty_key_slot( psa_key_handle_t handle,
                                             psa_key_slot_t **p_slot )
 {
@@ -884,11 +884,18 @@
     return( PSA_SUCCESS );
 }
 
+static void psa_abort_operations_using_key( psa_key_slot_t *slot )
+{
+    /*TODO*/
+    (void) slot;
+}
+
 /** Completely wipe a slot in memory, including its policy.
  * Persistent storage is not affected. */
 psa_status_t psa_wipe_key_slot( psa_key_slot_t *slot )
 {
     psa_status_t status = psa_remove_key_data_from_memory( slot );
+    psa_abort_operations_using_key( slot );
     /* At this point, key material and other type-specific content has
      * been wiped. Clear remaining metadata. We can call memset and not
      * zeroize because the metadata is not particularly sensitive. */
@@ -3309,7 +3316,7 @@
     uint8_t tag_length;
 } aead_operation_t;
 
-static void psa_aead_abort( aead_operation_t *operation )
+static void psa_aead_abort_internal( aead_operation_t *operation )
 {
     switch( operation->core_alg )
     {
@@ -3396,7 +3403,7 @@
     return( PSA_SUCCESS );
 
 cleanup:
-    psa_aead_abort( operation );
+    psa_aead_abort_internal( operation );
     return( status );
 }
 
@@ -3467,7 +3474,7 @@
         memset( ciphertext, 0, ciphertext_size );
 
 exit:
-    psa_aead_abort( &operation );
+    psa_aead_abort_internal( &operation );
     if( status == PSA_SUCCESS )
         *ciphertext_length = plaintext_length + operation.tag_length;
     return( status );
@@ -3564,7 +3571,7 @@
         memset( plaintext, 0, plaintext_size );
 
 exit:
-    psa_aead_abort( &operation );
+    psa_aead_abort_internal( &operation );
     if( status == PSA_SUCCESS )
         *plaintext_length = ciphertext_length - operation.tag_length;
     return( status );
@@ -3576,17 +3583,33 @@
 /* Generators */
 /****************************************************************/
 
+#define HKDF_STATE_INIT 0 /* no input yet */
+#define HKDF_STATE_STARTED 1 /* got salt */
+#define HKDF_STATE_KEYED 2 /* got key */
+#define HKDF_STATE_OUTPUT 3 /* output started */
+
+static psa_algorithm_t psa_generator_get_kdf_alg(
+    const psa_crypto_generator_t *generator )
+{
+    if ( PSA_ALG_IS_KEY_AGREEMENT( generator->alg ) )
+        return( PSA_ALG_KEY_AGREEMENT_GET_KDF( generator->alg ) );
+    else
+        return( generator->alg );
+}
+
+
 psa_status_t psa_generator_abort( psa_crypto_generator_t *generator )
 {
     psa_status_t status = PSA_SUCCESS;
-    if( generator->alg == 0 )
+    psa_algorithm_t kdf_alg = psa_generator_get_kdf_alg( generator );
+    if( kdf_alg == 0 )
     {
         /* The object has (apparently) been initialized but it is not
          * in use. It's ok to call abort on such an object, and there's
          * nothing to do. */
     }
     else
-    if( generator->alg == PSA_ALG_SELECT_RAW )
+    if( kdf_alg == PSA_ALG_SELECT_RAW )
     {
         if( generator->ctx.buffer.data != NULL )
         {
@@ -3597,14 +3620,14 @@
     }
     else
 #if defined(MBEDTLS_MD_C)
-    if( PSA_ALG_IS_HKDF( generator->alg ) )
+    if( PSA_ALG_IS_HKDF( kdf_alg ) )
     {
         mbedtls_free( generator->ctx.hkdf.info );
         status = psa_hmac_abort_internal( &generator->ctx.hkdf.hmac );
     }
-    else if( PSA_ALG_IS_TLS12_PRF( generator->alg ) ||
+    else if( PSA_ALG_IS_TLS12_PRF( kdf_alg ) ||
              /* TLS-1.2 PSK-to-MS KDF uses the same generator as TLS-1.2 PRF */
-             PSA_ALG_IS_TLS12_PSK_TO_MS( generator->alg ) )
+             PSA_ALG_IS_TLS12_PSK_TO_MS( kdf_alg ) )
     {
         if( generator->ctx.tls12_prf.key != NULL )
         {
@@ -3629,7 +3652,6 @@
     return( status );
 }
 
-
 psa_status_t psa_get_generator_capacity(const psa_crypto_generator_t *generator,
                                         size_t *capacity)
 {
@@ -3643,6 +3665,17 @@
     return( PSA_SUCCESS );
 }
 
+psa_status_t psa_set_generator_capacity( psa_crypto_generator_t *generator,
+                                         size_t capacity )
+{
+    if( generator->alg == 0 )
+        return( PSA_ERROR_BAD_STATE );
+    if( capacity > generator->capacity )
+        return( PSA_ERROR_INVALID_ARGUMENT );
+    generator->capacity = capacity;
+    return( PSA_SUCCESS );
+}
+
 #if defined(MBEDTLS_MD_C)
 /* Read some bytes from an HKDF-based generator. This performs a chunk
  * of the expand phase of the HKDF algorithm. */
@@ -3654,6 +3687,10 @@
     uint8_t hash_length = PSA_HASH_SIZE( hash_alg );
     psa_status_t status;
 
+    if( hkdf->state < HKDF_STATE_KEYED || ! hkdf->info_set )
+        return( PSA_ERROR_BAD_STATE );
+    hkdf->state = HKDF_STATE_OUTPUT;
+
     while( output_length != 0 )
     {
         /* Copy what remains of the current block */
@@ -3867,6 +3904,7 @@
                                  size_t output_length )
 {
     psa_status_t status;
+    psa_algorithm_t kdf_alg = psa_generator_get_kdf_alg( generator );
 
     if( generator->alg == 0 )
     {
@@ -3894,7 +3932,7 @@
     }
     generator->capacity -= output_length;
 
-    if( generator->alg == PSA_ALG_SELECT_RAW )
+    if( kdf_alg == PSA_ALG_SELECT_RAW )
     {
         /* Initially, the capacity of a selection generator is always
          * the size of the buffer, i.e. `generator->ctx.buffer.size`,
@@ -3912,17 +3950,17 @@
     }
     else
 #if defined(MBEDTLS_MD_C)
-    if( PSA_ALG_IS_HKDF( generator->alg ) )
+    if( PSA_ALG_IS_HKDF( kdf_alg ) )
     {
-        psa_algorithm_t hash_alg = PSA_ALG_HKDF_GET_HASH( generator->alg );
+        psa_algorithm_t hash_alg = PSA_ALG_HKDF_GET_HASH( kdf_alg );
         status = psa_generator_hkdf_read( &generator->ctx.hkdf, hash_alg,
                                           output, output_length );
     }
-    else if( PSA_ALG_IS_TLS12_PRF( generator->alg ) ||
-             PSA_ALG_IS_TLS12_PSK_TO_MS( generator->alg ) )
+    else if( PSA_ALG_IS_TLS12_PRF( kdf_alg ) ||
+             PSA_ALG_IS_TLS12_PSK_TO_MS( kdf_alg ) )
     {
         status = psa_generator_tls12_prf_read( &generator->ctx.tls12_prf,
-                                               generator->alg, output,
+                                               kdf_alg, output,
                                                output_length );
     }
     else
@@ -4035,6 +4073,8 @@
             return( PSA_ERROR_INSUFFICIENT_MEMORY );
         memcpy( hkdf->info, label, label_length );
     }
+    hkdf->state = HKDF_STATE_KEYED;
+    hkdf->info_set = 1;
     return( PSA_SUCCESS );
 }
 #endif /* MBEDTLS_MD_C */
@@ -4278,6 +4318,232 @@
     return( status );
 }
 
+static psa_status_t psa_key_derivation_setup_kdf(
+    psa_crypto_generator_t *generator,
+    psa_algorithm_t kdf_alg )
+{
+    /* Make sure that kdf_alg is a supported key derivation algorithm. */
+#if defined(MBEDTLS_MD_C)
+    if( PSA_ALG_IS_HKDF( kdf_alg ) ||
+        PSA_ALG_IS_TLS12_PRF( kdf_alg ) ||
+        PSA_ALG_IS_TLS12_PSK_TO_MS( kdf_alg ) )
+    {
+        psa_algorithm_t hash_alg = PSA_ALG_HKDF_GET_HASH( kdf_alg );
+        size_t hash_size = PSA_HASH_SIZE( hash_alg );
+        if( hash_size == 0 )
+            return( PSA_ERROR_NOT_SUPPORTED );
+        if( ( PSA_ALG_IS_TLS12_PRF( kdf_alg ) ||
+              PSA_ALG_IS_TLS12_PSK_TO_MS( kdf_alg ) ) &&
+            ! ( hash_alg == PSA_ALG_SHA_256 && hash_alg == PSA_ALG_SHA_384 ) )
+        {
+            return( PSA_ERROR_NOT_SUPPORTED );
+        }
+        generator->capacity = 255 * hash_size;
+        return( PSA_SUCCESS );
+    }
+#endif /* MBEDTLS_MD_C */
+    else
+        return( PSA_ERROR_NOT_SUPPORTED );
+}
+
+psa_status_t psa_key_derivation_setup( psa_crypto_generator_t *generator,
+                                       psa_algorithm_t alg )
+{
+    psa_status_t status;
+
+    if( generator->alg != 0 )
+        return( PSA_ERROR_BAD_STATE );
+
+    if( PSA_ALG_IS_RAW_KEY_AGREEMENT( alg ) )
+        return( PSA_ERROR_INVALID_ARGUMENT );
+    else if( PSA_ALG_IS_KEY_AGREEMENT( alg ) )
+    {
+        psa_algorithm_t kdf_alg = PSA_ALG_KEY_AGREEMENT_GET_KDF( alg );
+        status = psa_key_derivation_setup_kdf( generator, kdf_alg );
+    }
+    else if( PSA_ALG_IS_KEY_DERIVATION( alg ) )
+    {
+        status = psa_key_derivation_setup_kdf( generator, alg );
+    }
+    else
+        return( PSA_ERROR_INVALID_ARGUMENT );
+
+    if( status == PSA_SUCCESS )
+        generator->alg = alg;
+    return( status );
+}
+
+#if defined(MBEDTLS_MD_C)
+static psa_status_t psa_hkdf_input( psa_hkdf_generator_t *hkdf,
+                                    psa_algorithm_t hash_alg,
+                                    psa_key_derivation_step_t step,
+                                    const uint8_t *data,
+                                    size_t data_length )
+{
+    psa_status_t status;
+    switch( step )
+    {
+        case PSA_KDF_STEP_SALT:
+            if( hkdf->state == HKDF_STATE_INIT )
+            {
+                status = psa_hmac_setup_internal( &hkdf->hmac,
+                                                  data, data_length,
+                                                  hash_alg );
+                if( status != PSA_SUCCESS )
+                    return( status );
+                hkdf->state = HKDF_STATE_STARTED;
+                return( PSA_SUCCESS );
+            }
+            else
+                return( PSA_ERROR_BAD_STATE );
+            break;
+        case PSA_KDF_STEP_SECRET:
+            /* If no salt was provided, use an empty salt. */
+            if( hkdf->state == HKDF_STATE_INIT )
+            {
+                status = psa_hmac_setup_internal( &hkdf->hmac,
+                                                  NULL, 0,
+                                                  PSA_ALG_HMAC( hash_alg ) );
+                if( status != PSA_SUCCESS )
+                    return( status );
+                hkdf->state = HKDF_STATE_STARTED;
+            }
+            if( hkdf->state == HKDF_STATE_STARTED )
+            {
+                status = psa_hash_update( &hkdf->hmac.hash_ctx,
+                                          data, data_length );
+                if( status != PSA_SUCCESS )
+                    return( status );
+                status = psa_hmac_finish_internal( &hkdf->hmac,
+                                                   hkdf->prk,
+                                                   sizeof( hkdf->prk ) );
+                if( status != PSA_SUCCESS )
+                    return( status );
+                hkdf->offset_in_block = PSA_HASH_SIZE( hash_alg );
+                hkdf->block_number = 0;
+                hkdf->state = HKDF_STATE_KEYED;
+                return( PSA_SUCCESS );
+            }
+            else
+                return( PSA_ERROR_BAD_STATE );
+            break;
+        case PSA_KDF_STEP_INFO:
+            if( hkdf->state == HKDF_STATE_OUTPUT )
+                return( PSA_ERROR_BAD_STATE );
+            if( hkdf->info_set )
+                return( PSA_ERROR_BAD_STATE );
+            hkdf->info_length = data_length;
+            if( data_length != 0 )
+            {
+                hkdf->info = mbedtls_calloc( 1, data_length );
+                if( hkdf->info == NULL )
+                    return( PSA_ERROR_INSUFFICIENT_MEMORY );
+                memcpy( hkdf->info, data, data_length );
+            }
+            hkdf->info_set = 1;
+            return( PSA_SUCCESS );
+        default:
+            return( PSA_ERROR_INVALID_ARGUMENT );
+    }
+}
+#endif /* MBEDTLS_MD_C */
+
+static psa_status_t psa_key_derivation_input_raw(
+    psa_crypto_generator_t *generator,
+    psa_key_derivation_step_t step,
+    const uint8_t *data,
+    size_t data_length )
+{
+    psa_status_t status;
+    psa_algorithm_t kdf_alg = psa_generator_get_kdf_alg( generator );
+
+    if( kdf_alg == PSA_ALG_SELECT_RAW )
+    {
+        if( generator->capacity != 0 )
+            return( PSA_ERROR_INVALID_ARGUMENT );
+        generator->ctx.buffer.data = mbedtls_calloc( 1, data_length );
+        if( generator->ctx.buffer.data == NULL )
+            return( PSA_ERROR_INSUFFICIENT_MEMORY );
+        memcpy( generator->ctx.buffer.data, data, data_length );
+        generator->ctx.buffer.size = data_length;
+        generator->capacity = data_length;
+        status = PSA_SUCCESS;
+    }
+    else
+#if defined(MBEDTLS_MD_C)
+    if( PSA_ALG_IS_HKDF( kdf_alg ) )
+    {
+        status = psa_hkdf_input( &generator->ctx.hkdf,
+                                 PSA_ALG_HKDF_GET_HASH( kdf_alg ),
+                                 step, data, data_length );
+    }
+    else
+#endif /* MBEDTLS_MD_C */
+#if defined(MBEDTLS_MD_C)
+    /* TLS-1.2 PRF and TLS-1.2 PSK-to-MS are very similar, so share code. */
+    if( PSA_ALG_IS_TLS12_PRF( kdf_alg ) ||
+             PSA_ALG_IS_TLS12_PSK_TO_MS( kdf_alg ) )
+    {
+        // TODO
+        status = PSA_ERROR_NOT_SUPPORTED;
+    }
+    else
+#endif /* MBEDTLS_MD_C */
+    {
+        /* This can't happen unless the generator object was not initialized */
+        return( PSA_ERROR_BAD_STATE );
+    }
+
+    if( status != PSA_SUCCESS )
+        psa_generator_abort( generator );
+    return( status );
+}
+
+psa_status_t psa_key_derivation_input_bytes( psa_crypto_generator_t *generator,
+                                             psa_key_derivation_step_t step,
+                                             const uint8_t *data,
+                                             size_t data_length )
+{
+    switch( step )
+    {
+        case PSA_KDF_STEP_LABEL:
+        case PSA_KDF_STEP_SALT:
+        case PSA_KDF_STEP_INFO:
+            return( psa_key_derivation_input_raw( generator, step,
+                                                  data, data_length ) );
+        default:
+            return( PSA_ERROR_INVALID_ARGUMENT );
+    }
+}
+
+psa_status_t psa_key_derivation_input_key( psa_crypto_generator_t *generator,
+                                           psa_key_derivation_step_t step,
+                                           psa_key_handle_t handle )
+{
+    psa_key_slot_t *slot;
+    psa_status_t status;
+    status = psa_get_key_from_slot( handle, &slot,
+                                    PSA_KEY_USAGE_DERIVE,
+                                    generator->alg );
+    if( status != PSA_SUCCESS )
+        return( status );
+    // TODO: for a key agreement algorithm, allow the corresponding key type and step
+    if( slot->type != PSA_KEY_TYPE_DERIVE )
+        return( PSA_ERROR_INVALID_ARGUMENT );
+    /* Don't allow a key to be used as an input that is usually public.
+     * This is debatable. It's ok from a cryptographic perspective to
+     * use secret material as an input that is usually public. However
+     * the material should be dedicated to a particular input step,
+     * otherwise this may allow the key to be used in an unintended way
+     * and leak values derived from the key. So be conservative. */
+    if( step != PSA_KDF_STEP_SECRET )
+        return( PSA_ERROR_INVALID_ARGUMENT );
+    return( psa_key_derivation_input_raw( generator,
+                                          step,
+                                          slot->data.raw.data,
+                                          slot->data.raw.bytes ) );
+}
+
 
 
 /****************************************************************/
@@ -4334,10 +4600,10 @@
  * to potentially free embedded data structures and wipe confidential data.
  */
 static psa_status_t psa_key_agreement_internal( psa_crypto_generator_t *generator,
+                                                psa_key_derivation_step_t step,
                                                 psa_key_slot_t *private_key,
                                                 const uint8_t *peer_key,
-                                                size_t peer_key_length,
-                                                psa_algorithm_t alg )
+                                                size_t peer_key_length )
 {
     psa_status_t status;
     uint8_t shared_secret[PSA_KEY_AGREEMENT_MAX_SHARED_SECRET_SIZE];
@@ -4345,10 +4611,10 @@
 
     /* Step 1: run the secret agreement algorithm to generate the shared
      * secret. */
-    switch( PSA_ALG_KEY_AGREEMENT_GET_BASE( alg ) )
+    switch( PSA_ALG_KEY_AGREEMENT_GET_BASE( generator->alg ) )
     {
 #if defined(MBEDTLS_ECDH_C)
-        case PSA_ALG_ECDH_BASE:
+        case PSA_ALG_ECDH:
             if( ! PSA_KEY_TYPE_IS_ECC_KEYPAIR( private_key->type ) )
                 return( PSA_ERROR_INVALID_ARGUMENT );
             status = psa_key_agreement_ecdh( peer_key, peer_key_length,
@@ -4369,34 +4635,31 @@
 
     /* Step 2: set up the key derivation to generate key material from
      * the shared secret. */
-    status = psa_key_derivation_internal( generator,
-                                          shared_secret, shared_secret_length,
-                                          PSA_ALG_KEY_AGREEMENT_GET_KDF( alg ),
-                                          NULL, 0, NULL, 0,
-                                          PSA_GENERATOR_UNBRIDLED_CAPACITY );
+    status = psa_key_derivation_input_raw( generator, step,
+                                           shared_secret, shared_secret_length );
+
 exit:
     mbedtls_platform_zeroize( shared_secret, shared_secret_length );
     return( status );
 }
 
 psa_status_t psa_key_agreement( psa_crypto_generator_t *generator,
+                                psa_key_derivation_step_t step,
                                 psa_key_handle_t private_key,
                                 const uint8_t *peer_key,
-                                size_t peer_key_length,
-                                psa_algorithm_t alg )
+                                size_t peer_key_length )
 {
     psa_key_slot_t *slot;
     psa_status_t status;
-    if( ! PSA_ALG_IS_KEY_AGREEMENT( alg ) )
+    if( ! PSA_ALG_IS_KEY_AGREEMENT( generator->alg ) )
         return( PSA_ERROR_INVALID_ARGUMENT );
     status = psa_get_key_from_slot( private_key, &slot,
-                                    PSA_KEY_USAGE_DERIVE, alg );
+                                    PSA_KEY_USAGE_DERIVE, generator->alg );
     if( status != PSA_SUCCESS )
         return( status );
-    status = psa_key_agreement_internal( generator,
+    status = psa_key_agreement_internal( generator, step,
                                          slot,
-                                         peer_key, peer_key_length,
-                                         alg );
+                                         peer_key, peer_key_length );
     if( status != PSA_SUCCESS )
         psa_generator_abort( generator );
     return( status );
diff --git a/tests/suites/test_suite_psa_crypto.data b/tests/suites/test_suite_psa_crypto.data
index 65ac6d7..bd605ac 100644
--- a/tests/suites/test_suite_psa_crypto.data
+++ b/tests/suites/test_suite_psa_crypto.data
@@ -482,15 +482,15 @@
 
 PSA key policy: agreement, permitted
 depends_on:MBEDTLS_PK_PARSE_C:MBEDTLS_ECP_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED:MBEDTLS_ECDH_C
-agreement_key_policy:PSA_KEY_USAGE_DERIVE:PSA_ALG_ECDH(PSA_ALG_SELECT_RAW):PSA_KEY_TYPE_ECC_KEYPAIR(PSA_ECC_CURVE_SECP256R1):"49c9a8c18c4b885638c431cf1df1c994131609b580d4fd43a0cab17db2f13eee":PSA_ALG_ECDH(PSA_ALG_SELECT_RAW)
+agreement_key_policy:PSA_KEY_USAGE_DERIVE:PSA_ALG_ECDH:PSA_KEY_TYPE_ECC_KEYPAIR(PSA_ECC_CURVE_SECP256R1):"49c9a8c18c4b885638c431cf1df1c994131609b580d4fd43a0cab17db2f13eee":PSA_ALG_ECDH
 
 PSA key policy: agreement, not permitted
 depends_on:MBEDTLS_PK_PARSE_C:MBEDTLS_ECP_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED:MBEDTLS_ECDH_C
-agreement_key_policy:0:PSA_ALG_ECDH(PSA_ALG_SELECT_RAW):PSA_KEY_TYPE_ECC_KEYPAIR(PSA_ECC_CURVE_SECP256R1):"49c9a8c18c4b885638c431cf1df1c994131609b580d4fd43a0cab17db2f13eee":PSA_ALG_ECDH(PSA_ALG_SELECT_RAW)
+agreement_key_policy:0:PSA_ALG_ECDH:PSA_KEY_TYPE_ECC_KEYPAIR(PSA_ECC_CURVE_SECP256R1):"49c9a8c18c4b885638c431cf1df1c994131609b580d4fd43a0cab17db2f13eee":PSA_ALG_ECDH
 
 PSA key policy: agreement, wrong algorithm
 depends_on:MBEDTLS_PK_PARSE_C:MBEDTLS_ECP_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED:MBEDTLS_ECDH_C
-agreement_key_policy:PSA_KEY_USAGE_DERIVE:PSA_ALG_ECDH(PSA_ALG_SELECT_RAW):PSA_KEY_TYPE_ECC_KEYPAIR(PSA_ECC_CURVE_SECP256R1):"49c9a8c18c4b885638c431cf1df1c994131609b580d4fd43a0cab17db2f13eee":PSA_ALG_FFDH(PSA_ALG_SELECT_RAW)
+agreement_key_policy:PSA_KEY_USAGE_DERIVE:PSA_ALG_ECDH:PSA_KEY_TYPE_ECC_KEYPAIR(PSA_ECC_CURVE_SECP256R1):"49c9a8c18c4b885638c431cf1df1c994131609b580d4fd43a0cab17db2f13eee":PSA_ALG_FFDH
 
 Copy key: raw, 0 bytes
 copy_key_policy:0:0:PSA_KEY_TYPE_RAW_DATA:"":0:0:-1:-1:0:0
@@ -1379,7 +1379,7 @@
 
 PSA import/exercise: ECP SECP256R1 keypair, ECDH
 depends_on:MBEDTLS_PK_PARSE_C:MBEDTLS_ECP_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED:MBEDTLS_ECDH_C
-import_and_exercise_key:"49c9a8c18c4b885638c431cf1df1c994131609b580d4fd43a0cab17db2f13eee":PSA_KEY_TYPE_ECC_KEYPAIR(PSA_ECC_CURVE_SECP256R1):256:PSA_ALG_ECDH(PSA_ALG_SELECT_RAW)
+import_and_exercise_key:"49c9a8c18c4b885638c431cf1df1c994131609b580d4fd43a0cab17db2f13eee":PSA_KEY_TYPE_ECC_KEYPAIR(PSA_ECC_CURVE_SECP256R1):256:PSA_ALG_ECDH
 
 PSA sign: RSA PKCS#1 v1.5, raw
 depends_on:MBEDTLS_PK_PARSE_C:MBEDTLS_RSA_C:MBEDTLS_PKCS1_V15
diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index 4cec118..d70b7eb 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -529,11 +529,30 @@
 
     if( usage & PSA_KEY_USAGE_DERIVE )
     {
-        PSA_ASSERT( psa_key_derivation( &generator,
-                                        handle, alg,
-                                        label, label_length,
-                                        seed, seed_length,
-                                        sizeof( output ) ) );
+        if( PSA_ALG_IS_HKDF( alg ) )
+        {
+            PSA_ASSERT( psa_key_derivation_setup( &generator, alg ) );
+            PSA_ASSERT( psa_key_derivation_input_bytes( &generator,
+                                                        PSA_KDF_STEP_SALT,
+                                                        label,
+                                                        label_length ) );
+            PSA_ASSERT( psa_key_derivation_input_key( &generator,
+                                                      PSA_KDF_STEP_SECRET,
+                                                      handle ) );
+            PSA_ASSERT( psa_key_derivation_input_bytes( &generator,
+                                                        PSA_KDF_STEP_INFO,
+                                                        seed,
+                                                        seed_length ) );
+        }
+        else
+        {
+            // legacy
+            PSA_ASSERT( psa_key_derivation( &generator,
+                                            handle, alg,
+                                            label, label_length,
+                                            seed, seed_length,
+                                            sizeof( output ) ) );
+        }
         PSA_ASSERT( psa_generator_read( &generator,
                                         output,
                                         sizeof( output ) ) );
@@ -549,8 +568,7 @@
 /* We need two keys to exercise key agreement. Exercise the
  * private key against its own public key. */
 static psa_status_t key_agreement_with_self( psa_crypto_generator_t *generator,
-                                             psa_key_handle_t handle,
-                                             psa_algorithm_t alg )
+                                             psa_key_handle_t handle )
 {
     psa_key_type_t private_key_type;
     psa_key_type_t public_key_type;
@@ -572,9 +590,8 @@
                                        public_key, public_key_length,
                                        &public_key_length ) );
 
-    status = psa_key_agreement( generator, handle,
-                                public_key, public_key_length,
-                                alg );
+    status = psa_key_agreement( generator, PSA_KDF_STEP_SECRET, handle,
+                                public_key, public_key_length );
 exit:
     mbedtls_free( public_key );
     return( status );
@@ -592,7 +609,8 @@
     {
         /* We need two keys to exercise key agreement. Exercise the
          * private key against its own public key. */
-        PSA_ASSERT( key_agreement_with_self( &generator, handle, alg ) );
+        PSA_ASSERT( psa_key_derivation_setup( &generator, alg ) );
+        PSA_ASSERT( key_agreement_with_self( &generator, handle ) );
         PSA_ASSERT( psa_generator_read( &generator,
                                         output,
                                         sizeof( output ) ) );
@@ -1917,7 +1935,8 @@
     PSA_ASSERT( psa_import_key( handle, key_type,
                                 key_data->x, key_data->len ) );
 
-    status = key_agreement_with_self( &generator, handle, exercise_alg );
+    PSA_ASSERT( psa_key_derivation_setup( &generator, exercise_alg ) );
+    status = key_agreement_with_self( &generator, handle );
 
     if( policy_alg == exercise_alg &&
         ( policy_usage & PSA_KEY_USAGE_DERIVE ) != 0 )
@@ -4241,10 +4260,29 @@
                                 key_data->len ) );
 
     /* Extraction phase. */
-    PSA_ASSERT( psa_key_derivation( &generator, handle, alg,
-                                    salt->x, salt->len,
-                                    label->x, label->len,
-                                    requested_capacity ) );
+    if( PSA_ALG_IS_HKDF( alg ) )
+    {
+        PSA_ASSERT( psa_key_derivation_setup( &generator, alg ) );
+        PSA_ASSERT( psa_set_generator_capacity( &generator,
+                                                requested_capacity ) );
+        PSA_ASSERT( psa_key_derivation_input_bytes( &generator,
+                                                    PSA_KDF_STEP_SALT,
+                                                    salt->x, salt->len ) );
+        PSA_ASSERT( psa_key_derivation_input_key( &generator,
+                                                  PSA_KDF_STEP_SECRET,
+                                                  handle ) );
+        PSA_ASSERT( psa_key_derivation_input_bytes( &generator,
+                                                    PSA_KDF_STEP_INFO,
+                                                    label->x, label->len ) );
+    }
+    else
+    {
+        // legacy
+        PSA_ASSERT( psa_key_derivation( &generator, handle, alg,
+                                        salt->x, salt->len,
+                                        label->x, label->len,
+                                        requested_capacity ) );
+    }
     PSA_ASSERT( psa_get_generator_capacity( &generator,
                                             &current_capacity ) );
     TEST_EQUAL( current_capacity, requested_capacity );
@@ -4319,10 +4357,29 @@
                                 key_data->len ) );
 
     /* Extraction phase. */
-    PSA_ASSERT( psa_key_derivation( &generator, handle, alg,
-                                    salt->x, salt->len,
-                                    label->x, label->len,
-                                    requested_capacity ) );
+    if( PSA_ALG_IS_HKDF( alg ) )
+    {
+        PSA_ASSERT( psa_key_derivation_setup( &generator, alg ) );
+        PSA_ASSERT( psa_set_generator_capacity( &generator,
+                                                requested_capacity ) );
+        PSA_ASSERT( psa_key_derivation_input_bytes( &generator,
+                                                    PSA_KDF_STEP_SALT,
+                                                    salt->x, salt->len ) );
+        PSA_ASSERT( psa_key_derivation_input_key( &generator,
+                                                  PSA_KDF_STEP_SECRET,
+                                                  handle ) );
+        PSA_ASSERT( psa_key_derivation_input_bytes( &generator,
+                                                    PSA_KDF_STEP_INFO,
+                                                    label->x, label->len ) );
+    }
+    else
+    {
+        // legacy
+        PSA_ASSERT( psa_key_derivation( &generator, handle, alg,
+                                        salt->x, salt->len,
+                                        label->x, label->len,
+                                        requested_capacity ) );
+    }
     PSA_ASSERT( psa_get_generator_capacity( &generator,
                                             &current_capacity ) );
     TEST_EQUAL( current_capacity, expected_capacity );
@@ -4524,10 +4581,10 @@
                                 our_key_data->x,
                                 our_key_data->len ) );
 
-    TEST_EQUAL( psa_key_agreement( &generator,
+    PSA_ASSERT( psa_key_derivation_setup( &generator, alg ) );
+    TEST_EQUAL( psa_key_agreement( &generator, PSA_KDF_STEP_SECRET,
                                    our_key,
-                                   peer_key_data->x, peer_key_data->len,
-                                   alg ),
+                                   peer_key_data->x, peer_key_data->len ),
                 expected_status_arg );
 
 exit:
@@ -4560,10 +4617,10 @@
                                 our_key_data->x,
                                 our_key_data->len ) );
 
-    PSA_ASSERT( psa_key_agreement( &generator,
+    PSA_ASSERT( psa_key_derivation_setup( &generator, alg ) );
+    PSA_ASSERT( psa_key_agreement( &generator, PSA_KDF_STEP_SECRET,
                                    our_key,
-                                   peer_key_data->x, peer_key_data->len,
-                                   alg ) );
+                                   peer_key_data->x, peer_key_data->len ) );
 
     /* Test the advertized capacity. */
     PSA_ASSERT( psa_get_generator_capacity(
@@ -4614,10 +4671,10 @@
                                 our_key_data->x,
                                 our_key_data->len ) );
 
-    PSA_ASSERT( psa_key_agreement( &generator,
+    PSA_ASSERT( psa_key_derivation_setup( &generator, alg ) );
+    PSA_ASSERT( psa_key_agreement( &generator, PSA_KDF_STEP_SECRET,
                                    our_key,
-                                   peer_key_data->x, peer_key_data->len,
-                                   alg ) );
+                                   peer_key_data->x, peer_key_data->len ) );
 
     PSA_ASSERT( psa_generator_read( &generator,
                                     actual_output,
diff --git a/tests/suites/test_suite_psa_crypto_metadata.data b/tests/suites/test_suite_psa_crypto_metadata.data
index 1e7a996..edb09a8 100644
--- a/tests/suites/test_suite_psa_crypto_metadata.data
+++ b/tests/suites/test_suite_psa_crypto_metadata.data
@@ -266,24 +266,21 @@
 depends_on:MBEDTLS_SHA256_C
 key_derivation_algorithm:PSA_ALG_HKDF( PSA_ALG_SHA_256 ):ALG_IS_HKDF
 
-Key selection: raw
-key_selection_algorithm:PSA_ALG_SELECT_RAW:0
-
 Key agreement: FFDH, raw output
 depends_on:MBEDTLS_DHM_C
-key_agreement_algorithm:PSA_ALG_FFDH( PSA_ALG_SELECT_RAW ):ALG_IS_FFDH:PSA_ALG_SELECT_RAW
+key_agreement_algorithm:PSA_ALG_FFDH:ALG_IS_FFDH | ALG_IS_RAW_KEY_AGREEMENT:PSA_ALG_FFDH:PSA_ALG_CATEGORY_KEY_DERIVATION
 
 Key agreement: FFDH, HKDF using SHA-256
 depends_on:MBEDTLS_DHM_C
-key_agreement_algorithm:PSA_ALG_FFDH( PSA_ALG_HKDF( PSA_ALG_SHA_256 ) ):ALG_IS_FFDH:PSA_ALG_HKDF( PSA_ALG_SHA_256 )
+key_agreement_algorithm:PSA_ALG_KEY_AGREEMENT( PSA_ALG_FFDH, PSA_ALG_HKDF( PSA_ALG_SHA_256 ) ):ALG_IS_FFDH:PSA_ALG_FFDH:PSA_ALG_HKDF( PSA_ALG_SHA_256 )
 
 Key agreement: ECDH, raw output
 depends_on:MBEDTLS_ECDH_C
-key_agreement_algorithm:PSA_ALG_ECDH( PSA_ALG_SELECT_RAW ):ALG_IS_ECDH:PSA_ALG_SELECT_RAW
+key_agreement_algorithm:PSA_ALG_ECDH:ALG_IS_ECDH | ALG_IS_RAW_KEY_AGREEMENT:PSA_ALG_ECDH:PSA_ALG_CATEGORY_KEY_DERIVATION
 
 Key agreement: ECDH, HKDF using SHA-256
 depends_on:MBEDTLS_ECDH_C
-key_agreement_algorithm:PSA_ALG_ECDH( PSA_ALG_HKDF( PSA_ALG_SHA_256 ) ):ALG_IS_ECDH:PSA_ALG_HKDF( PSA_ALG_SHA_256 )
+key_agreement_algorithm:PSA_ALG_KEY_AGREEMENT( PSA_ALG_ECDH, PSA_ALG_HKDF( PSA_ALG_SHA_256 ) ):ALG_IS_ECDH:PSA_ALG_ECDH:PSA_ALG_HKDF( PSA_ALG_SHA_256 )
 
 Key type: raw data
 key_type:PSA_KEY_TYPE_RAW_DATA:KEY_TYPE_IS_UNSTRUCTURED
diff --git a/tests/suites/test_suite_psa_crypto_metadata.function b/tests/suites/test_suite_psa_crypto_metadata.function
index 01c8628..1bc8d64 100644
--- a/tests/suites/test_suite_psa_crypto_metadata.function
+++ b/tests/suites/test_suite_psa_crypto_metadata.function
@@ -35,6 +35,7 @@
 #define ALG_IS_FFDH                     ( 1u << 17 )
 #define ALG_IS_ECDH                     ( 1u << 18 )
 #define ALG_IS_WILDCARD                 ( 1u << 19 )
+#define ALG_IS_RAW_KEY_AGREEMENT        ( 1u << 20 )
 
 /* Flags for key type classification macros. There is a flag for every
  * key type classification macro PSA_KEY_TYPE_IS_xxx except for some that
@@ -73,6 +74,9 @@
     TEST_CLASSIFICATION_MACRO( ALG_IS_RSA_OAEP, alg, flags );
     TEST_CLASSIFICATION_MACRO( ALG_IS_HKDF, alg, flags );
     TEST_CLASSIFICATION_MACRO( ALG_IS_WILDCARD, alg, flags );
+    TEST_CLASSIFICATION_MACRO( ALG_IS_ECDH, alg, flags );
+    TEST_CLASSIFICATION_MACRO( ALG_IS_FFDH, alg, flags );
+    TEST_CLASSIFICATION_MACRO( ALG_IS_RAW_KEY_AGREEMENT, alg, flags );
 exit: ;
 }
 
@@ -113,7 +117,6 @@
     TEST_ASSERT( ! PSA_ALG_IS_ASYMMETRIC_ENCRYPTION( alg ) );
     TEST_ASSERT( ! PSA_ALG_IS_KEY_AGREEMENT( alg ) );
     TEST_ASSERT( ! PSA_ALG_IS_KEY_DERIVATION( alg ) );
-    TEST_ASSERT( ! PSA_ALG_IS_KEY_SELECTION( alg ) );
     algorithm_classification( alg, classification_flags );
 
     /* Length */
@@ -134,7 +137,6 @@
     TEST_ASSERT( ! PSA_ALG_IS_ASYMMETRIC_ENCRYPTION( alg ) );
     TEST_ASSERT( ! PSA_ALG_IS_KEY_AGREEMENT( alg ) );
     TEST_ASSERT( ! PSA_ALG_IS_KEY_DERIVATION( alg ) );
-    TEST_ASSERT( ! PSA_ALG_IS_KEY_SELECTION( alg ) );
     algorithm_classification( alg, classification_flags );
 
     /* Tag length */
@@ -174,7 +176,6 @@
     TEST_ASSERT( ! PSA_ALG_IS_ASYMMETRIC_ENCRYPTION( alg ) );
     TEST_ASSERT( ! PSA_ALG_IS_KEY_AGREEMENT( alg ) );
     TEST_ASSERT( ! PSA_ALG_IS_KEY_DERIVATION( alg ) );
-    TEST_ASSERT( ! PSA_ALG_IS_KEY_SELECTION( alg ) );
     algorithm_classification( alg, 0 );
 
     /* Dependent algorithms */
@@ -271,7 +272,6 @@
     TEST_ASSERT( ! PSA_ALG_IS_ASYMMETRIC_ENCRYPTION( alg ) );
     TEST_ASSERT( ! PSA_ALG_IS_KEY_AGREEMENT( alg ) );
     TEST_ASSERT( ! PSA_ALG_IS_KEY_DERIVATION( alg ) );
-    TEST_ASSERT( ! PSA_ALG_IS_KEY_SELECTION( alg ) );
     algorithm_classification( alg, classification_flags );
 }
 /* END_CASE */
@@ -320,7 +320,6 @@
     TEST_ASSERT( ! PSA_ALG_IS_ASYMMETRIC_ENCRYPTION( alg ) );
     TEST_ASSERT( ! PSA_ALG_IS_KEY_AGREEMENT( alg ) );
     TEST_ASSERT( ! PSA_ALG_IS_KEY_DERIVATION( alg ) );
-    TEST_ASSERT( ! PSA_ALG_IS_KEY_SELECTION( alg ) );
     algorithm_classification( alg, classification_flags );
 }
 /* END_CASE */
@@ -350,7 +349,6 @@
     TEST_ASSERT( PSA_ALG_IS_ASYMMETRIC_ENCRYPTION( alg ) );
     TEST_ASSERT( ! PSA_ALG_IS_KEY_AGREEMENT( alg ) );
     TEST_ASSERT( ! PSA_ALG_IS_KEY_DERIVATION( alg ) );
-    TEST_ASSERT( ! PSA_ALG_IS_KEY_SELECTION( alg ) );
     algorithm_classification( alg, classification_flags );
 }
 /* END_CASE */
@@ -359,6 +357,8 @@
 void key_derivation_algorithm( int alg_arg, int classification_flags )
 {
     psa_algorithm_t alg = alg_arg;
+    psa_algorithm_t ecdh_alg = PSA_ALG_KEY_AGREEMENT( PSA_ALG_ECDH, alg );
+    psa_algorithm_t ffdh_alg = PSA_ALG_KEY_AGREEMENT( PSA_ALG_FFDH, alg );
 
     /* Algorithm classification */
     TEST_ASSERT( ! PSA_ALG_IS_HASH( alg ) );
@@ -369,49 +369,25 @@
     TEST_ASSERT( ! PSA_ALG_IS_ASYMMETRIC_ENCRYPTION( alg ) );
     TEST_ASSERT( ! PSA_ALG_IS_KEY_AGREEMENT( alg ) );
     TEST_ASSERT( PSA_ALG_IS_KEY_DERIVATION( alg ) );
-    TEST_ASSERT( ! PSA_ALG_IS_KEY_SELECTION( alg ) );
     algorithm_classification( alg, classification_flags );
 
     /* Check combinations with key agreements */
-    TEST_ASSERT( PSA_ALG_IS_KEY_AGREEMENT( PSA_ALG_FFDH( alg ) ) );
-    TEST_ASSERT( PSA_ALG_IS_KEY_AGREEMENT( PSA_ALG_ECDH( alg ) ) );
-    TEST_EQUAL( PSA_ALG_KEY_AGREEMENT_GET_KDF( PSA_ALG_ECDH( alg ) ), alg );
-    TEST_EQUAL( PSA_ALG_KEY_AGREEMENT_GET_KDF( PSA_ALG_FFDH( alg ) ), alg );
-}
-/* END_CASE */
-
-/* BEGIN_CASE */
-void key_selection_algorithm( int alg_arg, int classification_flags )
-{
-    psa_algorithm_t alg = alg_arg;
-
-    /* Algorithm classification */
-    TEST_ASSERT( ! PSA_ALG_IS_HASH( alg ) );
-    TEST_ASSERT( ! PSA_ALG_IS_MAC( alg ) );
-    TEST_ASSERT( ! PSA_ALG_IS_CIPHER( alg ) );
-    TEST_ASSERT( ! PSA_ALG_IS_AEAD( alg ) );
-    TEST_ASSERT( ! PSA_ALG_IS_SIGN( alg ) );
-    TEST_ASSERT( ! PSA_ALG_IS_ASYMMETRIC_ENCRYPTION( alg ) );
-    TEST_ASSERT( ! PSA_ALG_IS_KEY_AGREEMENT( alg ) );
-    TEST_ASSERT( ! PSA_ALG_IS_KEY_DERIVATION( alg ) );
-    TEST_ASSERT( PSA_ALG_IS_KEY_SELECTION( alg ) );
-    algorithm_classification( alg, classification_flags );
-
-    /* Check combinations with key agreements */
-    TEST_ASSERT( PSA_ALG_IS_KEY_AGREEMENT( PSA_ALG_FFDH( alg ) ) );
-    TEST_ASSERT( PSA_ALG_IS_KEY_AGREEMENT( PSA_ALG_ECDH( alg ) ) );
-    TEST_EQUAL( PSA_ALG_KEY_AGREEMENT_GET_KDF( PSA_ALG_ECDH( alg ) ), alg );
-    TEST_EQUAL( PSA_ALG_KEY_AGREEMENT_GET_KDF( PSA_ALG_FFDH( alg ) ), alg );
+    TEST_ASSERT( PSA_ALG_IS_KEY_AGREEMENT( ecdh_alg ) );
+    TEST_ASSERT( PSA_ALG_IS_KEY_AGREEMENT( ffdh_alg ) );
+    TEST_EQUAL( PSA_ALG_KEY_AGREEMENT_GET_KDF( ecdh_alg ), alg );
+    TEST_EQUAL( PSA_ALG_KEY_AGREEMENT_GET_KDF( ffdh_alg ), alg );
 }
 /* END_CASE */
 
 /* BEGIN_CASE */
 void key_agreement_algorithm( int alg_arg, int classification_flags,
-                              int post_alg_arg )
+                              int ka_alg_arg, int kdf_alg_arg )
 {
     psa_algorithm_t alg = alg_arg;
-    psa_algorithm_t actual_post_alg = PSA_ALG_KEY_AGREEMENT_GET_KDF( alg );
-    psa_algorithm_t expected_post_alg = post_alg_arg;
+    psa_algorithm_t actual_ka_alg = PSA_ALG_KEY_AGREEMENT_GET_BASE( alg );
+    psa_algorithm_t expected_ka_alg = ka_alg_arg;
+    psa_algorithm_t actual_kdf_alg = PSA_ALG_KEY_AGREEMENT_GET_KDF( alg );
+    psa_algorithm_t expected_kdf_alg = kdf_alg_arg;
 
     /* Algorithm classification */
     TEST_ASSERT( ! PSA_ALG_IS_HASH( alg ) );
@@ -422,13 +398,11 @@
     TEST_ASSERT( ! PSA_ALG_IS_ASYMMETRIC_ENCRYPTION( alg ) );
     TEST_ASSERT( PSA_ALG_IS_KEY_AGREEMENT( alg ) );
     TEST_ASSERT( ! PSA_ALG_IS_KEY_DERIVATION( alg ) );
-    TEST_ASSERT( ! PSA_ALG_IS_KEY_SELECTION( alg ) );
     algorithm_classification( alg, classification_flags );
 
     /* Shared secret derivation properties */
-    TEST_ASSERT( PSA_ALG_IS_KEY_DERIVATION( actual_post_alg ) ||
-                 PSA_ALG_IS_KEY_SELECTION( actual_post_alg ) );
-    TEST_EQUAL( actual_post_alg, expected_post_alg );
+    TEST_EQUAL( actual_ka_alg, expected_ka_alg );
+    TEST_EQUAL( actual_kdf_alg, expected_kdf_alg );
 }
 /* END_CASE */