diff --git a/ChangeLog.d/ec_jpake_user_peer_2.txt b/ChangeLog.d/ec_jpake_user_peer_2.txt
new file mode 100644
index 0000000..9572ac7
--- /dev/null
+++ b/ChangeLog.d/ec_jpake_user_peer_2.txt
@@ -0,0 +1,3 @@
+Bugfix
+   * Fix the J-PAKE driver interface for user and peer to accept any values
+     (previously accepted values were limited to "client" or "server").
diff --git a/docs/proposed/psa-driver-interface.md b/docs/proposed/psa-driver-interface.md
index cd1b9fc..a3e3c76 100644
--- a/docs/proposed/psa-driver-interface.md
+++ b/docs/proposed/psa-driver-interface.md
@@ -390,10 +390,6 @@
     const psa_crypto_driver_pake_inputs_t *inputs,
     uint8_t *peer_id, size_t peer_id_size, size_t *peer_id_length);
 
-psa_status_t psa_crypto_driver_pake_get_role(
-    const psa_crypto_driver_pake_inputs_t *inputs,
-    psa_pake_role_t *role);
-
 psa_status_t psa_crypto_driver_pake_get_cipher_suite(
     const psa_crypto_driver_pake_inputs_t *inputs,
     psa_pake_cipher_suite_t *cipher_suite);
diff --git a/include/mbedtls/ecjpake.h b/include/mbedtls/ecjpake.h
index a63bb32..0008d73 100644
--- a/include/mbedtls/ecjpake.h
+++ b/include/mbedtls/ecjpake.h
@@ -54,6 +54,7 @@
 typedef enum {
     MBEDTLS_ECJPAKE_CLIENT = 0,         /**< Client                         */
     MBEDTLS_ECJPAKE_SERVER,             /**< Server                         */
+    MBEDTLS_ECJPAKE_NONE,               /**< Undefined                      */
 } mbedtls_ecjpake_role;
 
 #if !defined(MBEDTLS_ECJPAKE_ALT)
diff --git a/include/psa/crypto_builtin_composites.h b/include/psa/crypto_builtin_composites.h
index c280360..d9473ac 100644
--- a/include/psa/crypto_builtin_composites.h
+++ b/include/psa/crypto_builtin_composites.h
@@ -202,7 +202,7 @@
     uint8_t *MBEDTLS_PRIVATE(password);
     size_t MBEDTLS_PRIVATE(password_len);
 #if defined(MBEDTLS_PSA_BUILTIN_ALG_JPAKE)
-    uint8_t MBEDTLS_PRIVATE(role);
+    mbedtls_ecjpake_role MBEDTLS_PRIVATE(role);
     uint8_t MBEDTLS_PRIVATE(buffer[MBEDTLS_PSA_JPAKE_BUFFER_SIZE]);
     size_t MBEDTLS_PRIVATE(buffer_length);
     size_t MBEDTLS_PRIVATE(buffer_offset);
diff --git a/include/psa/crypto_extra.h b/include/psa/crypto_extra.h
index b858180..232a839 100644
--- a/include/psa/crypto_extra.h
+++ b/include/psa/crypto_extra.h
@@ -1328,20 +1328,6 @@
     const psa_crypto_driver_pake_inputs_t *inputs,
     uint8_t *buffer, size_t buffer_size, size_t *buffer_length);
 
-/** Get the role from given inputs.
- *
- * \param[in]  inputs           Operation inputs.
- * \param[out] role             Return buffer for role.
- *
- * \retval #PSA_SUCCESS
- *         Success.
- * \retval #PSA_ERROR_BAD_STATE
- *         Role hasn't been set yet.
- */
-psa_status_t psa_crypto_driver_pake_get_role(
-    const psa_crypto_driver_pake_inputs_t *inputs,
-    psa_pake_role_t *role);
-
 /** Get the length of the user id in bytes from given inputs.
  *
  * \param[in]  inputs           Operation inputs.
@@ -1560,7 +1546,6 @@
  *                              been set (psa_pake_set_user() hasn't been
  *                              called yet).
  * \param[in] user_id           The user ID to authenticate with.
- *                              (temporary limitation: "client" or "server" only)
  * \param user_id_len           Size of the \p user_id buffer in bytes.
  *
  * \retval #PSA_SUCCESS
@@ -1602,7 +1587,6 @@
  *                              been set (psa_pake_set_peer() hasn't been
  *                              called yet).
  * \param[in] peer_id           The peer's ID to authenticate.
- *                              (temporary limitation: "client" or "server" only)
  * \param peer_id_len           Size of the \p peer_id buffer in bytes.
  *
  * \retval #PSA_SUCCESS
@@ -2039,7 +2023,6 @@
 struct psa_crypto_driver_pake_inputs_s {
     uint8_t *MBEDTLS_PRIVATE(password);
     size_t MBEDTLS_PRIVATE(password_len);
-    psa_pake_role_t MBEDTLS_PRIVATE(role);
     uint8_t *MBEDTLS_PRIVATE(user);
     size_t MBEDTLS_PRIVATE(user_len);
     uint8_t *MBEDTLS_PRIVATE(peer);
diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 20918bc..f7e91d6 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -91,10 +91,6 @@
 #define BUILTIN_ALG_ANY_HKDF 1
 #endif
 
-/* The only two JPAKE user/peer identifiers supported for the time being. */
-static const uint8_t jpake_server_id[] = { 's', 'e', 'r', 'v', 'e', 'r' };
-static const uint8_t jpake_client_id[] = { 'c', 'l', 'i', 'e', 'n', 't' };
-
 /****************************************************************/
 /* Global data, support functions and library management */
 /****************************************************************/
