Handle edge-case in which we receive a partial WS header

me-no-dev#953

Author: nthd

src/AsyncWebSocket.cpp

@@ -483,6 +483,8 @@ AsyncWebSocketClient::AsyncWebSocketClient(AsyncWebServerRequest *request, Async
   _clientId = _server->_getNextId();
   _status = WS_CONNECTED;
   _pstate = 0;
+  _partialHeader = nullptr;
+  _partialHeaderLen = 0;
   _lastMessageTime = millis();
   _keepAlivePeriod = 0;
   _client->setRxTimeout(0);
@@ -670,116 +672,197 @@ void AsyncWebSocketClient::_onData(void *pbuf, size_t plen)
   {
     if (!_pstate)
     {
-      const uint8_t *fdata = data;
+      ssize_t dataPayloadOffset = 0;
+      const uint8_t *headerBuf = data;
+
+      // plen is backed up to initialPlen because, in case we receive a partial header, we would like to undo all of our
+      // parsing and copy all of what we have of the header into a buffer for later use.
+      // plen is modified during the parsing attempt, so if we don't back it up we won't know how much we need to copy.
+      // partialHeaderLen is also backed up for the same reason.
+      size_t initialPlen = plen;
+      size_t partialHeaderLen = 0;
+
+      if (_partialHeaderLen > 0)
+      {
+        // We previously received a truncated header. Recover it by doing the following:
+        // - Copy the new header chunk into the previous partial header, filling the buffer. It is allocated as a
+        //   buffer in a class field.
+        // - Change *headerBuf to point to said buffer
+        // - Update the length counters so that:
+        //   - The initialPlen and plen, which refer to the length of the remaining packet data, also accounts for the
+        //     previously received truncated header
+        //   - The dataPayloadOffset, which is the offset after the header at which the payload begins, so that it
+        //     refers to a point potentially before the beginning of the buffer. As we parse the header we increment it,
+        //     and we can pretty much guarantee it will go back to being positive unless there is a major bug.
+        //   - The class _partialHeaderLen is back to zero since we took ownership of the contained data.
+        memcpy(_partialHeader + _partialHeaderLen, data,
+               std::min(plen, (size_t)WS_MAX_HEADER_LEN - _partialHeaderLen));
+        headerBuf = _partialHeader;
+        initialPlen += _partialHeaderLen;
+        plen += _partialHeaderLen;
+        dataPayloadOffset -= _partialHeaderLen;
+        partialHeaderLen = _partialHeaderLen;
+
+        _partialHeaderLen = 0;
+      }
+
+      // The following series of gotos could have been a try-catch but we are likely being built with -fno-exceptions
+      if (plen < 2)
+        goto _exceptionHandleFailPartialHeader;
       _pinfo.index = 0;
-      _pinfo.final = (fdata[0] & 0x80) != 0;
-      _pinfo.opcode = fdata[0] & 0x0F;
-      _pinfo.masked = (fdata[1] & 0x80) != 0;
-      _pinfo.len = fdata[1] & 0x7F;
-      data += 2;
+      _pinfo.final = (headerBuf[0] & 0x80) != 0;
+      _pinfo.opcode = headerBuf[0] & 0x0F;
+      _pinfo.masked = (headerBuf[1] & 0x80) != 0;
+      _pinfo.len = headerBuf[1] & 0x7F;
+      dataPayloadOffset += 2;
       plen -= 2;
+
       if (_pinfo.len == 126)
       {
-        _pinfo.len = fdata[3] | (uint16_t)(fdata[2]) << 8;
-        data += 2;
-        plen -= 2;
-      }
-      else if (_pinfo.len == 127)
-      {
-        _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;
-      }
+        if (plen < 2)
+          goto _exceptionHandleFailPartialHeader;
 
-      if (_pinfo.masked)
-      {
-        memcpy(_pinfo.mask, data, 4);
-        data += 4;
-        plen -= 4;
+        _pinfo.len = headerBuf[3] | (uint16_t)(headerBuf[2]) << 8;
+        dataPayloadOffset += 2;
+        plen -= 2;
       }
     }
-
-    const size_t datalen = std::min((size_t)(_pinfo.len - _pinfo.index), plen);
-    const auto datalast = data[datalen];
+    else if (_pinfo.len == 127)
+    {
+      if (plen < 8)
+        goto _exceptionHandleFailPartialHeader;
+
+      _pinfo.len = headerBuf[9] | (uint16_t)(headerBuf[8]) << 8 | (uint32_t)(headerBuf[7]) << 16 |
+                   (uint32_t)(headerBuf[6]) << 24 | (uint64_t)(headerBuf[5]) << 32 | (uint64_t)(headerBuf[4]) << 40 |
+                   (uint64_t)(headerBuf[3]) << 48 | (uint64_t)(headerBuf[2]) << 56;
+      dataPayloadOffset += 8;
+      plen -= 8;
+    }
 
     if (_pinfo.masked)
     {
-      for (size_t i = 0; i < datalen; i++)
-        data[i] ^= _pinfo.mask[(_pinfo.index + i) % 4];
+      if (plen < 4)
+        goto _exceptionHandleFailPartialHeader;
+
+      memcpy(_pinfo.mask, headerBuf + dataPayloadOffset + partialHeaderLen, 4);
+      dataPayloadOffset += 4;
+      plen -= 4;
     }
 
-    if ((datalen + _pinfo.index) < _pinfo.len)
+    // Yes I know the control flow here isn't 100% legible but we must support -fno-exceptions.
+    // If we got to this point it means we did NOT receive a truncated header, therefore we can skip the exception
+    // handling.
+    // Control flow resumes after the following block.
+    goto _headerParsingSuccessful;
+
+  // We DID receive a truncated header:
+  // - We copy it to our buffer and set the _partialHeaderLen
+  // - We return early
+  // This will trigger the partial recovery at the next call of this method, once more data is received and we have
+  // a full header.
+  _exceptionHandleFailPartialHeader:
+  {
+    if (initialPlen <= WS_MAX_HEADER_LEN)
+    {
+      // If initialPlen > WS_MAX_HEADER_LEN there must be something wrong with this code. It should never happen but
+      // but it's better safe than sorry.
+      memcpy(_partialHeader, headerBuf, initialPlen * sizeof(uint8_t));
+      _partialHeaderLen = initialPlen;
+    }
+    else
     {
-      _pstate = 1;
+      DEBUGF("[AsyncWebSocketClient::_onData] initialPlen (= %d) > WS_MAX_HEADER_LEN (= %d)\n", initialPlen,
+             WS_MAX_HEADER_LEN);
+    }
+    return;
+  }
+
+  _headerParsingSuccessful:
+
+    data += dataPayloadOffset;
+  }
+
+  const size_t datalen = std::min((size_t)(_pinfo.len - _pinfo.index), plen);
+  const auto datalast = data[datalen];
+
+  if (_pinfo.masked)
+  {
+    for (size_t i = 0; i < datalen; i++)
+      data[i] ^= _pinfo.mask[(_pinfo.index + i) % 4];
+  }
+
+  if ((datalen + _pinfo.index) < _pinfo.len)
+  {
+    _pstate = 1;
 
-      if (_pinfo.index == 0)
+    if (_pinfo.index == 0)
+    {
+      if (_pinfo.opcode)
       {
-        if (_pinfo.opcode)
-        {
-          _pinfo.message_opcode = _pinfo.opcode;
-          _pinfo.num = 0;
-        }
-        else
-          _pinfo.num += 1;
+        _pinfo.message_opcode = _pinfo.opcode;
+        _pinfo.num = 0;
       }
-      _server->_handleEvent(this, WS_EVT_DATA, (void *)&_pinfo, (uint8_t *)data, datalen);
-
-      _pinfo.index += datalen;
+      else
+        _pinfo.num += 1;
     }
-    else if ((datalen + _pinfo.index) == _pinfo.len)
+    _server->_handleEvent(this, WS_EVT_DATA, (void *)&_pinfo, (uint8_t *)data, datalen);
+
+    _pinfo.index += datalen;
+  }
+  else if ((datalen + _pinfo.index) == _pinfo.len)
+  {
+    _pstate = 0;
+    if (_pinfo.opcode == WS_DISCONNECT)
     {
-      _pstate = 0;
-      if (_pinfo.opcode == WS_DISCONNECT)
+      if (datalen)
       {
-        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;
-          _client->close(true);
-        }
-        else
+        uint16_t reasonCode = (uint16_t)(data[0] << 8) + data[1];
+        char *reasonString = (char *)(data + 2);
+        if (reasonCode > 1001)
         {
-          _status = WS_DISCONNECTING;
-          _client->ackLater();
-          _queueControl(new AsyncWebSocketControl(WS_DISCONNECT, data, datalen));
+          _server->_handleEvent(this, WS_EVT_ERROR, (void *)&reasonCode, (uint8_t *)reasonString, strlen(reasonString));
         }
       }
-      else if (_pinfo.opcode == WS_PING)
+      if (_status == WS_DISCONNECTING)
       {
-        _queueControl(new AsyncWebSocketControl(WS_PONG, data, datalen));
+        _status = WS_DISCONNECTED;
+        _client->close(true);
       }
-      else if (_pinfo.opcode == WS_PONG)
+      else
       {
-        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 < 8)
-      { // continuation or text/binary frame
-        _server->_handleEvent(this, WS_EVT_DATA, (void *)&_pinfo, data, datalen);
+        _status = WS_DISCONNECTING;
+        _client->ackLater();
+        _queueControl(new AsyncWebSocketControl(WS_DISCONNECT, data, datalen));
       }
     }
-    else
+    else if (_pinfo.opcode == WS_PING)
     {
-      // os_printf("frame error: len: %u, index: %llu, total: %llu\n", datalen, _pinfo.index, _pinfo.len);
-      // what should we do?
-      break;
+      _queueControl(new AsyncWebSocketControl(WS_PONG, data, datalen));
+    }
+    else if (_pinfo.opcode == WS_PONG)
+    {
+      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 < 8)
+    { // continuation or text/binary frame
+      _server->_handleEvent(this, WS_EVT_DATA, (void *)&_pinfo, data, datalen);
+    }
+  }
+  else
+  {
+    // os_printf("frame error: len: %u, index: %llu, total: %llu\n", datalen, _pinfo.index, _pinfo.len);
+    // what should we do?
+    break;
+  }
 
-    // restore byte as _handleEvent may have added a null terminator i.e., data[len] = 0;
-    if (datalen > 0)
-      data[datalen] = datalast;
+  // restore byte as _handleEvent may have added a null terminator i.e., data[len] = 0;
+  if (datalen > 0)
+    data[datalen] = datalast;
 
-    data += datalen;
-    plen -= datalen;
-  }
+  data += datalen;
+  plen -= datalen;
+}
 }
 
 size_t AsyncWebSocketClient::printf(const char *format, ...)