From 44b09d2a67a4c33eb39d4eb4ebbe3e9877a7c22d Mon Sep 17 00:00:00 2001
From: itayzafrir <itay.zafrir@arm.com>
Date: Thu, 12 Jul 2018 13:06:41 +0300
Subject: [PATCH] psa: programs: Add cipher AES CTR multipart

---
 programs/psa/crypto_examples.c | 50 ++++++++++++++++++++++++++++++++++
 1 file changed, 50 insertions(+)

diff --git a/programs/psa/crypto_examples.c b/programs/psa/crypto_examples.c
index e2d2cf553..a7c36fa04 100644
--- a/programs/psa/crypto_examples.c
+++ b/programs/psa/crypto_examples.c
@@ -241,6 +241,51 @@ exit:
     return( status );
 }
 
+static psa_status_t cipher_example_encrypt_decrypt_aes_ctr_multi( void )
+{
+    enum {
+        block_size = PSA_BLOCK_CIPHER_BLOCK_SIZE( PSA_KEY_TYPE_AES ),
+        key_bits = 256,
+        input_size = 100,
+        part_size = 10,
+    };
+    const psa_algorithm_t alg = PSA_ALG_CTR;
+
+    psa_status_t status;
+    size_t output_len = 0;
+    uint8_t iv[block_size], input[input_size], encrypt[input_size],
+            decrypt[input_size];
+
+    status = psa_generate_random( input, sizeof( input ) );
+    ASSERT_STATUS( status, PSA_SUCCESS );
+
+    status = set_key_policy( key_slot_cipher,
+                             PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT,
+                             alg );
+    ASSERT_STATUS( status, PSA_SUCCESS );
+
+    status = psa_generate_key( key_slot_cipher, PSA_KEY_TYPE_AES, key_bits,
+                               NULL, 0 );
+    ASSERT_STATUS( status, PSA_SUCCESS );
+
+    status = cipher_encrypt( key_slot_cipher, alg, iv, sizeof( iv ),
+                             input, sizeof( input ), part_size,
+                             encrypt, sizeof( encrypt ), &output_len );
+    ASSERT_STATUS( status, PSA_SUCCESS );
+
+    status = cipher_decrypt( key_slot_cipher, alg, iv, sizeof( iv ),
+                             encrypt, output_len, part_size,
+                             decrypt, sizeof( decrypt ), &output_len );
+    ASSERT_STATUS( status, PSA_SUCCESS );
+
+    status = memcmp( input, decrypt, sizeof( input ) );
+    ASSERT_STATUS( status, PSA_SUCCESS );
+
+exit:
+    psa_destroy_key( key_slot_cipher );
+    return( status );
+}
+
 static void cipher_examples( void )
 {
     psa_status_t status;
@@ -254,6 +299,11 @@ static void cipher_examples( void )
     status = cipher_example_encrypt_decrypt_aes_cbc_pkcs7_multi( );
     if( status == PSA_SUCCESS )
         mbedtls_printf( "\tsuccess!\r\n" );
+
+    mbedtls_printf( "cipher encrypt/decrypt AES CTR multipart:\r\n" );
+    status = cipher_example_encrypt_decrypt_aes_ctr_multi( );
+    if( status == PSA_SUCCESS )
+        mbedtls_printf( "\tsuccess!\r\n" );
 }
 
 int main( void )