fix various issues

Signed-off-by: Jerry Yu <jerry.h.yu@arm.com>
diff --git a/library/ssl_tls13_server.c b/library/ssl_tls13_server.c
index 72ac8ad..cdbc8a7 100644
--- a/library/ssl_tls13_server.c
+++ b/library/ssl_tls13_server.c
@@ -24,184 +24,6 @@
 #include "mbedtls/debug.h"
 #include "mbedtls/error.h"
 #include "mbedtls/platform.h"
-#include <string.h>
-
-#include "ssl_misc.h"
-#include "ssl_tls13_keys.h"
-#include "ssl_debug_helpers.h"
-#include "ecdh_misc.h"
-#include "ssl_tls13_keys.h"
-
-#if defined(MBEDTLS_SSL_SERVER_NAME_INDICATION)
-/*
- * struct {
- *      ExtensionType extension_type;
- *      opaque extension_data<0..2^16-1>;
- *  } Extension;
- */
-static int ssl_tls13_write_sni_server_ext( mbedtls_ssl_context *ssl,
-                                           unsigned char *buf,
-                                           unsigned char *end,
-                                           size_t *out_len )
-{
-    unsigned char *p = buf;
-    *out_len = 0;
-
-    if( ( ssl->handshake->extensions_present &
-          MBEDTLS_SSL_EXT_SERVERNAME ) == 0 )
-    {
-        return( 0 );
-    }
-
-    MBEDTLS_SSL_DEBUG_MSG( 3, ( "adding server_name extension" ) );
-
-    /* Check if we have enough space:
-     * - ExtensionType      (2 bytes)
-     * - extensions length  (2 bytes)
-     */
-    MBEDTLS_SSL_CHK_BUF_PTR( p, end, 4);
-
-    /* Write extension type */
-    MBEDTLS_PUT_UINT16_BE( MBEDTLS_TLS_EXT_SERVERNAME, p, 0 );
-
-    /* Write extension length */
-    MBEDTLS_PUT_UINT16_BE( 0, p, 2 );
-
-    *out_len = 4;
-
-    return( 0 );
-}
-#endif /* MBEDTLS_SSL_SERVER_NAME_INDICATION */
-
-
-static int ssl_tls13_prepare_encrypted_extensions( mbedtls_ssl_context *ssl )
-{
-    int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
-    mbedtls_ssl_key_set traffic_keys;
-    mbedtls_ssl_transform *transform_handshake = NULL;
-
-    /* Compute handshake secret */
-    ret = mbedtls_ssl_tls13_key_schedule_stage_handshake( ssl );
-    if( ret != 0 )
-    {
-        MBEDTLS_SSL_DEBUG_RET( 1, "mbedtls_ssl_tls13_key_schedule_stage_handshake", ret );
-        return( ret );
-    }
-
-    /* Derive handshake key material */
-    ret = mbedtls_ssl_tls13_generate_handshake_keys( ssl, &traffic_keys );
-    if( ret != 0 )
-    {
-        MBEDTLS_SSL_DEBUG_RET( 1,
-                "mbedtls_ssl_tls13_generate_handshake_keys", ret );
-        return( ret );
-    }
-
-    transform_handshake = mbedtls_calloc( 1, sizeof( mbedtls_ssl_transform ) );
-    if( transform_handshake == NULL )
-        return( MBEDTLS_ERR_SSL_ALLOC_FAILED );
-
-    /* Setup transform from handshake key material */
-    ret = mbedtls_ssl_tls13_populate_transform(
-                               transform_handshake,
-                               ssl->conf->endpoint,
-                               ssl->session_negotiate->ciphersuite,
-                               &traffic_keys,
-                               ssl );
-    if( ret != 0 )
-    {
-        MBEDTLS_SSL_DEBUG_RET( 1, "mbedtls_ssl_tls13_populate_transform", ret );
-        mbedtls_free( transform_handshake );
-        return( ret );
-    }
-
-    ssl->handshake->transform_handshake = transform_handshake;
-    mbedtls_ssl_set_outbound_transform( ssl, ssl->handshake->transform_handshake );
-
-    /*
-     * Switch to our negotiated transform and session parameters for outbound
-     * data.
-     */
-    MBEDTLS_SSL_DEBUG_MSG( 3, ( "switching to new transform spec for outbound data" ) );
-    memset( ssl->out_ctr, 0, 8 );
-
-    return( 0 );
-}
-
-/*
- * struct {
- *    Extension extensions<0..2 ^ 16 - 1>;
- * } EncryptedExtensions;
- *
- */
-static int ssl_tls13_write_encrypted_extensions( mbedtls_ssl_context *ssl,
-                                                 unsigned char *buf,
-                                                 unsigned char *end,
-                                                 size_t *out_len )
-{
-    unsigned char *p = buf;
-    int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
-    size_t extensions_len = 0;
-    unsigned char *p_extensions_len;
-    size_t length = 0; // length of this extension
-    *out_len = 0;
-
-    /* Skip extension length; first write extensions, then update length */
-    p_extensions_len = p;
-    p += 2;
-
-#if defined(MBEDTLS_SSL_SERVER_NAME_INDICATION)
-    ret = ssl_tls13_write_sni_server_ext( ssl, p, end, &length );
-    if( ret != 0 )
-        return( ret );
-    p += length;
-#endif /* MBEDTLS_SSL_SERVER_NAME_INDICATION */
-
-    *out_len = (size_t)( p - buf );
-
-    /* write extensions length */
-    extensions_len = ( p - p_extensions_len ) - 2;
-    MBEDTLS_PUT_UINT16_BE( extensions_len, p_extensions_len, 0);
-
-    return( 0 );
-}
-
-static int ssl_tls13_process_encrypted_extensions( mbedtls_ssl_context *ssl )
-{
-    int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
-    unsigned char *buf;
-    size_t buf_len = 0;
-    size_t msg_len = 0;
-
-    MBEDTLS_SSL_DEBUG_MSG( 2, ( "=> write encrypted extension" ) );
-
-    if( ssl->handshake->state_local.encrypted_extensions_out.preparation_done == 0 )
-    {
-        MBEDTLS_SSL_PROC_CHK( ssl_tls13_prepare_encrypted_extensions( ssl ) );
-        ssl->handshake->state_local.encrypted_extensions_out.preparation_done = 1;
-    }
-
-    MBEDTLS_SSL_PROC_CHK( mbedtls_ssl_start_handshake_msg( ssl,
-                       MBEDTLS_SSL_HS_ENCRYPTED_EXTENSIONS, &buf, &buf_len ) );
-
-    MBEDTLS_SSL_PROC_CHK( ssl_tls13_write_encrypted_extensions(
-                              ssl, buf, buf + buf_len, &msg_len ) );
-
-    mbedtls_ssl_add_hs_msg_to_checksum(
-        ssl, MBEDTLS_SSL_HS_ENCRYPTED_EXTENSIONS, buf, msg_len );
-
-    /* Update state */
-    mbedtls_ssl_handshake_set_state( ssl, MBEDTLS_SSL_CERTIFICATE_REQUEST );
-
-    MBEDTLS_SSL_PROC_CHK( mbedtls_ssl_finish_handshake_msg(
-                              ssl, buf_len, msg_len ) );
-
-cleanup:
-
-    MBEDTLS_SSL_DEBUG_MSG( 2, ( "<= write encrypted extension" ) );
-    return( ret );
-}
-
 
 #if defined(MBEDTLS_ECP_C)
 #include "mbedtls/ecp.h"
