fix various issues

- Signature of
  - mbedtls_tls13_set_hs_sent_ext_mask
  - check_received_extension and issues
- Also fix comment issue.
- improve readablity.

Signed-off-by: Jerry Yu <jerry.h.yu@arm.com>
diff --git a/library/ssl_client.c b/library/ssl_client.c
index b0d2dcf..16cef02 100644
--- a/library/ssl_client.c
+++ b/library/ssl_client.c
@@ -107,8 +107,7 @@
     *olen = hostname_len + 9;
 
 #if defined(MBEDTLS_SSL_PROTO_TLS1_3)
-    mbedtls_tls13_set_sent_ext_mask( ssl,
-                                     MBEDTLS_TLS_EXT_SERVERNAME );
+    mbedtls_ssl_tls13_set_hs_sent_ext_mask( ssl, MBEDTLS_TLS_EXT_SERVERNAME );
     MBEDTLS_SSL_DEBUG_MSG(
         4, ( "sent %s extension",
              mbedtls_tls13_get_extension_name(
@@ -186,8 +185,7 @@
     MBEDTLS_PUT_UINT16_BE( *out_len - 4, buf, 2 );
 
 #if defined(MBEDTLS_SSL_PROTO_TLS1_3)
-    mbedtls_tls13_set_sent_ext_mask( ssl,
-                                     MBEDTLS_TLS_EXT_ALPN );
+    mbedtls_ssl_tls13_set_hs_sent_ext_mask( ssl, MBEDTLS_TLS_EXT_ALPN );
     MBEDTLS_SSL_DEBUG_MSG(
         4, ( "sent %s extension",
              mbedtls_tls13_get_extension_name(
@@ -312,8 +310,7 @@
     *out_len = p - buf;
 
 #if defined(MBEDTLS_SSL_PROTO_TLS1_3)
-    mbedtls_tls13_set_sent_ext_mask( ssl,
-                                     MBEDTLS_TLS_EXT_SUPPORTED_GROUPS );
+    mbedtls_ssl_tls13_set_hs_sent_ext_mask( ssl, MBEDTLS_TLS_EXT_SUPPORTED_GROUPS );
     MBEDTLS_SSL_DEBUG_MSG( 4, ( "sent %s extension",
                                 mbedtls_tls13_get_extension_name(
                                     MBEDTLS_TLS_EXT_SUPPORTED_GROUPS ) ) );
diff --git a/library/ssl_misc.h b/library/ssl_misc.h
index dad1c82..3aeab0c 100644
--- a/library/ssl_misc.h
+++ b/library/ssl_misc.h
@@ -103,11 +103,12 @@
 #define MBEDTLS_SSL_EXT_SIG_ALG_CERT                ( 1 << 20 )
 #define MBEDTLS_SSL_EXT_KEY_SHARE                   ( 1 << 21 )
 
-/* For request messages, we should just ignore unrecognized extension when
- * parsing messages. For response messages, we should not ignore unrecognized
- * extension when parsing messages. Request messages include ClientHello,
- * Certificate and NewSessionTicket. Response messages include ServerHello,
- * EncryptExtensions, Certificate and HelloRetryRequest.
+/* In messages containing extension requests, we should ignore unrecognized
+ * extensions. In messages containing extension responses, unrecognized
+ * extensions should result in handshake abortion. Messages containing
+ * extension requests include ClientHello, CertificateRequest and
+ * NewSessionTicket. Messages containing extension responses include
+ * ServerHello, HelloRetryRequest, EncryptedExtensions and Certificate.
  *
  * RFC 8446 section 4.1.3
  *
@@ -1932,13 +1933,14 @@
 uint32_t mbedtls_tls13_get_extension_mask( unsigned int extension_type );
 
 MBEDTLS_CHECK_RETURN_CRITICAL
-int mbedtls_tls13_check_received_extensions( mbedtls_ssl_context *ssl,
-                                             int hs_msg_type,
-                                             uint32_t extension_type,
-                                             uint32_t allowed_mask );
+int mbedtls_ssl_tls13_check_received_extension(
+        mbedtls_ssl_context *ssl,
+        int hs_msg_type,
+        unsigned int received_extension_type,
+        uint32_t hs_msg_allowed_extensions_mask );
 
-static inline void mbedtls_tls13_set_sent_ext_mask( mbedtls_ssl_context *ssl,
-                                                    uint16_t extension_type )
+static inline void mbedtls_ssl_tls13_set_hs_sent_ext_mask(
+                       mbedtls_ssl_context *ssl, unsigned int extension_type )
 {
     ssl->handshake->sent_extensions |=
         mbedtls_tls13_get_extension_mask( extension_type );
diff --git a/library/ssl_tls.c b/library/ssl_tls.c
index 9947d39..3678ec0 100644
--- a/library/ssl_tls.c
+++ b/library/ssl_tls.c
@@ -8713,8 +8713,7 @@
     *out_len = p - buf;
 
 #if defined(MBEDTLS_SSL_PROTO_TLS1_3)
-    mbedtls_tls13_set_sent_ext_mask( ssl,
-                                     MBEDTLS_TLS_EXT_SIG_ALG );
+    mbedtls_ssl_tls13_set_hs_sent_ext_mask( ssl, MBEDTLS_TLS_EXT_SIG_ALG );
     MBEDTLS_SSL_DEBUG_MSG(
         4, ( "sent %s extension",
              mbedtls_tls13_get_extension_name(
diff --git a/library/ssl_tls13_client.c b/library/ssl_tls13_client.c
index c29b90e..27747a2 100644
--- a/library/ssl_tls13_client.c
+++ b/library/ssl_tls13_client.c
@@ -89,8 +89,8 @@
     }
 
     *out_len = 5 + versions_len;
-    mbedtls_tls13_set_sent_ext_mask( ssl,
-                                     MBEDTLS_TLS_EXT_SUPPORTED_VERSIONS );
+    mbedtls_ssl_tls13_set_hs_sent_ext_mask(
+        ssl, MBEDTLS_TLS_EXT_SUPPORTED_VERSIONS );
     MBEDTLS_SSL_DEBUG_MSG(
         4, ( "sent %s extension",
              mbedtls_tls13_get_extension_name(
@@ -365,8 +365,7 @@
 
     MBEDTLS_SSL_DEBUG_BUF( 3, "client hello, key_share extension", buf, *out_len );
 
-    mbedtls_tls13_set_sent_ext_mask( ssl,
-                                     MBEDTLS_TLS_EXT_KEY_SHARE );
+    mbedtls_ssl_tls13_set_hs_sent_ext_mask( ssl, MBEDTLS_TLS_EXT_KEY_SHARE );
     MBEDTLS_SSL_DEBUG_MSG(
         4, ( "sent %s extension",
 
@@ -612,8 +611,7 @@
     *out_len = handshake->hrr_cookie_len + 6;
 
 
-    mbedtls_tls13_set_sent_ext_mask( ssl,
-                                     MBEDTLS_TLS_EXT_COOKIE );
+    mbedtls_ssl_tls13_set_hs_sent_ext_mask( ssl, MBEDTLS_TLS_EXT_COOKIE );
     MBEDTLS_SSL_DEBUG_MSG(
         4, ( "sent %s extension",
              mbedtls_tls13_get_extension_name(
@@ -688,8 +686,8 @@
 
     *out_len = p - buf;
 
-    mbedtls_tls13_set_sent_ext_mask( ssl,
-                                     MBEDTLS_TLS_EXT_PSK_KEY_EXCHANGE_MODES );
+    mbedtls_ssl_tls13_set_hs_sent_ext_mask(
+        ssl, MBEDTLS_TLS_EXT_PSK_KEY_EXCHANGE_MODES );
     MBEDTLS_SSL_DEBUG_MSG(
         4, ( "sent %s extension",
              mbedtls_tls13_get_extension_name(
@@ -1059,8 +1057,8 @@
 
     MBEDTLS_SSL_DEBUG_BUF( 3, "pre_shared_key binders", buf, p - buf );
 
-    mbedtls_tls13_set_sent_ext_mask( ssl,
-                                     MBEDTLS_TLS_EXT_PRE_SHARED_KEY );
+    mbedtls_ssl_tls13_set_hs_sent_ext_mask(
+        ssl, MBEDTLS_TLS_EXT_PRE_SHARED_KEY );
     MBEDTLS_SSL_DEBUG_MSG(
         4, ( "sent %s extension",
              mbedtls_tls13_get_extension_name(
@@ -1668,7 +1666,7 @@
 
     MBEDTLS_SSL_DEBUG_BUF( 3, "server hello extensions", p, extensions_len );
 
-    ssl->handshake->received_extensions = MBEDTLS_SSL_EXT_NONE;
+    handshake->received_extensions = MBEDTLS_SSL_EXT_NONE;
     allowed_extensions_mask = is_hrr ?
                                   MBEDTLS_SSL_TLS1_3_ALLOWED_EXTS_OF_HRR :
                                   MBEDTLS_SSL_TLS1_3_ALLOWED_EXTS_OF_SH;
@@ -1687,7 +1685,7 @@
         MBEDTLS_SSL_CHK_BUF_READ_PTR( p, extensions_end, extension_data_len );
         extension_data_end = p + extension_data_len;
 
-        ret = mbedtls_tls13_check_received_extensions(
+        ret = mbedtls_ssl_tls13_check_received_extension(
                   ssl,
                   is_hrr ?
                       -MBEDTLS_SSL_HS_SERVER_HELLO : MBEDTLS_SSL_HS_SERVER_HELLO,
@@ -1766,7 +1764,7 @@
 
     MBEDTLS_SSL_TLS1_3_PRINT_EXTS(
         3, is_hrr ? -MBEDTLS_SSL_HS_SERVER_HELLO : MBEDTLS_SSL_HS_SERVER_HELLO,
-        ssl->handshake->received_extensions );
+        handshake->received_extensions );
 
 cleanup:
 
@@ -1999,6 +1997,7 @@
     size_t extensions_len;
     const unsigned char *p = buf;
     const unsigned char *extensions_end;
+    mbedtls_ssl_handshake_params *handshake = ssl->handshake;
 
     MBEDTLS_SSL_CHK_BUF_READ_PTR( p, end, 2 );
     extensions_len = MBEDTLS_GET_UINT16_BE( p, 0 );
@@ -2008,7 +2007,7 @@
     MBEDTLS_SSL_CHK_BUF_READ_PTR( p, end, extensions_len );
     extensions_end = p + extensions_len;
 
-    ssl->handshake->received_extensions = MBEDTLS_SSL_EXT_NONE;
+    handshake->received_extensions = MBEDTLS_SSL_EXT_NONE;
 
     while( p < extensions_end )
     {
@@ -2028,7 +2027,7 @@
 
         MBEDTLS_SSL_CHK_BUF_READ_PTR( p, extensions_end, extension_data_len );
 
-        ret = mbedtls_tls13_check_received_extensions(
+        ret = mbedtls_ssl_tls13_check_received_extension(
                   ssl, MBEDTLS_SSL_HS_ENCRYPTED_EXTENSIONS, extension_type,
                   MBEDTLS_SSL_TLS1_3_ALLOWED_EXTS_OF_EE );
         if( ret != 0 )
@@ -2036,16 +2035,6 @@
 
         switch( extension_type )
         {
-            case MBEDTLS_TLS_EXT_SERVERNAME:
-                MBEDTLS_SSL_DEBUG_MSG( 3, ( "found server_name extension" ) );
-
-                /* The server_name extension should be an empty extension */
-
-                break;
-            case MBEDTLS_TLS_EXT_SUPPORTED_GROUPS:
-                MBEDTLS_SSL_DEBUG_MSG( 3, ( "found extensions supported groups" ) );
-                break;
-
 #if defined(MBEDTLS_SSL_ALPN)
             case MBEDTLS_TLS_EXT_ALPN:
                 MBEDTLS_SSL_DEBUG_MSG( 3, ( "found alpn extension" ) );
@@ -2069,7 +2058,7 @@
     }
 
     MBEDTLS_SSL_TLS1_3_PRINT_EXTS( 3, MBEDTLS_SSL_HS_ENCRYPTED_EXTENSIONS,
-                                   ssl->handshake->received_extensions );
+                                   handshake->received_extensions );
 
     /* Check that we consumed all the message. */
     if( p != end )
@@ -2176,6 +2165,7 @@
     size_t certificate_request_context_len = 0;
     size_t extensions_len = 0;
     const unsigned char *extensions_end;
+    mbedtls_ssl_handshake_params *handshake = ssl->handshake;
 
     /* ...
      * opaque certificate_request_context<0..2^8-1>
@@ -2191,7 +2181,6 @@
         MBEDTLS_SSL_DEBUG_BUF( 3, "Certificate Request Context",
                                p, certificate_request_context_len );
 
-        mbedtls_ssl_handshake_params *handshake = ssl->handshake;
         handshake->certificate_request_context =
                 mbedtls_calloc( 1, certificate_request_context_len );
         if( handshake->certificate_request_context == NULL )
@@ -2215,7 +2204,7 @@
     MBEDTLS_SSL_CHK_BUF_READ_PTR( p, end, extensions_len );
     extensions_end = p + extensions_len;
 
-    ssl->handshake->received_extensions = MBEDTLS_SSL_EXT_NONE;
+    handshake->received_extensions = MBEDTLS_SSL_EXT_NONE;
 
     while( p < extensions_end )
     {
@@ -2229,7 +2218,7 @@
 
         MBEDTLS_SSL_CHK_BUF_READ_PTR( p, extensions_end, extension_data_len );
 
-        ret = mbedtls_tls13_check_received_extensions(
+        ret = mbedtls_ssl_tls13_check_received_extension(
                   ssl, MBEDTLS_SSL_HS_CERTIFICATE_REQUEST, extension_type,
                   MBEDTLS_SSL_TLS1_3_ALLOWED_EXTS_OF_CR );
         if( ret != 0 )
@@ -2260,7 +2249,7 @@
 
     MBEDTLS_SSL_TLS1_3_PRINT_EXTS( 3,
                                    MBEDTLS_SSL_HS_CERTIFICATE_REQUEST,
-                                   ssl->handshake->received_extensions );
+                                   handshake->received_extensions );
 
     /* Check that we consumed all the message. */
     if( p != end )
@@ -2274,7 +2263,7 @@
      *
      * The "signature_algorithms" extension MUST be specified
      */
-    if( ( ssl->handshake->received_extensions & MBEDTLS_SSL_EXT_SIG_ALG ) == 0 )
+    if( ( handshake->received_extensions & MBEDTLS_SSL_EXT_SIG_ALG ) == 0 )
     {
         MBEDTLS_SSL_DEBUG_MSG( 3,
             ( "no signature algorithms extension found" ) );
@@ -2514,10 +2503,11 @@
                                                     const unsigned char *buf,
                                                     const unsigned char *end )
 {
+    mbedtls_ssl_handshake_params *handshake = ssl->handshake;
     const unsigned char *p = buf;
 
 
-    ssl->handshake->received_extensions = MBEDTLS_SSL_EXT_NONE;
+    handshake->received_extensions = MBEDTLS_SSL_EXT_NONE;
 
     while( p < end )
     {
@@ -2532,7 +2522,7 @@
 
         MBEDTLS_SSL_CHK_BUF_READ_PTR( p, end, extension_data_len );
 
-        ret = mbedtls_tls13_check_received_extensions(
+        ret = mbedtls_ssl_tls13_check_received_extension(
                   ssl, MBEDTLS_SSL_HS_CLIENT_HELLO, extension_type,
                   MBEDTLS_SSL_TLS1_3_ALLOWED_EXTS_OF_CH );
         if( ret != 0 )
@@ -2556,7 +2546,7 @@
     }
 
     MBEDTLS_SSL_TLS1_3_PRINT_EXTS(
-        3, MBEDTLS_SSL_HS_NEW_SESSION_TICKET, ssl->handshake->received_extensions );
+        3, MBEDTLS_SSL_HS_NEW_SESSION_TICKET, handshake->received_extensions );
 
     return( 0 );
 }
diff --git a/library/ssl_tls13_generic.c b/library/ssl_tls13_generic.c
index 2fbcdf0..21644ed 100644
--- a/library/ssl_tls13_generic.c
+++ b/library/ssl_tls13_generic.c
@@ -398,6 +398,7 @@
     size_t certificate_list_len = 0;
     const unsigned char *p = buf;
     const unsigned char *certificate_list_end;
+    mbedtls_ssl_handshake_params *handshake = ssl->handshake;
 
     MBEDTLS_SSL_CHK_BUF_READ_PTR( p, end, 4 );
     certificate_request_context_len = p[0];
@@ -507,7 +508,7 @@
         MBEDTLS_SSL_CHK_BUF_READ_PTR( p, certificate_list_end, extensions_len );
 
         extensions_end = p + extensions_len;
-        ssl->handshake->received_extensions = MBEDTLS_SSL_EXT_NONE;
+        handshake->received_extensions = MBEDTLS_SSL_EXT_NONE;
 
         while( p < extensions_end )
         {
@@ -527,9 +528,9 @@
 
             MBEDTLS_SSL_CHK_BUF_READ_PTR( p, extensions_end, extension_data_len );
 
-            ret = mbedtls_tls13_check_received_extensions(
-                  ssl, MBEDTLS_SSL_HS_CERTIFICATE, extension_type,
-                  MBEDTLS_SSL_TLS1_3_ALLOWED_EXTS_OF_CT );
+            ret = mbedtls_ssl_tls13_check_received_extension(
+                      ssl, MBEDTLS_SSL_HS_CERTIFICATE, extension_type,
+                      MBEDTLS_SSL_TLS1_3_ALLOWED_EXTS_OF_CT );
             if( ret != 0 )
                 return( ret );
 
@@ -547,7 +548,7 @@
         }
 
         MBEDTLS_SSL_TLS1_3_PRINT_EXTS(
-            3, MBEDTLS_SSL_HS_CERTIFICATE, ssl->handshake->received_extensions );
+            3, MBEDTLS_SSL_HS_CERTIFICATE, handshake->received_extensions );
     }
 
 exit:
@@ -555,7 +556,7 @@
     if( p != end )
     {
         MBEDTLS_SSL_DEBUG_MSG( 1, ( "bad Certificate message" ) );
-        MBEDTLS_SSL_PEND_FATAL_ALERT( MBEDTLS_SSL_ALERT_MSG_DECODE_ERROR, \
+        MBEDTLS_SSL_PEND_FATAL_ALERT( MBEDTLS_SSL_ALERT_MSG_DECODE_ERROR,
                                       MBEDTLS_ERR_SSL_DECODE_ERROR );
         return( MBEDTLS_ERR_SSL_DECODE_ERROR );
     }
@@ -1759,43 +1760,46 @@
  *
  */
 
-int mbedtls_tls13_check_received_extensions( mbedtls_ssl_context *ssl,
-                                             int hs_msg_type,
-                                             uint32_t extension_type,
-                                             uint32_t allowed_mask )
+int mbedtls_ssl_tls13_check_received_extension(
+        mbedtls_ssl_context *ssl,
+        int hs_msg_type,
+        unsigned int received_extension_type,
+        uint32_t hs_msg_allowed_extensions_mask )
 {
-    uint32_t extension_mask;
-
 #if defined(MBEDTLS_DEBUG_C)
     const char *hs_msg_name = ssl_tls13_get_hs_msg_name( hs_msg_type );
 #endif
-
-    extension_mask = mbedtls_tls13_get_extension_mask( extension_type );
+    uint32_t extension_mask = mbedtls_tls13_get_extension_mask( received_extension_type );
 
     MBEDTLS_SSL_DEBUG_MSG( 3,
                 ( "%s : received %s(%x) extension",
                   hs_msg_name,
-                  mbedtls_tls13_get_extension_name( extension_type ),
-                  (unsigned int)extension_type ) );
+                  mbedtls_tls13_get_extension_name( received_extension_type ),
+                  (unsigned int)received_extension_type ) );
 
-    if( ( extension_mask & allowed_mask ) == 0 )
+    if( ( extension_mask & hs_msg_allowed_extensions_mask ) == 0 )
     {
         MBEDTLS_SSL_DEBUG_MSG(
             3, ( "%s : forbidden extension received.", hs_msg_name ) );
         MBEDTLS_SSL_PEND_FATAL_ALERT(
             MBEDTLS_SSL_ALERT_MSG_ILLEGAL_PARAMETER,
-            MBEDTLS_ERR_SSL_HANDSHAKE_FAILURE );
-        return( MBEDTLS_ERR_SSL_HANDSHAKE_FAILURE );
+            MBEDTLS_ERR_SSL_ILLEGAL_PARAMETER );
+        return( MBEDTLS_ERR_SSL_ILLEGAL_PARAMETER );
     }
 
     ssl->handshake->received_extensions |= extension_mask;
+    /*
+     * If it is a message containing extension responses, check that we
+     * previously sent the extension.
+     */
     switch( hs_msg_type )
     {
         case MBEDTLS_SSL_HS_SERVER_HELLO:
         case -MBEDTLS_SSL_HS_SERVER_HELLO: // HRR does not have IANA value.
         case MBEDTLS_SSL_HS_ENCRYPTED_EXTENSIONS:
         case MBEDTLS_SSL_HS_CERTIFICATE:
-            if( ( ~ssl->handshake->sent_extensions & extension_mask ) == 0 )
+            /* Check if the received extension is sent by peer message.*/
+            if( ( ssl->handshake->sent_extensions & extension_mask ) != 0 )
                 return( 0 );
             break;
         default:
@@ -1803,11 +1807,11 @@
     }
 
     MBEDTLS_SSL_DEBUG_MSG(
-            3, ( "%s : forbidden extension received.", hs_msg_name ) );
+            3, ( "%s : unexpected extension received.", hs_msg_name ) );
     MBEDTLS_SSL_PEND_FATAL_ALERT(
         MBEDTLS_SSL_ALERT_MSG_UNSUPPORTED_EXT,
-        MBEDTLS_ERR_SSL_HANDSHAKE_FAILURE );
-    return( MBEDTLS_ERR_SSL_HANDSHAKE_FAILURE );
+        MBEDTLS_ERR_SSL_UNSUPPORTED_EXTENSION );
+    return( MBEDTLS_ERR_SSL_UNSUPPORTED_EXTENSION );
 }
 
 #endif /* MBEDTLS_SSL_TLS_C && MBEDTLS_SSL_PROTO_TLS1_3 */
