vdm/daemon/tests/ws_server_test.cpp

// Integration: a real WsServer on a loopback port, a hand-rolled WebSocket client.
// Covers the handshake, the pairing flow, the token gate (-32002), and the
// privileged-over-WS refusal (-32003).

#include <arpa/inet.h>
#include <netinet/in.h>
#include <sys/socket.h>
#include <unistd.h>

#include <cstdlib>
#include <string>
#include <thread>
#include <vector>

#include <nlohmann/json.hpp>

#include "check.hpp"
#include "rpc/dispatcher.hpp"
#include "rpc/event_hub.hpp"
#include "rpc/event_loop.hpp"
#include "rpc/pairing.hpp"
#include "rpc/runtime_dir.hpp"
#include "rpc/ws_frame.hpp"
#include "rpc/ws_server.hpp"
#include "store/migrations.hpp"
#include "store/sqlite.hpp"

using nlohmann::json;
namespace rpc = velox::daemon::rpc;
namespace store = velox::daemon::store;

namespace {

int dial(int port) {
    const int fd = ::socket(AF_INET, SOCK_STREAM | SOCK_CLOEXEC, 0);
    sockaddr_in a{};
    a.sin_family = AF_INET;
    a.sin_addr.s_addr = ::htonl(INADDR_LOOPBACK);
    a.sin_port = ::htons(static_cast<std::uint16_t>(port));
    if (::connect(fd, reinterpret_cast<sockaddr*>(&a), sizeof(a)) != 0) {
        ::close(fd);
        return -1;
    }
    return fd;
}

void write_all(int fd, std::string_view s) {
    std::size_t off = 0;
    while (off < s.size()) {
        const ssize_t n = ::write(fd, s.data() + off, s.size() - off);
        if (n <= 0) return;
        off += static_cast<std::size_t>(n);
    }
}

std::string read_some(int fd) {
    char buf[8192];
    const ssize_t n = ::read(fd, buf, sizeof(buf));
    return n > 0 ? std::string(buf, static_cast<std::size_t>(n)) : std::string{};
}

// A masked client text frame.
std::string client_text(std::string_view payload) {
    std::string f;
    f.push_back(static_cast<char>(0x81));  // FIN + text
    const std::size_t n = payload.size();
    if (n < 126) {
        f.push_back(static_cast<char>(0x80 | n));
    } else {
        f.push_back(static_cast<char>(0x80 | 126));
        f.push_back(static_cast<char>((n >> 8) & 0xFF));
        f.push_back(static_cast<char>(n & 0xFF));
    }
    const char k[4] = {0x0A, 0x0B, 0x0C, 0x0D};
    f.append(k, 4);
    for (std::size_t i = 0; i < n; ++i) f.push_back(static_cast<char>(payload[i] ^ k[i & 3]));
    return f;
}

// Decode one unmasked server frame from `buf`, consuming it. Returns payload; sets `op`.
std::string server_frame(std::string& buf, rpc::WsOpcode& op) {
    if (buf.size() < 2) return {};
    op = static_cast<rpc::WsOpcode>(buf[0] & 0x0F);
    std::size_t len = static_cast<std::uint8_t>(buf[1]) & 0x7F;
    std::size_t header = 2;
    if (len == 126) {
        len = (static_cast<std::size_t>(static_cast<std::uint8_t>(buf[2])) << 8) |
              static_cast<std::uint8_t>(buf[3]);
        header = 4;
    }
    if (buf.size() < header + len) return {};
    std::string payload = buf.substr(header, len);
    buf.erase(0, header + len);
    return payload;
}

// Send a request frame, wait for one text reply, return its parsed JSON.
json rpc_call(int fd, const json& req) {
    write_all(fd, client_text(req.dump()));
    std::string buf;
    for (;;) {
        buf += read_some(fd);
        rpc::WsOpcode op{};
        std::string save = buf;
        const std::string payload = server_frame(buf, op);
        if (payload.empty() && buf == save) continue;  // need more bytes
        if (op == rpc::WsOpcode::Text) return json::parse(payload, nullptr, false);
    }
}

}  // namespace

