Use PSA PAKE API when MBEDTLS_USE_PSA_CRYPTO is selected

Signed-off-by: Neil Armstrong <narmstrong@baylibre.com>
Signed-off-by: Valerio Setti <vsetti@baylibre.com>

library/ssl_misc.h

@@ -50,7 +50,8 @@
 #include "mbedtls/sha512.h"
 #endif
 
-#if defined(MBEDTLS_KEY_EXCHANGE_ECJPAKE_ENABLED)
+#if defined(MBEDTLS_KEY_EXCHANGE_ECJPAKE_ENABLED) && \
+    !defined(MBEDTLS_USE_PSA_CRYPTO)
 #include "mbedtls/ecjpake.h"
 #endif
 
@@ -663,7 +664,13 @@ struct mbedtls_ssl_handshake_params
 #endif /* MBEDTLS_ECDH_C || MBEDTLS_ECDSA_C */
 
 #if defined(MBEDTLS_KEY_EXCHANGE_ECJPAKE_ENABLED)
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_pake_operation_t psa_pake_ctx;        /*!< EC J-PAKE key exchange */
+    mbedtls_svc_key_id_t psa_pake_password;
+    uint8_t psa_pake_ctx_is_ok;
+#else
     mbedtls_ecjpake_context ecjpake_ctx;        /*!< EC J-PAKE key exchange */
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 #if defined(MBEDTLS_SSL_CLI_C)
     unsigned char *ecjpake_cache;               /*!< Cache for ClientHello ext */
     size_t ecjpake_cache_len;                   /*!< Length of cached data */

library/ssl_tls.c

@@ -668,7 +668,12 @@ static void ssl_handshake_params_init( mbedtls_ssl_handshake_params *handshake )
     mbedtls_ecdh_init( &handshake->ecdh_ctx );
 #endif
 #if defined(MBEDTLS_KEY_EXCHANGE_ECJPAKE_ENABLED)
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    handshake->psa_pake_ctx = psa_pake_operation_init();
+    handshake->psa_pake_password = MBEDTLS_SVC_KEY_ID_INIT;
+#else
     mbedtls_ecjpake_init( &handshake->ecjpake_ctx );
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 #if defined(MBEDTLS_SSL_CLI_C)
     handshake->ecjpake_cache = NULL;
     handshake->ecjpake_cache_len = 0;
@@ -1615,11 +1620,75 @@ int mbedtls_ssl_set_hs_ecjpake_password( mbedtls_ssl_context *ssl,
                                          const unsigned char *pw,
                                          size_t pw_len )
 {
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_pake_cipher_suite_t cipher_suite = psa_pake_cipher_suite_init();
+    psa_key_attributes_t attributes = PSA_KEY_ATTRIBUTES_INIT;
+    psa_pake_role_t psa_role;
+    psa_status_t status;
+#else
     mbedtls_ecjpake_role role;
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 
     if( ssl->handshake == NULL || ssl->conf == NULL )
         return( MBEDTLS_ERR_SSL_BAD_INPUT_DATA );
 
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    if( ssl->conf->endpoint == MBEDTLS_SSL_IS_SERVER )
+        psa_role = PSA_PAKE_ROLE_SERVER;
+    else
+        psa_role = PSA_PAKE_ROLE_CLIENT;
+
+
+    if( pw_len > 0 )
+    {
+        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 );
+
+        status = psa_import_key( &attributes, pw, pw_len,
+                                 &ssl->handshake->psa_pake_password );
+        if( status != PSA_SUCCESS )
+            return( MBEDTLS_ERR_SSL_HW_ACCEL_FAILED );
+    }
+
+    psa_pake_cs_set_algorithm( &cipher_suite, PSA_ALG_JPAKE );
+    psa_pake_cs_set_primitive( &cipher_suite,
+                               PSA_PAKE_PRIMITIVE( PSA_PAKE_PRIMITIVE_TYPE_ECC,
+                                                   PSA_ECC_FAMILY_SECP_R1,
+                                                   256) );
+    psa_pake_cs_set_hash( &cipher_suite, PSA_ALG_SHA_256 );
+
+    status = psa_pake_setup( &ssl->handshake->psa_pake_ctx, &cipher_suite );
+    if( status != PSA_SUCCESS )
+    {
+        psa_destroy_key( ssl->handshake->psa_pake_password );
+        return( MBEDTLS_ERR_SSL_HW_ACCEL_FAILED );
+    }
+
+    status = psa_pake_set_role( &ssl->handshake->psa_pake_ctx, psa_role );
+    if( status != PSA_SUCCESS )
+    {
+        psa_destroy_key( ssl->handshake->psa_pake_password );
+        psa_pake_abort( &ssl->handshake->psa_pake_ctx );
+        return( MBEDTLS_ERR_SSL_HW_ACCEL_FAILED );
+    }
+
+    if( pw_len > 0 )
+    {
+        psa_pake_set_password_key( &ssl->handshake->psa_pake_ctx,
+                                   ssl->handshake->psa_pake_password );
+        if( status != PSA_SUCCESS )
+        {
+            psa_destroy_key( ssl->handshake->psa_pake_password );
+            psa_pake_abort( &ssl->handshake->psa_pake_ctx );
+            return( MBEDTLS_ERR_SSL_HW_ACCEL_FAILED );
+        }
+    }
+
+    ssl->handshake->psa_pake_ctx_is_ok = 1;
+
+    return( 0 );
+#else
     if( ssl->conf->endpoint == MBEDTLS_SSL_IS_SERVER )
         role = MBEDTLS_ECJPAKE_SERVER;
     else
