Merge pull request #5180 from daverodgman/key_derivation_output_key_error_code

Improve PSA error return code for psa_key_derivation_output_key
diff --git a/library/ssl_misc.h b/library/ssl_misc.h
index 362117f..2408fd1 100644
--- a/library/ssl_misc.h
+++ b/library/ssl_misc.h
@@ -1185,6 +1185,7 @@
 int mbedtls_ssl_flush_output( mbedtls_ssl_context *ssl );
 
 int mbedtls_ssl_tls13_process_finished_message( mbedtls_ssl_context *ssl );
+int mbedtls_ssl_tls13_write_finished_message( mbedtls_ssl_context *ssl );
 
 int mbedtls_ssl_parse_certificate( mbedtls_ssl_context *ssl );
 int mbedtls_ssl_write_certificate( mbedtls_ssl_context *ssl );
diff --git a/library/ssl_tls13_client.c b/library/ssl_tls13_client.c
index 6deab2a..a2e5f33 100644
--- a/library/ssl_tls13_client.c
+++ b/library/ssl_tls13_client.c
@@ -1611,6 +1611,7 @@
     if( ret != 0 )
         return( ret );
 
+    mbedtls_ssl_set_outbound_transform( ssl, ssl->handshake->transform_handshake );
     mbedtls_ssl_handshake_set_state( ssl, MBEDTLS_SSL_CLIENT_FINISHED );
     return( 0 );
 }
@@ -1618,9 +1619,14 @@
 /*
  * Handler for MBEDTLS_SSL_CLIENT_FINISHED
  */
-static int ssl_tls1_3_write_client_finished( mbedtls_ssl_context *ssl )
+static int ssl_tls13_write_client_finished( mbedtls_ssl_context *ssl )
 {
-    MBEDTLS_SSL_DEBUG_MSG( 1, ( "%s hasn't been implemented", __func__ ) );
+    int ret;
+
+    ret = mbedtls_ssl_tls13_write_finished_message( ssl );
+    if( ret != 0 )
+        return( ret );
+
     mbedtls_ssl_handshake_set_state( ssl, MBEDTLS_SSL_FLUSH_BUFFERS );
     return( 0 );
 }
@@ -1689,7 +1695,7 @@
             break;
 
         case MBEDTLS_SSL_CLIENT_FINISHED:
-            ret = ssl_tls1_3_write_client_finished( ssl );
+            ret = ssl_tls13_write_client_finished( ssl );
             break;
 
         case MBEDTLS_SSL_FLUSH_BUFFERS:
diff --git a/library/ssl_tls13_generic.c b/library/ssl_tls13_generic.c
index b2a70f3..f17bf99 100644
--- a/library/ssl_tls13_generic.c
+++ b/library/ssl_tls13_generic.c
@@ -29,6 +29,7 @@
 #include "mbedtls/debug.h"
 #include "mbedtls/oid.h"
 #include "mbedtls/platform.h"
+#include <string.h>
 
 #include "ssl_misc.h"
 #include "ssl_tls13_keys.h"
@@ -974,7 +975,7 @@
 }
 #endif /* MBEDTLS_SSL_CLI_C */
 
-static int ssl_tls13_postprocess_finished_message( mbedtls_ssl_context* ssl )
+static int ssl_tls13_postprocess_finished_message( mbedtls_ssl_context *ssl )
 {
 
 #if defined(MBEDTLS_SSL_CLI_C)
@@ -1014,6 +1015,93 @@
     return( ret );
 }
 
