[ Web Proxy ]
URL:
Viewing: https://raw.githubusercontent.com/ESP32Async/ESPAsyncWebServer/main/src/AsyncWebSocket.cpp [Back]  [Original]

// SPDX-License-Identifier: LGPL-3.0-or-later
// Copyright 2016-2026 Hristo Gochkov, Mathieu Carbou, Emil Muratov, Will Miles

#include "AsyncWebSocket.h"
#include "AsyncWebServerLogging.h"

#include 

#if defined(ESP32)
#if ESP_IDF_VERSION_MAJOR < 5
#include "BackPort_SHA1Builder.h"
#else
#include 
#endif
#include 
#elif defined(TARGET_RP2040) || defined(TARGET_RP2350) || defined(PICO_RP2040) || defined(PICO_RP2350) || defined(ESP8266)
#include 
#elif defined(LIBRETINY)
#include 
#elif defined(HOST)
#include "BackPort_SHA1Builder.h"
#ifndef FPSTR
#define FPSTR (const char *)
#endif
#endif

#include 
#include 
#include 
#include 
#include 
#include 

#define STATE_FRAME_START 0
#define STATE_FRAME_MASK  1
#define STATE_FRAME_DATA  2

using namespace asyncsrv;

static AsyncWebSocketSharedBuffer makeSharedBuffer(const uint8_t *message, size_t len) {
  if (message) {
    return std::make_shared(message, message + len);
  } else {
    return std::make_shared(len);
  }
}

static size_t webSocketSendFrameWindow(AsyncClient *client) {
  if (!client || !client->canSend()) {
    return 0;
  }
  size_t space = client->space();
  if (space < 9) {
    return 0;
  }
  return space - 8;
}

// wire header length for a frame with this payload length/masking -- pure
// function of RFC 6455 framing rules, cheap enough to recompute on demand
// rather than cache
static uint8_t webSocketFrameHeaderLen(size_t payloadLen, bool mask) {
  return 2 + (mask ? 4 : 0) + ((payloadLen > 125) ? 2 : 0);
}

// Stack-allocated chunk size for masking payload data
static constexpr size_t WS_MASK_CHUNK_SIZE = 128;

// Commits any not-yet-committed bytes of one WS frame (header+payload) to
// the client, resuming from `frameSent` bytes already committed by a prior
// call for this exact frame. final/opcode/mask/maskKey/data/len must stay
// identical across calls for the same frame -- once any header byte is
// committed those values are fixed on the wire and can't change. Returns
// the number of *additional* bytes committed by this call (0 if none);
// never assumes an add() fully succeeds, since some AsyncClient backends
// (e.g. SSL) can commit a genuine partial amount.
static size_t
  webSocketAddFrame(AsyncClient *client, bool final, uint8_t opcode, bool mask, const uint8_t maskKey[4], const uint8_t *data, size_t len, size_t frameSent) {
  if (!client || !client->canSend()) {
    return 0;
  }

  size_t committed = 0;

  const uint8_t headLen = webSocketFrameHeaderLen(len, mask);
  if (frameSent < headLen) {
    uint8_t hdr[8] = {0, 0, 0, 0, 0, 0, 0, 0};
    hdr[0] = (opcode & 0x0F) | (final ? 0x80 : 0);
    if (len < 126) {
      hdr[1] = len & 0x7F;
    } else {
      hdr[1] = 126;
      hdr[2] = (uint8_t)((len >> 8) & 0xFF);
      hdr[3] = (uint8_t)(len & 0xFF);
    }
    if (mask) {
      hdr[1] |= 0x80;
      memcpy(hdr + (headLen - 4), maskKey, 4);
    }

    size_t added = client->add((const char *)(hdr + frameSent), headLen - frameSent, ASYNC_WRITE_FLAG_COPY | ((len > 0) ? ASYNC_WRITE_FLAG_MORE : 0));
    committed += added;
    frameSent += added;
    if (frameSent < headLen) {
      return committed;
    }
  }

  size_t payloadSent = frameSent - headLen;
  if (payloadSent < len) {
    size_t remaining = len - payloadSent;
    if (mask) {
      // The data buffer is shared, but each client masks with its own random key.
      // Mask in to a scratch buffer one chunk at a time to avoid malloc overhead.
      uint8_t masked_data[WS_MASK_CHUNK_SIZE];
      while (remaining > 0) {
        const size_t chunk = std::min(remaining, sizeof(masked_data));
        for (size_t i = 0; i < chunk; i++) {
          masked_data[i] = data[payloadSent + i] ^ maskKey[(payloadSent + i) % 4];
        }
        const size_t added = client->add((const char *)masked_data, chunk, ASYNC_WRITE_FLAG_COPY | ((chunk == remaining) ? 0 : ASYNC_WRITE_FLAG_MORE));
        committed += added;
        if (added < chunk) {
          break;  // out of space (or partial commit) -- resume here next time
        }
        payloadSent += added;
        remaining -= added;
      }
    } else {
      committed += client->add((const char *)(data + payloadSent), remaining, ASYNC_WRITE_FLAG_COPY);
    }
  }

  return committed;
}

