Revert "remove psk key when ephemeral selected"

This reverts commit 5c28e7aa0ebfa3d025e7802ae29eb86e43b0e92f.

Signed-off-by: Jerry Yu <jerry.h.yu@arm.com>
diff --git a/library/ssl_misc.h b/library/ssl_misc.h
index 33a6533..e76086a 100644
--- a/library/ssl_misc.h
+++ b/library/ssl_misc.h
@@ -2498,13 +2498,6 @@
 int mbedtls_ssl_tls13_write_binders_of_pre_shared_key_ext(
     mbedtls_ssl_context *ssl,
     unsigned char *buf, unsigned char *end );
-
-/**
- * \brief Remove psk from handshake context
- *
- * \param[in]   ssl     SSL context
- */
-void mbedtls_ssl_remove_psk( mbedtls_ssl_context *ssl );
 #endif /* MBEDTLS_KEY_EXCHANGE_SOME_PSK_ENABLED */
 
 #endif /* ssl_misc.h */
diff --git a/library/ssl_tls.c b/library/ssl_tls.c
index 639b18c..8aadeb1 100644
--- a/library/ssl_tls.c
+++ b/library/ssl_tls.c
@@ -1662,7 +1662,7 @@
     return( ret );
 }
 
-void mbedtls_ssl_remove_psk( mbedtls_ssl_context *ssl )
+static void ssl_remove_psk( mbedtls_ssl_context *ssl )
 {
 #if defined(MBEDTLS_USE_PSA_CRYPTO)
     if( ! mbedtls_svc_key_id_is_null( ssl->handshake->psk_opaque ) )
@@ -1682,7 +1682,6 @@
         mbedtls_platform_zeroize( ssl->handshake->psk,
                                   ssl->handshake->psk_len );
         mbedtls_free( ssl->handshake->psk );
-        ssl->handshake->psk = NULL;
         ssl->handshake->psk_len = 0;
     }
 #endif /* MBEDTLS_USE_PSA_CRYPTO */
@@ -1704,7 +1703,7 @@
     if( psk_len > MBEDTLS_PSK_MAX_LEN )
         return( MBEDTLS_ERR_SSL_BAD_INPUT_DATA );
 
-    mbedtls_ssl_remove_psk( ssl );
+    ssl_remove_psk( ssl );
 
 #if defined(MBEDTLS_USE_PSA_CRYPTO)
 #if defined(MBEDTLS_SSL_PROTO_TLS1_2)
@@ -1781,7 +1780,7 @@
         ( ssl->handshake == NULL ) )
         return( MBEDTLS_ERR_SSL_BAD_INPUT_DATA );
 
-    mbedtls_ssl_remove_psk( ssl );
+    ssl_remove_psk( ssl );
     ssl->handshake->psk_opaque = psk;
     return( 0 );
 }
@@ -3523,7 +3522,25 @@
 #endif
 
 #if defined(MBEDTLS_KEY_EXCHANGE_SOME_PSK_ENABLED)
-    mbedtls_ssl_remove_psk( ssl );
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    if( ! mbedtls_svc_key_id_is_null( ssl->handshake->psk_opaque ) )
+    {
+        /* The maintenance of the external PSK key slot is the
+         * user's responsibility. */
+        if( ssl->handshake->psk_opaque_is_internal )
+        {
+            psa_destroy_key( ssl->handshake->psk_opaque );
+            ssl->handshake->psk_opaque_is_internal = 0;
+        }
+        ssl->handshake->psk_opaque = MBEDTLS_SVC_KEY_ID_INIT;
+    }
+#else
+    if( handshake->psk != NULL )
+    {
+        mbedtls_platform_zeroize( handshake->psk, handshake->psk_len );
+        mbedtls_free( handshake->psk );
+    }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 #endif
 
 #if defined(MBEDTLS_X509_CRT_PARSE_C) && \
diff --git a/library/ssl_tls13_server.c b/library/ssl_tls13_server.c
index e9a4f99..b37fe5a 100644
--- a/library/ssl_tls13_server.c
+++ b/library/ssl_tls13_server.c
@@ -760,9 +760,6 @@
     else
     if( ssl_tls13_check_ephemeral_key_exchange( ssl ) )
     {
-#if defined(MBEDTLS_KEY_EXCHANGE_SOME_PSK_ENABLED)
-        mbedtls_ssl_remove_psk( ssl );
-#endif
         ssl->handshake->key_exchange_mode =
             MBEDTLS_SSL_TLS1_3_KEY_EXCHANGE_MODE_EPHEMERAL;
         MBEDTLS_SSL_DEBUG_MSG( 2, ( "key exchange mode: ephemeral" ) );