diff --git a/CHANGELOG.md b/CHANGELOG.md index 52e96415..58d5b12c 100644 --- a/CHANGELOG.md +++ b/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) diff --git a/docs/chatterino.d.ts b/docs/chatterino.d.ts index 608f6bfa..3c7690fd 100644 --- a/docs/chatterino.d.ts +++ b/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; + on_open?: () => void; on_close?: () => void; on_text?: (data: string) => void; on_binary?: (data: string) => void; diff --git a/docs/plugin-meta.lua b/docs/plugin-meta.lua index cac380ae..b9e0b7bf 100644 --- a/docs/plugin-meta.lua +++ b/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, on_close?: fun(), on_text?: fun(data: string), on_binary?: fun(data: string) } Additional options for the connection. +---@param options? { headers?: table, 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 diff --git a/src/common/websockets/WebSocketPool.hpp b/src/common/websockets/WebSocketPool.hpp index 93140f5c..f5349d45 100644 --- a/src/common/websockets/WebSocketPool.hpp +++ b/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. diff --git a/src/common/websockets/detail/WebSocketConnectionImpl.cpp b/src/common/websockets/detail/WebSocketConnectionImpl.cpp index 308a76f3..0fdd5a8f 100644 --- a/src/common/websockets/detail/WebSocketConnectionImpl.cpp +++ b/src/common/websockets/detail/WebSocketConnectionImpl.cpp @@ -194,6 +194,7 @@ void WebSocketConnectionHelper::onWsHandshake( qCDebug(chatterinoWebsocket) << *this << "WS handshake done"; + this->listener->onOpen(); this->trySend(); this->stream.async_read( this->readBuffer, diff --git a/src/controllers/plugins/api/WebSocket.cpp b/src/controllers/plugins/api/WebSocket.cpp index 90464770..d1a859b9 100644 --- a/src/controllers/plugins/api/WebSocket.cpp +++ b/src/controllers/plugins/api/WebSocket.cpp @@ -19,6 +19,7 @@ public: void onClose(std::unique_ptr self) override; void onBinaryMessage(QByteArray data) override; void onTextMessage(QByteArray data) override; + void onOpen() override; private: std::weak_ptr target; @@ -58,10 +59,15 @@ void WebSocket::createUserType(sol::table &c2, Plugin *plugin) v.as()); } } + sol::optional onOpen = luaOpts["on_open"]; sol::optional onText = luaOpts["on_text"]; sol::optional onBinary = luaOpts["on_binary"]; sol::optional 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 diff --git a/src/controllers/plugins/api/WebSocket.hpp b/src/controllers/plugins/api/WebSocket.hpp index 21730381..4cc6637c 100644 --- a/src/controllers/plugins/api/WebSocket.hpp +++ b/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, on_close?: fun(), on_text?: fun(data: string), on_binary?: fun(data: string) } Additional options for the connection. + * @lua@param options? { headers?: table, 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; diff --git a/tests/src/Plugins.cpp b/tests/src/Plugins.cpp index ab87d7ac..9bbd675e 100644 --- a/tests/src/Plugins.cpp +++ b/tests/src/Plugins.cpp @@ -645,12 +645,19 @@ TEST_F(PluginTest, testTcpWebSocket) RequestWaiter waiter; std::vector> 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 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> 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 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); diff --git a/tests/src/WebSocketPool.cpp b/tests/src/WebSocketPool.cpp index 4851049e..efbf5713 100644 --- a/tests/src/WebSocketPool.cpp +++ b/tests/src/WebSocketPool.cpp @@ -11,10 +11,11 @@ namespace { struct Listener : public WebSocketListener { Listener(std::vector> &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> &messages; OnceFlag &messageFlag; OnceFlag &closeFlag; + OnceFlag &openFlag; }; } // namespace @@ -50,6 +57,7 @@ TEST(WebSocketPool, tcpEcho) std::vector> messages; OnceFlag messageFlag; OnceFlag closeFlag; + OnceFlag openFlag; auto handle = pool.createSocket( { @@ -62,13 +70,14 @@ TEST(WebSocketPool, tcpEcho) {"User-Agent", "MyUserAgent"}, }, }, - std::make_unique(messages, messageFlag, closeFlag)); + std::make_unique(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> 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(messages, messageFlag, closeFlag)); + std::make_unique(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");