feat(WebSocket): add open event (#6315)

CHANGELOG.md

@@ -7,7 +7,7 @@
 - Minor: Add an option for the reduced opacity of message history. (#6121)
 - Minor: Make paused chat indicator more visible, and fix its zoom behavior. (#6123)
 - Minor: Added interactive REPL for plugins. (#6120)
-- Minor: Added WebSocket API for plugins. (#6076, #6186, #6314)
+- Minor: Added WebSocket API for plugins. (#6076, #6186, #6314, #6315)
 - Minor: Allow for themes to set transparent values for window background on Linux. (#6137)
 - Minor: Popup overlay now only draws an outline when being interacted with. (#6140)
 - Minor: Made filters searchable in the Settings dialog search bar. (#5890)

docs/chatterino.d.ts

@@ -133,6 +133,7 @@ declare namespace c2 {
         close(): void;
         send_text(data: string): void;
         send_binary(data: string): void;
+        on_open: null | (() => void);
         on_close: null | (() => void);
         on_text: null | ((data: string) => void);
         on_binary: null | ((data: string) => void);
@@ -143,6 +144,7 @@ declare namespace c2 {
             url: string,
             options?: {
                 headers?: Record<string, string>;
+                on_open?: () => void;
                 on_close?: () => void;
                 on_text?: (data: string) => void;
                 on_binary?: (data: string) => void;

docs/plugin-meta.lua

@@ -263,13 +263,14 @@ function c2.HTTPRequest.create(method, url) end
 ---@field on_close fun()|nil Handler called when the socket is closed.
 ---@field on_text fun(data: string)|nil Handler called when the socket receives a text message.
 ---@field on_binary fun(data: string)|nil Handler called when the socket receives a binary message.
+---@field on_open fun()|nil Handler called when the websocket handshake has been completed successfully.
 c2.WebSocket = {}
 
 --- Creates and connects to a WebSocket server. Upon calling this, a
 --- connection is made immediately.
 ---
 ---@param url string The URL to connect to. Must start with `wss://` or `ws://`.
----@param options? { headers?: table<string, string>, on_close?: fun(), on_text?: fun(data: string), on_binary?: fun(data: string) } Additional options for the connection.
+---@param options? { headers?: table<string, string>, on_close?: fun(), on_text?: fun(data: string), on_binary?: fun(data: string), on_open?: fun() } Additional options for the connection.
 ---@return c2.WebSocket
 ---@nodiscard
 function c2.WebSocket.new(url, options) end

src/common/websockets/WebSocketPool.hpp

@@ -42,6 +42,11 @@ private:
 struct WebSocketListener {
     virtual ~WebSocketListener() = default;
 
+    /// The WebSocket handshake completed successfully.
+    ///
+    /// This function is called from the websocket thread.
+    virtual void onOpen() = 0;
+
     /// A text message was received.
     ///
     /// This function is called from the websocket thread.

src/common/websockets/detail/WebSocketConnectionImpl.cpp

@@ -194,6 +194,7 @@ void WebSocketConnectionHelper<Derived, Inner>::onWsHandshake(
 
     qCDebug(chatterinoWebsocket) << *this << "WS handshake done";
 
+    this->listener->onOpen();
     this->trySend();
     this->stream.async_read(
         this->readBuffer,

src/controllers/plugins/api/WebSocket.cpp

@@ -19,6 +19,7 @@ public:
     void onClose(std::unique_ptr<WebSocketListener> self) override;
     void onBinaryMessage(QByteArray data) override;
     void onTextMessage(QByteArray data) override;
+    void onOpen() override;
 
 private:
     std::weak_ptr<WebSocket> target;
@@ -58,10 +59,15 @@ void WebSocket::createUserType(sol::table &c2, Plugin *plugin)
                                                   v.as<std::string>());
                     }
                 }
+                sol::optional<sol::main_function> onOpen = luaOpts["on_open"];
                 sol::optional<sol::main_function> onText = luaOpts["on_text"];
                 sol::optional<sol::main_function> onBinary =
                     luaOpts["on_binary"];
                 sol::optional<sol::main_function> onClose = luaOpts["on_close"];
+                if (onOpen)
+                {
+                    self->onOpen = std::move(*onOpen);
+                }
                 if (onText)
                 {
                     self->onText = std::move(*onText);
@@ -108,6 +114,14 @@ void WebSocket::createUserType(sol::table &c2, Plugin *plugin)
             [](WebSocket &ws, sol::main_function fn) {
                 ws.onBinary = std::move(fn);
             }),
+        "on_open",
+        sol::property(
+            [](WebSocket &ws) {
+                return ws.onOpen;
+            },
+            [](WebSocket &ws, sol::main_function fn) {
+                ws.onOpen = std::move(fn);
+            }),
         "close", &WebSocket::close,            //
         "send_text", &WebSocket::sendText,     //
         "send_binary", &WebSocket::sendBinary  //
@@ -178,6 +192,19 @@ void WebSocketListenerProxy::onBinaryMessage(QByteArray data)
     });
 }
 
+void WebSocketListenerProxy::onOpen()
+{
+    auto target = this->target;
+    runInGuiThread([target] {
+        auto strong = target.lock();
+        if (strong && strong->onOpen)
+        {
+            loggedVoidCall(strong->onOpen, u"WebSocket.on_open",
+                           strong->plugin);
+        }
+    });
+}
+
 }  // namespace chatterino::lua::api
 
 #endif

src/controllers/plugins/api/WebSocket.hpp

@@ -23,7 +23,7 @@ public:
      * connection is made immediately.
      *
      * @lua@param url string The URL to connect to. Must start with `wss://` or `ws://`.
-     * @lua@param options? { headers?: table<string, string>, on_close?: fun(), on_text?: fun(data: string), on_binary?: fun(data: string) } Additional options for the connection.
+     * @lua@param options? { headers?: table<string, string>, on_close?: fun(), on_text?: fun(data: string), on_binary?: fun(data: string), on_open?: fun() } Additional options for the connection.
      * @lua@return c2.WebSocket
      * @lua@nodiscard
      * @exposed c2.WebSocket.new
@@ -66,6 +66,10 @@ private:
      * @lua@field on_binary fun(data: string)|nil Handler called when the socket receives a binary message.
      */
     sol::main_function onBinary;
+    /**
+     * @lua@field on_open fun()|nil Handler called when the websocket handshake has been completed successfully.
+     */
+    sol::main_function onOpen;
     WebSocketHandle handle;
     // Note: this class lives inside the plugin -> this pointer will be valid.
     Plugin *plugin = nullptr;

tests/src/Plugins.cpp

@@ -645,12 +645,19 @@ TEST_F(PluginTest, testTcpWebSocket)
 
     RequestWaiter waiter;
     std::vector<std::pair<bool, QByteArray>> messages;
+    bool open = false;
     lua->set("done", [&] {
         waiter.requestDone();
     });
     lua->set("add", [&](bool isText, QByteArray data) {
+        EXPECT_TRUE(open);
         messages.emplace_back(isText, std::move(data));
     });
+    // On GCC in release mode, using set() would cause the done function to be called instead.
+    lua->set_function("open", [&] {
+        EXPECT_FALSE(open);
+        open = true;
+    });
 
     std::shared_ptr<lua::api::WebSocket> ws = lua->script(R"lua(
         local ws = c2.WebSocket.new("ws://127.0.0.1:9052/echo")
@@ -671,6 +678,9 @@ TEST_F(PluginTest, testTcpWebSocket)
         ws.on_close = function()
             done()
         end
+        ws.on_open = function()
+            open()
+        end
         ws:send_text("message1")
         ws:send_text("message2")
         ws:send_text("message3")
@@ -682,6 +692,7 @@ TEST_F(PluginTest, testTcpWebSocket)
     ws.reset();
 
     waiter.waitForRequest();
+    ASSERT_TRUE(open);
 
     ASSERT_EQ(messages.size(), 7);
     ASSERT_EQ(messages[0].first, true);
@@ -709,13 +720,20 @@ TEST_F(PluginTest, testTlsWebSocket)
     configure({PluginPermission{{{"type", "Network"}}}});
 
     RequestWaiter waiter;
+    bool open = false;
     std::vector<std::pair<bool, QByteArray>> messages;
     lua->set("done", [&] {
         waiter.requestDone();
     });
     lua->set("add", [&](bool isText, QByteArray data) {
+        EXPECT_TRUE(open);
         messages.emplace_back(isText, std::move(data));
     });
+    // On GCC in release mode, using set() would cause the done function to be called instead.
+    lua->set_function("open", [&] {
+        EXPECT_FALSE(open);
+        open = true;
+    });
 
     std::shared_ptr<lua::api::WebSocket> ws = lua->script(R"lua(
         local ws = c2.WebSocket.new("wss://127.0.0.1:9050/echo", { 
@@ -744,6 +762,9 @@ TEST_F(PluginTest, testTlsWebSocket)
         ws.on_close = function()
             done()
         end
+        ws.on_open = function()
+            open()
+        end
         ws:send_text("message1")
         ws:send_text("message2")
         ws:send_text("message3")
@@ -755,6 +776,7 @@ TEST_F(PluginTest, testTlsWebSocket)
     ws.reset();
 
     waiter.waitForRequest();
+    ASSERT_TRUE(open);
 
     ASSERT_EQ(messages.size(), 9);
     ASSERT_EQ(messages[0].first, true);

tests/src/WebSocketPool.cpp

@@ -11,10 +11,11 @@ namespace {
 
 struct Listener : public WebSocketListener {
     Listener(std::vector<std::pair<bool, QByteArray>> &messages,
-             OnceFlag &messageFlag, OnceFlag &closeFlag)
+             OnceFlag &messageFlag, OnceFlag &closeFlag, OnceFlag &openFlag)
         : messages(messages)
         , messageFlag(messageFlag)
         , closeFlag(closeFlag)
+        , openFlag(openFlag)
     {
     }
 
@@ -35,9 +36,15 @@ struct Listener : public WebSocketListener {
         messages.emplace_back(false, std::move(data));
     }
 
+    void onOpen() override
+    {
+        this->openFlag.set();
+    }
+
     std::vector<std::pair<bool, QByteArray>> &messages;
     OnceFlag &messageFlag;
     OnceFlag &closeFlag;
+    OnceFlag &openFlag;
 };
 
 }  // namespace
@@ -50,6 +57,7 @@ TEST(WebSocketPool, tcpEcho)
     std::vector<std::pair<bool, QByteArray>> messages;
     OnceFlag messageFlag;
     OnceFlag closeFlag;
+    OnceFlag openFlag;
 
     auto handle = pool.createSocket(
         {
@@ -62,13 +70,14 @@ TEST(WebSocketPool, tcpEcho)
                     {"User-Agent", "MyUserAgent"},
                 },
         },
-        std::make_unique<Listener>(messages, messageFlag, closeFlag));
+        std::make_unique<Listener>(messages, messageFlag, closeFlag, openFlag));
     handle.sendBinary("message1");
     handle.sendBinary("message2");
     handle.sendBinary("message3");
     handle.sendText("message4");
 
     ASSERT_TRUE(messageFlag.waitFor(1s));
+    ASSERT_TRUE(openFlag.isSet());
     QByteArray bigMsg(1 << 15, 'A');
     handle.sendBinary(bigMsg);
     handle.sendText("foo");
@@ -111,6 +120,7 @@ TEST(WebSocketPool, tlsEcho)
     std::vector<std::pair<bool, QByteArray>> messages;
     OnceFlag messageFlag;
     OnceFlag closeFlag;
+    OnceFlag openFlag;
 
     auto handle = pool.createSocket(
         {
@@ -121,13 +131,14 @@ TEST(WebSocketPool, tlsEcho)
                 {"Cookie", "xd"},  // "known" header
             },
         },
-        std::make_unique<Listener>(messages, messageFlag, closeFlag));
+        std::make_unique<Listener>(messages, messageFlag, closeFlag, openFlag));
     handle.sendBinary("message1");
     handle.sendBinary("message2");
     handle.sendBinary("message3");
     handle.sendText("message4");
 
     ASSERT_TRUE(messageFlag.waitFor(1s));
+    ASSERT_TRUE(openFlag.isSet());
     QByteArray bigMsg(1 << 15, 'A');
     handle.sendBinary(bigMsg);
     handle.sendText("foo");