Enforce maximum size of early data when rejected

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 1408361..0c1d606 100644
--- a/tests/suites/test_suite_ssl.function
+++ b/tests/suites/test_suite_ssl.function
@@ -4466,6 +4466,7 @@
     char pattern[128];
     unsigned char buf_write[64];
     size_t early_data_len = sizeof(buf_write);
+    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;
@@ -4503,6 +4504,14 @@
         case TEST_EARLY_DATA_ACCEPTED:
             break;
 
+        case TEST_EARLY_DATA_SERVER_REJECTS:
+            server_options.early_data = MBEDTLS_SSL_EARLY_DATA_DISABLED;
+            ret = mbedtls_snprintf(pattern, sizeof(pattern),
+                                   "EarlyData: deprotect and discard app data records.");
+            TEST_ASSERT(ret < (int) sizeof(pattern));
+            mbedtls_debug_set_threshold(3);
+            break;
+
         default:
             TEST_FAIL("Unknown scenario.");
     }
@@ -4552,7 +4561,7 @@
         uint32_t remaining = max_early_data_size -
                              server_ep.ssl.early_data_count;
 
-        /* Reach maximum early data exactly */
+        /* In case of accepted early data, reach max_early_data_size exactly. */
         if (early_data_len >= remaining) {
             early_data_len = remaining;
             write_early_data_flag = 0;
@@ -4585,13 +4594,43 @@
                                written_early_data_size);
                 }
                 break;
+
+            case TEST_EARLY_DATA_SERVER_REJECTS:
+                ret = mbedtls_ssl_handshake(&(server_ep.ssl));
+                /*
+                 * Can be the case if max_early_data_size is smaller then the
+                 * smallest inner content or protected record.
+                 */
+                if (ret == MBEDTLS_ERR_SSL_UNEXPECTED_MESSAGE) {
+                    /* Beyond 64 for max_early_data_size it is suspicious */
+                    TEST_ASSERT(max_early_data_size < 64);
+                    goto exit;
+                }
+
+                TEST_ASSERT(ret == MBEDTLS_ERR_SSL_WANT_READ);
+
+                TEST_EQUAL(server_pattern.counter, 1);
+                server_pattern.counter = 0;
+                if (expended_early_data_len == 0) {
+                    expended_early_data_len = server_ep.ssl.early_data_count;
+                }
+                remaining = max_early_data_size - server_ep.ssl.early_data_count;
+
+                if (expended_early_data_len > remaining) {
+                    write_early_data_flag = 0;
+                }
+                break;
         }
         TEST_ASSERT(server_ep.ssl.early_data_count <= max_early_data_size);
     }
 
     mbedtls_debug_set_threshold(3);
-    ret = write_early_data(&(client_ep.ssl), buf_write, 1);
-    TEST_EQUAL(ret, 1);
+
+    early_data_len = (scenario == TEST_EARLY_DATA_ACCEPTED) ?
+                     1 : sizeof(buf_write);
+
+    ret = write_early_data(&(client_ep.ssl), buf_write, early_data_len);
+    TEST_EQUAL(ret, early_data_len);
 
     ret = mbedtls_snprintf(pattern, sizeof(pattern),
                            "EarlyData: Too many early data received");