diff --git a/include/psa/crypto.h b/include/psa/crypto.h
index 089484f..3107952 100644
--- a/include/psa/crypto.h
+++ b/include/psa/crypto.h
@@ -321,6 +321,10 @@
 #define PSA_ALG_CTR                             ((psa_algorithm_t)0x04800001)
 #define PSA_ALG_ARC4                            ((psa_algorithm_t)0x04800002)
 
+#define PSA_ALG_IS_STREAM_CIPHER(alg)            \
+    (((alg) & (PSA_ALG_CATEGORY_MASK | PSA_ALG_CIPHER_SUBCATEGORY_MASK)) == \
+        PSA_ALG_STREAM_CIPHER)
+
 #define PSA_ALG_CCM                             ((psa_algorithm_t)0x06000001)
 #define PSA_ALG_GCM                             ((psa_algorithm_t)0x06000002)
 
diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index fbc5949..2672627 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -1466,13 +1466,31 @@
 {
     int ret = MBEDTLS_ERR_CIPHER_FEATURE_UNAVAILABLE;
     
+    uint8_t temp_output_buffer[ MBEDTLS_MAX_BLOCK_LENGTH ];
+
     if( ! operation->key_set )
         return( PSA_ERROR_BAD_STATE );
     if( ! operation->iv_set )
         return( PSA_ERROR_BAD_STATE );
+    if( operation->ctx.cipher.operation == MBEDTLS_ENCRYPT )
+    {
+        if (operation->ctx.cipher.unprocessed_len > operation->block_size)
+            return( PSA_ERROR_INVALID_ARGUMENT );
+        if ( ( ( ( operation->alg ) & PSA_ALG_BLOCK_CIPHER_PAD_NONE ) == PSA_ALG_BLOCK_CIPHER_PAD_NONE )
+            && ( operation->ctx.cipher.unprocessed_len != 0 ) )
+            return(PSA_ERROR_INVALID_ARGUMENT);
+        if ( ( ( ( operation->alg ) & PSA_ALG_BLOCK_CIPHER_PAD_PKCS7 ) == PSA_ALG_BLOCK_CIPHER_PAD_PKCS7 )
+            && ( output_size != operation->block_size ) )
+            return(PSA_ERROR_INVALID_ARGUMENT);
+    }
+    if ( operation->ctx.cipher.operation == MBEDTLS_DECRYPT )
+        if (operation->ctx.cipher.unprocessed_len != 0)
+            return( PSA_ERROR_INVALID_ARGUMENT );
 
-    ret = mbedtls_cipher_finish( &operation->ctx.cipher, output,
-                                output_length );
+    ret = mbedtls_cipher_finish(&operation->ctx.cipher, temp_output_buffer,
+                                output_length);
+    if ( output_size > *output_length )
+        memcpy( temp_output_buffer, output, *output_length );
     if( ret != 0 )
     {
         psa_cipher_abort( operation );