Address comments of various issues

Improve comments
Change coding style
Rename functions

Change-Id: Ia111aef303932cfeee693431c3d48f90342b32e5
Signed-off-by: XiaokangQian <xiaokang.qian@arm.com>
diff --git a/library/ssl_client.c b/library/ssl_client.c
index afe07e8..ee5a78a 100644
--- a/library/ssl_client.c
+++ b/library/ssl_client.c
@@ -645,6 +645,9 @@
 
 #if defined(MBEDTLS_SSL_PROTO_TLS1_3) && \
     defined(MBEDTLS_KEY_EXCHANGE_SOME_PSK_ENABLED)
+    /* The "pre_shared_key" extension (RFC 8446 Section 4.2.11)
+     * MUST be the last extension in the ClientHello.
+     */
     if( propose_tls13 && mbedtls_ssl_conf_tls13_some_psk_enabled( ssl ) )
     {
         ret = mbedtls_ssl_tls13_write_pre_shared_key_ext_without_binders(
@@ -907,11 +910,11 @@
         if( binders_len > 0 )
         {
             MBEDTLS_SSL_PROC_CHK(
-                mbedtls_ssl_tls13_write_pre_shared_key_ext_binders(
+                mbedtls_ssl_tls13_write_binders_of_pre_shared_key_ext(
                       ssl, buf + msg_len - binders_len, buf + msg_len ) );
+            ssl->handshake->update_checksum( ssl, buf + msg_len - binders_len,
+                                             binders_len );
         }
-        ssl->handshake->update_checksum( ssl, buf + msg_len - binders_len,
-                                         binders_len );
 #endif /* MBEDTLS_SSL_PROTO_TLS1_3 && MBEDTLS_KEY_EXCHANGE_SOME_PSK_ENABLED */
 
         MBEDTLS_SSL_PROC_CHK( mbedtls_ssl_finish_handshake_msg( ssl,
diff --git a/library/ssl_misc.h b/library/ssl_misc.h
index f55dede..295ff9f 100644
--- a/library/ssl_misc.h
+++ b/library/ssl_misc.h
@@ -1336,9 +1336,6 @@
 void mbedtls_ssl_optimize_checksum( mbedtls_ssl_context *ssl,
                             const mbedtls_ssl_ciphersuite_t *ciphersuite_info );
 
-void mbedtls_ssl_add_hs_hdr_to_checksum( mbedtls_ssl_context *ssl,
-                                         unsigned hs_type,
-                                         size_t total_hs_len );
 /*
  * Update checksum of handshake messages.
  */
@@ -1347,6 +1344,10 @@
                                          unsigned char const *msg,
                                          size_t msg_len );
 
+void mbedtls_ssl_add_hs_hdr_to_checksum( mbedtls_ssl_context *ssl,
+                                         unsigned hs_type,
+                                         size_t total_hs_len );
+
 #if defined(MBEDTLS_KEY_EXCHANGE_SOME_PSK_ENABLED)
 #if !defined(MBEDTLS_USE_PSA_CRYPTO)
 MBEDTLS_CHECK_RETURN_CRITICAL
@@ -2426,6 +2427,7 @@
 /* Check if we have any PSK to offer, returns 0 if PSK is available.
  * Assign the psk and ticket if pointers are present.
  */
+MBEDTLS_CHECK_RETURN_CRITICAL
 int mbedtls_ssl_get_psk_to_offer(
         const mbedtls_ssl_context *ssl,
         int *psk_type,
@@ -2445,6 +2447,7 @@
  * \param[out]  binders_len Length of the binders to be written at the end of
  *                          the extension.
  */
+MBEDTLS_CHECK_RETURN_CRITICAL
 int mbedtls_ssl_tls13_write_pre_shared_key_ext_without_binders(
     mbedtls_ssl_context *ssl,
     unsigned char *buf, unsigned char *end,
@@ -2459,7 +2462,8 @@
  * \param[in]   buf     Base address of the buffer where to write the binders
  * \param[in]   end     End address of the buffer where to write the binders
  */
-int mbedtls_ssl_tls13_write_pre_shared_key_ext_binders(
+MBEDTLS_CHECK_RETURN_CRITICAL
+int mbedtls_ssl_tls13_write_binders_of_pre_shared_key_ext(
     mbedtls_ssl_context *ssl,
     unsigned char *buf, unsigned char *end );
 #endif /* MBEDTLS_KEY_EXCHANGE_SOME_PSK_ENABLED */
diff --git a/library/ssl_tls13_client.c b/library/ssl_tls13_client.c
index 43a8486..37e57e8 100644
--- a/library/ssl_tls13_client.c
+++ b/library/ssl_tls13_client.c
@@ -605,15 +605,12 @@
  *     PskKeyExchangeMode ke_modes<1..255>;
  * } PskKeyExchangeModes;
  */
+MBEDTLS_CHECK_RETURN_CRITICAL
 static int ssl_tls13_write_psk_key_exchange_modes_ext( mbedtls_ssl_context *ssl,
                                                        unsigned char *buf,
                                                        unsigned char *end,
                                                        size_t *out_len )
 {
-    const unsigned char *psk;
-    size_t psk_len;
-    const unsigned char *psk_identity;
-    size_t psk_identity_len;
     unsigned char *p = buf;
     int ke_modes_len = 0;
 
@@ -623,9 +620,7 @@
     /* Skip writing extension if no PSK key exchange mode
      * is enabled in the config or there is no PSK to offer.
      */
-    if( !mbedtls_ssl_conf_tls13_some_psk_enabled( ssl ) ||
-         mbedtls_ssl_get_psk_to_offer( ssl, NULL, &psk, &psk_len,
-                                      &psk_identity, &psk_identity_len ) != 0 )
+    if( !mbedtls_ssl_conf_tls13_some_psk_enabled( ssl ) )
     {
         MBEDTLS_SSL_DEBUG_MSG( 3, ( "skip psk_key_exchange_modes extension" ) );
         return( 0 );
@@ -682,10 +677,15 @@
  * opaque PskBinderEntry<32..255>;
  *
  * struct {
+ *   PskIdentity identities<7..2^16-1>;
+ *   PskBinderEntry binders<33..2^16-1>;
+ * } OfferedPsks;
  *
- *     PskIdentity identities<7..2^16-1>;
- *     PskBinderEntry binders<33..2^16-1>;
- *
+ * struct {
+ *   select (Handshake.msg_type) {
+ *      case client_hello: OfferedPsks;
+ *      ...
+ *   };
  * } PreSharedKeyExtension;
  *
  */
@@ -702,6 +702,7 @@
     size_t psk_len;
     const unsigned char *psk_identity;
     size_t psk_identity_len;
+    int psk_type;
     const mbedtls_ssl_ciphersuite_t *ciphersuite_info = NULL;
     const int *ciphersuites;
     psa_algorithm_t psa_hash_alg;
@@ -725,7 +726,7 @@
      * - Otherwise, skip the PSK extension.
      */
 
-    if( mbedtls_ssl_get_psk_to_offer( ssl, NULL, &psk, &psk_len,
+    if( mbedtls_ssl_get_psk_to_offer( ssl, &psk_type, &psk, &psk_len,
                                       &psk_identity, &psk_identity_len ) != 0 )
     {
         MBEDTLS_SSL_DEBUG_MSG( 3, ( "skip pre_shared_key extensions" ) );
@@ -736,19 +737,25 @@
      * Ciphersuite list
      */
     ciphersuites = ssl->conf->ciphersuite_list;
-    for ( int i = 0; ciphersuites[i] != 0; i++ )
+    if( psk_type == MBEDTLS_SSL_TLS1_3_PSK_EXTERNAL )
     {
-        ciphersuite_info = mbedtls_ssl_ciphersuite_from_id( ciphersuites[i] );
+        for( int i = 0; ciphersuites[i] != 0; i++ )
+        {
+            ciphersuite_info = mbedtls_ssl_ciphersuite_from_id(
+                                    ciphersuites[i] );
 
-        if( mbedtls_ssl_validate_ciphersuite(
+            if( mbedtls_ssl_validate_ciphersuite(
                                 ssl, ciphersuite_info,
                                 MBEDTLS_SSL_VERSION_TLS1_3,
                                 MBEDTLS_SSL_VERSION_TLS1_3 ) != 0 )
-            continue;
+                continue;
 
-        /* In this implementation we only add one pre-shared-key extension. */
-        ssl->session_negotiate->ciphersuite = ciphersuites[i];
-        break;
+            /* In this implementation we only add one pre-shared-key
+             * extension.
+             */
+            ssl->session_negotiate->ciphersuite = ciphersuites[i];
+            break;
+        }
     }
 
     ciphersuite_info = mbedtls_ssl_ciphersuite_from_id(
@@ -804,7 +811,7 @@
     return( 0 );
 }
 
-int mbedtls_ssl_tls13_write_pre_shared_key_ext_binders(
+int mbedtls_ssl_tls13_write_binders_of_pre_shared_key_ext(
     mbedtls_ssl_context *ssl,
     unsigned char *buf, unsigned char *end )
 {
@@ -847,11 +854,6 @@
     /* 1 bytes length field for next psk binder */
     *p++ = MBEDTLS_BYTE_0( hash_len );
 
-    if( ssl->handshake->resume == 1 )
-        psk_type = MBEDTLS_SSL_TLS1_3_PSK_RESUMPTION;
-    else
-        psk_type = MBEDTLS_SSL_TLS1_3_PSK_EXTERNAL;
-
     /* Get current state of handshake transcript. */
     ret = mbedtls_ssl_get_handshake_transcript( ssl, ciphersuite_info->mac,
                                                 transcript, sizeof( transcript ),
@@ -1234,15 +1236,19 @@
  *
  * struct {
  *
- *   uint16 selected_identity;
+ *   select (Handshake.msg_type) {
+ *         ...
+ *         case server_hello: uint16 selected_identity;
+ *   };
  *
  * } PreSharedKeyExtension;
  *
  */
 
-static int ssl_tls13_parse_server_psk_identity_ext( mbedtls_ssl_context *ssl,
-                                                    const unsigned char *buf,
-                                                    const unsigned char *end )
+MBEDTLS_CHECK_RETURN_CRITICAL
+static int ssl_tls13_parse_server_pre_shared_key_ext( mbedtls_ssl_context *ssl,
+                                                      const unsigned char *buf,
+                                                      const unsigned char *end )
 {
     int ret = 0;
     size_t selected_identity;
@@ -1524,11 +1530,11 @@
                     goto cleanup;
                 }
 
-                if( ( ret = ssl_tls13_parse_server_psk_identity_ext(
+                if( ( ret = ssl_tls13_parse_server_pre_shared_key_ext(
                                 ssl, p, extension_data_end ) ) != 0 )
                 {
                     MBEDTLS_SSL_DEBUG_RET(
-                        1, ( "ssl_tls13_parse_server_psk_identity_ext" ), ret );
+                        1, ( "ssl_tls13_parse_server_pre_shared_key_ext" ), ret );
                     return( ret );
                 }
                 break;
diff --git a/library/ssl_tls13_generic.c b/library/ssl_tls13_generic.c
index 4cd0d0e..9fd246d 100644
--- a/library/ssl_tls13_generic.c
+++ b/library/ssl_tls13_generic.c
@@ -1509,6 +1509,7 @@
 /* Check if we have any PSK to offer, returns 0 if PSK is available.
  * Assign the psk and ticket if pointers are present.
  */
+MBEDTLS_CHECK_RETURN_CRITICAL
 int mbedtls_ssl_get_psk_to_offer(
         const mbedtls_ssl_context *ssl,
         int *psk_type,
@@ -1517,7 +1518,7 @@
 {
     int ptrs_present = 0;
 
-    if( psk != NULL && psk_len != NULL &&
+    if( psk_type != NULL && psk != NULL && psk_len != NULL &&
         psk_identity != NULL && psk_identity_len != NULL )
     {
         ptrs_present = 1;
@@ -1528,14 +1529,13 @@
     {
         if( ptrs_present )
         {
+            *psk_type = MBEDTLS_SSL_TLS1_3_PSK_EXTERNAL;
             *psk = ssl->conf->psk;
             *psk_len = ssl->conf->psk_len;
             *psk_identity = ssl->conf->psk_identity;
             *psk_identity_len = ssl->conf->psk_identity_len;
         }
 
-        if( psk_type != NULL )
-            *psk_type = MBEDTLS_SSL_TLS1_3_PSK_EXTERNAL;
         return( 0 );
     }