diff --git a/src/AsyncWebSocket.cpp b/src/AsyncWebSocket.cpp index 7e1d9e9f..ba0c9f1a 100644 --- a/src/AsyncWebSocket.cpp +++ b/src/AsyncWebSocket.cpp @@ -31,12 +31,29 @@ #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, + Error +}; + static AsyncWebSocketSharedBuffer makeSharedBuffer(const uint8_t *message, size_t len) { if (message) { return std::make_shared>(message, message + len); @@ -167,8 +184,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,126 +528,237 @@ 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) - ); - - 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 + "[%s][%" PRIu32 "] DATA plen: %" PRIu32 ", _pstate: %" PRIu8 ", _status: %" PRIu8, _server->url(), _clientId, static_cast(plen), + static_cast(_pstate), static_cast(_status) ); - // 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; + // 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 } - - // 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; + // 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; + } + 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 = 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 + } + consume_byte = false; + data += datalen; + plen -= datalen; + if (_pinfo.index >= _pinfo.len) { + _pstate = AwsParseState::Start; + } + break; } - // accumulate mask bytes - _pinfo.mask[_pinfo.masked - 1] = data[0]; + 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) { + // 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) { + if ((datalen + _pinfo.index) < _pinfo.len) { // more fragments to read for this frame + if (datalen > 0) { + 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, datalen == plen)) { // datalen == plen means that we are processing the last part of the current TCP packet - return; // stop processing on failure + if (!_handleDataEvent(data, datalen, last)) { + return false; // stop processing on failure + } + } else { + // 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); } + // Save data in buffer + memcpy(_pbuffer.get() + _pinfo.index, 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) { // 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 + ); - } else if ((datalen + _pinfo.index) == _pinfo.len) { // this is the last fragment for this frame - _pstate = STATE_FRAME_START; + if (!_handleDataEvent(data, datalen, last)) { + return false; // stop processing on failure + } + + if (_pinfo.final) { + _pinfo.num = 0; + } 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); - if (datalen) { + // 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 > 1001) { - _server->_handleEvent(this, WS_EVT_ERROR, (void *)&reasonCode, (uint8_t *)reasonString, strlen(reasonString)); + 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) { @@ -638,7 +766,7 @@ void AsyncWebSocketClient::_onData(void *pbuf, size_t plen) { if (_client) { _client->close(); } - return; // our object is now destroyed, so we must return immediately to avoid accessing any member + return false; // our object is now destroyed, so we must return immediately to avoid accessing any member } else { _status = WS_DISCONNECTING; if (_client) { @@ -657,43 +785,28 @@ void AsyncWebSocketClient::_onData(void *pbuf, size_t plen) { 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 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 (_pinfo.final) { - _pinfo.num = 0; - } else { - _pinfo.num += 1; - } - } - - } else { - // unexpected frame error, close connection - _pstate = STATE_FRAME_START; - - 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(); + } 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; } - _queueControl(WS_DISCONNECT, data, datalen); - break; } - data += datalen; - plen -= datalen; + _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 + ); + + close(WS_CLOSE_PROTOCOL_ERROR, nullptr); // Close the connection with a protocol error code + _pstate = AwsParseState::Error; + return false; } + + 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..07420b2c 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. @@ -86,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, @@ -161,7 +176,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,7 +195,10 @@ 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; + 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); @@ -193,6 +211,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;