diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 9c12863f0f..0fd0eff888 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -7483,6 +7483,12 @@ static psa_status_t psa_pake_complete_inputs(
         return PSA_ERROR_BAD_STATE;
     }
 
+    if (operation->alg == PSA_ALG_JPAKE &&
+        inputs.role != PSA_PAKE_ROLE_CLIENT &&
+        inputs.role != PSA_PAKE_ROLE_SERVER) {
+        return PSA_ERROR_NOT_SUPPORTED;
+    }
+
     /* Clear driver context */
     mbedtls_platform_zeroize(&operation->data, sizeof(operation->data));
 
diff --git a/library/psa_crypto_pake.c b/library/psa_crypto_pake.c
index 2d84f570b9..382f0214a4 100644
--- a/library/psa_crypto_pake.c
+++ b/library/psa_crypto_pake.c
@@ -221,13 +221,6 @@ psa_status_t mbedtls_psa_pake_setup(mbedtls_psa_pake_operation_t *operation,
             return PSA_ERROR_NOT_SUPPORTED;
         }
 
-        if (role != PSA_PAKE_ROLE_CLIENT &&
-            role != PSA_PAKE_ROLE_SERVER) {
-            return PSA_ERROR_NOT_SUPPORTED;
-        }
-
-
-
         operation->password = mbedtls_calloc(1, password_len);
         if (operation->password == NULL) {
             status = PSA_ERROR_INSUFFICIENT_MEMORY;