Implement aws_tls_handler_write

.gitignore

@@ -65,3 +65,4 @@ cmake-build*
 *.out
 *.app
 
+.worktrees/

include/aws/io/tls_channel_handler.h

@@ -858,6 +858,22 @@ AWS_IO_API struct aws_tls_ctx *aws_tls_ctx_acquire(struct aws_tls_ctx *ctx);
  */
 AWS_IO_API void aws_tls_ctx_release(struct aws_tls_ctx *ctx);
 
+/**
+ * Not necessary if you are installing more handlers into the channel, but if you just want to have TLS for arbitrary
+ * data and use the channel handler directly, this function allows you to write data to the channel and have it
+ * encrypted. The handler must have finished TLS negotiation before this is called, otherwise
+ * AWS_IO_TLS_ERROR_NOT_NEGOTIATED is raised.
+ *
+ * If buf is larger than a single channel message, the data is split across multiple messages. on_write_completed is
+ * invoked once the final message has been written.
+ */
+AWS_IO_API int aws_tls_handler_write(
+    struct aws_channel_handler *handler,
+    struct aws_channel_slot *slot,
+    struct aws_byte_buf *buf,
+    aws_channel_on_message_write_completed_fn *on_write_completed,
+    void *completion_user_data);
+
 /**
  * Returns a byte buffer by copy of the negotiated protocols. If there is no agreed upon protocol, len will be 0 and
  * buffer will be NULL.

source/tls_channel_handler.c

@@ -897,6 +897,48 @@ bool aws_tls_is_cipher_pref_supported(enum aws_tls_cipher_pref cipher_pref) {
 
 #endif /* BYO_CRYPTO */
 
