diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 14ef6e576c..fb74a0d852 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -3254,8 +3254,8 @@ psa_status_t psa_aead_encrypt_setup( psa_aead_operation_t *operation,
         goto exit;
     }
 
-    if( operation->nonce_set || operation->ad_started ||
-        operation->body_started )
+    if( operation->nonce_set || operation->lengths_set ||
+        operation->ad_started || operation->body_started )
     {
         status = PSA_ERROR_BAD_STATE;
         goto exit;
@@ -3316,9 +3316,8 @@ psa_status_t psa_aead_decrypt_setup( psa_aead_operation_t *operation,
         goto exit;
     }
 
-    if( operation->nonce_set || operation->ad_started ||
-
-        operation->body_started )
+   if( operation->nonce_set || operation->lengths_set ||
+        operation->ad_started || operation->body_started )
     {
         status = PSA_ERROR_BAD_STATE;
         goto exit;
@@ -3375,8 +3374,8 @@ psa_status_t psa_aead_generate_nonce( psa_aead_operation_t *operation,
         goto exit;
     }
 
-    if( operation->nonce_set || operation->ad_started ||
-        operation->body_started )
+    if( operation->nonce_set || operation->lengths_set ||
+        operation->ad_started || operation->body_started )
     {
         status = PSA_ERROR_BAD_STATE;
         goto exit;
@@ -3430,11 +3429,11 @@ psa_status_t psa_aead_set_nonce( psa_aead_operation_t *operation,
         goto exit;
     }
 
-    /* Not checking nonce size here as GCM spec allows almost abitrarily large
-     * nonces. Please note that we do not generally recommend the usage of
-     * nonces of greater length than PSA_AEAD_NONCE_MAX_SIZE, as large nonces
-     * are hashed to a shorter size, which can then lead to collisions if you
-       encrypt a very large number of messages. */
+    /* Not checking nonce size here as GCM spec allows almost arbitrarily
+     * large nonces. Please note that we do not generally recommend the usage
+     * of nonces of greater length than PSA_AEAD_NONCE_MAX_SIZE, as large
+     * nonces are hashed to a shorter size, which can then lead to collisions
+     * if you encrypt a very large number of messages.*/
 
     status = psa_driver_wrapper_aead_set_nonce( operation, nonce,
                                                 nonce_length );
@@ -3462,7 +3461,8 @@ psa_status_t psa_aead_set_lengths( psa_aead_operation_t *operation,
         goto exit;
     }
 
-    if( operation->lengths_set )
+    if( operation->lengths_set || operation->ad_started ||
+        operation->body_started)
     {
         status = PSA_ERROR_BAD_STATE;
         goto exit;
@@ -3497,7 +3497,7 @@ psa_status_t psa_aead_update_ad( psa_aead_operation_t *operation,
         goto exit;
     }
 
-    if( !operation->nonce_set )
+    if( !operation->nonce_set || operation->body_started )
     {
         status = PSA_ERROR_BAD_STATE;
         goto exit;
@@ -3505,7 +3505,7 @@ psa_status_t psa_aead_update_ad( psa_aead_operation_t *operation,
 
     if( operation->lengths_set )
     {
-        if ( operation->ad_remaining < input_length )
+        if( operation->ad_remaining < input_length )
         {
             status = PSA_ERROR_INVALID_ARGUMENT;
             goto exit;
@@ -3546,7 +3546,7 @@ psa_status_t psa_aead_update( psa_aead_operation_t *operation,
         goto exit;
     }
 
-    if( !operation->nonce_set || !operation->ad_started )
+    if( !operation->nonce_set )
     {
         status = PSA_ERROR_BAD_STATE;
         goto exit;
@@ -3606,8 +3606,7 @@ psa_status_t psa_aead_finish( psa_aead_operation_t *operation,
         goto exit;
     }
 
-    if( !operation->nonce_set || !operation->ad_started ||
-        !operation->body_started )
+    if( !operation->nonce_set )
     {
         status = PSA_ERROR_BAD_STATE;
         goto exit;
@@ -3616,7 +3615,7 @@ psa_status_t psa_aead_finish( psa_aead_operation_t *operation,
     if( operation->lengths_set && (operation->ad_remaining != 0 ||
                                    operation->body_remaining != 0 ) )
     {
-        status = PSA_ERROR_BAD_STATE;
+        status = PSA_ERROR_INVALID_ARGUMENT;
         goto exit;
     }
 
@@ -3651,8 +3650,7 @@ psa_status_t psa_aead_verify( psa_aead_operation_t *operation,
         goto exit;
     }
 
-    if( !operation->nonce_set || !operation->ad_started ||
-        !operation->body_started )
+    if( !operation->nonce_set )
     {
         status = PSA_ERROR_BAD_STATE;
         goto exit;