diff --git a/library/psa_crypto_pake.c b/library/psa_crypto_pake.c
index 9ac4c5f..659b712 100644
--- a/library/psa_crypto_pake.c
+++ b/library/psa_crypto_pake.c
@@ -274,19 +274,19 @@
     if( ( usage & PSA_KEY_USAGE_DERIVE ) == 0 )
         return( PSA_ERROR_NOT_PERMITTED );
 
+    if( operation->password != NULL )
+        return( PSA_ERROR_BAD_STATE );
+
     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( operation->password != NULL )
-        return( PSA_ERROR_BAD_STATE );
-
     operation->password = mbedtls_calloc( 1, slot->key.bytes );
     if( operation->password == NULL )
     {
-        status = psa_unlock_key_slot( slot );
+        psa_unlock_key_slot( slot );
         return( PSA_ERROR_INSUFFICIENT_MEMORY );
     }
     memcpy( operation->password, slot->key.data, slot->key.bytes );
@@ -388,14 +388,14 @@
                                  operation->password,
                                  operation->password_len );
 
-    if( ret != 0 )
-        return( mbedtls_ecjpake_to_psa_error( ret ) );
-
     mbedtls_platform_zeroize( operation->password, operation->password_len );
     mbedtls_free( operation->password );
     operation->password = NULL;
     operation->password_len = 0;
 
+    if( ret != 0 )
+        return( mbedtls_ecjpake_to_psa_error( ret ) );
+
     operation->state = PSA_PAKE_STATE_READY;
 
     return( PSA_SUCCESS );
@@ -445,13 +445,7 @@
         if( operation->state == PSA_PAKE_STATE_SETUP ) {
             status = psa_pake_ecjpake_setup( operation );
             if( status != PSA_SUCCESS )
-            {
-                mbedtls_platform_zeroize( operation->password, operation->password_len );
-                mbedtls_free( operation->password );
-                operation->password = NULL;
-                operation->password_len = 0;
                 return( status );
-            }
         }
 
         if( operation->state != PSA_PAKE_STATE_READY &&
@@ -659,13 +653,7 @@
         {
             status = psa_pake_ecjpake_setup( operation );
             if( status != PSA_SUCCESS )
-            {
-                mbedtls_platform_zeroize( operation->password, operation->password_len );
-                mbedtls_free( operation->password );
-                operation->password = NULL;
-                operation->password_len = 0;
                 return( status );
-            }
         }
 
         if( operation->state != PSA_PAKE_STATE_READY &&