Code optimizations and documentation fixes

Signed-off-by: Przemek Stekiel <przemyslaw.stekiel@mobica.com>
diff --git a/include/psa/crypto_extra.h b/include/psa/crypto_extra.h
index ba43c72..e491a86 100644
--- a/include/psa/crypto_extra.h
+++ b/include/psa/crypto_extra.h
@@ -1562,7 +1562,7 @@
  *                              been set (psa_pake_set_user() hasn't been
  *                              called yet).
  * \param[in] user_id           The user ID to authenticate with.
- *                              ("client" or "server")
+ *                              (temporary limitation: "client" or "server" only)
  * \param user_id_len           Size of the \p user_id buffer in bytes.
  *
  * \retval #PSA_SUCCESS
@@ -1604,7 +1604,7 @@
  *                              been set (psa_pake_set_peer() hasn't been
  *                              called yet).
  * \param[in] peer_id           The peer's ID to authenticate.
- *                              ("client" or "server")
+ *                              (temporary limitation: "client" or "server" only)
  * \param peer_id_len           Size of the \p peer_id buffer in bytes.
  *
  * \retval #PSA_SUCCESS
diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 767dbfc..b875412 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -91,9 +91,9 @@
 #define BUILTIN_ALG_ANY_HKDF 1
 #endif
 
-/* JPAKE user/peer ids. */
-#define JPAKE_SERVER_ID "server"
-#define JPAKE_CLIENT_ID "client"
+/* 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 */
@@ -7406,8 +7406,10 @@
     }
 
     /* Allow only "client" or "server" values (temporary restriction). */
-    if (memcmp(user_id, JPAKE_SERVER_ID, user_id_len) != 0 &&
-        memcmp(user_id, JPAKE_CLIENT_ID, user_id_len) != 0) {
+    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;
     }
@@ -7450,8 +7452,10 @@
     }
 
     /* Allow only "client" or "server" values (temporary restriction). */
-    if (memcmp(peer_id, JPAKE_SERVER_ID, peer_id_len) != 0 &&
-        memcmp(peer_id, JPAKE_CLIENT_ID, peer_id_len) != 0) {
+    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;
     }
@@ -7565,19 +7569,20 @@
        with the driver context which will be setup by the driver. */
     psa_crypto_driver_pake_inputs_t inputs = operation->data.inputs;
 
