From 330e460ef3c255216cb76b70f8d41d68b605c561 Mon Sep 17 00:00:00 2001
From: Przemyslaw Stekiel <przemyslaw.stekiel@mobica.com>
Date: Fri, 8 Oct 2021 12:26:21 +0200
Subject: [PATCH] Addapt psa_generate_key() tests

Signed-off-by: Przemyslaw Stekiel <przemyslaw.stekiel@mobica.com>
---
 tests/scripts/generate_psa_tests.py           | 17 +++++++++++------
 tests/suites/test_suite_psa_crypto.data       |  2 +-
 ...st_suite_psa_crypto_not_supported.function | 19 +++++++++++++++++++
 3 files changed, 31 insertions(+), 7 deletions(-)

diff --git a/tests/scripts/generate_psa_tests.py b/tests/scripts/generate_psa_tests.py
index 441b1c28d..ffb03b330 100755
--- a/tests/scripts/generate_psa_tests.py
+++ b/tests/scripts/generate_psa_tests.py
@@ -127,7 +127,7 @@ class Information:
         return constructors
 
 
-def test_case_for_key_type_not_supported(
+def test_case_for_key_type_not_supported_invalid_arg(
         verb: str, key_type: str, bits: int,
         dependencies: List[str],
         *args: str,
@@ -142,10 +142,15 @@ def test_case_for_key_type_not_supported(
     adverb = 'not' if dependencies else 'never'
     if param_descr:
         adverb = param_descr + ' ' + adverb
-    tc.set_description('PSA {} {} {}-bit {} supported'
-                       .format(verb, short_key_type, bits, adverb))
+    if (verb == "generate") and ("PUBLIC" in short_key_type):
+        tc.set_description('PSA {} {} {}-bit invalid argument'
+                    .format(verb, short_key_type, bits))
+        tc.set_function(verb + '_invalid_arg')
+    else:
+        tc.set_description('PSA {} {} {}-bit {} supported'
+                        .format(verb, short_key_type, bits, adverb))
+        tc.set_function(verb + '_not_supported')
     tc.set_dependencies(dependencies)
-    tc.set_function(verb + '_not_supported')
     tc.set_arguments([key_type] + list(args))
     return tc
 
@@ -186,7 +191,7 @@ class NotSupported:
         else:
             generate_dependencies = import_dependencies
         for bits in kt.sizes_to_test():
-            yield test_case_for_key_type_not_supported(
+            yield test_case_for_key_type_not_supported_invalid_arg(
                 'import', kt.expression, bits,
                 finish_family_dependencies(import_dependencies, bits),
                 test_case.hex_string(kt.key_material(bits)),
@@ -197,7 +202,7 @@ class NotSupported:
                 # supported or not depending on implementation capabilities,
                 # only generate the test case once.
                 continue
-            yield test_case_for_key_type_not_supported(
+            yield test_case_for_key_type_not_supported_invalid_arg(
                 'generate', kt.expression, bits,
                 finish_family_dependencies(generate_dependencies, bits),
                 str(bits),
diff --git a/tests/suites/test_suite_psa_crypto.data b/tests/suites/test_suite_psa_crypto.data
index 03e4a2760..580f90d45 100644
--- a/tests/suites/test_suite_psa_crypto.data
+++ b/tests/suites/test_suite_psa_crypto.data
@@ -3717,7 +3717,7 @@ generate_random:2 * MBEDTLS_CTR_DRBG_MAX_REQUEST + 1
 
 PSA generate key: bad type (RSA public key)
 depends_on:PSA_WANT_KEY_TYPE_RSA_PUBLIC_KEY
-generate_key:PSA_KEY_TYPE_RSA_PUBLIC_KEY:512:PSA_KEY_USAGE_EXPORT:0:PSA_ERROR_NOT_SUPPORTED:0
+generate_key:PSA_KEY_TYPE_RSA_PUBLIC_KEY:512:PSA_KEY_USAGE_EXPORT:0:PSA_ERROR_INVALID_ARGUMENT:0
 
 PSA generate key: raw data, 0 bits: invalid argument
 # The spec allows either INVALID_ARGUMENT or NOT_SUPPORTED
diff --git a/tests/suites/test_suite_psa_crypto_not_supported.function b/tests/suites/test_suite_psa_crypto_not_supported.function
index e3253d840..6b85fd75a 100644
--- a/tests/suites/test_suite_psa_crypto_not_supported.function
+++ b/tests/suites/test_suite_psa_crypto_not_supported.function
@@ -50,3 +50,22 @@ exit:
     PSA_DONE( );
 }
 /* END_CASE */
+
+/* BEGIN_CASE */
+void generate_invalid_arg( int key_type, int bits )
+{
+    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
+    mbedtls_svc_key_id_t key_id = INVALID_KEY_ID;
+
+    PSA_ASSERT( psa_crypto_init( ) );
+    psa_set_key_type( &attributes, key_type );
+    psa_set_key_bits( &attributes, bits );
+    TEST_EQUAL( psa_generate_key( &attributes, &key_id ),
+                PSA_ERROR_INVALID_ARGUMENT );
+    TEST_ASSERT( mbedtls_svc_key_id_equal( key_id, MBEDTLS_SVC_KEY_ID_INIT ) );
+
+exit:
+    psa_destroy_key( key_id );
+    PSA_DONE( );
+}
+/* END_CASE */