@@ -6481,6 +6477,27 @@
     return status;
 }
 
+static psa_status_t psa_key_derivation_input_integer_internal(
+    psa_key_derivation_operation_t *operation,
+    psa_key_derivation_step_t step,
+    uint64_t value)
+{
+    psa_status_t status;
+    psa_algorithm_t kdf_alg = psa_key_derivation_get_kdf_alg(operation);
+
+    {
+        (void) step;
+        (void) value;
+        (void) kdf_alg;
+        status = PSA_ERROR_INVALID_ARGUMENT;
+    }
+
+    if (status != PSA_SUCCESS) {
+        psa_key_derivation_abort(operation);
+    }
+    return status;
+}
+
 psa_status_t psa_key_derivation_input_bytes(
     psa_key_derivation_operation_t *operation,
     psa_key_derivation_step_t step,
@@ -6492,6 +6509,14 @@
                                              data, data_length);
 }
 
+psa_status_t psa_key_derivation_input_integer(
+    psa_key_derivation_operation_t *operation,
+    psa_key_derivation_step_t step,
+    uint64_t value)
+{
+    return psa_key_derivation_input_integer_internal(operation, step, value);
+}
+
 psa_status_t psa_key_derivation_input_key(
     psa_key_derivation_operation_t *operation,
     psa_key_derivation_step_t step,
@@ -7208,19 +7233,6 @@
     return PSA_SUCCESS;
 }
 
-psa_status_t psa_crypto_driver_pake_get_role(
-    const psa_crypto_driver_pake_inputs_t *inputs,
-    psa_pake_role_t *role)
-{
-    if (inputs->role == PSA_PAKE_ROLE_NONE) {
-        return PSA_ERROR_BAD_STATE;
-    }
-
-    *role = inputs->role;
-
-    return PSA_SUCCESS;
-}
-
 psa_status_t psa_crypto_driver_pake_get_user_len(
     const psa_crypto_driver_pake_inputs_t *inputs,
     size_t *user_len)
@@ -7415,15 +7427,6 @@
         goto exit;
     }
 
-    /* Allow only "client" or "server" values (temporary restriction). */
-    if ((user_id_len != sizeof(jpake_server_id) ||
-         memcmp(user_id, jpake_server_id, user_id_len) != 0) &&
-        (user_id_len != sizeof(jpake_client_id) ||
-         memcmp(user_id, jpake_client_id, user_id_len) != 0)) {
-        status = PSA_ERROR_NOT_SUPPORTED;
-        goto exit;
-    }
-
     operation->data.inputs.user = mbedtls_calloc(1, user_id_len);
     if (operation->data.inputs.user == NULL) {
         status = PSA_ERROR_INSUFFICIENT_MEMORY;
@@ -7461,15 +7464,6 @@
         goto exit;
     }
 
-    /* Allow only "client" or "server" values (temporary restriction). */
-    if ((peer_id_len != sizeof(jpake_server_id) ||
-         memcmp(peer_id, jpake_server_id, peer_id_len) != 0) &&
-        (peer_id_len != sizeof(jpake_client_id) ||
-         memcmp(peer_id, jpake_client_id, peer_id_len) != 0)) {
-        status = PSA_ERROR_NOT_SUPPORTED;
-        goto exit;
-    }
-
     operation->data.inputs.peer = mbedtls_calloc(1, peer_id_len);
     if (operation->data.inputs.peer == NULL) {
         status = PSA_ERROR_INSUFFICIENT_MEMORY;
@@ -7587,19 +7581,6 @@
         if (inputs.user_len == 0 || inputs.peer_len == 0) {
             return PSA_ERROR_BAD_STATE;
         }
-        if (memcmp(inputs.user, jpake_client_id, inputs.user_len) == 0 &&
-            memcmp(inputs.peer, jpake_server_id, inputs.peer_len) == 0) {
-            inputs.role = PSA_PAKE_ROLE_CLIENT;
-        } else
-        if (memcmp(inputs.user, jpake_server_id, inputs.user_len) == 0 &&
-            memcmp(inputs.peer, jpake_client_id, inputs.peer_len) == 0) {
-            inputs.role = PSA_PAKE_ROLE_SERVER;
-        }
-
-        if (inputs.role != PSA_PAKE_ROLE_CLIENT &&
-            inputs.role != PSA_PAKE_ROLE_SERVER) {
-            return PSA_ERROR_NOT_SUPPORTED;
-        }
     }
 
     /* Clear driver context */
