diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 82af920..4486242 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -1224,6 +1224,9 @@
     mbedtls_cipher_mode_t mode;
     mbedtls_cipher_id_t cipher_id_tmp;
 
+    if( PSA_ALG_IS_AEAD( alg ) )
+        alg &= ~PSA_ALG_AEAD_TAG_LENGTH_MASK;
+
     if( PSA_ALG_IS_CIPHER( alg ) || PSA_ALG_IS_AEAD( alg ) )
     {
         switch( alg )
@@ -1246,10 +1249,10 @@
             case PSA_ALG_CBC_PKCS7:
                 mode = MBEDTLS_MODE_CBC;
                 break;
-            case PSA_ALG_CCM:
+            case PSA_ALG_CCM & ~PSA_ALG_AEAD_TAG_LENGTH_MASK:
                 mode = MBEDTLS_MODE_CCM;
                 break;
-            case PSA_ALG_GCM:
+            case PSA_ALG_GCM & ~PSA_ALG_AEAD_TAG_LENGTH_MASK:
                 mode = MBEDTLS_MODE_GCM;
                 break;
             default:
@@ -2834,6 +2837,8 @@
         mbedtls_gcm_context gcm;
 #endif /* MBEDTLS_GCM_C */
     } ctx;
+    psa_algorithm_t core_alg;
+    uint8_t full_tag_length;
     uint8_t tag_length;
 } aead_operation_t;
 
@@ -2876,11 +2881,12 @@
     if( operation->cipher_info == NULL )
         return( PSA_ERROR_NOT_SUPPORTED );
 
-    switch( alg )
+    switch( PSA_ALG_AEAD_WITH_TAG_LENGTH( alg, 0 ) )
     {
 #if defined(MBEDTLS_CCM_C)
-        case PSA_ALG_CCM:
-            operation->tag_length = 16;
+        case PSA_ALG_AEAD_WITH_TAG_LENGTH( PSA_ALG_CCM, 0 ):
+            operation->core_alg = PSA_ALG_CCM;
+            operation->full_tag_length = 16;
             if( PSA_BLOCK_CIPHER_BLOCK_SIZE( operation->slot->type ) != 16 )
                 return( PSA_ERROR_INVALID_ARGUMENT );
             mbedtls_ccm_init( &operation->ctx.ccm );
@@ -2894,8 +2900,9 @@
 #endif /* MBEDTLS_CCM_C */
 
 #if defined(MBEDTLS_GCM_C)
-        case PSA_ALG_GCM:
-            operation->tag_length = 16;
+        case PSA_ALG_AEAD_WITH_TAG_LENGTH( PSA_ALG_GCM, 0 ):
+            operation->core_alg = PSA_ALG_GCM;
+            operation->full_tag_length = 16;
             if( PSA_BLOCK_CIPHER_BLOCK_SIZE( operation->slot->type ) != 16 )
                 return( PSA_ERROR_INVALID_ARGUMENT );
             mbedtls_gcm_init( &operation->ctx.gcm );
@@ -2910,6 +2917,16 @@
             return( PSA_ERROR_NOT_SUPPORTED );
     }
 
+    if( PSA_AEAD_TAG_LENGTH( alg ) > operation->full_tag_length )
+    {
+        status = PSA_ERROR_INVALID_ARGUMENT;
+        goto cleanup;
+    }
+    operation->tag_length = PSA_AEAD_TAG_LENGTH( alg );
+    /* CCM allows the following tag lengths: 4, 6, 8, 10, 12, 14, 16.
+     * GCM allows the following tag lengths: 4, 8, 12, 13, 14, 15, 16.
+     * In both cases, mbedtls_xxx will validate the tag length below. */
+
     return( PSA_SUCCESS );
 
 cleanup:
@@ -2948,7 +2965,7 @@
     }
     tag = ciphertext + plaintext_length;
 
-    if( alg == PSA_ALG_GCM )
+    if( operation.core_alg == PSA_ALG_GCM )
     {
         status = mbedtls_to_psa_error(
             mbedtls_gcm_crypt_and_tag( &operation.ctx.gcm,
@@ -2959,7 +2976,7 @@
                                        plaintext, ciphertext,
                                        operation.tag_length, tag ) );
     }
-    else if( alg == PSA_ALG_CCM )
+    else if( operation.core_alg == PSA_ALG_CCM )
     {
         status = mbedtls_to_psa_error(
             mbedtls_ccm_encrypt_and_tag( &operation.ctx.ccm,
@@ -3028,7 +3045,7 @@
     if( status != PSA_SUCCESS )
         return( status );
 
-    if( alg == PSA_ALG_GCM )
+    if( operation.core_alg == PSA_ALG_GCM )
     {
         status = psa_aead_unpadded_locate_tag( operation.tag_length,
                                                ciphertext, ciphertext_length,
@@ -3045,7 +3062,7 @@
                                       tag, operation.tag_length,
                                       ciphertext, plaintext ) );
     }
-    else if( alg == PSA_ALG_CCM )
+    else if( operation.core_alg == PSA_ALG_CCM )
     {
         status = psa_aead_unpadded_locate_tag( operation.tag_length,
                                                ciphertext, ciphertext_length,