test: pake: add tests for set password functions

Signed-off-by: Valerio Setti <vsetti@baylibre.com>
diff --git a/tests/suites/test_suite_ssl.data b/tests/suites/test_suite_ssl.data
index a35762d..0cec784 100644
--- a/tests/suites/test_suite_ssl.data
+++ b/tests/suites/test_suite_ssl.data
@@ -3567,3 +3567,31 @@
 
 TLS 1.3 srv Certificate msg - wrong vector lengths
 tls13_server_certificate_msg_invalid_vector_len
+
+EC-JPAKE set password
+depends_on:MBEDTLS_KEY_EXCHANGE_ECJPAKE_ENABLED
+ssl_ecjpake_set_password:0:ECJPAKE_ERR_NONE:0
+
+EC-JPAKE set password - uninitiazed SSL context
+depends_on:MBEDTLS_KEY_EXCHANGE_ECJPAKE_ENABLED
+ssl_ecjpake_set_password:0:ECJPAKE_ERR_UNITIALIZED_SSL_CONTEXT:MBEDTLS_ERR_SSL_BAD_INPUT_DATA
+
+EC-JPAKE set password - empty password
+depends_on:MBEDTLS_KEY_EXCHANGE_ECJPAKE_ENABLED
+ssl_ecjpake_set_password:0:ECJPAKE_ERR_EMPTY_PASSWORD:MBEDTLS_ERR_SSL_BAD_INPUT_DATA
+
+EC-JPAKE set opaque password
+depends_on:MBEDTLS_KEY_EXCHANGE_ECJPAKE_ENABLED:MBEDTLS_USE_PSA_CRYPTO
+ssl_ecjpake_set_password:1:ECJPAKE_ERR_NONE:0
+
+EC-JPAKE set opaque password - uninitiazed SSL context
+depends_on:MBEDTLS_KEY_EXCHANGE_ECJPAKE_ENABLED:MBEDTLS_USE_PSA_CRYPTO
+ssl_ecjpake_set_password:1:ECJPAKE_ERR_UNITIALIZED_SSL_CONTEXT:MBEDTLS_ERR_SSL_BAD_INPUT_DATA
+
+EC-JPAKE set opaque password - empty password
+depends_on:MBEDTLS_KEY_EXCHANGE_ECJPAKE_ENABLED:MBEDTLS_USE_PSA_CRYPTO
+ssl_ecjpake_set_password:1:ECJPAKE_ERR_EMPTY_PASSWORD:MBEDTLS_ERR_SSL_BAD_INPUT_DATA
+
+EC-JPAKE set opaque password - uninitalized password key
+depends_on:MBEDTLS_KEY_EXCHANGE_ECJPAKE_ENABLED:MBEDTLS_USE_PSA_CRYPTO
+ssl_ecjpake_set_password:1:ECJPAKE_ERR_UNINITIALIZED_PWD_KEY:MBEDTLS_ERR_SSL_BAD_INPUT_DATA
diff --git a/tests/suites/test_suite_ssl.function b/tests/suites/test_suite_ssl.function
index 7447a1d..bfc6a37 100644
--- a/tests/suites/test_suite_ssl.function
+++ b/tests/suites/test_suite_ssl.function
@@ -2582,6 +2582,15 @@
     return( 0 );
 }
 #endif /* MBEDTLS_TEST_HOOKS */
+
+typedef enum {
+    ECJPAKE_ERR_NONE,
+    ECJPAKE_ERR_UNITIALIZED_SSL_CONTEXT,
+    ECJPAKE_ERR_EMPTY_PASSWORD,
+    ECJPAKE_ERR_UNINITIALIZED_PWD_KEY,
+} ecjpake_err_inj_step_t;
+
+#define ECJPAKE_TEST_PWD        "bla"
 /* END_HEADER */
 
 /* BEGIN_DEPENDENCIES
@@ -6180,3 +6189,77 @@
     USE_PSA_DONE( );
 }
 /* END_CASE */
+
+/* BEGIN_CASE depends_on:MBEDTLS_KEY_EXCHANGE_ECJPAKE_ENABLED */
+void ssl_ecjpake_set_password( int use_opaque_arg,
+                               int err_injection_step_arg,
+                               int expected_error_arg )
+{
+    mbedtls_ssl_context ssl;
+    mbedtls_ssl_config conf;
+#if defined( MBEDTLS_USE_PSA_CRYPTO )
+    mbedtls_svc_key_id_t pwd_slot = MBEDTLS_SVC_KEY_ID_INIT;
+#else   /* MBEDTLS_USE_PSA_CRYPTO */
+    (void) use_opaque_arg;
+#endif  /* MBEDTLS_USE_PSA_CRYPTO */
+    const unsigned char pwd_string[ sizeof(ECJPAKE_TEST_PWD) ] = "";
+    size_t pwd_len = 0;
+    ecjpake_err_inj_step_t err_injection_step = err_injection_step_arg;
+    int ret;
+
+    USE_PSA_INIT( );
+
+    mbedtls_ssl_init( &ssl );
+
+    if( err_injection_step == ECJPAKE_ERR_UNITIALIZED_SSL_CONTEXT )
+        goto run_test;
+
+    mbedtls_ssl_config_init( &conf );
+
+    TEST_ASSERT( mbedtls_ssl_config_defaults( &conf,
+                                              MBEDTLS_SSL_IS_CLIENT,
+                                              MBEDTLS_SSL_TRANSPORT_STREAM,
+                                              MBEDTLS_SSL_PRESET_DEFAULT )
+                 == 0 );
+
+    TEST_ASSERT( mbedtls_ssl_setup( &ssl, &conf ) == 0 );
+
+    if( err_injection_step == ECJPAKE_ERR_EMPTY_PASSWORD )
+        goto run_test;
+
+    pwd_len = strlen( ECJPAKE_TEST_PWD );
+    memcpy( (void*) pwd_string, ECJPAKE_TEST_PWD, pwd_len );
+
+#if defined( MBEDTLS_USE_PSA_CRYPTO )
+    if( use_opaque_arg )
+    {
+        if( err_injection_step == ECJPAKE_ERR_UNINITIALIZED_PWD_KEY )
+            goto run_test;
+
+        psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
+
+        psa_set_key_usage_flags( &attributes, PSA_KEY_USAGE_DERIVE );
+        psa_set_key_algorithm( &attributes, PSA_ALG_JPAKE );
+        psa_set_key_type( &attributes, PSA_KEY_TYPE_PASSWORD );
+
+        TEST_ASSERT( psa_import_key( &attributes, pwd_string,
+                    pwd_len, &pwd_slot ) == PSA_SUCCESS );
+    }
+#endif  /* MBEDTLS_USE_PSA_CRYPTO */
+
+run_test:
+#if defined( MBEDTLS_USE_PSA_CRYPTO )
+    ret = ( use_opaque_arg ) ?
+            mbedtls_ssl_set_hs_ecjpake_password_opaque( &ssl, pwd_slot ) :
+            mbedtls_ssl_set_hs_ecjpake_password( &ssl, pwd_string, pwd_len );
+#else   /* MBEDTLS_USE_PSA_CRYPTO */
+    ret = mbedtls_ssl_set_hs_ecjpake_password( &ssl, pwd_string, pwd_len );
+#endif  /* MBEDTLS_USE_PSA_CRYPTO */
+    TEST_EQUAL( ret, expected_error_arg );
+
+    mbedtls_ssl_free( &ssl );
+    mbedtls_ssl_config_free( &conf );
+
+    USE_PSA_DONE( );
+}
+/* END_CASE */