+/*
+ *
+ * STATE HANDLING: Write and send Finished message.
+ *
+ */
+/*
+ * Implement
+ */
+
+static int ssl_tls13_prepare_finished_message( mbedtls_ssl_context *ssl )
+{
+    int ret;
+
+    /* Compute transcript of handshake up to now. */
+    ret = mbedtls_ssl_tls13_calculate_verify_data( ssl,
+                    ssl->handshake->state_local.finished_out.digest,
+                    sizeof( ssl->handshake->state_local.finished_out.digest ),
+                    &ssl->handshake->state_local.finished_out.digest_len,
+                    ssl->conf->endpoint );
+
+    if( ret != 0 )
+    {
+         MBEDTLS_SSL_DEBUG_RET( 1, "calculate_verify_data failed", ret );
+        return( ret );
+    }
+
+    return( 0 );
+}
+
+static int ssl_tls13_finalize_finished_message( mbedtls_ssl_context *ssl )
+{
+    // TODO: Add back resumption keys calculation after MVP.
+    ((void) ssl);
+
+    return( 0 );
+}
+
+static int ssl_tls13_write_finished_message_body( mbedtls_ssl_context *ssl,
+                                                  unsigned char *buf,
+                                                  unsigned char *end,
+                                                  size_t *olen )
+{
+    size_t verify_data_len = ssl->handshake->state_local.finished_out.digest_len;
+    /*
+     * struct {
+     *     opaque verify_data[Hash.length];
+     * } Finished;
+     */
+    MBEDTLS_SSL_CHK_BUF_PTR( buf, end, verify_data_len );
+
+    memcpy( buf, ssl->handshake->state_local.finished_out.digest,
+            verify_data_len );
+
+    *olen = verify_data_len;
+    return( 0 );
+}
+
+/* Main entry point: orchestrates the other functions */
+int mbedtls_ssl_tls13_write_finished_message( mbedtls_ssl_context *ssl )
+{
+    int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
+    unsigned char *buf;
+    size_t buf_len, msg_len;
+
+    MBEDTLS_SSL_DEBUG_MSG( 2, ( "=> write finished message" ) );
+
+    MBEDTLS_SSL_PROC_CHK( ssl_tls13_prepare_finished_message( ssl ) );
+
+    MBEDTLS_SSL_PROC_CHK( mbedtls_ssl_tls13_start_handshake_msg( ssl,
+                              MBEDTLS_SSL_HS_FINISHED, &buf, &buf_len ) );
+
+    MBEDTLS_SSL_PROC_CHK( ssl_tls13_write_finished_message_body(
+                              ssl, buf, buf + buf_len, &msg_len ) );
+
+    mbedtls_ssl_tls1_3_add_hs_msg_to_checksum( ssl, MBEDTLS_SSL_HS_FINISHED,
+                                               buf, msg_len );
+
+    MBEDTLS_SSL_PROC_CHK( ssl_tls13_finalize_finished_message( ssl ) );
+    MBEDTLS_SSL_PROC_CHK( mbedtls_ssl_tls13_finish_handshake_msg( ssl,
+                                              buf_len, msg_len ) );
+    MBEDTLS_SSL_PROC_CHK( mbedtls_ssl_flush_output( ssl ) );
+
+cleanup:
+
+    MBEDTLS_SSL_DEBUG_MSG( 2, ( "<= write finished message" ) );
+    return( ret );
+}
 
 #endif /* MBEDTLS_SSL_PROTO_TLS1_3_EXPERIMENTAL */
 
