Merge pull request #6608 from mprse/ecjpake_password_fix

Make a copy of the password key in operation object while setting j-pake password
diff --git a/include/psa/crypto_extra.h b/include/psa/crypto_extra.h
index 4f65398..33e2e77 100644
--- a/include/psa/crypto_extra.h
+++ b/include/psa/crypto_extra.h
@@ -1829,7 +1829,7 @@
  */
 #if defined(MBEDTLS_PSA_BUILTIN_PAKE)
 #define PSA_PAKE_OPERATION_INIT {PSA_ALG_NONE, 0, 0, 0, 0,              \
-                                 MBEDTLS_SVC_KEY_ID_INIT,               \
+                                 NULL, 0                ,               \
                                  PSA_PAKE_ROLE_NONE, {0}, 0, 0,         \
                                  {.dummy = 0}}
 #else
@@ -1920,7 +1920,8 @@
 #if defined(MBEDTLS_PSA_BUILTIN_PAKE)
     unsigned int MBEDTLS_PRIVATE(input_step);
     unsigned int MBEDTLS_PRIVATE(output_step);
-    mbedtls_svc_key_id_t MBEDTLS_PRIVATE(password);
+    uint8_t* MBEDTLS_PRIVATE(password);
+    size_t MBEDTLS_PRIVATE(password_len);
     psa_pake_role_t MBEDTLS_PRIVATE(role);
     uint8_t MBEDTLS_PRIVATE(buffer[MBEDTLS_PSA_PAKE_BUFFER_SIZE]);
     size_t MBEDTLS_PRIVATE(buffer_length);
diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 2ce5e43..8c9deff 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -877,20 +877,7 @@
     return( PSA_SUCCESS );
 }
 
