tls: psa_pake: make round two reading function symmatric to the writing one

Signed-off-by: Valerio Setti <vsetti@baylibre.com>
diff --git a/library/ssl_misc.h b/library/ssl_misc.h
index d022721..82a951a 100644
--- a/library/ssl_misc.h
+++ b/library/ssl_misc.h
@@ -2405,8 +2405,7 @@
  */
 int mbedtls_psa_ecjpake_read_round_two(
                                     psa_pake_operation_t *pake_ctx,
-                                    const unsigned char *buf,
-                                    size_t len, int role );
+                                    const unsigned char *buf, size_t len );
 
 /**
  * \brief       Write the first round of key exchange into the provided output
diff --git a/library/ssl_tls.c b/library/ssl_tls.c
index 06a5ec5..ae12c7e 100644
--- a/library/ssl_tls.c
+++ b/library/ssl_tls.c
@@ -8240,7 +8240,7 @@
 int mbedtls_psa_ecjpake_read_round_two(
                                     psa_pake_operation_t *pake_ctx,
                                     const unsigned char *buf,
-                                    size_t len, int role )
+                                    size_t len )
 {
     psa_status_t status;
     size_t input_offset = 0;
@@ -8251,25 +8251,6 @@
     {
         size_t length;
 
-        /*
-         * On its 2nd round, the server sends 3 extra bytes which identify the
-         * curve:
-         * - the 1st one is MBEDTLS_ECP_TLS_NAMED_CURVE
-         * - the 2nd and 3rd represent curve's TLS ID
-         * Validate this data before moving forward
-         */
-        if( ( step == PSA_PAKE_STEP_KEY_SHARE ) &&
-            ( role == MBEDTLS_SSL_IS_CLIENT ) )
-        {
-            uint16_t tls_id = MBEDTLS_GET_UINT16_BE( buf, 1 );
-
-            if( ( *buf != MBEDTLS_ECP_TLS_NAMED_CURVE ) ||
-                ( mbedtls_ecp_curve_info_from_tls_id( tls_id ) == NULL ) )
-                return( MBEDTLS_ERR_ECP_BAD_INPUT_DATA );
-
-            input_offset += 3;
-        }
-
         /* Length is stored at the first byte */
         length = buf[input_offset];
         input_offset += 1;
diff --git a/library/ssl_tls12_client.c b/library/ssl_tls12_client.c
index 4e986d1..6dd8ef5 100644
--- a/library/ssl_tls12_client.c
+++ b/library/ssl_tls12_client.c
@@ -2333,9 +2333,31 @@
     if( ciphersuite_info->key_exchange == MBEDTLS_KEY_EXCHANGE_ECJPAKE )
     {
 #if defined(MBEDTLS_USE_PSA_CRYPTO)
+        /*
+         * The first 3 bytes are:
+         * [0] MBEDTLS_ECP_TLS_NAMED_CURVE
+         * [1, 2] elliptic curve's TLS ID
+         *
+         * However since we only support secp256r1 for now, we check only
+         * that TLS ID here
+         */
+        uint16_t read_tls_id = MBEDTLS_GET_UINT16_BE( p, 1 );
+        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_SSL_FEATURE_UNAVAILABLE );
+        }
+
+        if( ( *p != MBEDTLS_ECP_TLS_NAMED_CURVE ) ||
+            ( read_tls_id != curve_info->tls_id ) )
+            return( MBEDTLS_ERR_ECP_BAD_INPUT_DATA );
+
+        p += 3;
+
         if( ( ret = mbedtls_psa_ecjpake_read_round_two(
-                        &ssl->handshake->psa_pake_ctx, p, end - p,
-                        ssl->conf->endpoint ) ) != 0 )
+                        &ssl->handshake->psa_pake_ctx, p, end - p ) ) != 0 )
         {
             psa_destroy_key( ssl->handshake->psa_pake_password );
             psa_pake_abort( &ssl->handshake->psa_pake_ctx );
diff --git a/library/ssl_tls12_server.c b/library/ssl_tls12_server.c
index 2e48063..3bc7217 100644
--- a/library/ssl_tls12_server.c
+++ b/library/ssl_tls12_server.c
@@ -4115,8 +4115,7 @@
     {
 #if defined(MBEDTLS_USE_PSA_CRYPTO)
         if( ( ret = mbedtls_psa_ecjpake_read_round_two(
-                        &ssl->handshake->psa_pake_ctx, p, end - p,
-                        ssl->conf->endpoint ) ) != 0 )
+                        &ssl->handshake->psa_pake_ctx, p, end - p ) ) != 0 )
         {
             psa_destroy_key( ssl->handshake->psa_pake_password );
             psa_pake_abort( &ssl->handshake->psa_pake_ctx );