Fix handling of websocket failed transform for MQTT311 connection (#383)

Author: sfod

source/client.c

@@ -1363,6 +1363,27 @@ static int s_websocket_connect(struct aws_mqtt_client_connection_311_impl *conne
     return AWS_OP_ERR;
 }
 
+struct mqtt_on_websocket_setup_task_arg {
+    struct aws_allocator *allocator;
+    struct aws_task task;
+    struct aws_mqtt_client_connection_311_impl *connection;
+    int error_code;
+};
+
+static void s_on_websocket_setup_task_fn(struct aws_task *task, void *userdata, enum aws_task_status status) {
+    (void)task;
+    (void)status;
+
+    struct mqtt_on_websocket_setup_task_arg *on_websocket_setup_task_arg = userdata;
+    struct aws_mqtt_client_connection_311_impl *connection = on_websocket_setup_task_arg->connection;
+    int error_code = on_websocket_setup_task_arg->error_code;
+
+    aws_mem_release(on_websocket_setup_task_arg->allocator, on_websocket_setup_task_arg);
+
+    struct aws_websocket_on_connection_setup_data websocket_setup = {.error_code = error_code};
+    s_on_websocket_setup(&websocket_setup, connection);
+}
+
 static void s_websocket_handshake_transform_complete(
     struct aws_http_message *handshake_request,
     int error_code,
@@ -1417,9 +1438,21 @@ static void s_websocket_handshake_transform_complete(
     return;
 
 error:;
-    /* Proceed to next step, telling it that we failed. */
-    struct aws_websocket_on_connection_setup_data websocket_setup = {.error_code = error_code};
-    s_on_websocket_setup(&websocket_setup, connection);
+    /* Proceed to next step, telling it that we failed.
+     * s_on_websocket_setup will shutdown MQTT connection, and this MUST happen on the MQTT connection's event loop. */
+    struct mqtt_on_websocket_setup_task_arg *on_websocket_setup_task_arg =
+        aws_mem_calloc(connection->allocator, 1, sizeof(struct mqtt_on_websocket_setup_task_arg));
+    on_websocket_setup_task_arg->allocator = connection->allocator;
+    /* NOTE: No need in acquiring MQTT connection ref counter, as it is already acquired at the start of connection
+     * process. s_on_websocket_setup will release it. */
+    on_websocket_setup_task_arg->connection = connection;
+    on_websocket_setup_task_arg->error_code = error_code;
+    aws_task_init(
+        &on_websocket_setup_task_arg->task,
+        s_on_websocket_setup_task_fn,
+        (void *)on_websocket_setup_task_arg,
+        "on_websocket_setup_task");
+    aws_event_loop_schedule_task_now(connection->loop, &on_websocket_setup_task_arg->task);
 }
 
 /*******************************************************************************

tests/CMakeLists.txt

@@ -108,6 +108,8 @@ add_test_case(mqtt_validation_failure_connect_invalid_client_id_utf8)
 add_test_case(mqtt_validation_failure_invalid_will_topic_utf8)
 add_test_case(mqtt_validation_failure_invalid_username_utf8)
 
+add_test_case(mqtt_websocket_failed_transform)
+
 # Operation statistics tests
 add_test_case(mqtt_operation_statistics_simple_publish)
 add_test_case(mqtt_operation_statistics_offline_publish)

tests/v3/connection_state_test.c

@@ -3968,3 +3968,45 @@ AWS_TEST_CASE_FIXTURE(
     s_test_mqtt_validation_failure_invalid_username_utf8_fn,
     s_clean_up_mqtt_server_fn,
     &test_data)
+
+static void s_mqtt_client_test_websocket_failed_transform(
+    struct aws_http_message *request,
+    void *user_data,
+    aws_mqtt_transform_websocket_handshake_complete_fn *complete_fn,
+    void *complete_ctx) {
+
+    struct mqtt_connection_state_test *state_test_data = user_data;
+    struct aws_mqtt_client_connection_311_impl *mqtt311_connection = state_test_data->mqtt_connection->impl;
+    /* This test should also perform a regression check: websocket transform failed on a non-eventloop thread does not
+     * cause crash. */
+    AWS_FATAL_ASSERT(!aws_event_loop_thread_is_callers_thread(mqtt311_connection->loop));
+
+    (*complete_fn)(request, AWS_ERROR_INVALID_STATE, complete_ctx);
+}
+
+/* Connection failure test where websocket MQTT channel establishment fails due to handshake transform failure */
+static int s_test_mqtt_websocket_failed_transform_fn(struct aws_allocator *allocator, void *ctx) {
+    (void)allocator;
+    struct mqtt_connection_state_test *state_test_data = ctx;
+
+    struct aws_mqtt_connection_options connection_options = {
+        .socket_options = &state_test_data->socket_options,
+    };
+
+    aws_mqtt_client_connection_use_websockets(
+        state_test_data->mqtt_connection, s_mqtt_client_test_websocket_failed_transform, state_test_data, NULL, NULL);
+
+    ASSERT_SUCCESS(aws_mqtt_client_connection_connect(state_test_data->mqtt_connection, &connection_options));
+    aws_test311_wait_for_connection_to_fail(state_test_data);
+
+    ASSERT_INT_EQUALS(AWS_ERROR_INVALID_STATE, state_test_data->error);
+
+    return AWS_OP_SUCCESS;
+}
+
+AWS_TEST_CASE_FIXTURE(
+    mqtt_websocket_failed_transform,
+    s_setup_mqtt_server_fn,
+    s_test_mqtt_websocket_failed_transform_fn,
+    s_clean_up_mqtt_server_fn,
+    &test_data)