diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index c0ef1c54f..2650fffe6 100755
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -1503,8 +1503,9 @@ psa_status_t psa_aead_encrypt( psa_key_slot_t key,
     if( cipher_info == NULL )
         return( PSA_ERROR_NOT_SUPPORTED );
 
-    if ( key_type != PSA_KEY_TYPE_RAW_DATA)
-        return( PSA_ERROR_BAD_STATE );
+    if ( ( key_type & PSA_KEY_TYPE_CATEGORY_MASK ) == PSA_KEY_TYPE_CATEGORY_CIPHER
+           && PSA_BLOCK_CIPHER_BLOCK_SIZE( key_type ) == 16 )
+        return( PSA_ERROR_INVALID_ARGUMENT );
 
     operation->block_size = cipher_info->block_size;
 
@@ -1586,7 +1587,7 @@ psa_status_t psa_aead_decrypt( psa_key_slot_t key,
     const mbedtls_cipher_info_t *cipher_info = NULL;
     unsigned char tag[16];
 
-    if( plaintext_size < ( ciphertext_length + sizeof( tag ) ) )
+    if( plaintext_size < ciphertext_length )
             return( PSA_ERROR_INVALID_ARGUMENT );
 
     status = psa_get_key_information( key, &key_type, &key_bits );