diff --git a/include/psa/crypto.h b/include/psa/crypto.h
index 089484f..97819b7 100644
--- a/include/psa/crypto.h
+++ b/include/psa/crypto.h
@@ -1046,8 +1046,7 @@
                                psa_key_slot_t key,
                                psa_algorithm_t alg);
 
-psa_status_t psa_encrypt_generate_iv(psa_cipher_operation_t *operation,
-                                     unsigned char *iv,
+psa_status_t psa_encrypt_generate_iv(unsigned char *iv,
                                      size_t iv_size,
                                      size_t *iv_length);
 
diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 428a237..d349d19 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -1295,7 +1295,7 @@
 
 psa_status_t psa_decrypt_setup(psa_cipher_operation_t *operation,
                                psa_key_slot_t key,
-                               psa_algorithm_t alg);
+                               psa_algorithm_t alg)
 {
     int ret = MBEDTLS_ERR_CIPHER_FEATURE_UNAVAILABLE;
     psa_status_t status;
@@ -1343,13 +1343,12 @@
     return ( PSA_SUCCESS );
 }
 
-psa_status_t psa_encrypt_generate_iv(psa_cipher_operation_t *operation,
-                                     unsigned char *iv,
+psa_status_t psa_encrypt_generate_iv(unsigned char *iv,
                                      size_t iv_size,
                                      size_t *iv_length)
 {
     int ret = MBEDTLS_ERR_CIPHER_FEATURE_UNAVAILABLE;
-
+    
     ret = mbedtls_ctr_drbg_random( &global_data.ctr_drbg, iv, iv_size);
     if (ret != 0)
     {
@@ -1357,7 +1356,7 @@
     }
     
     *iv_length = iv_size;
-    retuen ( PSA_SUCCESS );
+    return ( PSA_SUCCESS );
 }
 
 psa_status_t psa_encrypt_set_iv(psa_cipher_operation_t *operation,
@@ -1388,6 +1387,9 @@
 {
     int ret = MBEDTLS_ERR_CIPHER_FEATURE_UNAVAILABLE;
 
+    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)
@@ -1405,6 +1407,9 @@
                                size_t *output_length)
 {
     int ret = MBEDTLS_ERR_CIPHER_FEATURE_UNAVAILABLE;
+
+    if ( output_size < operation->block_size )
+        return ( PSA_ERROR_BUFFER_TOO_SMALL );
     
     if( ! operation->key_set )
         return( PSA_ERROR_BAD_STATE );