Add tests for user/peer input getters

Signed-off-by: Przemek Stekiel <przemyslaw.stekiel@mobica.com>
diff --git a/tests/suites/test_suite_psa_crypto_pake.data b/tests/suites/test_suite_psa_crypto_pake.data
index fb8bb9e..1cbd8ac 100644
--- a/tests/suites/test_suite_psa_crypto_pake.data
+++ b/tests/suites/test_suite_psa_crypto_pake.data
@@ -218,3 +218,9 @@
 
 PSA PAKE: input getters: role
 pake_input_getters_role
+
+PSA PAKE: input getters: user
+pake_input_getters_user
+
+PSA PAKE: input getters: peer
+pake_input_getters_peer
diff --git a/tests/suites/test_suite_psa_crypto_pake.function b/tests/suites/test_suite_psa_crypto_pake.function
index fae78ba..bea017f 100644
--- a/tests/suites/test_suite_psa_crypto_pake.function
+++ b/tests/suites/test_suite_psa_crypto_pake.function
@@ -1057,8 +1057,9 @@
     TEST_EQUAL(psa_crypto_driver_pake_get_role(&operation.data.inputs, &role_ret),
                PSA_ERROR_BAD_STATE);
 
-    PSA_ASSERT(psa_pake_set_role(&operation, PSA_PAKE_ROLE_SERVER));
-
+    /* 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. */
+    operation.data.inputs.role = PSA_PAKE_ROLE_SERVER;
     TEST_EQUAL(psa_crypto_driver_pake_get_role(&operation.data.inputs, &role_ret),
                PSA_SUCCESS);
 
@@ -1068,3 +1069,119 @@
     PSA_DONE();
 }
 /* END_CASE */
+
+/* BEGIN_CASE depends_on:PSA_WANT_ALG_JPAKE:PSA_ALG_SHA_256 */
+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[] = "server";
+    const size_t user_len = strlen("server");
+    uint8_t user_ret[20] = { 0 }; // max user length is 20 bytes
+    size_t user_len_ret = 0;
+    size_t buffer_len_ret = 0;
+
+    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_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_ERROR_BAD_STATE);
+
+    PSA_ASSERT(psa_pake_set_user(&operation, user, user_len));
+
+    TEST_EQUAL(psa_crypto_driver_pake_get_user_len(&operation.data.inputs, &user_len_ret),
+               PSA_SUCCESS);
+
+    TEST_EQUAL(user_len_ret, 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(&operation.data.inputs,
+                                                   (uint8_t *) &user_ret,
+                                                   user_len_ret,
+                                                   &buffer_len_ret),
+               PSA_SUCCESS);
+
+    TEST_EQUAL(buffer_len_ret, user_len);
+    PSA_ASSERT(memcmp(user_ret, user, buffer_len_ret));
+exit:
+    PSA_ASSERT(psa_pake_abort(&operation));
+    PSA_DONE();
+}
+/* END_CASE */
+
+/* BEGIN_CASE depends_on:PSA_WANT_ALG_JPAKE:PSA_ALG_SHA_256 */
+void pake_input_getters_peer()
+{
+    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[] = "server";
+    const size_t peer_len = strlen("server");
+    uint8_t peer_ret[20] = { 0 }; // max peer length is 20 bytes
+    size_t peer_len_ret = 0;
+    size_t buffer_len_ret = 0;
+
+    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_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_ERROR_BAD_STATE);
+
+    PSA_ASSERT(psa_pake_set_peer(&operation, peer, peer_len));
+
+    TEST_EQUAL(psa_crypto_driver_pake_get_peer_len(&operation.data.inputs, &peer_len_ret),
+               PSA_SUCCESS);
+
+    TEST_EQUAL(peer_len_ret, 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(&operation.data.inputs,
+                                                   (uint8_t *) &peer_ret,
+                                                   peer_len_ret,
+                                                   &buffer_len_ret),
+               PSA_SUCCESS);
+
+    TEST_EQUAL(buffer_len_ret, peer_len);
+    PSA_ASSERT(memcmp(peer_ret, peer, buffer_len_ret));
+exit:
+    PSA_ASSERT(psa_pake_abort(&operation));
+    PSA_DONE();
+}
+/* END_CASE */