diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index a9026e4..a5027f3 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -3749,8 +3749,11 @@
     return( status );
 }
 
-static psa_status_t psa_aead_final_checks( psa_aead_operation_t *operation )
+static psa_status_t psa_aead_final_checks( psa_aead_operation_t *operation,
+                                           size_t output_size )
 {
+    size_t finish_output_size;
+
     if( operation->id == 0 || !operation->nonce_set )
         return( PSA_ERROR_BAD_STATE );
 
@@ -3758,6 +3761,13 @@
                                    operation->body_remaining != 0 ) )
         return( PSA_ERROR_INVALID_ARGUMENT );
 
+    finish_output_size = operation->is_encrypt ?
+         PSA_AEAD_FINISH_OUTPUT_SIZE( operation->key_type, operation->alg ) :
+         PSA_AEAD_VERIFY_OUTPUT_SIZE( operation->key_type, operation->alg );
+
+    if( output_size < finish_output_size )
+        return( PSA_ERROR_BUFFER_TOO_SMALL );
+
     return( PSA_SUCCESS );
 }
 
@@ -3775,7 +3785,7 @@
     *ciphertext_length = 0;
     *tag_length = tag_size;
 
-    status = psa_aead_final_checks( operation );
+    status = psa_aead_final_checks( operation, ciphertext_size );
 
     if( status != PSA_SUCCESS )
         goto exit;
@@ -3816,7 +3826,7 @@
 
     *plaintext_length = 0;
 
-    status = psa_aead_final_checks( operation );
+    status = psa_aead_final_checks( operation, plaintext_size );
 
     if( status != PSA_SUCCESS )
         goto exit;
diff --git a/library/psa_crypto_aead.c b/library/psa_crypto_aead.c
index 1a515a1..f2096ce 100644
--- a/library/psa_crypto_aead.c
+++ b/library/psa_crypto_aead.c
@@ -603,21 +603,11 @@
    mbedtls_psa_aead_verify() */
 static psa_status_t mbedtls_psa_aead_finish_checks(
     mbedtls_psa_aead_operation_t *operation,
-    size_t output_size,
     size_t tag_size )
 {
-    size_t finish_output_size;
-
     if( tag_size < operation->tag_length )
         return ( PSA_ERROR_BUFFER_TOO_SMALL );
 
-    finish_output_size = operation->is_encrypt ?
-        PSA_AEAD_FINISH_OUTPUT_SIZE( operation->key_type, operation->alg ) :
-        PSA_AEAD_VERIFY_OUTPUT_SIZE( operation->key_type, operation->alg );
-
-    if( output_size < finish_output_size )
-        return ( PSA_ERROR_BUFFER_TOO_SMALL );
-
     return ( PSA_SUCCESS );
 }
 
@@ -634,8 +624,7 @@
     psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
     size_t finish_output_size = 0;
 
-    status = mbedtls_psa_aead_finish_checks( operation, ciphertext_size,
-                                             tag_size );
+    status = mbedtls_psa_aead_finish_checks( operation, tag_size );
 
     if( status != PSA_SUCCESS )
         return status;
@@ -690,8 +679,7 @@
     int do_tag_check = 1;
     uint8_t check_tag[PSA_AEAD_TAG_MAX_SIZE];
 
-    status = mbedtls_psa_aead_finish_checks( operation, plaintext_size,
-                                             tag_length );
+    status = mbedtls_psa_aead_finish_checks( operation, tag_length );
 
     if( status != PSA_SUCCESS )
         return status;
diff --git a/library/psa_crypto_aead.h b/library/psa_crypto_aead.h
index 57b1b74..c664f9f 100644
--- a/library/psa_crypto_aead.h
+++ b/library/psa_crypto_aead.h
@@ -477,12 +477,10 @@
  * \retval #PSA_SUCCESS
  *         Success.
  * \retval #PSA_ERROR_BUFFER_TOO_SMALL
- *         The size of the \p ciphertext or \p tag buffer is too small.
- *         #PSA_AEAD_FINISH_OUTPUT_SIZE(\c key_type, \c alg) or
- *         #PSA_AEAD_FINISH_OUTPUT_MAX_SIZE can be used to determine the
- *         required \p ciphertext buffer size. #PSA_AEAD_TAG_LENGTH(\c key_type,
- *         \c key_bits, \c alg) or #PSA_AEAD_TAG_MAX_SIZE can be used to
- *         determine the required \p tag buffer size.
+ *         The size of the \p tag buffer is too small.
+ *         #PSA_AEAD_TAG_LENGTH(\c key_type, key_bits, \c alg) or
+ *         #PSA_AEAD_TAG_MAX_SIZE can be used to determine the required \p tag
+ *         buffer size.
  */
 psa_status_t mbedtls_psa_aead_finish(
     mbedtls_psa_aead_operation_t *operation,
@@ -551,10 +549,10 @@
  *         The calculations were successful, but the authentication tag is
  *         not correct.
  * \retval #PSA_ERROR_BUFFER_TOO_SMALL
- *         The size of the \p plaintext buffer is too small.
- *         #PSA_AEAD_VERIFY_OUTPUT_SIZE(\c key_type, \c alg) or
- *         #PSA_AEAD_VERIFY_OUTPUT_MAX_SIZE can be used to determine the
- *         required buffer size.
+ *         The size of the \p tag buffer is too small.
+ *         #PSA_AEAD_TAG_LENGTH(\c key_type, key_bits, \c alg) or
+ *         #PSA_AEAD_TAG_MAX_SIZE can be used to determine the required \p tag
+ *         buffer size.
  */
 psa_status_t mbedtls_psa_aead_verify(
     mbedtls_psa_aead_operation_t *operation,