tests: ssl: Improve tls13_srv_max_early_data_size()

Signed-off-by: Ronald Cron <ronald.cron@arm.com>
diff --git a/tests/suites/test_suite_ssl.function b/tests/suites/test_suite_ssl.function
index 791386b..d89d8ac 100644
--- a/tests/suites/test_suite_ssl.function
+++ b/tests/suites/test_suite_ssl.function
@@ -4455,7 +4455,7 @@
  * an issue with mbedtls_vsnprintf().
  */
 /* BEGIN_CASE depends_on:!MBEDTLS_SSL_PROTO_TLS1_2: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_srv_max_early_data_size(int scenario, int max_early_data_size_arg, int reach_max)
+void tls13_srv_max_early_data_size(int scenario, int max_early_data_size_arg, int write_size_arg)
 {
     int ret = -1;
     mbedtls_test_ssl_endpoint client_ep, server_ep;
@@ -4469,11 +4469,12 @@
         MBEDTLS_SSL_IANA_TLS_GROUP_NONE
     };
     char pattern[128];
-    unsigned char buf_write[64];
-    size_t early_data_len;
+    unsigned char *buf_write = NULL;
+    uint32_t write_size = (uint32_t) write_size_arg;
+    unsigned char *buf_read = NULL;
+    uint32_t read_size;
     uint32_t expended_early_data_len = 0;
     uint32_t written_early_data_size = 0;
-    int write_early_data_flag = 1;
     uint32_t max_early_data_size;
 
     mbedtls_platform_zeroize(&client_ep, sizeof(client_ep));
@@ -4481,16 +4482,16 @@
     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_write, write_size);
 
     /*
-     * Reach maximum early data size exactly option only available in case of
-     * TEST_EARLY_DATA_ACCEPTED scenario.
+     * Allocate a smaller buffer for early data reading to exercise the reading
+     * of data in one record in multiple calls.
      */
-    if (reach_max) {
-        TEST_EQUAL(scenario, TEST_EARLY_DATA_ACCEPTED);
-    }
+    read_size = (write_size / 2) + 1;
+    TEST_CALLOC(buf_read, read_size);
+
+    PSA_INIT();
 
     /*
      * Run first handshake to get a ticket from the server.
@@ -4579,49 +4580,45 @@
     ret = mbedtls_ssl_handshake(&(server_ep.ssl));
     TEST_EQUAL(ret, MBEDTLS_ERR_SSL_WANT_READ);
 
-    while (write_early_data_flag) {
-        unsigned char buf_read[23];
+    /*
+     * Write and if possible read as much as possible chunks of write_size
+     * bytes data without getting over the max_early_data_size limit.
+     */
+    do {
         uint32_t read_early_data_size = 0;
-        uint32_t remaining = max_early_data_size -
-                             server_ep.ssl.total_early_data_size;
 
-        early_data_len = sizeof(buf_write);
-        /*
-         * Adjust the length of next data to write depending on the remaining
-         * number of early data bytes we are allowed to write and if we want
-         * to reach the maximum exactly or not.
-         */
-        if (early_data_len >= remaining) {
-            if (reach_max) {
-                early_data_len = remaining;
-                write_early_data_flag = 0;
-            } else {
-                if (early_data_len == remaining) {
-                    early_data_len /= 2;
-                } else {
-                    early_data_len = remaining + 1;
-                    break;
-                }
-            }
+        if ((written_early_data_size + write_size) > max_early_data_size) {
+            break;
         }
 
-        for (size_t i = 0; i < early_data_len; i++) {
+        /*
+         * If the server rejected early data, base the determination of when
+         * to stop the loop on the expended size (padding and encryption
+         * expansion) of early data on server side and the number of early data
+         * received so far by the server (multiple of the expended size).
+         */
+        if ((expended_early_data_len != 0) &&
+            ((server_ep.ssl.total_early_data_size +
+              expended_early_data_len) > max_early_data_size)) {
+            break;
+        }
+
+        for (size_t i = 0; i < write_size; i++) {
             buf_write[i] = (unsigned char) (written_early_data_size + i);
         }
 
-        ret = write_early_data(&(client_ep.ssl), buf_write, early_data_len);
-        TEST_EQUAL(ret, early_data_len);
-        written_early_data_size += early_data_len;
+        ret = write_early_data(&(client_ep.ssl), buf_write, write_size);
+        TEST_EQUAL(ret, write_size);
+        written_early_data_size += write_size;
 
         switch (scenario) {
             case TEST_EARLY_DATA_ACCEPTED:
-                while (read_early_data_size < early_data_len) {
+                while (read_early_data_size < write_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_read,
-                                                      sizeof(buf_read));
+                                                      buf_read, read_size);
                     TEST_ASSERT(ret > 0);
 
                     TEST_MEMORY_COMPARE(buf_read, ret,
@@ -4653,24 +4650,14 @@
                 if (expended_early_data_len == 0) {
                     expended_early_data_len = server_ep.ssl.total_early_data_size;
                 }
-                remaining = max_early_data_size - server_ep.ssl.total_early_data_size;
-
-                if (expended_early_data_len > remaining) {
-                    write_early_data_flag = 0;
-                }
                 break;
         }
         TEST_ASSERT(server_ep.ssl.total_early_data_size <= max_early_data_size);
-    }
+    } while (1);
 
     mbedtls_debug_set_threshold(3);
-
-    if (reach_max) {
-        TEST_EQUAL(server_ep.ssl.total_early_data_size, max_early_data_size);
-    }
-
-    ret = write_early_data(&(client_ep.ssl), buf_write, early_data_len);
-    TEST_EQUAL(ret, early_data_len);
+    ret = write_early_data(&(client_ep.ssl), buf_write, write_size);
+    TEST_EQUAL(ret, write_size);
 
     ret = mbedtls_snprintf(pattern, sizeof(pattern),
                            "EarlyData: Too much early data received");