diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 39a6b72..e824123 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -3555,8 +3555,8 @@
         goto exit;
     }
 
-    if( operation->nonce_set || operation->lengths_set ||
-        operation->ad_started || operation->body_started )
+    if( operation->nonce_set || operation->ad_started ||
+        operation->body_started )
     {
         status = PSA_ERROR_BAD_STATE;
         goto exit;
@@ -3845,7 +3845,7 @@
     if( operation->lengths_set && (operation->ad_remaining != 0 ||
                                    operation->body_remaining != 0 ) )
     {
-        status = PSA_ERROR_BAD_STATE;
+        status = PSA_ERROR_INVALID_ARGUMENT;
         goto exit;
     }