diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 5d55e4543..14ef6e576 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -3282,13 +3282,10 @@ psa_status_t psa_aead_encrypt_setup( psa_aead_operation_t *operation,
 
 exit:
 
-    if( slot )
-    {
-        unlock_status = psa_unlock_key_slot( slot );
+    unlock_status = psa_unlock_key_slot( slot );
 
-        if( unlock_status != PSA_SUCCESS )
-            status = unlock_status;
-    }
+    if( status == PSA_SUCCESS )
+        status = unlock_status;
 
     if( status == PSA_SUCCESS )
         operation->alg = psa_aead_get_base_algorithm( alg );
@@ -3320,6 +3317,7 @@ psa_status_t psa_aead_decrypt_setup( psa_aead_operation_t *operation,
     }
 
     if( operation->nonce_set || operation->ad_started ||
+
         operation->body_started )
     {
         status = PSA_ERROR_BAD_STATE;
@@ -3347,13 +3345,10 @@ psa_status_t psa_aead_decrypt_setup( psa_aead_operation_t *operation,
 
 exit:
 
-    if( slot )
-    {
-        unlock_status = psa_unlock_key_slot( slot );
+    unlock_status = psa_unlock_key_slot( slot );
 
-        if( unlock_status != PSA_SUCCESS )
-            status = unlock_status;
-    }
+    if( status == PSA_SUCCESS )
+        status = unlock_status;
 
     if( status == PSA_SUCCESS )
         operation->alg = psa_aead_get_base_algorithm( alg );