From 258ae07fb01909a5582001f4d29f4cf2eb5b7265 Mon Sep 17 00:00:00 2001
From: gabor-mezei-arm <gabor.mezei@arm.com>
Date: Fri, 25 Jun 2021 15:25:38 +0200
Subject: [PATCH] Add checks for buffer size

Signed-off-by: gabor-mezei-arm <gabor.mezei@arm.com>
---
 library/psa_crypto.c        | 21 +++++++++++++++++++++
 library/psa_crypto_cipher.c |  6 ++++++
 2 files changed, 27 insertions(+)

diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index c1517e5c1a..d3f7c2eca9 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -3513,6 +3513,13 @@ psa_status_t psa_cipher_encrypt( mbedtls_svc_key_id_t key,
     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 @@ psa_status_t psa_cipher_decrypt( mbedtls_svc_key_id_t key,
       .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 ed43085bb6..4d5bf209cc 100644
--- a/library/psa_crypto_cipher.c
+++ b/library/psa_crypto_cipher.c
@@ -532,6 +532,12 @@ static psa_status_t cipher_decrypt( const psa_key_attributes_t *attributes,
     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;