Add sub-protocol and optional header support to WebSocket connection (#…

…1565) * Add sub-protocol and optional header support to WebSocket connection Unsolicited WebSocket PONG messages received should be ignored * Set custom headers via setHeaders method for native WebSocket implementation. Ensure WASM WebSocket interface has same constructor as native interface * Fix temporary string reference Ensure passed url is passed as null terminated string
axmolengine · Jan 2, 2024 · 9ca7760 · 9ca7760
1 parent 99cf951
commit 9ca7760
Show file tree

Hide file tree

Showing 4 changed files with 58 additions and 20 deletions.
diff --git a/core/network/WebSocket-wasm.cpp b/core/network/WebSocket-wasm.cpp
@@ -54,7 +54,7 @@ WebSocket::WebSocket() {}
 
 WebSocket::~WebSocket() {}
 
-bool WebSocket::open(Delegate* delegate, std::string_view url, std::string_view caFilePath, const char* protocols)
+bool WebSocket::open(Delegate* delegate, std::string_view url, std::string_view caFilePath, std::string_view protocols)
 {
     if (url.empty())
     {
@@ -64,7 +64,11 @@ bool WebSocket::open(Delegate* delegate, std::string_view url, std::string_view
 
     _delegate = delegate;
 
-    EmscriptenWebSocketCreateAttributes ws_attrs = {url.data(), protocols, EM_TRUE};
+    _url          = url;
+    _subProtocols = protocols;
+
+    EmscriptenWebSocketCreateAttributes ws_attrs = {_url.c_str(),
+                                                    _subProtocols.empty() ? nullptr : _subProtocols.c_str(), EM_TRUE};
 
     AXLOG("ws open url: %s, protocols: %s", ws_attrs.url, ws_attrs.protocols);
 

diff --git a/core/network/WebSocket-wasm.h b/core/network/WebSocket-wasm.h
@@ -12,8 +12,8 @@ NS_AX_BEGIN
 namespace network
 {
 /**
- * WebSocket is wrapper of the libwebsockets-protocol, let the develop could call the websocket easily.
- * Please note that all public methods of WebSocket have to be invoked on Cocos Thread.
+ * WebSocket implementation using yasio.
+ * Please note that all public methods of WebSocket have to be invoked on Axmol Thread.
  */
 class AX_DLL WebSocket
 {
@@ -122,13 +122,14 @@ class AX_DLL WebSocket
      *  @param url The URL of websocket server.
      *  @param protocols The websocket protocols that agree with websocket server
      *  @param caFilePath The ca file path for wss connection
+     *  @param protocols Comma-separated list of sub-protocols that agree with websocket server
      *  @return true: Success, false: Failure.
      *  @lua NA
      */
     bool open(Delegate* delegate,
               std::string_view url,
               std::string_view caFilePath = "",
-              const char* protocols      = nullptr);
+              std::string_view protocols  = "");
 
     /**
      *  @brief Sends string data to websocket server.
@@ -141,7 +142,7 @@ class AX_DLL WebSocket
     /**
      *  @brief Sends binary data to websocket server.
      *
-     *  @param binaryMsg binary string data.
+     *  @param data binary string data.
      *  @param len the size of binary string data.
      *  @lua sendstring
      */
@@ -170,12 +171,12 @@ class AX_DLL WebSocket
     /**
      *  @brief Gets the URL of websocket connection.
      */
-    inline std::string_view getUrl() const { return _url; }
+    std::string_view getUrl() const { return _url; }
 
     /**
      *  @brief Gets the protocol selected by websocket server.
      */
-    inline std::string_view getProtocol() const { return _selectedProtocol; }
+    std::string_view getProtocol() const { return _selectedProtocol; }
 
 protected:
     static EM_BOOL em_ws_onopen(int eventType, const EmscriptenWebSocketOpenEvent *websocketEvent, void *userData);
@@ -187,6 +188,7 @@ class AX_DLL WebSocket
     int _wsfd = -1;
     State _state = State::CLOSED;
     std::string _url;
+    std::string _subProtocols;
     std::string _selectedProtocol;
 
     Delegate* _delegate{};

diff --git a/core/network/WebSocket.cpp b/core/network/WebSocket.cpp
@@ -27,6 +27,8 @@
 
 #include "network/WebSocket.h"
 
+#include "fmt/format.h"
+
 using namespace yasio;
 
 #define WS_MAX_PAYLOAD_LENGTH (1 << 24)  // 16M
@@ -227,12 +229,13 @@ WebSocket::~WebSocket()
 bool WebSocket::open(Delegate* delegate,
                      std::string_view url,
                      std::string_view caFilePath,
-                     const char* protocols)
+                     std::string_view protocols)
 {
     _delegate   = delegate;
     _url        = url;
     _caFilePath = FileUtils::getInstance()->fullPathForFilename(caFilePath);
     _requestUri = Uri::parse(url);
+    _protocols  = protocols;
 
     setupParsers();
     generateHandshakeSecKey();
@@ -393,7 +396,7 @@ int WebSocket::on_frame_end(websocket_parser* parser)
         case WS_OP_PONG:
             AXLOG("WS: control frame: PONG");
             if (ws->_receivedData.size() != 4 || 0 != memcmp(ws->_receivedData.data(), "WSWS", 4))
-                AXLOG("WS: PONG frame from server mismatch!\n");
+                AXLOG("WS: Unsolicited PONG frame from server (possible keep-alive)\n\n");
             break;
         }
     }
@@ -417,7 +420,7 @@ void WebSocket::send(std::string_view message)
 /**
  *  @brief Sends binary data to websocket server.
  *
- *  @param binaryMsg binary string data.
+ *  @param data binary data.
  *  @param len the size of binary string data.
  *  @lua sendstring
  */
@@ -533,9 +536,9 @@ void WebSocket::handleNetworkEvent(yasio::io_event* event)
 
                     // WebSocketProtocol::sendPing(*this);
 
-                    // starts websocket hartbeat timer
+                    // starts websocket heartbeat timer
                     //_heartbeatTimer->expires_from_now(std::chrono::seconds(30));
-                    //_heartbeatTimer->async_wait(*_service, [this](io_service& service) {
+                    //_heartbeatTimer->async_wait([this](io_service& service) {
                     //    WebSocketProtocol::sendFrame(*this, "WSWS", 4, ws::detail::opcode::ping);
                     //    return false;
                     //});
@@ -592,6 +595,19 @@ void WebSocket::handleNetworkEvent(yasio::io_event* event)
             obs.write_bytes(_handshakeSecKey);
             obs.write_bytes("\r\n");
 
+            if (!_protocols.empty())
+            {
+                obs.write_bytes("Sec-WebSocket-Protocol: ");
+                obs.write_bytes(_protocols);
+                obs.write_bytes("\r\n");
+            }
+
+            for (auto&& header : _headers)
+            {
+                obs.write_bytes(header);
+                obs.write_bytes("\r\n");
+            }
+
             obs.write_bytes("User-Agent: yasio-ws\r\n");
 
             obs.write_bytes("Accept: */*;q=0.8\r\n");

diff --git a/core/network/WebSocket.h b/core/network/WebSocket.h
@@ -65,8 +65,8 @@ namespace network
 struct WebSocketProtocol;
 
 /**
- * WebSocket is wrapper of the libwebsockets-protocol, let the develop could call the websocket easily.
- * Please note that all public methods of WebSocket have to be invoked on Cocos Thread.
+ * WebSocket implementation using yasio.
+ * Please note that all public methods of WebSocket have to be invoked on Axmol Thread.
  */
 class AX_DLL WebSocket
 {
@@ -244,14 +244,14 @@ class AX_DLL WebSocket
      *  @param delegate The delegate which want to receive event from websocket.
      *  @param url The URL of websocket server.
      *  @param caFilePath The ca file path for wss connection
-     *  @param protocols The websocket protocols that agree with websocket server
+     *  @param protocols Comma-separated list of sub-protocols that agree with websocket server
      *  @return true: Success, false: Failure.
      *  @lua NA
      */
     bool open(Delegate* delegate,
               std::string_view url,
               std::string_view caFilePath = "",
-              const char* protocols       = nullptr);
+              std::string_view protocols = "");
 
     /**
      *  @brief Sends string data to websocket server.
@@ -264,7 +264,7 @@ class AX_DLL WebSocket
     /**
      *  @brief Sends binary data to websocket server.
      *
-     *  @param binaryMsg binary string data.
+     *  @param data binary data.
      *  @param len the size of binary string data.
      *  @lua sendstring
      */
@@ -293,12 +293,26 @@ class AX_DLL WebSocket
     /**
      *  @brief Gets the URL of websocket connection.
      */
-    inline std::string_view getUrl() const { return _url; }
+    std::string_view getUrl() const { return _url; }
 
     /**
      *  @brief Gets the protocol selected by websocket server.
      */
-    inline std::string_view getProtocol() const { return _selectedProtocol; }
+    std::string_view getProtocol() const { return _selectedProtocol; }
+
+    /**
+     * Set custom-defined headers.
+     *
+     * @param headers The string vector of custom-defined headers.
+     */
+    void setHeaders(const std::vector<std::string>& headers) { _headers = headers; }
+
+    /**
+     * Get custom headers.
+     *
+     * @return std::vector<std::string> the string vector of custom-defined headers.
+     */
+    const std::vector<std::string>& getHeaders() const { return _headers; }
 
 protected:
     void purgePendingEvents();
@@ -418,6 +432,8 @@ class AX_DLL WebSocket
     EventListenerCustom* _resetDirectorListener;
 
     ConcurrentDeque<Event*> _eventQueue;
+    std::vector<std::string> _headers;  /// custom headers
+    std::string _protocols;
 };
 }  // namespace network