diff --git a/library/ssl_tls13_server.c b/library/ssl_tls13_server.c
index 4fdd6ad..f0f06b8 100644
--- a/library/ssl_tls13_server.c
+++ b/library/ssl_tls13_server.c
@@ -1239,6 +1239,7 @@
     const unsigned char *cipher_suites_end;
     size_t extensions_len;
     const unsigned char *extensions_end;
+    mbedtls_ssl_handshake_params *handshake = ssl->handshake;
     int hrr_required = 0;
 
 #if defined(MBEDTLS_SSL_TLS1_3_KEY_EXCHANGE_MODE_SOME_PSK_ENABLED)
@@ -1304,7 +1305,7 @@
     MBEDTLS_SSL_DEBUG_BUF( 3, "client hello, random bytes",
                            p, MBEDTLS_CLIENT_HELLO_RANDOM_LEN );
 
-    memcpy( &ssl->handshake->randbytes[0], p, MBEDTLS_CLIENT_HELLO_RANDOM_LEN );
+    memcpy( &handshake->randbytes[0], p, MBEDTLS_CLIENT_HELLO_RANDOM_LEN );
     p += MBEDTLS_CLIENT_HELLO_RANDOM_LEN;
 
     /* ...
@@ -1374,13 +1375,13 @@
             continue;
 
         ssl->session_negotiate->ciphersuite = cipher_suite;
-        ssl->handshake->ciphersuite_info = ciphersuite_info;
+        handshake->ciphersuite_info = ciphersuite_info;
         MBEDTLS_SSL_DEBUG_MSG( 2, ( "selected ciphersuite: %04x - %s",
                                     cipher_suite,
                                     ciphersuite_info->name ) );
     }
 
-    if( ssl->handshake->ciphersuite_info == NULL )
+    if( handshake->ciphersuite_info == NULL )
     {
         MBEDTLS_SSL_PEND_FATAL_ALERT( MBEDTLS_SSL_ALERT_MSG_HANDSHAKE_FAILURE,
                                       MBEDTLS_ERR_SSL_HANDSHAKE_FAILURE );
@@ -1416,7 +1417,7 @@
 
     MBEDTLS_SSL_DEBUG_BUF( 3, "client hello extensions", p, extensions_len );
 
-    ssl->handshake->received_extensions = MBEDTLS_SSL_EXT_NONE;
+    handshake->received_extensions = MBEDTLS_SSL_EXT_NONE;
 
     while( p < extensions_end )
     {
@@ -1431,8 +1432,7 @@
          * Servers MUST check that it is the last extension and otherwise fail
          * the handshake with an "illegal_parameter" alert.
          */
