diff --git a/.gitignore b/.gitignore index 2ef3eccf7..e36b4bacd 100644 --- a/.gitignore +++ b/.gitignore @@ -65,3 +65,4 @@ cmake-build* *.out *.app +.worktrees/ diff --git a/include/aws/io/tls_channel_handler.h b/include/aws/io/tls_channel_handler.h index a2d58f353..24ddeee7c 100644 --- a/include/aws/io/tls_channel_handler.h +++ b/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. diff --git a/source/tls_channel_handler.c b/source/tls_channel_handler.c index daa6ef226..28bac85f6 100644 --- a/source/tls_channel_handler.c +++ b/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) { diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 2a2472ebe..7e23409ee 100644 --- a/tests/CMakeLists.txt +++ b/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) diff --git a/tests/tls_handler_test.c b/tests/tls_handler_test.c index 1dcf100dd..9f9bcae35 100644 --- a/tests/tls_handler_test.c +++ b/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,