Move state change to finalize client hello

Signed-off-by: Xiaokang Qian <xiaokang.qian@arm.com>
diff --git a/library/ssl_client.c b/library/ssl_client.c
index 62af0f9..2ad69f9 100644
--- a/library/ssl_client.c
+++ b/library/ssl_client.c
@@ -963,21 +963,16 @@
                                                               buf_len,
                                                               msg_len));
 
-#if defined(MBEDTLS_SSL_PROTO_TLS1_3)
-        if ((ssl->handshake->min_tls_version == MBEDTLS_SSL_VERSION_TLS1_3) &&
-            (ssl->tls_version == MBEDTLS_SSL_VERSION_TLS1_3)) {
-#if defined(MBEDTLS_SSL_TLS1_3_COMPATIBILITY_MODE)
-            mbedtls_ssl_handshake_set_state(
-                ssl, MBEDTLS_SSL_CLIENT_CCS_AFTER_CLIENT_HELLO);
-#else
+#if defined(MBEDTLS_SSL_PROTO_TLS1_2)
+        if (mbedtls_ssl_conf_is_tls12_only(ssl->conf)) {
             mbedtls_ssl_handshake_set_state(ssl, MBEDTLS_SSL_SERVER_HELLO);
-#endif /* MBEDTLS_SSL_TLS1_3_COMPATIBILITY_MODE */
         } else
 #endif
-        mbedtls_ssl_handshake_set_state(ssl, MBEDTLS_SSL_SERVER_HELLO);
+        {
 #if defined(MBEDTLS_SSL_PROTO_TLS1_3)
-        mbedtls_ssl_tls13_finalize_write_client_hello(ssl);
+            mbedtls_ssl_tls13_finalize_write_client_hello(ssl);
 #endif
+        }
 
     }
 
diff --git a/library/ssl_tls13_client.c b/library/ssl_tls13_client.c
index 55e5665..0c4a912 100644
--- a/library/ssl_tls13_client.c
+++ b/library/ssl_tls13_client.c
@@ -1234,6 +1234,13 @@
 int mbedtls_ssl_tls13_finalize_write_client_hello(mbedtls_ssl_context *ssl)
 {
     ((void) ssl);
+#if defined(MBEDTLS_SSL_TLS1_3_COMPATIBILITY_MODE)
+    mbedtls_ssl_handshake_set_state(
+        ssl, MBEDTLS_SSL_CLIENT_CCS_AFTER_CLIENT_HELLO);
+#else
+    mbedtls_ssl_handshake_set_state(ssl, MBEDTLS_SSL_SERVER_HELLO);
+#endif /* MBEDTLS_SSL_TLS1_3_COMPATIBILITY_MODE */
+
 #if defined(MBEDTLS_SSL_EARLY_DATA)
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
     psa_algorithm_t hash_alg = PSA_ALG_NONE;