-/** Get the description of a key given its identifier and policy constraints
- *  and lock it.
- *
- * The key must have allow all the usage flags set in \p usage. If \p alg is
- * nonzero, the key must allow operations with this algorithm. If \p alg is
- * zero, the algorithm is not checked.
- *
- * In case of a persistent key, the function loads the description of the key
- * into a key slot if not already done.
- *
- * On success, the returned key slot is locked. It is the responsibility of
- * the caller to unlock the key slot when it does not access it anymore.
- */
-static psa_status_t psa_get_and_lock_key_slot_with_policy(
+psa_status_t psa_get_and_lock_key_slot_with_policy(
     mbedtls_svc_key_id_t key,
     psa_key_slot_t **p_slot,
     psa_key_usage_t usage,
diff --git a/library/psa_crypto_core.h b/library/psa_crypto_core.h
index 9863848..5cefa27 100644
--- a/library/psa_crypto_core.h
+++ b/library/psa_crypto_core.h
@@ -183,6 +183,24 @@
 }
 #endif
 
+/** Get the description of a key given its identifier and policy constraints
+ *  and lock it.
+ *
+ * The key must have allow all the usage flags set in \p usage. If \p alg is
+ * nonzero, the key must allow operations with this algorithm. If \p alg is
+ * zero, the algorithm is not checked.
+ *
+ * In case of a persistent key, the function loads the description of the key
+ * into a key slot if not already done.
+ *
+ * On success, the returned key slot is locked. It is the responsibility of
+ * the caller to unlock the key slot when it does not access it anymore.
+ */
+psa_status_t psa_get_and_lock_key_slot_with_policy( mbedtls_svc_key_id_t key,
+                                                    psa_key_slot_t **p_slot,
+                                                    psa_key_usage_t usage,
+                                                    psa_algorithm_t alg );
+
 /** Completely wipe a slot in memory, including its policy.
  *
  * Persistent storage is not affected.
diff --git a/library/psa_crypto_pake.c b/library/psa_crypto_pake.c
index 870b5b5..659b712 100644
--- a/library/psa_crypto_pake.c
+++ b/library/psa_crypto_pake.c
@@ -248,6 +248,7 @@
     psa_key_attributes_t attributes = psa_key_attributes_init();
     psa_key_type_t type;
     psa_key_usage_t usage;
+    psa_key_slot_t *slot = NULL;
 
     if( operation->alg == PSA_ALG_NONE ||
         operation->state != PSA_PAKE_STATE_SETUP )
@@ -273,7 +274,27 @@
     if( ( usage & PSA_KEY_USAGE_DERIVE ) == 0 )
         return( PSA_ERROR_NOT_PERMITTED );
 
-    operation->password = password;
+    if( operation->password != NULL )
+        return( PSA_ERROR_BAD_STATE );
+
+    status = psa_get_and_lock_key_slot_with_policy( password, &slot,
+                                                    PSA_KEY_USAGE_DERIVE,
+                                                    PSA_ALG_JPAKE );
+    if( status != PSA_SUCCESS )
+        return( status );
+
+    operation->password = mbedtls_calloc( 1, slot->key.bytes );
+    if( operation->password == NULL )
+    {
+        psa_unlock_key_slot( slot );
+        return( PSA_ERROR_INSUFFICIENT_MEMORY );
+    }
+    memcpy( operation->password, slot->key.data, slot->key.bytes );
+    operation->password_len = slot->key.bytes;
+
+    status = psa_unlock_key_slot( slot );
+    if( status != PSA_SUCCESS )
+        return( status );
 
     return( PSA_SUCCESS );
 }
@@ -348,9 +369,7 @@
 static psa_status_t psa_pake_ecjpake_setup( psa_pake_operation_t *operation )
 {
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
-    psa_status_t status = PSA_ERROR_CORRUPTION_DETECTED;
     mbedtls_ecjpake_role role;
-    psa_key_slot_t *slot = NULL;
 
     if( operation->role == PSA_PAKE_ROLE_CLIENT )
         role = MBEDTLS_ECJPAKE_CLIENT;
@@ -359,22 +378,20 @@
     else
         return( PSA_ERROR_BAD_STATE );
 
-    if( psa_is_valid_key_id( operation->password, 1 ) == 0 )
+    if( operation->password_len == 0 )
         return( PSA_ERROR_BAD_STATE );
 
-    status = psa_get_and_lock_key_slot( operation->password, &slot );
-    if( status != PSA_SUCCESS )
-        return( status );
-
-
     ret = mbedtls_ecjpake_setup( &operation->ctx.ecjpake,
                                  role,
                                  MBEDTLS_MD_SHA256,
                                  MBEDTLS_ECP_DP_SECP256R1,
-                                 slot->key.data, slot->key.bytes );
+                                 operation->password,
+                                 operation->password_len );
 
-    psa_unlock_key_slot( slot );
-    slot = NULL;
+    mbedtls_platform_zeroize( operation->password, operation->password_len );
+    mbedtls_free( operation->password );
+    operation->password = NULL;
+    operation->password_len = 0;
 
     if( ret != 0 )
         return( mbedtls_ecjpake_to_psa_error( ret ) );
@@ -840,7 +857,11 @@
     {
         operation->input_step = PSA_PAKE_STEP_INVALID;
         operation->output_step = PSA_PAKE_STEP_INVALID;
-        operation->password = MBEDTLS_SVC_KEY_ID_INIT;
+        if( operation->password_len > 0 )
+            mbedtls_platform_zeroize( operation->password, operation->password_len );
+        mbedtls_free( operation->password );
+        operation->password = NULL;
+        operation->password_len = 0;
         operation->role = PSA_PAKE_ROLE_NONE;
         mbedtls_platform_zeroize( operation->buffer, MBEDTLS_PSA_PAKE_BUFFER_SIZE );
         operation->buffer_length = 0;
diff --git a/tests/suites/test_suite_psa_crypto.data b/tests/suites/test_suite_psa_crypto.data
index cce3fd0..659205d 100644
--- a/tests/suites/test_suite_psa_crypto.data
+++ b/tests/suites/test_suite_psa_crypto.data
@@ -6549,11 +6549,16 @@
 
 PSA PAKE: ecjpake rounds
 depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256:PSA_WANT_ALG_TLS12_PSK_TO_MS
-ecjpake_rounds:PSA_ALG_JPAKE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_ALG_TLS12_PSK_TO_MS(PSA_ALG_SHA_256):"abcdef":0
+ecjpake_rounds:PSA_ALG_JPAKE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_ALG_TLS12_PSK_TO_MS(PSA_ALG_SHA_256):"abcdef":0:0
 
 PSA PAKE: ecjpake rounds, client input first
 depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256:PSA_WANT_ALG_TLS12_PSK_TO_MS
-ecjpake_rounds:PSA_ALG_JPAKE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_ALG_TLS12_PSK_TO_MS(PSA_ALG_SHA_256):"abcdef":1
+ecjpake_rounds:PSA_ALG_JPAKE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_ALG_TLS12_PSK_TO_MS(PSA_ALG_SHA_256):"abcdef":1:0
+
+# This test case relies on implementation (it may need to be adjusted in the future)
+PSA PAKE: ecjpake rounds - key is destroyed after being passed to set_password_key
+depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256:PSA_WANT_ALG_TLS12_PSK_TO_MS
+ecjpake_rounds:PSA_ALG_JPAKE:PSA_PAKE_PRIMITIVE(PSA_PAKE_PRIMITIVE_TYPE_ECC, PSA_ECC_FAMILY_SECP_R1, 256):PSA_ALG_SHA_256:PSA_ALG_TLS12_PSK_TO_MS(PSA_ALG_SHA_256):"abcdef":0:1
 
 PSA PAKE: ecjpake no input errors
 depends_on:PSA_WANT_KEY_TYPE_ECC_KEY_PAIR:PSA_WANT_ECC_SECP_R1_256:PSA_WANT_ALG_SHA_256
diff --git a/tests/suites/test_suite_psa_crypto.function b/tests/suites/test_suite_psa_crypto.function
index 779f594..ca1614b 100644
--- a/tests/suites/test_suite_psa_crypto.function
+++ b/tests/suites/test_suite_psa_crypto.function
@@ -31,6 +31,27 @@
 #define ASSERT_OPERATION_IS_ACTIVE(   operation ) TEST_ASSERT( operation.id != 0 )
 #define ASSERT_OPERATION_IS_INACTIVE( operation ) TEST_ASSERT( operation.id == 0 )
 
+#if defined(PSA_WANT_ALG_JPAKE)
+int ecjpake_operation_setup( psa_pake_operation_t *operation,
+                              psa_pake_cipher_suite_t *cipher_suite,
+                              psa_pake_role_t role,
+                              mbedtls_svc_key_id_t key,
+                              size_t key_available )
+{
+    PSA_ASSERT( psa_pake_abort( operation ) );
+
+    PSA_ASSERT( psa_pake_setup( operation, cipher_suite ) );
+
+    PSA_ASSERT( psa_pake_set_role( operation, role) );
+
+    if( key_available )
+        PSA_ASSERT( psa_pake_set_password_key( operation, key ) );
+    return 0;
+exit:
+    return 1;
+}
+#endif
+
 /** An invalid export length that will never be set by psa_export_key(). */
 static const size_t INVALID_EXPORT_LENGTH = ~0U;
 
@@ -8740,7 +8761,6 @@
 {
     psa_pake_cipher_suite_t cipher_suite = psa_pake_cipher_suite_init();
     psa_pake_operation_t operation = psa_pake_operation_init();
-    psa_pake_operation_t op_copy = psa_pake_operation_init();
     psa_algorithm_t alg = alg_arg;
     psa_pake_primitive_t primitive = primitive_arg;
     psa_key_type_t key_type_pw = key_type_pw_arg;
@@ -8839,22 +8859,25 @@
     if( input_first )
     {
         /* Invalid parameters (input) */
-        op_copy = operation;
-        TEST_EQUAL( psa_pake_input( &op_copy, PSA_PAKE_STEP_ZK_PROOF,
+        TEST_EQUAL( psa_pake_input( &operation, PSA_PAKE_STEP_ZK_PROOF,
                                     NULL, 0 ),
                     PSA_ERROR_INVALID_ARGUMENT );
         /* Invalid parameters (step) */
-        op_copy = operation;
-        TEST_EQUAL( psa_pake_input( &op_copy, PSA_PAKE_STEP_ZK_PROOF + 10,
+        TEST_EQUAL( ecjpake_operation_setup( &operation, &cipher_suite, role,
+                                 key, pw_data->len ) , 0 );
+        TEST_EQUAL( psa_pake_input( &operation, PSA_PAKE_STEP_ZK_PROOF + 10,
                                     output_buffer, size_zk_proof ),
                     PSA_ERROR_INVALID_ARGUMENT );
         /* Invalid first step */
-        op_copy = operation;
-        TEST_EQUAL( psa_pake_input( &op_copy, PSA_PAKE_STEP_ZK_PROOF,
+        TEST_EQUAL( ecjpake_operation_setup( &operation, &cipher_suite, role,
+                                 key, pw_data->len ), 0 );
+        TEST_EQUAL( psa_pake_input( &operation, PSA_PAKE_STEP_ZK_PROOF,
                                     output_buffer, size_zk_proof ),
                     PSA_ERROR_BAD_STATE );
 
         /* Possibly valid */
+        TEST_EQUAL( ecjpake_operation_setup( &operation, &cipher_suite, role,
+                                 key, pw_data->len ), 0 );
         TEST_EQUAL( psa_pake_input( &operation, PSA_PAKE_STEP_KEY_SHARE,
                                     output_buffer, size_key_share ),
                     expected_status_input_output);
@@ -8875,22 +8898,25 @@
     else
     {
         /* Invalid parameters (output) */
-        op_copy = operation;
-        TEST_EQUAL( psa_pake_output( &op_copy, PSA_PAKE_STEP_ZK_PROOF,
+        TEST_EQUAL( psa_pake_output( &operation, PSA_PAKE_STEP_ZK_PROOF,
                                      NULL, 0, NULL ),
                     PSA_ERROR_INVALID_ARGUMENT );
-        op_copy = operation;
         /* Invalid parameters (step) */
-        TEST_EQUAL( psa_pake_output( &op_copy, PSA_PAKE_STEP_ZK_PROOF + 10,
+        TEST_EQUAL( ecjpake_operation_setup( &operation, &cipher_suite, role,
+                                key, pw_data->len ), 0 );
+        TEST_EQUAL( psa_pake_output( &operation, PSA_PAKE_STEP_ZK_PROOF + 10,
                                      output_buffer, buf_size, &output_len ),
                     PSA_ERROR_INVALID_ARGUMENT );
         /* Invalid first step */
-        op_copy = operation;
-        TEST_EQUAL( psa_pake_output( &op_copy, PSA_PAKE_STEP_ZK_PROOF,
+        TEST_EQUAL( ecjpake_operation_setup( &operation, &cipher_suite, role,
+                                 key, pw_data->len ), 0 );
+        TEST_EQUAL( psa_pake_output( &operation, PSA_PAKE_STEP_ZK_PROOF,
                                      output_buffer, buf_size, &output_len ),
                     PSA_ERROR_BAD_STATE );
 
         /* Possibly valid */
+        TEST_EQUAL( ecjpake_operation_setup( &operation, &cipher_suite, role,
+                                 key, pw_data->len ), 0 );
         TEST_EQUAL( psa_pake_output( &operation, PSA_PAKE_STEP_KEY_SHARE,
                                      output_buffer, buf_size, &output_len ),
                     expected_status_input_output );
@@ -8974,7 +9000,7 @@
 /* BEGIN_CASE depends_on:PSA_WANT_ALG_JPAKE */
 void ecjpake_rounds( int alg_arg, int primitive_arg, int hash_arg,
                      int derive_alg_arg, data_t *pw_data,
-                     int client_input_first )
+                     int client_input_first, int destroy_key )
 {
     psa_pake_cipher_suite_t cipher_suite = psa_pake_cipher_suite_init();
     psa_pake_operation_t server = psa_pake_operation_init();
@@ -9025,6 +9051,9 @@
     PSA_ASSERT( psa_pake_set_password_key( &server, key ) );
     PSA_ASSERT( psa_pake_set_password_key( &client, key ) );
 
+    if( destroy_key == 1 )
+        psa_destroy_key( key );
+
     TEST_EQUAL( psa_pake_get_implicit_key( &server, &server_derive ),
                 PSA_ERROR_BAD_STATE );
     TEST_EQUAL( psa_pake_get_implicit_key( &client, &client_derive ),