+int aws_tls_handler_write(
+    struct aws_channel_handler *handler,
+    struct aws_channel_slot *slot,
+    struct aws_byte_buf *buf,
+    aws_channel_on_message_write_completed_fn *on_write_completed,
+    void *completion_user_data) {
+
+    AWS_PRECONDITION(handler);
+    AWS_PRECONDITION(slot);
+    AWS_PRECONDITION(buf);
+
+    /* Feed the plaintext through the TLS handler's own write path (process_write_message), which encrypts it and
+     * sends it downstream. A pooled message may be smaller than buf, so split across as many messages as needed and
+     * only attach the completion callback to the final message. */
+    struct aws_byte_cursor remaining = aws_byte_cursor_from_buf(buf);
+    do {
+        struct aws_io_message *message = aws_channel_acquire_message_from_pool(
+            slot->channel, AWS_IO_MESSAGE_APPLICATION_DATA, remaining.len);
+
+        const size_t chunk_len = aws_min_size(remaining.len, message->message_data.capacity);
+        struct aws_byte_cursor chunk = aws_byte_cursor_advance(&remaining, chunk_len);
+        if (!aws_byte_buf_write_from_whole_cursor(&message->message_data, chunk)) {
+            aws_mem_release(message->allocator, message);
+            return AWS_OP_ERR;
+        }
+
+        /* Only the last message carries the completion callback. */
+        if (remaining.len == 0) {
+            message->on_completion = on_write_completed;
+            message->user_data = completion_user_data;
+        }
+
+        if (handler->vtable->process_write_message(handler, slot, message)) {
+            /* On failure the handler does not take ownership of the message. */
+            aws_mem_release(message->allocator, message);
+            return AWS_OP_ERR;
+        }
+    } while (remaining.len > 0);
+
+    return AWS_OP_SUCCESS;
+}
+
 int aws_channel_setup_client_tls(
     struct aws_channel_slot *right_of_slot,
     struct aws_tls_connection_options *tls_options) {

tests/CMakeLists.txt

@@ -294,6 +294,7 @@ if(NOT BYO_CRYPTO)
     # Misc non-badssl tls tests
     add_net_test_case(test_duplicate_cert_import)
     add_net_test_case(tls_channel_echo_and_backpressure_test)
+    add_net_test_case(tls_channel_handler_write_test)
     add_net_test_case(tls_channel_shutdown_with_cache_test)
     add_net_test_case(tls_channel_shutdown_with_cache_window_update_after_shutdown_test)
     add_net_test_case(tls_client_channel_negotiation_success)

tests/tls_handler_test.c

@@ -790,6 +790,134 @@ static int s_tls_channel_echo_and_backpressure_test_fn(struct aws_allocator *all
 
 AWS_TEST_CASE(tls_channel_echo_and_backpressure_test, s_tls_channel_echo_and_backpressure_test_fn)
 
+struct tls_handler_write_completion {
+    struct aws_mutex *mutex;
+    struct aws_condition_variable *condition_variable;
+    bool completed;
+    int error_code;
+};
+
+static void s_tls_handler_write_completion_fn(
+    struct aws_channel *channel,
+    struct aws_io_message *message,
+    int err_code,
+    void *user_data) {
+
+    (void)channel;
+    (void)message;
+    struct tls_handler_write_completion *completion = user_data;
+    aws_mutex_lock(completion->mutex);
+    completion->completed = true;
+    completion->error_code = err_code;
+    aws_mutex_unlock(completion->mutex);
+    aws_condition_variable_notify_one(completion->condition_variable);
+}
+
+static bool s_tls_handler_write_completion_predicate(void *user_data) {
+    struct tls_handler_write_completion *completion = user_data;
+    return completion->completed;
+}
+
+struct tls_handler_write_task_args {
+    struct aws_channel_handler *tls_handler;
+    struct aws_channel_slot *tls_slot;
+    struct aws_byte_buf *buffer;
+    struct tls_handler_write_completion *completion;
+    struct aws_channel_task task;
+};
+
+static void s_tls_handler_write_task(struct aws_channel_task *task, void *arg, enum aws_task_status status) {
+    (void)task;
+    (void)status;
+    struct tls_handler_write_task_args *write_args = arg;
+    AWS_FATAL_ASSERT(
+        aws_tls_handler_write(
+            write_args->tls_handler,
+            write_args->tls_slot,
+            write_args->buffer,
+            s_tls_handler_write_completion_fn,
+            write_args->completion) == AWS_OP_SUCCESS);
+}
+
+/* Verify that aws_tls_handler_write() encrypts plaintext directly through the TLS handler and that the peer receives
+ * the decrypted data. This exercises the write path without an upstream handler feeding the TLS handler. */
+static int s_tls_channel_handler_write_test_fn(struct aws_allocator *allocator, void *ctx) {
+    (void)ctx;
+    ASSERT_SUCCESS(s_tls_channel_server_client_tester_init(allocator));
+    struct tls_test_rw_args *client_rw_args = &s_server_client_tester.client_rw_args;
+    struct tls_test_rw_args *server_rw_args = &s_server_client_tester.server_rw_args;
+    struct tls_test_args *client_args = &s_server_client_tester.client_args;
+    struct tls_test_args *server_args = &s_server_client_tester.server_args;
+
+    struct aws_byte_buf write_tag = aws_byte_buf_from_c_str("I'm a little teapot.");
+
+    struct aws_channel_handler *client_rw_handler =
+        rw_handler_new(allocator, s_tls_test_handle_read, s_tls_test_handle_write, true, 10000, client_rw_args);
+    ASSERT_NOT_NULL(client_rw_handler);
+    struct aws_channel_handler *server_rw_handler =
+        rw_handler_new(allocator, s_tls_test_handle_read, s_tls_test_handle_write, true, 10000, server_rw_args);
+    ASSERT_NOT_NULL(server_rw_handler);
+    server_args->rw_handler = server_rw_handler;
+    client_args->rw_handler = client_rw_handler;
+
+    g_aws_channel_max_fragment_size = 4096;
+    ASSERT_SUCCESS(s_set_socket_channel(&s_server_client_tester));
+
+    /* The TLS handler sits immediately to the left of the client's read/write test handler. */
+    struct aws_channel_slot *tls_slot = client_args->rw_slot->adj_left;
+    ASSERT_NOT_NULL(tls_slot);
+
+    struct tls_handler_write_completion completion = {
+        .mutex = &c_tester.mutex,
+        .condition_variable = &c_tester.condition_variable,
+        .completed = false,
+        .error_code = -1,
+    };
+    struct tls_handler_write_task_args write_args = {
+        .tls_handler = tls_slot->handler,
+        .tls_slot = tls_slot,
+        .buffer = &write_tag,
+        .completion = &completion,
+    };
+    aws_channel_task_init(
+        &write_args.task, s_tls_handler_write_task, &write_args, "tls_handler_write_test");
+    aws_channel_schedule_task_now(client_args->channel, &write_args.task);
+
+    /* Server should receive the plaintext after the TLS handler decrypts it. */
+    ASSERT_SUCCESS(aws_mutex_lock(&s_server_client_tester.server_mutex));
+    ASSERT_SUCCESS(aws_condition_variable_wait_pred(
+        &s_server_client_tester.server_condition_variable,
+        &s_server_client_tester.server_mutex,
+        s_tls_test_read_predicate,
+        server_rw_args));
+    ASSERT_SUCCESS(aws_mutex_unlock(&s_server_client_tester.server_mutex));
+
+    ASSERT_BIN_ARRAYS_EQUALS(
+        write_tag.buffer, write_tag.len, server_rw_args->received_message.buffer, server_rw_args->received_message.len);
+
+    /* The write completion callback should have fired with no error. */
+    ASSERT_SUCCESS(aws_mutex_lock(&c_tester.mutex));
+    ASSERT_SUCCESS(aws_condition_variable_wait_pred(
+        &c_tester.condition_variable, &c_tester.mutex, s_tls_handler_write_completion_predicate, &completion));
+    ASSERT_SUCCESS(aws_mutex_unlock(&c_tester.mutex));
+    ASSERT_INT_EQUALS(AWS_OP_SUCCESS, completion.error_code);
+
+    aws_channel_shutdown(server_args->channel, AWS_OP_SUCCESS);
+    ASSERT_SUCCESS(aws_mutex_lock(&s_server_client_tester.server_mutex));
+    ASSERT_SUCCESS(aws_condition_variable_wait_pred(
+        &s_server_client_tester.server_condition_variable,
+        &s_server_client_tester.server_mutex,
+        s_tls_channel_shutdown_predicate,
+        &s_server_client_tester.server_args));
+    ASSERT_SUCCESS(aws_mutex_unlock(&s_server_client_tester.server_mutex));
+
+    ASSERT_SUCCESS(s_tls_channel_server_client_tester_cleanup());
+
+    return AWS_OP_SUCCESS;
+}
+
+AWS_TEST_CASE(tls_channel_handler_write_test, s_tls_channel_handler_write_test_fn)
+
 static struct aws_byte_buf s_on_client_recive_shutdown_with_cache_data(
     struct aws_channel_handler *handler,
     struct aws_channel_slot *slot,