/*
 *    AsyncWebSocketMessageBuffer
 */

AsyncWebSocketMessageBuffer::AsyncWebSocketMessageBuffer(const uint8_t *data, size_t size) : _buffer(std::make_shared(size)) {
  if (_buffer->capacity() < size) {
    _buffer->reserve(size);
  } else {
    std::memcpy(_buffer->data(), data, size);
  }
}

AsyncWebSocketMessageBuffer::AsyncWebSocketMessageBuffer(size_t size) : _buffer(std::make_shared(size)) {
  if (_buffer->capacity() < size) {
    _buffer->reserve(size);
  }
}

bool AsyncWebSocketMessageBuffer::reserve(size_t size) {
  if (_buffer->capacity() >= size) {
    return true;
  }
  _buffer->reserve(size);
  return _buffer->capacity() >= size;
}

/*
 * Async WebSocket Client
 */
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->setRxTimeout(0);
  _client->onError(
    [](void *r, AsyncClient *c, int8_t error) {
      (void)c;
      ((AsyncWebSocketClient *)(r))->_onError(error);
    },
    this
  );
  _client->onAck(
    [](void *r, AsyncClient *c, size_t len, uint32_t time) {
      (void)c;
      ((AsyncWebSocketClient *)(r))->_onAck(len, time);
    },
    this
  );
  _client->onDisconnect(
    [](void *r, AsyncClient *c) {
      ((AsyncWebSocketClient *)(r))->_onDisconnect();
      delete c;
    },
    this
  );
  _client->onTimeout(
    [](void *r, AsyncClient *c, uint32_t time) {
      (void)c;
      ((AsyncWebSocketClient *)(r))->_onTimeout(time);
    },
    this
  );
  _client->onData(
    [](void *r, AsyncClient *c, void *buf, size_t len) {
      (void)c;
      ((AsyncWebSocketClient *)(r))->_onData(buf, len);
    },
    this
  );
  _client->onPoll(
    [](void *r, AsyncClient *c) {
      (void)c;
      ((AsyncWebSocketClient *)(r))->_onPoll();
    },
    this
  );
  memset(&_pinfo, 0, sizeof(_pinfo));
}

AsyncWebSocketClient::~AsyncWebSocketClient() {
  {
    asyncsrv::lock_guard_type lock(_queue_lock);
    _messageQueue.clear();
    _controlQueue.clear();
  }
  _server->_handleEvent(this, WS_EVT_DISCONNECT, NULL, NULL, 0);
}

void AsyncWebSocketClient::_onAck(size_t len, uint32_t time) {
  (void)time;
  _lastMessageTime = millis();

  asyncsrv::unique_lock_type lock(_queue_lock);

  async_ws_log_v("[%s][%" PRIu32 "] ACK(%u)", _server->url(), _clientId, len);

  // Messages are dequeued as soon as they're fully add()'ed, not on ack --
  // an ack only ever matters here as a sign that space() may have grown.
  _runQueue(lock);
}

void AsyncWebSocketClient::_onPoll() {
  asyncsrv::unique_lock_type lock(_queue_lock);

  if (!_client) {
    return;
  }

  if (_client && _client->canSend() && (!_controlQueue.empty() || !_messageQueue.empty())) {
    _runQueue(lock);
  } else if (_keepAlivePeriod > 0 && (millis() - _lastMessageTime) >= _keepAlivePeriod && (_controlQueue.empty() && _messageQueue.empty())) {
    lock.unlock();
    ping((uint8_t *)AWSC_PING_PAYLOAD, AWSC_PING_PAYLOAD_LEN);
  }
}

