vdm/daemon/src/rpc/ws_server.cpp

#include "rpc/ws_server.hpp"

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

#include <cerrno>
#include <cstring>
#include <ctime>
#include <random>
#include <string>

#include <nlohmann/json.hpp>

#include "rpc/event_loop.hpp"
#include "rpc/ws_handshake.hpp"
#include "util/time.hpp"
#include "store/pairings.hpp"
#include "store/sqlite.hpp"
#include "version.hpp"

namespace velox::daemon::rpc {

namespace proto = velox::proto;
using nlohmann::json;

namespace {

std::error_code errc(int e) { return std::error_code(e, std::generic_category()); }

constexpr std::size_t kMaxOutBytes = 16 * 1024 * 1024;
constexpr std::size_t kMaxHandshakeBytes = 16 * 1024;

std::string uuid4() {
    std::random_device rd;
    std::uniform_int_distribution<std::uint32_t> d;
    std::uint32_t a = d(rd), b = d(rd), c = d(rd), e = d(rd);
    b = (b & 0xFFFF0FFFu) | 0x00004000u;
    c = (c & 0x3FFFFFFFu) | 0x80000000u;
    char s[37];
    std::snprintf(s, sizeof(s), "%08x-%04x-%04x-%04x-%04x%08x", a, (b >> 16), (b & 0xFFFF),
                  (c >> 16), (c & 0xFFFF), e);
    return std::string(s);
}

int major_of(const std::string& semver) {
    try {
        return std::stoi(semver.substr(0, semver.find('.')));
    } catch (...) {
        return -1;
    }
}

json rpc_error(const json& id, proto::ErrorCode code, std::string_view msg, json data = nullptr) {
    return proto::make_error(id, code, msg, std::move(data));
}

}  // namespace

WsServer::WsServer(EventLoop& loop, proto::Dispatcher& dispatcher, store::Db& db,
                   PairingApprover& approver, EventHub& hub, RuntimeDir runtime,
                   TaskActionPort* actions)
    : loop_(loop),
      dispatcher_(dispatcher),
      db_(db),
      approver_(approver),
      hub_(hub),
      actions_(actions),
      runtime_(std::move(runtime)) {}

WsServer::~WsServer() {
    for (auto& [fd, c] : conns_) {
        loop_.del_fd(fd);
        ::close(fd);
    }
    if (listen_fd_ >= 0) {
        loop_.del_fd(listen_fd_);
        ::close(listen_fd_);
    }
    if (wrote_port_file_) ::unlink(runtime_.ws_port_path().c_str());
}

std::error_code WsServer::start() {
    int fd = -1;
    for (int p = kPortLo; p <= kPortHi; ++p) {
        fd = ::socket(AF_INET, SOCK_STREAM | SOCK_NONBLOCK | SOCK_CLOEXEC, 0);
        if (fd < 0) return errc(errno);

        sockaddr_in addr{};
        addr.sin_family = AF_INET;
        addr.sin_addr.s_addr = ::htonl(INADDR_LOOPBACK);  // 127.0.0.1 only — never INADDR_ANY
        addr.sin_port = ::htons(static_cast<std::uint16_t>(p));

        if (::bind(fd, reinterpret_cast<sockaddr*>(&addr), sizeof(addr)) == 0) {
            port_ = p;
            break;
        }
        ::close(fd);
        fd = -1;
        if (errno != EADDRINUSE) return errc(errno);
    }
    if (fd < 0) return errc(EADDRINUSE);  // 52000-52016 all taken

    if (::listen(fd, SOMAXCONN) != 0) {
        const int e = errno;
        ::close(fd);
        return errc(e);
    }

    // Publish the port for the extension, which cannot read $XDG_RUNTIME_DIR itself but
    // can be told where to look by a connected GUI. 0600, same as the socket.
    const std::string pf = runtime_.ws_port_path();
    const int pfd = ::open(pf.c_str(), O_WRONLY | O_CREAT | O_TRUNC | O_CLOEXEC, 0600);
    if (pfd < 0) {
        const int e = errno;
        ::close(fd);
        return errc(e);
    }
    const std::string line = std::to_string(port_) + "\n";
    [[maybe_unused]] ssize_t w = ::write(pfd, line.data(), line.size());
    ::close(pfd);
    wrote_port_file_ = true;

    listen_fd_ = fd;
    loop_.add_fd(listen_fd_, kRead, [this](int, unsigned) { on_listener_readable(); });
    return {};
}

void WsServer::on_listener_readable() {
    for (;;) {
        const int cfd = ::accept4(listen_fd_, nullptr, nullptr, SOCK_NONBLOCK | SOCK_CLOEXEC);
        if (cfd < 0) {
            if (errno == EAGAIN || errno == EWOULDBLOCK) break;
            if (errno == EINTR || errno == ECONNABORTED) continue;
            break;
        }
        auto conn = std::make_unique<Conn>();
        conn->fd = cfd;
        conns_.emplace(cfd, std::move(conn));
        loop_.add_fd(cfd, kRead, [this](int fd, unsigned ev) { on_conn_event(fd, ev); });
    }
}

void WsServer::on_conn_event(int fd, unsigned events) {
    const auto it = conns_.find(fd);
    if (it == conns_.end()) return;
    Conn& c = *it->second;

    if (events & kWrite) {
        flush(c);
        if (conns_.find(fd) == conns_.end()) return;
    }
    if (!(events & kRead)) return;

    char buf[64 * 1024];
    for (;;) {
        const ssize_t n = ::read(fd, buf, sizeof(buf));
        if (n > 0) {
            const std::string_view chunk(buf, static_cast<std::size_t>(n));
            if (c.phase == Phase::Handshake) {
                c.in_raw.append(chunk);
                if (c.in_raw.size() > kMaxHandshakeBytes) {
                    close_conn(fd);
                    return;
                }
                progress_handshake(c);
            } else {
                on_ws_bytes(c, chunk);
            }
            if (conns_.find(fd) == conns_.end()) return;
            continue;
        }
        if (n == 0) {
            close_conn(fd);
            return;
        }
        if (errno == EAGAIN || errno == EWOULDBLOCK) break;
        if (errno == EINTR) continue;
        close_conn(fd);
        return;
    }
}

void WsServer::progress_handshake(Conn& c) {
    const HandshakeResult hs = ws_try_handshake(c.in_raw);
    if (!hs.complete) return;

    c.outbuf.append(hs.response);
    if (!hs.ok) {
        c.close_after_flush = true;
        flush(c);
        return;
    }

    c.origin = hs.origin;
    c.phase = Phase::Open;
    std::string leftover = c.in_raw.substr(hs.consumed);
    c.in_raw.clear();
    flush(c);
    if (conns_.count(c.fd) && !leftover.empty()) on_ws_bytes(c, leftover);
}

void WsServer::on_ws_bytes(Conn& c, std::string_view bytes) {
    std::vector<WsMessage> msgs;
    const auto st = c.frames.feed(bytes, msgs);

    for (auto& m : msgs) {
        switch (m.opcode) {
        case WsOpcode::Text:
            handle_rpc(c, m.payload);
            break;
        case WsOpcode::Binary:
            begin_close(c, 1003, "binary frames are not accepted");  // 1003: unacceptable data
            break;
        case WsOpcode::Ping:
            send_frame(c, WsOpcode::Pong, m.payload);
            break;
        case WsOpcode::Pong:
            break;
        case WsOpcode::Close:
            send_frame(c, WsOpcode::Close, m.payload);
            c.close_after_flush = true;
            flush(c);
            return;
        default:
            break;
        }
        if (conns_.find(c.fd) == conns_.end()) return;
    }

    if (st == WsFrameReader::Status::ProtocolError) {
        begin_close(c, 1002, c.frames.error());  // 1002: protocol error
    } else if (st == WsFrameReader::Status::MessageTooBig) {
        begin_close(c, 1009, "message too big");  // 1009: message too big
    }
}

void WsServer::handle_rpc(Conn& c, const std::string& text) {
    json req = json::parse(text, nullptr, false);
    if (req.is_discarded()) {
        send_text(c, rpc_error(nullptr, proto::ErrorCode::ParseError, "invalid JSON"));
        return;
    }
    const json id = req.is_object() && req.contains("id") ? req.at("id") : json(nullptr);
    const std::string method =
        req.is_object() && req.contains("method") && req.at("method").is_string()
            ? req.at("method").get<std::string>()
            : std::string{};

    if (!method.empty()) {
        json reply;
        if (handle_session_ws(c, method, req, reply)) {
            if (!reply.is_null()) send_text(c, reply);
            return;
        }
    }

    // Any non-session method requires an authenticated connection.
    if (!c.authed) {
        send_text(c, rpc_error(id, proto::ErrorCode::NotPaired, "not paired: call session.pair first"));
        return;
    }

    if (method == "download.probe") {
        handle_download_probe(c, req);
        return;
    }
    if (method == "download.refreshUrl") {
        handle_download_refreshUrl(c, req);
        return;
    }

    json reply = proto::dispatch(dispatcher_, proto::Transport::Ws, req);
    if (!reply.is_null()) send_text(c, reply);
}

void WsServer::handle_download_probe(Conn& c, const json& request) {
    const json id = request.contains("id") ? request.at("id") : json(nullptr);
    const json params_json = request.contains("params") ? request.at("params") : json::object();

    auto parsed = proto::parse<proto::DownloadProbeParams>(params_json, "params");
    if (!parsed) {
        send_text(c, rpc_error(id, proto::ErrorCode::InvalidParams, parsed.error().message,
                               json{{"path", parsed.error().path}}));
        return;
    }
    if (!actions_) {
        send_text(c, rpc_error(id, proto::ErrorCode::InternalError,
                               "not implemented in this build: download.probe"));
        return;
    }

    const int fd = c.fd;
    actions_->probe_now(
        *parsed, [this, fd, id](proto::HandlerResult<proto::DownloadProbeResult> r) {
            auto it = conns_.find(fd);
            if (it == conns_.end()) return;  // client gone while the probe was outstanding
            if (r) {
                send_text(*it->second, proto::make_result(id, *r));
            } else {
                send_text(*it->second,
                         rpc_error(id, r.error().code, r.error().message, r.error().data));
            }
        });
}

void WsServer::handle_download_refreshUrl(Conn& c, const json& request) {
    const json id = request.contains("id") ? request.at("id") : json(nullptr);
    const json params_json = request.contains("params") ? request.at("params") : json::object();

    auto parsed = proto::parse<proto::DownloadRefreshUrlParams>(params_json, "params");
    if (!parsed) {
        send_text(c, rpc_error(id, proto::ErrorCode::InvalidParams, parsed.error().message,
                               json{{"path", parsed.error().path}}));
        return;
    }
    if (!actions_) {
        send_text(c, rpc_error(id, proto::ErrorCode::InternalError,
                               "not implemented in this build: download.refreshUrl"));
        return;
    }

    const int fd = c.fd;
    actions_->refresh_url(
        parsed->taskId, parsed->url, parsed->headers, parsed->cookies,
        [this, fd, id](proto::HandlerResult<proto::DownloadRefreshUrlResult> r) {
            auto it = conns_.find(fd);
            if (it == conns_.end()) return;  // client gone while the probe was outstanding
            if (r) {
                send_text(*it->second, proto::make_result(id, *r));
            } else {
                send_text(*it->second,
                         rpc_error(id, r.error().code, r.error().message, r.error().data));
            }
        });
}

bool WsServer::handle_session_ws(Conn& c, const std::string& method, const json& request,
                                 json& reply) {
    const json id = request.contains("id") ? request.at("id") : json(nullptr);
    const json params = request.contains("params") ? request.at("params") : json::object();

    if (method == "session.pair") {
        const auto dec = rate_limiter_.check(c.origin);
        if (!dec.allowed) {
            reply = rpc_error(id, proto::ErrorCode::RateLimited,
                              "too many pairing attempts; try again later",
                              json{{"retryAfterSec", dec.retry_after_sec}});
            return true;
        }
        auto p = proto::parse<proto::SessionPairParams>(params, "params");
        if (!p) {
            reply = rpc_error(id, proto::ErrorCode::InvalidParams, p.error().message,
                              json{{"path", p.error().path}});
            return true;
        }

        PairingRequest pr{c.origin, p->clientName, make_pairing_code()};
        if (!approver_.approve(pr)) {
            rate_limiter_.record_failure(c.origin);
            reply = rpc_error(id, proto::ErrorCode::NotPaired, "pairing was not approved");
            return true;
        }

        store::Pairings pairings(db_);
        auto created = pairings.create(c.origin, p->clientName, now_iso());
        if (!created) {
            reply = rpc_error(id, proto::ErrorCode::InternalError,
                              "could not store the pairing: " + created.error().message);
            return true;
        }
        rate_limiter_.record_success(c.origin);

        proto::SessionPairResult r;
        r.token = created->token;
        reply = proto::make_result(id, r);
        return true;
    }

    if (method == "session.hello") {
        auto p = proto::parse<proto::SessionHelloParams>(params, "params");
        if (!p) {
            reply = rpc_error(id, proto::ErrorCode::InvalidParams, p.error().message,
                              json{{"path", p.error().path}});
            return true;
        }
        const int want = major_of(std::string(proto::kProtocolVersion));
        const int got = major_of(p->protocolVersion);
        if (got != want) {
            reply = rpc_error(id, proto::ErrorCode::VersionMismatch,
                              "protocol major version mismatch",
                              json{{"expected", std::string(proto::kProtocolVersion)},
                                   {"actual", p->protocolVersion}});
            c.close_after_flush = true;
            return true;
        }

        store::Pairings pairings(db_);
        auto found = p->token ? pairings.find_active_by_token(*p->token)
                              : store::DbResult<std::optional<store::Pairing>>(std::nullopt);
        if (!found) {
            reply = rpc_error(id, proto::ErrorCode::InternalError, found.error().message);
            return true;
        }
        if (!found->has_value()) {
            // Absent, malformed or wrong — all reported the same, and all count toward the
            // pairing rate limit (fixture session.hello.not-paired).
            rate_limiter_.record_failure(c.origin);
            reply = rpc_error(id, proto::ErrorCode::NotPaired, "not paired: call session.pair first");
            return true;
        }

        c.authed = true;
        c.pairing_id = (*found)->pairing_id;
        if (c.session_id.empty()) c.session_id = uuid4();
        (void)pairings.touch(c.pairing_id, now_iso());

        proto::SessionHelloResult r;
        r.daemonVersion = std::string(velox::daemon::kDaemonVersion);
        r.protocolVersion = std::string(proto::kProtocolVersion);
        r.sessionId = c.session_id;
        r.transport = proto::SessionHelloResultTransport::Ws;
        reply = proto::make_result(id, r);
        return true;
    }

    if (method == "session.subscribe") {
        if (!c.authed) {
            reply = rpc_error(id, proto::ErrorCode::NotPaired, "not paired: call session.pair first");
            return true;
        }
        auto p = proto::parse<proto::SessionSubscribeParams>(params, "params");
        if (!p) {
            reply = rpc_error(id, proto::ErrorCode::InvalidParams, p.error().message,
                              json{{"path", p.error().path}});
            return true;
        }
        if (!c.sub_id) {
            const int fd = c.fd;
            c.sub_id = hub_.subscribe([this, fd](const json& n) {
                if (const auto it = conns_.find(fd); it != conns_.end()) send_text(*it->second, n);
            });
        }
        std::vector<proto::Event> events;
        proto::SessionSubscribeResult r;
        r.ok = true;
        for (const auto& ev : p->events) {
            const auto name = proto::to_string(ev);
            r.events.emplace_back(name);
            if (auto e = proto::event_from_string(name)) events.push_back(*e);
        }
        hub_.set_filter(*c.sub_id, std::move(events), p->taskIds);
        reply = proto::make_result(id, r);
        return true;
    }

    return false;
}

void WsServer::send_text(Conn& c, const json& value) {
    send_frame(c, WsOpcode::Text, value.dump());
}

void WsServer::send_frame(Conn& c, WsOpcode op, std::string_view payload) {
    c.outbuf += ws_encode(op, payload);
    if (c.outbuf.size() - c.out_off > kMaxOutBytes) {
        close_conn(c.fd);
        return;
    }
    flush(c);
}

void WsServer::begin_close(Conn& c, std::uint16_t code, std::string_view reason) {
    if (c.phase == Phase::Closing) return;
    c.phase = Phase::Closing;
    send_frame(c, WsOpcode::Close, ws_close_payload(code, reason));
    if (conns_.count(c.fd)) {
        c.close_after_flush = true;
        flush(c);
    }
}

void WsServer::flush(Conn& c) {
    while (c.out_off < c.outbuf.size()) {
        const ssize_t n = ::write(c.fd, c.outbuf.data() + c.out_off, c.outbuf.size() - c.out_off);
        if (n > 0) {
            c.out_off += static_cast<std::size_t>(n);
            continue;
        }
        if (n < 0 && (errno == EAGAIN || errno == EWOULDBLOCK)) {
            loop_.mod_fd(c.fd, kRead | kWrite);
            return;
        }
        if (n < 0 && errno == EINTR) continue;
        close_conn(c.fd);
        return;
    }
    c.outbuf.clear();
    c.out_off = 0;
    if (c.close_after_flush) {
        close_conn(c.fd);
        return;
    }
    loop_.mod_fd(c.fd, kRead);
}

void WsServer::close_conn(int fd) {
    if (const auto it = conns_.find(fd); it != conns_.end()) {
        if (it->second->sub_id) hub_.unsubscribe(*it->second->sub_id);
        loop_.del_fd(fd);
        ::close(fd);
        conns_.erase(it);
    }
}

}  // namespace velox::daemon::rpc