diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 0d309c0..0761978 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -286,41 +286,6 @@
     }
 }
 
-static void psa_operation_init(void *operation,
-                                psa_algorithm_t alg)
-{
-    if( PSA_ALG_IS_MAC(alg) )
-    {
-        if ( ((psa_mac_operation_t*)operation)->alg != 0 ) //restart
-        {
-            ((psa_mac_operation_t*)operation)->alg = 0;
-            ((psa_mac_operation_t*)operation)->iv_required = 0;
-        }
-        else
-        {
-            ((psa_mac_operation_t*)operation)->alg = alg;
-            ((psa_mac_operation_t*)operation)->iv_required = 1;
-        }
-
-        ((psa_mac_operation_t*)operation)->key_set = 0;
-        ((psa_mac_operation_t*)operation)->iv_set = 0;
-        ((psa_mac_operation_t*)operation)->has_input = 0;
-        ((psa_mac_operation_t*)operation)->mac_size = 0;
-    }
-    else if( PSA_ALG_IS_CIPHER(alg) )
-    {
-        if ( ((psa_cipher_operation_t*)operation)->alg != 0 ) //restart
-            ((psa_cipher_operation_t*)operation)->alg = 0;
-        else
-            ((psa_cipher_operation_t*)operation)->alg = alg;
-
-        ((psa_cipher_operation_t*)operation)->key_set = 0;
-        ((psa_cipher_operation_t*)operation)->iv_set = 0;
-        ((psa_cipher_operation_t*)operation)->iv_size = 0;
-        ((psa_cipher_operation_t*)operation)->block_size = 0;
-    }
-}
-
 /****************************************************************/
 /* Key management */
 /****************************************************************/
@@ -992,7 +957,13 @@
 #endif /* MBEDTLS_MD_C */
                 return( PSA_ERROR_NOT_SUPPORTED );
     }
-    psa_operation_init(operation, 0);
+
+    operation->alg = 0;
+    operation->key_set = 0;	
+    operation->iv_set = 0;	
+    operation->iv_required = 0;	
+    operation->has_input = 0;
+
     return( PSA_SUCCESS );
 }
 
@@ -1007,7 +978,11 @@
     size_t key_bits;
     const mbedtls_cipher_info_t *cipher_info = NULL;
 
-    psa_operation_init(operation, alg);
+    operation->alg = 0;
+    operation->key_set = 0;	
+    operation->iv_set = 0;	
+    operation->iv_required = 1;	
+    operation->has_input = 0;
 
     status = psa_get_key_information( key, &key_type, &key_bits );
     if( status != PSA_SUCCESS )
@@ -1333,7 +1308,11 @@
     psa_algorithm_t padding_mode = PSA_ALG_BLOCK_CIPHER_PAD_NONE;
     mbedtls_cipher_padding_t mode = MBEDTLS_PADDING_NONE;
 
-    psa_operation_init(operation, alg);
+    operation->alg = alg;
+    operation->key_set = 0;	
+    operation->iv_set = 0;	
+    operation->iv_size = 0;	
+    operation->block_size = 0;
 
     status = psa_get_key_information( key, &key_type, &key_bits );
     if( status != PSA_SUCCESS )
@@ -1348,7 +1327,7 @@
 
     mbedtls_cipher_init( &operation->ctx.cipher );
     ret = mbedtls_cipher_setup( &operation->ctx.cipher, cipher_info );
-    if (ret != 0)
+    if( ret != 0 )
     {
         psa_cipher_abort( operation );
         return( mbedtls_to_psa_error( ret ) );
@@ -1356,14 +1335,14 @@
 
     ret = mbedtls_cipher_setkey( &operation->ctx.cipher, slot->data.raw.data,
                    key_bits, cipher_operation );
-    if (ret != 0)
+    if( ret != 0 )
     {
         psa_cipher_abort( operation );
         return( mbedtls_to_psa_error( ret ) );
     }
 
 #if defined(MBEDTLS_CIPHER_MODE_WITH_PADDING)
-    if (( alg & PSA_ALG_CBC_BASE) == PSA_ALG_CBC_BASE)
+    if( ( alg & PSA_ALG_CBC_BASE) == PSA_ALG_CBC_BASE )
     {
         padding_mode = alg & PSA_ALG_BLOCK_CIPHER_PADDING_MASK;
 
@@ -1376,10 +1355,10 @@
                 mode = MBEDTLS_PADDING_NONE;
                 break;
             default:
-                return ( PSA_ERROR_INVALID_PADDING );
+                return ( PSA_ERROR_INVALID_ARGUMENT );
         }
         ret = mbedtls_cipher_set_padding_mode( &operation->ctx.cipher, mode );
-        if (ret != 0)
+        if( ret != 0 )
             return( mbedtls_to_psa_error( ret ) );
     }
 #endif //MBEDTLS_CIPHER_MODE_WITH_PADDING
