From bc16308dbe3b4d91faed381048ec6ef1cab65644 Mon Sep 17 00:00:00 2001 From: Will Miles Date: Fri, 28 Aug 2026 20:36:58 -0400 Subject: [PATCH 1/3] fix(websocket): Fix handling of torn headers Use a state machine to process headers byte-by-byte so we can handle partial reception at any point. --- src/AsyncWebSocket.cpp | 384 ++++++++++++++++++++++++----------------- src/AsyncWebSocket.h | 9 +- 2 files changed, 235 insertions(+), 158 deletions(-) diff --git a/src/AsyncWebSocket.cpp b/src/AsyncWebSocket.cpp index 7e1d9e9f..65c28e58 100644 --- a/src/AsyncWebSocket.cpp +++ b/src/AsyncWebSocket.cpp @@ -31,12 +31,28 @@ #include #include -#define STATE_FRAME_START 0 -#define STATE_FRAME_MASK 1 -#define STATE_FRAME_DATA 2 - using namespace asyncsrv; +enum class AwsParseState : uint8_t { + Start = 0, + Length, + Length2_1, + Length2_2, + Length8_1, + Length8_2, + Length8_3, + Length8_4, + Length8_5, + Length8_6, + Length8_7, + Length8_8, + Mask_0, + Mask_1, + Mask_2, + Mask_3, + Payload, +}; + static AsyncWebSocketSharedBuffer makeSharedBuffer(const uint8_t *message, size_t len) { if (message) { return std::make_shared>(message, message + len); @@ -167,8 +183,8 @@ const char *AWSC_PING_PAYLOAD = "ESPAsyncWebServer-PING"; const size_t AWSC_PING_PAYLOAD_LEN = 22; AsyncWebSocketClient::AsyncWebSocketClient(AsyncClient *client, AsyncWebSocket *server) - : _client(client), _server(server), _clientId(_server->_getNextId()), _status(WS_CONNECTED), _pstate(STATE_FRAME_START), _lastMessageTime(millis()), - _keepAlivePeriod(0), _tempObject(NULL) { + : _client(client), _server(server), _clientId(_server->_getNextId()), _status(WS_CONNECTED), _lastMessageTime(millis()), _keepAlivePeriod(0), + _pstate(AwsParseState::Start), _tempObject(NULL) { _client->setRxTimeout(0); _client->onError( @@ -511,189 +527,243 @@ void AsyncWebSocketClient::_onData(void *pbuf, size_t plen) { while (plen > 0) { async_ws_log_v( - "[%s][%" PRIu32 "] DATA plen: %" PRIu32 ", _pstate: %" PRIu8 ", _status: %" PRIu8, _server->url(), _clientId, static_cast(plen), _pstate, - static_cast(_status) + "[%s][%" PRIu32 "] DATA plen: %" PRIu32 ", _pstate: %" PRIu8 ", _status: %" PRIu8, _server->url(), _clientId, static_cast(plen), + static_cast(_pstate), static_cast(_status) ); - if (_pstate == STATE_FRAME_START) { - const uint8_t *fdata = data; - - _pinfo.index = 0; - _pinfo.final = (fdata[0] & 0x80) != 0; - _pinfo.opcode = fdata[0] & 0x0F; - _pinfo.masked = ((fdata[1] & 0x80) != 0) ? 1 : 0; - _pinfo.len = fdata[1] & 0x7F; - - data += 2; - plen -= 2; - - if (_pinfo.len == 126 && plen >= 2) { - _pinfo.len = fdata[3] | (uint16_t)(fdata[2]) << 8; - data += 2; - plen -= 2; - - } else if (_pinfo.len == 127 && plen >= 8) { - _pinfo.len = fdata[9] | (uint16_t)(fdata[8]) << 8 | (uint32_t)(fdata[7]) << 16 | (uint32_t)(fdata[6]) << 24 | (uint64_t)(fdata[5]) << 32 - | (uint64_t)(fdata[4]) << 40 | (uint64_t)(fdata[3]) << 48 | (uint64_t)(fdata[2]) << 56; - data += 8; - plen -= 8; - } - } - - async_ws_log_v( - "[%s][%" PRIu32 "] DATA _pinfo: index: %" PRIu64 ", final: %" PRIu8 ", opcode: %" PRIu8 ", masked: %" PRIu8 ", len: %" PRIu64, _server->url(), _clientId, - _pinfo.index, _pinfo.final, _pinfo.opcode, _pinfo.masked, _pinfo.len - ); - - // Handle fragmented mask data - Safari may split the 4-byte mask across multiple packets - // _pinfo.masked is 1 if we need to start reading mask bytes - // _pinfo.masked is 2, 3, or 4 if we have partially read the mask - // _pinfo.masked is 5 if the mask is complete - while (_pinfo.masked && _pstate <= STATE_FRAME_MASK && _pinfo.masked < 5) { - // check if we have some data - if (plen == 0) { - // Safari close frame edge case: masked bit set but no mask data - if (_pinfo.opcode == WS_DISCONNECT) { - async_ws_log_v("[%s][%" PRIu32 "] DATA close frame with incomplete mask, treating as unmasked", _server->url(), _clientId); - _pinfo.masked = 0; - _pinfo.index = 0; - _pinfo.len = 0; - _pstate = STATE_FRAME_START; - break; + bool consume_byte = true; + + switch (_pstate) { + case AwsParseState::Start: + // First header byte + _pinfo.index = 0; + _pinfo.final = (data[0] & 0x80) != 0; + _pinfo.opcode = data[0] & 0x0F; + _pstate = AwsParseState::Length; + break; + case AwsParseState::Length: + // Second header byte + _pinfo.masked = ((data[0] & 0x80) != 0) ? 1 : 0; + _pinfo.len = data[0] & 0x7F; + // Select length type + if (_pinfo.len == 126) { + _pstate = AwsParseState::Length2_1; + } else if (_pinfo.len == 127) { + _pstate = AwsParseState::Length8_1; + } else { + _pstate = (_pinfo.masked) ? AwsParseState::Mask_0 : AwsParseState::Payload; } - - // wait for more data - _pstate = STATE_FRAME_MASK; - async_ws_log_v("[%s][%" PRIu32 "] DATA waiting for more mask data: read: %" PRIu8 "/4", _server->url(), _clientId, _pinfo.masked - 1); - return; + break; + + // 2-byte length form + case AwsParseState::Length2_1: + _pinfo.len = (data[0] << 8); + _pstate = AwsParseState::Length2_2; + break; + case AwsParseState::Length2_2: + _pinfo.len += data[0]; + _pstate = (_pinfo.masked) ? AwsParseState::Mask_0 : AwsParseState::Payload; + break; + + // 8-byte length form, tediously unrolled + case AwsParseState::Length8_1: + _pinfo.len = (uint64_t)data[0] << 56; + _pstate = AwsParseState::Length8_2; + break; + case AwsParseState::Length8_2: + _pinfo.len += (uint64_t)data[0] << 48; + _pstate = AwsParseState::Length8_3; + break; + case AwsParseState::Length8_3: + _pinfo.len += (uint64_t)data[0] << 40; + _pstate = AwsParseState::Length8_4; + break; + case AwsParseState::Length8_4: + _pinfo.len += (uint64_t)data[0] << 32; + _pstate = AwsParseState::Length8_5; + break; + case AwsParseState::Length8_5: + _pinfo.len += (uint64_t)data[0] << 24; + _pstate = AwsParseState::Length8_6; + break; + case AwsParseState::Length8_6: + _pinfo.len += (uint64_t)data[0] << 16; + _pstate = AwsParseState::Length8_7; + break; + case AwsParseState::Length8_7: + _pinfo.len += (uint64_t)data[0] << 8; + _pstate = AwsParseState::Length8_8; + break; + case AwsParseState::Length8_8: + _pinfo.len += (uint64_t)data[0]; + _pstate = (_pinfo.masked) ? AwsParseState::Mask_0 : AwsParseState::Payload; + break; + + // Mask bytes + case AwsParseState::Mask_0: + _pinfo.mask[0] = data[0]; + _pstate = AwsParseState::Mask_1; + break; + case AwsParseState::Mask_1: + _pinfo.mask[1] = data[0]; + _pstate = AwsParseState::Mask_2; + break; + case AwsParseState::Mask_2: + _pinfo.mask[2] = data[0]; + _pstate = AwsParseState::Mask_3; + break; + case AwsParseState::Mask_3: + _pinfo.mask[3] = data[0]; + _pstate = AwsParseState::Payload; + break; + + // And finally, payload processing + case AwsParseState::Payload: + { + async_ws_log_v( + "[%s][%" PRIu32 "] DATA _pinfo: index: %" PRIu64 ", final: %" PRIu8 ", opcode: %" PRIu8 ", masked: %" PRIu8 ", len: %" PRIu64, _server->url(), + _clientId, _pinfo.index, _pinfo.final, _pinfo.opcode, _pinfo.masked, _pinfo.len + ); + const size_t datalen = std::min((size_t)(_pinfo.len - _pinfo.index), plen); + if (!_handleClientFrame(data, datalen, datalen == plen)) { // datalen == plen means that we are processing the last part of the current TCP packet + return; // client is now destroyed, so we must return immediately to avoid accessing any member + } + consume_byte = false; + data += datalen; + plen -= datalen; + if (_pinfo.index >= _pinfo.len) { + _pstate = AwsParseState::Start; + } + break; } + } // end switch over _pstate - // accumulate mask bytes - _pinfo.mask[_pinfo.masked - 1] = data[0]; + if (consume_byte) { + // Advance the buffer by one byte. Centralized to save copy-and-paste in so many states. data += 1; plen -= 1; - _pinfo.masked++; } + } // end while(plen > 0) - // all mask bytes read if we were reading them - _pstate = STATE_FRAME_DATA; + // data completely consumed + if ((_pinfo.opcode == WS_DISCONNECT) && (_pstate == AwsParseState::Mask_0) && (_pinfo.len == 0)) { + // Safari close frame edge case: masked bit set but no mask data + async_ws_log_v("[%s][%" PRIu32 "] DATA close frame with incomplete mask, treating as unmasked", _server->url(), _clientId); + _pinfo.masked = 0; + _pstate = AwsParseState::Payload; + } - // restore masked to 1 for backward compatibility - if (_pinfo.masked >= 5) { - async_ws_log_v("[%s][%" PRIu32 "] DATA mask read complete", _server->url(), _clientId); - _pinfo.masked = 1; + // A zero-length frame could be at the end of the packet; dispatch it now. + if ((_pstate == AwsParseState::Payload) && (_pinfo.len == 0)) { + if (!_handleClientFrame(data, 0, true)) { + return; } + _pstate = AwsParseState::Start; + } +} - const size_t datalen = std::min((size_t)(_pinfo.len - _pinfo.index), plen); +bool AsyncWebSocketClient::_handleClientFrame(uint8_t *data, size_t datalen, bool last) { + // Process a frame from the client. Header information is stored in _pinfo. + // Returns true on successful handling, false on error. - if (_pinfo.masked) { - for (size_t i = 0; i < datalen; i++) { - data[i] ^= _pinfo.mask[(_pinfo.index + i) % 4]; - } + if (_pinfo.masked) { + for (size_t i = 0; i < datalen; i++) { + data[i] ^= _pinfo.mask[(_pinfo.index + i) % 4]; } + } - if (_pinfo.index == 0) { // first fragment of the frame - // init message_opcode for this frame - // note: For next WS_CONTINUATION frames, they have opcode 0, so message_opcode will stay like the first frame - if (_pinfo.opcode == WS_TEXT || _pinfo.opcode == WS_BINARY) { - _pinfo.message_opcode = _pinfo.opcode; - } - // init frame number to 0 if only 1 frame or if this is the first frame of a fragmented message - if (_pinfo.final || datalen < _pinfo.len) { - _pinfo.num = 0; - } + if (_pinfo.index == 0) { // first fragment of the frame + // init message_opcode for this frame + // note: For next WS_CONTINUATION frames, they have opcode 0, so message_opcode will stay like the first frame + if (_pinfo.opcode == WS_TEXT || _pinfo.opcode == WS_BINARY) { + _pinfo.message_opcode = _pinfo.opcode; } + // init frame number to 0 if only 1 frame or if this is the first frame of a fragmented message + if (_pinfo.final || datalen < _pinfo.len) { + _pinfo.num = 0; + } + } - if ((datalen + _pinfo.index) < _pinfo.len) { // more fragments to read for this frame - _pstate = STATE_FRAME_DATA; - - if (datalen > 0) { - async_ws_log_v( - "[%s][%" PRIu32 "] DATA processing next fragment of %s frame %" PRIu32 ", index: %" PRIu64 ", len: %" PRIu32 "", _server->url(), _clientId, - (_pinfo.message_opcode == WS_TEXT) ? "text" : "binary", _pinfo.num, _pinfo.index, (uint32_t)datalen - ); - if (!_handleDataEvent(data, datalen, datalen == plen)) { // datalen == plen means that we are processing the last part of the current TCP packet - return; // stop processing on failure - } + if ((datalen + _pinfo.index) < _pinfo.len) { // more fragments to read for this frame + if (datalen > 0) { + async_ws_log_v( + "[%s][%" PRIu32 "] DATA processing next fragment of %s frame %" PRIu32 ", index: %" PRIu64 ", len: %" PRIu32 "", _server->url(), _clientId, + (_pinfo.message_opcode == WS_TEXT) ? "text" : "binary", _pinfo.num, _pinfo.index, (uint32_t)datalen + ); + if (!_handleDataEvent(data, datalen, last)) { + return false; // stop processing on failure } + } - // track index for next fragment - _pinfo.index += datalen; - - } else if ((datalen + _pinfo.index) == _pinfo.len) { // this is the last fragment for this frame - _pstate = STATE_FRAME_START; - - if (_pinfo.opcode == WS_DISCONNECT) { - async_ws_log_v("[%s][%" PRIu32 "] DATA WS_DISCONNECT", _server->url(), _clientId); - - if (datalen) { - uint16_t reasonCode = (uint16_t)(data[0] << 8) + data[1]; - char *reasonString = (char *)(data + 2); - if (reasonCode > 1001) { - _server->_handleEvent(this, WS_EVT_ERROR, (void *)&reasonCode, (uint8_t *)reasonString, strlen(reasonString)); - } - } - if (_status == WS_DISCONNECTING) { - _status = WS_DISCONNECTED; - if (_client) { - _client->close(); - } - return; // our object is now destroyed, so we must return immediately to avoid accessing any member - } else { - _status = WS_DISCONNECTING; - if (_client) { - _client->ackLater(); - } - _queueControl(WS_DISCONNECT, data, datalen); - } - - } else if (_pinfo.opcode == WS_PING) { - async_ws_log_v("[%s][%" PRIu32 "] DATA PING", _server->url(), _clientId); - _server->_handleEvent(this, WS_EVT_PING, NULL, NULL, 0); - _queueControl(WS_PONG, data, datalen); - - } else if (_pinfo.opcode == WS_PONG) { - async_ws_log_v("[%s][%" PRIu32 "] DATA PONG", _server->url(), _clientId); - if (datalen != AWSC_PING_PAYLOAD_LEN || memcmp(AWSC_PING_PAYLOAD, data, AWSC_PING_PAYLOAD_LEN) != 0) { - _server->_handleEvent(this, WS_EVT_PONG, NULL, data, datalen); + // track index for next fragment + _pinfo.index += datalen; + } else if ((datalen + _pinfo.index) == _pinfo.len) { // this is the last fragment for this frame + if (_pinfo.opcode == WS_DISCONNECT) { + async_ws_log_v("[%s][%" PRIu32 "] DATA WS_DISCONNECT", _server->url(), _clientId); + + if (datalen) { + uint16_t reasonCode = (uint16_t)(data[0] << 8) + data[1]; + char *reasonString = (char *)(data + 2); + if (reasonCode > 1001) { + _server->_handleEvent(this, WS_EVT_ERROR, (void *)&reasonCode, (uint8_t *)reasonString, strlen(reasonString)); } - - } else if (_pinfo.opcode < WS_DISCONNECT) { // continuation or text/binary frame - async_ws_log_v( - "[%s][%" PRIu32 "] DATA processing final fragment of %s frame %" PRIu32 ", index: %" PRIu64 ", len: %" PRIu32 "", _server->url(), _clientId, - (_pinfo.message_opcode == WS_TEXT) ? "text" : "binary", _pinfo.num, _pinfo.index, (uint32_t)datalen - ); - - if (!_handleDataEvent(data, datalen, datalen == plen)) { // datalen == plen means that we are processing the last part of the current TCP packet - return; // stop processing on failure + } + if (_status == WS_DISCONNECTING) { + _status = WS_DISCONNECTED; + if (_client) { + _client->close(); } - - if (_pinfo.final) { - _pinfo.num = 0; - } else { - _pinfo.num += 1; + return false; // our object is now destroyed, so we must return immediately to avoid accessing any member + } else { + _status = WS_DISCONNECTING; + if (_client) { + _client->ackLater(); } + _queueControl(WS_DISCONNECT, data, datalen); } - } else { - // unexpected frame error, close connection - _pstate = STATE_FRAME_START; + } else if (_pinfo.opcode == WS_PING) { + async_ws_log_v("[%s][%" PRIu32 "] DATA PING", _server->url(), _clientId); + _server->_handleEvent(this, WS_EVT_PING, NULL, NULL, 0); + _queueControl(WS_PONG, data, datalen); + } else if (_pinfo.opcode == WS_PONG) { + async_ws_log_v("[%s][%" PRIu32 "] DATA PONG", _server->url(), _clientId); + if (datalen != AWSC_PING_PAYLOAD_LEN || memcmp(AWSC_PING_PAYLOAD, data, AWSC_PING_PAYLOAD_LEN) != 0) { + _server->_handleEvent(this, WS_EVT_PONG, NULL, data, datalen); + } + } else if (_pinfo.opcode < WS_DISCONNECT) { // continuation or text/binary frame async_ws_log_v( - "[%s][%" PRIu32 "] DATA frame error: len: %u, index: %" PRIu64 ", total: %" PRIu64 "\n", _server->url(), _clientId, datalen, _pinfo.index, _pinfo.len + "[%s][%" PRIu32 "] DATA processing final fragment of %s frame %" PRIu32 ", index: %" PRIu64 ", len: %" PRIu32 "", _server->url(), _clientId, + (_pinfo.message_opcode == WS_TEXT) ? "text" : "binary", _pinfo.num, _pinfo.index, (uint32_t)datalen ); - _status = WS_DISCONNECTING; - if (_client) { - _client->ackLater(); + if (!_handleDataEvent(data, datalen, last)) { + return false; // stop processing on failure + } + + if (_pinfo.final) { + _pinfo.num = 0; + } else { + _pinfo.num += 1; } - _queueControl(WS_DISCONNECT, data, datalen); - break; } - data += datalen; - plen -= datalen; + _pinfo.index = _pinfo.len; // mark packet as complete + } else { + // unexpected frame protocol error - how is this possible? + async_ws_log_v( + "[%s][%" PRIu32 "] DATA frame error: len: %u, index: %" PRIu64 ", total: %" PRIu64 "\n", _server->url(), _clientId, datalen, _pinfo.index, _pinfo.len + ); + + _status = WS_DISCONNECTING; + if (_client) { + _client->ackLater(); + } + _queueControl(WS_DISCONNECT, data, datalen); } + + return true; } bool AsyncWebSocketClient::_handleDataEvent(uint8_t *data, size_t len, bool endOfPaquet) { diff --git a/src/AsyncWebSocket.h b/src/AsyncWebSocket.h index b03b1f78..e2d80c79 100644 --- a/src/AsyncWebSocket.h +++ b/src/AsyncWebSocket.h @@ -57,6 +57,8 @@ class AsyncWebSocket; class AsyncWebSocketResponse; class AsyncWebSocketClient; +enum class AwsParseState : uint8_t; + typedef struct { /** Message type as defined by enum AwsFrameType. * Note: Applications will only see WS_TEXT and WS_BINARY. @@ -161,7 +163,7 @@ class AsyncWebSocketClient { AsyncWebSocket *_server; uint32_t _clientId; AwsClientStatus _status; - uint8_t _pstate; + uint32_t _lastMessageTime; uint32_t _keepAlivePeriod; mutable asyncsrv::mutex_type _queue_lock; @@ -180,6 +182,8 @@ class AsyncWebSocketClient { size_t _framePayloadLen{0}; // payload length committed to the header of the in-flight frame size_t _frameSent{0}; // bytes of (header+in-flight payload) committed so far for the in-flight frame; 0 when idle + // The following fields are used to parse incoming frames. They are reset when a frame is fully received. + AwsParseState _pstate; AwsFrameInfo _pinfo; bool _queueControl(uint8_t opcode, const uint8_t *data = NULL, size_t len = 0, bool mask = false); @@ -193,6 +197,9 @@ class AsyncWebSocketClient { // Returns true on success, false on failure (e.g. memory allocation failure) bool _handleDataEvent(uint8_t *data, size_t len, bool endOfPaquet); + // Internal function to handle a frame from the client. Header information is stored in _pinfo. + bool _handleClientFrame(uint8_t *data, size_t datalen, bool last); + public: void *_tempObject; From d1fefc983c83b0c54547da8ecea2de167bee2241 Mon Sep 17 00:00:00 2001 From: Will Miles Date: Fri, 28 Aug 2026 20:39:37 -0400 Subject: [PATCH 2/3] fix(websockets): Fix handling of very large frames --- src/AsyncWebSocket.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/AsyncWebSocket.cpp b/src/AsyncWebSocket.cpp index 65c28e58..84dcbb69 100644 --- a/src/AsyncWebSocket.cpp +++ b/src/AsyncWebSocket.cpp @@ -624,7 +624,7 @@ void AsyncWebSocketClient::_onData(void *pbuf, size_t plen) { "[%s][%" PRIu32 "] DATA _pinfo: index: %" PRIu64 ", final: %" PRIu8 ", opcode: %" PRIu8 ", masked: %" PRIu8 ", len: %" PRIu64, _server->url(), _clientId, _pinfo.index, _pinfo.final, _pinfo.opcode, _pinfo.masked, _pinfo.len ); - const size_t datalen = std::min((size_t)(_pinfo.len - _pinfo.index), plen); + const size_t datalen = static_cast(std::min(_pinfo.len - _pinfo.index, plen)); if (!_handleClientFrame(data, datalen, datalen == plen)) { // datalen == plen means that we are processing the last part of the current TCP packet return; // client is now destroyed, so we must return immediately to avoid accessing any member } From 6cca9248bf97f02723becf988fb405def7f186e4 Mon Sep 17 00:00:00 2001 From: Will Miles Date: Fri, 28 Aug 2026 22:23:23 -0400 Subject: [PATCH 3/3] fix(websockets): Handle torn control frames If a control frame spans multiple TCP packets, buffer the data so that the frame can be processed once fully received. This ensures that the frame can be correctly handled instead of generating invalid PONG responses or overrunning the buffer with a disconnect reason. --- src/AsyncWebSocket.cpp | 137 +++++++++++++++++++++++++++-------------- src/AsyncWebSocket.h | 28 ++++++--- 2 files changed, 111 insertions(+), 54 deletions(-) diff --git a/src/AsyncWebSocket.cpp b/src/AsyncWebSocket.cpp index 84dcbb69..ba0c9f1a 100644 --- a/src/AsyncWebSocket.cpp +++ b/src/AsyncWebSocket.cpp @@ -51,6 +51,7 @@ enum class AwsParseState : uint8_t { Mask_2, Mask_3, Payload, + Error }; static AsyncWebSocketSharedBuffer makeSharedBuffer(const uint8_t *message, size_t len) { @@ -545,6 +546,15 @@ void AsyncWebSocketClient::_onData(void *pbuf, size_t plen) { // Second header byte _pinfo.masked = ((data[0] & 0x80) != 0) ? 1 : 0; _pinfo.len = data[0] & 0x7F; + // Validate length for control frames (must be <= 125) + if ((_pinfo.opcode & 0x08) != 0 && _pinfo.len > 125) { + async_ws_log_v( + "[%s][%" PRIu32 "] DATA control frame length error: opcode: %" PRIu8 ", len: %" PRIu64 "\n", _server->url(), _clientId, _pinfo.opcode, _pinfo.len + ); + close(WS_CLOSE_PROTOCOL_ERROR, nullptr); // Send disconnect message with protocol error code + _pstate = AwsParseState::Error; + return; // Abort processing this frame + } // Select length type if (_pinfo.len == 126) { _pstate = AwsParseState::Length2_1; @@ -636,6 +646,10 @@ void AsyncWebSocketClient::_onData(void *pbuf, size_t plen) { } break; } + + case AwsParseState::Error: + async_ws_log_v("[%s][%" PRIu32 "] DATA error state, ignoring data len: %" PRIu64, _server->url(), _clientId, plen); + return; // ignore any further data } // end switch over _pstate if (consume_byte) { @@ -686,53 +700,35 @@ bool AsyncWebSocketClient::_handleClientFrame(uint8_t *data, size_t datalen, boo if ((datalen + _pinfo.index) < _pinfo.len) { // more fragments to read for this frame if (datalen > 0) { - async_ws_log_v( - "[%s][%" PRIu32 "] DATA processing next fragment of %s frame %" PRIu32 ", index: %" PRIu64 ", len: %" PRIu32 "", _server->url(), _clientId, - (_pinfo.message_opcode == WS_TEXT) ? "text" : "binary", _pinfo.num, _pinfo.index, (uint32_t)datalen - ); - if (!_handleDataEvent(data, datalen, last)) { - return false; // stop processing on failure - } - } - - // track index for next fragment - _pinfo.index += datalen; - } else if ((datalen + _pinfo.index) == _pinfo.len) { // this is the last fragment for this frame - if (_pinfo.opcode == WS_DISCONNECT) { - async_ws_log_v("[%s][%" PRIu32 "] DATA WS_DISCONNECT", _server->url(), _clientId); - - if (datalen) { - uint16_t reasonCode = (uint16_t)(data[0] << 8) + data[1]; - char *reasonString = (char *)(data + 2); - if (reasonCode > 1001) { - _server->_handleEvent(this, WS_EVT_ERROR, (void *)&reasonCode, (uint8_t *)reasonString, strlen(reasonString)); - } - } - if (_status == WS_DISCONNECTING) { - _status = WS_DISCONNECTED; - if (_client) { - _client->close(); + if (_pinfo.opcode < WS_DISCONNECT) { // continuation or text/binary frame + async_ws_log_v( + "[%s][%" PRIu32 "] DATA processing next fragment of %s frame %" PRIu32 ", index: %" PRIu64 ", len: %" PRIu32 "", _server->url(), _clientId, + (_pinfo.message_opcode == WS_TEXT) ? "text" : "binary", _pinfo.num, _pinfo.index, (uint32_t)datalen + ); + if (!_handleDataEvent(data, datalen, last)) { + return false; // stop processing on failure } - return false; // our object is now destroyed, so we must return immediately to avoid accessing any member } else { - _status = WS_DISCONNECTING; - if (_client) { - _client->ackLater(); + // Control frame fragmented across TCP packets. We must buffer the data until we have the complete frame. + if (!_pbuffer) { + uint8_t *pbuf = new (std::nothrow) uint8_t[(size_t)_pinfo.len]; // cast is safe because _pinfo.len is guaranteed to be <= 125 for control frames + if (!pbuf) { + async_ws_log_e("[%s][%" PRIu32 "] DATA failed to allocate buffer for control frame", _server->url(), _clientId); + close(WS_CLOSE_INTERNAL_ERROR, nullptr); // Close the connection with a protocol error code + _pstate = AwsParseState::Error; + return false; + } + _pbuffer.reset(pbuf); } - _queueControl(WS_DISCONNECT, data, datalen); + // Save data in buffer + memcpy(_pbuffer.get() + _pinfo.index, data, datalen); } - } else if (_pinfo.opcode == WS_PING) { - async_ws_log_v("[%s][%" PRIu32 "] DATA PING", _server->url(), _clientId); - _server->_handleEvent(this, WS_EVT_PING, NULL, NULL, 0); - _queueControl(WS_PONG, data, datalen); - - } else if (_pinfo.opcode == WS_PONG) { - async_ws_log_v("[%s][%" PRIu32 "] DATA PONG", _server->url(), _clientId); - if (datalen != AWSC_PING_PAYLOAD_LEN || memcmp(AWSC_PING_PAYLOAD, data, AWSC_PING_PAYLOAD_LEN) != 0) { - _server->_handleEvent(this, WS_EVT_PONG, NULL, data, datalen); - } - } else if (_pinfo.opcode < WS_DISCONNECT) { // continuation or text/binary frame + // track index for next fragment + _pinfo.index += datalen; + } + } else if ((datalen + _pinfo.index) == _pinfo.len) { // this is the last fragment for this frame + if (_pinfo.opcode < WS_DISCONNECT) { // most likely case: continuation or text/binary frame async_ws_log_v( "[%s][%" PRIu32 "] DATA processing final fragment of %s frame %" PRIu32 ", index: %" PRIu64 ", len: %" PRIu32 "", _server->url(), _clientId, (_pinfo.message_opcode == WS_TEXT) ? "text" : "binary", _pinfo.num, _pinfo.index, (uint32_t)datalen @@ -747,20 +743,67 @@ bool AsyncWebSocketClient::_handleClientFrame(uint8_t *data, size_t datalen, boo } else { _pinfo.num += 1; } + } else { // control frame + if (_pbuffer) { + memcpy(_pbuffer.get() + _pinfo.index, data, datalen); + data = _pbuffer.get(); + datalen = _pinfo.len; + } + + if (_pinfo.opcode == WS_DISCONNECT) { + async_ws_log_v("[%s][%" PRIu32 "] DATA WS_DISCONNECT", _server->url(), _clientId); + + // Pass up the close frame error information + if (datalen >= 2) { + uint16_t reasonCode = (uint16_t)(data[0] << 8) + data[1]; + char *reasonString = (char *)(data + 2); + if (reasonCode > WS_CLOSE_GOING_AWAY) { + _server->_handleEvent(this, WS_EVT_ERROR, (void *)&reasonCode, (uint8_t *)reasonString, strnlen(reasonString, datalen - 2)); + } + } + if (_status == WS_DISCONNECTING) { + _status = WS_DISCONNECTED; + if (_client) { + _client->close(); + } + return false; // our object is now destroyed, so we must return immediately to avoid accessing any member + } else { + _status = WS_DISCONNECTING; + if (_client) { + _client->ackLater(); + } + _queueControl(WS_DISCONNECT, data, datalen); + } + + } else if (_pinfo.opcode == WS_PING) { + async_ws_log_v("[%s][%" PRIu32 "] DATA PING", _server->url(), _clientId); + _server->_handleEvent(this, WS_EVT_PING, NULL, NULL, 0); + _queueControl(WS_PONG, data, datalen); + + } else if (_pinfo.opcode == WS_PONG) { + async_ws_log_v("[%s][%" PRIu32 "] DATA PONG", _server->url(), _clientId); + if (datalen != AWSC_PING_PAYLOAD_LEN || memcmp(AWSC_PING_PAYLOAD, data, AWSC_PING_PAYLOAD_LEN) != 0) { + _server->_handleEvent(this, WS_EVT_PONG, NULL, data, datalen); + } + } else { + async_ws_log_v("[%s][%" PRIu32 "] DATA unknown control frame: %" PRIu8, _server->url(), _clientId, _pinfo.opcode); + close(WS_CLOSE_PROTOCOL_ERROR, nullptr); // Close the connection with a protocol error code + _pstate = AwsParseState::Error; + return false; + } } _pinfo.index = _pinfo.len; // mark packet as complete + _pbuffer.reset(); // free any control frame buffer } else { // unexpected frame protocol error - how is this possible? async_ws_log_v( "[%s][%" PRIu32 "] DATA frame error: len: %u, index: %" PRIu64 ", total: %" PRIu64 "\n", _server->url(), _clientId, datalen, _pinfo.index, _pinfo.len ); - _status = WS_DISCONNECTING; - if (_client) { - _client->ackLater(); - } - _queueControl(WS_DISCONNECT, data, datalen); + close(WS_CLOSE_PROTOCOL_ERROR, nullptr); // Close the connection with a protocol error code + _pstate = AwsParseState::Error; + return false; } return true; diff --git a/src/AsyncWebSocket.h b/src/AsyncWebSocket.h index e2d80c79..07420b2c 100644 --- a/src/AsyncWebSocket.h +++ b/src/AsyncWebSocket.h @@ -88,14 +88,27 @@ typedef enum { WS_CONNECTED, WS_DISCONNECTING } AwsClientStatus; -typedef enum { - WS_CONTINUATION, - WS_TEXT, - WS_BINARY, - WS_DISCONNECT = 0x08, - WS_PING, - WS_PONG +typedef enum { // RFC6455 frame types, section 5.2 + WS_CONTINUATION = 0x0, + WS_TEXT = 0x1, + WS_BINARY = 0x2, + WS_DISCONNECT = 0x8, + WS_PING = 0x9, + WS_PONG = 0xA } AwsFrameType; +typedef enum { // RFC6455 close reason codes, section 7.4.1 + WS_CLOSE_NORMAL = 1000, + WS_CLOSE_GOING_AWAY = 1001, + WS_CLOSE_PROTOCOL_ERROR = 1002, + WS_CLOSE_UNSUPPORTED_DATA = 1003, + WS_CLOSE_NO_STATUS_RECEIVED = 1005, + WS_CLOSE_ABNORMAL_CLOSURE = 1006, + WS_CLOSE_INVALID_PAYLOAD = 1007, + WS_CLOSE_POLICY_VIOLATION = 1008, + WS_CLOSE_MESSAGE_TOO_BIG = 1009, + WS_CLOSE_MANDATORY_EXTENSION = 1010, + WS_CLOSE_INTERNAL_ERROR = 1011, +} AwsCloseCode; typedef enum { WS_EVT_CONNECT, WS_EVT_DISCONNECT, @@ -185,6 +198,7 @@ class AsyncWebSocketClient { // The following fields are used to parse incoming frames. They are reset when a frame is fully received. AwsParseState _pstate; AwsFrameInfo _pinfo; + std::unique_ptr _pbuffer; // Payload buffer for torn control frames bool _queueControl(uint8_t opcode, const uint8_t *data = NULL, size_t len = 0, bool mask = false); bool _queueMessage(AsyncWebSocketSharedBuffer buffer, uint8_t opcode = WS_TEXT, bool mask = false);