From 0ee0d52967ce17b66364246e29a57530a6e07632 Mon Sep 17 00:00:00 2001
From: Steven Cooreman <steven.cooreman@silabs.com>
Date: Mon, 5 Oct 2020 16:03:42 +0200
Subject: [PATCH] Fix PSA crypto inconsistencies in agreement+derivation

* #3742 After input of a key as SECRET in the derivation, allow the
        derivation result to be used as key.

Signed-off-by: Steven Cooreman <steven.cooreman@silabs.com>
---
 library/psa_crypto.c                        |  5 +++++
 tests/suites/test_suite_psa_crypto.data     |  8 ++++++++
 tests/suites/test_suite_psa_crypto.function | 17 ++++++++++++++---
 3 files changed, 27 insertions(+), 3 deletions(-)

diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 71a505c70..69852fa56 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -5877,6 +5877,11 @@ static psa_status_t psa_key_agreement_internal( psa_key_derivation_operation_t *
                                                 shared_secret,
                                                 shared_secret_length );
 
+    /* If a private key has been added as SECRET, we allow the derived
+     * key material to be used as a key in PSA Crypto. */
+    if( step == PSA_KEY_DERIVATION_INPUT_SECRET )
+        operation->can_output_key = 1;
+
 exit:
     mbedtls_platform_zeroize( shared_secret, shared_secret_length );
     return( status );
diff --git a/tests/suites/test_suite_psa_crypto.data b/tests/suites/test_suite_psa_crypto.data
index b82e67f41..0626097bc 100644
--- a/tests/suites/test_suite_psa_crypto.data
+++ b/tests/suites/test_suite_psa_crypto.data
@@ -2221,6 +2221,14 @@ PSA key derivation: TLS 1.2 PSK-to-MS, SHA-256, PSK too long (160 Bytes)
 depends_on:MBEDTLS_SHA256_C
 derive_input:PSA_ALG_TLS12_PSK_TO_MS(PSA_ALG_SHA_256):PSA_KEY_DERIVATION_INPUT_SEED:PSA_KEY_TYPE_NONE:"":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_SECRET:PSA_KEY_TYPE_DERIVE:"01020304050607080102030405060708010203040506070801020304050607080102030405060708010203040506070801020304050607080102030405060708010203040506070801020304050607080102030405060708010203040506070801020304050607080102030405060708010203040506070801020304050607080102030405060708010203040506070801020304050607080102030405060708":PSA_ERROR_INVALID_ARGUMENT:PSA_KEY_DERIVATION_INPUT_LABEL:PSA_KEY_TYPE_NONE:"":PSA_ERROR_BAD_STATE:PSA_KEY_TYPE_NONE:PSA_ERROR_BAD_STATE
 
+PSA key derivation: ECDH on P256 with HKDF-SHA256, raw output
+depends_on:MBEDTLS_PK_PARSE_C:MBEDTLS_ECP_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED:MBEDTLS_ECDH_C:MBEDTLS_SHA256_C
+derive_input:PSA_ALG_KEY_AGREEMENT(PSA_ALG_ECDH, PSA_ALG_HKDF(PSA_ALG_SHA_256)):PSA_KEY_DERIVATION_INPUT_SALT:PSA_KEY_TYPE_NONE:"":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_SECRET:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1):"c88f01f510d9ac3f70a292daa2316de544e9aab8afe84049c62a9c57862d1433":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_INFO:PSA_KEY_TYPE_NONE:"":PSA_SUCCESS:PSA_KEY_TYPE_NONE:PSA_SUCCESS
+
+PSA key derivation: ECDH on P256 with HKDF-SHA256, key output
+depends_on:MBEDTLS_PK_PARSE_C:MBEDTLS_ECP_C:MBEDTLS_ECP_DP_SECP256R1_ENABLED:MBEDTLS_ECDH_C:MBEDTLS_SHA256_C
+derive_input:PSA_ALG_KEY_AGREEMENT(PSA_ALG_ECDH, PSA_ALG_HKDF(PSA_ALG_SHA_256)):PSA_KEY_DERIVATION_INPUT_SALT:PSA_KEY_TYPE_NONE:"":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_SECRET:PSA_KEY_TYPE_ECC_KEY_PAIR(PSA_ECC_FAMILY_SECP_R1):"c88f01f510d9ac3f70a292daa2316de544e9aab8afe84049c62a9c57862d1433":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_INFO:PSA_KEY_TYPE_NONE:"":PSA_SUCCESS:PSA_KEY_TYPE_RAW_DATA:PSA_SUCCESS
+
 PSA key derivation: HKDF invalid state (double generate + read past capacity)
 depends_on:MBEDTLS_SHA256_C
 test_derive_invalid_key_derivation_state:PSA_ALG_HKDF(PSA_ALG_SHA_256)
diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index fe9c1ccf9..328508dd2 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -4729,9 +4729,20 @@ void derive_input( int alg_arg,
             PSA_ASSERT( psa_import_key( &attributes,
                                         inputs[i]->x, inputs[i]->len,
                                         &handles[i] ) );
-            TEST_EQUAL( psa_key_derivation_input_key( &operation, steps[i],
-                                                      handles[i] ),
-                        expected_statuses[i] );
+            if( PSA_KEY_TYPE_IS_KEY_PAIR( key_types[i] ) &&
+                steps[i] == PSA_KEY_DERIVATION_INPUT_SECRET )
+            {
+                // When taking a private key as secret input, use key agreement
+                // to add the shared secret to the derivation
+                TEST_EQUAL( key_agreement_with_self( &operation, handles[i] ),
+                            expected_statuses[i] );
+            }
+            else
+            {
+                TEST_EQUAL( psa_key_derivation_input_key( &operation, steps[i],
+                                                          handles[i] ),
+                            expected_statuses[i] );
+            }
         }
         else
         {