From b4eedf7a234f9ee5a0e16e4743ca1ef8fa0e8eb9 Mon Sep 17 00:00:00 2001
From: Andrzej Kurek <andrzej.kurek@arm.com>
Date: Fri, 15 Apr 2022 05:41:14 -0400
Subject: [PATCH] Test failing raw_key_agreement in ssl mock tests

Force a bitflip in server key to make the raw key
agreement fail, and then verify that no key slots
are left open at the end. Use a Weierstrass curve
to have a high chance of failure upon encountering
such bitflip.
Signed-off-by: Andrzej Kurek <andrzej.kurek@arm.com>
---
 tests/suites/test_suite_ssl.data     |  3 ++
 tests/suites/test_suite_ssl.function | 76 ++++++++++++++++++++++++++++
 2 files changed, 79 insertions(+)

diff --git a/tests/suites/test_suite_ssl.data b/tests/suites/test_suite_ssl.data
index fe7d97825..15d9c5f85 100644
--- a/tests/suites/test_suite_ssl.data
+++ b/tests/suites/test_suite_ssl.data
@@ -10683,3 +10683,6 @@ ssl_cf_memcpy_offset:0:255:32
 # we could get this with 255-bytes plaintext and untruncated SHA-384
 Constant-flow memcpy from offset: large
 ssl_cf_memcpy_offset:100:339:48
+
+Raw key agreement fail
+raw_key_agreement_fail:
diff --git a/tests/suites/test_suite_ssl.function b/tests/suites/test_suite_ssl.function
index 4717264d3..8bf1cacd9 100644
--- a/tests/suites/test_suite_ssl.function
+++ b/tests/suites/test_suite_ssl.function
@@ -4555,3 +4555,79 @@ exit:
     mbedtls_free( src );
 }
 /* END_CASE */
+
+/* BEGIN_CASE depends_on:MBEDTLS_USE_PSA_CRYPTO */
+void raw_key_agreement_fail( )
+{
+    enum { BUFFSIZE = 17000 };
+    mbedtls_endpoint client, server;
+    mbedtls_psa_stats_t stats;
+
+#if defined(MBEDTLS_TIMING_C)
+    mbedtls_timing_delay_context timer_client, timer_server;
+#endif
+    mbedtls_ecp_group_id curve_list[] = { MBEDTLS_ECP_DP_SECP256R1,
+                                          MBEDTLS_ECP_DP_NONE };
+
+    mbedtls_test_message_queue server_queue, client_queue;
+    mbedtls_test_message_socket_context server_context, client_context;
+
+    mbedtls_message_socket_init( &server_context );
+    mbedtls_message_socket_init( &client_context );
+
+    USE_PSA_INIT( );
+
+    /* Client side, force SECP256R1 to make one key bitflip fail
+     * the raw key agreement. Flipping one bit with a Weierstrass
+     * curve (as opposed to a Montgomery curve) has a high chance of
+     * making it invalid. */
+    TEST_ASSERT( mbedtls_endpoint_init( &client, MBEDTLS_SSL_IS_CLIENT,
+                                        MBEDTLS_PK_RSA, &client_context,
+                                        &client_queue,
+                                        &server_queue, curve_list ) == 0 );
+#if defined(MBEDTLS_TIMING_C)
+    mbedtls_ssl_set_timer_cb( &client.ssl, &timer_client,
+                              mbedtls_timing_set_delay,
+                              mbedtls_timing_get_delay );
+#endif
+
+    /* Server side */
+    TEST_ASSERT( mbedtls_endpoint_init( &server, MBEDTLS_SSL_IS_SERVER,
+                                        MBEDTLS_PK_RSA, &server_context,
+                                        &server_queue,
+                                        &client_queue, NULL ) == 0 );
+#if defined(MBEDTLS_TIMING_C)
+    mbedtls_ssl_set_timer_cb( &server.ssl, &timer_server,
+                              mbedtls_timing_set_delay,
+                              mbedtls_timing_get_delay );
+#endif
+
+    TEST_ASSERT( mbedtls_mock_socket_connect( &(client.socket),
+                                              &(server.socket),
+                                              BUFFSIZE ) == 0 );
+
+    TEST_ASSERT( mbedtls_move_handshake_to_state( &(client.ssl),
+                                                  &(server.ssl),
+                                                  MBEDTLS_SSL_CLIENT_KEY_EXCHANGE )
+                 ==  0 );
+
+    /* Force a simulated bitflip in the server key. to make the
+     * raw key agreement in ssl_write_client_key_exchange fail. */
+    (client.ssl).handshake->ecdh_psa_peerkey[5] ^= 0x02;
+
+    TEST_ASSERT( mbedtls_move_handshake_to_state( &(client.ssl),
+                                                  &(server.ssl),
+                                                  MBEDTLS_SSL_HANDSHAKE_OVER )
+                 !=  0 );
+
+    mbedtls_psa_get_stats( &stats );
+
+    /* Make sure that the key slot is destroyed properly in case of failure. */
+    TEST_ASSERT( stats.empty_slots == MBEDTLS_PSA_KEY_SLOT_COUNT );
+
+exit:
+    mbedtls_endpoint_free( &client, &client_context );
+    mbedtls_endpoint_free( &server, &server_context );
+    USE_PSA_DONE( );
+}
+/* END_CASE */