diff --git a/include/mbedtls/ssl.h b/include/mbedtls/ssl.h
index 97caad881..6d453ad1c 100644
--- a/include/mbedtls/ssl.h
+++ b/include/mbedtls/ssl.h
@@ -3864,8 +3864,7 @@ int mbedtls_ssl_set_hs_ecjpake_password( mbedtls_ssl_context *ssl,
 /**
  * \brief          Set the EC J-PAKE opaque password for current handshake.
  *
- * \note           The input key in not copied, so the caller must not destroy
- *                 it before the handshake is over.
+ * \note           The key must remain valid until the handshake is over.
  *
  * \note           The SSL context needs to be already set up. The right place
  *                 to call this function is between \c mbedtls_ssl_setup() or
diff --git a/library/ssl_tls.c b/library/ssl_tls.c
index 1438124dd..810db243d 100644
--- a/library/ssl_tls.c
+++ b/library/ssl_tls.c
@@ -1932,11 +1932,6 @@ int mbedtls_ssl_set_hs_ecjpake_password_opaque( mbedtls_ssl_context *ssl,
     if( ssl->handshake == NULL || ssl->conf == NULL )
         return( MBEDTLS_ERR_SSL_BAD_INPUT_DATA );
 
-    if( ssl->conf->endpoint == MBEDTLS_SSL_IS_SERVER )
-        psa_role = PSA_PAKE_ROLE_SERVER;
-    else
-        psa_role = PSA_PAKE_ROLE_CLIENT;
-
     if( mbedtls_svc_key_id_is_null( pwd ) )
         return( MBEDTLS_ERR_SSL_BAD_INPUT_DATA );
     ssl->handshake->psa_pake_password = pwd;
@@ -1952,6 +1947,11 @@ int mbedtls_ssl_set_hs_ecjpake_password_opaque( mbedtls_ssl_context *ssl,
     if( status != PSA_SUCCESS )
         goto error;
 
+    if( ssl->conf->endpoint == MBEDTLS_SSL_IS_SERVER )
+        psa_role = PSA_PAKE_ROLE_SERVER;
+    else
+        psa_role = PSA_PAKE_ROLE_CLIENT;
+
     status = psa_pake_set_role( &ssl->handshake->psa_pake_ctx, psa_role );
     if( status != PSA_SUCCESS )
         goto error;
@@ -1979,15 +1979,15 @@ int mbedtls_ssl_set_hs_ecjpake_password( mbedtls_ssl_context *ssl,
     if( ssl->handshake == NULL || ssl->conf == NULL )
         return( MBEDTLS_ERR_SSL_BAD_INPUT_DATA );
 
+    /* Empty password is not valid  */
+    if( ( pw == NULL) || ( pw_len == 0 ) )
+        return( MBEDTLS_ERR_SSL_BAD_INPUT_DATA );
+
     if( ssl->conf->endpoint == MBEDTLS_SSL_IS_SERVER )
         role = MBEDTLS_ECJPAKE_SERVER;
     else
         role = MBEDTLS_ECJPAKE_CLIENT;
 
-    /* Empty password is not valid  */
-    if( ( pw == NULL) || ( pw_len == 0 ) )
-        return( MBEDTLS_ERR_SSL_BAD_INPUT_DATA );
-
     return( mbedtls_ecjpake_setup( &ssl->handshake->ecjpake_ctx,
                                    role,
                                    MBEDTLS_MD_SHA256,