void run() {
    ::unsetenv("VELOX_PAIR_AUTO");

    auto db = store::Db::open(":memory:");
    CHECK(db.has_value());
    if (!db) return;
    CHECK(store::migrate_to_head(*db).has_value());

    char tmpl[] = "/tmp/velox-ws-test-XXXXXX";
    const char* dir = ::mkdtemp(tmpl);
    CHECK(dir != nullptr);
    rpc::RuntimeDir rt{dir ? dir : "/tmp"};

    rpc::EventLoop loop;
    rpc::EventHub hub;
    rpc::VeloxDispatcher dispatcher(*db, hub);
    rpc::EnvAutoApprover approver;
    rpc::WsServer server(loop, dispatcher, *db, approver, hub, rt);
    const auto ec = server.start();
    CHECK(!ec);
    if (ec) return;
    CHECK(server.port() >= rpc::WsServer::kPortLo);
    CHECK(server.port() <= rpc::WsServer::kPortHi);

    std::thread th([&loop] { loop.run(); });

    const std::string origin = "moz-extension://11111111-2222-3333-4444-555555555555";

    // --- handshake ---------------------------------------------------------------
    const int fd = dial(server.port());
    CHECK(fd >= 0);
    write_all(fd,
              "GET / HTTP/1.1\r\nHost: 127.0.0.1\r\nUpgrade: websocket\r\nConnection: Upgrade\r\n"
              "Sec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==\r\nSec-WebSocket-Version: 13\r\n"
              "Origin: " + origin + "\r\n\r\n");
    std::string hs;
    while (hs.find("\r\n\r\n") == std::string::npos) hs += read_some(fd);
    CHECK(hs.find("101 Switching Protocols") != std::string::npos);

    // --- session.hello with no token -> -32002 --------------------------------
    {
        const json r = rpc_call(fd, {{"jsonrpc", "2.0"}, {"id", 1}, {"method", "session.hello"},
                                     {"params", {{"clientType", "extension"},
                                                 {"clientName", "Velox for Firefox"},
                                                 {"protocolVersion",
                                                  std::string(velox::proto::kProtocolVersion)}}}});
        CHECK(r.contains("error"));
        CHECK_EQ(r["error"]["code"].get<int>(), -32002);
    }

    // --- session.pair with approval off -> not approved ----------------------
    {
        const json r = rpc_call(fd, {{"jsonrpc", "2.0"}, {"id", 2}, {"method", "session.pair"},
                                     {"params", {{"clientName", "Velox for Firefox"},
                                                 {"extensionId", "11111111-2222-3333-4444-555555555555"}}}});
        CHECK(r.contains("error"));
    }

    // --- approval on -> a token, then hello with it succeeds ----------------
    ::setenv("VELOX_PAIR_AUTO", "1", 1);
    std::string token;
    {
        const json r = rpc_call(fd, {{"jsonrpc", "2.0"}, {"id", 3}, {"method", "session.pair"},
                                     {"params", {{"clientName", "Velox for Firefox"},
                                                 {"extensionId", "11111111-2222-3333-4444-555555555555"}}}});
        CHECK(r.contains("result"));
        if (r.contains("result")) {
            token = r["result"]["token"].get<std::string>();
            CHECK(token.size() >= 40);
        }
    }
    {
        const json r = rpc_call(fd, {{"jsonrpc", "2.0"}, {"id", 4}, {"method", "session.hello"},
                                     {"params", {{"clientType", "extension"},
                                                 {"clientName", "Velox for Firefox"},
                                                 {"protocolVersion",
                                                  std::string(velox::proto::kProtocolVersion)},
                                                 {"token", token}}}});
        CHECK(r.contains("result"));
        if (r.contains("result"))
            CHECK_EQ(r["result"]["transport"].get<std::string>(), std::string("ws"));
    }

    // --- privileged method over WS -> -32003 -------------------------------
    {
        const json r = rpc_call(fd, {{"jsonrpc", "2.0"}, {"id", 5}, {"method", "settings.get"},
                                     {"params", {{"keys", nullptr}}}});
        CHECK(r.contains("error"));
        CHECK_EQ(r["error"]["code"].get<int>(), -32003);
    }

    // --- a non-privileged method while authed -> a real result -----------
    {
        const json r = rpc_call(fd, {{"jsonrpc", "2.0"}, {"id", 6}, {"method", "download.list"},
                                     {"params", json::object()}});
        CHECK(r.contains("result"));
        if (r.contains("result")) CHECK_EQ(r["result"]["total"].get<int>(), 0);
    }

    ::close(fd);
    loop.stop();
    th.join();
    ::unlink(rt.ws_port_path().c_str());
}

TEST_MAIN()