@@ -215,6 +37,10 @@
 #define mbedtls_free       free
 #endif /* MBEDTLS_PLATFORM_C */
 
+#include "ssl_misc.h"
+#include "ssl_tls13_keys.h"
+#include "ssl_debug_helpers.h"
+
 /* From RFC 8446:
  *   struct {
  *          ProtocolVersion versions<2..254>;
@@ -1231,6 +1057,129 @@
 }
 
 /*
+ * Handler for MBEDTLS_SSL_ENCRYPTED_EXTENSIONS
+ */
+static int ssl_tls13_prepare_encrypted_extensions( mbedtls_ssl_context *ssl )
+{
+    int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
+    mbedtls_ssl_key_set traffic_keys;
+    mbedtls_ssl_transform *transform_handshake = NULL;
+
+    /* Compute handshake secret */
+    ret = mbedtls_ssl_tls13_key_schedule_stage_handshake( ssl );
+    if( ret != 0 )
+    {
+        MBEDTLS_SSL_DEBUG_RET(
+            1, "mbedtls_ssl_tls13_key_schedule_stage_handshake", ret );
+        return( ret );
+    }
+
+    /* Derive handshake key material */
+    ret = mbedtls_ssl_tls13_generate_handshake_keys( ssl, &traffic_keys );
+    if( ret != 0 )
+    {
+        MBEDTLS_SSL_DEBUG_RET(
+            1, "mbedtls_ssl_tls13_generate_handshake_keys", ret );
+        return( ret );
+    }
+
+    transform_handshake = mbedtls_calloc( 1, sizeof( mbedtls_ssl_transform ) );
+    if( transform_handshake == NULL )
+        return( MBEDTLS_ERR_SSL_ALLOC_FAILED );
+
+    /* Setup transform from handshake key material */
+    ret = mbedtls_ssl_tls13_populate_transform(
+                               transform_handshake,
+                               ssl->conf->endpoint,
+                               ssl->session_negotiate->ciphersuite,
+                               &traffic_keys,
+                               ssl );
+    if( ret != 0 )
+    {
+        MBEDTLS_SSL_DEBUG_RET( 1, "mbedtls_ssl_tls13_populate_transform", ret );
+        mbedtls_free( transform_handshake );
+        return( ret );
+    }
+
+    ssl->handshake->transform_handshake = transform_handshake;
+    mbedtls_ssl_set_outbound_transform( ssl, ssl->handshake->transform_handshake );
+
+    /*
+     * Switch to our negotiated transform and session parameters for outbound
+     * data.
+     */
+    MBEDTLS_SSL_DEBUG_MSG(
+        3, ( "switching to new transform spec for outbound data" ) );
+    memset( ssl->out_ctr, 0, 8 );
+
+    return( 0 );
+}
+
+/*
+ * struct {
+ *    Extension extensions<0..2 ^ 16 - 1>;
+ * } EncryptedExtensions;
+ *
+ */
+static int ssl_tls13_write_encrypted_extensions_body( mbedtls_ssl_context *ssl,
+                                                      unsigned char *buf,
+                                                      unsigned char *end,
+                                                      size_t *out_len )
+{
+    ((void) ssl);
+    unsigned char *p = buf;
+    size_t extensions_len = 0;
+    unsigned char *extensions_start;
+    *out_len = 0;
+
+    MBEDTLS_SSL_CHK_BUF_PTR( p, end, 2 );
+    extensions_start = p;
+    p += 2;
+
+    *out_len = (size_t)( p - buf );
+
+    /* write extensions length */
+    extensions_len = ( p - extensions_start ) - 2;
+    MBEDTLS_PUT_UINT16_BE( extensions_len, extensions_start, 0);
+
+    MBEDTLS_SSL_DEBUG_BUF( 4, "encrypted extensions", buf, *out_len );
+
+    return( 0 );
+}
+
+static int ssl_tls13_write_encrypted_extensions( mbedtls_ssl_context *ssl )
+{
+    int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
+    unsigned char *buf;
+    size_t buf_len = 0;
+    size_t msg_len = 0;
+
+    MBEDTLS_SSL_DEBUG_MSG( 2, ( "=> write encrypted extension" ) );
+
+    MBEDTLS_SSL_PROC_CHK( ssl_tls13_prepare_encrypted_extensions( ssl ) );
+
+    MBEDTLS_SSL_PROC_CHK( mbedtls_ssl_start_handshake_msg( ssl,
+                       MBEDTLS_SSL_HS_ENCRYPTED_EXTENSIONS, &buf, &buf_len ) );
+
+    MBEDTLS_SSL_PROC_CHK( ssl_tls13_write_encrypted_extensions_body(
+                              ssl, buf, buf + buf_len, &msg_len ) );
+
+    mbedtls_ssl_add_hs_msg_to_checksum(
+        ssl, MBEDTLS_SSL_HS_ENCRYPTED_EXTENSIONS, buf, msg_len );
+
+    /* Update state */
+    mbedtls_ssl_handshake_set_state( ssl, MBEDTLS_SSL_SERVER_CERTIFICATE );
+
+    MBEDTLS_SSL_PROC_CHK( mbedtls_ssl_finish_handshake_msg(
+                              ssl, buf_len, msg_len ) );
+
+cleanup:
+
+    MBEDTLS_SSL_DEBUG_MSG( 2, ( "<= write encrypted extension" ) );
+    return( ret );
+}
+
+/*
  * TLS 1.3 State Machine -- server side
  */
 int mbedtls_ssl_tls13_handshake_server_step( mbedtls_ssl_context *ssl )
@@ -1266,10 +1215,11 @@
             break;
 
         case MBEDTLS_SSL_ENCRYPTED_EXTENSIONS:
-            ret = ssl_tls13_process_encrypted_extensions( ssl );
+            ret = ssl_tls13_write_encrypted_extensions( ssl );
             if( ret != 0 )
             {
-                MBEDTLS_SSL_DEBUG_RET( 1, "ssl_tls13_process_encrypted_extensions", ret );
+                MBEDTLS_SSL_DEBUG_RET( 1, "ssl_tls13_write_encrypted_extensions", ret );
+                return( ret );
             }
             break;