Return the same error in multipart and single shot AEAD

psa_aead_encrypt_setup() and psa_aead_decrypt_setup() were returning
PSA_ERROR_INVALID_ARGUMENT, while the same failed checks were producing
PSA_ERROR_NOT_SUPPORTED if they happened in psa_aead_encrypt() or
psa_aead_decrypt().

The PSA Crypto API 1.1 spec will specify PSA_ERROR_INVALID_ARGUMENT
in the case that the supplied algorithm is not an AEAD one.

Also move these shared checks to a helper function, to reduce code
duplication and ensure that the functions remain in sync.

Signed-off-by: Bence Szépkúti <bence.szepkuti@arm.com>
diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 829ed45..dbff133 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -3719,6 +3719,13 @@
     return( PSA_ERROR_INVALID_ARGUMENT );
 }
 
+static psa_status_t psa_aead_initial_checks( psa_algorithm_t alg ) {
+    if( !PSA_ALG_IS_AEAD( alg ) || PSA_ALG_IS_WILDCARD( alg ) )
+        return( PSA_ERROR_INVALID_ARGUMENT );
+
+    return( PSA_SUCCESS );
+}
+
 psa_status_t psa_aead_encrypt( mbedtls_svc_key_id_t key,
                                psa_algorithm_t alg,
                                const uint8_t *nonce,
@@ -3736,8 +3743,9 @@
 
     *ciphertext_length = 0;
 
-    if( !PSA_ALG_IS_AEAD( alg ) || PSA_ALG_IS_WILDCARD( alg ) )
-        return( PSA_ERROR_NOT_SUPPORTED );
+    status = psa_aead_initial_checks( alg );
+    if( status != PSA_SUCCESS )
+        return( status );
 
     status = psa_get_and_lock_key_slot_with_policy(
                  key, &slot, PSA_KEY_USAGE_ENCRYPT, alg );
@@ -3786,8 +3794,9 @@
 
     *plaintext_length = 0;
 
-    if( !PSA_ALG_IS_AEAD( alg ) || PSA_ALG_IS_WILDCARD( alg ) )
-        return( PSA_ERROR_NOT_SUPPORTED );
+    status = psa_aead_initial_checks( alg );
+    if( status != PSA_SUCCESS )
+        return( status );
 
     status = psa_get_and_lock_key_slot_with_policy(
                  key, &slot, PSA_KEY_USAGE_DECRYPT, alg );
@@ -3830,11 +3839,9 @@
     psa_key_slot_t *slot = NULL;
     psa_key_usage_t key_usage = 0;
 
-    if( !PSA_ALG_IS_AEAD( alg ) || PSA_ALG_IS_WILDCARD( alg ) )
-    {
-        status = PSA_ERROR_INVALID_ARGUMENT;
+    status = psa_aead_initial_checks( alg );
+    if( status != PSA_SUCCESS )
         goto exit;
-    }
 
     if( operation->id != 0 )
     {