diff --git a/src/AsyncWebSocket.cpp b/src/AsyncWebSocket.cpp index 9268eb19..641bb69b 100644 --- a/src/AsyncWebSocket.cpp +++ b/src/AsyncWebSocket.cpp @@ -35,6 +35,12 @@ #define STATE_FRAME_MASK 1 #define STATE_FRAME_DATA 2 +// Sentinel returned by webSocketSendFrame() when a torn frame is detected +// (partial add() after header committed, or header add() / malloc failure). +// Distinct from 0, which means "busy, try again later". The caller must not +// roll back its send offset and must schedule a connection close. +constexpr size_t WS_FRAME_TORN = static_cast(-1); + using namespace asyncsrv; size_t webSocketSendFrameWindow(AsyncClient *client) { @@ -83,8 +89,9 @@ size_t webSocketSendFrame(AsyncClient *client, bool final, uint8_t opcode, bool uint8_t *buf = (uint8_t *)malloc(headLen); if (buf == NULL) { async_ws_log_e("Failed to allocate"); - client->abort(); - return 0; + // Nothing is on the wire yet. Signal hard error so the caller closes + // via the safe onDisconnect path (abort() leaks: no onDisconnect fires). + return WS_FRAME_TORN; } buf[0] = opcode & 0x0F; @@ -103,10 +110,11 @@ size_t webSocketSendFrame(AsyncClient *client, bool final, uint8_t opcode, bool memcpy(buf + (headLen - 4), mbuf, 4); } if (client->add((const char *)buf, headLen) != headLen) { - // os_printf("error adding %lu header bytes\n", headLen); free(buf); - // Serial.println("SF 4"); - return 0; + // Header could not be fully committed; nothing usable is on the wire. + // Signal hard error so the caller closes via the safe onDisconnect path + // (abort() leaks: no onDisconnect fires on ESP32 or ESP8266). + return WS_FRAME_TORN; } free(buf); @@ -118,17 +126,19 @@ size_t webSocketSendFrame(AsyncClient *client, bool final, uint8_t opcode, bool } } if (client->add((const char *)data, len) != len) { - // os_printf("error adding %lu data bytes\n", len); - // Serial.println("SF 5"); - return 0; + // The header is already committed but the payload was only partially + // accepted (lwIP segment exhaustion under load). The stream is now + // unrecoverable. Returning 0 would make the caller roll back its offset + // and re-send with a NEW frame header injected mid-payload. Signal hard + // error so the caller closes via the safe onDisconnect path (abort() + // leaks: no onDisconnect fires on ESP32 or ESP8266). + return WS_FRAME_TORN; } } - if (!client->send()) { - // os_printf("error sending frame: %lu\n", headLen+len); - // Serial.println("SF 6"); - return 0; - } - // Serial.println("SF"); + // A failing send() is not an error: the frame is fully queued in lwIP and + // will be flushed by the next poll/ack cycle. Returning 0 here caused a + // rollback and a duplicate re-send of already-queued bytes. + client->send(); return len; } @@ -228,6 +238,14 @@ size_t AsyncWebSocketMessage::send(AsyncClient *client) { size_t sent = webSocketSendFrame(client, final, opCode, _mask, dPtr, toSend); _status = WS_MSG_SENDING; + if (sent == WS_FRAME_TORN) { + // Torn frame: header (and possibly partial payload) is already on the wire. + // Do NOT roll back _sent/_ack — re-sending from a rolled-back offset would + // inject a new frame header mid-payload and corrupt the stream. Mark the + // message errored so _runQueue schedules a safe connection close. + _status = WS_MSG_ERROR; + return 0; + } if (toSend && sent != toSend) { _sent -= (toSend - sent); _ack -= (toSend - sent); @@ -361,6 +379,28 @@ void AsyncWebSocketClient::_onPoll() { return; } + // A torn frame was detected during a previous _runQueue() call. We cannot + // close from inside _runQueue: its callers hold _queue_lock, and close() + // fires _onDisconnect which re-acquires _queue_lock (non-recursive → deadlock) + // and would delete *this mid-iteration. Defer to here: unlock, then abort() + // to send an RST (immediate reset, no TIME_WAIT). + // + // abort() synchronously fires the onDisconnect discard callback, which runs + // _onDisconnect() (nulls _client and erases *this from the server's _clients + // list, destroying this AsyncWebSocketClient) and then `delete c`s the + // AsyncClient. We must NOT `delete c` ourselves here: that would double-free + // the AsyncClient and corrupt the heap (esp-idf abort() always invokes the + // discard cb when err != ERR_CONN; ERR_CONN means the pcb was already gone, + // in which case the discard cb has already run and *this is already dead). + if (_abortPending) { + AsyncClient *c = _client; + if (c) { + lock.unlock(); + c->abort(); // fires onDisconnect → _onDisconnect() + delete c (full teardown) + } + return; + } + if (_client && _client->canSend() && (!_controlQueue.empty() || !_messageQueue.empty())) { _runQueue(); } else if (_keepAlivePeriod > 0 && (millis() - _lastMessageTime) >= _keepAlivePeriod && (_controlQueue.empty() && _messageQueue.empty())) { @@ -374,6 +414,9 @@ void AsyncWebSocketClient::_runQueue() { if (!_client) { return; } + if (_abortPending) { + return; // connection is being torn down; do not attempt further sends + } _clearQueue(); @@ -392,14 +435,20 @@ void AsyncWebSocketClient::_runQueue() { } if (space > (size_t)(ctrl.len() - 1)) { async_ws_log_v("[%s][%" PRIu32 "] SEND CTRL %" PRIu8, _server->url(), _clientId, ctrl.opcode()); - ctrl.send(_client); + size_t sent = ctrl.send(_client); + if (sent == WS_FRAME_TORN) { + async_ws_log_e("[%s][%" PRIu32 "] TORN CTRL FRAME: scheduling abort", _server->url(), _clientId); + _abortPending = true; + _status = WS_DISCONNECTING; + break; + } space = webSocketSendFrameWindow(_client); } } } - // then we can send message frames if there is space - if (space) { + // then we can send message frames if there is space and no abort is pending + if (!_abortPending && space) { for (auto &msg : _messageQueue) { if (msg._remainingBytesToSend()) { async_ws_log_v( @@ -411,6 +460,13 @@ void AsyncWebSocketClient::_runQueue() { msg.send(_client); space = webSocketSendFrameWindow(_client); + if (msg._status == WS_MSG_ERROR) { + async_ws_log_e("[%s][%" PRIu32 "] TORN FRAME: scheduling close", _server->url(), _clientId); + _abortPending = true; + _status = WS_DISCONNECTING; + break; + } + // If we haven't finished sending this message, we must stop here to preserve WebSocket ordering. // We can only pipeline subsequent messages if the current one is fully passed to TCP buffer. if (msg._remainingBytesToSend()) { @@ -526,15 +582,18 @@ void AsyncWebSocketClient::close(uint16_t code, const char *message) { free(buf); return; } else { + // Could not allocate the close-frame payload. We can't send a proper WS + // close frame, so abort() to send an RST. abort() synchronously fires the + // onDisconnect discard callback, which runs _onDisconnect() (erases *this + // from the server's _clients, destroying this object) and `delete c`s the + // AsyncClient. Do NOT `delete c` here: that would double-free it. + // close() is unsuitable because it may be called with _queue_lock held. async_ws_log_e("Failed to allocate"); - // Reads _client, then dereference it without any lock. - // A concurrent _onDisconnect could null + delete the client between the check and the use. - // Local capture ensures the pointer is read exactly once, eliminating the null-dereference. - // (TOCTOU) AsyncClient *c = _client; if (c) { - c->abort(); + c->abort(); // fires onDisconnect → _onDisconnect() + delete c } + return; } } _queueControl(WS_DISCONNECT); @@ -790,8 +849,17 @@ void AsyncWebSocketClient::_handleDataEvent(uint8_t *data, size_t len, bool endO _server->_handleEvent(this, WS_EVT_DATA, (void *)&_pinfo, copy.get(), len); } else { async_ws_log_e("Failed to allocate"); - if (_client) { - _client->abort(); + // abort() to send an RST. It synchronously fires the onDisconnect + // discard callback, which runs _onDisconnect() (erases *this from the + // server's _clients, destroying this object) and `delete c`s the + // AsyncClient. Do NOT `delete c` here: that would double-free it. + // close() is unsafe here because _handleDataEvent is called from _onData + // and close() would fire _onDisconnect → _handleDisconnect → erase + // *this from _clients mid-call (same outcome, but close() may also + // flush a FIN, which we don't want on an allocation failure). + AsyncClient *c = _client; + if (c) { + c->abort(); // fires onDisconnect → _onDisconnect() + delete c } } } else { diff --git a/src/AsyncWebSocket.h b/src/AsyncWebSocket.h index 16cd8fc2..94968b8a 100644 --- a/src/AsyncWebSocket.h +++ b/src/AsyncWebSocket.h @@ -226,6 +226,7 @@ class AsyncWebSocketClient { std::deque _controlQueue; std::deque _messageQueue; bool _closeWhenFull = false; + bool _abortPending = false; // set when a torn frame is detected; _onPoll closes the connection safely AwsFrameInfo _pinfo;