-    if (inputs.password_len == 0 ||
-        inputs.user_len == 0 ||
-        inputs.peer_len == 0) {
+    if (inputs.password_len == 0) {
         return PSA_ERROR_BAD_STATE;
     }
 
     if (operation->alg == PSA_ALG_JPAKE) {
-        if (memcmp(inputs.user, JPAKE_CLIENT_ID, inputs.user_len) == 0 &&
-            memcmp(inputs.peer, JPAKE_SERVER_ID, inputs.peer_len) == 0) {
+        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) {
+        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;
         }
 
@@ -7599,8 +7604,6 @@
     /* User and peer are translated to role. */
     mbedtls_free(inputs.user);
     mbedtls_free(inputs.peer);
-    inputs.user = NULL; inputs.user_len = 0;
-    inputs.peer = NULL; inputs.peer_len = 0;
 
     if (status == PSA_SUCCESS) {
 #if defined(PSA_WANT_ALG_JPAKE)
diff --git a/library/ssl_tls.c b/library/ssl_tls.c
index 4a351f3..f43bb73 100644
--- a/library/ssl_tls.c
+++ b/library/ssl_tls.c
@@ -61,10 +61,6 @@
                                                       psa_generic_status_to_mbedtls)
 #endif
 
-/* JPAKE user/peer ids. */
-#define JPAKE_SERVER_ID "server"
-#define JPAKE_CLIENT_ID "client"
-
 #if defined(MBEDTLS_TEST_HOOKS)
 static mbedtls_ssl_chk_buf_ptr_args chk_buf_ptr_fail_args;
 
@@ -1953,15 +1949,19 @@
 #if defined(MBEDTLS_KEY_EXCHANGE_ECJPAKE_ENABLED)
 
 #if defined(MBEDTLS_USE_PSA_CRYPTO)
+/* 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' };
+
 static psa_status_t mbedtls_ssl_set_hs_ecjpake_password_common(
     mbedtls_ssl_context *ssl,
     mbedtls_svc_key_id_t pwd)
 {
     psa_status_t status;
     psa_pake_cipher_suite_t cipher_suite = psa_pake_cipher_suite_init();
-    uint8_t *user = NULL;
+    const uint8_t *user = NULL;
     size_t user_len = 0;
-    uint8_t *peer = NULL;
+    const uint8_t *peer = NULL;
     size_t peer_len = 0;
     psa_pake_cs_set_algorithm(&cipher_suite, PSA_ALG_JPAKE);
     psa_pake_cs_set_primitive(&cipher_suite,
@@ -1976,15 +1976,15 @@
     }
 
     if (ssl->conf->endpoint == MBEDTLS_SSL_IS_SERVER) {
-        user = (uint8_t *) JPAKE_SERVER_ID;
-        user_len = strlen(JPAKE_SERVER_ID);
-        peer = (uint8_t *) JPAKE_CLIENT_ID;
-        peer_len = strlen(JPAKE_CLIENT_ID);
+        user = jpake_server_id;
+        user_len = sizeof(jpake_server_id);
+        peer = jpake_client_id;
+        peer_len = sizeof(jpake_client_id);
     } else {
-        user = (uint8_t *) JPAKE_CLIENT_ID;
-        user_len = strlen(JPAKE_CLIENT_ID);
-        peer = (uint8_t *) JPAKE_SERVER_ID;
-        peer_len = strlen(JPAKE_SERVER_ID);
+        user = jpake_client_id;
+        user_len = sizeof(jpake_client_id);
+        peer = jpake_server_id;
+        peer_len = sizeof(jpake_server_id);
     }
 
     status = psa_pake_set_user(&ssl->handshake->psa_pake_ctx, user, user_len);
diff --git a/tests/suites/test_suite_psa_crypto_driver_wrappers.function b/tests/suites/test_suite_psa_crypto_driver_wrappers.function
index c235ff6..56f4d1d 100644
--- a/tests/suites/test_suite_psa_crypto_driver_wrappers.function
+++ b/tests/suites/test_suite_psa_crypto_driver_wrappers.function
@@ -6,9 +6,9 @@
 size_t pake_expected_hit_count = 0;
 int pake_in_driver = 0;
 
-/* JPAKE user/peer ids. */
-#define JPAKE_SERVER_ID "server"
-#define JPAKE_CLIENT_ID "client"
+/* 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' };
 
 #if defined(PSA_WANT_ALG_JPAKE) && defined(PSA_WANT_KEY_TYPE_ECC_KEY_PAIR) && \
     defined(PSA_WANT_ECC_SECP_R1_256) && defined(PSA_WANT_ALG_SHA_256)
@@ -2999,10 +2999,6 @@
         PSA_ECC_FAMILY_SECP_R1, 256);
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
     unsigned char *input_buffer = NULL;
-    const uint8_t server_id[] = JPAKE_SERVER_ID;
-    const uint8_t client_id[] = JPAKE_CLIENT_ID;
-    const size_t server_id_len = strlen(JPAKE_SERVER_ID);
-    const size_t client_id_len = strlen(JPAKE_CLIENT_ID);
     const size_t size_key_share = PSA_PAKE_INPUT_SIZE(PSA_ALG_JPAKE, primitive,
                                                       PSA_PAKE_STEP_KEY_SHARE);
     unsigned char *output_buffer = NULL;
@@ -3044,8 +3040,8 @@
     TEST_EQUAL(psa_pake_setup(&operation, &cipher_suite),
                PSA_SUCCESS);
 
-    PSA_ASSERT(psa_pake_set_user(&operation, server_id, server_id_len));
-    PSA_ASSERT(psa_pake_set_peer(&operation, client_id, client_id_len));
+    PSA_ASSERT(psa_pake_set_user(&operation, jpake_server_id, sizeof(jpake_server_id)));
+    PSA_ASSERT(psa_pake_set_peer(&operation, jpake_client_id, sizeof(jpake_client_id)));
 
     TEST_EQUAL(psa_pake_set_password_key(&operation, key),
                PSA_SUCCESS);
@@ -3193,10 +3189,6 @@
         PSA_KEY_DERIVATION_OPERATION_INIT;
     psa_key_derivation_operation_t client_derive =
         PSA_KEY_DERIVATION_OPERATION_INIT;
-    const uint8_t server_id[] = JPAKE_SERVER_ID;
-    const uint8_t client_id[] = JPAKE_CLIENT_ID;
-    const size_t server_id_len = strlen(JPAKE_SERVER_ID);
-    const size_t client_id_len = strlen(JPAKE_CLIENT_ID);
     pake_in_driver = in_driver;
     /* driver setup is called indirectly through pake_output/pake_input */
     if (pake_in_driver) {
@@ -3242,11 +3234,11 @@
     TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total, 0);
 
 
-    PSA_ASSERT(psa_pake_set_user(&server, server_id, server_id_len));
-    PSA_ASSERT(psa_pake_set_peer(&server, client_id, client_id_len));
+    PSA_ASSERT(psa_pake_set_user(&server, jpake_server_id, sizeof(jpake_server_id)));
+    PSA_ASSERT(psa_pake_set_peer(&server, jpake_client_id, sizeof(jpake_client_id)));
     TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total, 0);
-    PSA_ASSERT(psa_pake_set_user(&client, client_id, client_id_len));
-    PSA_ASSERT(psa_pake_set_peer(&client, server_id, server_id_len));
+    PSA_ASSERT(psa_pake_set_user(&client, jpake_client_id, sizeof(jpake_client_id)));
+    PSA_ASSERT(psa_pake_set_peer(&client, jpake_server_id, sizeof(jpake_server_id)));
     TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total, 0);
     PSA_ASSERT(psa_pake_set_password_key(&server, key));
     TEST_EQUAL(mbedtls_test_driver_pake_hooks.hits.total, 0);
