Detect invalid tag lengths in psa_aead_setup

Read tag lengths from the driver and validate against preset values.
Signed-off-by: Andrzej Kurek <andrzej.kurek@arm.com>
diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index cbac109..0781a72 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -4409,6 +4409,7 @@
                             data_t *input_data,
                             data_t *tag,
                             int tag_usage_arg,
+                            int expected_setup_status_arg,
                             int expected_status_arg )
 {
     mbedtls_svc_key_id_t key = MBEDTLS_SVC_KEY_ID_INIT;
@@ -4418,6 +4419,7 @@
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
     psa_status_t status = PSA_ERROR_GENERIC_ERROR;
     psa_status_t expected_status = expected_status_arg;
+    psa_status_t expected_setup_status = expected_setup_status_arg;
     unsigned char *plaintext = NULL;
     unsigned char *finish_plaintext = NULL;
     size_t plaintext_size = 0;
@@ -4457,6 +4459,10 @@
         MBEDTLS_TEST_PSA_SKIP_IF_ALT_AES_192( key_type, key_data->len * 8 );
         MBEDTLS_TEST_PSA_SKIP_IF_ALT_GCM_NOT_12BYTES_NONCE( alg, nonce->len );
     }
+    TEST_EQUAL( status, expected_setup_status );
+
+    if( status != PSA_SUCCESS )
+        goto exit;
 
     PSA_ASSERT( status );
 
@@ -4464,17 +4470,7 @@
 
     status = psa_aead_set_lengths( &operation, additional_data->len,
                                       input_data->len );
-
-    if( status != PSA_SUCCESS )
-    {
-        /* Invalid tag lengths are detected in CCM at this point, as they
-         * would be written into the first block. They should really be
-         * detected in psa_aead_encrypt/decrypt_setup, and will be fixed
-         * to do so in the future, until that point, this is a
-         * workaround.*/
-        TEST_EQUAL( status, expected_status );
-        goto exit;
-    }
+    PSA_ASSERT( status );
 
     PSA_ASSERT( psa_aead_update_ad( &operation, additional_data->x,
                                     additional_data->len ) );