void AsyncWebSocketClient::_runQueue(asyncsrv::unique_lock_type &lock) {
  // all calls to this method MUST be protected by a mutex lock, passed in as `lock`
  if (!_client) {
    return;
  }

  bool needs_send = false;
  while (true) {
    if (_frameSent == 0) {
      // idle: pick the next target by priority (control first)
      if (!_controlQueue.empty()) {
        _sendingControl = true;
      } else if (!_messageQueue.empty()) {
        _sendingControl = false;
      } else {
        break;  // nothing queued
      }

      AsyncWebSocketMessage &target = _sendingControl ? _controlQueue.front() : _messageQueue.front();
      if (_sendingControl) {
        _framePayloadLen = target.size();
      } else {
        const size_t window = webSocketSendFrameWindow(_client);
        if (!window) {
          break;  // no space to send right now
        }
        const size_t remaining = target.size() - _sent;
        _framePayloadLen = std::min(remaining, window);
      }
      if (target.mask()) {
        _maskKey[0] = rand() % 0xFF;  // NOLINT(runtime/threadsafe_fn)
        _maskKey[1] = rand() % 0xFF;  // NOLINT(runtime/threadsafe_fn)
        _maskKey[2] = rand() % 0xFF;  // NOLINT(runtime/threadsafe_fn)
        _maskKey[3] = rand() % 0xFF;  // NOLINT(runtime/threadsafe_fn)
      }
    }

    AsyncWebSocketMessage &target = _sendingControl ? _controlQueue.front() : _messageQueue.front();
    const bool final = _sendingControl || (_sent + _framePayloadLen == target.size());
    const uint8_t frameOpcode = (_sendingControl || _sent == 0) ? target.opcode() : (uint8_t)WS_CONTINUATION;
    uint8_t *payload = target.data();

    // Avoid UB pointer arithmetic if payload is nullptr
    if (payload && !_sendingControl) {
      payload += _sent;
    }

    const size_t added = webSocketAddFrame(_client, final, frameOpcode, target.mask(), _maskKey, payload, _framePayloadLen, _frameSent);

    async_ws_log_v(
      "[%s][%" PRIu32 "][%" PRIu8 "] SEND ctrl:%d %u/%u added %u", _server->url(), _clientId, frameOpcode, _sendingControl, _frameSent,
      webSocketFrameHeaderLen(_framePayloadLen, target.mask()) + _framePayloadLen, added
    );

    _frameSent += added;
    needs_send |= (added > 0);

    const size_t frameLen = webSocketFrameHeaderLen(_framePayloadLen, target.mask()) + _framePayloadLen;
    if (_frameSent < frameLen) {
      break;  // stalled; resume on the next trigger
    }

    // frame fully committed to TCP -- done with it, regardless of ack
    _frameSent = 0;
    const size_t framePayloadLen = _framePayloadLen;
    _framePayloadLen = 0;

    if (_sendingControl) {
      const uint8_t opcode = target.opcode();
      _controlQueue.pop_front();
      if (opcode == WS_DISCONNECT && _status == WS_DISCONNECTING) {
        _status = WS_DISCONNECTED;
        async_ws_log_v("[%s][%" PRIu32 "] DISCONNECT SENT", _server->url(), _clientId);
        // Capture _client before unlocking: close() may synchronously run
        // _onDisconnect() -> AsyncWebSocket::_handleDisconnect() -> list::erase(),
        // destroying *this -- nothing after unlock() may touch any member.
        // close() itself (not abort()) hands the pcb to lwIP's graceful
        // close, which flushes/retransmits already-add()'ed data (this
        // close frame included) independent of our lifetime from here on.
        AsyncClient *c = _client;
        lock.unlock();
        c->send();  // Needed on some backends (e.g. SSL) to flush the close frame out before closing
        c->close();
        return;
      }
    } else {
      _sent += framePayloadLen;
      if (_sent >= target.size()) {
        _messageQueue.pop_front();
        _sent = 0;
      }
    }
  }

  if (needs_send) {
    _client->send();
  }
}

bool AsyncWebSocketClient::queueIsFull() const {
  asyncsrv::lock_guard_type lock(_queue_lock);
  return (_messageQueue.size() >= WS_MAX_QUEUED_MESSAGES) || (_status != WS_CONNECTED);
}

size_t AsyncWebSocketClient::queueLen() const {
  asyncsrv::lock_guard_type lock(_queue_lock);
  return _messageQueue.size();
}

bool AsyncWebSocketClient::canSend() const {
  asyncsrv::lock_guard_type lock(_queue_lock);
  return _messageQueue.size() < WS_MAX_QUEUED_MESSAGES;
}

bool AsyncWebSocketClient::_queueControl(uint8_t opcode, const uint8_t *data, size_t len, bool mask) {
  asyncsrv::unique_lock_type lock(_queue_lock);

  if (!_client) {
    return false;
  }

  if (!data) {
    len = 0;
  } else if (len > 125) {
    len = 125;
  }
  AsyncWebSocketSharedBuffer buffer = len ? makeSharedBuffer(data, len) : AsyncWebSocketSharedBuffer{};
  _controlQueue.emplace_back(buffer, opcode, len && mask);
  async_ws_log_v("[%s][%" PRIu32 "] QUEUE CTRL (%u) 

Web Proxy Viewer  |  New URL  |  Original Page