diff --git a/library/psa_crypto_pake.c b/library/psa_crypto_pake.c
index a537184..4136614 100644
--- a/library/psa_crypto_pake.c
+++ b/library/psa_crypto_pake.c
@@ -168,13 +168,11 @@
 static psa_status_t psa_pake_ecjpake_setup(mbedtls_psa_pake_operation_t *operation)
 {
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
-    mbedtls_ecjpake_role role = (operation->role == PSA_PAKE_ROLE_CLIENT) ?
-                                MBEDTLS_ECJPAKE_CLIENT : MBEDTLS_ECJPAKE_SERVER;
 
     mbedtls_ecjpake_init(&operation->ctx.jpake);
 
     ret = mbedtls_ecjpake_setup(&operation->ctx.jpake,
-                                role,
+                                operation->role,
                                 MBEDTLS_MD_SHA256,
                                 MBEDTLS_ECP_DP_SECP256R1,
                                 operation->password,
@@ -190,21 +188,30 @@
 }
 #endif
 
+/* The only two JPAKE user/peer identifiers supported in built-in implementation. */
+static const uint8_t jpake_server_id[] = { 's', 'e', 'r', 'v', 'e', 'r' };
+static const uint8_t jpake_client_id[] = { 'c', 'l', 'i', 'e', 'n', 't' };
+
 psa_status_t mbedtls_psa_pake_setup(mbedtls_psa_pake_operation_t *operation,
                                     const psa_crypto_driver_pake_inputs_t *inputs)
 {
     psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
-    size_t password_len = 0;
-    psa_pake_role_t role = PSA_PAKE_ROLE_NONE;
+    size_t user_len = 0, peer_len = 0, password_len = 0;
+    uint8_t *peer = NULL, *user = NULL;
+    size_t actual_user_len = 0, actual_peer_len = 0, actual_password_len = 0;
     psa_pake_cipher_suite_t cipher_suite = psa_pake_cipher_suite_init();
-    size_t actual_password_len = 0;
 
     status = psa_crypto_driver_pake_get_password_len(inputs, &password_len);
     if (status != PSA_SUCCESS) {
         return status;
     }
 
-    status = psa_crypto_driver_pake_get_role(inputs, &role);
+    psa_crypto_driver_pake_get_user_len(inputs, &user_len);
+    if (status != PSA_SUCCESS) {
+        return status;
+    }
+
+    psa_crypto_driver_pake_get_peer_len(inputs, &peer_len);
     if (status != PSA_SUCCESS) {
         return status;
     }
@@ -216,7 +223,20 @@
 
     operation->password = mbedtls_calloc(1, password_len);
     if (operation->password == NULL) {
-        return PSA_ERROR_INSUFFICIENT_MEMORY;
+        status = PSA_ERROR_INSUFFICIENT_MEMORY;
+        goto error;
+    }
+
+    user = mbedtls_calloc(1, user_len);
+    if (user == NULL) {
+        status = PSA_ERROR_INSUFFICIENT_MEMORY;
+        goto error;
+    }
+
+    peer = mbedtls_calloc(1, peer_len);
+    if (peer == NULL) {
+        status = PSA_ERROR_INSUFFICIENT_MEMORY;
+        goto error;
     }
 
     status = psa_crypto_driver_pake_get_password(inputs, operation->password,
@@ -225,6 +245,18 @@
         goto error;
     }
 
+    status = psa_crypto_driver_pake_get_user(inputs, user,
+                                             user_len, &actual_user_len);
+    if (status != PSA_SUCCESS) {
+        goto error;
+    }
+
+    status = psa_crypto_driver_pake_get_peer(inputs, peer,
+                                             peer_len, &actual_peer_len);
+    if (status != PSA_SUCCESS) {
+        goto error;
+    }
+
     operation->password_len = actual_password_len;
     operation->alg = cipher_suite.algorithm;
 
@@ -238,7 +270,24 @@
             goto error;
         }
 
-        operation->role = role;
+        const size_t user_peer_len = sizeof(jpake_client_id); // client and server have the same length
+        if (actual_user_len != user_peer_len ||
+            actual_peer_len != user_peer_len) {
+            status = PSA_ERROR_NOT_SUPPORTED;
+            goto error;
+        }
+
+        if (memcmp(user, jpake_client_id, actual_user_len) == 0 &&
+            memcmp(peer, jpake_server_id, actual_peer_len) == 0) {
+            operation->role = MBEDTLS_ECJPAKE_CLIENT;
+        } else
+        if (memcmp(user, jpake_server_id, actual_user_len) == 0 &&
+            memcmp(peer, jpake_client_id, actual_peer_len) == 0) {
+            operation->role = MBEDTLS_ECJPAKE_SERVER;
+        } else {
+            status = PSA_ERROR_NOT_SUPPORTED;
+            goto error;
+        }
 
         operation->buffer_length = 0;
         operation->buffer_offset = 0;
@@ -248,6 +297,9 @@
             goto error;
         }
 
