Add IV generation to one-shot cipher operation

The functions psa_cipher_encrypt and psa_cipher_decrypt are
one-shot operations that can take an arbitrary sized input.
These operations are implemented as client-side functions
that use multi-part cipher operations to allow large inputs
to be handled. The existing implementations were missing the
generation and setting of the IV at the start of the data.
This was leading to PSA Arch test failures (248 & 249). This
commit adds the missing IV handling and resolves the test
failures.

Signed-off-by: Julian Hall <julian.hall@arm.com>
Change-Id: I959980a8a83560c97f6c8d3840965b01ba4e434c
diff --git a/components/service/crypto/client/psa/psa_cipher.c b/components/service/crypto/client/psa/psa_cipher.c
index 3ab8ea2..24a8c28 100644
--- a/components/service/crypto/client/psa/psa_cipher.c
+++ b/components/service/crypto/client/psa/psa_cipher.c
@@ -8,7 +8,6 @@
 #include "psa_crypto_client.h"
 #include "crypto_caller_selector.h"
 
-
 psa_status_t psa_cipher_encrypt_setup(psa_cipher_operation_t *operation,
 	psa_key_id_t key,
 	psa_algorithm_t alg)
@@ -112,7 +111,13 @@
 		max_update_size = input_length;
 	}
 
-	while ((bytes_input < input_length) && (bytes_output < output_size)) {
+	while (bytes_input < input_length) {
+
+		if (bytes_output >= output_size) {
+
+			psa_status = PSA_ERROR_BUFFER_TOO_SMALL;
+			break;
+		}
 
 		size_t update_output_len = 0;
 		size_t bytes_remaining = input_length - bytes_input;
@@ -124,11 +129,8 @@
 			&input[bytes_input], update_len,
 			&output[bytes_output], output_size - bytes_output, &update_output_len);
 
-		if (psa_status != PSA_SUCCESS) {
-
-			psa_cipher_abort(operation);
+		if (psa_status != PSA_SUCCESS)
 			break;
-		}
 
 		bytes_input += update_len;
 		bytes_output += update_output_len;
@@ -150,7 +152,6 @@
 		}
 		else {
 
-			psa_cipher_abort(operation);
 			psa_status = PSA_ERROR_BUFFER_TOO_SMALL;
 		}
 	}
@@ -171,9 +172,31 @@
 
 	if (psa_status == PSA_SUCCESS) {
 
-		psa_status = multi_cipher_update(&operation,
-			input, input_length,
-			output, output_size, output_length);
+		size_t ciphertext_len = 0;
+		size_t iv_len = 0;
+
+		psa_status = psa_cipher_generate_iv(&operation, output, output_size, &iv_len);
+
+		if (psa_status == PSA_SUCCESS) {
+
+			if (iv_len <= output_size) {
+
+				psa_status = multi_cipher_update(&operation,
+					input, input_length,
+					&output[iv_len], output_size - iv_len, &ciphertext_len);
+
+				*output_length = iv_len + ciphertext_len;
+			}
+			else {
+
+				psa_status = PSA_ERROR_BUFFER_TOO_SMALL;
+			}
+		}
+
+		if (psa_status != PSA_SUCCESS) {
+
+			psa_cipher_abort(&operation);
+		}
 	}
 
 	return psa_status;
@@ -187,14 +210,41 @@
 	size_t output_size,
 	size_t *output_length)
 {
-	psa_cipher_operation_t operation = psa_cipher_operation_init();
-	psa_status_t psa_status = psa_cipher_decrypt_setup(&operation, key, alg);
+	psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
+	psa_status_t psa_status = psa_get_key_attributes(key, &attributes);
 
 	if (psa_status == PSA_SUCCESS) {
 
-		psa_status = multi_cipher_update(&operation,
-			input, input_length,
-			output, output_size, output_length);
+		psa_cipher_operation_t operation = psa_cipher_operation_init();
+		psa_status = psa_cipher_decrypt_setup(&operation, key, alg);
+
+		if (psa_status == PSA_SUCCESS) {
+
+			size_t iv_len = PSA_CIPHER_IV_LENGTH(psa_get_key_type(&attributes), alg);
+
+			if (input_length >= iv_len) {
+
+				psa_status = psa_cipher_set_iv(&operation, input, iv_len);
+
+				if (psa_status == PSA_SUCCESS) {
+
+					psa_status = multi_cipher_update(&operation,
+						&input[iv_len], input_length - iv_len,
+						output, output_size, output_length);
+				}
+			}
+			else {
+
+				psa_status = PSA_ERROR_INVALID_ARGUMENT;
+			}
+
+			if (psa_status != PSA_SUCCESS) {
+
+				psa_cipher_abort(&operation);
+			}
+		}
+
+		psa_reset_key_attributes(&attributes);
 	}
 
 	return psa_status;