@@ -1387,9 +1366,9 @@
     operation->key_set = 1;
     operation->alg = alg;
     operation->block_size = PSA_BLOCK_CIPHER_BLOCK_SIZE(key_type);
-    if ( PSA_ALG_IS_BLOCK_CIPHER(alg) )
+    if( PSA_ALG_IS_BLOCK_CIPHER( alg ) || ( alg == PSA_ALG_CTR ) )
     {
-        operation->iv_size = operation->block_size;
+        operation->iv_size = PSA_BLOCK_CIPHER_BLOCK_SIZE(key_type);
     }
 
     return ( PSA_SUCCESS );
@@ -1414,28 +1393,39 @@
                                      size_t iv_size,
                                      size_t *iv_length)
 {
-    int ret = MBEDTLS_ERR_CIPHER_FEATURE_UNAVAILABLE;
-    if (iv_size < operation->iv_size)
-        return ( PSA_ERROR_BUFFER_TOO_SMALL );
-    
-    ret = mbedtls_ctr_drbg_random( &global_data.ctr_drbg, iv, operation->iv_size);
-    if (ret != 0)
+    int ret = PSA_SUCCESS;
+    if( operation->iv_set )
+        return( PSA_ERROR_BAD_STATE );
+    if( iv_size < operation->iv_size )
     {
-        return( mbedtls_to_psa_error( ret ) );       
+        ret = PSA_ERROR_BUFFER_TOO_SMALL;
+        goto exit;
+    }
+    ret = mbedtls_ctr_drbg_random( &global_data.ctr_drbg, iv, operation->iv_size);
+    if( ret != 0 )
+    {
+        ret = mbedtls_to_psa_error( ret );
+        goto exit;       
     }
     
     *iv_length = operation->iv_size;
-    return psa_encrypt_set_iv( operation, iv, *iv_length);
+    ret = psa_encrypt_set_iv( operation, iv, *iv_length);
+
+    exit:
+        if( ret != PSA_SUCCESS)
+        psa_cipher_abort( operation );
+        return( ret );
 }
 
 psa_status_t psa_encrypt_set_iv(psa_cipher_operation_t *operation,
                                 const unsigned char *iv,
                                 size_t iv_length)
 {
-    int ret = MBEDTLS_ERR_CIPHER_FEATURE_UNAVAILABLE;
-
+    int ret = PSA_SUCCESS;
+    if( operation->iv_set )
+        return( PSA_ERROR_BAD_STATE );
     ret =  mbedtls_cipher_set_iv( &operation->ctx.cipher, iv, iv_length );
-    if (ret != 0)
+    if( ret != 0 )
     {
         psa_cipher_abort( operation );
         return( mbedtls_to_psa_error( ret ) );
@@ -1455,12 +1445,12 @@
 {
     int ret = MBEDTLS_ERR_CIPHER_FEATURE_UNAVAILABLE;
 
-    if ( output_size < input_length )
+    if( output_size < input_length )
         return ( PSA_ERROR_BUFFER_TOO_SMALL );
 
     ret = mbedtls_cipher_update( &operation->ctx.cipher, input,
                    input_length, output, output_length );
-    if (ret != 0)
+    if( ret != 0 )
     {
         psa_cipher_abort( operation );
         return( mbedtls_to_psa_error( ret ) );
@@ -1482,7 +1472,7 @@
 
     ret = mbedtls_cipher_finish( &operation->ctx.cipher, output,
                                 output_length );
-    if (ret != 0)
+    if( ret != 0 )
     {
         psa_cipher_abort( operation );
         return( mbedtls_to_psa_error( ret ) );
@@ -1495,8 +1485,12 @@
 {   
     mbedtls_cipher_free( &operation->ctx.cipher );
     
-    psa_operation_init(operation, 0);
-    
+    operation->alg = 0;
+    operation->key_set = 0;	
+    operation->iv_set = 0;	
+    operation->iv_size = 0;	
+    operation->block_size = 0;
+
     return ( PSA_SUCCESS );
 }