Memory optimization

tests/unit/s2n_ktls_io_test.c

@@ -911,9 +911,12 @@ int main(int argc, char **argv)
         };
         s2n_blocked_status blocked = S2N_NOT_BLOCKED;
 
-        const size_t max_frag_len = S2N_DEFAULT_FRAGMENT_LENGTH;
+        const size_t max_frag_len = S2N_KTLS_CONTROL_MESSAGE_MAX_FRAG_LEN;
+        const size_t app_data_max_frag_len = S2N_DEFAULT_FRAGMENT_LENGTH;
         /* Our test assumptions are wrong if this isn't true */
         EXPECT_TRUE(max_frag_len < sizeof(test_data));
+        EXPECT_TRUE(app_data_max_frag_len < sizeof(test_data));
+        EXPECT_TRUE(app_data_max_frag_len + max_frag_len < sizeof(test_data));
 
         /* Safety */
         {
@@ -954,7 +957,7 @@ int main(int argc, char **argv)
 
         /* Test: Receive does not completely fill the output buffer */
         {
-            const size_t small_frag_len = 10;
+            const size_t small_frag_len = 3;
             EXPECT_TRUE(small_frag_len < max_frag_len);
             EXPECT_TRUE(small_frag_len < sizeof(test_data));
             struct iovec small_test_iovec = test_iovec;
@@ -985,6 +988,44 @@ int main(int argc, char **argv)
             uint8_t *read = s2n_stuffer_raw_read(&conn->in, small_frag_len);
             EXPECT_BYTEARRAY_EQUAL(read, test_data, small_frag_len);
         };
+
+        /* Test: Receive resizes to handle application data */
+        {
+            DEFER_CLEANUP(struct s2n_connection *conn = s2n_connection_new(S2N_CLIENT),
+                    s2n_connection_ptr_free);
+            EXPECT_NOT_NULL(conn);
+
+            DEFER_CLEANUP(struct s2n_test_ktls_io_stuffer_pair pair = { 0 },
+                    s2n_ktls_io_stuffer_pair_free);
+            EXPECT_OK(s2n_test_init_ktls_io_stuffer(conn, conn, &pair));
+            struct s2n_test_ktls_io_stuffer *ctx = &pair.client_in;
+
+            size_t written = 0;
+            EXPECT_OK(s2n_ktls_sendmsg(ctx, TLS_APPLICATION_DATA,
+                    &test_iovec, 1, &blocked, &written));
+            EXPECT_EQUAL(written, sizeof(test_data));
+
+            uint8_t record_type = 0;
+            uint8_t *read = NULL;
+
+            /* First read returns first fragment */
+            EXPECT_SUCCESS(s2n_ktls_read_full_record(conn, &record_type));
+            EXPECT_EQUAL(record_type, TLS_APPLICATION_DATA);
+            EXPECT_EQUAL(s2n_stuffer_data_available(&conn->in), max_frag_len);
+            read = s2n_stuffer_raw_read(&conn->in, max_frag_len);
+            EXPECT_BYTEARRAY_EQUAL(read, test_data, max_frag_len);
+
+            /* Verify that conn->in was resized for future reads */
+            EXPECT_EQUAL(conn->in.blob.allocated, app_data_max_frag_len);
+            EXPECT_SUCCESS(s2n_stuffer_wipe(&conn->in));
+
+            /* Second read returns larger fragment */
+            EXPECT_SUCCESS(s2n_ktls_read_full_record(conn, &record_type));
+            EXPECT_EQUAL(record_type, TLS_APPLICATION_DATA);
+            EXPECT_EQUAL(s2n_stuffer_data_available(&conn->in), app_data_max_frag_len);
+            read = s2n_stuffer_raw_read(&conn->in, app_data_max_frag_len);
+            EXPECT_BYTEARRAY_EQUAL(read, test_data + max_frag_len, app_data_max_frag_len);
+        };
     };
 
     END_TEST();

tls/s2n_ktls.h

@@ -26,6 +26,12 @@
  */
 #include "tls/s2n_ktls_parameters.h"
 
+/* ktls currently only handles alerts and HelloRequests.
+ * HelloRequests are empty handshake messages (just the header).
+ */
+#define S2N_KTLS_CONTROL_MESSAGE_MAX_FRAG_LEN \
+    MAX(S2N_ALERT_LENGTH, TLS_HANDSHAKE_HEADER_LENGTH)
+
 /* A set of kTLS configurations representing the combination of sending
  * and receiving.
  */

tls/s2n_ktls_io.c

@@ -413,12 +413,10 @@ int s2n_ktls_read_full_record(struct s2n_connection *conn, uint8_t *record_type)
     POSIX_ENSURE_REF(record_type);
 
     /* This method copies data into conn->in, so is intended for control messages
-     * rather than application data. However in some cases-- such as when attempting
-     * to read the close_notify alert during s2n_shutdown-- it may encounter application
-     * data. Set a reasonable conn->in size to avoid an excessive number of calls
-     * to recvmsg when reading a larger record.
+     * rather than application data. Set a fragment size that should result in
+     * only one recvmsg call for any reasonably-sized control message.
      */
-    POSIX_GUARD(s2n_stuffer_resize_if_empty(&conn->in, S2N_DEFAULT_FRAGMENT_LENGTH));
+    POSIX_GUARD(s2n_stuffer_resize_if_empty(&conn->in, S2N_KTLS_CONTROL_MESSAGE_MAX_FRAG_LEN));
 
     struct s2n_stuffer record_stuffer = conn->in;
     size_t len = s2n_stuffer_space_remaining(&record_stuffer);
@@ -435,6 +433,15 @@ int s2n_ktls_read_full_record(struct s2n_connection *conn, uint8_t *record_type)
             buf, len, &blocked, &bytes_read);
     WITH_ERROR_BLINDING(conn, POSIX_GUARD_RESULT(result));
 
+    /* While this method normally only handles control messages, during
+     * s2n_shutdown it may encounter unprocessed application data which needs to
+     * be discarded. In that case, resize the input buffer to reduce the number
+     * of recvmsg calls required to finish reading the application data.
+     */
+    if (*record_type == TLS_APPLICATION_DATA) {
+        POSIX_GUARD(s2n_stuffer_resize(&conn->in, S2N_DEFAULT_FRAGMENT_LENGTH));
+    }
+
     POSIX_GUARD(s2n_stuffer_skip_write(&conn->in, bytes_read));
     return S2N_SUCCESS;
 }