Move erase handshake secrets

Signed-off-by: Jerry Yu <jerry.h.yu@arm.com>
diff --git a/library/ssl_tls13_client.c b/library/ssl_tls13_client.c
index 88f240e..1874d4f 100644
--- a/library/ssl_tls13_client.c
+++ b/library/ssl_tls13_client.c
@@ -1646,7 +1646,6 @@
  */
 static int ssl_tls13_handshake_wrapup( mbedtls_ssl_context *ssl )
 {
-
     MBEDTLS_SSL_DEBUG_MSG( 1, ( "Switch to application keys for inbound traffic" ) );
     mbedtls_ssl_set_inbound_transform ( ssl, ssl->transform_application );
 
diff --git a/library/ssl_tls13_generic.c b/library/ssl_tls13_generic.c
index d15296f..4543d74 100644
--- a/library/ssl_tls13_generic.c
+++ b/library/ssl_tls13_generic.c
@@ -1060,11 +1060,7 @@
                     sizeof( ssl->handshake->state_local.finished_out.digest ),
                     &ssl->handshake->state_local.finished_out.digest_len,
                     ssl->conf->endpoint );
-    if( ssl->conf->endpoint == MBEDTLS_SSL_IS_CLIENT )
-    {
-        mbedtls_platform_zeroize( &ssl->handshake->tls13_hs_secrets,
-                                  sizeof( ssl->handshake->tls13_hs_secrets ) );
-    }
+
     if( ret != 0 )
     {
         MBEDTLS_SSL_DEBUG_RET( 1, "calculate_verify_data failed", ret );
diff --git a/library/ssl_tls13_keys.c b/library/ssl_tls13_keys.c
index 3306d62..ea0c55b 100644
--- a/library/ssl_tls13_keys.c
+++ b/library/ssl_tls13_keys.c
@@ -654,7 +654,8 @@
     unsigned char transcript[MBEDTLS_TLS1_3_MD_MAX_SIZE];
     size_t transcript_len;
 
-    unsigned char const *base_key = NULL;
+    unsigned char *base_key = NULL;
+    size_t base_key_len;
 
     mbedtls_md_type_t const md_type = ssl->handshake->ciphersuite_info->mac;
     const mbedtls_md_info_t* const md_info =
@@ -677,9 +678,15 @@
     MBEDTLS_SSL_DEBUG_BUF( 4, "handshake hash", transcript, transcript_len );
 
     if( from == MBEDTLS_SSL_IS_CLIENT )
+    {
         base_key = ssl->handshake->tls13_hs_secrets.client_handshake_traffic_secret;
+        base_key_len = sizeof( ssl->handshake->tls13_hs_secrets.client_handshake_traffic_secret );
+    }
     else
+    {
         base_key = ssl->handshake->tls13_hs_secrets.server_handshake_traffic_secret;
+        base_key_len = sizeof( ssl->handshake->tls13_hs_secrets.server_handshake_traffic_secret );
+    }
 
     ret = ssl_tls13_calc_finished_core( md_type, base_key, transcript, dst );
     if( ret != 0 )
@@ -690,7 +697,8 @@
     MBEDTLS_SSL_DEBUG_MSG( 2, ( "<= mbedtls_ssl_tls13_calculate_verify_data" ) );
 
 exit:
-
+    /* Erase handshake secrets */
+    mbedtls_platform_zeroize( base_key, base_key_len );
     mbedtls_platform_zeroize( transcript, sizeof( transcript ) );
     return( ret );
 }
@@ -1164,7 +1172,7 @@
                                    handshake->tls13_master_secrets.app,
                                    transcript, transcript_len,
                                    app_secrets );
-    /* Erase master secrets*/
+    /* Erase master secrets */
     mbedtls_platform_zeroize( &ssl->handshake->tls13_master_secrets,
                               sizeof( ssl->handshake->tls13_master_secrets ) );
     if( ret != 0 )