1、ws-flv直播支持客户端主动关闭请求:#430

2、兼容CONTINUATION类型的websocket包
3、修复websocket客户端在处理Content-Length时的相关bug

src/Http/HttpSession.cpp

@@ -132,36 +132,37 @@ void HttpSession::onManager() {
 
 bool HttpSession::checkWebSocket(){
     auto Sec_WebSocket_Key = _parser["Sec-WebSocket-Key"];
-    if(Sec_WebSocket_Key.empty()){
+    if (Sec_WebSocket_Key.empty()) {
         return false;
     }
-    auto Sec_WebSocket_Accept = encodeBase64(SHA1::encode_bin(Sec_WebSocket_Key + "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"));
+    auto Sec_WebSocket_Accept = encodeBase64(
+            SHA1::encode_bin(Sec_WebSocket_Key + "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"));
 
     KeyValue headerOut;
     headerOut["Upgrade"] = "websocket";
     headerOut["Connection"] = "Upgrade";
     headerOut["Sec-WebSocket-Accept"] = Sec_WebSocket_Accept;
-    if(!_parser["Sec-WebSocket-Protocol"].empty()){
+    if (!_parser["Sec-WebSocket-Protocol"].empty()) {
         headerOut["Sec-WebSocket-Protocol"] = _parser["Sec-WebSocket-Protocol"];
     }
 
-    auto res_cb = [this,headerOut](){
+    auto res_cb = [this, headerOut]() {
         _flv_over_websocket = true;
-        sendResponse("101 Switching Protocols",false,nullptr,headerOut,nullptr, true);
+        sendResponse("101 Switching Protocols", false, nullptr, headerOut, nullptr, true);
     };
 
     //判断是否为websocket-flv
-    if(checkLiveFlvStream(res_cb)){
+    if (checkLiveFlvStream(res_cb)) {
         //这里是websocket-flv直播请求
         return true;
     }
 
     //如果checkLiveFlvStream返回false,则代表不是websocket-flv，而是普通的websocket连接
-    if(!onWebSocketConnect(_parser)){
+    if (!onWebSocketConnect(_parser)) {
-        sendResponse("501 Not Implemented",true, nullptr, headerOut);
+        sendResponse("501 Not Implemented", true, nullptr, headerOut);
         return true;
     }
-    sendResponse("101 Switching Protocols",false, nullptr,headerOut);
+    sendResponse("101 Switching Protocols", false, nullptr, headerOut, nullptr, true);
     return true;
 }
 
@@ -389,7 +390,7 @@ void HttpSession::sendResponse(const char *pcStatus,
                                const char *pcContentType,
                                const HttpSession::KeyValue &header,
                                const HttpBody::Ptr &body,
-                               bool is_http_flv ){
+                               bool no_content_length ){
     GET_CONFIG(string,charSet,Http::kCharSet);
     GET_CONFIG(uint32_t,keepAliveSec,Http::kKeepAliveSecond);
 
@@ -400,7 +401,7 @@ void HttpSession::sendResponse(const char *pcStatus,
         size = body->remainSize();
     }
 
-    if(is_http_flv){
+    if(no_content_length){
         //http-flv直播是Keep-Alive类型
         bClose = false;
     }else if(size >= INT64_MAX){
@@ -425,7 +426,7 @@ void HttpSession::sendResponse(const char *pcStatus,
         headerOut.emplace(kAccessControlAllowCredentials, "true");
     }
 
-    if(!is_http_flv && size >= 0 && size < INT64_MAX){
+    if(!no_content_length && size >= 0 && size < INT64_MAX){
         //文件长度为固定值,且不是http-flv强制设置Content-Length
         headerOut[kContentLength] = to_string(size);
     }
@@ -645,6 +646,21 @@ void HttpSession::onWebSocketEncodeData(const Buffer::Ptr &buffer){
     send(buffer);
 }
 
+void HttpSession::onWebSocketDecodeComplete(const WebSocketHeader &header_in){
+    WebSocketHeader& header = const_cast<WebSocketHeader&>(header_in);
+    header._mask_flag = false;
+
+    switch (header._opcode) {
+        case WebSocketHeader::CLOSE: {
+            encode(header, nullptr);
+            shutdown(SockException(Err_shutdown, "recv close request from client"));
+            break;
+        }
+
+        default : break;
+    }
+}
+
 void HttpSession::onDetach() {
     shutdown(SockException(Err_shutdown,"rtmp ring buffer detached"));
 }

src/Http/HttpSession.h

@@ -47,6 +47,7 @@ public:
     void onError(const SockException &err) override;
     void onManager() override;
     static string urlDecode(const string &str);
+
 protected:
     //FlvMuxer override
     void onWrite(const Buffer::Ptr &data, bool flush) override ;
@@ -90,6 +91,13 @@ protected:
      * @param buffer websocket协议数据
      */
     void onWebSocketEncodeData(const Buffer::Ptr &buffer) override;
+
+    /**
+     * 接收到完整的一个webSocket数据包后回调
+     * @param header 数据包包头
+     */
+    void onWebSocketDecodeComplete(const WebSocketHeader &header_in) override;
+
 private:
     void Handle_Req_GET(int64_t &content_len);
     void Handle_Req_GET_l(int64_t &content_len, bool sendBody);
@@ -103,10 +111,11 @@ private:
     void sendNotFound(bool bClose);
     void sendResponse(const char *pcStatus, bool bClose, const char *pcContentType = nullptr,
                       const HttpSession::KeyValue &header = HttpSession::KeyValue(),
-                      const HttpBody::Ptr &body = nullptr,bool is_http_flv = false);
+                      const HttpBody::Ptr &body = nullptr, bool no_content_length = false);
 
     //设置socket标志
     void setSocketFlags();
+
 private:
     string _origin;
     Parser _parser;

src/Http/WebSocketClient.h

@@ -38,11 +38,10 @@ public:
     template<typename ...ArgsType>
     ClientTypeImp(ArgsType &&...args): ClientType(std::forward<ArgsType>(args)...){}
     ~ClientTypeImp() override {};
+
 protected:
     /**
      * 发送前拦截并打包为websocket协议
-     * @param buf
-     * @return
      */
     int send(const Buffer::Ptr &buf) override{
         if(_beforeSendCB){
@@ -50,6 +49,7 @@ protected:
         }
         return ClientType::send(buf);
     }
+
     /**
      * 设置发送数据截取回调函数
      * @param cb 截取回调函数
@@ -57,6 +57,7 @@ protected:
     void setOnBeforeSendCB(const onBeforeSendCB &cb){
         _beforeSendCB = cb;
     }
+
 private:
     onBeforeSendCB _beforeSendCB;
 };
@@ -108,6 +109,7 @@ public:
         header._mask_flag = true;
         WebSocketSplitter::encode(header, nullptr);
     }
+
 protected:
     //HttpClientImp override
 
@@ -124,6 +126,8 @@ protected:
             if(Sec_WebSocket_Accept == const_cast<HttpHeader &>(headers)["Sec-WebSocket-Accept"]){
                 //success
                 onWebSocketException(SockException());
+                //防止ws服务器返回Content-Length
+                const_cast<HttpHeader &>(headers).erase("Content-Length");
                 //后续全是websocket负载数据
                 return -1;
             }
@@ -180,7 +184,6 @@ protected:
 
     /**
      * tcp连接结果
-     * @param ex
      */
     void onConnect(const SockException &ex) override{
         if(ex){
@@ -194,7 +197,6 @@ protected:
 
     /**
      * tcp连接断开
-     * @param ex
      */
     void onErr(const SockException &ex) override{
         //tcp断开或者shutdown导致的断开
@@ -208,7 +210,7 @@ protected:
      * @param header 数据包头
      */
     void onWebSocketDecodeHeader(const WebSocketHeader &header) override{
-        _payload.clear();
+        _payload_section.clear();
     }
 
     /**
@@ -219,10 +221,9 @@ protected:
      * @param recved 已接收数据长度(包含本次数据长度)，等于header._payload_len时则接受完毕
      */
     void onWebSocketDecodePayload(const WebSocketHeader &header, const uint8_t *ptr, uint64_t len, uint64_t recved) override{
-        _payload.append((char *)ptr,len);
+        _payload_section.append((char *)ptr,len);
     }
 
-
     /**
      * 接收到完整的一个webSocket数据包后回调
      * @param header 数据包包头
@@ -238,28 +239,46 @@ protected:
                 //服务器主动关闭
                 WebSocketSplitter::encode(header,nullptr);
                 shutdown(SockException(Err_eof,"websocket server close the connection"));
-            }
                 break;
+            }
+
             case WebSocketHeader::PING:{
                 //心跳包
                 header._opcode = WebSocketHeader::PONG;
-                WebSocketSplitter::encode(header,std::make_shared<BufferString>(std::move(_payload)));
+                WebSocketSplitter::encode(header,std::make_shared<BufferString>(std::move(_payload_section)));
-            }
                 break;
-            case WebSocketHeader::CONTINUATION:{
+            }
 
-            }
+            case WebSocketHeader::CONTINUATION:
-                break;
             case WebSocketHeader::TEXT:
             case WebSocketHeader::BINARY:{
-                //接收完毕websocket数据包，触发onRecv事件
+                if (!header._fin) {
-                _delegate.onRecv(std::make_shared<BufferString>(std::move(_payload)));
+                    //还有后续分片数据, 我们先缓存数据，所有分片收集完成才一次性输出
+                    _payload_cache.append(std::move(_payload_section));
+                    if (_payload_cache.size() < MAX_WS_PACKET) {
+                        //还有内存容量缓存分片数据
+                        break;
+                    }
+                    //分片缓存太大，需要清空
+                }
+
+                //最后一个包
+                if (_payload_cache.empty()) {
+                    //这个包是唯一个分片
+                    _delegate.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)));
+                _payload_cache.clear();
+                break;
             }
-                break;
+
-            default:
+            default: break;
-                break;
         }
-        _payload.clear();
+        _payload_section.clear();
         header._mask_flag = flag;
     }
 
@@ -271,6 +290,7 @@ protected:
     void onWebSocketEncodeData(const Buffer::Ptr &buffer) override{
         HttpClientImp::send(buffer);
     }
+
 private:
     void onWebSocketException(const SockException &ex){
         if(!ex){
@@ -319,10 +339,10 @@ private:
     string _Sec_WebSocket_Key;
     function<void(const char *data, int len)> _onRecv;
     ClientTypeImp<ClientType,DataType> &_delegate;
-    string _payload;
+    string _payload_section;
+    string _payload_cache;
 };
 
-
 /**
  * Tcp客户端转WebSocket客户端模板，
  * 通过该模板，开发者再不修改TcpClient派生类任何代码的情况下快速实现WebSocket协议的包装
@@ -365,6 +385,7 @@ public:
     void startWebSocket(const string &ws_url,float fTimeOutSec = 3){
         _wsClient->startWsClient(ws_url,fTimeOutSec);
     }
+
 private:
     typename HttpWsClient<ClientType,DataType>::Ptr _wsClient;
 };

src/Http/WebSocketSession.h

@@ -78,7 +78,6 @@ public:
     }
 };
 
-
 /**
 * 通过该模板类可以透明化WebSocket协议，
 * 用户只要实现WebSock协议下的具体业务协议，譬如基于WebSocket协议的Rtmp协议等
@@ -107,8 +106,9 @@ public:
 
     void attachServer(const TcpServer &server) override{
         HttpSessionType::attachServer(server);
-        _weakServer = const_cast<TcpServer &>(server).shared_from_this();
+        _weak_server = const_cast<TcpServer &>(server).shared_from_this();
     }
+
 protected:
     /**
      * websocket客户端连接上事件
@@ -122,7 +122,7 @@ protected:
             //此url不允许创建websocket连接
             return false;
         }
-        auto strongServer = _weakServer.lock();
+        auto strongServer = _weak_server.lock();
         if(strongServer){
             _session->attachServer(*strongServer);
         }
@@ -145,24 +145,20 @@ protected:
         //允许websocket客户端
         return true;
     }
+
     /**
      * 开始收到一个webSocket数据包
-     * @param packet
      */
     void onWebSocketDecodeHeader(const WebSocketHeader &packet) override{
         //新包，原来的包残余数据清空掉
-        _remian_data.clear();
+        _payload_section.clear();
     }
 
     /**
      * 收到websocket数据包负载
-     * @param packet
-     * @param ptr
-     * @param len
-     * @param recved
      */
     void onWebSocketDecodePayload(const WebSocketHeader &packet,const uint8_t *ptr,uint64_t len,uint64_t recved) override {
-        _remian_data.append((char *)ptr,len);
+        _payload_section.append((char *)ptr,len);
     }
 
     /**
@@ -178,39 +174,59 @@ protected:
             case WebSocketHeader::CLOSE:{
                 HttpSessionType::encode(header,nullptr);
                 HttpSessionType::shutdown(SockException(Err_shutdown, "recv close request from client"));
-            }
                 break;
+            }
+            
             case WebSocketHeader::PING:{
                 header._opcode = WebSocketHeader::PONG;
-                HttpSessionType::encode(header,std::make_shared<BufferString>(_remian_data));
+                HttpSessionType::encode(header,std::make_shared<BufferString>(_payload_section));
-            }
                 break;
-            case WebSocketHeader::CONTINUATION:{
-
             }
-                break;
+            
+            case WebSocketHeader::CONTINUATION:
             case WebSocketHeader::TEXT:
             case WebSocketHeader::BINARY:{
-                _session->onRecv(std::make_shared<BufferString>(_remian_data));
+                if (!header._fin) {
+                    //还有后续分片数据, 我们先缓存数据，所有分片收集完成才一次性输出
+                    _payload_cache.append(std::move(_payload_section));
+                    if (_payload_cache.size() < MAX_WS_PACKET) {
+                        //还有内存容量缓存分片数据
+                        break;
+                    }
+                    //分片缓存太大，需要清空
+                }
+
+                //最后一个包
+                if (_payload_cache.empty()) {
+                    //这个包是唯一个分片
+                    _session->onRecv(std::make_shared<WebSocketBuffer>(header._opcode, header._fin, std::move(_payload_section)));
+                    break;
+                }
+
+                //这个包由多个分片组成
+                _payload_cache.append(std::move(_payload_section));
+                _session->onRecv(std::make_shared<WebSocketBuffer>(header._opcode, header._fin, std::move(_payload_cache)));
+                _payload_cache.clear();
+                break;
             }
-                break;
+            
-            default:
+            default: break;
-                break;
         }
-        _remian_data.clear();
+        _payload_section.clear();
         header._mask_flag = flag;
     }
 
     /**
-    * 发送数据进行websocket协议打包后回调
+     * 发送数据进行websocket协议打包后回调
-    * @param buffer
     */
     void onWebSocketEncodeData(const Buffer::Ptr &buffer) override{
         HttpSessionType::send(buffer);
     }
+
 private:
-    string _remian_data;
+    string _payload_cache;
-    weak_ptr<TcpServer> _weakServer;
+    string _payload_section;
+    weak_ptr<TcpServer> _weak_server;
     TcpSession::Ptr _session;
     Creator _creator;
 };

src/Http/WebSocketSplitter.h

@@ -16,10 +16,12 @@
 #include <vector>
 #include <memory>
 #include "Network/Buffer.h"
-
 using namespace std;
 using namespace toolkit;
 
+//websocket组合包最大不得超过4MB(防止内存爆炸)
+#define MAX_WS_PACKET (4 * 1024 * 1024)
+
 namespace mediakit {
 
 class WebSocketHeader {
@@ -44,6 +46,7 @@ public:
         CONTROL_RSVF = 0xF
     } Type;
 public:
+
     WebSocketHeader() : _mask(4){
         //获取_mask内部buffer的内存地址，该内存是malloc开辟的，地址为随机
         uint64_t ptr = (uint64_t)(&_mask[0]);
@@ -51,6 +54,7 @@ public:
         _mask.assign((uint8_t*)(&ptr), (uint8_t*)(&ptr) + 4);
     }
     virtual ~WebSocketHeader(){}
+
 public:
     bool _fin;
     uint8_t _reserved;
@@ -60,6 +64,26 @@ public:
     vector<uint8_t > _mask;
 };
 
+//websocket协议收到的字符串类型缓存，用户协议层获取该数据传输的方式
+class WebSocketBuffer : public BufferString {
+public:
+    typedef std::shared_ptr<WebSocketBuffer> Ptr;
+
+    template<typename ...ARGS>
+    WebSocketBuffer(WebSocketHeader::Type headType, bool fin, ARGS &&...args)
+            : _head_type(headType), _fin(fin), BufferString(std::forward<ARGS>(args)...) {}
+
+    ~WebSocketBuffer() override {}
+
+    WebSocketHeader::Type headType() const { return _head_type; }
+
+    bool isFinished() const { return _fin; };
+
+private:
+    WebSocketHeader::Type _head_type;
+    bool _fin;
+};
+
 class WebSocketSplitter : public WebSocketHeader{
 public:
     WebSocketSplitter(){}
@@ -80,6 +104,7 @@ public:
      * @param buffer 负载数据
      */
     void encode(const WebSocketHeader &header,const Buffer::Ptr &buffer);
+
 protected:
     /**
      * 收到一个webSocket数据包包头，后续将继续触发onWebSocketDecodePayload回调
@@ -96,7 +121,6 @@ protected:
      */
     virtual void onWebSocketDecodePayload(const WebSocketHeader &header, const uint8_t *ptr, uint64_t len, uint64_t recved) {};
 
-
     /**
      * 接收到完整的一个webSocket数据包后回调
      * @param header 数据包包头
@@ -109,8 +133,10 @@ protected:
      * @param len 数据指针长度
      */
     virtual void onWebSocketEncodeData(const Buffer::Ptr &buffer){};
+
 private:
     void onPayloadData(uint8_t *data, uint64_t len);
+
 private:
     string _remain_data;
     int _mask_offset = 0;