Make a copy of the key in operation while setting pake password

Additionally use psa_get_and_lock_key_slot_with_policy() to obtain key.
This requires making this function public. This will have to be solved while adding driver dipatch for EC-JPAKE.

Signed-off-by: Przemek Stekiel <przemyslaw.stekiel@mobica.com>

include/psa/crypto_extra.h

@@ -1829,7 +1829,7 @@ psa_status_t psa_pake_abort( psa_pake_operation_t * operation );
  */
 #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 @@ struct psa_pake_operation_s
 #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_data);
+    size_t MBEDTLS_PRIVATE(password_bytes);
     psa_pake_role_t MBEDTLS_PRIVATE(role);
     uint8_t MBEDTLS_PRIVATE(buffer[MBEDTLS_PSA_PAKE_BUFFER_SIZE]);
     size_t MBEDTLS_PRIVATE(buffer_length);

library/psa_crypto.c

@@ -890,7 +890,7 @@ static psa_status_t psa_restrict_key_policy(
  * 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,

library/psa_crypto_pake.c

@@ -33,6 +33,11 @@
 #include <mbedtls/error.h>
 #include <string.h>
 
+extern 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 );
 /*
  * State sequence:
  *
@@ -248,6 +253,7 @@ psa_status_t psa_pake_set_password_key( psa_pake_operation_t *operation,
     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 )
@@ -255,6 +261,9 @@ psa_status_t psa_pake_set_password_key( psa_pake_operation_t *operation,
         return( PSA_ERROR_BAD_STATE );
     }
 
+    if( psa_is_valid_key_id( password, 1 ) == 0 )
+        return( PSA_ERROR_BAD_STATE );
+
     status = psa_get_key_attributes( password, &attributes );
     if( status != PSA_SUCCESS )
         return( status );
@@ -273,7 +282,33 @@ psa_status_t psa_pake_set_password_key( psa_pake_operation_t *operation,
     if( ( usage & PSA_KEY_USAGE_DERIVE ) == 0 )
         return( PSA_ERROR_NOT_PERMITTED );
 
-    operation->password = password;
+    status = psa_get_and_lock_key_slot_with_policy( password, &slot,
+                                                    PSA_KEY_USAGE_DERIVE,
+                                                    PSA_ALG_JPAKE );
+    if( status != PSA_SUCCESS )
+        return( status );
+
+    if( slot->key.data == NULL || slot->key.bytes == 0 )
+        return( PSA_ERROR_INVALID_ARGUMENT );
+
+    if( operation->password_data != NULL )
+    {
+        mbedtls_free( operation->password_data );
+        operation->password_bytes = 0;
+    }
+
+    operation->password_data = mbedtls_calloc( 1, slot->key.bytes );
+    if( operation->password_data == NULL )
+    {
+        status = psa_unlock_key_slot( slot );
+        return( PSA_ERROR_INSUFFICIENT_MEMORY );
+    }
+    memcpy( operation->password_data, slot->key.data, slot->key.bytes );
+    operation->password_bytes = slot->key.bytes;
+
+    status = psa_unlock_key_slot( slot );
+    if( status != PSA_SUCCESS )
+        return( status );
 
     return( PSA_SUCCESS );
 }
@@ -348,9 +383,7 @@ psa_status_t psa_pake_set_role( psa_pake_operation_t *operation,
 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 +392,18 @@ static psa_status_t psa_pake_ecjpake_setup( psa_pake_operation_t *operation )
     else
         return( PSA_ERROR_BAD_STATE );
 
-    if( psa_is_valid_key_id( operation->password, 1 ) == 0 )
+    if (operation->password_data == NULL ||
+        operation->password_bytes == 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_data,
-
+                                 operation->password_bytes );
-    psa_unlock_key_slot( slot );
-    slot = NULL;
 
     if( ret != 0 )
         return( mbedtls_ecjpake_to_psa_error( ret ) );
@@ -840,7 +869,9 @@ psa_status_t psa_pake_abort(psa_pake_operation_t * operation)
     {
         operation->input_step = PSA_PAKE_STEP_INVALID;
         operation->output_step = PSA_PAKE_STEP_INVALID;
-        operation->password = MBEDTLS_SVC_KEY_ID_INIT;
+        mbedtls_free( operation->password_data );
+        operation->password_data = NULL;
+        operation->password_bytes = 0;
         operation->role = PSA_PAKE_ROLE_NONE;
         mbedtls_platform_zeroize( operation->buffer, MBEDTLS_PSA_PAKE_BUFFER_SIZE );
         operation->buffer_length = 0;