diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 2583735..6598cf4 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -3214,6 +3214,255 @@
     return( status );
 }
 
+/* Set the key for a multipart authenticated encryption operation. */
+psa_status_t psa_aead_encrypt_setup( psa_aead_operation_t *operation,
+                                     mbedtls_svc_key_id_t key,
+                                     psa_algorithm_t alg )
+{
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+    psa_status_t unlock_status = PSA_ERROR_CORRUPTION_DETECTED;
+    psa_key_slot_t *slot;
+
+    if( !PSA_ALG_IS_AEAD( alg ) || PSA_ALG_IS_WILDCARD( alg ) )
+        return( PSA_ERROR_NOT_SUPPORTED );
+
+    status = psa_get_and_lock_key_slot_with_policy(
+                 key, &slot, PSA_KEY_USAGE_DECRYPT, alg );
+
+    if( status != PSA_SUCCESS )
+    {
+        return( status );
+    }
+
+    psa_key_attributes_t attributes = {
+      .core = slot->attr
+    };
+
+    status = psa_driver_wrapper_aead_encrypt_setup( operation,
+                                                    &attributes, slot->key.data,
+                                                    slot->key.bytes, alg );
+
+
+    unlock_status = psa_unlock_key_slot( slot );
+
+    if( unlock_status != PSA_SUCCESS )
+    {
+        return( unlock_status );
+    }
+
+    return( status );
+}
+
+/* Set the key for a multipart authenticated decryption operation. */
+psa_status_t psa_aead_decrypt_setup( psa_aead_operation_t *operation,
+                                     mbedtls_svc_key_id_t key,
+                                     psa_algorithm_t alg )
+{
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+    psa_status_t unlock_status = PSA_ERROR_CORRUPTION_DETECTED;
+    psa_key_slot_t *slot;
+
+    if( !PSA_ALG_IS_AEAD( alg ) || PSA_ALG_IS_WILDCARD( alg ) )
+        return( PSA_ERROR_NOT_SUPPORTED );
+
+    status = psa_get_and_lock_key_slot_with_policy(
+                 key, &slot, PSA_KEY_USAGE_DECRYPT, alg );
+
+    if( status != PSA_SUCCESS )
+    {
+        return( status );
+    }
+
+    psa_key_attributes_t attributes = {
+      .core = slot->attr
+    };
+
+    status = psa_driver_wrapper_aead_decrypt_setup( operation,
+                                                    &attributes, slot->key.data,
+                                                    slot->key.bytes, alg );
+
+
+    unlock_status = psa_unlock_key_slot( slot );
+
+    if( unlock_status != PSA_SUCCESS )
+    {
+        return( unlock_status );
+    }
+
+    return( status );
+}
+
+/* Generate a random nonce / IV for multipart AEAD operation */
+psa_status_t psa_aead_generate_nonce( psa_aead_operation_t *operation,
+                                      uint8_t *nonce,
+                                      size_t nonce_size,
+                                      size_t *nonce_length )
+{
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+    size_t required_nonce_size = nonce_size;
+
+     *nonce_length = 0;
+
+    if( !operation->key_set || operation->nonce_set ||
+        operation->ad_started || operation->body_started )
+    {
+        return( PSA_ERROR_BAD_STATE );
+    }
+
+    required_nonce_size = PSA_AEAD_NONCE_LENGTH(operation->key_type, operation->alg);
+
+    if( nonce_size == 0 || nonce_size < required_nonce_size )
+    {
+        return( PSA_ERROR_BUFFER_TOO_SMALL );
+    }
+
+    status = psa_generate_random( nonce, required_nonce_size );
+
+    if( status != PSA_SUCCESS )
+    {
+        return status;
+    }
+
+    status = psa_driver_wrapper_aead_set_nonce( operation, nonce, required_nonce_size );
+
+    if( status == PSA_SUCCESS )
+    {
+        *nonce_length = required_nonce_size;
+    }
+
+    return status;
+}
+
+/* Set the nonce for a multipart authenticated encryption or decryption
+   operation.*/
+psa_status_t psa_aead_set_nonce( psa_aead_operation_t *operation,
+                                 const uint8_t *nonce,
+                                 size_t nonce_length )
+{
+    if( !operation->key_set || operation->nonce_set ||
+        operation->ad_started || operation->body_started )
+    {
+        return( PSA_ERROR_BAD_STATE );
+    }
+
+    return( psa_driver_wrapper_aead_set_nonce( operation, nonce, nonce_length ) );
+}
+
+/* Declare the lengths of the message and additional data for multipart AEAD. */
+psa_status_t psa_aead_set_lengths( psa_aead_operation_t *operation,
+                                   size_t ad_length,
+                                   size_t plaintext_length )
+{
+    if( !operation->key_set || operation->lengths_set )
+    {
+        return( PSA_ERROR_BAD_STATE );
+    }
+
+    return( psa_driver_wrapper_aead_set_lengths( operation, ad_length, plaintext_length ) );
+}
+ /* Pass additional data to an active multipart AEAD operation. */
+psa_status_t psa_aead_update_ad( psa_aead_operation_t *operation,
+                                 const uint8_t *input,
+                                 size_t input_length )
+{
+    if( !operation->nonce_set || !operation->key_set )
+    {
+        return( PSA_ERROR_BAD_STATE );
+    }
+
+    return( psa_driver_wrapper_aead_update_ad( operation, input, input_length ) );
+}
+
+/* Encrypt or decrypt a message fragment in an active multipart AEAD
+   operation.*/
+psa_status_t psa_aead_update( psa_aead_operation_t *operation,
+                              const uint8_t *input,
+                              size_t input_length,
+                              uint8_t *output,
+                              size_t output_size,
+                              size_t *output_length )
+{
+
+    *output_length = 0;
+
+    if( !operation->nonce_set || !operation->key_set || !operation->ad_started )
+    {
+        return( PSA_ERROR_BAD_STATE );
+    }
+
+    return( psa_driver_wrapper_aead_update( operation, input, input_length, output, output_size,
+                                            output_length ) );
+}
+
+/* Finish encrypting a message in a multipart AEAD operation. */
+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 )
+{
+    *ciphertext_length = 0;
+    *tag_length = 0;
+
+    if( !operation->key_set || !operation->nonce_set ||
+        !operation->ad_started || !operation->body_started )
+    {
+        return( PSA_ERROR_BAD_STATE );
+    }
+
+    return( psa_driver_wrapper_aead_finish( operation, ciphertext, ciphertext_size,
+                                            ciphertext_length, tag, tag_size, tag_length ) );
+}
+
+/* Finish authenticating and decrypting a message in a multipart AEAD
+   operation.*/
+psa_status_t psa_aead_verify( psa_aead_operation_t *operation,
+                              uint8_t *plaintext,
+                              size_t plaintext_size,
+                              size_t *plaintext_length,
+                              const uint8_t *tag,
+                              size_t tag_length )
+{
+    *plaintext_length = 0;
+
+    if( !operation->key_set || !operation->nonce_set ||
+        !operation->ad_started || !operation->body_started )
+    {
+        return( PSA_ERROR_BAD_STATE );
+    }
+
+    return( psa_driver_wrapper_aead_verify( operation, plaintext, plaintext_size, plaintext_length,
+                                             tag, tag_length ) );
+}
+
+/* Abort an AEAD operation. */
+psa_status_t psa_aead_abort(psa_aead_operation_t *operation)
+{
+    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
+
+    if( operation->id == 0 )
+    {
+        /* The object has (apparently) been initialized but it is not (yet)
+         * in use. It's ok to call abort on such an object, and there's
+         * nothing to do. */
+        return( PSA_SUCCESS );
+    }
+
+    status = psa_driver_wrapper_aead_abort( operation );
+
+    operation->id = 0;
+    operation->key_set = 0;
+    operation->nonce_set = 0;
+    operation->lengths_set = 0;
+    operation->is_encrypt = 0;
+    operation->ad_started = 0;
+    operation->body_started = 0;
+
+    return( status );
+}
+
 /****************************************************************/
 /* Generators */
 /****************************************************************/
