diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 1c066ce13b..e2cb06f427 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -7494,12 +7494,13 @@ static psa_status_t psa_pake_complete_inputs(
 
     status = psa_driver_wrapper_pake_setup(operation, &inputs);
 
+    operation->stage = PSA_PAKE_OPERATION_STAGE_COMPUTATION;
+
     /* Driver is responsible for creating its own copy of the password. */
     mbedtls_platform_zeroize(inputs.password, inputs.password_len);
     mbedtls_free(inputs.password);
 
     if (status == PSA_SUCCESS) {
-        operation->stage = PSA_PAKE_OPERATION_STAGE_COMPUTATION;
         if (operation->alg == PSA_ALG_JPAKE) {
             psa_jpake_computation_stage_t *computation_stage =
                 &operation->computation_stage.jpake;
diff --git a/library/psa_crypto_pake.c b/library/psa_crypto_pake.c
index 150270c6c3..a6798111de 100644
--- a/library/psa_crypto_pake.c
+++ b/library/psa_crypto_pake.c
@@ -214,38 +214,38 @@ psa_status_t mbedtls_psa_pake_setup(mbedtls_psa_pake_operation_t *operation,
         return status;
     }
 
+    operation->password = mbedtls_calloc(1, password_len);
+    if (operation->password == NULL) {
+        return PSA_ERROR_INSUFFICIENT_MEMORY;
+    }
+
+    status = psa_crypto_driver_pake_get_password(inputs, operation->password,
+                                                password_len, &actual_password_len);
+    if (status != PSA_SUCCESS) {
+        goto error;
+    }
+
+    operation->password_len = actual_password_len;
+    operation->alg = cipher_suite.algorithm;
+
 #if defined(MBEDTLS_PSA_BUILTIN_ALG_JPAKE)
     if (cipher_suite.algorithm == PSA_ALG_JPAKE) {
         if (cipher_suite.type != PSA_PAKE_PRIMITIVE_TYPE_ECC ||
             cipher_suite.family != PSA_ECC_FAMILY_SECP_R1 ||
             cipher_suite.bits != 256 ||
             cipher_suite.hash != PSA_ALG_SHA_256) {
-            return PSA_ERROR_NOT_SUPPORTED;
+            status = PSA_ERROR_NOT_SUPPORTED;
+            goto error;
         }
 
-        operation->password = mbedtls_calloc(1, password_len);
-        if (operation->password == NULL) {
-            status = PSA_ERROR_INSUFFICIENT_MEMORY;
-            return status;
-        }
-
-        status = psa_crypto_driver_pake_get_password(inputs, operation->password,
-                                                     password_len, &actual_password_len);
-        if (status != PSA_SUCCESS) {
-            return status;
-        }
-
-        operation->password_len = actual_password_len;
         operation->role = role;
-        operation->alg = cipher_suite.algorithm;
 
-        mbedtls_platform_zeroize(operation->buffer, MBEDTLS_PSA_JPAKE_BUFFER_SIZE);
         operation->buffer_length = 0;
         operation->buffer_offset = 0;
 
         status = psa_pake_ecjpake_setup(operation);
         if (status != PSA_SUCCESS) {
-            return status;
+            goto error;
         }
 
         return PSA_SUCCESS;
@@ -254,8 +254,11 @@ psa_status_t mbedtls_psa_pake_setup(mbedtls_psa_pake_operation_t *operation,
     (void) operation;
     (void) inputs;
 #endif
-    { status = PSA_ERROR_NOT_SUPPORTED; }
+    { return PSA_ERROR_NOT_SUPPORTED; }
 
+error:
+    mbedtls_platform_zeroize(operation->password, operation->password_len);
+    mbedtls_free(operation->password);
     return status;
 }