Fix websocket client lifetime races during broadcast

Author: MitchBradley

src/AsyncWebSocket.cpp

@@ -37,6 +37,15 @@
 
 using namespace asyncsrv;
 
+namespace {
+AsyncWebSocketClient *find_connected_client_locked(std::list<AsyncWebSocketClient> &clients, uint32_t id) {
+  const auto iter = std::find_if(clients.begin(), clients.end(), [id](const AsyncWebSocketClient &client) {
+    return client.id() == id && client.status() == WS_CONNECTED;
+  });
+  return iter == clients.end() ? nullptr : &(*iter);
+}
+}  // namespace
+
 size_t webSocketSendFrameWindow(AsyncClient *client) {
   if (!client || !client->canSend()) {
     return 0;
@@ -357,11 +366,11 @@ void AsyncWebSocketClient::_onAck(size_t len, uint32_t time) {
 }
 
 void AsyncWebSocketClient::_onPoll() {
+  asyncsrv::unique_lock_type lock(_lock);
   if (!_client) {
     return;
   }
 
-  asyncsrv::unique_lock_type lock(_lock);
   if (_client && _client->canSend() && (!_controlQueue.empty() || !_messageQueue.empty())) {
     _runQueue();
   } else if (_keepAlivePeriod > 0 && (millis() - _lastMessageTime) >= _keepAlivePeriod && (_controlQueue.empty() && _messageQueue.empty())) {
@@ -444,12 +453,11 @@ bool AsyncWebSocketClient::canSend() const {
 }
 
 bool AsyncWebSocketClient::_queueControl(uint8_t opcode, const uint8_t *data, size_t len, bool mask) {
+  asyncsrv::lock_guard_type lock(_lock);
   if (!_client) {
     return false;
   }
 
-  asyncsrv::lock_guard_type lock(_lock);
-
   _controlQueue.emplace_back(opcode, data, len, mask);
   async_ws_log_v("[%s][%" PRIu32 "] QUEUE CTRL (%u) << %" PRIu8, _server->url(), _clientId, _controlQueue.size(), opcode);
 
@@ -461,12 +469,11 @@ bool AsyncWebSocketClient::_queueControl(uint8_t opcode, const uint8_t *data, si
 }
 
 bool AsyncWebSocketClient::_queueMessage(AsyncWebSocketSharedBuffer buffer, uint8_t opcode, bool mask) {
-  if (!_client || buffer->size() == 0 || _status != WS_CONNECTED) {
+  asyncsrv::unique_lock_type lock(_lock);
+  if (!_client || !buffer || buffer->empty() || _status != WS_CONNECTED) {
     return false;
   }
 
-  asyncsrv::unique_lock_type lock(_lock);
-
   if (_messageQueue.size() >= WS_MAX_QUEUED_MESSAGES) {
     if (closeWhenFull) {
       _status = WS_DISCONNECTED;
@@ -545,6 +552,7 @@ void AsyncWebSocketClient::_onError(int8_t err) {
 }
 
 void AsyncWebSocketClient::_onTimeout(uint32_t time) {
+  asyncsrv::lock_guard_type lock(_lock);
   if (!_client) {
     return;
   }
@@ -553,7 +561,9 @@ void AsyncWebSocketClient::_onTimeout(uint32_t time) {
 }
 
 void AsyncWebSocketClient::_onDisconnect() {
+  asyncsrv::lock_guard_type lock(_lock);
   async_ws_log_v("[%s][%" PRIu32 "] DISCONNECT", _server->url(), _clientId);
+  _status = WS_DISCONNECTED;
   _client = nullptr;
   _server->_handleDisconnect(this);
 }
@@ -947,6 +957,7 @@ bool AsyncWebSocketClient::binary(const __FlashStringHelper *data, size_t len) {
 #endif
 
 IPAddress AsyncWebSocketClient::remoteIP() const {
+  asyncsrv::lock_guard_type lock(_lock);
   if (!_client) {
     return IPAddress((uint32_t)0U);
   }
@@ -955,6 +966,7 @@ IPAddress AsyncWebSocketClient::remoteIP() const {
 }
 
 uint16_t AsyncWebSocketClient::remotePort() const {
+  asyncsrv::lock_guard_type lock(_lock);
   if (!_client) {
     return 0;
   }
@@ -983,14 +995,10 @@ AsyncWebSocketClient *AsyncWebSocket::_newClient(AsyncWebServerRequest *request)
 }
 
 void AsyncWebSocket::_handleDisconnect(AsyncWebSocketClient *client) {
-  asyncsrv::lock_guard_type lock(_lock);
-  const auto client_id = client->id();
-  const auto iter = std::find_if(std::begin(_clients), std::end(_clients), [client_id](const AsyncWebSocketClient &c) {
-    return c.id() == client_id;
-  });
-  if (iter != std::end(_clients)) {
-    _clients.erase(iter);
-  }
+  (void)client;
+  // Defer removal to cleanupClients(). Disconnect callbacks can fire while
+  // iterating _clients for broadcast sends, and erasing here invalidates the
+  // active iterator in the caller.
 }
 
 bool AsyncWebSocket::availableForWriteAll() {
@@ -1031,7 +1039,8 @@ AsyncWebSocketClient *AsyncWebSocket::client(uint32_t id) {
 }
 
 void AsyncWebSocket::close(uint32_t id, uint16_t code, const char *message) {
-  if (AsyncWebSocketClient *c = client(id)) {
+  asyncsrv::lock_guard_type lock(_lock);
+  if (AsyncWebSocketClient *c = find_connected_client_locked(_clients, id)) {
     c->close(code, message);
   }
 }
@@ -1047,22 +1056,32 @@ void AsyncWebSocket::closeAll(uint16_t code, const char *message) {
 
 void AsyncWebSocket::cleanupClients(uint16_t maxClients) {
   asyncsrv::lock_guard_type lock(_lock);
-  const size_t c = count();
-  if (c > maxClients) {
-    async_ws_log_v("[%s] CLEANUP %" PRIu32 " (%u/%" PRIu16 ")", _url.c_str(), _clients.front().id(), c, maxClients);
-    _clients.front().close();
+  const size_t connected = std::count_if(std::begin(_clients), std::end(_clients), [](const AsyncWebSocketClient &c) {
+    return c.status() == WS_CONNECTED;
+  });
+
+  if (connected > maxClients) {
+    const auto connected_iter = std::find_if(std::begin(_clients), std::end(_clients), [](const AsyncWebSocketClient &c) {
+      return c.status() == WS_CONNECTED;
+    });
+    if (connected_iter != std::end(_clients)) {
+      async_ws_log_v("[%s] CLEANUP %" PRIu32 " (%u/%" PRIu16 ")", _url.c_str(), connected_iter->id(), connected, maxClients);
+      connected_iter->close();
+    }
   }
 
-  for (auto i = _clients.begin(); i != _clients.end(); ++i) {
-    if (i->shouldBeDeleted()) {
-      _clients.erase(i);
-      break;
+  for (auto iter = _clients.begin(); iter != _clients.end();) {
+    if (iter->shouldBeDeleted()) {
+      iter = _clients.erase(iter);
+    } else {
+      ++iter;
     }
   }
 }
 
 bool AsyncWebSocket::ping(uint32_t id, const uint8_t *data, size_t len) {
-  AsyncWebSocketClient *c = client(id);
+  asyncsrv::lock_guard_type lock(_lock);
+  AsyncWebSocketClient *c = find_connected_client_locked(_clients, id);
   return c && c->ping(data, len);
 }
 
@@ -1081,7 +1100,8 @@ AsyncWebSocket::SendStatus AsyncWebSocket::pingAll(const uint8_t *data, size_t l
 }
 
 bool AsyncWebSocket::text(uint32_t id, const uint8_t *message, size_t len) {
-  AsyncWebSocketClient *c = client(id);
+  asyncsrv::lock_guard_type lock(_lock);
+  AsyncWebSocketClient *c = find_connected_client_locked(_clients, id);
   return c && c->text(makeSharedBuffer(message, len));
 }
 bool AsyncWebSocket::text(uint32_t id, const char *message, size_t len) {
@@ -1127,7 +1147,8 @@ bool AsyncWebSocket::text(uint32_t id, AsyncWebSocketMessageBuffer *buffer) {
   return enqueued;
 }
 bool AsyncWebSocket::text(uint32_t id, AsyncWebSocketSharedBuffer buffer) {
-  AsyncWebSocketClient *c = client(id);
+  asyncsrv::lock_guard_type lock(_lock);
+  AsyncWebSocketClient *c = find_connected_client_locked(_clients, id);
   return c && c->text(buffer);
 }
 
@@ -1190,7 +1211,8 @@ AsyncWebSocket::SendStatus AsyncWebSocket::textAll(AsyncWebSocketSharedBuffer bu
 }
 
 bool AsyncWebSocket::binary(uint32_t id, const uint8_t *message, size_t len) {
-  AsyncWebSocketClient *c = client(id);
+  asyncsrv::lock_guard_type lock(_lock);
+  AsyncWebSocketClient *c = find_connected_client_locked(_clients, id);
   return c && c->binary(makeSharedBuffer(message, len));
 }
 bool AsyncWebSocket::binary(uint32_t id, const char *message, size_t len) {
@@ -1226,7 +1248,8 @@ bool AsyncWebSocket::binary(uint32_t id, AsyncWebSocketMessageBuffer *buffer) {
   return enqueued;
 }
 bool AsyncWebSocket::binary(uint32_t id, AsyncWebSocketSharedBuffer buffer) {
-  AsyncWebSocketClient *c = client(id);
+  asyncsrv::lock_guard_type lock(_lock);
+  AsyncWebSocketClient *c = find_connected_client_locked(_clients, id);
   return c && c->binary(buffer);
 }
 

src/AsyncWebSocket.h

@@ -303,6 +303,7 @@ class AsyncWebSocketClient {
   uint16_t remotePort() const;
 
   bool shouldBeDeleted() const {
+    asyncsrv::lock_guard_type lock(_lock);
     return !_client;
   }