diff --git a/library/psa_crypto_aead.c b/library/psa_crypto_aead.c
index 47b0e7b..f8cceae 100644
--- a/library/psa_crypto_aead.c
+++ b/library/psa_crypto_aead.c
@@ -394,45 +394,6 @@
     return ( status );
 }
 
-/* Generate a random nonce / IV for multipart AEAD operation */
-psa_status_t mbedtls_psa_aead_generate_nonce( psa_aead_operation_t *operation,
-                                              uint8_t *nonce,
-                                              size_t nonce_size,
-                                              size_t *nonce_length )
-{
-    psa_status_t status;
-    size_t required_nonce_size = nonce_size;
-
-    if( !operation->key_set || operation->nonce_set ||
-        operation->ad_started || operation->body_started )
-    {
-        return( PSA_ERROR_BAD_STATE );
-    }
-
-    required_nonce_size = PSA_AEAD_NONCE_LENGTH(operation->key_type, operation->alg);
-
-    if( nonce_size == 0 || nonce_size < required_nonce_size )
-    {
-        return( PSA_ERROR_BUFFER_TOO_SMALL );
-    }
-
-    status = psa_generate_random( nonce, required_nonce_size );
-
-    if( status != PSA_SUCCESS )
-    {
-        return status;
-    }
-
-    status = mbedtls_psa_aead_set_nonce( operation, nonce, required_nonce_size );
-
-    if( status == PSA_SUCCESS )
-    {
-        *nonce_length = required_nonce_size;
-    }
-
-    return status;
-}
-
 /* Set a nonce for the multipart AEAD operation*/
 psa_status_t mbedtls_psa_aead_set_nonce( psa_aead_operation_t *operation,
                                          const uint8_t *nonce,
@@ -440,19 +401,6 @@
 {
     psa_status_t status;
 
-    if( !operation->key_set || operation->nonce_set ||
-        operation->ad_started || operation->body_started )
-    {
-        return( PSA_ERROR_BAD_STATE );
-    }
-
-    /* Restricting to a nominal safe length for nonces even though some
-       algorithms can handle longer nonces, but not without collision */
-    if( nonce_length > PSA_AEAD_NONCE_MAX_SIZE )
-    {
-        return( PSA_ERROR_INVALID_ARGUMENT );
-    }
-
     #if defined(MBEDTLS_PSA_BUILTIN_ALG_GCM)
     if( operation->alg == PSA_ALG_GCM )
     {
@@ -514,11 +462,6 @@
                                            size_t plaintext_length )
 {
 
-    if( !operation->key_set || operation->lengths_set )
-    {
-        return( PSA_ERROR_BAD_STATE );
-    }
-
 #if defined(MBEDTLS_PSA_BUILTIN_ALG_GCM)
     if( operation->alg == PSA_ALG_GCM )
     {
@@ -570,11 +513,6 @@
 {
     psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
 
-    if( !operation->nonce_set || !operation->key_set )
-    {
-        return( PSA_ERROR_BAD_STATE );
-    }
-
     if( operation->lengths_set )
     {
         if ( operation->ad_remaining < input_length )
@@ -675,11 +613,6 @@
     size_t update_output_size;
     psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
 
-    if( !operation->nonce_set || !operation->key_set || !operation->ad_started )
-    {
-        return( PSA_ERROR_BAD_STATE );
-    }
-
     update_output_size = PSA_AEAD_UPDATE_OUTPUT_SIZE(operation->key_type,
                                                      operation->alg, input_length);
 
@@ -791,12 +724,6 @@
                                                     size_t *finish_output_size,
                                                     size_t *output_tag_length )
 {
-    if( !operation->key_set || !operation->nonce_set
-        || !operation->ad_started || !operation->body_started)
-    {
-        return( PSA_ERROR_BAD_STATE );
-    }
-
     if( operation->lengths_set )
     {
         if( operation->ad_remaining != 0 || operation->body_remaining != 0 )
diff --git a/library/psa_crypto_aead.h b/library/psa_crypto_aead.h
index d7aac24..a9d2687 100644
--- a/library/psa_crypto_aead.h
+++ b/library/psa_crypto_aead.h
@@ -167,9 +167,9 @@
  *    the inputs to the subsequent calls to mbedtls_psa_aead_update_ad() and
  *    mbedtls_psa_aead_update(). See the documentation of mbedtls_psa_aead_set_lengths()
  *    for details.
- * -# Call either mbedtls_psa_aead_generate_nonce() or mbedtls_psa_aead_set_nonce() to
- *    generate or set the nonce. You should use
- *    mbedtls_psa_aead_generate_nonce() unless the protocol you are implementing
+ * -# Call either psa_aead_generate_nonce() or
+ *    mbedtls_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 mbedtls_psa_aead_update_ad() zero, one or more times, passing a fragment
  *    of the non-encrypted additional authenticated data each time.
@@ -297,52 +297,6 @@
                                             const uint8_t *key_buffer, size_t key_buffer_size,
                                             psa_algorithm_t alg);
 
-/** Generate a random nonce for an authenticated encryption operation.
- *
- * \note The signature of this function is that of a PSA driver
- *       aead_generate_nonce entry point. This function behaves as an
- *       aead_generate_nonce entry point as defined in the PSA driver interface
- *       specification for transparent drivers.
- *
- * 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 mbedtls_psa_aead_encrypt_setup() before
- * calling this function.
- *
- * If this function returns an error status, the operation enters an error
- * state and must be aborted by calling mbedtls_psa_aead_abort().
- *
- * \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 (it must be an active aead encrypt
- *         operation, with no nonce 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_CORRUPTION_DETECTED
- * \retval #PSA_ERROR_STORAGE_FAILURE
- * \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 mbedtls_psa_aead_generate_nonce(psa_aead_operation_t *operation,
-                                            uint8_t *nonce,
-                                            size_t nonce_size,
-                                            size_t *nonce_length);
-
 /** Set the nonce for an authenticated encryption or decryption operation.
  *
  * \note The signature of this function is that of a PSA driver
@@ -402,7 +356,7 @@
  * then the implementation must enforce the lengths.
  *
  * You may call this function before or after setting the nonce with
- * mbedtls_psa_aead_set_nonce() or mbedtls_psa_aead_generate_nonce().
+ * mbedtls_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
@@ -454,7 +408,7 @@
  *
  * Before calling this function, you must:
  * 1. Call either mbedtls_psa_aead_encrypt_setup() or mbedtls_psa_aead_decrypt_setup().
- * 2. Set the nonce with mbedtls_psa_aead_generate_nonce() or
+ * 2. Set the nonce with psa_aead_generate_nonce() or
  *    mbedtls_psa_aead_set_nonce().
  *
  * If this function returns an error status, the operation enters an error
@@ -509,8 +463,9 @@
  * 1. Call either mbedtls_psa_aead_encrypt_setup() or mbedtls_psa_aead_decrypt_setup().
  *    The choice of setup function determines whether this function
  *    encrypts or decrypts its input.
- * 2. Set the nonce with mbedtls_psa_aead_generate_nonce() or mbedtls_psa_aead_set_nonce().
- * 3. Call mbedtls_psa_aead_update_ad() to pass all the additional data.
+ * 2. Set the nonce with psa_aead_generate_nonce() or
+ * mbedtls_psa_aead_set_nonce(). 3. Call mbedtls_psa_aead_update_ad() to pass
+ * all the additional data.
  *
  * If this function returns an error status, the operation enters an error
  * state and must be aborted by calling mbedtls_psa_aead_abort().
diff --git a/library/psa_crypto_driver_wrappers.c b/library/psa_crypto_driver_wrappers.c
index 59a00a6..5e09fd2 100644
--- a/library/psa_crypto_driver_wrappers.c
+++ b/library/psa_crypto_driver_wrappers.c
@@ -1394,39 +1394,6 @@
     }
 }
 
-psa_status_t psa_driver_wrapper_aead_generate_nonce(
-   psa_aead_operation_t *operation,
-   uint8_t *nonce,
-   size_t nonce_size,
-   size_t *nonce_length )
-{
-    switch( operation->id )
-    {
-#if defined(MBEDTLS_PSA_BUILTIN_CIPHER)
-        case PSA_CRYPTO_MBED_TLS_DRIVER_ID:
-            return( mbedtls_psa_aead_generate_nonce( operation, nonce, nonce_size,
-                                                     nonce_length ) );
-#endif /* MBEDTLS_PSA_BUILTIN_CIPHER */
-
-#if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
-#if defined(PSA_CRYPTO_DRIVER_TEST)
-        case PSA_CRYPTO_TRANSPARENT_TEST_DRIVER_ID:
-//          return( test_transparent_aead_generate_nonce(
-//                      operation, nonce, nonce_size, nonce_length ) );
-
-        /* Add cases for opaque driver here */
-
-#endif /* PSA_CRYPTO_DRIVER_TEST */
-#endif /* PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT */
-    }
-
-    (void)nonce;
-    (void)nonce_size;
-    (void)nonce_length;
-
-    return( PSA_ERROR_INVALID_ARGUMENT );
-}
-
 psa_status_t psa_driver_wrapper_aead_set_nonce(
    psa_aead_operation_t *operation,
    const uint8_t *nonce,
diff --git a/library/psa_crypto_driver_wrappers.h b/library/psa_crypto_driver_wrappers.h
index bdb2eba..05adb53 100644
--- a/library/psa_crypto_driver_wrappers.h
+++ b/library/psa_crypto_driver_wrappers.h
@@ -195,12 +195,6 @@
     const uint8_t *key_buffer, size_t key_buffer_size,
     psa_algorithm_t alg );
 
-psa_status_t psa_driver_wrapper_aead_generate_nonce(
-    psa_aead_operation_t *operation,
-    uint8_t *nonce,
-    size_t nonce_size,
-    size_t *nonce_length );
-
 psa_status_t psa_driver_wrapper_aead_set_nonce(
     psa_aead_operation_t *operation,
     const uint8_t *nonce,