diff --git a/library/ssl_tls12_client.c b/library/ssl_tls12_client.c
index 3d25e40..c90ed2e 100644
--- a/library/ssl_tls12_client.c
+++ b/library/ssl_tls12_client.c
@@ -130,13 +130,9 @@
                                        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;
+    size_t kkpp_len = 0;
 
     *olen = 0;
 
@@ -168,41 +164,15 @@
         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 )
+        ret = psa_tls12_write_ecjpake_round_one(&ssl->handshake->psa_pake_ctx,
+                                                p + 2, end - p - 2, &kkpp_len );
+        if ( ret != 0 )
         {
-            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;
-            }
+            psa_destroy_key( ssl->handshake->psa_pake_password );
+            psa_pake_abort( &ssl->handshake->psa_pake_ctx );
+            MBEDTLS_SSL_DEBUG_RET( 1 , "psa_pake_output", ret );
+            return( ret );
         }
-
-        kkpp_len = output_offset;
 #else
         ret = mbedtls_ecjpake_write_round_one( &ssl->handshake->ecjpake_ctx,
                                                p + 2, end - p - 2, &kkpp_len,
@@ -924,9 +894,6 @@
                                    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 )
@@ -941,50 +908,21 @@
     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 )
+    if( ( ret = psa_tls12_parse_ecjpake_round_one(
+                            &ssl->handshake->psa_pake_ctx, buf, len ) ) != 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 = buf[input_offset];
-            input_offset += 1;
+        psa_destroy_key( ssl->handshake->psa_pake_password );
+        psa_pake_abort( &ssl->handshake->psa_pake_ctx );
 
-            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;
-        }
+        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 );
     }
 
     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 )
@@ -2395,48 +2333,9 @@
     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 )
+        if( ( ret = psa_tls12_parse_ecjpake_round_two(
+                        &ssl->handshake->psa_pake_ctx, p, end - p,
+                        ssl->conf->endpoint ) ) != 0 )
         {
             psa_destroy_key( ssl->handshake->psa_pake_password );
             psa_pake_abort( &ssl->handshake->psa_pake_ctx );
@@ -3393,37 +3292,15 @@
         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 )
+        ret = psa_tls12_write_ecjpake_round_two( &ssl->handshake->psa_pake_ctx,
+                                    out_p, end_p - out_p, &content_len );
+        if ( ret != 0 )
         {
-            /* 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;
+            psa_destroy_key( ssl->handshake->psa_pake_password );
+            psa_pake_abort( &ssl->handshake->psa_pake_ctx );
+            MBEDTLS_SSL_DEBUG_RET( 1 , "psa_pake_output", ret );
+            return( ret );
         }
-
-        content_len = output_offset;
 #else
         ret = mbedtls_ecjpake_write_round_two( &ssl->handshake->ecjpake_ctx,
                 ssl->out_msg + header_len,