derive_output mix-psk test: add more cases for derivation of output key

Signed-off-by: Przemek Stekiel <przemyslaw.stekiel@mobica.com>
diff --git a/tests/suites/test_suite_psa_crypto.data b/tests/suites/test_suite_psa_crypto.data
index c5354f0..637188d 100644
--- a/tests/suites/test_suite_psa_crypto.data
+++ b/tests/suites/test_suite_psa_crypto.data
@@ -5424,10 +5424,26 @@
 depends_on:PSA_WANT_ALG_SHA_256:PSA_WANT_ALG_TLS12_PSK_TO_MS
 derive_output:PSA_ALG_TLS12_PSK_TO_MS(PSA_ALG_SHA_256):PSA_KEY_DERIVATION_INPUT_SEED:"5bc0b19b4a8b24b07afe7ec65c471e94a7d518fcef06c3574315255c52afe21b5bc0b19b872b9b26508458f03603744d575f463a11ae7f1b090c012606fd3e9f":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_OTHER_SECRET:"c4eb02cb10175ab8a33aeeb068ba23df08206b0e":PSA_ERROR_INVALID_ARGUMENT:PSA_KEY_DERIVATION_INPUT_SECRET:"01020304":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_LABEL:"6d617374657220736563726574":PSA_SUCCESS:"":48:"2d6cba7ad474ac3dfcc01516ed7161a9b00c2e8a35b3e921b9eb09b546a54c98491b21d1baafb659c9094b760144ea1f":"":11:1:0
 
-PSA key derivation: TLS 1.2 Mix-PSK-to-MS, output key not permitted
+# output key test: secret passed as key, other secret passed as key
+PSA key derivation: TLS 1.2 Mix-PSK-to-MS, output key ok #1
+depends_on:PSA_WANT_ALG_SHA_256:PSA_WANT_ALG_TLS12_PSK_TO_MS
+derive_output:PSA_ALG_TLS12_PSK_TO_MS(PSA_ALG_SHA_256):PSA_KEY_DERIVATION_INPUT_SEED:"5bc0b19b4a8b24b07afe7ec65c471e94a7d518fcef06c3574315255c52afe21b5bc0b19b872b9b26508458f03603744d575f463a11ae7f1b090c012606fd3e9f":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_OTHER_SECRET:"c4eb02cb10175ab8a33aeeb068ba23df08206b0e":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_SECRET:"01020304":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_LABEL:"6d617374657220736563726574":PSA_SUCCESS:"":48:"2d6cba7ad474ac3dfcc01516ed7161a9b00c2e8a35b3e921b9eb09b546a54c98491b21d1baafb659c9094b760144ea1f":"":1:1:1
+
+# output key test: secret passed as key, other secret passed as bytes
+PSA key derivation: TLS 1.2 Mix-PSK-to-MS, output key ok #2
+depends_on:PSA_WANT_ALG_SHA_256:PSA_WANT_ALG_TLS12_PSK_TO_MS
+derive_output:PSA_ALG_TLS12_PSK_TO_MS(PSA_ALG_SHA_256):PSA_KEY_DERIVATION_INPUT_SEED:"5bc0b19b4a8b24b07afe7ec65c471e94a7d518fcef06c3574315255c52afe21b5bc0b19b872b9b26508458f03603744d575f463a11ae7f1b090c012606fd3e9f":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_OTHER_SECRET:"c4eb02cb10175ab8a33aeeb068ba23df08206b0e":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_SECRET:"01020304":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_LABEL:"6d617374657220736563726574":PSA_SUCCESS:"":48:"2d6cba7ad474ac3dfcc01516ed7161a9b00c2e8a35b3e921b9eb09b546a54c98491b21d1baafb659c9094b760144ea1f":"":0:1:1
+
+# output key test: secret passed as bytes, other secret passed as key
+PSA key derivation: TLS 1.2 Mix-PSK-to-MS, output key not permitted #1
 depends_on:PSA_WANT_ALG_SHA_256:PSA_WANT_ALG_TLS12_PSK_TO_MS
 derive_output:PSA_ALG_TLS12_PSK_TO_MS(PSA_ALG_SHA_256):PSA_KEY_DERIVATION_INPUT_SEED:"5bc0b19b4a8b24b07afe7ec65c471e94a7d518fcef06c3574315255c52afe21b5bc0b19b872b9b26508458f03603744d575f463a11ae7f1b090c012606fd3e9f":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_OTHER_SECRET:"c4eb02cb10175ab8a33aeeb068ba23df08206b0e":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_SECRET:"01020304":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_LABEL:"6d617374657220736563726574":PSA_SUCCESS:"":48:"2d6cba7ad474ac3dfcc01516ed7161a9b00c2e8a35b3e921b9eb09b546a54c98491b21d1baafb659c9094b760144ea1f":"":1:0:1
 
