diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index bc19ed0..049edfc 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -7323,6 +7323,8 @@
     memset(&operation->data.inputs, 0, sizeof(operation->data.inputs));
 
     operation->alg = cipher_suite->algorithm;
+    operation->primitive = PSA_PAKE_PRIMITIVE(cipher_suite->type,
+                                              cipher_suite->family, cipher_suite->bits);
     operation->data.inputs.cipher_suite = *cipher_suite;
 
 #if defined(PSA_WANT_ALG_JPAKE)
@@ -7920,7 +7922,9 @@
         goto exit;
     }
 
-    if (input_length == 0 || input_length > PSA_PAKE_INPUT_MAX_SIZE) {
+    if (input_length == 0 || input_length > PSA_PAKE_INPUT_SIZE(operation->alg,
+                                                                operation->primitive,
+                                                                step)) {
         status = PSA_ERROR_INVALID_ARGUMENT;
         goto exit;
     }