From 2e46e9cf2117d5e92f78c8bf57f3dbf0b74847e1 Mon Sep 17 00:00:00 2001
From: Gilles Peskine <Gilles.Peskine@arm.com>
Date: Thu, 11 Apr 2019 21:24:55 +0200
Subject: [PATCH] Add exercise_key for raw key agreement

---
 tests/suites/test_suite_psa_crypto.function | 56 +++++++++++++++++++++
 1 file changed, 56 insertions(+)

diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index d70b7eb7f..2659081fd 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -597,6 +597,60 @@ exit:
     return( status );
 }
 
+/* We need two keys to exercise key agreement. Exercise the
+ * private key against its own public key. */
+static psa_status_t raw_key_agreement_with_self( psa_algorithm_t alg,
+                                                 psa_key_handle_t handle )
+{
+    psa_key_type_t private_key_type;
+    psa_key_type_t public_key_type;
+    size_t key_bits;
+    uint8_t *public_key = NULL;
+    size_t public_key_length;
+    uint8_t output[1024];
+    size_t output_length;
+    /* Return GENERIC_ERROR if something other than the final call to
+     * psa_key_agreement fails. This isn't fully satisfactory, but it's
+     * good enough: callers will report it as a failed test anyway. */
+    psa_status_t status = PSA_ERROR_GENERIC_ERROR;
+
+    PSA_ASSERT( psa_get_key_information( handle,
+                                         &private_key_type,
+                                         &key_bits ) );
+    public_key_type = PSA_KEY_TYPE_PUBLIC_KEY_OF_KEYPAIR( private_key_type );
+    public_key_length = PSA_KEY_EXPORT_MAX_SIZE( public_key_type, key_bits );
+    ASSERT_ALLOC( public_key, public_key_length );
+    PSA_ASSERT( psa_export_public_key( handle,
+                                       public_key, public_key_length,
+                                       &public_key_length ) );
+
+    status = psa_key_agreement_raw_shared_secret(
+        alg, handle,
+        public_key, public_key_length,
+        output, sizeof( output ), &output_length );
+exit:
+    mbedtls_free( public_key );
+    return( status );
+}
+
+static int exercise_raw_key_agreement_key( psa_key_handle_t handle,
+                                           psa_key_usage_t usage,
+                                           psa_algorithm_t alg )
+{
+    int ok = 0;
+
+    if( usage & PSA_KEY_USAGE_DERIVE )
+    {
+        /* We need two keys to exercise key agreement. Exercise the
+         * private key against its own public key. */
+        PSA_ASSERT( raw_key_agreement_with_self( alg, handle ) );
+    }
+    ok = 1;
+
+exit:
+    return( ok );
+}
+
 static int exercise_key_agreement_key( psa_key_handle_t handle,
                                        psa_key_usage_t usage,
                                        psa_algorithm_t alg )
@@ -973,6 +1027,8 @@ static int exercise_key( psa_key_handle_t handle,
         ok = exercise_asymmetric_encryption_key( handle, usage, alg );
     else if( PSA_ALG_IS_KEY_DERIVATION( alg ) )
         ok = exercise_key_derivation_key( handle, usage, alg );
+    else if( PSA_ALG_IS_RAW_KEY_AGREEMENT( alg ) )
+        ok = exercise_raw_key_agreement_key( handle, usage, alg );
     else if( PSA_ALG_IS_KEY_AGREEMENT( alg ) )
         ok = exercise_key_agreement_key( handle, usage, alg );
     else