diff --git a/programs/ssl/ssl_client2.c b/programs/ssl/ssl_client2.c
index 29bda7f..cb299b1 100644
--- a/programs/ssl/ssl_client2.c
+++ b/programs/ssl/ssl_client2.c
@@ -21,6 +21,10 @@
 
 #include "ssl_test_lib.h"
 
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+#include "test/psa_crypto_helpers.h"
+#endif
+
 #if defined(MBEDTLS_SSL_TEST_IMPOSSIBLE)
 int main( void )
 {
@@ -3015,6 +3019,19 @@
 
     mbedtls_net_free( &server_fd );
 
+    mbedtls_ssl_free( &ssl );
+    mbedtls_ssl_config_free( &conf );
+    mbedtls_ssl_session_free( &saved_session );
+
+    if( session_data != NULL )
+        mbedtls_platform_zeroize( session_data, session_data_len );
+    mbedtls_free( session_data );
+#if defined(MBEDTLS_SSL_CONTEXT_SERIALIZATION)
+    if( context_buf != NULL )
+        mbedtls_platform_zeroize( context_buf, context_buf_len );
+    mbedtls_free( context_buf );
+#endif
+
 #if defined(MBEDTLS_X509_CRT_PARSE_C)
     mbedtls_x509_crt_free( &clicert );
     mbedtls_x509_crt_free( &cacert );
@@ -3045,23 +3062,25 @@
 #endif /* MBEDTLS_KEY_EXCHANGE_SOME_PSK_ENABLED &&
           MBEDTLS_USE_PSA_CRYPTO */
 
-    mbedtls_ssl_session_free( &saved_session );
-    mbedtls_ssl_free( &ssl );
-    mbedtls_ssl_config_free( &conf );
-    rng_free( &rng );
-    if( session_data != NULL )
-        mbedtls_platform_zeroize( session_data, session_data_len );
-    mbedtls_free( session_data );
-#if defined(MBEDTLS_SSL_CONTEXT_SERIALIZATION)
-    if( context_buf != NULL )
-        mbedtls_platform_zeroize( context_buf, context_buf_len );
-    mbedtls_free( context_buf );
-#endif
-
 #if defined(MBEDTLS_USE_PSA_CRYPTO)
+    const char* message = mbedtls_test_helper_is_psa_leaking();
+    if( message )
+    {
+        if( ret == 0 )
+            ret = 1;
+        mbedtls_printf( "PSA memory leak detected: %s\n",  message);
+    }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
+
+    /* For builds with MBEDTLS_TEST_USE_PSA_CRYPTO_RNG psa crypto
+     * resources are freed by rng_free(). */
+#if defined(MBEDTLS_USE_PSA_CRYPTO) && \
+    !defined(MBEDTLS_TEST_USE_PSA_CRYPTO_RNG)
     mbedtls_psa_crypto_free( );
 #endif
 
+    rng_free( &rng );
+
 #if defined(MBEDTLS_TEST_HOOKS)
     if( test_hooks_failure_detected( ) )
     {
diff --git a/programs/ssl/ssl_server2.c b/programs/ssl/ssl_server2.c
index 4d785d7..0e6f78b 100644
--- a/programs/ssl/ssl_server2.c
+++ b/programs/ssl/ssl_server2.c
@@ -65,6 +65,10 @@
 #include <windows.h>
 #endif
 
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+#include "test/psa_crypto_helpers.h"
+#endif
+
 /* Size of memory to be allocated for the heap, when using the library's memory
  * management and MBEDTLS_MEMORY_BUFFER_ALLOC_C is enabled. */
 #define MEMORY_HEAP_SIZE        120000
@@ -3947,9 +3951,35 @@
     mbedtls_net_free( &client_fd );
     mbedtls_net_free( &listen_fd );
 
-#if defined(MBEDTLS_DHM_C) && defined(MBEDTLS_FS_IO)
-    mbedtls_dhm_free( &dhm );
+    mbedtls_ssl_free( &ssl );
+    mbedtls_ssl_config_free( &conf );
+
+#if defined(MBEDTLS_SSL_CACHE_C)
+    mbedtls_ssl_cache_free( &cache );
 #endif
+#if defined(MBEDTLS_SSL_SESSION_TICKETS)
+    mbedtls_ssl_ticket_free( &ticket_ctx );
+#endif
+#if defined(MBEDTLS_SSL_COOKIE_C)
+    mbedtls_ssl_cookie_free( &cookie_ctx );
+#endif
+
+#if defined(MBEDTLS_SSL_CONTEXT_SERIALIZATION)
+    if( context_buf != NULL )
+        mbedtls_platform_zeroize( context_buf, context_buf_len );
+    mbedtls_free( context_buf );
+#endif
+
+#if defined(SNI_OPTION)
+    sni_free( sni_info );
+#endif
+
+#if defined(MBEDTLS_KEY_EXCHANGE_SOME_PSK_ENABLED)
+    ret = psk_free( psk_info );
+    if( ( ret != 0 ) && ( opt.query_config_mode == DFL_QUERY_CONFIG_MODE ) )
+        mbedtls_printf( "Failed to list of opaque PSKs - error was %d\n", ret );
+#endif
+
 #if defined(MBEDTLS_X509_CRT_PARSE_C)
     mbedtls_x509_crt_free( &cacert );
     mbedtls_x509_crt_free( &srvcert );
@@ -3961,6 +3991,11 @@
     psa_destroy_key( key_slot2 );
 #endif
 #endif
+
+#if defined(MBEDTLS_DHM_C) && defined(MBEDTLS_FS_IO)
+    mbedtls_dhm_free( &dhm );
+#endif
+
 #if defined(MBEDTLS_SSL_ASYNC_PRIVATE)
     for( i = 0; (size_t) i < ssl_async_keys.slots_used; i++ )
     {
@@ -3972,17 +4007,6 @@
         }
     }
 #endif
-#if defined(SNI_OPTION)
-    sni_free( sni_info );
-#endif
-#if defined(MBEDTLS_KEY_EXCHANGE_SOME_PSK_ENABLED)
-    ret = psk_free( psk_info );
-    if( ( ret != 0 ) && ( opt.query_config_mode == DFL_QUERY_CONFIG_MODE ) )
-        mbedtls_printf( "Failed to list of opaque PSKs - error was %d\n", ret );
-#endif
-#if defined(MBEDTLS_DHM_C) && defined(MBEDTLS_FS_IO)
-    mbedtls_dhm_free( &dhm );
-#endif
 
 #if defined(MBEDTLS_KEY_EXCHANGE_SOME_PSK_ENABLED) && \
     defined(MBEDTLS_USE_PSA_CRYPTO)
@@ -4003,32 +4027,27 @@
 #endif /* MBEDTLS_KEY_EXCHANGE_SOME_PSK_ENABLED &&
           MBEDTLS_USE_PSA_CRYPTO */
 
-    mbedtls_ssl_free( &ssl );
-    mbedtls_ssl_config_free( &conf );
-    rng_free( &rng );
-
-#if defined(MBEDTLS_SSL_CACHE_C)
-    mbedtls_ssl_cache_free( &cache );
-#endif
-#if defined(MBEDTLS_SSL_SESSION_TICKETS)
-    mbedtls_ssl_ticket_free( &ticket_ctx );
-#endif
-#if defined(MBEDTLS_SSL_COOKIE_C)
-    mbedtls_ssl_cookie_free( &cookie_ctx );
-#endif
-
-    mbedtls_free( buf );
-
-#if defined(MBEDTLS_SSL_CONTEXT_SERIALIZATION)
-    if( context_buf != NULL )
-        mbedtls_platform_zeroize( context_buf, context_buf_len );
-    mbedtls_free( context_buf );
-#endif
-
 #if defined(MBEDTLS_USE_PSA_CRYPTO)
+    const char* message = mbedtls_test_helper_is_psa_leaking();
+    if( message )
+    {
+        if( ret == 0 )
+            ret = 1;
+        mbedtls_printf( "PSA memory leak detected: %s\n",  message);
+    }
+#endif
+
+    /* For builds with MBEDTLS_TEST_USE_PSA_CRYPTO_RNG psa crypto
+     * resources are freed by rng_free(). */
+#if defined(MBEDTLS_USE_PSA_CRYPTO) && \
+    !defined(MBEDTLS_TEST_USE_PSA_CRYPTO_RNG)
     mbedtls_psa_crypto_free( );
 #endif
 
+    rng_free( &rng );
+
+    mbedtls_free( buf );
+
 #if defined(MBEDTLS_TEST_HOOKS)
     /* Let test hooks detect errors such as resource leaks.
      * Don't do it in query_config mode, because some test code prints
diff --git a/tests/include/test/psa_crypto_helpers.h b/tests/include/test/psa_crypto_helpers.h
index 8a8c37e..f5622e2 100644
--- a/tests/include/test/psa_crypto_helpers.h
+++ b/tests/include/test/psa_crypto_helpers.h
@@ -28,7 +28,6 @@
 #include "test/psa_helpers.h"
 
 #include <psa/crypto.h>
-#include <psa_crypto_slot_management.h>
 
 #if defined(MBEDTLS_USE_PSA_CRYPTO)
 #include "mbedtls/psa_util.h"
diff --git a/tests/src/psa_crypto_helpers.c b/tests/src/psa_crypto_helpers.c
index d9d841a..299b6d1 100644
--- a/tests/src/psa_crypto_helpers.c
+++ b/tests/src/psa_crypto_helpers.c
@@ -22,6 +22,7 @@
 
 #include <test/helpers.h>
 #include <test/macros.h>
+#include <psa_crypto_slot_management.h>
 #include <test/psa_crypto_helpers.h>
 
 #if defined(MBEDTLS_PSA_CRYPTO_C)
diff --git a/tests/src/psa_exercise_key.c b/tests/src/psa_exercise_key.c
index 91bac67..fc58fbd 100644
--- a/tests/src/psa_exercise_key.c
+++ b/tests/src/psa_exercise_key.c
@@ -29,6 +29,7 @@
 #include <psa/crypto.h>
 
 #include <test/asn1_helpers.h>
+#include <psa_crypto_slot_management.h>
 #include <test/psa_crypto_helpers.h>
 
 #if defined(MBEDTLS_PSA_CRYPTO_SE_C)