Add get-data functions for inputs + tests

Signed-off-by: Przemek Stekiel <przemyslaw.stekiel@mobica.com>
diff --git a/include/psa/crypto_extra.h b/include/psa/crypto_extra.h
index 57b1733..79fb263 100644
--- a/include/psa/crypto_extra.h
+++ b/include/psa/crypto_extra.h
@@ -1305,6 +1305,64 @@
  */
 static psa_pake_operation_t psa_pake_operation_init(void);
 
+/** Get the lengths of the password in bytes from given inputs.
+ *
+ * \param[in]  inputs           Operation inputs.
+ * \param[out] password_len     Return buffer for password length.
+ *
+ * \retval #PSA_SUCCESS
+ *         Success.
+ * \retval #PSA_ERROR_BAD_STATE
+ *         Password hasn't been set yet.
+ */
+psa_status_t psa_crypto_driver_pake_get_password_len(
+    const psa_crypto_driver_pake_inputs_t *inputs,
+    size_t *password_len);
+
+/** Get the password from given inputs.
+ *
+ * \param[in]  inputs           Operation inputs.
+ * \param[out] buffer           Return buffer for password.
+ * \param[in]  buffer_size      Size of the return buffer in bytes.
+ * \param[in]  buffer_length    Actual size of the password in bytes.
+ *
+ * \retval #PSA_SUCCESS
+ *         Success.
+ * \retval #PSA_ERROR_BAD_STATE
+ *         Password hasn't been set yet.
+ */
+psa_status_t psa_crypto_driver_pake_get_password(
+    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 cipher suite from given inputs.
+ *
+ * \param[in]  inputs           Operation inputs.
+ * \param[out] cipher_suite     Return buffer for role.
+ *
+ * \retval #PSA_SUCCESS
+ *         Success.
+ * \retval #PSA_ERROR_BAD_STATE
+ *         Cipher_suite hasn't been set yet.
+ */
+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);
+
 /** Set the session information for a password-authenticated key exchange.
  *
  * The sequence of operations to set up a password-authenticated key exchange
diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 46d62b0..0630885 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -7176,6 +7176,63 @@
     return status;
 }
 
+psa_status_t psa_crypto_driver_pake_get_password_len(
+    const psa_crypto_driver_pake_inputs_t *inputs,
+    size_t *password_len)
+{
+    if (inputs->password_len == 0) {
+        return PSA_ERROR_BAD_STATE;
+    }
+
+    *password_len = inputs->password_len;
+
+    return PSA_SUCCESS;
+}
+
+psa_status_t psa_crypto_driver_pake_get_password(
+    const psa_crypto_driver_pake_inputs_t *inputs,
+    uint8_t *buffer, size_t buffer_size, size_t *buffer_length)
+{
+    if (inputs->password_len == 0) {
+        return PSA_ERROR_BAD_STATE;
+    }
+
+    if (buffer_size < inputs->password_len) {
+        return PSA_ERROR_BUFFER_TOO_SMALL;
+    }
+
+    memcpy(buffer, inputs->password, inputs->password_len);
+    *buffer_length = inputs->password_len;
+
+    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_cipher_suite(
+    const psa_crypto_driver_pake_inputs_t *inputs,
+    psa_pake_cipher_suite_t *cipher_suite)
+{
+    if (inputs->cipher_suite.algorithm == PSA_ALG_NONE) {
+        return PSA_ERROR_BAD_STATE;
+    }
+
+    *cipher_suite = inputs->cipher_suite;
+
+    return PSA_SUCCESS;
+}
+
 psa_status_t psa_pake_setup(
     psa_pake_operation_t *operation,
     const psa_pake_cipher_suite_t *cipher_suite)
diff --git a/tests/suites/test_suite_psa_crypto_pake.data b/tests/suites/test_suite_psa_crypto_pake.data
index e4bb92b..3be249f 100644
--- a/tests/suites/test_suite_psa_crypto_pake.data
+++ b/tests/suites/test_suite_psa_crypto_pake.data
@@ -193,3 +193,15 @@
 PSA PAKE: ecjpake size macros
 depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256
 ecjpake_size_macros:
+
+PSA PAKE: input getters: ok #1
+pake_input_getters:"aabbccddee":PSA_PAKE_ROLE_SERVER:5:PSA_ALG_JPAKE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_SUCCESS:PSA_SUCCESS:PSA_SUCCESS:PSA_SUCCESS
+
+PSA PAKE: input getters: ok #2
+pake_input_getters:"ddccbbaa":PSA_PAKE_ROLE_CLIENT:5:PSA_ALG_JPAKE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_512:PSA_SUCCESS:PSA_SUCCESS:PSA_SUCCESS:PSA_SUCCESS
+
+PSA PAKE: input getters: buffer for password to small
+pake_input_getters:"aabbccddee":PSA_PAKE_ROLE_SERVER:4:PSA_ALG_JPAKE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_ERROR_BUFFER_TOO_SMALL:PSA_SUCCESS:PSA_SUCCESS:PSA_SUCCESS
+
+PSA PAKE: input getters: inputs not ready
+pake_input_getters:"":0:5:0:0:0:PSA_ERROR_BAD_STATE:PSA_ERROR_BAD_STATE:PSA_ERROR_BAD_STATE:PSA_ERROR_BAD_STATE
diff --git a/tests/suites/test_suite_psa_crypto_pake.function b/tests/suites/test_suite_psa_crypto_pake.function
index 4dffa3b..5af41f7 100644
--- a/tests/suites/test_suite_psa_crypto_pake.function
+++ b/tests/suites/test_suite_psa_crypto_pake.function
@@ -728,6 +728,7 @@
     psa_set_key_usage_flags(&attributes, PSA_KEY_USAGE_DERIVE);
     psa_set_key_algorithm(&attributes, alg);
     psa_set_key_type(&attributes, PSA_KEY_TYPE_PASSWORD);
+
     PSA_ASSERT(psa_import_key(&attributes, pw_data->x, pw_data->len,
                               &key));
 
@@ -905,3 +906,77 @@
               PSA_PAKE_INPUT_MAX_SIZE);
 }
 /* END_CASE */
