diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index c1071b0..714e556b 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -3760,6 +3760,18 @@
     return( status );
 }
 
+static psa_status_t psa_aead_final_checks( psa_aead_operation_t *operation )
+{
+    if( operation->id == 0 || operation->nonce_set == 0 )
+        return( PSA_ERROR_BAD_STATE );
+
+    if( operation->lengths_set && (operation->ad_remaining != 0 ||
+                                   operation->body_remaining != 0 ) )
+        return( PSA_ERROR_INVALID_ARGUMENT );
+
+    return( PSA_SUCCESS );
+}
+
 /* Finish encrypting a message in a multipart AEAD operation. */
 psa_status_t psa_aead_finish( psa_aead_operation_t *operation,
                               uint8_t *ciphertext,
@@ -3774,25 +3786,17 @@
     *ciphertext_length = 0;
     *tag_length = tag_size;
 
-    if( operation->id == 0 )
+    status = psa_aead_final_checks( operation );
+
+    if( status != PSA_SUCCESS )
+        goto exit;
+
+    if( operation->is_encrypt == 0 )
     {
         status = PSA_ERROR_BAD_STATE;
         goto exit;
     }
 
-    if( !operation->nonce_set || operation->is_encrypt == 0 )
-    {
-        status = PSA_ERROR_BAD_STATE;
-        goto exit;
-    }
-
-    if( operation->lengths_set && (operation->ad_remaining != 0 ||
-                                   operation->body_remaining != 0 ) )
-    {
-        status = PSA_ERROR_INVALID_ARGUMENT;
-        goto exit;
-    }
-
     status = psa_driver_wrapper_aead_finish( operation, ciphertext,
                                              ciphertext_size,
                                              ciphertext_length,
@@ -3823,24 +3827,21 @@
 
     *plaintext_length = 0;
 
-    if( operation->id == 0 )
+    status = psa_aead_final_checks( operation );
+
+    if( status != PSA_SUCCESS )
+        goto exit;
+
+    if( operation->is_encrypt == 1 )
     {
         status = PSA_ERROR_BAD_STATE;
         goto exit;
     }
 
-    if( !operation->nonce_set || operation->is_encrypt == 1 )
-    {
-        status = PSA_ERROR_BAD_STATE;
-        goto exit;
-    }
+    status = psa_aead_final_checks( operation );
 
-    if( operation->lengths_set && (operation->ad_remaining != 0 ||
-                                   operation->body_remaining != 0 ) )
-    {
-        status = PSA_ERROR_INVALID_ARGUMENT;
+    if( status != PSA_SUCCESS )
         goto exit;
-    }
 
     status = psa_driver_wrapper_aead_verify( operation, plaintext,
                                              plaintext_size,