diff --git a/tests/suites/test_suite_psa_crypto_pake.function b/tests/suites/test_suite_psa_crypto_pake.function
index d8b1035..00b5c2f 100644
--- a/tests/suites/test_suite_psa_crypto_pake.function
+++ b/tests/suites/test_suite_psa_crypto_pake.function
@@ -53,9 +53,9 @@
     PAKE_ROUND_TWO
 } pake_round_t;
 
-/* JPAKE user/peer ids. */
-#define JPAKE_SERVER_ID "server"
-#define JPAKE_CLIENT_ID "client"
+/* 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' };
 
 /*
  * Inject an error on the specified buffer ONLY it this is the correct stage.
@@ -737,10 +737,6 @@
     mbedtls_svc_key_id_t key = MBEDTLS_SVC_KEY_ID_INIT;
     psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
     ecjpake_error_stage_t err_stage = err_stage_arg;
-    const uint8_t server_id[] = JPAKE_SERVER_ID;
-    const uint8_t client_id[] = JPAKE_CLIENT_ID;
-    const size_t server_id_len = strlen(JPAKE_SERVER_ID);
-    const size_t client_id_len = strlen(JPAKE_CLIENT_ID);
 
     PSA_INIT();
 
@@ -758,10 +754,10 @@
     PSA_ASSERT(psa_pake_setup(&server, &cipher_suite));
     PSA_ASSERT(psa_pake_setup(&client, &cipher_suite));
 
-    PSA_ASSERT(psa_pake_set_user(&server, server_id, server_id_len));
-    PSA_ASSERT(psa_pake_set_peer(&server, client_id, client_id_len));
-    PSA_ASSERT(psa_pake_set_user(&client, client_id, client_id_len));
-    PSA_ASSERT(psa_pake_set_peer(&client, server_id, server_id_len));
+    PSA_ASSERT(psa_pake_set_user(&server, jpake_server_id, sizeof(jpake_server_id)));
+    PSA_ASSERT(psa_pake_set_peer(&server, jpake_client_id, sizeof(jpake_client_id)));
+    PSA_ASSERT(psa_pake_set_user(&client, jpake_client_id, sizeof(jpake_client_id)));
+    PSA_ASSERT(psa_pake_set_peer(&client, jpake_server_id, sizeof(jpake_server_id)));
 
     PSA_ASSERT(psa_pake_set_password_key(&server, key));
     PSA_ASSERT(psa_pake_set_password_key(&client, key));
@@ -805,10 +801,6 @@
     psa_key_derivation_operation_t client_derive =
         PSA_KEY_DERIVATION_OPERATION_INIT;
     ecjpake_error_stage_t err_stage = err_stage_arg;
-    const uint8_t server_id[] = JPAKE_SERVER_ID;
-    const uint8_t client_id[] = JPAKE_CLIENT_ID;
-    const size_t server_id_len = strlen(JPAKE_SERVER_ID);
-    const size_t client_id_len = strlen(JPAKE_CLIENT_ID);
 
     PSA_INIT();
 
@@ -839,10 +831,10 @@
     PSA_ASSERT(psa_pake_setup(&server, &cipher_suite));
     PSA_ASSERT(psa_pake_setup(&client, &cipher_suite));
 
-    PSA_ASSERT(psa_pake_set_user(&server, server_id, server_id_len));
-    PSA_ASSERT(psa_pake_set_peer(&server, client_id, client_id_len));
-    PSA_ASSERT(psa_pake_set_user(&client, client_id, client_id_len));
-    PSA_ASSERT(psa_pake_set_peer(&client, server_id, server_id_len));
+    PSA_ASSERT(psa_pake_set_user(&server, jpake_server_id, sizeof(jpake_server_id)));
+    PSA_ASSERT(psa_pake_set_peer(&server, jpake_client_id, sizeof(jpake_client_id)));
+    PSA_ASSERT(psa_pake_set_user(&client, jpake_client_id, sizeof(jpake_client_id)));
+    PSA_ASSERT(psa_pake_set_peer(&client, jpake_server_id, sizeof(jpake_server_id)));
 
     PSA_ASSERT(psa_pake_set_password_key(&server, key));
     PSA_ASSERT(psa_pake_set_password_key(&client, key));
@@ -1056,7 +1048,7 @@
                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. Simulate that Role is already set. */
+       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);