@@ -1630,6 +1699,7 @@ int mbedtls_ssl_set_hs_ecjpake_password( mbedtls_ssl_context *ssl,
                                    MBEDTLS_MD_SHA256,
                                    MBEDTLS_ECP_DP_SECP256R1,
                                    pw, pw_len ) );
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 }
 #endif /* MBEDTLS_KEY_EXCHANGE_ECJPAKE_ENABLED */
 
@@ -3665,7 +3735,13 @@ void mbedtls_ssl_handshake_free( mbedtls_ssl_context *ssl )
     mbedtls_ecdh_free( &handshake->ecdh_ctx );
 #endif
 #if defined(MBEDTLS_KEY_EXCHANGE_ECJPAKE_ENABLED)
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_pake_abort( &handshake->psa_pake_ctx );
+    psa_destroy_key( handshake->psa_pake_password );
+    handshake->psa_pake_password = MBEDTLS_SVC_KEY_ID_INIT;
+#else
     mbedtls_ecjpake_free( &handshake->ecjpake_ctx );
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 #if defined(MBEDTLS_SSL_CLI_C)
     mbedtls_free( handshake->ecjpake_cache );
     handshake->ecjpake_cache = NULL;
@@ -5879,6 +5955,55 @@ static int ssl_compute_master( mbedtls_ssl_handshake_params *handshake,
     else
 #endif
     {
+#if defined(MBEDTLS_USE_PSA_CRYPTO) &&                              \
+    defined(MBEDTLS_KEY_EXCHANGE_ECJPAKE_ENABLED)
+        if( handshake->ciphersuite_info->key_exchange == MBEDTLS_KEY_EXCHANGE_ECJPAKE )
+        {
+            psa_status_t status;
+            psa_algorithm_t alg = PSA_ALG_TLS12_ECJPAKE_TO_PMS;
+            psa_key_derivation_operation_t derivation =
+                PSA_KEY_DERIVATION_OPERATION_INIT;
+
+            MBEDTLS_SSL_DEBUG_MSG( 2, ( "perform PSA-based PMS KDF for ECJPAKE" ) );
+
+            handshake->pmslen = PSA_TLS12_ECJPAKE_TO_PMS_DATA_SIZE;
+
+            status = psa_key_derivation_setup( &derivation, alg );
+            if( status != PSA_SUCCESS )
+                return( MBEDTLS_ERR_SSL_HW_ACCEL_FAILED );
+
+            status = psa_key_derivation_set_capacity( &derivation,
+                                            PSA_TLS12_ECJPAKE_TO_PMS_DATA_SIZE );
+            if( status != PSA_SUCCESS )
+            {
+                psa_key_derivation_abort( &derivation );
+                return( MBEDTLS_ERR_SSL_HW_ACCEL_FAILED );
+            }
+
+            status = psa_pake_get_implicit_key( &handshake->psa_pake_ctx,
+                                                &derivation );
+            if( status != PSA_SUCCESS )
+            {
+                psa_key_derivation_abort( &derivation );
+                return( MBEDTLS_ERR_SSL_HW_ACCEL_FAILED );
+            }
+
+            status = psa_key_derivation_output_bytes( &derivation,
+                                                      handshake->premaster,
+                                                      handshake->pmslen );
+            if( status != PSA_SUCCESS )
+            {
+                psa_key_derivation_abort( &derivation );
+                return( MBEDTLS_ERR_SSL_HW_ACCEL_FAILED );
+            }
+
+            status = psa_key_derivation_abort( &derivation );
+            if( status != PSA_SUCCESS )
+            {
+                return( MBEDTLS_ERR_SSL_HW_ACCEL_FAILED );
+            }
+        }
+#endif
         ret = handshake->tls_prf( handshake->premaster, handshake->pmslen,
                                   lbl, seed, seed_len,
                                   master,
@@ -5917,6 +6042,7 @@ int mbedtls_ssl_derive_keys( mbedtls_ssl_context *ssl )
         return( ret );
     }
 
+
     /* Compute master secret if needed */
     ret = ssl_compute_master( ssl->handshake,
                               ssl->session_negotiate->master,
@@ -8620,8 +8746,13 @@ int mbedtls_ssl_validate_ciphersuite(
 
 #if defined(MBEDTLS_SSL_PROTO_TLS1_2) && defined(MBEDTLS_SSL_CLI_C)
 #if defined(MBEDTLS_KEY_EXCHANGE_ECJPAKE_ENABLED)
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    if( suite_info->key_exchange == MBEDTLS_KEY_EXCHANGE_ECJPAKE &&
+        ssl->handshake->psa_pake_ctx_is_ok != 1 )
+#else
     if( suite_info->key_exchange == MBEDTLS_KEY_EXCHANGE_ECJPAKE &&
         mbedtls_ecjpake_check( &ssl->handshake->ecjpake_ctx ) != 0 )
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
     {
         return( -1 );
     }

library/ssl_tls12_client.c

@@ -130,15 +130,24 @@ static int ssl_write_ecjpake_kkpp_ext( mbedtls_ssl_context *ssl,
                                        const unsigned char *end,
                                        size_t *olen )
 {
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_status_t status;
+#else
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
     unsigned char *p = buf;
     size_t kkpp_len;
 
     *olen = 0;
 
     /* Skip costly extension if we can't use EC J-PAKE anyway */
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    if( ssl->handshake->psa_pake_ctx_is_ok != 1 )
+        return( 0 );
+#else
     if( mbedtls_ecjpake_check( &ssl->handshake->ecjpake_ctx ) != 0 )
         return( 0 );
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 
     MBEDTLS_SSL_DEBUG_MSG( 3,
         ( "client hello, adding ecjpake_kkpp extension" ) );
@@ -158,6 +167,43 @@ static int ssl_write_ecjpake_kkpp_ext( mbedtls_ssl_context *ssl,
     {
         MBEDTLS_SSL_DEBUG_MSG( 3, ( "generating new ecjpake parameters" ) );
 
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+        size_t output_offset = 0;
+        size_t output_len;
+
+        /* Repeat the KEY_SHARE, ZK_PUBLIC & ZF_PROOF twice */
+        for( unsigned int x = 1 ; x <= 2 ; ++x )
+        {
+            for( psa_pake_step_t step = PSA_PAKE_STEP_KEY_SHARE ;
+                 step <= PSA_PAKE_STEP_ZK_PROOF ;
+                 ++step )
+            {
+                /* For each step, prepend 1 byte with the length of the data */
+                if (step != PSA_PAKE_STEP_ZK_PROOF) {
+                    *(p + 2 + output_offset) = 65;
+                } else {
+                    *(p + 2 + output_offset) = 32;
+                }
+                output_offset += 1;
+
+                status = psa_pake_output( &ssl->handshake->psa_pake_ctx,
+                                          step, p + 2 + output_offset,
+                                          end - p - output_offset - 2,
+                                          &output_len );
+                if( status != PSA_SUCCESS )
+                {
+                    psa_destroy_key( ssl->handshake->psa_pake_password );
+                    psa_pake_abort( &ssl->handshake->psa_pake_ctx );
+                    MBEDTLS_SSL_DEBUG_RET( 1 , "psa_pake_output", status );
+                    return( psa_ssl_status_to_mbedtls( status ) );
+                }
+
+                output_offset += output_len;
+            }
+        }
+
+        kkpp_len = output_offset;
+#else
         ret = mbedtls_ecjpake_write_round_one( &ssl->handshake->ecjpake_ctx,
                                                p + 2, end - p - 2, &kkpp_len,
                                                ssl->conf->f_rng, ssl->conf->p_rng );
@@ -167,6 +213,7 @@ static int ssl_write_ecjpake_kkpp_ext( mbedtls_ssl_context *ssl,
                 "mbedtls_ecjpake_write_round_one", ret );
             return( ret );
         }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 
         ssl->handshake->ecjpake_cache = mbedtls_calloc( 1, kkpp_len );
         if( ssl->handshake->ecjpake_cache == NULL )
@@ -849,10 +896,11 @@ static int ssl_parse_supported_point_formats_ext( mbedtls_ssl_context *ssl,
             ssl->handshake->ecdh_ctx.point_format = p[0];
 #endif /* !MBEDTLS_USE_PSA_CRYPTO &&
           ( MBEDTLS_ECDH_C || MBEDTLS_ECDSA_C ) */
-#if defined(MBEDTLS_KEY_EXCHANGE_ECJPAKE_ENABLED)
+#if !defined(MBEDTLS_USE_PSA_CRYPTO) &&                             \
+    defined(MBEDTLS_KEY_EXCHANGE_ECJPAKE_ENABLED)
             mbedtls_ecjpake_set_point_format( &ssl->handshake->ecjpake_ctx,
                                               p[0] );
-#endif
+#endif /* !MBEDTLS_USE_PSA_CRYPTO && MBEDTLS_KEY_EXCHANGE_ECJPAKE_ENABLED */
             MBEDTLS_SSL_DEBUG_MSG( 4, ( "point format selected: %d", p[0] ) );
             return( 0 );
         }
@@ -876,6 +924,9 @@ static int ssl_parse_ecjpake_kkpp( mbedtls_ssl_context *ssl,
                                    size_t len )
 {
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_status_t status;
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 
     if( ssl->handshake->ciphersuite_info->key_exchange !=
         MBEDTLS_KEY_EXCHANGE_ECJPAKE )
@@ -889,6 +940,52 @@ static int ssl_parse_ecjpake_kkpp( mbedtls_ssl_context *ssl,
     ssl->handshake->ecjpake_cache = NULL;
     ssl->handshake->ecjpake_cache_len = 0;
 
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    size_t input_offset = 0;
+
+    /* Repeat the KEY_SHARE, ZK_PUBLIC & ZF_PROOF twice */
+    for( unsigned int x = 1 ; x <= 2 ; ++x )
+    {
+        for( psa_pake_step_t step = PSA_PAKE_STEP_KEY_SHARE ;
+             step <= PSA_PAKE_STEP_ZK_PROOF ;
+             ++step )
+        {
+            /* Length is stored at the first byte */
+            size_t length = buf[input_offset];
+            input_offset += 1;
+
+            if( input_offset + length > len )
+            {
+                ret = MBEDTLS_ERR_SSL_HANDSHAKE_FAILURE;
+                goto psa_pake_error;
+            }
+
+            status = psa_pake_input( &ssl->handshake->psa_pake_ctx, step,
+                                     buf + input_offset, length );
+            if( status != PSA_SUCCESS)
+            {
+                ret = psa_ssl_status_to_mbedtls( status );
+                goto psa_pake_error;
+            }
+
+            input_offset += length;
+        }
+    }
+
+    return( 0 );
+
+psa_pake_error:
+    psa_destroy_key( ssl->handshake->psa_pake_password );
+    psa_pake_abort( &ssl->handshake->psa_pake_ctx );
+
+    MBEDTLS_SSL_DEBUG_RET( 1, "psa_pake_input round one", ret );
+    mbedtls_ssl_send_alert_message(
+            ssl,
+            MBEDTLS_SSL_ALERT_LEVEL_FATAL,
+            MBEDTLS_SSL_ALERT_MSG_HANDSHAKE_FAILURE );
+
+    return( ret );
+#else
     if( ( ret = mbedtls_ecjpake_read_round_one( &ssl->handshake->ecjpake_ctx,
                                                 buf, len ) ) != 0 )
     {
@@ -901,6 +998,7 @@ static int ssl_parse_ecjpake_kkpp( mbedtls_ssl_context *ssl,
     }
 
     return( 0 );
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 }
 #endif /* MBEDTLS_KEY_EXCHANGE_ECJPAKE_ENABLED */
 
@@ -2296,6 +2394,61 @@ start_processing:
 #if defined(MBEDTLS_KEY_EXCHANGE_ECJPAKE_ENABLED)
     if( ciphersuite_info->key_exchange == MBEDTLS_KEY_EXCHANGE_ECJPAKE )
     {
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+        psa_status_t status;
+        size_t len = end - p;
+        size_t input_offset = 0;
+
+        for( psa_pake_step_t step = PSA_PAKE_STEP_KEY_SHARE ;
+             step <= PSA_PAKE_STEP_ZK_PROOF ;
+             ++step )
+        {
+            size_t length;
+
+            if( step == PSA_PAKE_STEP_KEY_SHARE )
+            {
+                /* Length is stored after 3bytes curve */
+                length = p[input_offset + 3];
+                input_offset += 3 + 1;
+            }
+            else
+            {
+                /* Length is stored at the first byte */
+                length = p[input_offset];
+                input_offset += 1;
+            }
+
+            if( input_offset + length > len )
+            {
+                ret = MBEDTLS_ERR_SSL_BAD_INPUT_DATA;
+                goto psa_pake_out;
+            }
+
+            status = psa_pake_input( &ssl->handshake->psa_pake_ctx, step,
+                                     p + input_offset, length );
+            if( status != PSA_SUCCESS)
+            {
+                ret = psa_ssl_status_to_mbedtls( status );
+                goto psa_pake_out;
+            }
+
+            input_offset += length;
+        }
+
+psa_pake_out:
+        if( ret != 0 )
+        {
+            psa_destroy_key( ssl->handshake->psa_pake_password );
+            psa_pake_abort( &ssl->handshake->psa_pake_ctx );
+
+            MBEDTLS_SSL_DEBUG_RET( 1, "psa_pake_input round two", ret );
+            mbedtls_ssl_send_alert_message(
+                ssl,
+                MBEDTLS_SSL_ALERT_LEVEL_FATAL,
+                MBEDTLS_SSL_ALERT_MSG_HANDSHAKE_FAILURE );
+            return( MBEDTLS_ERR_SSL_HANDSHAKE_FAILURE );
+        }
+#else
         ret = mbedtls_ecjpake_read_round_two( &ssl->handshake->ecjpake_ctx,
                                               p, end - p );
         if( ret != 0 )
@@ -2307,6 +2460,7 @@ start_processing:
                 MBEDTLS_SSL_ALERT_MSG_HANDSHAKE_FAILURE );
             return( MBEDTLS_ERR_SSL_HANDSHAKE_FAILURE );
         }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
     }
     else
 #endif /* MBEDTLS_KEY_EXCHANGE_ECJPAKE_ENABLED */
@@ -3235,6 +3389,42 @@ ecdh_calc_secret:
     {
         header_len = 4;
 
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+        unsigned char *out_p = ssl->out_msg + header_len;
+        unsigned char *end_p = ssl->out_msg + MBEDTLS_SSL_OUT_CONTENT_LEN -
+                               header_len;
+        psa_status_t status;
+        size_t output_offset = 0;
+        size_t output_len;
+
+        for( psa_pake_step_t step = PSA_PAKE_STEP_KEY_SHARE ;
+             step <= PSA_PAKE_STEP_ZK_PROOF ;
+             ++step )
+        {
+            /* For each step, prepend 1 byte with the length of the data */
+            if (step != PSA_PAKE_STEP_ZK_PROOF) {
+                *(out_p + output_offset) = 65;
+            } else {
+                *(out_p + output_offset) = 32;
+            }
+            output_offset += 1;
+            status = psa_pake_output( &ssl->handshake->psa_pake_ctx,
+                                      step, out_p + output_offset,
+                                      end_p - out_p - output_offset,
+                                      &output_len );
+            if( status != PSA_SUCCESS )
+            {
+                psa_destroy_key( ssl->handshake->psa_pake_password );
+                psa_pake_abort( &ssl->handshake->psa_pake_ctx );
+                MBEDTLS_SSL_DEBUG_RET( 1 , "psa_pake_output", status );
+                return( psa_ssl_status_to_mbedtls( status ) );
+            }
+
+            output_offset += output_len;
+        }
+
+        content_len = output_offset;
+#else
         ret = mbedtls_ecjpake_write_round_two( &ssl->handshake->ecjpake_ctx,
                 ssl->out_msg + header_len,
                 MBEDTLS_SSL_OUT_CONTENT_LEN - header_len,
@@ -3254,6 +3444,7 @@ ecdh_calc_secret:
             MBEDTLS_SSL_DEBUG_RET( 1, "mbedtls_ecjpake_derive_secret", ret );
             return( ret );
         }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
     }
     else
 #endif /* MBEDTLS_KEY_EXCHANGE_RSA_ENABLED */

library/ssl_tls12_server.c

@@ -268,10 +268,11 @@ static int ssl_parse_supported_point_formats( mbedtls_ssl_context *ssl,
             ssl->handshake->ecdh_ctx.point_format = p[0];
 #endif /* !MBEDTLS_USE_PSA_CRYPTO &&
           ( MBEDTLS_ECDH_C || MBEDTLS_ECDSA_C ) */
-#if defined(MBEDTLS_KEY_EXCHANGE_ECJPAKE_ENABLED)
+#if !defined(MBEDTLS_USE_PSA_CRYPTO) &&                             \
+    defined(MBEDTLS_KEY_EXCHANGE_ECJPAKE_ENABLED)
             mbedtls_ecjpake_set_point_format( &ssl->handshake->ecjpake_ctx,
                                               p[0] );
-#endif
+#endif /* !MBEDTLS_USE_PSA_CRYPTO && MBEDTLS_KEY_EXCHANGE_ECJPAKE_ENABLED */
             MBEDTLS_SSL_DEBUG_MSG( 4, ( "point format selected: %d", p[0] ) );
             return( 0 );
         }
@@ -292,13 +293,52 @@ static int ssl_parse_ecjpake_kkpp( mbedtls_ssl_context *ssl,
                                    size_t len )
 {
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_status_t status;
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    if( ssl->handshake->psa_pake_ctx_is_ok != 1 )
+#else
     if( mbedtls_ecjpake_check( &ssl->handshake->ecjpake_ctx ) != 0 )
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
     {
         MBEDTLS_SSL_DEBUG_MSG( 3, ( "skip ecjpake kkpp extension" ) );
         return( 0 );
     }
 
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    size_t input_offset = 0;
+
+    /* Repeat the KEY_SHARE, ZK_PUBLIC & ZF_PROOF twice */
+    for( unsigned int x = 1 ; x <= 2 ; ++x )
+    {
+        for( psa_pake_step_t step = PSA_PAKE_STEP_KEY_SHARE ;
+             step <= PSA_PAKE_STEP_ZK_PROOF ;
+             ++step )
+        {
+            /* Length is stored at the first byte */
+            size_t length = buf[input_offset];
+            input_offset += 1;
+
+            if( input_offset + length > len )
+            {
+                ret = MBEDTLS_ERR_SSL_HANDSHAKE_FAILURE;
+                goto psa_pake_error;
+            }
+
+            status = psa_pake_input( &ssl->handshake->psa_pake_ctx, step,
+                                     buf + input_offset, length );
+            if( status != PSA_SUCCESS)
+            {
+                ret = psa_ssl_status_to_mbedtls( status );
+                goto psa_pake_error;
+            }
+
+            input_offset += length;
+        }
+    }
+#else
     if( ( ret = mbedtls_ecjpake_read_round_one( &ssl->handshake->ecjpake_ctx,
                                                 buf, len ) ) != 0 )
     {
@@ -307,11 +347,26 @@ static int ssl_parse_ecjpake_kkpp( mbedtls_ssl_context *ssl,
                                         MBEDTLS_SSL_ALERT_MSG_ILLEGAL_PARAMETER );
         return( ret );
     }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 
     /* Only mark the extension as OK when we're sure it is */
     ssl->handshake->cli_exts |= MBEDTLS_TLS_EXT_ECJPAKE_KKPP_OK;
 
     return( 0 );
+
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+psa_pake_error:
+    psa_destroy_key( ssl->handshake->psa_pake_password );
+    psa_pake_abort( &ssl->handshake->psa_pake_ctx );
+
+    MBEDTLS_SSL_DEBUG_RET( 1, "psa_pake_input round one", ret );
+    mbedtls_ssl_send_alert_message(
+            ssl,
+            MBEDTLS_SSL_ALERT_LEVEL_FATAL,
+            MBEDTLS_SSL_ALERT_MSG_HANDSHAKE_FAILURE );
+
+    return( ret );
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 }
 #endif /* MBEDTLS_KEY_EXCHANGE_ECJPAKE_ENABLED */
 
@@ -1973,7 +2028,11 @@ static void ssl_write_ecjpake_kkpp_ext( mbedtls_ssl_context *ssl,
                                         unsigned char *buf,
                                         size_t *olen )
 {
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    psa_status_t status;
+#else
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
     unsigned char *p = buf;
     const unsigned char *end = ssl->out_msg + MBEDTLS_SSL_OUT_CONTENT_LEN;
     size_t kkpp_len;
@@ -1996,6 +2055,42 @@ static void ssl_write_ecjpake_kkpp_ext( mbedtls_ssl_context *ssl,
     MBEDTLS_PUT_UINT16_BE( MBEDTLS_TLS_EXT_ECJPAKE_KKPP, p, 0 );
     p += 2;
 
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+    size_t output_offset = 0;
+    size_t output_len;
+
+    /* Repeat the KEY_SHARE, ZK_PUBLIC & ZF_PROOF twice */
+    for( unsigned int x = 1 ; x <= 2 ; ++x )
+    {
+        for( psa_pake_step_t step = PSA_PAKE_STEP_KEY_SHARE ;
+             step <= PSA_PAKE_STEP_ZK_PROOF ;
+             ++step )
+        {
+            /* For each step, prepend 1 byte with the length of the data */
+            if (step != PSA_PAKE_STEP_ZK_PROOF) {
+                *(p + 2 + output_offset) = 65;
+            } else {
+                *(p + 2 + output_offset) = 32;
+            }
+            output_offset += 1;
+            status = psa_pake_output( &ssl->handshake->psa_pake_ctx,
+                                      step, p + 2 + output_offset,
+                                      end - p - output_offset - 2,
+                                      &output_len );
+            if( status != PSA_SUCCESS )
+            {
+                psa_destroy_key( ssl->handshake->psa_pake_password );
+                psa_pake_abort( &ssl->handshake->psa_pake_ctx );
+                MBEDTLS_SSL_DEBUG_RET( 1 , "psa_pake_output", status );
+                return;
+            }
+
+            output_offset += output_len;
+        }
+    }
+
+    kkpp_len = output_offset;
+#else
     ret = mbedtls_ecjpake_write_round_one( &ssl->handshake->ecjpake_ctx,
                                         p + 2, end - p - 2, &kkpp_len,
                                         ssl->conf->f_rng, ssl->conf->p_rng );
@@ -2004,6 +2099,7 @@ static void ssl_write_ecjpake_kkpp_ext( mbedtls_ssl_context *ssl,
         MBEDTLS_SSL_DEBUG_RET( 1 , "mbedtls_ecjpake_write_round_one", ret );
         return;
     }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
 
     MBEDTLS_PUT_UINT16_BE( kkpp_len, p, 0 );
     p += 2;
@@ -2807,6 +2903,61 @@ static int ssl_prepare_server_key_exchange( mbedtls_ssl_context *ssl,
 #if defined(MBEDTLS_KEY_EXCHANGE_ECJPAKE_ENABLED)
     if( ciphersuite_info->key_exchange == MBEDTLS_KEY_EXCHANGE_ECJPAKE )
     {
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+        unsigned char *out_p = ssl->out_msg + ssl->out_msglen;
+        unsigned char *end_p = ssl->out_msg + MBEDTLS_SSL_OUT_CONTENT_LEN -
+                               ssl->out_msglen;
+        psa_status_t status;
+        size_t output_offset = 0;
+        size_t output_len;
+        size_t ec_len;
+
+#if !defined(MBEDTLS_ECJPAKE_ALT)
+        psa_pake_operation_t* pake_op = &(ssl->handshake->psa_pake_ctx);
+
+        mbedtls_ecp_tls_write_group( &(pake_op->ctx.ecjpake.grp),
+                                    &ec_len, out_p + output_offset,
+                                    end_p - out_p);
+#else
+        const mbedtls_ecp_curve_info *curve_info;
+
+        if( ( curve_info = mbedtls_ecp_curve_info_from_grp_id( MBEDTLS_ECP_DP_SECP256R1 ) ) == NULL )
+            return( MBEDTLS_ERR_ECP_BAD_INPUT_DATA );
+
+        *out_p = MBEDTLS_ECP_TLS_NAMED_CURVE;
+
+        MBEDTLS_PUT_UINT16_BE( curve_info->tls_id, out_p + 1, 0 );
+        ec_len = 3;
+#endif //MBEDTLS_PSA_BUILTIN_ALG_JPAKE
+        output_offset += ec_len;
+
+        for( psa_pake_step_t step = PSA_PAKE_STEP_KEY_SHARE ;
+             step <= PSA_PAKE_STEP_ZK_PROOF ;
+             ++step )
+        {
+            if (step != PSA_PAKE_STEP_ZK_PROOF) {
+                *(out_p + output_offset) = 65;
+            } else {
+                *(out_p + output_offset) = 32;
+            }
+            output_offset += 1;
+            status = psa_pake_output( &ssl->handshake->psa_pake_ctx,
+                                      step, out_p + output_offset,
+                                      end_p - out_p - output_offset,
+                                      &output_len );
+            if( status != PSA_SUCCESS )
+            {
+                psa_destroy_key( ssl->handshake->psa_pake_password );
+                psa_pake_abort( &ssl->handshake->psa_pake_ctx );
+                MBEDTLS_SSL_DEBUG_RET( 1 , "psa_pake_output", status );
+                return( psa_ssl_status_to_mbedtls( status ) );
+            }
+
+            output_offset += output_len;
+        }
+
+        ssl->out_msglen += output_offset;
+#else
         int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
         size_t len = 0;
 
@@ -2822,6 +2973,7 @@ static int ssl_prepare_server_key_exchange( mbedtls_ssl_context *ssl,
         }
 
         ssl->out_msglen += len;
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
     }
 #endif /* MBEDTLS_KEY_EXCHANGE_ECJPAKE_ENABLED */
 
@@ -4039,6 +4191,46 @@ static int ssl_parse_client_key_exchange( mbedtls_ssl_context *ssl )
 #if defined(MBEDTLS_KEY_EXCHANGE_ECJPAKE_ENABLED)
     if( ciphersuite_info->key_exchange == MBEDTLS_KEY_EXCHANGE_ECJPAKE )
     {
+#if defined(MBEDTLS_USE_PSA_CRYPTO)
+        size_t len = end - p;
+        psa_status_t status;
+        size_t input_offset = 0;
+
+        for( psa_pake_step_t step = PSA_PAKE_STEP_KEY_SHARE ;
+             step <= PSA_PAKE_STEP_ZK_PROOF ;
+             ++step )
+        {
+            /* Length is stored at the first byte */
+            size_t length = p[input_offset];
+            input_offset += 1;
+
+            if( input_offset + length > len )
+            {
+                ret = MBEDTLS_ERR_SSL_HANDSHAKE_FAILURE;
+                goto psa_pake_out;
+            }
+
+            status = psa_pake_input( &ssl->handshake->psa_pake_ctx, step,
+                                     p + input_offset, length );
+            if( status != PSA_SUCCESS)
+            {
+                ret = psa_ssl_status_to_mbedtls( status );
+                goto psa_pake_out;
+            }
+
+            input_offset += length;
+        }
+
+psa_pake_out:
+        if( ret != 0 )
+        {
+            psa_destroy_key( ssl->handshake->psa_pake_password );
+            psa_pake_abort( &ssl->handshake->psa_pake_ctx );
+
+            MBEDTLS_SSL_DEBUG_RET( 1, "psa_pake_input round two", ret );
+            return( ret );
+        }
+#else
         ret = mbedtls_ecjpake_read_round_two( &ssl->handshake->ecjpake_ctx,
                                               p, end - p );
         if( ret != 0 )
@@ -4055,6 +4247,7 @@ static int ssl_parse_client_key_exchange( mbedtls_ssl_context *ssl )
             MBEDTLS_SSL_DEBUG_RET( 1, "mbedtls_ecjpake_derive_secret", ret );
             return( ret );
         }
+#endif /* MBEDTLS_USE_PSA_CRYPTO */
     }
     else
 #endif /* MBEDTLS_KEY_EXCHANGE_ECJPAKE_ENABLED */