复现已有算法
This commit is contained in:
cloud
@@ -0,0 +1,497 @@
|
||||
/**
|
||||
*
|
||||
* @file WebSocketConnectionImpl.cc
|
||||
* @author An Tao
|
||||
*
|
||||
* Copyright 2018, An Tao. All rights reserved.
|
||||
* https://github.com/an-tao/drogon
|
||||
* Use of this source code is governed by a MIT license
|
||||
* that can be found in the License file.
|
||||
*
|
||||
* Drogon
|
||||
*
|
||||
*/
|
||||
|
||||
#include "WebSocketConnectionImpl.h"
|
||||
#include "HttpAppFrameworkImpl.h"
|
||||
#include <json/value.h>
|
||||
#include <json/writer.h>
|
||||
#include <thread>
|
||||
#include <limits>
|
||||
|
||||
using namespace drogon;
|
||||
|
||||
WebSocketConnectionImpl::WebSocketConnectionImpl(
|
||||
const trantor::TcpConnectionPtr &conn,
|
||||
bool isServer)
|
||||
: tcpConnectionPtr_(conn),
|
||||
localAddr_(conn->localAddr()),
|
||||
peerAddr_(conn->peerAddr()),
|
||||
isServer_(isServer),
|
||||
usingMask_(false)
|
||||
{
|
||||
}
|
||||
|
||||
WebSocketConnectionImpl::~WebSocketConnectionImpl()
|
||||
{
|
||||
shutdown();
|
||||
}
|
||||
|
||||
void WebSocketConnectionImpl::send(const char *msg,
|
||||
uint64_t len,
|
||||
const WebSocketMessageType type)
|
||||
{
|
||||
unsigned char opcode;
|
||||
if (type == WebSocketMessageType::Text)
|
||||
opcode = 1;
|
||||
else if (type == WebSocketMessageType::Binary)
|
||||
opcode = 2;
|
||||
else if (type == WebSocketMessageType::Close)
|
||||
{
|
||||
assert(len <= 125);
|
||||
opcode = 8;
|
||||
}
|
||||
else if (type == WebSocketMessageType::Ping)
|
||||
{
|
||||
assert(len <= 125);
|
||||
opcode = 9;
|
||||
}
|
||||
else if (type == WebSocketMessageType::Pong)
|
||||
{
|
||||
assert(len <= 125);
|
||||
opcode = 10;
|
||||
}
|
||||
else
|
||||
{
|
||||
opcode = 0;
|
||||
assert(0);
|
||||
}
|
||||
sendWsData(msg, len, opcode);
|
||||
}
|
||||
|
||||
void WebSocketConnectionImpl::sendWsData(const char *msg,
|
||||
uint64_t len,
|
||||
unsigned char opcode)
|
||||
{
|
||||
LOG_TRACE << "send " << len << " bytes";
|
||||
|
||||
// Format the frame
|
||||
std::string bytesFormatted;
|
||||
bytesFormatted.resize(len + 10);
|
||||
bytesFormatted[0] = char(0x80 | (opcode & 0x0f));
|
||||
|
||||
int indexStartRawData = -1;
|
||||
|
||||
if (len <= 125)
|
||||
{
|
||||
bytesFormatted[1] = static_cast<char>(len);
|
||||
indexStartRawData = 2;
|
||||
}
|
||||
else if (len <= 65535)
|
||||
{
|
||||
bytesFormatted[1] = 126;
|
||||
bytesFormatted[2] = ((len >> 8) & 255);
|
||||
bytesFormatted[3] = ((len) & 255);
|
||||
LOG_TRACE << "bytes[2]=" << (size_t)bytesFormatted[2];
|
||||
LOG_TRACE << "bytes[3]=" << (size_t)bytesFormatted[3];
|
||||
indexStartRawData = 4;
|
||||
}
|
||||
else
|
||||
{
|
||||
bytesFormatted[1] = 127;
|
||||
bytesFormatted[2] = ((len >> 56) & 255);
|
||||
bytesFormatted[3] = ((len >> 48) & 255);
|
||||
bytesFormatted[4] = ((len >> 40) & 255);
|
||||
bytesFormatted[5] = ((len >> 32) & 255);
|
||||
bytesFormatted[6] = ((len >> 24) & 255);
|
||||
bytesFormatted[7] = ((len >> 16) & 255);
|
||||
bytesFormatted[8] = ((len >> 8) & 255);
|
||||
bytesFormatted[9] = ((len) & 255);
|
||||
|
||||
indexStartRawData = 10;
|
||||
}
|
||||
if (!isServer_)
|
||||
{
|
||||
int random;
|
||||
// Use the cached randomness if no one else is also using it. Otherwise
|
||||
// generate one from scratch.
|
||||
if (!usingMask_.exchange(true, std::memory_order_acq_rel))
|
||||
{
|
||||
if (masks_.empty())
|
||||
{
|
||||
masks_.resize(16);
|
||||
bool status =
|
||||
utils::secureRandomBytes(masks_.data(),
|
||||
masks_.size() * sizeof(uint32_t));
|
||||
if (status == false)
|
||||
{
|
||||
LOG_ERROR << "Failed to generate random numbers for "
|
||||
"WebSocket mask";
|
||||
abort();
|
||||
}
|
||||
}
|
||||
random = masks_.back();
|
||||
masks_.pop_back();
|
||||
usingMask_.store(false, std::memory_order_release);
|
||||
}
|
||||
else
|
||||
{
|
||||
bool status = utils::secureRandomBytes(&random, sizeof(random));
|
||||
if (status == false)
|
||||
{
|
||||
LOG_ERROR
|
||||
<< "Failed to generate random numbers for WebSocket mask";
|
||||
abort();
|
||||
}
|
||||
}
|
||||
|
||||
bytesFormatted[1] = (bytesFormatted[1] | 0x80);
|
||||
bytesFormatted.resize(indexStartRawData + 4 + len);
|
||||
memcpy(&bytesFormatted[indexStartRawData], &random, sizeof(random));
|
||||
for (size_t i = 0; i < len; ++i)
|
||||
{
|
||||
bytesFormatted[indexStartRawData + 4 + i] =
|
||||
(msg[i] ^ bytesFormatted[indexStartRawData + (i % 4)]);
|
||||
}
|
||||
}
|
||||
else
|
||||
{
|
||||
bytesFormatted.resize(indexStartRawData);
|
||||
bytesFormatted.append(msg, len);
|
||||
}
|
||||
tcpConnectionPtr_->send(std::move(bytesFormatted));
|
||||
}
|
||||
|
||||
void WebSocketConnectionImpl::send(const std::string_view msg,
|
||||
const WebSocketMessageType type)
|
||||
{
|
||||
send(msg.data(), msg.length(), type);
|
||||
}
|
||||
|
||||
void WebSocketConnectionImpl::sendJson(const Json::Value &json,
|
||||
const WebSocketMessageType type)
|
||||
{
|
||||
static std::once_flag once;
|
||||
static Json::StreamWriterBuilder builder;
|
||||
std::call_once(once, []() {
|
||||
builder["commentStyle"] = "None";
|
||||
builder["indentation"] = "";
|
||||
if (!app().isUnicodeEscapingUsedInJson())
|
||||
{
|
||||
builder["emitUTF8"] = true;
|
||||
}
|
||||
auto &precision = app().getFloatPrecisionInJson();
|
||||
if (precision.first != 0)
|
||||
{
|
||||
builder["precision"] = precision.first;
|
||||
builder["precisionType"] = precision.second;
|
||||
}
|
||||
});
|
||||
auto msg = writeString(builder, json);
|
||||
send(msg.data(), msg.length(), type);
|
||||
}
|
||||
|
||||
const trantor::InetAddress &WebSocketConnectionImpl::localAddr() const
|
||||
{
|
||||
return localAddr_;
|
||||
}
|
||||
|
||||
const trantor::InetAddress &WebSocketConnectionImpl::peerAddr() const
|
||||
{
|
||||
return peerAddr_;
|
||||
}
|
||||
|
||||
bool WebSocketConnectionImpl::connected() const
|
||||
{
|
||||
return tcpConnectionPtr_->connected();
|
||||
}
|
||||
|
||||
bool WebSocketConnectionImpl::disconnected() const
|
||||
{
|
||||
return tcpConnectionPtr_->disconnected();
|
||||
}
|
||||
|
||||
void WebSocketConnectionImpl::WebSocketConnectionImpl::shutdown(
|
||||
const CloseCode code,
|
||||
const std::string &reason)
|
||||
{
|
||||
tcpConnectionPtr_->getLoop()->invalidateTimer(pingTimerId_);
|
||||
if (!tcpConnectionPtr_->connected())
|
||||
return;
|
||||
std::string message;
|
||||
message.resize(reason.length() + 2);
|
||||
auto c = htons(static_cast<unsigned short>(code));
|
||||
memcpy(&message[0], &c, 2);
|
||||
if (!reason.empty())
|
||||
memcpy(&message[2], reason.data(), reason.length());
|
||||
send(message, WebSocketMessageType::Close);
|
||||
tcpConnectionPtr_->shutdown();
|
||||
}
|
||||
|
||||
void WebSocketConnectionImpl::WebSocketConnectionImpl::forceClose()
|
||||
{
|
||||
tcpConnectionPtr_->forceClose();
|
||||
}
|
||||
|
||||
void WebSocketConnectionImpl::setPingMessage(
|
||||
const std::string &message,
|
||||
const std::chrono::duration<double> &interval)
|
||||
{
|
||||
auto loop = tcpConnectionPtr_->getLoop();
|
||||
if (loop->isInLoopThread())
|
||||
{
|
||||
setPingMessageInLoop(std::string{message}, interval);
|
||||
}
|
||||
else
|
||||
{
|
||||
loop->queueInLoop(
|
||||
[msg = message, interval, thisPtr = shared_from_this()]() mutable {
|
||||
thisPtr->setPingMessageInLoop(std::move(msg), interval);
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
void WebSocketConnectionImpl::disablePing()
|
||||
{
|
||||
auto loop = tcpConnectionPtr_->getLoop();
|
||||
if (loop->isInLoopThread())
|
||||
{
|
||||
disablePingInLoop();
|
||||
}
|
||||
else
|
||||
{
|
||||
loop->queueInLoop(
|
||||
[thisPtr = shared_from_this()]() { thisPtr->disablePingInLoop(); });
|
||||
}
|
||||
}
|
||||
|
||||
bool WebSocketMessageParser::parse(trantor::MsgBuffer *buffer)
|
||||
{
|
||||
// According to the rfc6455
|
||||
gotAll_ = false;
|
||||
while (buffer->readableBytes() >= 2)
|
||||
{
|
||||
unsigned char opcode = (*buffer)[0] & 0x0f;
|
||||
bool isControlFrame = false;
|
||||
switch (opcode)
|
||||
{
|
||||
case 0:
|
||||
LOG_TRACE << "continuation frame";
|
||||
break;
|
||||
case 1:
|
||||
type_ = WebSocketMessageType::Text;
|
||||
break;
|
||||
case 2:
|
||||
type_ = WebSocketMessageType::Binary;
|
||||
break;
|
||||
case 8:
|
||||
type_ = WebSocketMessageType::Close;
|
||||
isControlFrame = true;
|
||||
break;
|
||||
case 9:
|
||||
type_ = WebSocketMessageType::Ping;
|
||||
isControlFrame = true;
|
||||
break;
|
||||
case 10:
|
||||
type_ = WebSocketMessageType::Pong;
|
||||
isControlFrame = true;
|
||||
break;
|
||||
default:
|
||||
LOG_ERROR << "Unknown frame type";
|
||||
return false;
|
||||
break;
|
||||
}
|
||||
|
||||
bool isFin = (((*buffer)[0] & 0x80) == 0x80);
|
||||
if (!isFin && isControlFrame)
|
||||
{
|
||||
// rfc6455-5.5
|
||||
LOG_ERROR << "Bad frame: all control frames MUST NOT be fragmented";
|
||||
return false;
|
||||
}
|
||||
auto secondByte = (*buffer)[1];
|
||||
size_t length = secondByte & 127;
|
||||
int isMasked = (secondByte & 0x80);
|
||||
if (isMasked != 0)
|
||||
{
|
||||
LOG_TRACE << "data encoded!";
|
||||
}
|
||||
else
|
||||
LOG_TRACE << "plain data";
|
||||
size_t indexFirstMask = 2;
|
||||
|
||||
if (length == 126)
|
||||
{
|
||||
indexFirstMask = 4;
|
||||
}
|
||||
else if (length == 127)
|
||||
{
|
||||
indexFirstMask = 10;
|
||||
}
|
||||
if (indexFirstMask > 2)
|
||||
{
|
||||
if (buffer->readableBytes() < indexFirstMask)
|
||||
{
|
||||
// Not enough data yet, wait for more.
|
||||
return true;
|
||||
}
|
||||
if (isControlFrame)
|
||||
{
|
||||
// rfc6455-5.5
|
||||
LOG_ERROR << "Bad frame: all control frames MUST have a "
|
||||
"payload length "
|
||||
"of 125 bytes or less";
|
||||
return false;
|
||||
}
|
||||
if (indexFirstMask == 4)
|
||||
{
|
||||
length = (unsigned char)(*buffer)[2];
|
||||
length = (length << 8) + (unsigned char)(*buffer)[3];
|
||||
}
|
||||
else if (indexFirstMask == 10)
|
||||
{
|
||||
length = 0;
|
||||
for (int i = 2; i <= 9; ++i)
|
||||
{
|
||||
if (length > ((std::numeric_limits<size_t>::max)() >> 8))
|
||||
{
|
||||
LOG_ERROR
|
||||
<< "Payload length too large to handle safely";
|
||||
return false;
|
||||
}
|
||||
length = (length << 8) + (unsigned char)(*buffer)[i];
|
||||
}
|
||||
}
|
||||
else
|
||||
{
|
||||
LOG_ERROR << "Websock parsing failed!";
|
||||
return false;
|
||||
}
|
||||
}
|
||||
if (isMasked != 0)
|
||||
{
|
||||
// The message is sent by the client, check the length
|
||||
if (length > HttpAppFrameworkImpl::instance()
|
||||
.getClientMaxWebSocketMessageSize())
|
||||
{
|
||||
LOG_ERROR << "The size of the WebSocket message is too large!";
|
||||
buffer->retrieveAll();
|
||||
return false;
|
||||
}
|
||||
if (buffer->readableBytes() >= (indexFirstMask + 4 + length))
|
||||
{
|
||||
auto masks = buffer->peek() + indexFirstMask;
|
||||
auto indexFirstDataByte = indexFirstMask + 4;
|
||||
auto rawData = buffer->peek() + indexFirstDataByte;
|
||||
auto oldLen = message_.length();
|
||||
message_.resize(oldLen + length);
|
||||
for (size_t i = 0; i < length; ++i)
|
||||
{
|
||||
message_[oldLen + i] = (rawData[i] ^ masks[i % 4]);
|
||||
}
|
||||
buffer->retrieve(indexFirstMask + 4 + length);
|
||||
if (isFin)
|
||||
{
|
||||
gotAll_ = true;
|
||||
return true;
|
||||
}
|
||||
}
|
||||
else
|
||||
{
|
||||
// Not enough data yet, wait for more.
|
||||
return true;
|
||||
}
|
||||
}
|
||||
else
|
||||
{
|
||||
if (buffer->readableBytes() >= (indexFirstMask + length))
|
||||
{
|
||||
auto rawData = buffer->peek() + indexFirstMask;
|
||||
message_.append(rawData, length);
|
||||
buffer->retrieve(indexFirstMask + length);
|
||||
if (isFin)
|
||||
{
|
||||
gotAll_ = true;
|
||||
return true;
|
||||
}
|
||||
}
|
||||
else
|
||||
{
|
||||
// Not enough data yet, wait for more.
|
||||
return true;
|
||||
}
|
||||
}
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
void WebSocketConnectionImpl::onNewMessage(
|
||||
const trantor::TcpConnectionPtr &connPtr,
|
||||
trantor::MsgBuffer *buffer)
|
||||
{
|
||||
auto self = shared_from_this();
|
||||
while (buffer->readableBytes() > 0)
|
||||
{
|
||||
auto success = parser_.parse(buffer);
|
||||
if (success)
|
||||
{
|
||||
std::string message;
|
||||
WebSocketMessageType type;
|
||||
if (parser_.gotAll(message, type))
|
||||
{
|
||||
if (type == WebSocketMessageType::Ping)
|
||||
{
|
||||
// ping
|
||||
send(message, WebSocketMessageType::Pong);
|
||||
}
|
||||
else if (type == WebSocketMessageType::Close)
|
||||
{
|
||||
// close
|
||||
connPtr->shutdown();
|
||||
}
|
||||
else if (type == WebSocketMessageType::Unknown)
|
||||
{
|
||||
return;
|
||||
}
|
||||
// LOG_TRACE << "new message received: " << message
|
||||
// << "\n(type=" << (int)type << ")";
|
||||
messageCallback_(std::move(message), self, type);
|
||||
}
|
||||
else
|
||||
{
|
||||
return;
|
||||
}
|
||||
}
|
||||
else
|
||||
{
|
||||
// Websock error!
|
||||
connPtr->shutdown();
|
||||
return;
|
||||
}
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
void WebSocketConnectionImpl::disablePingInLoop()
|
||||
{
|
||||
if (pingTimerId_ != trantor::InvalidTimerId)
|
||||
{
|
||||
tcpConnectionPtr_->getLoop()->invalidateTimer(pingTimerId_);
|
||||
}
|
||||
}
|
||||
|
||||
void WebSocketConnectionImpl::setPingMessageInLoop(
|
||||
std::string &&message,
|
||||
const std::chrono::duration<double> &interval)
|
||||
{
|
||||
std::weak_ptr<WebSocketConnectionImpl> weakPtr = shared_from_this();
|
||||
disablePingInLoop();
|
||||
pingTimerId_ = tcpConnectionPtr_->getLoop()->runEvery(
|
||||
interval.count(), [weakPtr, message = std::move(message)]() {
|
||||
auto thisPtr = weakPtr.lock();
|
||||
if (thisPtr)
|
||||
{
|
||||
thisPtr->send(message, WebSocketMessageType::Ping);
|
||||
}
|
||||
});
|
||||
}
|
||||
Reference in New Issue
Block a user