diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 6575ec4..eb760d4 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -4102,10 +4102,7 @@
                                                           alg );
 
     if( status == PSA_SUCCESS )
-    {
         operation->accelerator_set = 1;
-        operation->key_set = 1;
-    }
 
     if( status != PSA_ERROR_NOT_SUPPORTED ||
         psa_key_lifetime_is_external( slot->attr.lifetime ) )
@@ -4172,7 +4169,6 @@
         goto exit;
 #endif //MBEDTLS_CIPHER_MODE_WITH_PADDING
 
-    operation->key_set = 1;
     operation->block_size = ( PSA_ALG_IS_STREAM_CIPHER( alg ) ? 1 :
                               PSA_BLOCK_CIPHER_BLOCK_SIZE( slot->attr.type ) );
     if( ( alg & PSA_ALG_CIPHER_FROM_BLOCK_FLAG ) != 0 &&
@@ -4186,10 +4182,17 @@
         operation->iv_size = 12;
 #endif
 
+    status = PSA_SUCCESS;
+
 exit:
-    if( status == 0 )
+    if( ret != 0 )
         status = mbedtls_to_psa_error( ret );
-    if( status != 0 )
+    if( status == PSA_SUCCESS )
+    {
+        /* Update operation flags for both driver and software implementations */
+        operation->key_set = 1;
+    }
+    else
         psa_cipher_abort( operation );
     return( status );
 }
@@ -4215,6 +4218,10 @@
 {
     psa_status_t status;
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
+    if( operation->iv_set || ! operation->iv_required )
+    {
+        return( PSA_ERROR_BAD_STATE );
+    }
 
     if( operation->accelerator_set == 1 )
     {
@@ -4225,10 +4232,6 @@
         goto exit;
     }
 
-    if( operation->iv_set || ! operation->iv_required )
-    {
-        return( PSA_ERROR_BAD_STATE );
-    }
     if( iv_size < operation->iv_size )
     {
         status = PSA_ERROR_BUFFER_TOO_SMALL;
@@ -4246,7 +4249,9 @@
     status = psa_cipher_set_iv( operation, iv, *iv_length );
 
 exit:
-    if( status != PSA_SUCCESS )
+    if( status == PSA_SUCCESS )
+        operation->iv_set = 1;
+    else
         psa_cipher_abort( operation );
     return( status );
 }
@@ -4257,6 +4262,10 @@
 {
     psa_status_t status;
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
+    if( operation->iv_set || ! operation->iv_required )
+    {
+        return( PSA_ERROR_BAD_STATE );
+    }
 
     if( operation->accelerator_set == 1 )
     {
@@ -4266,10 +4275,6 @@
         goto exit;
     }
 
-    if( operation->iv_set || ! operation->iv_required )
-    {
-        return( PSA_ERROR_BAD_STATE );
-    }
     if( iv_length != operation->iv_size )
     {
         status = PSA_ERROR_INVALID_ARGUMENT;
@@ -4382,6 +4387,14 @@
 {
     psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
     size_t expected_output_size;
+    if( operation->alg == 0 )
+    {
+        return( PSA_ERROR_BAD_STATE );
+    }
+    if( operation->iv_required && ! operation->iv_set )
+    {
+        return( PSA_ERROR_BAD_STATE );
+    }
 
     if( operation->accelerator_set == 1 )
     {
@@ -4394,15 +4407,6 @@
         goto exit;
     }
 
-    if( operation->alg == 0 )
-    {
-        return( PSA_ERROR_BAD_STATE );
-    }
-    if( operation->iv_required && ! operation->iv_set )
-    {
-        return( PSA_ERROR_BAD_STATE );
-    }
-
     if( ! PSA_ALG_IS_STREAM_CIPHER( operation->alg ) )
     {
         /* Take the unprocessed partial block left over from previous
@@ -4456,6 +4460,14 @@
     psa_status_t status = PSA_ERROR_GENERIC_ERROR;
     int cipher_ret = MBEDTLS_ERR_CIPHER_FEATURE_UNAVAILABLE;
     uint8_t temp_output_buffer[MBEDTLS_MAX_BLOCK_LENGTH];
+    if( operation->alg == 0 )
+    {
+        return( PSA_ERROR_BAD_STATE );
+    }
+    if( operation->iv_required && ! operation->iv_set )
+    {
+        return( PSA_ERROR_BAD_STATE );
+    }
 
     if( operation->accelerator_set == 1 )
     {
@@ -4470,11 +4482,6 @@
         return( status );
     }
 
-    if( operation->iv_required && ! operation->iv_set )
-    {
-        return( PSA_ERROR_BAD_STATE );
-    }
-
     if( operation->ctx.cipher.unprocessed_len != 0 )
     {
         if( operation->alg == PSA_ALG_ECB_NO_PADDING ||