214 lines
8.6 KiB
C++
214 lines
8.6 KiB
C++
/*
|
|
* Authored by Alex Hultman, 2018-2019.
|
|
* Intellectual property of third-party.
|
|
|
|
* Licensed under the Apache License, Version 2.0 (the "License");
|
|
* you may not use this file except in compliance with the License.
|
|
* You may obtain a copy of the License at
|
|
|
|
* http://www.apache.org/licenses/LICENSE-2.0
|
|
|
|
* Unless required by applicable law or agreed to in writing, software
|
|
* distributed under the License is distributed on an "AS IS" BASIS,
|
|
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
* See the License for the specific language governing permissions and
|
|
* limitations under the License.
|
|
*/
|
|
|
|
#ifndef UWS_WEBSOCKET_H
|
|
#define UWS_WEBSOCKET_H
|
|
|
|
#include "WebSocketData.h"
|
|
#include "WebSocketProtocol.h"
|
|
#include "AsyncSocket.h"
|
|
#include "WebSocketContextData.h"
|
|
|
|
#include <string_view>
|
|
|
|
namespace uWS {
|
|
|
|
template <bool SSL, bool isServer>
|
|
struct WebSocket : AsyncSocket<SSL> {
|
|
template <bool> friend struct TemplatedApp;
|
|
private:
|
|
typedef AsyncSocket<SSL> Super;
|
|
|
|
void *init(bool perMessageDeflate, bool slidingCompression, std::string &&backpressure) {
|
|
new (us_socket_ext(SSL, (us_socket_t *) this)) WebSocketData(perMessageDeflate, slidingCompression, std::move(backpressure));
|
|
return this;
|
|
}
|
|
public:
|
|
|
|
/* Returns pointer to the per socket user data */
|
|
void *getUserData() {
|
|
WebSocketData *webSocketData = (WebSocketData *) us_socket_ext(SSL, (us_socket_t *) this);
|
|
/* We just have it overallocated by sizeof type */
|
|
return (webSocketData + 1);
|
|
}
|
|
|
|
/* See AsyncSocket */
|
|
using Super::getBufferedAmount;
|
|
using Super::getRemoteAddress;
|
|
|
|
/* Simple, immediate close of the socket. Emits close event */
|
|
using Super::close;
|
|
|
|
/* Send or buffer a WebSocket frame, compressed or not. Returns false on increased user space backpressure. */
|
|
bool send(std::string_view message, uWS::OpCode opCode = uWS::OpCode::BINARY, bool compress = false) {
|
|
/* Transform the message to compressed domain if requested */
|
|
if (compress) {
|
|
WebSocketData *webSocketData = (WebSocketData *) Super::getAsyncSocketData();
|
|
|
|
/* Check and correct the compress hint */
|
|
if (opCode < 3 && webSocketData->compressionStatus == WebSocketData::ENABLED) {
|
|
LoopData *loopData = Super::getLoopData();
|
|
/* Compress using either shared or dedicated deflationStream */
|
|
if (webSocketData->deflationStream) {
|
|
message = webSocketData->deflationStream->deflate(loopData->zlibContext, message, false);
|
|
} else {
|
|
message = loopData->deflationStream->deflate(loopData->zlibContext, message, true);
|
|
}
|
|
} else {
|
|
compress = false;
|
|
}
|
|
}
|
|
|
|
/* Check to see if we can cork for the user */
|
|
bool automaticallyCorked = false;
|
|
if (!Super::isCorked() && Super::canCork()) {
|
|
automaticallyCorked = true;
|
|
Super::cork();
|
|
}
|
|
|
|
/* Get size, alloate size, write if needed */
|
|
size_t messageFrameSize = protocol::messageFrameSize(message.length());
|
|
auto[sendBuffer, requiresWrite] = Super::getSendBuffer(messageFrameSize);
|
|
protocol::formatMessage<isServer>(sendBuffer, message.data(), message.length(), opCode, message.length(), compress);
|
|
/* This is the slow path, when we couldn't cork for the user */
|
|
if (requiresWrite) {
|
|
auto[written, failed] = Super::write(sendBuffer, (int) messageFrameSize);
|
|
|
|
/* For now, we are slow here */
|
|
free(sendBuffer);
|
|
|
|
if (failed) {
|
|
/* Return false for failure, skipping to reset the timeout below */
|
|
return false;
|
|
}
|
|
}
|
|
|
|
/* Uncork here if we automatically corked for the user */
|
|
if (automaticallyCorked) {
|
|
auto [written, failed] = Super::uncork();
|
|
if (failed) {
|
|
return false;
|
|
}
|
|
}
|
|
|
|
/* Every successful send resets the timeout */
|
|
WebSocketContextData<SSL> *webSocketContextData = (WebSocketContextData<SSL> *) us_socket_context_ext(SSL,
|
|
(us_socket_context_t *) us_socket_context(SSL, (us_socket_t *) this)
|
|
);
|
|
AsyncSocket<SSL>::timeout(webSocketContextData->idleTimeout);
|
|
|
|
/* Return success */
|
|
return true;
|
|
}
|
|
|
|
/* Send websocket close frame, emit close event, send FIN if successful */
|
|
void end(int code, std::string_view message = {}) {
|
|
/* Check if we already called this one */
|
|
WebSocketData *webSocketData = (WebSocketData *) us_socket_ext(SSL, (us_socket_t *) this);
|
|
if (webSocketData->isShuttingDown) {
|
|
return;
|
|
}
|
|
|
|
/* We postpone any FIN sending to either drainage or uncorking */
|
|
webSocketData->isShuttingDown = true;
|
|
|
|
/* Format and send the close frame */
|
|
static const int MAX_CLOSE_PAYLOAD = 123;
|
|
int length = (int) std::min<size_t>(MAX_CLOSE_PAYLOAD, message.length());
|
|
char closePayload[MAX_CLOSE_PAYLOAD + 2];
|
|
int closePayloadLength = (int) protocol::formatClosePayload(closePayload, (uint16_t) code, message.data(), length);
|
|
bool ok = send(std::string_view(closePayload, closePayloadLength), OpCode::CLOSE);
|
|
|
|
/* FIN if we are ok and not corked */
|
|
WebSocket<SSL, true> *webSocket = (WebSocket<SSL, true> *) this;
|
|
if (!webSocket->isCorked()) {
|
|
if (ok) {
|
|
/* If we are not corked, and we just sent off everything, we need to FIN right here.
|
|
* In all other cases, we need to fin either if uncork was successful, or when drainage is complete. */
|
|
webSocket->shutdown();
|
|
}
|
|
}
|
|
|
|
/* Emit close event */
|
|
WebSocketContextData<SSL> *webSocketContextData = (WebSocketContextData<SSL> *) us_socket_context_ext(SSL,
|
|
(us_socket_context_t *) us_socket_context(SSL, (us_socket_t *) this)
|
|
);
|
|
if (webSocketContextData->closeHandler) {
|
|
webSocketContextData->closeHandler(this, code, message);
|
|
}
|
|
|
|
/* Make sure to unsubscribe from any pub/sub node at exit */
|
|
webSocketContextData->topicTree.unsubscribeAll(webSocketData->subscriber);
|
|
delete webSocketData->subscriber;
|
|
webSocketData->subscriber = nullptr;
|
|
}
|
|
|
|
/* Corks the response if possible. Leaves already corked socket be. */
|
|
void cork(fu2::unique_function<void()> &&handler) {
|
|
if (!Super::isCorked() && Super::canCork()) {
|
|
Super::cork();
|
|
handler();
|
|
|
|
/* There is no timeout when failing to uncork for WebSockets,
|
|
* as that is handled by idleTimeout */
|
|
auto [written, failed] = Super::uncork();
|
|
} else {
|
|
/* We are already corked, or can't cork so let's just call the handler */
|
|
handler();
|
|
}
|
|
}
|
|
|
|
/* Subscribe to a topic according to MQTT rules and syntax */
|
|
void subscribe(std::string_view topic) {
|
|
WebSocketContextData<SSL> *webSocketContextData = (WebSocketContextData<SSL> *) us_socket_context_ext(SSL,
|
|
(us_socket_context_t *) us_socket_context(SSL, (us_socket_t *) this)
|
|
);
|
|
|
|
/* Make us a subscriber if we aren't yet */
|
|
WebSocketData *webSocketData = (WebSocketData *) us_socket_ext(SSL, (us_socket_t *) this);
|
|
if (!webSocketData->subscriber) {
|
|
webSocketData->subscriber = new Subscriber(this);
|
|
}
|
|
|
|
webSocketContextData->topicTree.subscribe(topic, webSocketData->subscriber);
|
|
}
|
|
|
|
/* Unsubscribe from a topic, returns true if we were subscribed */
|
|
bool unsubscribe(std::string_view topic) {
|
|
WebSocketContextData<SSL> *webSocketContextData = (WebSocketContextData<SSL> *) us_socket_context_ext(SSL,
|
|
(us_socket_context_t *) us_socket_context(SSL, (us_socket_t *) this)
|
|
);
|
|
|
|
WebSocketData *webSocketData = (WebSocketData *) us_socket_ext(SSL, (us_socket_t *) this);
|
|
|
|
return webSocketContextData->topicTree.unsubscribe(topic, webSocketData->subscriber);
|
|
}
|
|
|
|
/* Publish a message to a topic according to MQTT rules and syntax */
|
|
void publish(std::string_view topic, std::string_view message, OpCode opCode = OpCode::TEXT, bool compress = false) {
|
|
WebSocketContextData<SSL> *webSocketContextData = (WebSocketContextData<SSL> *) us_socket_context_ext(SSL,
|
|
(us_socket_context_t *) us_socket_context(SSL, (us_socket_t *) this)
|
|
);
|
|
/* Is the same as publishing per websocket context */
|
|
webSocketContextData->publish(topic, message, opCode, compress);
|
|
}
|
|
};
|
|
|
|
}
|
|
|
|
#endif // UWS_WEBSOCKET_H
|