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/library/psa_crypto.c b/library/psa_crypto.c
index e021f8a..642fc13 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -3829,6 +3829,47 @@
     return( status );
 }
 
+static psa_status_t psa_validate_tag_length( psa_aead_operation_t *operation,
+                                             psa_algorithm_t alg ) {
+    uint8_t tag_len = 0;
+    if( psa_driver_get_tag_len( operation, &tag_len ) != PSA_SUCCESS )
+    {
+        return( PSA_ERROR_INVALID_ARGUMENT );
+    }
+
+    switch( PSA_ALG_AEAD_WITH_SHORTENED_TAG( alg, 0 ) )
+    {
+#if defined(MBEDTLS_PSA_BUILTIN_ALG_CCM)
+        case PSA_ALG_AEAD_WITH_SHORTENED_TAG( PSA_ALG_CCM, 0 ):
+            /* CCM allows the following tag lengths: 4, 6, 8, 10, 12, 14, 16.*/
+            if( tag_len < 4 || tag_len > 16 || tag_len % 2 )
+                return( PSA_ERROR_INVALID_ARGUMENT );
+            break;
+#endif /* MBEDTLS_PSA_BUILTIN_ALG_CCM */
+
+#if defined(MBEDTLS_PSA_BUILTIN_ALG_GCM)
+        case PSA_ALG_AEAD_WITH_SHORTENED_TAG( PSA_ALG_GCM, 0 ):
+            /* GCM allows the following tag lengths: 4, 8, 12, 13, 14, 15, 16. */
+            if( tag_len != 4 && tag_len != 8 && ( tag_len < 12 || tag_len > 16 ) )
+                return( PSA_ERROR_INVALID_ARGUMENT );
+            break;
+#endif /* MBEDTLS_PSA_BUILTIN_ALG_GCM */
+
+#if defined(MBEDTLS_PSA_BUILTIN_ALG_CHACHA20_POLY1305)
+        case PSA_ALG_AEAD_WITH_SHORTENED_TAG( PSA_ALG_CHACHA20_POLY1305, 0 ):
+            /* We only support the default tag length. */
+            if( tag_len != 16 )
+                return( PSA_ERROR_INVALID_ARGUMENT );
+            break;
+#endif /* MBEDTLS_PSA_BUILTIN_ALG_CHACHA20_POLY1305 */
+
+        default:
+            (void) tag_len;
+            return( PSA_ERROR_NOT_SUPPORTED );
+    }
+    return( PSA_SUCCESS );
+}
+
 /* Set the key for a multipart authenticated operation. */
 static psa_status_t psa_aead_setup( psa_aead_operation_t *operation,
                                     int is_encrypt,
@@ -3886,6 +3927,9 @@
     if( status != PSA_SUCCESS )
         goto exit;
 
+    if( ( status = psa_validate_tag_length( operation, alg ) ) != PSA_SUCCESS )
+        goto exit;
+
     operation->key_type = psa_get_key_type( &attributes );
 
 exit: