feat(plugins): add c2.WebSocket (#6076)

2025-04-07 19:38:10 +02:00
parent d3bab9132f
commit ab66be21b3
31 changed files with 1859 additions and 29 deletions
@@ -35,6 +35,7 @@ public:
    void waitForRequest()
    {
        using namespace std::chrono_literals;
+        auto start = std::chrono::system_clock::now();

        while (true)
        {
@@ -50,6 +51,11 @@ public:
            }
            QCoreApplication::processEvents(QEventLoop::AllEvents);
            QCoreApplication::sendPostedEvents(nullptr, QEvent::DeferredDelete);
+
+            if (std::chrono::system_clock::now() - start > 2min)
+            {
+                throw std::runtime_error("Timeout");
+            }
        }

        ASSERT_TRUE(this->requestDone_);
@@ -5,6 +5,7 @@
 #    include "controllers/commands/Command.hpp"  // IWYU pragma: keep
 #    include "controllers/commands/CommandController.hpp"
 #    include "controllers/plugins/api/ChannelRef.hpp"
+#    include "controllers/plugins/api/WebSocket.hpp"
 #    include "controllers/plugins/Plugin.hpp"
 #    include "controllers/plugins/PluginController.hpp"
 #    include "controllers/plugins/PluginPermission.hpp"
@@ -638,4 +639,176 @@ TEST_F(PluginTest, tryCallTest)
    }
 }

+TEST_F(PluginTest, testTcpWebSocket)
+{
+    configure({PluginPermission{{{"type", "Network"}}}});
+
+    RequestWaiter waiter;
+    std::vector<std::pair<bool, QByteArray>> messages;
+    lua->set("done", [&] {
+        waiter.requestDone();
+    });
+    lua->set("add", [&](bool isText, QByteArray data) {
+        messages.emplace_back(isText, std::move(data));
+    });
+
+    std::shared_ptr<lua::api::WebSocket> ws = lua->script(R"lua(
+        local ws = c2.WebSocket.new("ws://127.0.0.1:9052/echo")
+        ws.on_text = function(data)
+            add(true, data)
+        end
+        local any_msg = false
+        ws.on_binary = function(data)
+            if not any_msg then
+                any_msg = true
+                ws:send_text(string.rep("a", 1 << 15))
+                ws:send_binary("wow")
+                ws:send_text("/HEADER user-agent")
+                ws:send_binary("/CLOSE")
+            end
+            add(false, data)
+        end
+        ws.on_close = function()
+            done()
+        end
+        ws:send_text("message1")
+        ws:send_text("message2")
+        ws:send_text("message3")
+        ws:send_binary("message4")
+
+        return ws
+    )lua");
+    std::weak_ptr<lua::api::WebSocket> weakWs{ws};
+    ws.reset();
+
+    waiter.waitForRequest();
+
+    ASSERT_EQ(messages.size(), 7);
+    ASSERT_EQ(messages[0].first, true);
+    ASSERT_EQ(messages[0].second, "message1");
+    ASSERT_EQ(messages[1].first, true);
+    ASSERT_EQ(messages[1].second, "message2");
+    ASSERT_EQ(messages[2].first, true);
+    ASSERT_EQ(messages[2].second, "message3");
+    ASSERT_EQ(messages[3].first, false);
+    ASSERT_EQ(messages[3].second, "message4");
+    ASSERT_EQ(messages[4].first, true);
+    ASSERT_EQ(messages[4].second, QByteArray(1 << 15, 'a'));
+    ASSERT_EQ(messages[5].first, false);
+    ASSERT_EQ(messages[5].second, "wow");
+    ASSERT_EQ(messages[6].first, true);
+    ASSERT_TRUE(messages[6].second.startsWith("Chatterino"));
+
+    ASSERT_FALSE(weakWs.expired());
+    lua->collect_garbage();
+    ASSERT_TRUE(weakWs.expired());
+}
+
+TEST_F(PluginTest, testTlsWebSocket)
+{
+    configure({PluginPermission{{{"type", "Network"}}}});
+
+    RequestWaiter waiter;
+    std::vector<std::pair<bool, QByteArray>> messages;
+    lua->set("done", [&] {
+        waiter.requestDone();
+    });
+    lua->set("add", [&](bool isText, QByteArray data) {
+        messages.emplace_back(isText, std::move(data));
+    });
+
+    std::shared_ptr<lua::api::WebSocket> ws = lua->script(R"lua(
+        local ws = c2.WebSocket.new("wss://127.0.0.1:9050/echo", { 
+            headers = {
+                ["User-Agent"] = "Lua",
+                ["A-Header"] = "A value",
+                ["Referer"] = "https://chatterino.com",
+            },
+        })
+        ws.on_text = function(data)
+            add(true, data)
+        end
+        local any_msg = false
+        ws.on_binary = function(data)
+            if not any_msg then
+                any_msg = true
+                ws:send_text(string.rep("a", 1 << 15))
+                ws:send_binary("wow")
+                ws:send_text("/HEADER user-agent")
+                ws:send_text("/HEADER a-header")
+                ws:send_text("/HEADER referer")
+                ws:send_binary("/CLOSE")
+            end
+            add(false, data)
+        end
+        ws.on_close = function()
+            done()
+        end
+        ws:send_text("message1")
+        ws:send_text("message2")
+        ws:send_text("message3")
+        ws:send_binary("message4")
+
+        return ws
+    )lua");
+    std::weak_ptr<lua::api::WebSocket> weakWs{ws};
+    ws.reset();
+
+    waiter.waitForRequest();
+
+    ASSERT_EQ(messages.size(), 9);
+    ASSERT_EQ(messages[0].first, true);
+    ASSERT_EQ(messages[0].second, "message1");
+    ASSERT_EQ(messages[1].first, true);
+    ASSERT_EQ(messages[1].second, "message2");
+    ASSERT_EQ(messages[2].first, true);
+    ASSERT_EQ(messages[2].second, "message3");
+    ASSERT_EQ(messages[3].first, false);
+    ASSERT_EQ(messages[3].second, "message4");
+    ASSERT_EQ(messages[4].first, true);
+    ASSERT_EQ(messages[4].second, QByteArray(1 << 15, 'a'));
+    ASSERT_EQ(messages[5].first, false);
+    ASSERT_EQ(messages[5].second, "wow");
+    ASSERT_EQ(messages[6].first, true);
+    ASSERT_EQ(messages[6].second, "Lua");
+    ASSERT_EQ(messages[7].first, true);
+    ASSERT_EQ(messages[7].second, "A value");
+    ASSERT_EQ(messages[8].first, true);
+    ASSERT_EQ(messages[8].second, "https://chatterino.com");
+
+    ASSERT_FALSE(weakWs.expired());
+    lua->collect_garbage();
+    ASSERT_TRUE(weakWs.expired());
+}
+
+TEST_F(PluginTest, testWebSocketNoPerms)
+{
+    configure();
+
+    bool res = lua->script(R"lua(
+        return c2["WebSocket"] == nil
+    )lua");
+    ASSERT_TRUE(res);
+}
+
+TEST_F(PluginTest, testWebSocketApi)
+{
+    configure({PluginPermission{{{"type", "Network"}}}});
+
+    bool ok = lua->script(R"lua(
+        local t = function () end
+        local b = function () end
+        local c = function () end
+        local ws = c2.WebSocket.new("wss://127.0.0.1:9050/echo", { 
+            on_text = t,
+            on_binary = b,
+            on_close = c,
+        })
+
+        return ws.on_text == t and ws.on_binary == b and ws.on_close == c
+    )lua");
+
+    ASSERT_TRUE(ok);
+}
+
 #endif
@@ -0,0 +1,163 @@
+#include "common/websockets/WebSocketPool.hpp"
+
+#include "mocks/BaseApplication.hpp"
+#include "Test.hpp"
+#include "util/OnceFlag.hpp"
+
+using namespace chatterino;
+using namespace std::chrono_literals;
+
+namespace {
+
+struct Listener : public WebSocketListener {
+    Listener(std::vector<std::pair<bool, QByteArray>> &messages,
+             OnceFlag &messageFlag, OnceFlag &closeFlag)
+        : messages(messages)
+        , messageFlag(messageFlag)
+        , closeFlag(closeFlag)
+    {
+    }
+
+    void onClose(std::unique_ptr<WebSocketListener> /*self*/) override
+    {
+        this->closeFlag.set();
+    }
+
+    void onTextMessage(QByteArray data) override
+    {
+        this->messageFlag.set();
+        messages.emplace_back(true, std::move(data));
+    }
+
+    void onBinaryMessage(QByteArray data) override
+    {
+        // no flag to know when the initial queue is empty
+        messages.emplace_back(false, std::move(data));
+    }
+
+    std::vector<std::pair<bool, QByteArray>> &messages;
+    OnceFlag &messageFlag;
+    OnceFlag &closeFlag;
+};
+
+}  // namespace
+
+TEST(WebSocketPool, tcpEcho)
+{
+    mock::BaseApplication app;
+    WebSocketPool pool;
+
+    std::vector<std::pair<bool, QByteArray>> messages;
+    OnceFlag messageFlag;
+    OnceFlag closeFlag;
+
+    auto handle = pool.createSocket(
+        {
+            .url = QUrl("ws://127.0.0.1:9052/echo"),
+            .headers =
+                {
+                    {"My-Header", "my-header-VALUE"},
+                    {"Another-Header", "other-header"},
+                    {"Cookie", "xd"},  // "known" header
+                    {"User-Agent", "MyUserAgent"},
+                },
+        },
+        std::make_unique<Listener>(messages, messageFlag, closeFlag));
+    handle.sendBinary("message1");
+    handle.sendBinary("message2");
+    handle.sendBinary("message3");
+    handle.sendText("message4");
+
+    ASSERT_TRUE(messageFlag.waitFor(1s));
+    QByteArray bigMsg(1 << 15, 'A');
+    handle.sendBinary(bigMsg);
+    handle.sendText("foo");
+    handle.sendText("/HEADER my-header");
+    handle.sendText("/HEADER another-header");
+    handle.sendText("/HEADER cookie");
+    handle.sendText("/HEADER user-agent");
+    handle.sendText("/CLOSE");
+
+    ASSERT_TRUE(closeFlag.waitFor(1s));
+
+    ASSERT_EQ(messages.size(), 10);
+    ASSERT_EQ(messages[0].first, false);
+    ASSERT_EQ(messages[0].second, "message1");
+    ASSERT_EQ(messages[1].first, false);
+    ASSERT_EQ(messages[1].second, "message2");
+    ASSERT_EQ(messages[2].first, false);
+    ASSERT_EQ(messages[2].second, "message3");
+    ASSERT_EQ(messages[3].first, true);
+    ASSERT_EQ(messages[3].second, "message4");
+    ASSERT_EQ(messages[4].first, false);
+    ASSERT_EQ(messages[4].second, bigMsg);
+    ASSERT_EQ(messages[5].first, true);
+    ASSERT_EQ(messages[5].second, "foo");
+    ASSERT_EQ(messages[6].first, true);
+    ASSERT_EQ(messages[6].second, "my-header-VALUE");
+    ASSERT_EQ(messages[7].first, true);
+    ASSERT_EQ(messages[7].second, "other-header");
+    ASSERT_EQ(messages[8].first, true);
+    ASSERT_EQ(messages[8].second, "xd");
+    ASSERT_EQ(messages[9].first, true);
+    ASSERT_EQ(messages[9].second, "MyUserAgent");
+}
+
+TEST(WebSocketPool, tlsEcho)
+{
+    mock::BaseApplication app;
+    WebSocketPool pool;
+
+    std::vector<std::pair<bool, QByteArray>> messages;
+    OnceFlag messageFlag;
+    OnceFlag closeFlag;
+
+    auto handle = pool.createSocket(
+        {
+            .url = QUrl("wss://127.0.0.1:9050/echo"),
+            .headers{
+                {"My-Header", "my-header-VALUE"},
+                {"Another-Header", "other-header"},
+                {"Cookie", "xd"},  // "known" header
+            },
+        },
+        std::make_unique<Listener>(messages, messageFlag, closeFlag));
+    handle.sendBinary("message1");
+    handle.sendBinary("message2");
+    handle.sendBinary("message3");
+    handle.sendText("message4");
+
+    ASSERT_TRUE(messageFlag.waitFor(1s));
+    QByteArray bigMsg(1 << 15, 'A');
+    handle.sendBinary(bigMsg);
+    handle.sendText("foo");
+    handle.sendText("/HEADER my-header");
+    handle.sendText("/HEADER another-header");
+    handle.sendText("/HEADER cookie");
+    handle.sendText("/HEADER user-agent");
+    handle.sendText("/CLOSE");
+
+    ASSERT_TRUE(closeFlag.waitFor(1s));
+
+    ASSERT_EQ(messages.size(), 10);
+    ASSERT_EQ(messages[0].first, false);
+    ASSERT_EQ(messages[0].second, "message1");
+    ASSERT_EQ(messages[1].first, false);
+    ASSERT_EQ(messages[1].second, "message2");
+    ASSERT_EQ(messages[2].first, false);
+    ASSERT_EQ(messages[2].second, "message3");
+    ASSERT_EQ(messages[3].first, true);
+    ASSERT_EQ(messages[3].second, "message4");
+    ASSERT_EQ(messages[4].first, false);
+    ASSERT_EQ(messages[4].second, bigMsg);
+    ASSERT_EQ(messages[5].first, true);
+    ASSERT_EQ(messages[5].second, "foo");
+    ASSERT_EQ(messages[6].first, true);
+    ASSERT_EQ(messages[6].second, "my-header-VALUE");
+    ASSERT_EQ(messages[7].first, true);
+    ASSERT_EQ(messages[7].second, "other-header");
+    ASSERT_EQ(messages[8].first, true);
+    ASSERT_EQ(messages[8].second, "xd");
+    ASSERT_EQ(messages[9].first, true);
+    ASSERT_TRUE(messages[9].second.startsWith("Chatterino"));
+}