Fix wrongly kex mode fallback issue in psk cases

Signed-off-by: Xiaokang Qian <xiaokang.qian@arm.com>
diff --git a/library/ssl_tls13_client.c b/library/ssl_tls13_client.c
index 714acb9..8af7035 100644
--- a/library/ssl_tls13_client.c
+++ b/library/ssl_tls13_client.c
@@ -1687,19 +1687,16 @@
         /* Only the pre_shared_key extension was received */
         case MBEDTLS_SSL_EXT_PRE_SHARED_KEY:
             handshake->key_exchange_mode = MBEDTLS_SSL_TLS1_3_KEY_EXCHANGE_MODE_PSK;
-            MBEDTLS_SSL_DEBUG_MSG( 2, ( "key exchange mode: psk" ) );
             break;
 
         /* Only the key_share extension was received */
         case MBEDTLS_SSL_EXT_KEY_SHARE:
             handshake->key_exchange_mode = MBEDTLS_SSL_TLS1_3_KEY_EXCHANGE_MODE_EPHEMERAL;
-            MBEDTLS_SSL_DEBUG_MSG( 2, ( "key exchange mode: ephemeral" ) );
             break;
 
         /* Both the pre_shared_key and key_share extensions were received */
         case ( MBEDTLS_SSL_EXT_PRE_SHARED_KEY | MBEDTLS_SSL_EXT_KEY_SHARE ):
             handshake->key_exchange_mode = MBEDTLS_SSL_TLS1_3_KEY_EXCHANGE_MODE_PSK_EPHEMERAL;
-            MBEDTLS_SSL_DEBUG_MSG( 2, ( "key exchange mode: psk_ephemeral" ) );
             break;
 
         /* Neither pre_shared_key nor key_share extension was received */
@@ -1711,10 +1708,25 @@
 
     MBEDTLS_SSL_DEBUG_MSG( 3,
             ( "Server selected key exchange mode: %s",
-              handshake->key_exchange_mode == 1 ? "psk" :
-              (handshake->key_exchange_mode == 2 ? "ephemeral" :
+              handshake->key_exchange_mode ==
+              MBEDTLS_SSL_TLS1_3_KEY_EXCHANGE_MODE_PSK ? "psk" :
+              (handshake->key_exchange_mode ==
+              MBEDTLS_SSL_TLS1_3_KEY_EXCHANGE_MODE_EPHEMERAL ? "ephemeral" :
               "psk_ephemeral")) );
 
+    if( !mbedtls_ssl_conf_tls13_check_kex_modes( ssl, handshake->key_exchange_mode ) )
+    {
+        ret = MBEDTLS_ERR_SSL_HANDSHAKE_FAILURE;
+        MBEDTLS_SSL_DEBUG_MSG( 2,
+                ( "Not supported kex mode in client: %s",
+              handshake->key_exchange_mode ==
+              MBEDTLS_SSL_TLS1_3_KEY_EXCHANGE_MODE_PSK ? "psk" :
+              (handshake->key_exchange_mode ==
+              MBEDTLS_SSL_TLS1_3_KEY_EXCHANGE_MODE_EPHEMERAL ? "ephemeral" :
+              "psk_ephemeral")) );
+        goto cleanup;
+    }
+
     /* Start the TLS 1.3 key schedule: Set the PSK and derive early secret.
      *
      * TODO: We don't have to do this in case we offered 0-RTT and the