share write_early_data_ext function

Signed-off-by: Jerry Yu <jerry.h.yu@arm.com>
diff --git a/library/ssl_misc.h b/library/ssl_misc.h
index 36f332f..8c3da49 100644
--- a/library/ssl_misc.h
+++ b/library/ssl_misc.h
@@ -2115,7 +2115,8 @@
 int mbedtls_ssl_tls13_write_early_data_ext(mbedtls_ssl_context *ssl,
                                            unsigned char *buf,
                                            const unsigned char *end,
-                                           size_t *out_len);
+                                           size_t *out_len,
+                                           const mbedtls_ssl_session *session);
 
 #if defined(MBEDTLS_SSL_SRV_C)
 #define MBEDTLS_SSL_EARLY_DATA_STATUS_NOT_RECEIVED \
diff --git a/library/ssl_tls13_client.c b/library/ssl_tls13_client.c
index bc8b161..fa6c4c6 100644
--- a/library/ssl_tls13_client.c
+++ b/library/ssl_tls13_client.c
@@ -1174,7 +1174,9 @@
     if (mbedtls_ssl_conf_tls13_some_psk_enabled(ssl) &&
         ssl_tls13_early_data_has_valid_ticket(ssl) &&
         ssl->conf->early_data_enabled == MBEDTLS_SSL_EARLY_DATA_ENABLED) {
-        ret = mbedtls_ssl_tls13_write_early_data_ext(ssl, p, end, &ext_len);
+
+        ret = mbedtls_ssl_tls13_write_early_data_ext(
+            ssl, p, end, &ext_len, NULL);
         if (ret != 0) {
             return ret;
         }
diff --git a/library/ssl_tls13_generic.c b/library/ssl_tls13_generic.c
index cc77a94..938bf80 100644
--- a/library/ssl_tls13_generic.c
+++ b/library/ssl_tls13_generic.c
@@ -1402,28 +1402,48 @@
  *
  * struct {
  *   select ( Handshake.msg_type ) {
- *     ...
+ *     case new_session_ticket:   uint32 max_early_data_size;
  *     case client_hello:         Empty;
  *     case encrypted_extensions: Empty;
  *   };
  * } EarlyDataIndication;
+ *
+ * We use `mbedtls_ssl_is_handshake_over()` to decide if `max_early_data_size`
+ * should be sent for `new_session_ticket` is post-handshake message.
  */
 #if defined(MBEDTLS_SSL_EARLY_DATA)
 int mbedtls_ssl_tls13_write_early_data_ext(mbedtls_ssl_context *ssl,
                                            unsigned char *buf,
                                            const unsigned char *end,
-                                           size_t *out_len)
+                                           size_t *out_len,
+                                           const mbedtls_ssl_session *session)
 {
     unsigned char *p = buf;
-    *out_len = 0;
-    ((void) ssl);
 
-    MBEDTLS_SSL_CHK_BUF_PTR(p, end, 4);
+#if defined(MBEDTLS_SSL_SRV_C)
+    const size_t needed = session != NULL ? 8 : 4;
+#else
+    const size_t needed = 4;
+    ((void) session);
+#endif
+
+    *out_len = 0;
+
+    MBEDTLS_SSL_CHK_BUF_PTR(p, end, needed);
 
     MBEDTLS_PUT_UINT16_BE(MBEDTLS_TLS_EXT_EARLY_DATA, p, 0);
-    MBEDTLS_PUT_UINT16_BE(0, p, 2);
+    MBEDTLS_PUT_UINT16_BE(needed - 4, p, 2);
 
-    *out_len = 4;
+#if defined(MBEDTLS_SSL_SRV_C)
+    if (session != NULL) {
+        MBEDTLS_PUT_UINT32_BE(session->max_early_data_size, p, 4);
+        MBEDTLS_SSL_DEBUG_MSG(
+            4, ("Sent max_early_data_size=%u",
+                (unsigned int) session->max_early_data_size));
+    }
+#endif
+
+    *out_len = needed;
 
     mbedtls_ssl_tls13_set_hs_sent_ext_mask(ssl, MBEDTLS_TLS_EXT_EARLY_DATA);
 
diff --git a/library/ssl_tls13_server.c b/library/ssl_tls13_server.c
index 133245b..addbbe1 100644
--- a/library/ssl_tls13_server.c
+++ b/library/ssl_tls13_server.c
@@ -2524,7 +2524,8 @@
 
 #if defined(MBEDTLS_SSL_EARLY_DATA)
     if (ssl->early_data_status == MBEDTLS_SSL_EARLY_DATA_STATUS_ACCEPTED) {
-        ret = mbedtls_ssl_tls13_write_early_data_ext(ssl, p, end, &output_len);
+        ret = mbedtls_ssl_tls13_write_early_data_ext(
+            ssl, p, end, &output_len, NULL);
         if (ret != 0) {
             return ret;
         }
@@ -3202,49 +3203,6 @@
     return 0;
 }
 
-#if defined(MBEDTLS_SSL_EARLY_DATA)
-/* RFC 8446 section 4.2.10
- *
- * struct {
- *     select (Handshake.msg_type) {
- *         case new_session_ticket:   uint32 max_early_data_size;
- *         ...
- *     };
- * } EarlyDataIndication;
- */
-MBEDTLS_CHECK_RETURN_CRITICAL
-static int ssl_tls13_write_nst_early_data_ext(mbedtls_ssl_context *ssl,
-                                              unsigned char *buf,
-                                              const unsigned char *end,
-                                              size_t *out_len)
-{
-    unsigned char *p = buf;
-    *out_len = 0;
-
-    if (!mbedtls_ssl_session_ticket_allow_early_data(ssl->session)) {
-        MBEDTLS_SSL_DEBUG_MSG(
-            4, ("early_data not allowed, skip early_data extension in "
-                "NewSessionTicket"));
-        return 0;
-    }
-
-    MBEDTLS_SSL_CHK_BUF_PTR(p, end, 8);
-
-    MBEDTLS_PUT_UINT16_BE(MBEDTLS_TLS_EXT_EARLY_DATA, p, 0);
-    MBEDTLS_PUT_UINT16_BE(4, p, 2);
-    MBEDTLS_PUT_UINT32_BE(ssl->conf->max_early_data_size, p, 4);
-    MBEDTLS_SSL_DEBUG_MSG(
-        4, ("Sent max_early_data_size=%u",
-            (unsigned int) ssl->conf->max_early_data_size));
-
-    *out_len = 8;
-
-    mbedtls_ssl_tls13_set_hs_sent_ext_mask(ssl, MBEDTLS_TLS_EXT_EARLY_DATA);
-
-    return 0;
-}
-#endif /* MBEDTLS_SSL_EARLY_DATA */
-
 /* This function creates a NewSessionTicket message in the following format:
  *
  * struct {
@@ -3371,7 +3329,7 @@
     if (ssl->conf->early_data_enabled == MBEDTLS_SSL_EARLY_DATA_ENABLED &&
         ssl->conf->max_early_data_size > 0) {
         if ((ret = mbedtls_ssl_tls13_write_early_data_ext(
-                 ssl, p, end, &output_len)) != 0) {
+                 ssl, p, end, &output_len, session)) != 0) {
             MBEDTLS_SSL_DEBUG_RET(
                 1, "mbedtls_ssl_tls13_write_early_data_ext", ret);
             return ret;