-        if( ssl->handshake->received_extensions &
-            mbedtls_tls13_get_extension_mask( MBEDTLS_TLS_EXT_PRE_SHARED_KEY ) )
+        if( handshake->received_extensions & MBEDTLS_SSL_EXT_PRE_SHARED_KEY )
         {
             MBEDTLS_SSL_DEBUG_MSG(
                 3, ( "pre_shared_key is not last extension." ) );
@@ -1450,7 +1450,7 @@
         MBEDTLS_SSL_CHK_BUF_READ_PTR( p, extensions_end, extension_data_len );
         extension_data_end = p + extension_data_len;
 
-        ret = mbedtls_tls13_check_received_extensions(
+        ret = mbedtls_ssl_tls13_check_received_extension(
                   ssl, MBEDTLS_SSL_HS_CLIENT_HELLO, extension_type,
                   MBEDTLS_SSL_TLS1_3_ALLOWED_EXTS_OF_CH );
         if( ret != 0 )
@@ -1554,13 +1554,13 @@
 
             case MBEDTLS_TLS_EXT_PRE_SHARED_KEY:
                 MBEDTLS_SSL_DEBUG_MSG( 3, ( "found pre_shared_key extension" ) );
-                if( ( ssl->handshake->received_extensions &
+                if( ( handshake->received_extensions &
                       MBEDTLS_SSL_EXT_PSK_KEY_EXCHANGE_MODES ) == 0 )
                 {
                     MBEDTLS_SSL_PEND_FATAL_ALERT(
                         MBEDTLS_SSL_ALERT_MSG_ILLEGAL_PARAMETER,
-                        MBEDTLS_ERR_SSL_HANDSHAKE_FAILURE );
-                    return( MBEDTLS_ERR_SSL_HANDSHAKE_FAILURE );
+                        MBEDTLS_ERR_SSL_ILLEGAL_PARAMETER );
+                    return( MBEDTLS_ERR_SSL_ILLEGAL_PARAMETER );
                 }
 #if defined(MBEDTLS_SSL_TLS1_3_KEY_EXCHANGE_MODE_SOME_PSK_ENABLED)
                 /* Delay processing of the PSK identity once we have
@@ -1614,7 +1614,7 @@
     }
 
     MBEDTLS_SSL_TLS1_3_PRINT_EXTS(
-        3, MBEDTLS_SSL_HS_CLIENT_HELLO, ssl->handshake->received_extensions );
+        3, MBEDTLS_SSL_HS_CLIENT_HELLO, handshake->received_extensions );
 
     mbedtls_ssl_add_hs_hdr_to_checksum( ssl,
                                         MBEDTLS_SSL_HS_CLIENT_HELLO,
@@ -1628,10 +1628,9 @@
     /* If we've settled on a PSK-based exchange, parse PSK identity ext */
     if( mbedtls_ssl_tls13_some_psk_enabled( ssl ) &&
         mbedtls_ssl_conf_tls13_some_psk_enabled( ssl ) &&
-        ( ssl->handshake->received_extensions &
-          MBEDTLS_SSL_EXT_PRE_SHARED_KEY ) )
+        ( handshake->received_extensions & MBEDTLS_SSL_EXT_PRE_SHARED_KEY ) )
     {
-        ssl->handshake->update_checksum( ssl, buf,
+        handshake->update_checksum( ssl, buf,
                                          pre_shared_key_ext - buf );
         ret = ssl_tls13_parse_pre_shared_key_ext( ssl,
                                                   pre_shared_key_ext,
@@ -1640,8 +1639,7 @@
                                                   cipher_suites_end );
         if( ret == MBEDTLS_ERR_SSL_UNKNOWN_IDENTITY )
         {
-            ssl->handshake->received_extensions &=
-                                ~MBEDTLS_SSL_EXT_PRE_SHARED_KEY;
+            handshake->received_extensions &= ~MBEDTLS_SSL_EXT_PRE_SHARED_KEY;
         }
         else if( ret != 0 )
         {
@@ -1653,14 +1651,14 @@
     else
 #endif /* MBEDTLS_SSL_TLS1_3_KEY_EXCHANGE_MODE_SOME_PSK_ENABLED */
     {
-        ssl->handshake->update_checksum( ssl, buf, p - buf );
+        handshake->update_checksum( ssl, buf, p - buf );
     }
 
     ret = ssl_tls13_determine_key_exchange_mode( ssl );
     if( ret < 0 )
         return( ret );
 
-    mbedtls_ssl_optimize_checksum( ssl, ssl->handshake->ciphersuite_info );
+    mbedtls_ssl_optimize_checksum( ssl, handshake->ciphersuite_info );
 
     return( hrr_required ? SSL_CLIENT_HELLO_HRR_REQUIRED : SSL_CLIENT_HELLO_OK );
 }