Merge pull request #8857 from ronald-cron-arm/tls13-cli-max-early-data-size

TLS 1.3: Enforce max_early_data_size on client
diff --git a/library/ssl_msg.c b/library/ssl_msg.c
index 0c71157..4949306 100644
--- a/library/ssl_msg.c
+++ b/library/ssl_msg.c
@@ -6074,7 +6074,7 @@
 {
     int ret = MBEDTLS_ERR_ERROR_CORRUPTION_DETECTED;
     const struct mbedtls_ssl_config *conf;
-    int written_data_len = 0;
+    uint32_t remaining;
 
     MBEDTLS_SSL_DEBUG_MSG(2, ("=> write early_data"));
 
@@ -6124,18 +6124,27 @@
                 return ret;
             }
         }
+        remaining = ssl->session_negotiate->max_early_data_size;
     } else {
         /*
-         * If we are past the point where we can send early data, return
-         * immediatly. Otherwise, progress the handshake as much as possible to
-         * not delay it too much. If we reach a point where we can still send
-         * early data, then we will send some.
+         * If we are past the point where we can send early data or we have
+         * already reached the maximum early data size, return immediatly.
+         * Otherwise, progress the handshake as much as possible to not delay
+         * it too much. If we reach a point where we can still send early data,
+         * then we will send some.
          */
         if ((ssl->early_data_status != MBEDTLS_SSL_EARLY_DATA_STATUS_CAN_WRITE) &&
             (ssl->early_data_status != MBEDTLS_SSL_EARLY_DATA_STATUS_ACCEPTED)) {
             return MBEDTLS_ERR_SSL_CANNOT_WRITE_EARLY_DATA;
         }
 
+        remaining = ssl->session_negotiate->max_early_data_size -
+                    ssl->total_early_data_size;
+
+        if (remaining == 0) {
+            return MBEDTLS_ERR_SSL_CANNOT_WRITE_EARLY_DATA;
+        }
+
         ret = mbedtls_ssl_handshake(ssl);
         if ((ret != 0) && (ret != MBEDTLS_ERR_SSL_WANT_READ)) {
             MBEDTLS_SSL_DEBUG_RET(1, "mbedtls_ssl_handshake", ret);
@@ -6143,16 +6152,24 @@
         }
     }
 
-    if ((ssl->early_data_status != MBEDTLS_SSL_EARLY_DATA_STATUS_CAN_WRITE) &&
-        (ssl->early_data_status != MBEDTLS_SSL_EARLY_DATA_STATUS_ACCEPTED)) {
+    if (((ssl->early_data_status != MBEDTLS_SSL_EARLY_DATA_STATUS_CAN_WRITE) &&
+         (ssl->early_data_status != MBEDTLS_SSL_EARLY_DATA_STATUS_ACCEPTED))
+        || (remaining == 0)) {
         return MBEDTLS_ERR_SSL_CANNOT_WRITE_EARLY_DATA;
     }
 
-    written_data_len = ssl_write_real(ssl, buf, len);
+    if (len > remaining) {
+        len = remaining;
+    }
 
-    MBEDTLS_SSL_DEBUG_MSG(2, ("<= write early_data, len=%d", written_data_len));
+    ret = ssl_write_real(ssl, buf, len);
+    if (ret >= 0) {
+        ssl->total_early_data_size += ret;
+    }
 
-    return written_data_len;
+    MBEDTLS_SSL_DEBUG_MSG(2, ("<= write early_data, ret=%d", ret));
+
+    return ret;
 }
 #endif /* MBEDTLS_SSL_EARLY_DATA && MBEDTLS_SSL_CLI_C */
 
diff --git a/tests/suites/test_suite_ssl.data b/tests/suites/test_suite_ssl.data
index c0c816c..e95aada 100644
--- a/tests/suites/test_suite_ssl.data
+++ b/tests/suites/test_suite_ssl.data
@@ -3310,6 +3310,18 @@
 TLS 1.3 write early data, hello retry request
 tls13_write_early_data:TEST_EARLY_DATA_HRR
 
