vdm/daemon/src/rpc/ws_server.hpp

#pragma once

// The loopback WebSocket transport (docs/05 §4). Binds 127.0.0.1 only, on the first free
// port in 52000-52016, and writes the chosen port to <runtime>/ws.port. Any local process
// can connect, so a token is mandatory: session.pair mints one (behind a user prompt,
// rate-limited), session.hello must present it, and every other method is refused -32002
// until it does. Privileged methods are refused -32003 by the generated dispatch().

#include <cstdint>
#include <memory>
#include <optional>
#include <string>
#include <system_error>
#include <unordered_map>

#include <nlohmann/json_fwd.hpp>

#include "rpc/event_hub.hpp"
#include "rpc/pairing.hpp"
#include "rpc/runtime_dir.hpp"
#include "rpc/task_action_port.hpp"
#include "rpc/ws_frame.hpp"
#include "velox_proto.hpp"

namespace velox::daemon::store {
class Db;
}

namespace velox::daemon::rpc {

class EventLoop;

class WsServer {
public:
    // `actions` is optional (nullptr in tests that don't need download.probe) — see
    // handle_download_probe's own comment for why this method can't go through the
    // generic dispatch() path like everything else.
    WsServer(EventLoop& loop, velox::proto::Dispatcher& dispatcher, velox::daemon::store::Db& db,
             PairingApprover& approver, EventHub& hub, RuntimeDir runtime,
             TaskActionPort* actions = nullptr);
    ~WsServer();

    WsServer(const WsServer&) = delete;
    WsServer& operator=(const WsServer&) = delete;

    // Pick a port, bind loopback, write ws.port, listen, register with the loop.
    std::error_code start();

    int port() const noexcept { return port_; }
    std::size_t connection_count() const noexcept { return conns_.size(); }

    static constexpr int kPortLo = 52000;
    static constexpr int kPortHi = 52016;

private:
    enum class Phase { Handshake, Open, Closing };

    struct Conn {
        int fd;
        Phase phase = Phase::Handshake;
        std::string in_raw;            // bytes before the upgrade completes
        WsFrameReader frames;
        std::string outbuf;
        std::size_t out_off = 0;
        bool close_after_flush = false;
        std::string origin;
        bool authed = false;
        std::string pairing_id;
        std::string session_id;
        std::optional<EventHub::SubId> sub_id;
    };

    void on_listener_readable();
    void on_conn_event(int fd, unsigned events);
    void progress_handshake(Conn& c);
    void on_ws_bytes(Conn& c, std::string_view bytes);
    void handle_rpc(Conn& c, const std::string& text);
    bool handle_session_ws(Conn& c, const std::string& method, const nlohmann::json& request,
                           nlohmann::json& reply);

    // See UdsServer::handle_download_probe — same reasoning, same pattern, duplicated per
    // transport because each owns its own Conn/send mechanics.
    void handle_download_probe(Conn& c, const nlohmann::json& request);
    void handle_download_refreshUrl(Conn& c, const nlohmann::json& request);

    void send_text(Conn& c, const nlohmann::json& value);
    void send_frame(Conn& c, WsOpcode op, std::string_view payload);
    void begin_close(Conn& c, std::uint16_t code, std::string_view reason);
    void flush(Conn& c);
    void close_conn(int fd);

    EventLoop& loop_;
    velox::proto::Dispatcher& dispatcher_;
    velox::daemon::store::Db& db_;
    PairingApprover& approver_;
    EventHub& hub_;
    TaskActionPort* actions_;
    RuntimeDir runtime_;
    PairingRateLimiter rate_limiter_;

    int listen_fd_ = -1;
    int port_ = 0;
    bool wrote_port_file_ = false;
    std::unordered_map<int, std::unique_ptr<Conn>> conns_;
};

}  // namespace velox::daemon::rpc