vacuum-wall/webui/static/hoover/websocket.js

/**
 * Hoover — websocket.js
 *
 * WebSocket connection manager with auto-reconnect. WS messages are routed
 * to model-based refresh and direct onMessage handlers.
 * Page-level subscribe/unsubscribe is replaced by the model layer.
 */

import { refreshByTopic } from './model.js?v=9';
import { clearAuthTokens } from './api.js?v=12';

let _wsConn = null;
let _wsReconnectMs = 0;
let _wsFailCount = 0;

/** Direct onMessage handlers — { topics, handler, unsubscribed }[] */
const _directHandlers = [];

/**
 * Refresh the access token. On failure, clears all tokens to prevent
 * an infinite reconnection loop with a stale token.
 *
 * @returns {Promise<boolean>} true if token was refreshed
 */
async function _tryRefreshToken() {
    const refresh = sessionStorage.getItem('vw:refresh');
    if (!refresh) return false;
    try {
        const res = await fetch('/api/auth/refresh', {
            method: 'POST',
            headers: { 'Content-Type': 'application/json', Accept: 'application/json' },
            body: JSON.stringify({ refresh_token: refresh }),
            credentials: 'same-origin',
        });
        if (res.status !== 200) {
            clearAuthTokens();
            return false;
        }
        const json = await res.json();
        if (!json.ok || !json.data?.tokens) {
            clearAuthTokens();
            return false;
        }
        const tokens = json.data.tokens;
        window.__auth_token__ = tokens.access_token;
        sessionStorage.setItem('vw:refresh', tokens.refresh_token);
        sessionStorage.setItem('vw:access_ttl', String((json.data.access_ttl || 300) * 1000));
        if (tokens.session_id) {
            sessionStorage.setItem('vw:session_id', tokens.session_id);
        }
        if (json.data.user) {
            sessionStorage.setItem('vw:user', JSON.stringify(json.data.user));
            localStorage.setItem('vw:user', JSON.stringify(json.data.user));
        }
        return true;
    } catch {
        clearAuthTokens();
        return false;
    }
}

/**
 * Build the WebSocket URL.  Supports an override via `window.__WS_URL__`
 * (useful for proxy setups).  Falls back to port 9091 when the current
 * origin has no port (nginx fronting the WS on a different port).
 */
function _wsUrl() {
    if (window.__WS_URL__) return window.__WS_URL__;
    const proto = location.protocol === 'https:' ? 'wss:' : 'ws:';
    return proto + '//' + location.host + '/ws';
}

/** Attempt a WebSocket connection.
 *  Passes the JWT in the WebSocket subprotocol header (Sec-WebSocket-Protocol)
 *  instead of a query parameter, keeping it out of logs and browser history.
 */
function _wsConnect() {
    if (_wsConn && _wsConn.readyState <= 1) return;

    const token = window.__auth_token__;
    if (token) {
        _wsConn = new WebSocket(_wsUrl(), ['Bearer ' + token]);
    } else {
        _wsConn = new WebSocket(_wsUrl());
    }

    _wsConn.onopen = () => {
        _wsReconnectMs = 0;
        _wsFailCount = 0;
    };

    _wsConn.onclose = () => {
        if (!window.__auth_token__) return;
        _wsFailCount++;

        if (_wsFailCount >= 3) {
            // Attempt token refresh after repeated failures.  No redirect
            // on failure — the reconnect loop continues.
            (async () => {
                const ok = await _tryRefreshToken();
                if (ok) {
                    _wsFailCount = 0;
                    _wsReconnectMs = 0;
                    _wsConn = null;
                    setTimeout(_wsConnect, 100);
                }
            })();
        }

        _wsReconnectMs = Math.min(_wsReconnectMs * 2 + 1000, 15000);
        setTimeout(_wsConnect, _wsReconnectMs);
    };

    _wsConn.onerror = () => {
        _wsConn.close();
    };

    _wsConn.onmessage = (ev) => {
        try {
            const msg = typeof ev.data === 'string' ? JSON.parse(ev.data) : ev.data;
            handleMessage(msg);
        } catch (_) {}
    };
}

/**
 * Route an incoming WS message to model refresh and direct handlers.
 *
 * Expected message shapes:
 *   { type: 'versions', updated: ['firewall', 'dnsmasq', …] }
 *   { type: 'tick',    subsystems: ['firewall', 'wireguard', …] }
 *   { type: 'notify',  topic: 'firewall' }
 *   { type: 'status',  topic: 'firewall', … }
 */
function handleMessage(msg) {
    const topics = [];

    if (msg.type === 'versions' || msg.type === 'refresh' || msg.type === 'tick') {
        topics.push(...(msg.updated || msg.subsystems || msg.topics || []));
    } else if (msg.type === 'notify') {
        topics.push(msg.topic);
    } else if (msg.type === 'status') {
        topics.push(msg.topic || '*');
    }

    // Refresh models for each topic
    for (const topic of topics) {
        refreshByTopic(topic);
    }

    // Notify direct onMessage handlers
    for (const h of _directHandlers) {
        if (h.unsubscribed) continue;
        if (topics.some(t => h.topics.includes(t) || h.topics.includes('*'))) {
            try { h.handler(msg); } catch (_) {}
        }
    }
}

/**
 * Public subscribe API for direct one-off usage (e.g. from page code).
 * Handler receives the raw parsed message when a matching topic arrives.
 * @param {string|string[]} topics
 * @param {function} handler
 * @returns {function} unsubscribe
 */
export function onMessage(topics, handler) {
    const tArray = Array.isArray(topics) ? topics : [topics];
    const entry = { topics: tArray, handler, unsubscribed: false };
    _directHandlers.push(entry);
    return () => {
        entry.unsubscribed = true;
        const idx = _directHandlers.indexOf(entry);
        if (idx !== -1) _directHandlers.splice(idx, 1);
    };
}

/** Start the WebSocket connection. */
export function connect() {
    _wsConnect();
}