Split multipart AEAD contexts into two parts

Split to data required for internal implementation and data required for
driver implementation with data left over for the PSA layer.

Signed-off-by: Paul Elliott <paul.elliott@arm.com>
diff --git a/library/psa_crypto_driver_wrappers.c b/library/psa_crypto_driver_wrappers.c
index 7faedb3..7590800 100644
--- a/library/psa_crypto_driver_wrappers.c
+++ b/library/psa_crypto_driver_wrappers.c
@@ -1310,10 +1310,9 @@
 
 #if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
 #if defined(PSA_CRYPTO_DRIVER_TEST)
-            status = PSA_ERROR_NOT_SUPPORTED;
             status = mbedtls_test_transparent_aead_encrypt_setup(
-                        operation, attributes,
-                        key_buffer, key_buffer_size,
+                        &operation->ctx.transparent_test_driver_ctx,
+                        attributes, key_buffer, key_buffer_size,
                         alg );
             /* Declared with fallback == true */
             operation->id = PSA_CRYPTO_TRANSPARENT_TEST_DRIVER_ID;
@@ -1325,7 +1324,7 @@
 
             /* Fell through, meaning no accelerator supports this operation */
             status = mbedtls_psa_aead_encrypt_setup(
-                        operation, attributes,
+                        &operation->ctx.mbedtls_ctx, attributes,
                         key_buffer, key_buffer_size,
                         alg );
 
@@ -1360,9 +1359,9 @@
 
 #if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
 #if defined(PSA_CRYPTO_DRIVER_TEST)
-            status = PSA_ERROR_NOT_SUPPORTED;
             status = mbedtls_test_transparent_aead_decrypt_setup(
-                        operation, attributes,
+                        &operation->ctx.transparent_test_driver_ctx,
+                         attributes,
                         key_buffer, key_buffer_size,
                         alg );
             /* Declared with fallback == true */
@@ -1375,7 +1374,8 @@
 
             /* Fell through, meaning no accelerator supports this operation */
             status = mbedtls_psa_aead_decrypt_setup(
-                        operation, attributes,
+                        &operation->ctx.mbedtls_ctx,
+                         attributes,
                         key_buffer, key_buffer_size,
                         alg );
 
