From e38575434ec35262caebd38a710f2e3caf83bcbe Mon Sep 17 00:00:00 2001
From: Igor Abdrakhimov <igorabd@amazon.com>
Date: Thu, 20 Feb 2025 14:00:45 -0800
Subject: [PATCH] PR feedback

---
 source/client.c                  | 43 +++++++++++++-------------------
 tests/v3/connection_state_test.c |  9 +++++--
 2 files changed, 25 insertions(+), 27 deletions(-)

diff --git a/source/client.c b/source/client.c
index b9934022..5a3b309c 100644
--- a/source/client.c
+++ b/source/client.c
@@ -1364,6 +1364,7 @@ static int s_websocket_connect(struct aws_mqtt_client_connection_311_impl *conne
 }
 
 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;
@@ -1377,33 +1378,12 @@ static void s_on_websocket_setup_task_fn(struct aws_task *task, void *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(connection->allocator, on_websocket_setup_task_arg);
+    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);
 }
 
-void s_on_websocket_handshake_failure(struct aws_mqtt_client_connection_311_impl *connection, int error_code) {
-    AWS_FATAL_ASSERT(connection);
-
-    /* s_on_websocket_setup shutdowns mqtt connection, this must happen on the MQTT connection's event loop. */
-    if (aws_event_loop_thread_is_callers_thread(connection->loop)) {
-        struct aws_websocket_on_connection_setup_data websocket_setup = {.error_code = error_code};
-        s_on_websocket_setup(&websocket_setup, connection);
-    } else {
-        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->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);
-    }
-}
-
 static void s_websocket_handshake_transform_complete(
     struct aws_http_message *handshake_request,
     int error_code,
@@ -1457,9 +1437,22 @@ static void s_websocket_handshake_transform_complete(
     /* Success */
     return;
 
-error:
-    /* Proceed to next step, telling it that we failed. */
-    s_on_websocket_handshake_failure(connection, error_code);
+error:;
+    /* 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);
 }
 
 /*******************************************************************************
diff --git a/tests/v3/connection_state_test.c b/tests/v3/connection_state_test.c
index 989d9307..1088d53a 100644
--- a/tests/v3/connection_state_test.c
+++ b/tests/v3/connection_state_test.c
@@ -3975,11 +3975,16 @@ static void s_mqtt_client_test_websocket_failed_transform(
     aws_mqtt_transform_websocket_handshake_complete_fn *complete_fn,
     void *complete_ctx) {
 
-    (void)user_data;
+    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;
@@ -3989,7 +3994,7 @@ static int s_test_mqtt_websocket_failed_transform_fn(struct aws_allocator *alloc
     };
 
     aws_mqtt_client_connection_use_websockets(
-        state_test_data->mqtt_connection, s_mqtt_client_test_websocket_failed_transform, NULL, NULL, NULL);
+        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);