void WebSocket::onReadable(uv_poll_t *p, int status, int events)
{
SocketData *socketData = (SocketData *) p->data;
// this one is not needed, read will do this!
if (status < 0) {
WebSocket(p).close(true, 1006);
return;
}
char *src = socketData->server->recvBuffer;
memcpy(src, socketData->spill, socketData->spillLength);
uv_os_fd_t fd;
uv_fileno((uv_handle_t *) p, &fd);
ssize_t received;
if (socketData->ssl) {
received = SSL_read(socketData->ssl, src + socketData->spillLength, Server::LARGE_BUFFER_SIZE - socketData->spillLength);
} else {
received = recv(fd, src + socketData->spillLength, Server::LARGE_BUFFER_SIZE - socketData->spillLength, 0);
}
if (received == -1 || received == 0) {
// do we have a close frame in our buffer, and did we already set the state as CLOSING?
if (socketData->state == CLOSING && socketData->controlBuffer.length()) {
std::tuple<unsigned short, char *, size_t> closeFrame = Parser::parseCloseFrame(socketData->controlBuffer);
if (!std::get<0>(closeFrame)) {
std::get<0>(closeFrame) = 1006;
}
WebSocket(p).close(true, std::get<0>(closeFrame), std::get<1>(closeFrame), std::get<2>(closeFrame));
} else {
WebSocket(p).close(true, 1006);
}
return;
}
// do not parse any data once in closing state
if (socketData->state == CLOSING) {
return;
}
// cork sends into one large package
#ifdef __linux
int cork = 1;
setsockopt(fd, IPPROTO_TCP, TCP_CORK, &cork, sizeof(int));
#endif
Parser::consume(socketData->spillLength + received, src, socketData, p);
#ifdef __linux
cork = 0;
setsockopt(fd, IPPROTO_TCP, TCP_CORK, &cork, sizeof(int));
#endif
}
bool HttpServer::initWebSocket(HttpServerConnection& connection, HttpRequest& request, HttpResponse& response)
{
if (!wsEnabled)
return false;
auto sock = WebSocket(&connection);
if (!sock.initialize(request, response))
return false;
connection.setDisconnectionHandler(HttpServerConnectionDelegate(&HttpServer::onCloseWebSocket, this)); // auto remove on close
response.sendHeader(connection); // Will push header before user data
wsocks.add(sock);
if (wsConnect) wsConnect(sock);
return true;
}
// async Unix send (has a Message struct in the start if transferOwnership)
void WebSocket::write(char *data, size_t length, bool transferOwnership, void(*callback)(WebSocket webSocket, void *data, bool cancelled), void *callbackData)
{
uv_os_fd_t fd;
uv_fileno((uv_handle_t *) p, &fd);
ssize_t sent = 0;
SocketData *socketData = (SocketData *) p->data;
if (!socketData->messageQueue.empty()) {
goto queueIt;
}
if (socketData->ssl) {
sent = SSL_write(socketData->ssl, data, length);
} else {
sent = ::send(fd, data, length, MSG_NOSIGNAL);
}
if (sent == (int) length) {
// everything was sent in one go!
if (transferOwnership) {
delete [] (data - sizeof(SocketData::Queue::Message));
}
if (callback) {
callback(p, callbackData, false);
}
} else {
// not everything was sent
if (sent == -1) {
// check to see if any error occurred
if (socketData->ssl) {
int error = SSL_get_error(socketData->ssl, sent);
if (error == SSL_ERROR_WANT_READ || error == SSL_ERROR_WANT_WRITE) {
goto queueIt;
}
} else {
#ifdef _WIN32
if (WSAGetLastError() == WSAENOBUFS || WSAGetLastError() == WSAEWOULDBLOCK) {
#else
if (errno == EAGAIN || errno == EWOULDBLOCK) {
#endif
goto queueIt;
}
}
// error sending!
if (transferOwnership) {
delete [] (data - sizeof(SocketData::Queue::Message));
}
return;
} else {
queueIt:
sent = std::max<ssize_t>(sent, 0);
// queue the rest of the message!
SocketData::Queue::Message *messagePtr;
if (transferOwnership) {
messagePtr = (SocketData::Queue::Message *) (data - sizeof(SocketData::Queue::Message));
messagePtr->data = data + sent;
messagePtr->length = length - sent;
messagePtr->nextMessage = nullptr;
} else {
// we need to copy the buffer
messagePtr = (SocketData::Queue::Message *) new char[sizeof(SocketData::Queue::Message) + length - sent];
messagePtr->length = length - sent;
messagePtr->data = ((char *) messagePtr) + sizeof(SocketData::Queue::Message);
messagePtr->nextMessage = nullptr;
memcpy(messagePtr->data, data + sent, messagePtr->length);
}
messagePtr->callback = callback;
messagePtr->callbackData = callbackData;
((SocketData *) p->data)->messageQueue.push(messagePtr);
// only start this if we just broke the 0 queue size!
uv_poll_start(p, UV_WRITABLE | UV_READABLE, [](uv_poll_t *handle, int status, int events) {
// handle all poll errors with forced disconnection
if (status < 0) {
WebSocket(handle).close(true, 1006);
return;
}
// handle reads if available
if (events & UV_READABLE) {
onReadable(handle, status, events);
if (!(events & UV_WRITABLE)) {
return;
}
}
SocketData *socketData = (SocketData *) handle->data;
if (socketData->state == CLOSING) {
if (uv_is_closing((uv_handle_t *) handle)) {
return;
} else {
uv_poll_start(handle, UV_READABLE, onReadable);
}
}
uv_os_fd_t fd;
uv_fileno((uv_handle_t *) handle, &fd);
do {
SocketData::Queue::Message *messagePtr = socketData->messageQueue.front();
ssize_t sent;
if (socketData->ssl) {
sent = SSL_write(socketData->ssl, messagePtr->data, messagePtr->length);
} else {
sent = ::send(fd, messagePtr->data, messagePtr->length, MSG_NOSIGNAL);
}
if (sent == (int) messagePtr->length) {
if (messagePtr->callback) {
messagePtr->callback(handle, messagePtr->callbackData, false);
}
socketData->messageQueue.pop();
} else {
if (sent == -1) {
// check to see if any error occurred
if (socketData->ssl) {
int error = SSL_get_error(socketData->ssl, sent);
if (error == SSL_ERROR_WANT_READ || error == SSL_ERROR_WANT_WRITE) {
return;
}
} else {
#ifdef _WIN32
if (WSAGetLastError() == WSAENOBUFS || WSAGetLastError() == WSAEWOULDBLOCK) {
#else
if (errno == EAGAIN || errno == EWOULDBLOCK) {
#endif
return;
}
}
// error sending!
uv_poll_start(handle, UV_READABLE, onReadable);
return;
} else {
// update the Message
messagePtr->data += sent;
messagePtr->length -= sent;
return;
}
}
} while (!socketData->messageQueue.empty());
// only receive when we have fully sent everything
uv_poll_start(handle, UV_READABLE, onReadable);
});
}
}
void WebSocket::close(bool force, unsigned short code, char *data, size_t length)
{
uv_os_fd_t fd;
uv_fileno((uv_handle_t *) p, &fd);
SocketData *socketData = (SocketData *) p->data;
if (socketData->state != CLOSING) {
socketData->state = CLOSING;
if (socketData->prev == socketData->next) {
socketData->server->clients = nullptr;
} else {
if (socketData->prev) {
((SocketData *) socketData->prev->data)->next = socketData->next;
} else {
socketData->server->clients = socketData->next;
}
if (socketData->next) {
((SocketData *) socketData->next->data)->prev = socketData->prev;
}
}
// reuse prev as timer, mark no timer set
socketData->prev = nullptr;
// call disconnection callback on first close (graceful or force)
socketData->server->disconnectionCallback(p, code, data, length);
} else if (!force) {
std::cerr << "WARNING: Already gracefully closed: " << p << std::endl;
return;
}
if (force) {
// delete all messages in queue
while (!socketData->messageQueue.empty()) {
SocketData::Queue::Message *message = socketData->messageQueue.front();
if (message->callback) {
message->callback(nullptr, message->callbackData, true);
}
socketData->messageQueue.pop();
}
uv_poll_stop(p);
uv_close((uv_handle_t *) p, [](uv_handle_t *handle) {
delete (uv_poll_t *) handle;
});
::close(fd);
SSL_free(socketData->ssl);
socketData->controlBuffer.clear();
// cancel force close timer
if (socketData->prev) {
uv_timer_stop((uv_timer_t *) socketData->prev);
uv_close((uv_handle_t *) socketData->prev, [](uv_handle_t *handle) {
delete (uv_timer_t *) handle;
});
}
delete socketData->pmd;
delete socketData;
} else {
// force close after 15 seconds
socketData->prev = (uv_poll_t *) new uv_timer_t;
uv_timer_init(socketData->server->loop, (uv_timer_t *) socketData->prev);
((uv_timer_t *) socketData->prev)->data = p;
uv_timer_start((uv_timer_t *) socketData->prev, [](uv_timer_t *timer) {
WebSocket((uv_poll_t *) timer->data).close(true, 1006);
}, 15000, 0);
char *sendBuffer = socketData->server->sendBuffer;
if (code) {
length = std::min<size_t>(1024, length) + 2;
*((uint16_t *) &sendBuffer[length + 2]) = htons(code);
memcpy(&sendBuffer[length + 4], data, length - 2);
}
write(sendBuffer, formatMessage(sendBuffer, &sendBuffer[length + 2], length, CLOSE, length), false, [](WebSocket webSocket, void *data, bool cancelled) {
uv_os_fd_t fd;
uv_fileno((uv_handle_t *) webSocket.p, &fd);
SocketData *socketData = (SocketData *) webSocket.p->data;
if (socketData->ssl) {
SSL_shutdown(socketData->ssl);
}
shutdown(fd, SHUT_WR);
});
}
}