@@ -1401,16 +1401,18 @@
     {
 #if defined(MBEDTLS_PSA_BUILTIN_CIPHER)
         case PSA_CRYPTO_MBED_TLS_DRIVER_ID:
-            return( mbedtls_psa_aead_set_nonce( operation, nonce,
+            return( mbedtls_psa_aead_set_nonce( &operation->ctx.mbedtls_ctx,
+                                                nonce,
                                                 nonce_length ) );
 
-#endif /* MBEDTLS_PSA_BUILTIN_CIPHER */
+#endif /* MBEDTLS_PSA_BUILTIN_AEAD */
 
 #if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
 #if defined(PSA_CRYPTO_DRIVER_TEST)
         case PSA_CRYPTO_TRANSPARENT_TEST_DRIVER_ID:
             return( mbedtls_test_transparent_aead_set_nonce(
-                        operation, nonce, nonce_length ) );
+                         &operation->ctx.transparent_test_driver_ctx,
+                         nonce, nonce_length ) );
 
         /* Add cases for opaque driver here */
 
@@ -1431,18 +1433,20 @@
 {
     switch( operation->id )
     {
-#if defined(MBEDTLS_PSA_BUILTIN_CIPHER)
+#if defined(MBEDTLS_PSA_BUILTIN_AEAD)
         case PSA_CRYPTO_MBED_TLS_DRIVER_ID:
-            return( mbedtls_psa_aead_set_lengths( operation, ad_length,
+            return( mbedtls_psa_aead_set_lengths( &operation->ctx.mbedtls_ctx,
+                                                  ad_length,
                                                   plaintext_length ) );
 
-#endif /* MBEDTLS_PSA_BUILTIN_CIPHER */
+#endif /* MBEDTLS_PSA_BUILTIN_AEAD */
 
 #if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
 #if defined(PSA_CRYPTO_DRIVER_TEST)
         case PSA_CRYPTO_TRANSPARENT_TEST_DRIVER_ID:
             return( mbedtls_test_transparent_aead_set_lengths(
-                        operation, ad_length, plaintext_length ) );
+                        &operation->ctx.transparent_test_driver_ctx,
+                        ad_length, plaintext_length ) );
 
         /* Add cases for opaque driver here */
 
@@ -1463,18 +1467,20 @@
 {
     switch( operation->id )
     {
-#if defined(MBEDTLS_PSA_BUILTIN_CIPHER)
+#if defined(MBEDTLS_PSA_BUILTIN_AEAD)
         case PSA_CRYPTO_MBED_TLS_DRIVER_ID:
-            return( mbedtls_psa_aead_update_ad( operation, input,
+            return( mbedtls_psa_aead_update_ad( &operation->ctx.mbedtls_ctx,
+                                                input,
                                                 input_length ) );
 
-#endif /* MBEDTLS_PSA_BUILTIN_CIPHER */
+#endif /* MBEDTLS_PSA_BUILTIN_AEAD */
 
 #if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
 #if defined(PSA_CRYPTO_DRIVER_TEST)
         case PSA_CRYPTO_TRANSPARENT_TEST_DRIVER_ID:
             return( mbedtls_test_transparent_aead_update_ad(
-                        operation, input, input_length ) );
+                        &operation->ctx.transparent_test_driver_ctx,
+                        input, input_length ) );
 
         /* Add cases for opaque driver here */
 
@@ -1498,19 +1504,21 @@
 {
     switch( operation->id )
     {
-#if defined(MBEDTLS_PSA_BUILTIN_CIPHER)
+#if defined(MBEDTLS_PSA_BUILTIN_AEAD)
         case PSA_CRYPTO_MBED_TLS_DRIVER_ID:
-            return( mbedtls_psa_aead_update( operation, input, input_length,
+            return( mbedtls_psa_aead_update( &operation->ctx.mbedtls_ctx,
+                                             input, input_length,
                                              output, output_size,
                                              output_length ) );
 
-#endif /* MBEDTLS_PSA_BUILTIN_CIPHER */
+#endif /* MBEDTLS_PSA_BUILTIN_AEAD */
 
 #if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
 #if defined(PSA_CRYPTO_DRIVER_TEST)
         case PSA_CRYPTO_TRANSPARENT_TEST_DRIVER_ID:
             return( mbedtls_test_transparent_aead_update(
-                        operation, input, input_length, output, output_size,
+                        &operation->ctx.transparent_test_driver_ctx,
+                        input, input_length, output, output_size,
                         output_length ) );
 
         /* Add cases for opaque driver here */
@@ -1539,20 +1547,22 @@
 {
     switch( operation->id )
     {
-#if defined(MBEDTLS_PSA_BUILTIN_CIPHER)
+#if defined(MBEDTLS_PSA_BUILTIN_AEAD)
         case PSA_CRYPTO_MBED_TLS_DRIVER_ID:
-            return( mbedtls_psa_aead_finish( operation, ciphertext,
+            return( mbedtls_psa_aead_finish( &operation->ctx.mbedtls_ctx,
+                                             ciphertext,
                                              ciphertext_size,
                                              ciphertext_length, tag,
                                              tag_size, tag_length ) );
 
-#endif /* MBEDTLS_PSA_BUILTIN_CIPHER */
+#endif /* MBEDTLS_PSA_BUILTIN_AEAD */
 
 #if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
 #if defined(PSA_CRYPTO_DRIVER_TEST)
         case PSA_CRYPTO_TRANSPARENT_TEST_DRIVER_ID:
             return( mbedtls_test_transparent_aead_finish(
-                        operation, ciphertext, ciphertext_size,
+                        &operation->ctx.transparent_test_driver_ctx,
+                        ciphertext, ciphertext_size,
                         ciphertext_length, tag, tag_size, tag_length ) );
 
         /* Add cases for opaque driver here */
@@ -1581,19 +1591,22 @@
 {
     switch( operation->id )
     {
-#if defined(MBEDTLS_PSA_BUILTIN_CIPHER)
+#if defined(MBEDTLS_PSA_BUILTIN_AEAD)
         case PSA_CRYPTO_MBED_TLS_DRIVER_ID:
-            return( mbedtls_psa_aead_verify( operation, plaintext,
-                                             plaintext_size, plaintext_length,
+            return( mbedtls_psa_aead_verify( &operation->ctx.mbedtls_ctx,
+                                             plaintext,
+                                             plaintext_size,
+                                             plaintext_length,
                                              tag, tag_length ) );
 
-#endif /* MBEDTLS_PSA_BUILTIN_CIPHER */
+#endif /* MBEDTLS_PSA_BUILTIN_AEAD */
 
 #if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
 #if defined(PSA_CRYPTO_DRIVER_TEST)
         case PSA_CRYPTO_TRANSPARENT_TEST_DRIVER_ID:
             return( mbedtls_test_transparent_aead_verify(
-                        operation, plaintext, plaintext_size,
+                        &operation->ctx.transparent_test_driver_ctx,
+                        plaintext, plaintext_size,
                         plaintext_length, tag, tag_length ) );
 
         /* Add cases for opaque driver here */
@@ -1616,16 +1629,17 @@
 {
     switch( operation->id )
     {
-#if defined(MBEDTLS_PSA_BUILTIN_CIPHER)
+#if defined(MBEDTLS_PSA_BUILTIN_AEAD)
         case PSA_CRYPTO_MBED_TLS_DRIVER_ID:
-            return( mbedtls_psa_aead_abort( operation ) );
+            return( mbedtls_psa_aead_abort( &operation->ctx.mbedtls_ctx ) );
 
-#endif /* MBEDTLS_PSA_BUILTIN_CIPHER */
+#endif /* MBEDTLS_PSA_BUILTIN_AEAD */
 
 #if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
 #if defined(PSA_CRYPTO_DRIVER_TEST)
         case PSA_CRYPTO_TRANSPARENT_TEST_DRIVER_ID:
-            return( mbedtls_test_transparent_aead_abort( operation ) );
+            return( mbedtls_test_transparent_aead_abort(
+               &operation->ctx.transparent_test_driver_ctx ) );
 
         /* Add cases for opaque driver here */