amass/ZLMediaKit
1
0
Fork 0
mirror of https://github.com/ZLMediaKit/ZLMediaKit.git synced 2024-11-22 10:40:05 +08:00
Code Issues Actions 6 Packages Projects Releases Wiki Activity

修复WebSocket线程安全问题，同时新增内置客户端、服务端心跳机制。

Browse Source
This commit is contained in:
xia-chu 2023-03-11 11:08:14 +08:00
parent c2a8d46a64
commit 1bab0b8e31
2 changed files with 163 additions and 140 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

100
src/Http/WebSocketClient.h
View File

@ -54,9 +54,7 @@ protected:
* 设置发送数据截取回调函数
* @param cb 截取回调函数
*/
void setOnBeforeSendCB(const onBeforeSendCB &cb){
_beforeSendCB = cb;
}
void setOnBeforeSendCB(const onBeforeSendCB &cb) { _beforeSendCB = cb; }
private:
onBeforeSendCB _beforeSendCB;
@ -72,12 +70,11 @@ class HttpWsClient : public HttpClientImp , public WebSocketSplitter{
public:
using Ptr = std::shared_ptr<HttpWsClient>;
HttpWsClient(const std::shared_ptr<ClientTypeImp<ClientType, DataType> > &delegate) : _weak_delegate(delegate),
_delegate(*delegate) {
HttpWsClient(const std::shared_ptr<ClientTypeImp<ClientType, DataType>> &delegate) : _weak_delegate(delegate) {
_Sec_WebSocket_Key = encodeBase64(toolkit::makeRandStr(16, false));
setPoller(_delegate.getPoller());
setPoller(delegate->getPoller());
}
~HttpWsClient(){}
~HttpWsClient() = default;
/**
* 发起ws握手
@ -155,7 +152,6 @@ protected:
// TcpClient override
void onRecv(const toolkit::Buffer::Ptr &buf) override {
auto strong_ref = _weak_delegate.lock();;
HttpClientImp::onRecv(buf);
}
@ -163,24 +159,43 @@ protected:
* 定时触发
*/
void onManager() override {
auto strong_ref = _weak_delegate.lock();;
if (_onRecv) {
// websocket连接成功了
_delegate.onManager();
if (auto strong_ref = _weak_delegate.lock()) {
strong_ref->onManager();
}
} else {
// websocket连接中...
HttpClientImp::onManager();
}
if (!_onRecv) {
// websocket尚未链接
return;
}
if (_recv_ticker.elapsedTime() > 30 * 1000) {
shutdown(toolkit::SockException(toolkit::Err_timeout, "websocket timeout"));
} else if (_recv_ticker.elapsedTime() > 10 * 1000) {
// 没收到回复，每10秒发送次ping 包
WebSocketHeader header;
header._fin = true;
header._reserved = 0;
header._opcode = PING;
header._mask_flag = true;
WebSocketSplitter::encode(header, nullptr);
}
}
/**
* 数据全部发送完毕后回调
*/
void onFlush() override {
auto strong_ref = _weak_delegate.lock();;
if (_onRecv) {
// websocket连接成功了
_delegate.onFlush();
if (auto strong_ref = _weak_delegate.lock()) {
strong_ref->onFlush();
}
} else {
// websocket连接中...
HttpClientImp::onFlush();
@ -191,7 +206,6 @@ protected:
* tcp连接结果
*/
void onConnect(const toolkit::SockException &ex) override {
auto strong_ref = _weak_delegate.lock();;
if (ex) {
// tcp连接失败，直接返回失败
onWebSocketException(ex);
@ -205,7 +219,6 @@ protected:
* tcp连接断开
*/
void onErr(const toolkit::SockException &ex) override {
auto strong_ref = _weak_delegate.lock();;
// tcp断开或者shutdown导致的断开
onWebSocketException(ex);
}
@ -216,9 +229,7 @@ protected:
* 收到一个webSocket数据包包头，后续将继续触发onWebSocketDecodePayload回调
* @param header 数据包头
*/
void onWebSocketDecodeHeader(const WebSocketHeader &header) override{
_payload_section.clear();
}
void onWebSocketDecodeHeader(const WebSocketHeader &header) override { _payload_section.clear(); }
/**
* 收到webSocket数据包负载
@ -240,7 +251,7 @@ protected:
auto flag = header._mask_flag;
// websocket客户端发送数据需要加密
header._mask_flag = true;
_recv_ticker.resetTime();
switch (header._opcode) {
case WebSocketHeader::CLOSE: {
// 服务器主动关闭
@ -272,13 +283,17 @@ protected:
// 最后一个包
if (_payload_cache.empty()) {
// 这个包是唯一个分片
_delegate.onRecv(std::make_shared<WebSocketBuffer>(header._opcode, header._fin, std::move(_payload_section)));
if (auto strong_ref = _weak_delegate.lock()) {
strong_ref->onRecv(std::make_shared<WebSocketBuffer>(header._opcode, header._fin, std::move(_payload_section)));
}
break;
}
// 这个包由多个分片组成
_payload_cache.append(std::move(_payload_section));
_delegate.onRecv(std::make_shared<WebSocketBuffer>(header._opcode, header._fin, std::move(_payload_cache)));
if (auto strong_ref = _weak_delegate.lock()) {
strong_ref->onRecv(std::make_shared<WebSocketBuffer>(header._opcode, header._fin, std::move(_payload_cache)));
}
_payload_cache.clear();
break;
}
@ -294,9 +309,7 @@ protected:
* @param ptr 数据指针
* @param len 数据指针长度
*/
void onWebSocketEncodeData(toolkit::Buffer::Ptr buffer) override{
HttpClientImp::send(std::move(buffer));
}
void onWebSocketEncodeData(toolkit::Buffer::Ptr buffer) override { HttpClientImp::send(std::move(buffer)); }
private:
void onWebSocketException(const toolkit::SockException &ex) {
@ -304,24 +317,26 @@ private:
// websocket握手成功
// 此处截取TcpClient派生类发送的数据并进行websocket协议打包
std::weak_ptr<HttpWsClient> weakSelf = std::dynamic_pointer_cast<HttpWsClient>(shared_from_this());
_delegate.setOnBeforeSendCB([weakSelf](const toolkit::Buffer::Ptr &buf){
auto strongSelf = weakSelf.lock();
if(strongSelf){
if (auto strong_ref = _weak_delegate.lock()) {
strong_ref->setOnBeforeSendCB([weakSelf](const toolkit::Buffer::Ptr &buf) {
auto strong_self = weakSelf.lock();
if (strong_self) {
WebSocketHeader header;
header._fin = true;
header._reserved = 0;
header._opcode = DataType;
// 客户端需要加密
header._mask_flag = true;
strongSelf->WebSocketSplitter::encode(header,buf);
strong_self->WebSocketSplitter::encode(header, buf);
}
return buf->size();
});
// 设置sock，否则shutdown等接口都无效
_delegate.setSock(HttpClientImp::getSock());
strong_ref->setSock(HttpClientImp::getSock());
// 触发连接成功事件
_delegate.onConnect(ex);
strong_ref->onConnect(ex);
}
// 拦截websocket数据接收
_onRecv = [this](const char *data, size_t len) {
// 解析websocket数据包
@ -334,21 +349,25 @@ private:
if (_onRecv) {
// 握手成功之后的中途断开
_onRecv = nullptr;
_delegate.onErr(ex);
if (auto strong_ref = _weak_delegate.lock()) {
strong_ref->onErr(ex);
}
return;
}
// websocket握手失败或者tcp连接失败
_delegate.onConnect(ex);
if (auto strong_ref = _weak_delegate.lock()) {
strong_ref->onConnect(ex);
}
}
private:
std::string _Sec_WebSocket_Key;
std::function<void(const char *data, size_t len)> _onRecv;
std::weak_ptr<ClientTypeImp<ClientType, DataType>> _weak_delegate;
ClientTypeImp<ClientType, DataType> &_delegate;
std::string _payload_section;
std::string _payload_cache;
toolkit::Ticker _recv_ticker;
};
/**
@ -364,11 +383,8 @@ public:
using Ptr = std::shared_ptr<WebSocketClient>;
template <typename... ArgsType>
WebSocketClient(ArgsType &&...args) : ClientTypeImp<ClientType,DataType>(std::forward<ArgsType>(args)...){
}
~WebSocketClient() override {
_wsClient->closeWsClient();
}
WebSocketClient(ArgsType &&...args) : ClientTypeImp<ClientType, DataType>(std::forward<ArgsType>(args)...) {}
~WebSocketClient() override { _wsClient->closeWsClient(); }
/**
* 重载startConnect方法，
@ -392,15 +408,11 @@ public:
void startWebSocket(const std::string &ws_url, float fTimeOutSec = 3) {
_wsClient = std::make_shared<HttpWsClient<ClientType, DataType>>(std::static_pointer_cast<WebSocketClient>(this->shared_from_this()));
_wsClient->setOnCreateSocket([this](const toolkit::EventPoller::Ptr &){
return this->createSocket();
});
_wsClient->setOnCreateSocket([this](const toolkit::EventPoller::Ptr &) { return this->createSocket(); });
_wsClient->startWsClient(ws_url, fTimeOutSec);
}
HttpClient &getHttpClient() {
return *_wsClient;
}
HttpClient &getHttpClient() { return *_wsClient; }
private:
typename HttpWsClient<ClientType, DataType>::Ptr _wsClient;

25
src/Http/WebSocketSession.h
View File

@ -36,7 +36,7 @@ public:
using Ptr = std::shared_ptr<SessionTypeImp>;
SessionTypeImp(const mediakit::Parser &header, const mediakit::HttpSession &parent, const toolkit::Socket::Ptr &pSock) :
SessionType(pSock), _identifier(parent.getIdentifier()) {}
SessionType(pSock) {}
~SessionTypeImp() = default;
@ -61,12 +61,7 @@ protected:
return SessionType::send(std::move(buf));
}
std::string getIdentifier() const override {
return _identifier;
}
private:
std::string _identifier;
onBeforeSendCB _beforeSendCB;
};
@ -103,6 +98,21 @@ public:
} else {
HttpSessionType::onManager();
}
if (!_session) {
// websocket尚未链接
return;
}
if (_recv_ticker.elapsedTime() > 30 * 1000) {
HttpSessionType::shutdown(toolkit::SockException(toolkit::Err_timeout, "websocket timeout"));
} else if (_recv_ticker.elapsedTime() > 10 * 1000) {
// 没收到回复，每10秒发送次ping 包
mediakit::WebSocketHeader header;
header._fin = true;
header._reserved = 0;
header._opcode = mediakit::WebSocketHeader::PING;
header._mask_flag = false;
HttpSessionType::encode(header, nullptr);
}
}
void attachServer(const toolkit::Server &server) override{
@ -170,7 +180,7 @@ protected:
auto header = const_cast<mediakit::WebSocketHeader&>(header_in);
auto flag = header._mask_flag;
header._mask_flag = false;
_recv_ticker.resetTime();
switch (header._opcode){
case mediakit::WebSocketHeader::CLOSE:{
HttpSessionType::encode(header,nullptr);
@ -230,6 +240,7 @@ private:
std::weak_ptr<toolkit::Server> _weak_server;
toolkit::Session::Ptr _session;
Creator _creator;
toolkit::Ticker _recv_ticker;
};