From b46e7ca16bd1f06d3162d9ba8576cb44751f79f6 Mon Sep 17 00:00:00 2001
From: Nir Sonnenschein <nir.sonnenschein@arm.com>
Date: Thu, 25 Oct 2018 14:46:09 +0300
Subject: [PATCH] add additional generator tests and generalize key derivation
 test

Key derivation test now uses an indirect way to test generator validity
as the direct way previously used isn't compatible with the PSA IPC
implementation. Additional bad path test for the generator added
to check basic bad-path scenarios.
---
 tests/suites/test_suite_psa_crypto.data     |  8 ++-
 tests/suites/test_suite_psa_crypto.function | 80 +++++++++++++++++----
 2 files changed, 72 insertions(+), 16 deletions(-)

diff --git a/tests/suites/test_suite_psa_crypto.data b/tests/suites/test_suite_psa_crypto.data
index 1a93a8929..39ac88839 100644
--- a/tests/suites/test_suite_psa_crypto.data
+++ b/tests/suites/test_suite_psa_crypto.data
@@ -1209,9 +1209,13 @@ PSA key derivation: unsupported key derivation algorithm
 depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C
 derive_setup:PSA_KEY_TYPE_DERIVE:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_ALG_CATEGORY_KEY_DERIVATION:"":"":42:PSA_ERROR_NOT_SUPPORTED
 
-PSA key derivation: bad arguments test
+PSA key derivation: invalid generator state ( double generate + read past capacity )
 depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C
-test_derive_invalid_generator:
+test_derive_invalid_generator_state:PSA_KEY_TYPE_DERIVE:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b"
+
+PSA key derivation:  invalid generator state ( call read/get_capacity after init and abort )
+depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C
+test_derive_invalid_generator_tests:
 
 PSA key derivation: HKDF SHA-256, RFC5869 #1, output 42+0
 depends_on:MBEDTLS_MD_C:MBEDTLS_SHA256_C
diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index c6f49c007..65bec58c3 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -3086,25 +3086,77 @@ exit:
 /* END_CASE */
 
 /* BEGIN_CASE */
-void test_derive_invalid_generator()
+void test_derive_invalid_generator_state( int key_type_arg, data_t *key_data)
 {
-    psa_crypto_generator_t generator = PSA_CRYPTO_GENERATOR_INIT;
     psa_key_slot_t base_key = 1;
+    size_t key_type = key_type_arg;
+    psa_crypto_generator_t generator = PSA_CRYPTO_GENERATOR_INIT;
     psa_algorithm_t alg = PSA_ALG_HKDF(PSA_ALG_SHA_256);
-    data_t salt;
-    data_t label;
-    size_t capacity = 0;
-    salt.x = NULL;
-    salt.len = 0;
-    label.x = NULL;
-    label.len = 0;
+    size_t capacity = 42;
+    uint8_t buffer[42];
+    psa_key_policy_t policy;
 
-    generator.alg = alg;
-    /* invalid generator.alg */
+    TEST_ASSERT( psa_crypto_init( ) == PSA_SUCCESS );
+
+    psa_key_policy_init( &policy );
+    psa_key_policy_set_usage( &policy, PSA_KEY_USAGE_DERIVE, alg );
+    TEST_ASSERT( psa_set_key_policy( base_key, &policy ) == PSA_SUCCESS );
+
+    TEST_ASSERT( psa_import_key( base_key, key_type,
+                                 key_data->x,
+                                 key_data->len ) == PSA_SUCCESS );
+
+    /* valid key derivation */
     TEST_ASSERT(  psa_key_derivation( &generator, base_key, alg,
-                                     salt.x, salt.len,
-                                     label.x, label.len,
-                                     capacity ) == PSA_ERROR_BAD_STATE );
+                                      NULL, 0,
+                                      NULL, 0,
+                                      capacity ) == PSA_SUCCESS );
+
+    /* state of generator shouldn't allow additional generation */
+    TEST_ASSERT(  psa_key_derivation( &generator, base_key, alg,
+                                      NULL, 0,
+                                      NULL, 0,
+                                      capacity ) == PSA_ERROR_BAD_STATE );
+
+    TEST_ASSERT( psa_generator_read( &generator, buffer, capacity )
+                 == PSA_SUCCESS );
+
+    TEST_ASSERT( psa_generator_read( &generator, buffer, capacity )
+                 == PSA_ERROR_INSUFFICIENT_CAPACITY );
+
+
+exit:
+    psa_generator_abort( &generator );
+    psa_destroy_key( base_key );
+    mbedtls_psa_crypto_free( );
+}
+/* END_CASE */
+
+
+/* BEGIN_CASE */
+void test_derive_invalid_generator_tests( )
+{
+    uint8_t output_buffer[16];
+    size_t buffer_size = 16;
+    size_t capacity = 0;
+    psa_crypto_generator_t generator = PSA_CRYPTO_GENERATOR_INIT;
+
+    TEST_ASSERT( psa_generator_read( &generator, output_buffer, buffer_size)
+                 == PSA_ERROR_INSUFFICIENT_CAPACITY );
+
+    TEST_ASSERT( psa_get_generator_capacity( &generator, &capacity )
+                 == PSA_ERROR_BAD_STATE );
+
+    TEST_ASSERT( psa_generator_abort(&generator) == PSA_SUCCESS );
+
+    TEST_ASSERT( psa_generator_read( &generator, output_buffer, buffer_size)
+                 == PSA_ERROR_INSUFFICIENT_CAPACITY );
+
+    TEST_ASSERT( psa_get_generator_capacity( &generator, &capacity)
+                 == PSA_ERROR_BAD_STATE );
+
+exit:
+    psa_generator_abort( &generator );
 }
 /* END_CASE */