+# output key test: secret passed as bytes, other secret passed as bytes
+PSA key derivation: TLS 1.2 Mix-PSK-to-MS, output key not permitted #2
+depends_on:PSA_WANT_ALG_SHA_256:PSA_WANT_ALG_TLS12_PSK_TO_MS
+derive_output:PSA_ALG_TLS12_PSK_TO_MS(PSA_ALG_SHA_256):PSA_KEY_DERIVATION_INPUT_SEED:"5bc0b19b4a8b24b07afe7ec65c471e94a7d518fcef06c3574315255c52afe21b5bc0b19b872b9b26508458f03603744d575f463a11ae7f1b090c012606fd3e9f":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_OTHER_SECRET:"c4eb02cb10175ab8a33aeeb068ba23df08206b0e":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_SECRET:"01020304":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_LABEL:"6d617374657220736563726574":PSA_SUCCESS:"":48:"2d6cba7ad474ac3dfcc01516ed7161a9b00c2e8a35b3e921b9eb09b546a54c98491b21d1baafb659c9094b760144ea1f":"":0:0:1
+
 PSA key derivation: HKDF SHA-256, request maximum capacity
 depends_on:PSA_WANT_ALG_HKDF:PSA_WANT_ALG_SHA_256
 derive_output:PSA_ALG_HKDF(PSA_ALG_SHA_256):PSA_KEY_DERIVATION_INPUT_SALT:"000102030405060708090a0b0c":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_SECRET:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_INFO:"f0f1f2f3f4f5f6f7f8f9":PSA_SUCCESS:0:"":PSA_SUCCESS:"":255 * PSA_HASH_LENGTH(PSA_ALG_SHA_256):"3cb25f25faacd57a90434f64d0362f2a2d2d0a90cf1a5a4c5db02d56ecc4c5bf34007208d5b887185865":"":0:1:0
diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index b7ffcf6..118b1c7 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -6932,6 +6932,7 @@
     psa_key_attributes_t attributes2 = PSA_KEY_ATTRIBUTES_INIT;
     psa_key_attributes_t attributes3 = PSA_KEY_ATTRIBUTES_INIT;
     psa_key_attributes_t attributes4 = PSA_KEY_ATTRIBUTES_INIT;
+    mbedtls_svc_key_id_t derived_key = MBEDTLS_SVC_KEY_ID_INIT;
     psa_status_t status;
     size_t i;
 
@@ -7055,9 +7056,12 @@
 
     if( derive_type == 1 ) // output key
     {
-        /* Test that output key derivation is not permitted when secret is
-         * passed using input bytes and other secret is passed using input key. */
-        mbedtls_svc_key_id_t derived_key = MBEDTLS_SVC_KEY_ID_INIT;
+        psa_status_t expected_status = PSA_ERROR_NOT_PERMITTED;
+
+        /* For output key derivation secret must be provided using
+           input key, otherwise operation is not permitted. */
+        if ( key_input_type == 1 )
+            expected_status = PSA_SUCCESS;
 
         psa_set_key_usage_flags( &attributes4, PSA_KEY_USAGE_EXPORT );
         psa_set_key_algorithm( &attributes4, alg );
@@ -7065,7 +7069,7 @@
         psa_set_key_bits( &attributes4, 48 );
 
         TEST_EQUAL( psa_key_derivation_output_key( &attributes4, &operation,
-                                        &derived_key ), PSA_ERROR_NOT_PERMITTED );
+                                        &derived_key ), expected_status );
     }
     else // output bytes
     {
@@ -7109,6 +7113,7 @@
     psa_key_derivation_abort( &operation );
     for( i = 0; i < ARRAY_LENGTH( keys ); i++ )
         psa_destroy_key( keys[i] );
+    psa_destroy_key( derived_key );
     PSA_DONE( );
 }
 /* END_CASE */