Implement multipart AEAD PSA interface

Signed-off-by: Paul Elliott <paul.elliott@arm.com>
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 */
 /****************************************************************/