+
+/* BEGIN_CASE */
+void pake_input_getters(data_t *password, int role_arg, int password_buffer_size,
+                        int alg_arg, int primitive_arg, int hash_arg,
+                        int expected_status_pass,  int expected_status_pass_len,
+                        int expected_status_role,  int expected_status_cs)
+{
+    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 = role_arg;
+    psa_algorithm_t alg = alg_arg;
+    mbedtls_svc_key_id_t key = MBEDTLS_SVC_KEY_ID_INIT;
+    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
+    psa_pake_role_t role_ret = PSA_PAKE_ROLE_NONE;
+    uint8_t password_ret[20] = { 0 }; // max key length is 20 bytes
+    size_t password_len_ret = 0;
+    psa_pake_cipher_suite_t cipher_suite_ret = psa_pake_cipher_suite_init();
+    size_t buffer_len_ret = 0;
+
+    PSA_INIT();
+
+    /* alg equal to 0 indicates case when inputs are not set yet. */
+    if (alg != 0) {
+        psa_pake_cs_set_algorithm(&cipher_suite, alg);
+        psa_pake_cs_set_primitive(&cipher_suite, primitive_arg);
+        psa_pake_cs_set_hash(&cipher_suite, hash_arg);
+
+        psa_set_key_usage_flags(&attributes, PSA_KEY_USAGE_DERIVE);
+        psa_set_key_algorithm(&attributes, alg);
+        psa_set_key_type(&attributes, PSA_KEY_TYPE_PASSWORD);
+
+        PSA_ASSERT(psa_import_key(&attributes, password->x, password->len, &key));
+
+        PSA_ASSERT(psa_pake_setup(&operation, &cipher_suite));
+        PSA_ASSERT(psa_pake_set_password_key(&operation, key));
+        PSA_ASSERT(psa_pake_set_role(&operation, role));
+    }
+
+    TEST_EQUAL(psa_crypto_driver_pake_get_password_len(&operation.data.inputs, &password_len_ret),
+               expected_status_pass_len);
+
+    TEST_EQUAL(psa_crypto_driver_pake_get_password(&operation.data.inputs,
+                                                   (uint8_t *) &password_ret,
+                                                   password_buffer_size, &buffer_len_ret),
+               expected_status_pass);
+
+    TEST_EQUAL(psa_crypto_driver_pake_get_role(&operation.data.inputs, &role_ret),
+               expected_status_role);
+
+    TEST_EQUAL(psa_crypto_driver_pake_get_cipher_suite(&operation.data.inputs, &cipher_suite_ret),
+               expected_status_cs);
+
+    if (expected_status_pass_len == PSA_SUCCESS) {
+        TEST_EQUAL(password_len_ret, password->len);
+    }
+
+    if (expected_status_pass == PSA_SUCCESS) {
+        PSA_ASSERT(memcmp(password_ret, password->x, password->len));
+    }
+
+    if (expected_status_role == PSA_SUCCESS) {
+        TEST_EQUAL(role_ret, role);
+    }
+
+    if (expected_status_pass == PSA_SUCCESS) {
+        PSA_ASSERT(memcmp(&cipher_suite_ret, &cipher_suite, sizeof(cipher_suite)));
+    }
+
+exit:
+    PSA_ASSERT(psa_destroy_key(key));
+    PSA_ASSERT(psa_pake_abort(&operation));
+    PSA_DONE();
+}
+/* END_CASE */