tls: pake: do not destroy password key in TLS

Signed-off-by: Valerio Setti <vsetti@baylibre.com>
diff --git a/library/ssl_tls.c b/library/ssl_tls.c
index 57cfe42..259d088 100644
--- a/library/ssl_tls.c
+++ b/library/ssl_tls.c
@@ -1934,7 +1934,6 @@
 
     if( mbedtls_svc_key_id_is_null( pwd ) )
         return( MBEDTLS_ERR_SSL_BAD_INPUT_DATA );
-    ssl->handshake->psa_pake_password = pwd;
 
     psa_pake_cs_set_algorithm( &cipher_suite, PSA_ALG_JPAKE );
     psa_pake_cs_set_primitive( &cipher_suite,
@@ -1956,8 +1955,7 @@
     if( status != PSA_SUCCESS )
         goto error;
 
-    status = psa_pake_set_password_key( &ssl->handshake->psa_pake_ctx,
-                                ssl->handshake->psa_pake_password );
+    status = psa_pake_set_password_key( &ssl->handshake->psa_pake_ctx, pwd );
     if( status != PSA_SUCCESS )
         goto error;
 
@@ -4037,7 +4035,15 @@
 #if defined(MBEDTLS_KEY_EXCHANGE_ECJPAKE_ENABLED)
 #if defined(MBEDTLS_USE_PSA_CRYPTO)
     psa_pake_abort( &handshake->psa_pake_ctx );
-    psa_destroy_key( handshake->psa_pake_password );
+    /*
+     * Opaque keys are not stored in the handshake's data and it's the user
+     * responsibility to destroy them. Clear ones, instead, are created by
+     * the TLS library and should be destroyed at the same level
+     */
+    if( ! mbedtls_svc_key_id_is_null( handshake->psa_pake_password ) )
+    {
+        psa_destroy_key( handshake->psa_pake_password );
+    }
     handshake->psa_pake_password = MBEDTLS_SVC_KEY_ID_INIT;
 #else
     mbedtls_ecjpake_free( &handshake->ecjpake_ctx );
diff --git a/programs/ssl/ssl_client2.c b/programs/ssl/ssl_client2.c
index 54e9861..13ffa3d 100644
--- a/programs/ssl/ssl_client2.c
+++ b/programs/ssl/ssl_client2.c
@@ -3329,18 +3329,14 @@
 
 #if defined(MBEDTLS_KEY_EXCHANGE_ECJPAKE_ENABLED) && \
     defined(MBEDTLS_USE_PSA_CRYPTO)
-    if( opt.ecjpake_pw_opaque != DFL_ECJPAKE_PW_OPAQUE )
+    /*
+     * In case opaque keys it's the user responsibility to keep the key valid
+     * for the duration of the handshake and destroy it at the end
+     */
+    if( ( opt.ecjpake_pw_opaque != DFL_ECJPAKE_PW_OPAQUE ) &&
+        ( ! mbedtls_svc_key_id_is_null( ecjpake_pw_slot ) ) )
     {
-        psa_key_attributes_t key_attr = PSA_KEY_ATTRIBUTES_INIT;
-
-        /* Ensure the key is still valid before destroying it */
-        status = psa_get_key_attributes( ecjpake_pw_slot, &key_attr );
-        if( status == PSA_SUCCESS &&
-            PSA_ALG_IS_PAKE( psa_get_key_algorithm( &key_attr ) ) )
-        {
-            psa_destroy_key( ecjpake_pw_slot );
-        }
-        psa_reset_key_attributes( &key_attr );
+        psa_destroy_key( ecjpake_pw_slot );
     }
 #endif  /* MBEDTLS_KEY_EXCHANGE_ECJPAKE_ENABLED && MBEDTLS_USE_PSA_CRYPTO */
 
diff --git a/programs/ssl/ssl_server2.c b/programs/ssl/ssl_server2.c
index 004616f..4b195c4 100644
--- a/programs/ssl/ssl_server2.c
+++ b/programs/ssl/ssl_server2.c
@@ -4439,18 +4439,14 @@
 
 #if defined(MBEDTLS_KEY_EXCHANGE_ECJPAKE_ENABLED) && \
     defined(MBEDTLS_USE_PSA_CRYPTO)
-    if( opt.ecjpake_pw_opaque != DFL_ECJPAKE_PW_OPAQUE )
+    /*
+     * In case opaque keys it's the user responsibility to keep the key valid
+     * for the duration of the handshake and destroy it at the end
+     */
+    if( ( opt.ecjpake_pw_opaque != DFL_ECJPAKE_PW_OPAQUE ) &&
+        ( ! mbedtls_svc_key_id_is_null( ecjpake_pw_slot ) ) )
     {
-        psa_key_attributes_t key_attr = PSA_KEY_ATTRIBUTES_INIT;
-
-        /* Ensure the key is still valid before destroying it */
-        status = psa_get_key_attributes( ecjpake_pw_slot, &key_attr );
-        if( status == PSA_SUCCESS &&
-            PSA_ALG_IS_PAKE( psa_get_key_algorithm( &key_attr ) ) )
-        {
-            psa_destroy_key( ecjpake_pw_slot );
-        }
-        psa_reset_key_attributes( &key_attr );
+        psa_destroy_key( ecjpake_pw_slot );
     }
 #endif  /* MBEDTLS_KEY_EXCHANGE_ECJPAKE_ENABLED && MBEDTLS_USE_PSA_CRYPTO */