From 665ece1e128edff1afa508641b448cfa234572f9 Mon Sep 17 00:00:00 2001 From: agessaman Date: Mon, 21 Sep 2026 16:05:39 -0700 Subject: [PATCH] fix(mqtt): never hand back a half-initialized client or transport esp_mqtt_client_init() succeeded without a registered event sink when registration failed, and the WebSocket wrap returned a transport with the unpadded SDK buffer when the padded allocation failed. Both paths now fail cleanly instead, the IDF 4 WebSocket layout is pinned to the validated patch level, and the connection flags esp-mqtt tasks write become atomic. --- lib/PsychicMqttClient/src/MqttClientInit.h | 17 +++++ .../src/PsychicMqttClient.cpp | 21 +++--- lib/PsychicMqttClient/src/PsychicMqttClient.h | 5 +- src/helpers/ESP32WsTransportFix.cpp | 28 ++++--- src/helpers/RequiredTransportBuffer.h | 18 +++++ .../test_observer_init/test_observer_init.cpp | 73 +++++++++++++++++++ 6 files changed, 135 insertions(+), 27 deletions(-) create mode 100644 lib/PsychicMqttClient/src/MqttClientInit.h create mode 100644 src/helpers/RequiredTransportBuffer.h create mode 100644 test/test_observer_init/test_observer_init.cpp diff --git a/lib/PsychicMqttClient/src/MqttClientInit.h b/lib/PsychicMqttClient/src/MqttClientInit.h new file mode 100644 index 00000000..148b8e22 --- /dev/null +++ b/lib/PsychicMqttClient/src/MqttClientInit.h @@ -0,0 +1,17 @@ +#pragma once + +// Initialization is incomplete until lifecycle events have a registered sink. +// The candidate has never started, so registration failure can destroy it. +template +int initializeMqttClient(Handle& target, Init init, Register register_events, + Destroy destroy, int no_memory) { + Handle candidate = init(); + if (!candidate) return no_memory; + const int result = register_events(candidate); + if (result != 0) { + destroy(candidate); + return result; + } + target = candidate; + return 0; +} diff --git a/lib/PsychicMqttClient/src/PsychicMqttClient.cpp b/lib/PsychicMqttClient/src/PsychicMqttClient.cpp index 01efb194..b4567071 100644 --- a/lib/PsychicMqttClient/src/PsychicMqttClient.cpp +++ b/lib/PsychicMqttClient/src/PsychicMqttClient.cpp @@ -1,3 +1,4 @@ +#include "MqttClientInit.h" #include "PsychicMqttClient.h" #include @@ -386,17 +387,17 @@ esp_err_t PsychicMqttClient::applyConfig() if (_client == nullptr) { - // esp_mqtt_client_init() takes the whole configuration, including the - // buffer sizes it allocates once and never resizes. - _client = esp_mqtt_client_init(&_mqtt_cfg); - if (_client == nullptr) - { - ESP_LOGE(TAG, "esp_mqtt_client_init failed"); - return ESP_ERR_NO_MEM; + const esp_err_t result = initializeMqttClient(_client, + [&]() { return esp_mqtt_client_init(&_mqtt_cfg); }, + [&](esp_mqtt_client_handle_t client) { + return esp_mqtt_client_register_event(client, MQTT_EVENT_ANY, _onMqttEventStatic, this); + }, + [](esp_mqtt_client_handle_t client) { esp_mqtt_client_destroy(client); }, + ESP_ERR_NO_MEM); + if (result != ESP_OK) { + ESP_LOGE(TAG, "MQTT initialization failed: %s", esp_err_to_name(result)); + return result; } - // Register event handler only once when client is first created - // to avoid memory leak from repeated registrations - esp_mqtt_client_register_event(_client, MQTT_EVENT_ANY, _onMqttEventStatic, this); _config_dirty = false; return ESP_OK; } diff --git a/lib/PsychicMqttClient/src/PsychicMqttClient.h b/lib/PsychicMqttClient/src/PsychicMqttClient.h index 93b85a43..99f34833 100644 --- a/lib/PsychicMqttClient/src/PsychicMqttClient.h +++ b/lib/PsychicMqttClient/src/PsychicMqttClient.h @@ -1,4 +1,5 @@ #pragma once +#include /** * PsychicMqttClient * @@ -545,8 +546,8 @@ private: esp_mqtt_client_handle_t _client = nullptr; esp_mqtt_client_config_t _mqtt_cfg; esp_mqtt_error_codes_t _lastError; - bool _connected = false; - bool _stopMqttClient = false; + std::atomic _connected{false}; + std::atomic _stopMqttClient{false}; bool _config_dirty = true; bool _started = false; diff --git a/src/helpers/ESP32WsTransportFix.cpp b/src/helpers/ESP32WsTransportFix.cpp index a69e0b08..2febce06 100644 --- a/src/helpers/ESP32WsTransportFix.cpp +++ b/src/helpers/ESP32WsTransportFix.cpp @@ -62,7 +62,12 @@ #if ESP_IDF_VERSION_MAJOR == 4 +#if ESP_IDF_VERSION_MINOR != 4 || ESP_IDF_VERSION_PATCH != 7 +#error "Revalidate the private WebSocket layout before using another IDF 4 SDK" +#endif + #include +#include "RequiredTransportBuffer.h" #include "sdkconfig.h" #include "esp_transport.h" #include "esp_transport_ws.h" @@ -97,22 +102,15 @@ extern "C" { esp_transport_handle_t __real_esp_transport_ws_init(esp_transport_handle_t parent_handle); esp_transport_handle_t __wrap_esp_transport_ws_init(esp_transport_handle_t parent_handle) { - esp_transport_handle_t t = __real_esp_transport_ws_init(parent_handle); - if (t != nullptr) { - transport_ws_t* ws = (transport_ws_t*)esp_transport_get_context_data(t); - if (ws != nullptr && ws->buffer != nullptr) { - // The buffer is untouched at this point (allocated moments ago inside - // __real_esp_transport_ws_init), so a swap is safe. - char* padded = (char*)malloc(CONFIG_WS_BUFFER_SIZE + 1); - if (padded != nullptr) { + return initWithRequiredTransportBuffer( + CONFIG_WS_BUFFER_SIZE + 1, malloc, + [&]() { return __real_esp_transport_ws_init(parent_handle); }, + [](esp_transport_handle_t transport, void* buffer) { + // Successful init in the pinned SDK guarantees this context and buffer. + transport_ws_t* ws = (transport_ws_t*)esp_transport_get_context_data(transport); free(ws->buffer); - ws->buffer = padded; - } - // On alloc failure keep the original buffer: same behavior as before - // this fix, which is still strictly better than failing init here. - } - } - return t; + ws->buffer = static_cast(buffer); + }, free); } } // extern "C" diff --git a/src/helpers/RequiredTransportBuffer.h b/src/helpers/RequiredTransportBuffer.h new file mode 100644 index 00000000..dba5df7e --- /dev/null +++ b/src/helpers/RequiredTransportBuffer.h @@ -0,0 +1,18 @@ +#pragma once +#include + +// Allocate the protection buffer first. A successful init must install it; +// neither allocation failure is allowed to return an unprotected transport. +template +Handle initWithRequiredTransportBuffer(size_t size, Allocate allocate, Init init, + Install install, Free release) { + void* buffer = allocate(size); + if (!buffer) return nullptr; + Handle transport = init(); + if (!transport) { + release(buffer); + return nullptr; + } + install(transport, buffer); + return transport; +} diff --git a/test/test_observer_init/test_observer_init.cpp b/test/test_observer_init/test_observer_init.cpp new file mode 100644 index 00000000..75f0b5f3 --- /dev/null +++ b/test/test_observer_init/test_observer_init.cpp @@ -0,0 +1,73 @@ +#include +#include "helpers/RequiredTransportBuffer.h" +#include "../../lib/PsychicMqttClient/src/MqttClientInit.h" + +TEST(ObserverInit, MissingPaddingPreventsSdkInitialization) { + bool initialized = false, installed = false, freed = false; + int transport; + int* result = initWithRequiredTransportBuffer(1025, + [](size_t) -> void* { return nullptr; }, + [&] { initialized = true; return &transport; }, + [&](int*, void*) { installed = true; }, + [&](void*) { freed = true; }); + EXPECT_EQ(nullptr, result); + EXPECT_FALSE(initialized); + EXPECT_FALSE(installed); + EXPECT_FALSE(freed); +} + +TEST(ObserverInit, FailedSdkInitReleasesReservedPadding) { + char buffer[1025]; + bool freed = false, installed = false; + int* result = initWithRequiredTransportBuffer(sizeof(buffer), + [&](size_t n) -> void* { EXPECT_EQ(sizeof(buffer), n); return buffer; }, + []() -> int* { return nullptr; }, + [&](int*, void*) { installed = true; }, + [&](void* p) { EXPECT_EQ(buffer, p); freed = true; }); + EXPECT_EQ(nullptr, result); + EXPECT_TRUE(freed); + EXPECT_FALSE(installed); +} + +TEST(ObserverInit, SuccessfulTransportOwnsProtectionBuffer) { + char buffer[1025]; + int transport; + bool installed = false, freed = false; + int* result = initWithRequiredTransportBuffer(sizeof(buffer), + [&](size_t) -> void* { return buffer; }, + [&] { return &transport; }, + [&](int* t, void* p) { EXPECT_EQ(&transport, t); EXPECT_EQ(buffer, p); installed = true; }, + [&](void*) { freed = true; }); + EXPECT_EQ(&transport, result); + EXPECT_TRUE(installed); + EXPECT_FALSE(freed); +} + +TEST(ObserverInit, FailedRegistrationDestroysCandidateAndCanRetry) { + int candidate; + int* client = nullptr; + int registrations = 0, destroys = 0; + auto create = [&] { return &candidate; }; + auto register_events = [&](int* p) { EXPECT_EQ(&candidate, p); return ++registrations == 1 ? -7 : 0; }; + auto destroy = [&](int* p) { EXPECT_EQ(&candidate, p); ++destroys; }; + EXPECT_EQ(-7, initializeMqttClient(client, create, register_events, destroy, -2)); + EXPECT_EQ(nullptr, client); + EXPECT_EQ(1, destroys); + EXPECT_EQ(0, initializeMqttClient(client, create, register_events, destroy, -2)); + EXPECT_EQ(&candidate, client); + EXPECT_EQ(1, destroys); +} + +TEST(ObserverInit, FailedClientAllocationDoesNotRegisterOrDestroy) { + int* client = nullptr; + int effects = 0; + EXPECT_EQ(-2, initializeMqttClient(client, []() -> int* { return nullptr; }, + [&](int*) { ++effects; return 0; }, [&](int*) { ++effects; }, -2)); + EXPECT_EQ(nullptr, client); + EXPECT_EQ(0, effects); +} + +int main(int argc, char** argv) { + ::testing::InitGoogleTest(&argc, argv); + return RUN_ALL_TESTS(); +}