Add checks for buffer size

Signed-off-by: gabor-mezei-arm <gabor.mezei@arm.com>
diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index c1517e5..d3f7c2e 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -3513,6 +3513,13 @@
     key_type = slot->attr.type;
     iv_length = PSA_CIPHER_IV_LENGTH( key_type, alg );
 
+    if( output_size < PSA_CIPHER_ENCRYPT_OUTPUT_SIZE( key_type, alg,
+                                                      input_length ) )
+    {
+        status = PSA_ERROR_BUFFER_TOO_SMALL;
+        goto exit;
+    }
+
     if( iv_length > 0 )
     {
         if( output_size < iv_length )
@@ -3564,11 +3571,25 @@
       .core = slot->attr
     };
 
+    if( input_length < PSA_CIPHER_IV_LENGTH( slot->attr.type, alg ) )
+    {
+        status = PSA_ERROR_INVALID_ARGUMENT;
+        goto exit;
+    }
+
+    if( output_size < PSA_CIPHER_DECRYPT_OUTPUT_SIZE( slot->attr.type, alg,
+                                                      input_length ) )
+    {
+        status = PSA_ERROR_BUFFER_TOO_SMALL;
+        goto exit;
+    }
+
     status = psa_driver_wrapper_cipher_decrypt(
         &attributes, slot->key.data, slot->key.bytes,
         alg, input, input_length,
         output, output_size, output_length );
 
+exit:
     unlock_status = psa_unlock_key_slot( slot );
 
     return( ( status == PSA_SUCCESS ) ? unlock_status : status );
diff --git a/library/psa_crypto_cipher.c b/library/psa_crypto_cipher.c
index ed43085..4d5bf20 100644
--- a/library/psa_crypto_cipher.c
+++ b/library/psa_crypto_cipher.c
@@ -532,6 +532,12 @@
     if( status != PSA_SUCCESS )
         goto exit;
 
+    if( output_size < accumulated_length )
+    {
+        status = PSA_ERROR_BUFFER_TOO_SMALL;
+        goto exit;
+    }
+
     status = cipher_finish( &operation, output + accumulated_length,
                             output_size - accumulated_length, &olength );
     accumulated_length += olength;