+        /* Role has been set, release user/peer buffers. */
+        mbedtls_free(user); mbedtls_free(peer);
+
         return PSA_SUCCESS;
     } else
 #else
@@ -257,6 +309,7 @@
     { status = PSA_ERROR_NOT_SUPPORTED; }
 
 error:
+    mbedtls_free(user); mbedtls_free(peer);
     /* In case of failure of the setup of a multipart operation, the PSA driver interface
      * specifies that the core does not call any other driver entry point thus does not
      * call mbedtls_psa_pake_abort(). Therefore call it here to do the needed clean
@@ -332,7 +385,7 @@
          * information is already available.
          */
         if (step == PSA_JPAKE_X2S_STEP_KEY_SHARE &&
-            operation->role == PSA_PAKE_ROLE_SERVER) {
+            operation->role == MBEDTLS_ECJPAKE_SERVER) {
             /* Skip ECParameters, with is 3 bytes (RFC 8422) */
             operation->buffer_offset += 3;
         }
@@ -423,7 +476,7 @@
          * we're a client.
          */
         if (step == PSA_JPAKE_X4S_STEP_KEY_SHARE &&
-            operation->role == PSA_PAKE_ROLE_CLIENT) {
+            operation->role == MBEDTLS_ECJPAKE_CLIENT) {
             /* We only support secp256r1. */
             /* This is the ECParameters structure defined by RFC 8422. */
             unsigned char ecparameters[3] = {
@@ -541,7 +594,7 @@
 
 #if defined(MBEDTLS_PSA_BUILTIN_ALG_JPAKE)
     if (operation->alg == PSA_ALG_JPAKE) {
-        operation->role = PSA_PAKE_ROLE_NONE;
+        operation->role = MBEDTLS_ECJPAKE_NONE;
         mbedtls_platform_zeroize(operation->buffer, sizeof(operation->buffer));
         operation->buffer_length = 0;
         operation->buffer_offset = 0;
diff --git a/tests/suites/test_suite_psa_crypto.data b/tests/suites/test_suite_psa_crypto.data
index f0b3574..fd35c87 100644
--- a/tests/suites/test_suite_psa_crypto.data
+++ b/tests/suites/test_suite_psa_crypto.data
@@ -5031,6 +5031,9 @@
 depends_on:PSA_WANT_ALG_SHA_256
 derive_setup:PSA_ALG_CATEGORY_KEY_DERIVATION:PSA_ERROR_NOT_SUPPORTED
 
+Parse binary string
+parse_binary_string_test:"123456":0x123456
+
 PSA key derivation: HKDF-SHA-256, good case, direct output
 depends_on:PSA_WANT_ALG_HKDF:PSA_WANT_ALG_SHA_256
 derive_input:PSA_ALG_HKDF(PSA_ALG_SHA_256):PSA_KEY_DERIVATION_INPUT_SALT:PSA_KEY_TYPE_NONE:"":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_SECRET:PSA_KEY_TYPE_DERIVE:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_INFO:PSA_KEY_TYPE_NONE:"":PSA_SUCCESS:PSA_KEY_TYPE_NONE:PSA_SUCCESS
@@ -5159,6 +5162,23 @@
 depends_on:PSA_WANT_ALG_HKDF:PSA_WANT_ALG_SHA_256
 derive_input:PSA_ALG_HKDF(PSA_ALG_SHA_256):PSA_KEY_DERIVATION_INPUT_SEED:PSA_KEY_TYPE_NONE:"":PSA_ERROR_INVALID_ARGUMENT:PSA_KEY_DERIVATION_INPUT_SECRET:PSA_KEY_TYPE_DERIVE:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_ERROR_BAD_STATE:PSA_KEY_DERIVATION_INPUT_INFO:PSA_KEY_TYPE_NONE:"":PSA_ERROR_BAD_STATE:PSA_KEY_TYPE_NONE:PSA_ERROR_BAD_STATE
 
+PSA key derivation: HKDF-SHA-256, reject using input integer with direct secret
+depends_on:PSA_WANT_ALG_HKDF:PSA_WANT_ALG_SHA_256
+derive_input:PSA_ALG_HKDF(PSA_ALG_SHA_256):PSA_KEY_DERIVATION_INPUT_SALT:PSA_KEY_TYPE_NONE:"":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_SECRET:INPUT_INTEGER:"0b0b0b0b0b0b0b0b":PSA_ERROR_INVALID_ARGUMENT:PSA_KEY_DERIVATION_INPUT_INFO:PSA_KEY_TYPE_NONE:"":PSA_ERROR_BAD_STATE:PSA_KEY_TYPE_NONE:PSA_ERROR_BAD_STATE
+
+PSA key derivation: HKDF-SHA-256, reject input cost step using input_bytes
+depends_on:PSA_WANT_ALG_HKDF:PSA_WANT_ALG_SHA_256
+derive_input:PSA_ALG_HKDF(PSA_ALG_SHA_256):PSA_KEY_DERIVATION_INPUT_COST:PSA_KEY_TYPE_NONE:"100000":PSA_ERROR_INVALID_ARGUMENT:PSA_KEY_DERIVATION_INPUT_SECRET:PSA_KEY_TYPE_NONE:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_ERROR_BAD_STATE:PSA_KEY_DERIVATION_INPUT_INFO:PSA_KEY_TYPE_NONE:"":PSA_ERROR_BAD_STATE:PSA_KEY_TYPE_NONE:PSA_ERROR_BAD_STATE
+
+PSA key derivation: HKDF-SHA-256, input cost using input_integer after secret
+depends_on:PSA_WANT_ALG_HKDF:PSA_WANT_ALG_SHA_256
+derive_input:PSA_ALG_HKDF(PSA_ALG_SHA_256):PSA_KEY_DERIVATION_INPUT_SECRET:PSA_KEY_TYPE_NONE:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_COST:INPUT_INTEGER:"100000":PSA_ERROR_INVALID_ARGUMENT:PSA_KEY_DERIVATION_INPUT_INFO:PSA_KEY_TYPE_NONE:"":PSA_ERROR_BAD_STATE:PSA_KEY_TYPE_NONE:PSA_ERROR_BAD_STATE
+
+PSA key derivation: HKDF-SHA-256, reject input cost using input_integer after secret and info
+depends_on:PSA_WANT_ALG_HKDF:PSA_WANT_ALG_SHA_256
+derive_input:PSA_ALG_HKDF(PSA_ALG_SHA_256):PSA_KEY_DERIVATION_INPUT_SECRET:PSA_KEY_TYPE_NONE:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_INFO:PSA_KEY_TYPE_NONE:"":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_COST:INPUT_INTEGER:"100000":PSA_ERROR_INVALID_ARGUMENT:PSA_KEY_TYPE_NONE:PSA_ERROR_BAD_STATE
+
+
 PSA key derivation: TLS 1.2 PRF SHA-256, good case
 depends_on:PSA_WANT_ALG_SHA_256:PSA_WANT_ALG_TLS12_PRF
 derive_input:PSA_ALG_TLS12_PRF(PSA_ALG_SHA_256):PSA_KEY_DERIVATION_INPUT_SEED:PSA_KEY_TYPE_NONE:"":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_SECRET:PSA_KEY_TYPE_DERIVE:"0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b":PSA_SUCCESS:PSA_KEY_DERIVATION_INPUT_LABEL:PSA_KEY_TYPE_NONE:"":PSA_SUCCESS:PSA_KEY_TYPE_DERIVE:PSA_SUCCESS
diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index cd8a7b5..fc8e6eb 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -294,6 +294,19 @@
     ((void) 0)
 #endif /* MBEDTLS_PSA_CRYPTO_SE_C */
 
+#define INPUT_INTEGER 0x10000   /* Out of range of psa_key_type_t */
+
+uint64_t parse_binary_string(data_t *bin_string)
+{
+    uint64_t result = 0;
+    TEST_LE_U(bin_string->len, 8);
+    for (size_t i = 0; i < bin_string->len; i++) {
+        result = result << 8 | bin_string->x[i];
+    }
+exit:
+    return result; /* returns 0 if len > 8 */
+}
+
 /* An overapproximation of the amount of storage needed for a key of the
  * given type and with the given content. The API doesn't make it easy
  * to find a good value for the size. The current implementation doesn't
@@ -318,6 +331,7 @@
     USE_GIVEN_TAG = 1,
 } tag_usage_method_t;
 
+
 /*!
  * \brief                           Internal Function for AEAD multipart tests.
  * \param key_type_arg              Type of key passed in
@@ -8446,6 +8460,15 @@
 /* END_CASE */
 
 /* BEGIN_CASE */
+void parse_binary_string_test(data_t *input, int output)
+{
+    uint64_t value;
+    value = parse_binary_string(input);
+    TEST_EQUAL(value, output);
+}
+/* END_CASE */
+
+/* BEGIN_CASE */
 void derive_input(int alg_arg,
                   int step_arg1, int key_type_arg1, data_t *input1,
                   int expected_status_arg1,
@@ -8457,7 +8480,7 @@
 {
     psa_algorithm_t alg = alg_arg;
     psa_key_derivation_step_t steps[] = { step_arg1, step_arg2, step_arg3 };
-    psa_key_type_t key_types[] = { key_type_arg1, key_type_arg2, key_type_arg3 };
+    uint32_t key_types[] = { key_type_arg1, key_type_arg2, key_type_arg3 };
     psa_status_t expected_statuses[] = { expected_status_arg1,
                                          expected_status_arg2,
                                          expected_status_arg3 };
@@ -8484,12 +8507,13 @@
         mbedtls_test_set_step(i);
         if (steps[i] == 0) {
             /* Skip this step */
-        } else if (key_types[i] != PSA_KEY_TYPE_NONE) {
-            psa_set_key_type(&attributes, key_types[i]);
+        } else if (((psa_key_type_t) key_types[i]) != PSA_KEY_TYPE_NONE &&
+                   key_types[i] != INPUT_INTEGER) {
+            psa_set_key_type(&attributes, ((psa_key_type_t) key_types[i]));
             PSA_ASSERT(psa_import_key(&attributes,
                                       inputs[i]->x, inputs[i]->len,
                                       &keys[i]));
-            if (PSA_KEY_TYPE_IS_KEY_PAIR(key_types[i]) &&
+            if (PSA_KEY_TYPE_IS_KEY_PAIR((psa_key_type_t) key_types[i]) &&
                 steps[i] == PSA_KEY_DERIVATION_INPUT_SECRET) {
                 // When taking a private key as secret input, use key agreement
                 // to add the shared secret to the derivation
@@ -8502,10 +8526,17 @@
                            expected_statuses[i]);
             }
         } else {
-            TEST_EQUAL(psa_key_derivation_input_bytes(
-                           &operation, steps[i],
-                           inputs[i]->x, inputs[i]->len),
-                       expected_statuses[i]);
+            if (key_types[i] == INPUT_INTEGER) {
+                TEST_EQUAL(psa_key_derivation_input_integer(
+                               &operation, steps[i],
+                               parse_binary_string(inputs[i])),
+                           expected_statuses[i]);
+            } else {
+                TEST_EQUAL(psa_key_derivation_input_bytes(
+                               &operation, steps[i],
+                               inputs[i]->x, inputs[i]->len),
+                           expected_statuses[i]);
+            }
         }
     }
 
diff --git a/tests/suites/test_suite_psa_crypto_pake.data b/tests/suites/test_suite_psa_crypto_pake.data
index c467d01..b9f68e1 100644
--- a/tests/suites/test_suite_psa_crypto_pake.data
+++ b/tests/suites/test_suite_psa_crypto_pake.data
@@ -48,11 +48,11 @@
 
 PSA PAKE: set invalid user
 depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
-ecjpake_setup:PSA_ALG_JPAKE:PSA_KEY_TYPE_PASSWORD:PSA_KEY_USAGE_DERIVE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:"aaaa":"server":0:ERR_IN_SET_USER:PSA_ERROR_NOT_SUPPORTED
+ecjpake_setup:PSA_ALG_JPAKE:PSA_KEY_TYPE_PASSWORD:PSA_KEY_USAGE_DERIVE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:"something":"server":0:ERR_IN_OUTPUT:PSA_ERROR_NOT_SUPPORTED
 
 PSA PAKE: set invalid peer
 depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
-ecjpake_setup:PSA_ALG_JPAKE:PSA_KEY_TYPE_PASSWORD:PSA_KEY_USAGE_DERIVE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:"client":"aaaa":0:ERR_IN_SET_PEER:PSA_ERROR_NOT_SUPPORTED
+ecjpake_setup:PSA_ALG_JPAKE:PSA_KEY_TYPE_PASSWORD:PSA_KEY_USAGE_DERIVE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:"client":"something":0:ERR_IN_OUTPUT:PSA_ERROR_NOT_SUPPORTED
 
 PSA PAKE: user already set
 depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
@@ -220,9 +220,6 @@
 PSA PAKE: input getters: cipher suite
 pake_input_getters_cipher_suite
 
-PSA PAKE: input getters: role
-pake_input_getters_role
-
 PSA PAKE: input getters: user
 pake_input_getters_user
 
diff --git a/tests/suites/test_suite_psa_crypto_pake.function b/tests/suites/test_suite_psa_crypto_pake.function
index ecbd363..52380de 100644
--- a/tests/suites/test_suite_psa_crypto_pake.function
+++ b/tests/suites/test_suite_psa_crypto_pake.function
@@ -989,8 +989,7 @@
                                                    &buffer_len_ret),
                PSA_SUCCESS);
 
-    TEST_EQUAL(buffer_len_ret, strlen(password));
-    PSA_ASSERT(memcmp(password_ret, password, buffer_len_ret));
+    ASSERT_COMPARE(password_ret, buffer_len_ret, password, strlen(password));
 exit:
     PSA_ASSERT(psa_destroy_key(key));
     PSA_ASSERT(psa_pake_abort(&operation));
@@ -1023,7 +1022,8 @@
     TEST_EQUAL(psa_crypto_driver_pake_get_cipher_suite(&operation.data.inputs, &cipher_suite_ret),
                PSA_SUCCESS);
 
-    PSA_ASSERT(memcmp(&cipher_suite_ret, &cipher_suite, sizeof(cipher_suite)));
+    ASSERT_COMPARE(&cipher_suite_ret, sizeof(cipher_suite_ret),
+                   &cipher_suite, sizeof(cipher_suite));
 
 exit:
     PSA_ASSERT(psa_pake_abort(&operation));
@@ -1032,47 +1032,11 @@
 /* END_CASE */
 
 /* BEGIN_CASE depends_on:PSA_WANT_ALG_JPAKE */
-void pake_input_getters_role()
-{
-    psa_pake_cipher_suite_t cipher_suite = psa_pake_cipher_suite_init();
-    psa_pake_operation_t operation = psa_pake_operation_init();
-    psa_pake_role_t role_ret = PSA_PAKE_ROLE_NONE;
-
-    psa_pake_primitive_t primitive = PSA_PAKE_PRIMITIVE(
-        PSA_PAKE_PRIMITIVE_TYPE_ECC,
-        PSA_ECC_FAMILY_SECP_R1, 256);
-
-    PSA_INIT();
-
-    psa_pake_cs_set_algorithm(&cipher_suite, PSA_ALG_JPAKE);
-    psa_pake_cs_set_primitive(&cipher_suite, primitive);
-    psa_pake_cs_set_hash(&cipher_suite, PSA_ALG_SHA_256);
-
-    PSA_ASSERT(psa_pake_setup(&operation, &cipher_suite));
-
-    TEST_EQUAL(psa_crypto_driver_pake_get_role(&operation.data.inputs, &role_ret),
-               PSA_ERROR_BAD_STATE);
-
-    /* Role can not be set directly using psa_pake_set_role(). It is set by the core
-       based on given user/peer identifiers. Simulate that Role is already set. */
-    operation.data.inputs.role = PSA_PAKE_ROLE_SERVER;
-    TEST_EQUAL(psa_crypto_driver_pake_get_role(&operation.data.inputs, &role_ret),
-               PSA_SUCCESS);
-
-    TEST_EQUAL(role_ret, PSA_PAKE_ROLE_SERVER);
-exit:
-    PSA_ASSERT(psa_pake_abort(&operation));
-    PSA_DONE();
-}
-/* END_CASE */
-
-/* BEGIN_CASE depends_on:PSA_WANT_ALG_JPAKE */
 void pake_input_getters_user()
 {
     psa_pake_cipher_suite_t cipher_suite = psa_pake_cipher_suite_init();
     psa_pake_operation_t operation = psa_pake_operation_init();
-    const uint8_t user[] = { 's', 'e', 'r', 'v', 'e', 'r' };
-    const size_t user_len = sizeof(user);
+    const char *users[] = { "client", "server", "other" };
     uint8_t user_ret[20] = { 0 }; // max user length is 20 bytes
     size_t user_len_ret = 0;
     size_t buffer_len_ret = 0;
@@ -1087,37 +1051,43 @@
     psa_pake_cs_set_primitive(&cipher_suite, primitive);
     psa_pake_cs_set_hash(&cipher_suite, PSA_ALG_SHA_256);
 
-    PSA_ASSERT(psa_pake_setup(&operation, &cipher_suite));
+    for (size_t i = 0; i < ARRAY_LENGTH(users); i++) {
+        uint8_t *user = (uint8_t *) users[i];
+        uint8_t user_len = strlen(users[i]);
 
-    TEST_EQUAL(psa_crypto_driver_pake_get_user(&operation.data.inputs,
-                                               (uint8_t *) &user_ret,
-                                               10, &buffer_len_ret),
-               PSA_ERROR_BAD_STATE);
+        PSA_ASSERT(psa_pake_abort(&operation));
 
-    TEST_EQUAL(psa_crypto_driver_pake_get_user_len(&operation.data.inputs, &user_len_ret),
-               PSA_ERROR_BAD_STATE);
+        PSA_ASSERT(psa_pake_setup(&operation, &cipher_suite));
 
-    PSA_ASSERT(psa_pake_set_user(&operation, user, user_len));
+        TEST_EQUAL(psa_crypto_driver_pake_get_user(&operation.data.inputs,
+                                                   (uint8_t *) &user_ret,
+                                                   10, &buffer_len_ret),
+                   PSA_ERROR_BAD_STATE);
 
-    TEST_EQUAL(psa_crypto_driver_pake_get_user_len(&operation.data.inputs, &user_len_ret),
-               PSA_SUCCESS);
+        TEST_EQUAL(psa_crypto_driver_pake_get_user_len(&operation.data.inputs, &user_len_ret),
+                   PSA_ERROR_BAD_STATE);
 
-    TEST_EQUAL(user_len_ret, user_len);
+        PSA_ASSERT(psa_pake_set_user(&operation, user, user_len));
 
-    TEST_EQUAL(psa_crypto_driver_pake_get_user(&operation.data.inputs,
-                                               (uint8_t *) &user_ret,
-                                               user_len_ret - 1,
-                                               &buffer_len_ret),
-               PSA_ERROR_BUFFER_TOO_SMALL);
+        TEST_EQUAL(psa_crypto_driver_pake_get_user_len(&operation.data.inputs, &user_len_ret),
+                   PSA_SUCCESS);
 
-    TEST_EQUAL(psa_crypto_driver_pake_get_user(&operation.data.inputs,
-                                               (uint8_t *) &user_ret,
-                                               user_len_ret,
-                                               &buffer_len_ret),
-               PSA_SUCCESS);
+        TEST_EQUAL(user_len_ret, user_len);
 
-    TEST_EQUAL(buffer_len_ret, user_len);
-    PSA_ASSERT(memcmp(user_ret, user, buffer_len_ret));
+        TEST_EQUAL(psa_crypto_driver_pake_get_user(&operation.data.inputs,
+                                                   (uint8_t *) &user_ret,
+                                                   user_len_ret - 1,
+                                                   &buffer_len_ret),
+                   PSA_ERROR_BUFFER_TOO_SMALL);
+
+        TEST_EQUAL(psa_crypto_driver_pake_get_user(&operation.data.inputs,
+                                                   (uint8_t *) &user_ret,
+                                                   user_len_ret,
+                                                   &buffer_len_ret),
+                   PSA_SUCCESS);
+
+        ASSERT_COMPARE(user_ret, buffer_len_ret, user, user_len);
+    }
 exit:
     PSA_ASSERT(psa_pake_abort(&operation));
     PSA_DONE();
@@ -1129,8 +1099,7 @@
 {
     psa_pake_cipher_suite_t cipher_suite = psa_pake_cipher_suite_init();
     psa_pake_operation_t operation = psa_pake_operation_init();
-    const uint8_t peer[] = { 's', 'e', 'r', 'v', 'e', 'r' };
-    const size_t peer_len = sizeof(peer);
+    const char *peers[] = { "client", "server", "other" };
     uint8_t peer_ret[20] = { 0 }; // max peer length is 20 bytes
     size_t peer_len_ret = 0;
     size_t buffer_len_ret = 0;
@@ -1145,37 +1114,43 @@
     psa_pake_cs_set_primitive(&cipher_suite, primitive);
     psa_pake_cs_set_hash(&cipher_suite, PSA_ALG_SHA_256);
 
-    PSA_ASSERT(psa_pake_setup(&operation, &cipher_suite));
+    for (size_t i = 0; i < ARRAY_LENGTH(peers); i++) {
+        uint8_t *peer = (uint8_t *) peers[i];
+        uint8_t peer_len = strlen(peers[i]);
 
-    TEST_EQUAL(psa_crypto_driver_pake_get_peer(&operation.data.inputs,
-                                               (uint8_t *) &peer_ret,
-                                               10, &buffer_len_ret),
-               PSA_ERROR_BAD_STATE);
+        PSA_ASSERT(psa_pake_abort(&operation));
 
-    TEST_EQUAL(psa_crypto_driver_pake_get_peer_len(&operation.data.inputs, &peer_len_ret),
-               PSA_ERROR_BAD_STATE);
+        PSA_ASSERT(psa_pake_setup(&operation, &cipher_suite));
 
-    PSA_ASSERT(psa_pake_set_peer(&operation, peer, peer_len));
+        TEST_EQUAL(psa_crypto_driver_pake_get_peer(&operation.data.inputs,
+                                                   (uint8_t *) &peer_ret,
+                                                   10, &buffer_len_ret),
+                   PSA_ERROR_BAD_STATE);
 
-    TEST_EQUAL(psa_crypto_driver_pake_get_peer_len(&operation.data.inputs, &peer_len_ret),
-               PSA_SUCCESS);
+        TEST_EQUAL(psa_crypto_driver_pake_get_peer_len(&operation.data.inputs, &peer_len_ret),
+                   PSA_ERROR_BAD_STATE);
 
-    TEST_EQUAL(peer_len_ret, peer_len);
+        PSA_ASSERT(psa_pake_set_peer(&operation, peer, peer_len));
 
-    TEST_EQUAL(psa_crypto_driver_pake_get_peer(&operation.data.inputs,
-                                               (uint8_t *) &peer_ret,
-                                               peer_len_ret - 1,
-                                               &buffer_len_ret),
-               PSA_ERROR_BUFFER_TOO_SMALL);
+        TEST_EQUAL(psa_crypto_driver_pake_get_peer_len(&operation.data.inputs, &peer_len_ret),
+                   PSA_SUCCESS);
 
-    TEST_EQUAL(psa_crypto_driver_pake_get_peer(&operation.data.inputs,
-                                               (uint8_t *) &peer_ret,
-                                               peer_len_ret,
-                                               &buffer_len_ret),
-               PSA_SUCCESS);
+        TEST_EQUAL(peer_len_ret, peer_len);
 
-    TEST_EQUAL(buffer_len_ret, peer_len);
-    PSA_ASSERT(memcmp(peer_ret, peer, buffer_len_ret));
+        TEST_EQUAL(psa_crypto_driver_pake_get_peer(&operation.data.inputs,
+                                                   (uint8_t *) &peer_ret,
+                                                   peer_len_ret - 1,
+                                                   &buffer_len_ret),
+                   PSA_ERROR_BUFFER_TOO_SMALL);
+
+        TEST_EQUAL(psa_crypto_driver_pake_get_peer(&operation.data.inputs,
+                                                   (uint8_t *) &peer_ret,
+                                                   peer_len_ret,
+                                                   &buffer_len_ret),
+                   PSA_SUCCESS);
+
+        ASSERT_COMPARE(peer_ret, buffer_len_ret, peer, peer_len);
+    }
 exit:
     PSA_ASSERT(psa_pake_abort(&operation));
     PSA_DONE();