+TLS 1.3 cli, maximum early data size, default size
+tls13_cli_max_early_data_size:-1
+
+TLS 1.3 cli, maximum early data size, zero
+tls13_cli_max_early_data_size:0
+
+TLS 1.3 cli, maximum early data size, very small but not 0
+tls13_cli_max_early_data_size:3
+
+TLS 1.3 cli, maximum early data size, 93
+tls13_cli_max_early_data_size:93
+
 TLS 1.3 srv, max early data size, dflt, wsz=96
 tls13_srv_max_early_data_size:TEST_EARLY_DATA_ACCEPTED:-1:96
 
diff --git a/tests/suites/test_suite_ssl.function b/tests/suites/test_suite_ssl.function
index 6e95817..0c9ccf1 100644
--- a/tests/suites/test_suite_ssl.function
+++ b/tests/suites/test_suite_ssl.function
@@ -4503,6 +4503,154 @@
 }
 /* END_CASE */
 
+/* BEGIN_CASE depends_on:MBEDTLS_SSL_EARLY_DATA:MBEDTLS_SSL_CLI_C:MBEDTLS_SSL_SRV_C:MBEDTLS_DEBUG_C:MBEDTLS_TEST_AT_LEAST_ONE_TLS1_3_CIPHERSUITE:MBEDTLS_SSL_TLS1_3_KEY_EXCHANGE_MODE_EPHEMERAL_ENABLED:MBEDTLS_SSL_TLS1_3_KEY_EXCHANGE_MODE_PSK_EPHEMERAL_ENABLED:MBEDTLS_MD_CAN_SHA256:MBEDTLS_ECP_HAVE_SECP256R1:MBEDTLS_ECP_HAVE_SECP384R1:MBEDTLS_PK_CAN_ECDSA_VERIFY:MBEDTLS_SSL_SESSION_TICKETS */
+void tls13_cli_max_early_data_size(int max_early_data_size_arg)
+{
+    int ret = -1;
+    mbedtls_test_ssl_endpoint client_ep, server_ep;
+    mbedtls_test_handshake_test_options client_options;
+    mbedtls_test_handshake_test_options server_options;
+    mbedtls_ssl_session saved_session;
+    unsigned char *buf = NULL;
+    uint32_t buf_size = 64;
+    uint32_t max_early_data_size;
+    uint32_t written_early_data_size = 0;
+    uint32_t read_early_data_size = 0;
+
+    mbedtls_platform_zeroize(&client_ep, sizeof(client_ep));
+    mbedtls_platform_zeroize(&server_ep, sizeof(server_ep));
+    mbedtls_test_init_handshake_options(&client_options);
+    mbedtls_test_init_handshake_options(&server_options);
+    mbedtls_ssl_session_init(&saved_session);
+
+    PSA_INIT();
+    TEST_CALLOC(buf, buf_size);
+
+    /*
+     * Run first handshake to get a ticket from the server.
+     */
+
+    client_options.pk_alg = MBEDTLS_PK_ECDSA;
+    client_options.early_data = MBEDTLS_SSL_EARLY_DATA_ENABLED;
+    server_options.pk_alg = MBEDTLS_PK_ECDSA;
+    server_options.early_data = MBEDTLS_SSL_EARLY_DATA_ENABLED;
+    server_options.max_early_data_size = max_early_data_size_arg;
+
+    ret = mbedtls_test_get_tls13_ticket(&client_options, &server_options,
+                                        &saved_session);
+    TEST_EQUAL(ret, 0);
+
+    /*
+     * Prepare for handshake with the ticket.
+     */
+    ret = mbedtls_test_ssl_endpoint_init(&client_ep, MBEDTLS_SSL_IS_CLIENT,
+                                         &client_options, NULL, NULL, NULL);
+    TEST_EQUAL(ret, 0);
+
+    ret = mbedtls_test_ssl_endpoint_init(&server_ep, MBEDTLS_SSL_IS_SERVER,
+                                         &server_options, NULL, NULL, NULL);
+    TEST_EQUAL(ret, 0);
+
+    mbedtls_ssl_conf_session_tickets_cb(&server_ep.conf,
+                                        mbedtls_test_ticket_write,
+                                        mbedtls_test_ticket_parse,
+                                        NULL);
+
+    max_early_data_size = saved_session.max_early_data_size;
+    /*
+     * (max_early_data_size + 1024) for the size of the socket buffers for the
+     * server one to be able to contain the maximum number of early data bytes
+     * plus the first flight of client messages. Needed because we cannot
+     * initiate the handshake on server side before doing all the calls to
+     * mbedtls_ssl_write_early_data() we want to test. See below for more
+     * information.
+     */
+    ret = mbedtls_test_mock_socket_connect(&(client_ep.socket),
+                                           &(server_ep.socket),
+                                           max_early_data_size + 1024);
+    TEST_EQUAL(ret, 0);
+
+    /* If our server is configured with max_early_data_size equal to zero, it
+     * does not set the MBEDTLS_SSL_TLS1_3_TICKET_ALLOW_EARLY_DATA flag for
+     * the tickets it creates. To be able to test early data with a ticket
+     * allowing early data in its flags but with max_early_data_size equal to
+     * zero (case supported by our client) tweak the ticket flags here.
+     */
+    if (max_early_data_size == 0) {
+        saved_session.ticket_flags |= MBEDTLS_SSL_TLS1_3_TICKET_ALLOW_EARLY_DATA;
+    }
+
+    ret = mbedtls_ssl_set_session(&(client_ep.ssl), &saved_session);
+    TEST_EQUAL(ret, 0);
+
+    while (written_early_data_size < max_early_data_size) {
+        uint32_t remaining = max_early_data_size - written_early_data_size;
+
+        for (size_t i = 0; i < buf_size; i++) {
+            buf[i] = (unsigned char) (written_early_data_size + i);
+        }
+
+        ret = mbedtls_ssl_write_early_data(&(client_ep.ssl),
+                                           buf,
+                                           buf_size);
+
+        if (buf_size <= remaining) {
+            TEST_EQUAL(ret, buf_size);
+        } else {
+            TEST_EQUAL(ret, remaining);
+        }
+        written_early_data_size += buf_size;
+    }
+    TEST_EQUAL(client_ep.ssl.total_early_data_size, max_early_data_size);
+
+    ret = mbedtls_ssl_write_early_data(&(client_ep.ssl), buf, 1);
+    TEST_EQUAL(ret, MBEDTLS_ERR_SSL_CANNOT_WRITE_EARLY_DATA);
+    TEST_EQUAL(client_ep.ssl.total_early_data_size, max_early_data_size);
+    TEST_EQUAL(client_ep.ssl.early_data_status,
+               MBEDTLS_SSL_EARLY_DATA_STATUS_CAN_WRITE);
+
+    /*
+     * Now, check data on server side. It is not done in the previous loop as
+     * in the first call to mbedtls_ssl_handshake(), the server ends up sending
+     * its Finished message and then in the following call to
+     * mbedtls_ssl_write_early_data() we go past the early data writing window
+     * and we cannot test multiple calls to the API is this writing window.
+     */
+    while (read_early_data_size < max_early_data_size) {
+        ret = mbedtls_ssl_handshake(&(server_ep.ssl));
+        TEST_EQUAL(ret, MBEDTLS_ERR_SSL_RECEIVED_EARLY_DATA);
+
+        ret = mbedtls_ssl_read_early_data(&(server_ep.ssl),
+                                          buf,
+                                          buf_size);
+        TEST_ASSERT(ret > 0);
+
+        for (size_t i = 0; i < (size_t) ret; i++) {
+            TEST_EQUAL(buf[i], (unsigned char) (read_early_data_size + i));
+        }
+
+        read_early_data_size += ret;
+    }
+    TEST_EQUAL(read_early_data_size, max_early_data_size);
+
+    ret = mbedtls_ssl_handshake(&(server_ep.ssl));
+    TEST_EQUAL(ret, MBEDTLS_ERR_SSL_WANT_READ);
+
+    TEST_ASSERT(mbedtls_test_move_handshake_to_state(
+                    &(client_ep.ssl), &(server_ep.ssl), MBEDTLS_SSL_HANDSHAKE_OVER)
+                ==  0);
+
+exit:
+    mbedtls_test_ssl_endpoint_free(&client_ep, NULL);
+    mbedtls_test_ssl_endpoint_free(&server_ep, NULL);
+    mbedtls_test_free_handshake_options(&client_options);
+    mbedtls_test_free_handshake_options(&server_options);
+    mbedtls_ssl_session_free(&saved_session);
+    mbedtls_free(buf);
+    PSA_DONE();
+}
+/* END_CASE */
+
 /*
  * The !MBEDTLS_SSL_PROTO_TLS1_2 dependency of tls13_early_data() below is
  * a temporary workaround to not run the test in Windows-2013 where there is