diff --git a/library/psa_crypto_pake.c b/library/psa_crypto_pake.c
index 870b5b5..1deb488 100644
--- a/library/psa_crypto_pake.c
+++ b/library/psa_crypto_pake.c
@@ -33,6 +33,11 @@
 #include <mbedtls/error.h>
 #include <string.h>
 
+extern psa_status_t psa_get_and_lock_key_slot_with_policy(
+    mbedtls_svc_key_id_t key,
+    psa_key_slot_t **p_slot,
+    psa_key_usage_t usage,
+    psa_algorithm_t alg );
 /*
  * State sequence:
  *
@@ -248,6 +253,7 @@
     psa_key_attributes_t attributes = psa_key_attributes_init();
     psa_key_type_t type;
     psa_key_usage_t usage;
+    psa_key_slot_t *slot = NULL;
 
     if( operation->alg == PSA_ALG_NONE ||
         operation->state != PSA_PAKE_STATE_SETUP )
@@ -255,6 +261,9 @@
         return( PSA_ERROR_BAD_STATE );
     }
 
+    if( psa_is_valid_key_id( password, 1 ) == 0 )
+        return( PSA_ERROR_BAD_STATE );
+
     status = psa_get_key_attributes( password, &attributes );
     if( status != PSA_SUCCESS )
         return( status );
@@ -273,7 +282,33 @@
     if( ( usage & PSA_KEY_USAGE_DERIVE ) == 0 )
         return( PSA_ERROR_NOT_PERMITTED );
 
-    operation->password = password;
+    status = psa_get_and_lock_key_slot_with_policy( password, &slot,
+                                                    PSA_KEY_USAGE_DERIVE,
+                                                    PSA_ALG_JPAKE );
+    if( status != PSA_SUCCESS )
+        return( status );
+
+    if( slot->key.data == NULL || slot->key.bytes == 0 )
+        return( PSA_ERROR_INVALID_ARGUMENT );
+
+    if( operation->password_data != NULL )
+    {
+        mbedtls_free( operation->password_data );
+        operation->password_bytes = 0;
+    }
+
+    operation->password_data = mbedtls_calloc( 1, slot->key.bytes );
+    if( operation->password_data == NULL )
+    {
+        status = psa_unlock_key_slot( slot );
+        return( PSA_ERROR_INSUFFICIENT_MEMORY );
+    }
+    memcpy( operation->password_data, slot->key.data, slot->key.bytes );
+    operation->password_bytes = slot->key.bytes;
+
+    status = psa_unlock_key_slot( slot );
+    if( status != PSA_SUCCESS )
+        return( status );
 
     return( PSA_SUCCESS );
 }
@@ -348,9 +383,7 @@
 static psa_status_t psa_pake_ecjpake_setup( psa_pake_operation_t *operation )
 {
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
-    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
     mbedtls_ecjpake_role role;
-    psa_key_slot_t *slot = NULL;
 
     if( operation->role == PSA_PAKE_ROLE_CLIENT )
         role = MBEDTLS_ECJPAKE_CLIENT;
@@ -359,22 +392,18 @@
     else
         return( PSA_ERROR_BAD_STATE );
 
-    if( psa_is_valid_key_id( operation->password, 1 ) == 0 )
+    if (operation->password_data == NULL ||
+        operation->password_bytes == 0 )
+    {
         return( PSA_ERROR_BAD_STATE );
-
-    status = psa_get_and_lock_key_slot( operation->password, &slot );
-    if( status != PSA_SUCCESS )
-        return( status );
-
+    }
 
     ret = mbedtls_ecjpake_setup( &operation->ctx.ecjpake,
                                  role,
                                  MBEDTLS_MD_SHA256,
                                  MBEDTLS_ECP_DP_SECP256R1,
-                                 slot->key.data, slot->key.bytes );
-
-    psa_unlock_key_slot( slot );
-    slot = NULL;
+                                 operation->password_data,
+                                 operation->password_bytes );
 
     if( ret != 0 )
         return( mbedtls_ecjpake_to_psa_error( ret ) );
@@ -840,7 +869,9 @@
     {
         operation->input_step = PSA_PAKE_STEP_INVALID;
         operation->output_step = PSA_PAKE_STEP_INVALID;
-        operation->password = MBEDTLS_SVC_KEY_ID_INIT;
+        mbedtls_free( operation->password_data );
+        operation->password_data = NULL;
+        operation->password_bytes = 0;
         operation->role = PSA_PAKE_ROLE_NONE;
         mbedtls_platform_zeroize( operation->buffer, MBEDTLS_PSA_PAKE_BUFFER_SIZE );
         operation->buffer_length = 0;